@lobehub/chat 0.126.5 → 0.127.0

package/src/libs/agent-runtime/ollama/index.ts

@@ -0,0 +1,80 @@
+import { OpenAIStream, StreamingTextResponse } from 'ai';
+import OpenAI from 'openai';
+
+import { LobeRuntimeAI } from '../BaseAI';
+import { AgentRuntimeErrorType } from '../error';
+import { ChatStreamPayload, ModelProvider } from '../types';
+import { AgentRuntimeError } from '../utils/createError';
+import { debugStream } from '../utils/debugStream';
+import { desensitizeUrl } from '../utils/desensitizeUrl';
+import { DEBUG_CHAT_COMPLETION } from '../utils/env';
+import { handleOpenAIError } from '../utils/handleOpenAIError';
+
+const DEFAULT_BASE_URL = 'http://127.0.0.1:11434/v1';
+
+export class LobeOllamaAI implements LobeRuntimeAI {
+  private _llm: OpenAI;
+
+  baseURL: string;
+
+  constructor(baseURL: string = DEFAULT_BASE_URL) {
+    if (!baseURL) throw AgentRuntimeError.createError(AgentRuntimeErrorType.InvalidOllamaArgs);
+
+    this._llm = new OpenAI({ apiKey: 'ollama', baseURL });
+    this.baseURL = baseURL;
+  }
+
+  async chat(payload: ChatStreamPayload) {
+    try {
+      const response = await this._llm.chat.completions.create(
+        payload as unknown as OpenAI.ChatCompletionCreateParamsStreaming,
+      );
+
+      const stream = OpenAIStream(response);
+
+      const [debug, returnStream] = stream.tee();
+
+      if (DEBUG_CHAT_COMPLETION) {
+        debugStream(debug).catch(console.error);
+      }
+
+      return new StreamingTextResponse(returnStream);
+    } catch (error) {
+      let desensitizedEndpoint = this.baseURL;
+
+      if (this.baseURL !== DEFAULT_BASE_URL) {
+        desensitizedEndpoint = desensitizeUrl(this.baseURL);
+      }
+
+      if ('status' in (error as any)) {
+        switch ((error as Response).status) {
+          case 401: {
+            throw AgentRuntimeError.chat({
+              endpoint: desensitizedEndpoint,
+              error: error as any,
+              errorType: AgentRuntimeErrorType.InvalidOllamaArgs,
+              provider: ModelProvider.Ollama,
+            });
+          }
+
+          default: {
+            break;
+          }
+        }
+      }
+
+      const { errorResult, RuntimeError } = handleOpenAIError(error);
+
+      const errorType = RuntimeError || AgentRuntimeErrorType.OllamaBizError;
+
+      throw AgentRuntimeError.chat({
+        endpoint: desensitizedEndpoint,
+        error: errorResult,
+        errorType,
+        provider: ModelProvider.Ollama,
+      });
+    }
+  }
+}
+
+export default LobeOllamaAI;

package/src/libs/agent-runtime/types/type.ts

