feat: support martian nvidia engine (#956)

marknguyen1302 · web-flow · commit 9f48d1df152e · 2024-08-01T11:36:27.000+07:00
diff --git a/cortex-js/src/domain/abstracts/oai.abstract.ts b/cortex-js/src/domain/abstracts/oai.abstract.ts
@@ -2,7 +2,7 @@ import { HttpService } from '@nestjs/axios';
 import { EngineExtension } from './engine.abstract';
 import stream, { Transform } from 'stream';
 import { firstValueFrom } from 'rxjs';
-import _ from 'lodash';
+import { omit } from 'lodash';
 
 export abstract class OAIEngineExtension extends EngineExtension {
   abstract apiUrl: string;
@@ -22,7 +22,7 @@ export abstract class OAIEngineExtension extends EngineExtension {
       ? this.transformPayload(createChatDto)
       : createChatDto;
     const { stream: isStream } = payload;
-    const additionalHeaders = _.omit(headers, [
+    const additionalHeaders = omit(headers, [
       'content-type',
       'authorization',
       'content-length',
@@ -54,7 +54,7 @@ export abstract class OAIEngineExtension extends EngineExtension {
           const transformedLines = [];
           for (const line of lines) {
             if (line.trim().length > 0) {
-              const transformedLine = transformResponse(line);
+              const transformedLine = transformResponse(line, true);
               if (transformedLine) {
                 transformedLines.push(`data: ${transformedLine}\n\n`);
               }
@@ -65,6 +65,6 @@ export abstract class OAIEngineExtension extends EngineExtension {
       });
       return response.data.pipe(lineStream);
     }
-    return this.transformResponse(response.data);
+    return this.transformResponse(response.data, false);
   }
 }
diff --git a/cortex-js/src/extensions/anthropic.engine.ts b/cortex-js/src/extensions/anthropic.engine.ts
@@ -3,7 +3,7 @@ import { HttpService } from '@nestjs/axios';
 import { OAIEngineExtension } from '../domain/abstracts/oai.abstract';
 import { ConfigsUsecases } from '@/usecases/configs/configs.usecase';
 import { EventEmitter2 } from '@nestjs/event-emitter';
-import _ from 'lodash';
+import { pick } from 'lodash';
 import { EngineStatus } from '@/domain/abstracts/engine.abstract';
 
 /**
@@ -59,7 +59,13 @@ export default class AnthropicEngineExtension extends OAIEngineExtension {
   }
 
   transformPayload = (data: any): any => {
-    return _.pick(data, ['messages', 'model', 'stream', 'max_tokens']);
+    const system = data.messages.find((m: any) => m.role === 'system');
+    const messages = data.messages.filter((m: any) => m.role !== 'system');
+    return {
+      system: system.content,
+      messages,
+      ...pick(data, ['model', 'stream', 'max_tokens']),
+    };
   };
 
   transformResponse = (data: any) => {
diff --git a/cortex-js/src/extensions/cohere.engine.ts b/cortex-js/src/extensions/cohere.engine.ts
@@ -1,9 +1,7 @@
-import stream from 'stream';
 import { HttpService } from '@nestjs/axios';
 import { OAIEngineExtension } from '../domain/abstracts/oai.abstract';
 import { ConfigsUsecases } from '@/usecases/configs/configs.usecase';
 import { EventEmitter2 } from '@nestjs/event-emitter';
-import _ from 'lodash';
 import { EngineStatus } from '@/domain/abstracts/engine.abstract';
 import { ChatCompletionMessage } from '@/infrastructure/dtos/chat/chat-completion-message.dto';
 
@@ -98,17 +96,27 @@ export default class CoHereEngineExtension extends OAIEngineExtension {
     return convertedData;
   };
 
-  transformResponse = (data: any) => {
+  transformResponse = (data: any, stream: boolean) => {
     const text =
       typeof data === 'object' ? data.text : (JSON.parse(data).text ?? '');
-    return JSON.stringify({
-      choices: [
-        {
-          delta: {
-            content: text,
-          },
-        },
-      ],
-    });
+    return stream
+      ? JSON.stringify({
+          choices: [
+            {
+              delta: {
+                content: text,
+              },
+            },
+          ],
+        })
+      : {
+          choices: [
+            {
+              message: {
+                content: text,
+              },
+            },
+          ],
+        };
   };
 }
diff --git a/cortex-js/src/extensions/extensions.module.ts b/cortex-js/src/extensions/extensions.module.ts
@@ -9,6 +9,8 @@ import { EventEmitter2 } from '@nestjs/event-emitter';
 import AnthropicEngineExtension from './anthropic.engine';
 import OpenRouterEngineExtension from './openrouter.engine';
 import CoHereEngineExtension from './cohere.engine';
+import MartianEngineExtension from './martian.engine';
+import NvidiaEngineExtension from './nvidia.engine';
 
 const provider = {
   provide: 'EXTENSIONS_PROVIDER',
@@ -24,6 +26,8 @@ const provider = {
     new AnthropicEngineExtension(httpService, configUsecases, eventEmitter),
     new OpenRouterEngineExtension(httpService, configUsecases, eventEmitter),
     new CoHereEngineExtension(httpService, configUsecases, eventEmitter),
+    new MartianEngineExtension(httpService, configUsecases, eventEmitter),
+    new NvidiaEngineExtension(httpService, configUsecases, eventEmitter),
   ],
 };
 
diff --git a/cortex-js/src/extensions/martian.engine.ts b/cortex-js/src/extensions/martian.engine.ts
@@ -0,0 +1,49 @@
+import { HttpService } from '@nestjs/axios';
+import { OAIEngineExtension } from '../domain/abstracts/oai.abstract';
+import { ConfigsUsecases } from '@/usecases/configs/configs.usecase';
+import { EventEmitter2 } from '@nestjs/event-emitter';
+import { EngineStatus } from '@/domain/abstracts/engine.abstract';
+
+/**
+ * A class that implements the InferenceExtension interface from the @janhq/core package.
+ * The class provides methods for initializing and stopping a model, and for making inference requests.
+ * It also subscribes to events emitted by the @janhq/core package and handles new message requests.
+ */
+export default class MartianEngineExtension extends OAIEngineExtension {
+  apiUrl = 'https://withmartian.com/api/openai/v1/chat/completions';
+  name = 'martian';
+  productName = 'Martian Inference Engine';
+  description = 'This extension enables Martian chat completion API calls';
+  version = '0.0.1';
+  apiKey?: string;
+
+  constructor(
+    protected readonly httpService: HttpService,
+    protected readonly configsUsecases: ConfigsUsecases,
+    protected readonly eventEmmitter: EventEmitter2,
+  ) {
+    super(httpService);
+
+    eventEmmitter.on('config.updated', async (data) => {
+      if (data.engine === this.name) {
+        this.apiKey = data.value;
+        this.status =
+          (this.apiKey?.length ?? 0) > 0
+            ? EngineStatus.READY
+            : EngineStatus.MISSING_CONFIGURATION;
+      }
+    });
+  }
+
+  async onLoad() {
+    const configs = (await this.configsUsecases.getGroupConfigs(
+      this.name,
+    )) as unknown as { apiKey: string };
+
+    this.apiKey = configs?.apiKey;
+    this.status =
+      (this.apiKey?.length ?? 0) > 0
+        ? EngineStatus.READY
+        : EngineStatus.MISSING_CONFIGURATION;
+  }
+}
diff --git a/cortex-js/src/extensions/nvidia.engine.ts b/cortex-js/src/extensions/nvidia.engine.ts
@@ -0,0 +1,49 @@
+import { HttpService } from '@nestjs/axios';
+import { OAIEngineExtension } from '../domain/abstracts/oai.abstract';
+import { ConfigsUsecases } from '@/usecases/configs/configs.usecase';
+import { EventEmitter2 } from '@nestjs/event-emitter';
+import { EngineStatus } from '@/domain/abstracts/engine.abstract';
+
+/**
+ * A class that implements the InferenceExtension interface from the @janhq/core package.
+ * The class provides methods for initializing and stopping a model, and for making inference requests.
+ * It also subscribes to events emitted by the @janhq/core package and handles new message requests.
+ */
+export default class NvidiaEngineExtension extends OAIEngineExtension {
+  apiUrl = 'https://integrate.api.nvidia.com/v1/chat/completions';
+  name = 'nvidia';
+  productName = 'Nvidia Inference Engine';
+  description = 'This extension enables Nvidia chat completion API calls';
+  version = '0.0.1';
+  apiKey?: string;
+
+  constructor(
+    protected readonly httpService: HttpService,
+    protected readonly configsUsecases: ConfigsUsecases,
+    protected readonly eventEmmitter: EventEmitter2,
+  ) {
+    super(httpService);
+
+    eventEmmitter.on('config.updated', async (data) => {
+      if (data.engine === this.name) {
+        this.apiKey = data.value;
+        this.status =
+          (this.apiKey?.length ?? 0) > 0
+            ? EngineStatus.READY
+            : EngineStatus.MISSING_CONFIGURATION;
+      }
+    });
+  }
+
+  async onLoad() {
+    const configs = (await this.configsUsecases.getGroupConfigs(
+      this.name,
+    )) as unknown as { apiKey: string };
+
+    this.apiKey = configs?.apiKey;
+    this.status =
+      (this.apiKey?.length ?? 0) > 0
+        ? EngineStatus.READY
+        : EngineStatus.MISSING_CONFIGURATION;
+  }
+}
diff --git a/cortex-js/src/extensions/openrouter.engine.ts b/cortex-js/src/extensions/openrouter.engine.ts
@@ -1,9 +1,7 @@
-import stream from 'stream';
 import { HttpService } from '@nestjs/axios';
 import { OAIEngineExtension } from '../domain/abstracts/oai.abstract';
 import { ConfigsUsecases } from '@/usecases/configs/configs.usecase';
 import { EventEmitter2 } from '@nestjs/event-emitter';
-import _ from 'lodash';
 import { EngineStatus } from '@/domain/abstracts/engine.abstract';
 
 /**
diff --git a/cortex-js/src/infrastructure/commanders/types/engine.interface.ts b/cortex-js/src/infrastructure/commanders/types/engine.interface.ts
@@ -10,6 +10,8 @@ export enum Engines {
   anthropic = 'anthropic',
   openrouter = 'openrouter',
   cohere = 'cohere',
+  martian = 'martian',
+  nvidia = 'nvidia',
 }
 
 export const EngineNamesMap: {
@@ -27,4 +29,6 @@ export const RemoteEngines: Engines[] = [
   Engines.anthropic,
   Engines.openrouter,
   Engines.cohere,
+  Engines.martian,
+  Engines.nvidia,
 ];
diff --git a/cortex-js/src/infrastructure/controllers/chat.controller.ts b/cortex-js/src/infrastructure/controllers/chat.controller.ts
@@ -33,7 +33,7 @@ export class ChatController {
     @Body() createChatDto: CreateChatCompletionDto,
     @Res() res: Response,
   ) {
-    let { stream } = createChatDto;
+    const { stream } = createChatDto;
     this.chatService
       .inference(createChatDto, extractCommonHeaders(headers))
       .then((response) => {
@@ -50,6 +50,9 @@ export class ChatController {
         let errorMessage;
         if (!stream) {
           const data = error.response?.data;
+          if (!data) {
+            return res.status(500).send(error.message || 'An error occurred');
+          }
           return res
             .status(statusCode)
             .send(
@@ -60,8 +63,10 @@ export class ChatController {
             );
         }
         const streamResponse = error.response?.data;
+        if (!streamResponse) {
+          return res.status(500).send(error.message || 'An error occurred');
+        }
         let data = '';
-
         streamResponse.on('data', (chunk: any) => {
           data += chunk;
         });
diff --git a/cortex-js/src/infrastructure/services/resources-manager/resources-manager.service.ts b/cortex-js/src/infrastructure/services/resources-manager/resources-manager.service.ts
@@ -21,7 +21,7 @@ export class ResourcesManagerService {
     return {
       mem: memInfo,
       cpu: {
-        usage: cpuUsage.avgLoad,
+        usage: Number(cpuUsage.currentLoad.toFixed(2)),
       },
     };
   }

Original file line number	Diff line number	Diff line change
`@@ -21,7 +21,7 @@ export class ResourcesManagerService {`
`21`	`21`	`return {`
`22`	`22`	`mem: memInfo,`
`23`	`23`	`cpu: {`
`24`		`- usage: cpuUsage.avgLoad,`
	`24`	`+ usage: Number(cpuUsage.currentLoad.toFixed(2)),`
`25`	`25`	`},`
`26`	`26`	`};`
`27`	`27`	`}`