feat: cortex extensions

.gitignore

@@ -569,3 +569,5 @@ uploads/**
 
 # cortex-js
 cortex-js/cortex.db
+dist
+*.lock

core/package.json

@@ -0,0 +1,25 @@
+{
+  "name": "@janhq/core",
+  "version": "1.0.0",
+  "description": "Cortex Core SDK",
+  "type": "commonjs",
+  "scripts": {
+    "start": "node --experimental-specifier-resolution=node dist/index.js",
+    "dev": "tsc && concurrently \"tsc -w\" \"nodemon --experimental-specifier-resolution=node dist/index.js\"",
+    "build": "tsc"
+  },
+  "exports": {
+    ".": "./dist/index.js"
+  },
+  "types": "./dist/index.d.ts",
+  "dependencies": {
+    "node-fetch": "2"
+  },
+  "devDependencies": {
+    "@types/node": "20.12.8",
+    "concurrently": "5.2.0",
+    "nodemon": "2.0.4",
+    "typescript": "^5.4.5"
+  },
+  "license": "AGPL-3.0"
+}

core/src/engine.ts

@@ -0,0 +1,7 @@
+import { Extension } from "./extension";
+
+export abstract class EngineExtension extends Extension {
+  abstract provider: string;
+  abstract inference(completion: any, req: any, res: any): void;
+  abstract loadModel(loadModel: any): Promise<void>;
+}

core/src/extension.ts

@@ -0,0 +1,35 @@
+/**
+ * Represents a base extension.
+ * This class should be extended by any class that represents an extension.
+ */
+export abstract class Extension {
+  /** @type {string} Name of the extension. */
+  name?: string;
+
+  /** @type {string} Product Name of the extension. */
+  productName?: string;
+
+  /** @type {string} The URL of the extension to load. */
+  url?: string;
+
+  /** @type {boolean} Whether the extension is activated or not. */
+  active?: boolean;
+
+  /** @type {string} Extension's description. */
+  description?: string;
+
+  /** @type {string} Extension's version. */
+  version?: string;
+
+  /**
+   * Called when the extension is loaded.
+   * Any initialization logic for the extension should be put here.
+   */
+  onLoad(): void {}
+
+  /**
+   * Called when the extension is unloaded.
+   * Any cleanup logic for the extension should be put here.
+   */
+  onUnload(): void {}
+}

core/src/index.ts

@@ -0,0 +1,4 @@
+export * from "./types";
+export * from "./engine";
+export * from "./extension";
+export * from "./oai";

core/src/oai.ts

@@ -0,0 +1,36 @@
+import { EngineExtension } from "./engine";
+
+export abstract class OAIEngineExtension extends EngineExtension {
+  abstract apiUrl: string;
+
+  async inference(
+    createChatDto: any,
+    headers: Record<string, string>,
+    res: any
+  ) {
+    // eslint-disable-next-line @typescript-eslint/no-var-requires
+    const fetch = require("node-fetch");
+    const response = await fetch(this.apiUrl, {
+      method: "POST",
+      headers: {
+        "Content-Type": headers["content-type"] ?? "application/json",
+        Authorization: headers["authorization"],
+      },
+      body: JSON.stringify(createChatDto),
+    });
+
+    res.writeHead(200, {
+      "Content-Type":
+        createChatDto.stream === true
+          ? "text/event-stream"
+          : "application/json",
+      "Cache-Control": "no-cache",
+      Connection: "keep-alive",
+      "Access-Control-Allow-Origin": "*",
+    });
+
+    response.body.pipe(res);
+  }
+
+  async loadModel(loadModel: any): Promise<void> {}
+}

core/src/types.ts

@@ -0,0 +1,6 @@
+export type PromptTemplate = {
+  system_prompt?: string;
+  ai_prompt?: string;
+  user_prompt?: string;
+  error?: string;
+};

core/tsconfig.json

@@ -0,0 +1,14 @@
+{
+  "compilerOptions": {
+    "module": "CommonJS",
+    "target": "ES2017",
+    "strict": true,
+    "sourceMap": true,
+    "outDir": "dist",
+    "moduleResolution": "node",
+    "declaration": true,
+    "esModuleInterop": true,
+    "forceConsistentCasingInFileNames": true
+  },
+  "include": ["src/**/*"]
+}

cortex-js/.gitignore

@@ -57,3 +57,4 @@ report.[0-9]*.[0-9]*.[0-9]*.[0-9]*.json
 
 # cortex-js
 /cortex-js/cortex.db
+dist

cortex-js/package.json

