🤖 feat: Support new o1 model

danny-avila · Dec 18, 2024 · 0b05325 · 0b05325
1 parent f873587
commit 0b05325
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 4 deletions.
diff --git a/api/app/clients/OpenAIClient.js b/api/app/clients/OpenAIClient.js
@@ -107,7 +107,8 @@ class OpenAIClient extends BaseClient {
       this.checkVisionRequest(this.options.attachments);
     }
 
-    this.isO1Model = /\bo1\b/i.test(this.modelOptions.model);
+    const o1Pattern = /\bo1\b/i;
+    this.isO1Model = o1Pattern.test(this.modelOptions.model);
 
     const { OPENROUTER_API_KEY, OPENAI_FORCE_PROMPT } = process.env ?? {};
     if (OPENROUTER_API_KEY && !this.azure) {
@@ -147,7 +148,7 @@ class OpenAIClient extends BaseClient {
     const { model } = this.modelOptions;
 
     this.isChatCompletion =
-      /\bo1\b/i.test(model) || model.includes('gpt') || this.useOpenRouter || !!reverseProxy;
+      o1Pattern.test(model) || model.includes('gpt') || this.useOpenRouter || !!reverseProxy;
     this.isChatGptModel = this.isChatCompletion;
     if (
       model.includes('text-davinci') ||
@@ -1325,7 +1326,11 @@ ${convo}
       /** @type {(value: void | PromiseLike<void>) => void} */
       let streamResolve;
 
-      if (this.isO1Model === true && this.azure && modelOptions.stream) {
+      if (
+        this.isO1Model === true &&
+        (this.azure || /o1(?!-(?:mini|preview)).*$/.test(modelOptions.model)) &&
+        modelOptions.stream
+      ) {
         delete modelOptions.stream;
         delete modelOptions.stop;
       }

diff --git a/api/utils/tokens.js b/api/utils/tokens.js
@@ -2,7 +2,7 @@ const z = require('zod');
 const { EModelEndpoint } = require('librechat-data-provider');
 
 const openAIModels = {
-  o1: 127500, // -500 from max
+  o1: 195000, // -5000 from max
   'o1-mini': 127500, // -500 from max
   'o1-preview': 127500, // -500 from max
   'gpt-4': 8187, // -5 from max

diff --git a/api/utils/tokens.spec.js b/api/utils/tokens.spec.js
@@ -248,6 +248,32 @@ describe('getModelMaxTokens', () => {
   test('should return undefined for a model when using an unsupported endpoint', () => {
     expect(getModelMaxTokens('azure-gpt-3', 'unsupportedEndpoint')).toBeUndefined();
   });
+
+  test('should return correct max context tokens for o1-series models', () => {
+    // Standard o1 variations
+    const o1Tokens = maxTokensMap[EModelEndpoint.openAI]['o1'];
+    expect(getModelMaxTokens('o1')).toBe(o1Tokens);
+    expect(getModelMaxTokens('o1-latest')).toBe(o1Tokens);
+    expect(getModelMaxTokens('o1-2024-12-17')).toBe(o1Tokens);
+    expect(getModelMaxTokens('o1-something-else')).toBe(o1Tokens);
+    expect(getModelMaxTokens('openai/o1-something-else')).toBe(o1Tokens);
+
+    // Mini variations
+    const o1MiniTokens = maxTokensMap[EModelEndpoint.openAI]['o1-mini'];
+    expect(getModelMaxTokens('o1-mini')).toBe(o1MiniTokens);
+    expect(getModelMaxTokens('o1-mini-latest')).toBe(o1MiniTokens);
+    expect(getModelMaxTokens('o1-mini-2024-09-12')).toBe(o1MiniTokens);
+    expect(getModelMaxTokens('o1-mini-something')).toBe(o1MiniTokens);
+    expect(getModelMaxTokens('openai/o1-mini-something')).toBe(o1MiniTokens);
+
+    // Preview variations
+    const o1PreviewTokens = maxTokensMap[EModelEndpoint.openAI]['o1-preview'];
+    expect(getModelMaxTokens('o1-preview')).toBe(o1PreviewTokens);
+    expect(getModelMaxTokens('o1-preview-latest')).toBe(o1PreviewTokens);
+    expect(getModelMaxTokens('o1-preview-2024-09-12')).toBe(o1PreviewTokens);
+    expect(getModelMaxTokens('o1-preview-something')).toBe(o1PreviewTokens);
+    expect(getModelMaxTokens('openai/o1-preview-something')).toBe(o1PreviewTokens);
+  });
 });
 
 describe('matchModelName', () => {