Migration of mistral ai client

OpenCTI-Platform · Jan 2, 2025 · 5da46e3 · 5da46e3
1 parent 7f1e42f
commit 5da46e3
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 10 deletions.
diff --git a/opencti-platform/opencti-graphql/package.json b/opencti-platform/opencti-graphql/package.json
@@ -54,7 +54,7 @@
     "@graphql-tools/schema": "10.0.6",
     "@graphql-tools/utils": "10.5.4",
     "@jorgeferrero/stream-to-buffer": "2.0.6",
-    "@mistralai/mistralai": "1.3.5",
+    "@mistralai/mistralai": "^1.3.5",
     "@node-saml/passport-saml": "5.0.0",
     "@opensearch-project/opensearch": "2.12.0",
     "@opentelemetry/api": "1.9.0",
@@ -159,7 +159,8 @@
     "winston": "3.17.0",
     "winston-daily-rotate-file": "5.0.0",
     "ws": "8.18.0",
-    "xml2js": "0.6.2"
+    "xml2js": "0.6.2",
+    "zod": "3.24.1"
   },
   "devDependencies": {
     "@graphql-codegen/cli": "5.0.2",

diff --git a/opencti-platform/opencti-graphql/src/database/ai-llm.ts b/opencti-platform/opencti-graphql/src/database/ai-llm.ts
@@ -1,4 +1,4 @@
-import MistralClient from '@mistralai/mistralai';
+import { Mistral } from '@mistralai/mistralai';
 import OpenAI from 'openai';
 import conf, { BUS_TOPICS, logApp } from '../config/conf';
 import { isEmptyField } from './utils';
@@ -13,11 +13,14 @@ const AI_ENDPOINT = conf.get('ai:endpoint');
 const AI_TOKEN = conf.get('ai:token');
 const AI_MODEL = conf.get('ai:model');
 
-let client: MistralClient | OpenAI | null = null;
+let client: Mistral | OpenAI | null = null;
 if (AI_ENABLED && AI_TOKEN) {
   switch (AI_TYPE) {
     case 'mistralai':
-      client = new MistralClient(AI_TOKEN, isEmptyField(AI_ENDPOINT) ? undefined : AI_ENDPOINT);
+      client = new Mistral({
+        serverURL: isEmptyField(AI_ENDPOINT) ? undefined : AI_ENDPOINT,
+        apiKey: AI_TOKEN,
+      });
       break;
     case 'openai':
       client = new OpenAI({
@@ -36,16 +39,16 @@ export const queryMistralAi = async (busId: string | null, question: string, use
   }
   try {
     logApp.debug('[AI] Querying MistralAI with prompt', { questionStart: question.substring(0, 100) });
-    const response = (client as MistralClient)?.chatStream({
+    const response = await (client as Mistral)?.chat.stream({
       model: AI_MODEL,
       messages: [{ role: 'user', content: question }],
     });
     let content = '';
     if (response) {
       // eslint-disable-next-line no-restricted-syntax
       for await (const chunk of response) {
-        if (chunk.choices[0].delta.content !== undefined) {
-          const streamText = chunk.choices[0].delta.content;
+        if (chunk.data.choices[0].delta.content !== undefined) {
+          const streamText = chunk.data.choices[0].delta.content;
           content += streamText;
           if (busId !== null) {
             await notify(BUS_TOPICS[AI_BUS].EDIT_TOPIC, { bus_id: busId, content }, user);

diff --git a/opencti-platform/opencti-graphql/yarn.lock b/opencti-platform/opencti-graphql/yarn.lock
@@ -3094,7 +3094,7 @@ __metadata:
   languageName: node
   linkType: hard
 
-"@mistralai/mistralai@npm:1.3.5":
+"@mistralai/mistralai@npm:^1.3.5":
   version: 1.3.5
   resolution: "@mistralai/mistralai@npm:1.3.5"
   peerDependencies:
@@ -10972,7 +10972,7 @@ __metadata:
     "@graphql-tools/utils": "npm:10.5.4"
     "@jorgeferrero/stream-to-buffer": "npm:2.0.6"
     "@luckycatfactory/esbuild-graphql-loader": "npm:3.8.1"
-    "@mistralai/mistralai": "npm:1.3.5"
+    "@mistralai/mistralai": "npm:^1.3.5"
     "@node-saml/passport-saml": "npm:5.0.0"
     "@opensearch-project/opensearch": "npm:2.12.0"
     "@opentelemetry/api": "npm:1.9.0"
@@ -11112,6 +11112,7 @@ __metadata:
     winston-daily-rotate-file: "npm:5.0.0"
     ws: "npm:8.18.0"
     xml2js: "npm:0.6.2"
+    zod: "npm:^3.24.1"
   languageName: unknown
   linkType: soft
 
@@ -14250,3 +14251,10 @@ __metadata:
   checksum: 10/aa5abd6a89590eadeba040afbc375f53337f12637e5e98330012a12d9886cde7a3ccc28bd91aafab50576035bbb1de39a9a316eecf2411c8b9009c9f94f0db27
   languageName: node
   linkType: hard
+
+"zod@npm:^3.24.1":
+  version: 3.24.1
+  resolution: "zod@npm:3.24.1"
+  checksum: 10/54e25956495dec22acb9399c168c6ba657ff279801a7fcd0530c414d867f1dcca279335e160af9b138dd70c332e17d548be4bc4d2f7eaf627dead50d914fec27
+  languageName: node
+  linkType: hard