@@ -21,6 +21,7 @@
     "typeorm": "typeorm-ts-node-esm"
   },
   "dependencies": {
+    "@janhq/core": "link:../core",
     "@nestjs/common": "^10.0.0",
     "@nestjs/core": "^10.0.0",
     "@nestjs/devtools-integration": "^0.1.6",

cortex-js/src/app.module.ts

@@ -10,6 +10,7 @@ import { DatabaseModule } from './infrastructure/database/database.module';
 import { ChatModule } from './usecases/chat/chat.module';
 import { AssistantsModule } from './usecases/assistants/assistants.module';
 import { InferenceSettingsModule } from './usecases/inference-settings/inference-settings.module';
+import { ExtensionModule } from './infrastructure/repositories/extensions/extension.module';
 import { CortexModule } from './usecases/cortex/cortex.module';
 
 @Module({
@@ -25,6 +26,7 @@ import { CortexModule } from './usecases/cortex/cortex.module';
     AssistantsModule,
     InferenceSettingsModule,
     CortexModule,
+    ExtensionModule,
   ],
   controllers: [AppController, ThreadsController, ModelsController],
 })

cortex-js/src/domain/models/message.interface.ts

@@ -79,10 +79,3 @@ export interface Message {
  * The additional metadata of this message.
  */
 export interface MessageMetadata {}
-
-export type PromptTemplate = {
-  system_prompt?: string;
-  ai_prompt?: string;
-  user_prompt?: string;
-  error?: string;
-};

cortex-js/src/domain/models/model.interface.ts

@@ -6,35 +6,9 @@ export interface ModelInfo {
   id: string;
   settings: ModelSettingParams;
   parameters: ModelRuntimeParams;
-  engine?: InferenceEngine;
+  engine?: string;
 }
 
-/**
- * Represents the remote inference engine.
- * @stored
- */
-export const RemoteInferenceEngines = ['openai', 'groq'];
-
-/**
- * Represents the local inference engine.
- */
-export const LocalInferenceEngines = [
-  'nitro',
-  'triton_trtllm',
-  'nitro_tensorrt_llm',
-];
-
-/**
- * Represents all supported inference engine.
- * @stored
- */
-const AllInferenceEngines = [
-  ...RemoteInferenceEngines,
-  ...LocalInferenceEngines,
-] as const;
-export type InferenceEngineTuple = typeof AllInferenceEngines;
-export type InferenceEngine = InferenceEngineTuple[number];
-
 export interface ModelArtifact {
   url: string;
 }

cortex-js/src/domain/repositories/extension.interface.ts

@@ -0,0 +1,4 @@
+import { Extension } from '@janhq/core';
+import { Repository } from './repository.interface';
+
+export abstract class ExtensionRepository extends Repository<Extension> {}

cortex-js/src/domain/repositories/repository.interface.ts

@@ -1,11 +1,11 @@
-export interface Repository<T> {
-  create(object: Partial<T>): Promise<T>;
+export abstract class Repository<T> {
+  abstract create(object: Partial<T>): Promise<T>;
 
-  findAll(): Promise<T[]>;
+  abstract findAll(): Promise<T[]>;
 
-  findOne(id: string): Promise<T | null>;
+  abstract findOne(id: string): Promise<T | null>;
 
-  update(id: string, object: Partial<T>): Promise<void>;
+  abstract update(id: string, object: Partial<T>): Promise<void>;
 
-  remove(id: string): Promise<void>;
+  abstract remove(id: string): Promise<void>;
 }

cortex-js/src/infrastructure/controllers/chat.controller.ts

@@ -1,4 +1,4 @@
-import { Body, Controller, Post, Res } from '@nestjs/common';
+import { Body, Controller, Post, Headers, Res } from '@nestjs/common';
 import { CreateChatCompletionDto } from '../dtos/chat/create-chat-completion.dto';
 import { ChatUsecases } from '../../usecases/chat/chat.usecases';
 import { Response } from 'express';
@@ -11,9 +11,10 @@ export class ChatController {
 
   @Post('completions')
   async create(
+    @Headers() headers: Record<string, string>,
     @Body() createChatDto: CreateChatCompletionDto,
     @Res() res: Response,
   ) {
-    return this.chatService.createChatCompletions(createChatDto, res);
+    this.chatService.createChatCompletions(createChatDto, headers, res);
   }
 }

cortex-js/src/infrastructure/dtos/models/create-model.dto.ts

@@ -1,10 +1,6 @@
 import { Type } from 'class-transformer';
 import { IsArray, IsEnum, IsString, ValidateNested } from 'class-validator';
-import {
-  InferenceEngine,
-  Model,
-  ModelFormat,
-} from 'src/domain/models/model.interface';
+import { Model, ModelFormat } from 'src/domain/models/model.interface';
 import { ModelArtifactDto } from './model-artifact.dto';
 import { ModelSettingParamsDto } from './model-setting-params.dto';
 import { ModelRuntimeParamsDto } from './model-runtime-params.dto';
@@ -42,5 +38,5 @@ export class CreateModelDto implements Partial<Model> {
   metadata: ModelMetadataDto;
 
   @IsString()
-  engine: InferenceEngine;
+  engine: string;
 }

cortex-js/src/infrastructure/dtos/threads/create-thread-model-info.dto.ts

@@ -1,5 +1,5 @@
 import { IsOptional, IsString, ValidateNested } from 'class-validator';
-import { InferenceEngine, ModelInfo } from 'src/domain/models/model.interface';
+import { ModelInfo } from 'src/domain/models/model.interface';
 import { ModelRuntimeParamsDto } from 'src/infrastructure/dtos/models/model-runtime-params.dto';
 import { ModelSettingParamsDto } from 'src/infrastructure/dtos/models/model-setting-params.dto';
 
@@ -15,5 +15,5 @@ export class CreateThreadModelInfoDto implements ModelInfo {
 
   @IsOptional()
   @IsString()
-  engine?: InferenceEngine;
+  engine?: string;
 }

cortex-js/src/infrastructure/repositories/extensions/extension.module.ts

@@ -0,0 +1,14 @@
+import { Module } from '@nestjs/common';
+import { ExtensionRepositoryImpl } from './extension.repository';
+import { ExtensionRepository } from 'src/domain/repositories/extension.interface';
+
+@Module({
+  providers: [
+    {
+      provide: ExtensionRepository,
+      useClass: ExtensionRepositoryImpl,
+    },
+  ],
+  exports: [ExtensionRepository],
+})
+export class ExtensionModule {}

cortex-js/src/infrastructure/repositories/extensions/extension.repository.ts

@@ -0,0 +1,46 @@
+import { Injectable } from '@nestjs/common';
+import { LazyModuleLoader } from '@nestjs/core';
+import { ExtensionRepository } from 'src/domain/repositories/extension.interface';
+import { Extension } from '@janhq/core';
+import { readdir, lstat } from 'fs/promises';
+
+@Injectable()
+export class ExtensionRepositoryImpl implements ExtensionRepository {
+  extensions = new Map<string, Extension>();
+
+  constructor(private lazyModuleLoader: LazyModuleLoader) {
+    this.loadExtensions();
+  }
+  create(object: Extension): Promise<Extension> {
+    this.extensions.set(object.name ?? '', object);
+    return Promise.resolve(object);
+  }
+  findAll(): Promise<Extension[]> {
+    return Promise.resolve(Array.from(this.extensions.values()));
+  }
+  findOne(id: string): Promise<Extension | null> {
+    return Promise.resolve(this.extensions.get(id) ?? null);
+  }
+  update(): Promise<void> {
+    throw new Error('Method not implemented.');
+  }
+  remove(id: string): Promise<void> {
+    this.extensions.delete(id);
+    return Promise.resolve();
+  }
+
+  loadExtensions(): void {
+    const extensionsPath = process.env.EXTENSIONS_PATH ?? 'extensions';
+    readdir(extensionsPath).then((files) => {
+      files.forEach(async (extension) => {
+        if ((await lstat(`${extensionsPath}/${extension}`)).isDirectory())
+          return;
+
+        import(`${extensionsPath}/${extension}`).then((extensionClass) => {
+          const newExtension = new extensionClass.default();
+          this.extensions.set(extension, newExtension);
+        });
+      });
+    });
+  }
+}

cortex-js/src/usecases/chat/chat.module.ts

@@ -1,11 +1,11 @@
 import { Module } from '@nestjs/common';
 import { ChatController } from '../../infrastructure/controllers/chat.controller';
 import { ChatUsecases } from './chat.usecases';
-import { InferenceSettingsModule } from '../inference-settings/inference-settings.module';
-import { ModelsModule } from '../models/models.module';
+import { DatabaseModule } from 'src/infrastructure/database/database.module';
+import { ExtensionModule } from '@/infrastructure/repositories/extensions/extension.module';
 
 @Module({
-  imports: [ModelsModule, InferenceSettingsModule],
+  imports: [DatabaseModule, ExtensionModule],
   controllers: [ChatController],
   providers: [ChatUsecases],
 })