feat: o1 models support streaming now.

logancyang · Dec 2, 2024 · da9c4fe · da9c4fe
1 parent 8324d1d
commit da9c4fe
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 22 deletions.
diff --git a/package-lock.json b/package-lock.json
diff --git a/package.json b/package.json
@@ -69,6 +69,7 @@
     "@langchain/core": "^0.3.3",
     "@langchain/google-genai": "^0.1.2",
     "@langchain/groq": "^0.1.2",
+    "@langchain/openai": "^0.3.14",
     "@orama/orama": "^3.0.0-rc-2",
     "@radix-ui/react-dropdown-menu": "^2.1.2",
     "@radix-ui/react-tooltip": "^1.1.3",

diff --git a/src/LLMProviders/chainManager.ts b/src/LLMProviders/chainManager.ts
@@ -276,7 +276,7 @@ export default class ChainManager {
       if (isO1Model) {
         //  Temporary fix：for o1-xx model need to covert systemMessage to aiMessage
         effectivePrompt = ChatPromptTemplate.fromMessages([
-          [AI_SENDER, this.getLangChainParams().systemMessage || ""],
+          [AI_SENDER, getSystemPrompt() || ""],
           effectivePrompt,
         ]);
       }

diff --git a/src/LLMProviders/chatModelManager.ts b/src/LLMProviders/chatModelManager.ts
@@ -68,9 +68,6 @@ export default class ChatModelManager {
   }
 
   private getModelConfig(customModel: CustomModel): ModelConfig {
-    const decrypt = (key: string) => this.encryptionService.getDecryptedKey(key);
-    const params = this.getLangChainParams();
-
     const settings = getSettings();
 
     // Check if the model starts with "o1"
@@ -97,7 +94,7 @@ export default class ChatModelManager {
         },
         // @ts-ignore
         openAIOrgId: getDecryptedKey(settings.openAIOrgId),
-        ...this.handleOpenAIExtraArgs(isO1Model, settings.maxTokens, settings.temperature, true),
+        ...this.handleOpenAIExtraArgs(isO1Model, settings.maxTokens, settings.temperature),
       },
       [ChatModelProviders.ANTHROPIC]: {
         anthropicApiKey: getDecryptedKey(customModel.apiKey || settings.anthropicApiKey),
@@ -118,7 +115,7 @@ export default class ChatModelManager {
           baseURL: customModel.baseUrl,
           fetch: customModel.enableCors ? safeFetch : undefined,
         },
-        ...this.handleOpenAIExtraArgs(isO1Model, settings.maxTokens, settings.temperature, true),
+        ...this.handleOpenAIExtraArgs(isO1Model, settings.maxTokens, settings.temperature),
       },
       [ChatModelProviders.COHEREAI]: {
         apiKey: getDecryptedKey(customModel.apiKey || settings.cohereApiKey),
@@ -183,7 +180,7 @@ export default class ChatModelManager {
           fetch: customModel.enableCors ? safeFetch : undefined,
           dangerouslyAllowBrowser: true,
         },
-        ...this.handleOpenAIExtraArgs(isO1Model, settings.maxTokens, settings.temperature, true),
+        ...this.handleOpenAIExtraArgs(isO1Model, settings.maxTokens, settings.temperature),
       },
     };
 
@@ -193,22 +190,15 @@ export default class ChatModelManager {
     return { ...baseConfig, ...selectedProviderConfig };
   }
 
-  private handleOpenAIExtraArgs(
-    isO1Model: boolean,
-    maxTokens: number,
-    temperature: number,
-    streaming: boolean
-  ) {
+  private handleOpenAIExtraArgs(isO1Model: boolean, maxTokens: number, temperature: number) {
     return isO1Model
       ? {
           maxCompletionTokens: maxTokens,
           temperature: 1,
-          streaming: false,
         }
       : {
           maxTokens: maxTokens,
           temperature: temperature,
-          streaming: streaming,
         };
   }