@@ -29,6 +29,7 @@ export enum ModelProvider {
   Google = 'google',
   Mistral = 'mistral',
   Moonshot = 'moonshot',
+  Ollama = 'ollama',
   OpenAI = 'openai',
   Tongyi = 'tongyi',
   ZhiPu = 'zhipu',

package/src/locales/default/common.ts

@@ -103,6 +103,7 @@ export default {
     bedrock: 'AWS Bedrock',
     google: 'Google',
     moonshot: 'Moonshot AI',
+    ollama: 'Ollama',
     oneapi: 'One API',
     openai: 'OpenAI',
     zhipu: '智谱AI',

package/src/locales/default/error.ts

@@ -73,6 +73,9 @@ export default {
     InvalidAzureAPIKey: 'Azure API Key 不正确或为空，请检查 Azure API Key 后重试',
     AzureBizError: '请求 Azure AI 服务出错，请根据以下信息排查或重试',
 
+    InvalidOllamaArgs: 'Ollama 配置不正确，请检查 Ollama 配置后重试',
+    OllamaBizError: '请求 Ollama 服务出错，请根据以下信息排查或重试',
+
     AgentRuntimeError: 'Lobe 语言模型运行时执行出错，请根据以下信息排查或重试',
     /* eslint-enable */
   },

package/src/locales/default/setting.ts

@@ -85,6 +85,22 @@ export default {
         title: 'API Key',
       },
     },
+    Ollama: {
+      checker: {
+        desc: '测试代理地址是否正确填写',
+      },
+      customModelName: {
+        desc: '增加自定义模型，多个模型使用逗号（,）隔开',
+        placeholder: 'vicuna,llava,codellama,llama2:13b-text',
+        title: '自定义模型名称',
+      },
+      endpoint: {
+        desc: '填入 Ollama 接口代理地址，本地未额外指定可留空',
+        placeholder: 'http://127.0.0.1:11434/v1',
+        title: '接口代理地址',
+      },
+      title: 'Ollama',
+    },
     OpenAI: {
       azureApiVersion: {
         desc: 'Azure 的 API 版本，遵循 YYYY-MM-DD 格式，查阅[最新版本](https://learn.microsoft.com/zh-cn/azure/ai-services/openai/reference#chat-completions)',

package/src/services/_auth.ts

@@ -40,6 +40,14 @@ const getProviderAuthPayload = (provider: string) => {
       };
     }
 
+    case ModelProvider.Ollama: {
+      const endpoint = modelProviderSelectors.ollamaProxyUrl(useGlobalStore.getState());
+
+      return {
+        endpoint,
+      };
+    }
+
     default:
     case ModelProvider.OpenAI: {
       const openai = modelProviderSelectors.openAIConfig(useGlobalStore.getState());

package/src/store/global/slices/settings/selectors/modelProvider.ts

@@ -5,6 +5,7 @@ import {
   GoogleProvider,
   LOBE_DEFAULT_MODEL_LIST,
   MoonshotProvider,
+  OllamaProvider,
   OpenAIProvider,
   ZhiPuProvider,
 } from '@/config/modelProviders';
@@ -37,6 +38,12 @@ const azureConfig = (s: GlobalStore) => modelProvider(s).azure;
 const enableMoonshot = (s: GlobalStore) => modelProvider(s).moonshot.enabled;
 const moonshotAPIKey = (s: GlobalStore) => modelProvider(s).moonshot.apiKey;
 
+const enableOllamaConfigInSettings = (s: GlobalStore) =>
+  s.serverConfig.languageModel?.ollama?.enabled || false;
+
+const enableOllama = (s: GlobalStore) => modelProvider(s).ollama.enabled;
+const ollamaProxyUrl = (s: GlobalStore) => modelProvider(s).ollama.endpoint;
+
 // const azureModelList = (s: GlobalStore): ModelProviderCard => {
 //   const azure = azureConfig(s);
 //   return {
@@ -46,8 +53,11 @@ const moonshotAPIKey = (s: GlobalStore) => modelProvider(s).moonshot.apiKey;
 // };
 
 // 提取处理 chatModels 的专门方法
-const processChatModels = (modelConfig: ReturnType<typeof parseModelString>): ChatModelCard[] => {
-  let chatModels = modelConfig.removeAll ? [] : OpenAIProvider.chatModels;
+const processChatModels = (
+  modelConfig: ReturnType<typeof parseModelString>,
+  defaultChartModels = OpenAIProvider.chatModels,
+): ChatModelCard[] => {
+  let chatModels = modelConfig.removeAll ? [] : defaultChartModels;
 
   // 处理移除逻辑
   if (!modelConfig.removeAll) {
@@ -102,6 +112,12 @@ const modelSelectList = (s: GlobalStore): ModelProviderCard[] => {
 
   const chatModels = processChatModels(modelConfig);
 
+  const ollamaModelConfig = parseModelString(
+    currentSettings(s).languageModel.ollama.customModelName,
+  );
+
+  const ollamaChatModels = processChatModels(ollamaModelConfig, OllamaProvider.chatModels);
+
   return [
     {
       ...OpenAIProvider,
@@ -112,6 +128,7 @@ const modelSelectList = (s: GlobalStore): ModelProviderCard[] => {
     { ...MoonshotProvider, enabled: enableMoonshot(s) },
     { ...GoogleProvider, enabled: enableGoogle(s) },
     { ...BedrockProvider, enabled: enableBedrock(s) },
+    { ...OllamaProvider, chatModels: ollamaChatModels, enabled: enableOllama(s) },
   ];
 };
 
@@ -174,4 +191,9 @@ export const modelProviderSelectors = {
   // Moonshot
   enableMoonshot,
   moonshotAPIKey,
+
+  // Ollama
+  enableOllamaConfigInSettings,
+  enableOllama,
+  ollamaProxyUrl,
 };

package/src/types/settings/modelProvider.ts

@@ -47,11 +47,18 @@ export interface AWSBedrockConfig {
   secretAccessKey?: string;
 }
 
+export interface OllamaConfig {
+  customModelName?: string;
+  enabled?: boolean;
+  endpoint?: string;
+}
+
 export interface GlobalLLMConfig {
   azure: AzureOpenAIConfig;
   bedrock: AWSBedrockConfig;
   google: GoogleConfig;
   moonshot: MoonshotConfig;
+  ollama: OllamaConfig;
   openAI: OpenAIConfig;
   zhipu: ZhiPuConfig;
 }