✨ feat: Add Taichu model provider (#3129)

* ✨ feat: Add Taichu model provider * 💄 style: fix model tag icon missing * 💄 style: update Taichu icon * 💄 style: resize Taichu icon
lobehub · Jul 9, 2024 · a4580e9 · a4580e9
1 parent e6dc16a
commit a4580e9
Show file tree

Hide file tree

Showing 16 changed files with 340 additions and 0 deletions.
diff --git a/src/app/(main)/settings/llm/ProviderList/providers.tsx b/src/app/(main)/settings/llm/ProviderList/providers.tsx
@@ -1,4 +1,5 @@
 import {
+  AiMass,
   Anthropic,
   Baichuan,
   Claude,
@@ -36,6 +37,7 @@ import {
   PerplexityProviderCard,
   QwenProviderCard,
   StepfunProviderCard,
+  TaichuProviderCard,
   TogetherAIProviderCard,
   ZeroOneProviderCard,
   ZhiPuProviderCard,
@@ -177,6 +179,11 @@ export const useProviderList = (): ProviderItem[] => {
         docUrl: urlJoin(BASE_DOC_URL, 'baichuan'),
         title: <Baichuan.Combine size={ 20 } type={ 'color' } />,
       },
+      {
+        ...TaichuProviderCard,
+        docUrl: urlJoin(BASE_DOC_URL, 'taichu'),
+        title: <AiMass.Combine size={ 28 } type={ 'color' } />,
+      },
     ],
     [azureProvider, ollamaProvider, ollamaProvider, bedrockProvider],
   );

diff --git a/src/app/api/chat/agentRuntime.ts b/src/app/api/chat/agentRuntime.ts
@@ -177,6 +177,13 @@ const getLlmOptionsFromPayload = (provider: string, payload: JWTPayload) => {
 
       const apiKey = apiKeyManager.pick(payload?.apiKey || BAICHUAN_API_KEY);
 
+      return { apiKey };
+    }
+    case ModelProvider.Taichu: {
+      const { TAICHU_API_KEY } = getLLMConfig();
+
+      const apiKey = apiKeyManager.pick(payload?.apiKey || TAICHU_API_KEY);
+
       return { apiKey };
     }
   }

