From 61bfeb2c99425e9a888480a34a40aed451a9252a Mon Sep 17 00:00:00 2001
From: Arvin Xu <arvinx@foxmail.com>
Date: Fri, 13 Sep 2024 12:54:21 +0800
Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20feat:=20support=20openai=20new=20Op?=
 =?UTF-8?q?enAI=20o1-preview/o1-mini=20models=20(#3943)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* support o1

* Update openai.ts

* fix test

* fix test
---
 src/config/modelProviders/openai.ts           | 28 +++++++++++++++++++
 src/libs/agent-runtime/openai/index.ts        | 28 ++++++++++++++++++-
 .../modelList/selectors/modelProvider.test.ts |  9 +++++-
 3 files changed, 63 insertions(+), 2 deletions(-)

diff --git a/src/config/modelProviders/openai.ts b/src/config/modelProviders/openai.ts
index 75b5d34b0eee..7784b1290424 100644
--- a/src/config/modelProviders/openai.ts
+++ b/src/config/modelProviders/openai.ts
@@ -3,6 +3,34 @@ import { ModelProviderCard } from '@/types/llm';
 // ref: https://platform.openai.com/docs/deprecations
 const OpenAI: ModelProviderCard = {
   chatModels: [
+    {
+      description:
+        'o1-mini是一款针对编程、数学和科学应用场景而设计的快速、经济高效的推理模型。该模型具有128K上下文和2023年10月的知识截止日期。',
+      displayName: 'OpenAI o1-mini',
+      enabled: true,
+      id: 'o1-mini',
+      maxOutput: 65_536,
+      pricing: {
+        input: 3,
+        output: 12,
+      },
+      releasedAt: '2024-09-12',
+      tokens: 128_000,
+    },
+    {
+      description:
+        'o1是OpenAI新的推理模型，适用于需要广泛通用知识的复杂任务。该模型具有128K上下文和2023年10月的知识截止日期。',
+      displayName: 'OpenAI o1-preview',
+      enabled: true,
+      id: 'o1-preview',
+      maxOutput: 32_768,
+      pricing: {
+        input: 15,
+        output: 60,
+      },
+      releasedAt: '2024-09-12',
+      tokens: 128_000,
+    },
     {
       description:
         'GPT-4o mini是OpenAI在GPT-4 Omni之后推出的最新模型，支持图文输入并输出文本。作为他们最先进的小型模型，它比其他近期的前沿模型便宜很多，并且比GPT-3.5 Turbo便宜超过60%。它保持了最先进的智能，同时具有显著的性价比。GPT-4o mini在MMLU测试中获得了 82% 的得分，目前在聊天偏好上排名高于 GPT-4。',
diff --git a/src/libs/agent-runtime/openai/index.ts b/src/libs/agent-runtime/openai/index.ts
index df222bb76a2c..61ca33aeb03f 100644
--- a/src/libs/agent-runtime/openai/index.ts
+++ b/src/libs/agent-runtime/openai/index.ts
@@ -1,8 +1,34 @@
-import { ModelProvider } from '../types';
+import { ChatStreamPayload, ModelProvider, OpenAIChatMessage } from '../types';
 import { LobeOpenAICompatibleFactory } from '../utils/openaiCompatibleFactory';
 
+// TODO: 临时写法，后续要重构成 model card 展示配置
+const o1Models = new Set(['o1-preview', 'o1-mini']);
+const truneO1Payload = (payload: ChatStreamPayload) => ({
+  ...payload,
+  frequency_penalty: 0,
+  messages: payload.messages.map((message: OpenAIChatMessage) => ({
+    ...message,
+    role: message.role === 'system' ? 'user' : message.role,
+  })),
+  presence_penalty: 0,
+  stream: false,
+  temperature: 1,
+  top_p: 1,
+});
+
 export const LobeOpenAI = LobeOpenAICompatibleFactory({
   baseURL: 'https://api.openai.com/v1',
+  chatCompletion: {
+    handlePayload: (payload) => {
+      const { model } = payload;
+
+      if (o1Models.has(model)) {
+        return truneO1Payload(payload) as any;
+      }
+
+      return { ...payload, stream: payload.stream ?? true };
+    },
+  },
   debug: {
     chatCompletion: () => process.env.DEBUG_OPENAI_CHAT_COMPLETION === '1',
   },
diff --git a/src/store/user/slices/modelList/selectors/modelProvider.test.ts b/src/store/user/slices/modelList/selectors/modelProvider.test.ts
index 38e755741e0a..0aeaa7a67f93 100644
--- a/src/store/user/slices/modelList/selectors/modelProvider.test.ts
+++ b/src/store/user/slices/modelList/selectors/modelProvider.test.ts
@@ -49,7 +49,14 @@ describe('modelProviderSelectors', () => {
       const s = merge(initialState, {}) as unknown as UserStore;
 
       const result = modelProviderSelectors.getDefaultEnabledModelsById('openai')(s);
-      expect(result).toEqual(['gpt-4o-mini', 'gpt-4o', 'gpt-4o-2024-08-06', 'chatgpt-4o-latest']);
+      expect(result).toEqual([
+        'o1-mini',
+        'o1-preview',
+        'gpt-4o-mini',
+        'gpt-4o',
+        'gpt-4o-2024-08-06',
+        'chatgpt-4o-latest',
+      ]);
     });
 
     it('should return undefined for a non-existing provider', () => {