diff --git a/src/components/ModelIcon/index.tsx b/src/components/ModelIcon/index.tsx
@@ -1,4 +1,5 @@
 import {
+  AiMass,
   Adobe,
   Ai21,
   Aws,
@@ -66,6 +67,7 @@ const ModelIcon = memo<ModelProviderIconProps>(({ model: originModel, size = 12
   if (model.includes('command')) return <Cohere.Avatar size={size} />;
   if (model.includes('dbrx')) return <Dbrx.Avatar size={size} />;
   if (model.includes('step')) return <Stepfun.Avatar size={size} />;
+  if (model.includes('taichu')) return <AiMass.Avatar size={size} />;
 
   // below: To be supported in providers, move up if supported
   if (model.includes('baichuan'))

diff --git a/src/components/ModelProviderIcon/index.tsx b/src/components/ModelProviderIcon/index.tsx
@@ -1,4 +1,5 @@
 import {
+  AiMass,
   Anthropic,
   Azure,
   Baichuan,
@@ -119,6 +120,10 @@ const ModelProviderIcon = memo<ModelProviderIconProps>(({ provider }) => {
       return <Baichuan size={20} />;
     }
 
+    case ModelProvider.Taichu: {
+      return <AiMass size={20} />;
+    }
+
     default: {
       return null;
     }

diff --git a/src/components/ModelTag/ModelIcon.tsx b/src/components/ModelTag/ModelIcon.tsx
@@ -1,6 +1,7 @@
 import {
   AdobeFirefly,
   Ai21,
+  AiMass,
   Aws,
   Azure,
   Baichuan,
@@ -61,6 +62,7 @@ const ModelIcon = memo<ModelIconProps>(({ model, size = 12 }) => {
   if (model.startsWith('openchat')) return <OpenChat size={size} />;
   if (model.includes('command')) return <Cohere size={size} />;
   if (model.includes('dbrx')) return <Dbrx size={size} />;
+  if (model.includes('taichu')) return <AiMass size={size} />;
 
   // below: To be supported in providers, move up if supported
   if (model.includes('baichuan')) return <Baichuan size={size} />;

diff --git a/src/config/llm.ts b/src/config/llm.ts
@@ -85,6 +85,9 @@ export const getLLMConfig = () => {
 
       ENABLED_BAICHUAN: z.boolean(),
       BAICHUAN_API_KEY: z.string().optional(),
+
+      ENABLED_TAICHU: z.boolean(),
+      TAICHU_API_KEY: z.string().optional(),
     },
     runtimeEnv: {
       API_KEY_SELECT_MODE: process.env.API_KEY_SELECT_MODE,
@@ -161,6 +164,9 @@ export const getLLMConfig = () => {
 
       ENABLED_BAICHUAN: !!process.env.BAICHUAN_API_KEY,
       BAICHUAN_API_KEY: process.env.BAICHUAN_API_KEY,
+
+      ENABLED_TAICHU: !!process.env.TAICHU_API_KEY,
+      TAICHU_API_KEY: process.env.TAICHU_API_KEY,
     },
   });
 };

diff --git a/src/config/modelProviders/index.ts b/src/config/modelProviders/index.ts
@@ -16,6 +16,7 @@ import OpenRouterProvider from './openrouter';
 import PerplexityProvider from './perplexity';
 import QwenProvider from './qwen';
 import StepfunProvider from './stepfun';
+import TaichuProvider from './taichu';
 import TogetherAIProvider from './togetherai';
 import ZeroOneProvider from './zeroone';
 import ZhiPuProvider from './zhipu';
@@ -39,6 +40,7 @@ export const LOBE_DEFAULT_MODEL_LIST: ChatModelCard[] = [
   ZeroOneProvider.chatModels,
   StepfunProvider.chatModels,
   BaichuanProvider.chatModels,
+  TaichuProvider.chatModels,
 ].flat();
 
 export const DEFAULT_MODEL_PROVIDER_LIST = [
@@ -61,6 +63,7 @@ export const DEFAULT_MODEL_PROVIDER_LIST = [
   ZhiPuProvider,
   StepfunProvider,
   BaichuanProvider,
+  TaichuProvider,
 ];
 
 export const filterEnabledModels = (provider: ModelProviderCard) => {
@@ -83,6 +86,7 @@ export { default as OpenRouterProviderCard } from './openrouter';
 export { default as PerplexityProviderCard } from './perplexity';
 export { default as QwenProviderCard } from './qwen';
 export { default as StepfunProviderCard } from './stepfun';
+export { default as TaichuProviderCard } from './taichu';
 export { default as TogetherAIProviderCard } from './togetherai';
 export { default as ZeroOneProviderCard } from './zeroone';
 export { default as ZhiPuProviderCard } from './zhipu';
diff --git a/src/config/modelProviders/taichu.ts b/src/config/modelProviders/taichu.ts
@@ -0,0 +1,21 @@
+import { ModelProviderCard } from '@/types/llm';
+
+// ref https://ai-maas.wair.ac.cn/#/doc
+const Taichu: ModelProviderCard = {
+  chatModels: [
+    {
+      description: '语言大模型',
+      displayName: 'Taichu-2.0',
+      enabled: true,
+      functionCall: false,
+      id: 'taichu_llm',
+      tokens: 32_768,
+    },
+  ],
+  checkModel: 'taichu_llm',
+  id: 'taichu',
+  modelList: { showModelFetcher: true },
+  name: 'Taichu',
+};
+
+export default Taichu;
diff --git a/src/const/settings/llm.ts b/src/const/settings/llm.ts
@@ -14,6 +14,7 @@ import {
   PerplexityProviderCard,
   QwenProviderCard,
   StepfunProviderCard,
+  TaichuProviderCard,
   TogetherAIProviderCard,
   ZeroOneProviderCard,
   ZhiPuProviderCard,
@@ -87,6 +88,10 @@ export const DEFAULT_LLM_CONFIG: UserModelProviderConfig = {
     enabled: false,
     enabledModels: filterEnabledModels(StepfunProviderCard),
   },
+  taichu: {
+    enabled: false,
+    enabledModels: filterEnabledModels(TaichuProviderCard),
+  },
   togetherai: {
     enabled: false,
     enabledModels: filterEnabledModels(TogetherAIProviderCard),

diff --git a/src/features/Conversation/Error/APIKeyForm/ProviderAvatar.tsx b/src/features/Conversation/Error/APIKeyForm/ProviderAvatar.tsx
@@ -1,4 +1,5 @@
 import {
+  AiMass,
   Anthropic,
   Baichuan,
   DeepSeek,
@@ -81,6 +82,10 @@ const ProviderAvatar = memo<ProviderAvatarProps>(({ provider }) => {
       return <Stepfun color={Stepfun.colorPrimary} size={56} />;
     }
 
+    case ModelProvider.Taichu: {
+      return <AiMass.Color size={56} />;
+    }
+
     case ModelProvider.TogetherAI: {
       return <Together color={Together.colorPrimary} size={56} />;
     }

diff --git a/src/libs/agent-runtime/AgentRuntime.ts b/src/libs/agent-runtime/AgentRuntime.ts
@@ -19,6 +19,7 @@ import { LobeOpenRouterAI } from './openrouter';
 import { LobePerplexityAI } from './perplexity';
 import { LobeQwenAI } from './qwen';
 import { LobeStepfunAI } from './stepfun';
+import { LobeTaichuAI } from './taichu';
 import { LobeTogetherAI } from './togetherai';
 import {
   ChatCompetitionOptions,
@@ -118,6 +119,7 @@ class AgentRuntime {
       perplexity: Partial<ClientOptions>;
       qwen: Partial<ClientOptions>;
       stepfun: Partial<ClientOptions>;
+      taichu: Partial<ClientOptions>;
       togetherai: Partial<ClientOptions>;
       zeroone: Partial<ClientOptions>;
       zhipu: Partial<ClientOptions>;
@@ -226,6 +228,11 @@ class AgentRuntime {
         runtimeModel = new LobeBaichuanAI(params.baichuan ?? {});
         break
       }
+
+      case ModelProvider.Taichu: {
+        runtimeModel = new LobeTaichuAI(params.taichu ?? {});
+        break
+      }
     }
 
     return new AgentRuntime(runtimeModel);