@lobehub/chat 0.144.0 → 0.145.0

package/src/libs/agent-runtime/togetherai/index.test.ts

@@ -0,0 +1,347 @@
+// @vitest-environment node
+import OpenAI from 'openai';
+import { Mock, afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+
+import { ChatStreamCallbacks } from '@/libs/agent-runtime';
+
+import * as debugStreamModule from '../utils/debugStream';
+import { LobeTogetherAI } from './index';
+
+const provider = 'togetherai';
+const defaultBaseURL = 'https://api.together.xyz/v1';
+const bizErrorType = 'TogetherAIBizError';
+const invalidErrorType = 'InvalidTogetherAIAPIKey';
+
+// Mock the console.error to avoid polluting test output
+vi.spyOn(console, 'error').mockImplementation(() => {});
+
+let instance: LobeTogetherAI;
+
+beforeEach(() => {
+  instance = new LobeTogetherAI({ apiKey: 'test' });
+
+  // 使用 vi.spyOn 来模拟 chat.completions.create 方法
+  vi.spyOn(instance['client'].chat.completions, 'create').mockResolvedValue(
+    new ReadableStream() as any,
+  );
+});
+
+afterEach(() => {
+  vi.clearAllMocks();
+});
+
+describe('LobeTogetherAI', () => {
+  describe('init', () => {
+    it('should correctly initialize with an API key', async () => {
+      const instance = new LobeTogetherAI({ apiKey: 'test_api_key' });
+      expect(instance).toBeInstanceOf(LobeTogetherAI);
+      expect(instance.baseURL).toEqual(defaultBaseURL);
+    });
+  });
+
+  describe('chat', () => {
+    it('should return a StreamingTextResponse on successful API call', async () => {
+      // Arrange
+      const mockStream = new ReadableStream();
+      const mockResponse = Promise.resolve(mockStream);
+
+      (instance['client'].chat.completions.create as Mock).mockResolvedValue(mockResponse);
+
+      // Act
+      const result = await instance.chat({
+        messages: [{ content: 'Hello', role: 'user' }],
+        model: 'mistralai/mistral-7b-instruct:free',
+        temperature: 0,
+      });
+
+      // Assert
+      expect(result).toBeInstanceOf(Response);
+    });
+
+    it('should call TogetherAI API with corresponding options', async () => {
+      // Arrange
+      const mockStream = new ReadableStream();
+      const mockResponse = Promise.resolve(mockStream);
+
+      (instance['client'].chat.completions.create as Mock).mockResolvedValue(mockResponse);
+
+      // Act
+      const result = await instance.chat({
+        max_tokens: 1024,
+        messages: [{ content: 'Hello', role: 'user' }],
+        model: 'mistralai/mistral-7b-instruct:free',
+        temperature: 0.7,
+        top_p: 1,
+      });
+
+      // Assert
+      expect(instance['client'].chat.completions.create).toHaveBeenCalledWith({
+        max_tokens: 1024,
+        messages: [{ content: 'Hello', role: 'user' }],
+        model: 'mistralai/mistral-7b-instruct:free',
+        temperature: 0.7,
+        top_p: 1,
+      });
+      expect(result).toBeInstanceOf(Response);
+    });
+
+    describe('Error', () => {
+      it('should return TogetherAIBizError with an openai error response when OpenAI.APIError is thrown', async () => {
+        // Arrange
+        const apiError = new OpenAI.APIError(
+          400,
+          {
+            status: 400,
+            error: {
+              message: 'Bad Request',
+            },
+          },
+          'Error message',
+          {},
+        );
+
+        vi.spyOn(instance['client'].chat.completions, 'create').mockRejectedValue(apiError);
+
+        // Act
+        try {
+          await instance.chat({
+            messages: [{ content: 'Hello', role: 'user' }],
+            model: 'mistralai/mistral-7b-instruct:free',
+            temperature: 0,
+          });
+        } catch (e) {
+          expect(e).toEqual({
+            endpoint: defaultBaseURL,
+            error: {
+              error: { message: 'Bad Request' },
+              status: 400,
+            },
+            errorType: bizErrorType,
+            provider,
+          });
+        }
+      });
+
+      it('should throw AgentRuntimeError with InvalidTogetherAIAPIKey if no apiKey is provided', async () => {
+        try {
+          new LobeTogetherAI({});
+        } catch (e) {
+          expect(e).toEqual({ errorType: invalidErrorType });
+        }
+      });
+
+      it('should return TogetherAIBizError with the cause when OpenAI.APIError is thrown with cause', async () => {
+        // Arrange
+        const errorInfo = {
+          stack: 'abc',
+          cause: {
+            message: 'api is undefined',
+          },
+        };
+        const apiError = new OpenAI.APIError(400, errorInfo, 'module error', {});
+
+        vi.spyOn(instance['client'].chat.completions, 'create').mockRejectedValue(apiError);
+
+        // Act
+        try {
+          await instance.chat({
+            messages: [{ content: 'Hello', role: 'user' }],
+            model: 'mistralai/mistral-7b-instruct:free',
+            temperature: 0,
+          });
+        } catch (e) {
+          expect(e).toEqual({
+            endpoint: defaultBaseURL,
+            error: {
+              cause: { message: 'api is undefined' },
+              stack: 'abc',
+            },
+            errorType: bizErrorType,
+            provider,
+          });
+        }
+      });
+
+      it('should return TogetherAIBizError with an cause response with desensitize Url', async () => {
+        // Arrange
+        const errorInfo = {
+          stack: 'abc',
+          cause: { message: 'api is undefined' },
+        };
+        const apiError = new OpenAI.APIError(400, errorInfo, 'module error', {});
+
+        instance = new LobeTogetherAI({
+          apiKey: 'test',
+
+          baseURL: 'https://api.abc.com/v1',
+        });
+
+        vi.spyOn(instance['client'].chat.completions, 'create').mockRejectedValue(apiError);
+
+        // Act
+        try {
+          await instance.chat({
+            messages: [{ content: 'Hello', role: 'user' }],
+            model: 'mistralai/mistral-7b-instruct:free',
+            temperature: 0,
+          });
+        } catch (e) {
+          expect(e).toEqual({
+            endpoint: 'https://api.***.com/v1',
+            error: {
+              cause: { message: 'api is undefined' },
+              stack: 'abc',
+            },
+            errorType: bizErrorType,
+            provider,
+          });
+        }
+      });
+
+      it('should throw an InvalidTogetherAIAPIKey error type on 401 status code', async () => {
+        // Mock the API call to simulate a 401 error
+        const error = new Error('Unauthorized') as any;
+        error.status = 401;
+        vi.mocked(instance['client'].chat.completions.create).mockRejectedValue(error);
+
+        try {
+          await instance.chat({
+            messages: [{ content: 'Hello', role: 'user' }],
+            model: 'mistralai/mistral-7b-instruct:free',
+            temperature: 0,
+          });
+        } catch (e) {
+          // Expect the chat method to throw an error with InvalidTogetherAIAPIKey
+          expect(e).toEqual({
+            endpoint: defaultBaseURL,
+            error: new Error('Unauthorized'),
+            errorType: invalidErrorType,
+            provider,
+          });
+        }
+      });
+
+      it('should return AgentRuntimeError for non-OpenAI errors', async () => {
+        // Arrange
+        const genericError = new Error('Generic Error');
+
+        vi.spyOn(instance['client'].chat.completions, 'create').mockRejectedValue(genericError);
+
+        // Act
+        try {
+          await instance.chat({
+            messages: [{ content: 'Hello', role: 'user' }],
+            model: 'mistralai/mistral-7b-instruct:free',
+            temperature: 0,
+          });
+        } catch (e) {
+          expect(e).toEqual({
+            endpoint: defaultBaseURL,
+            errorType: 'AgentRuntimeError',
+            provider,
+            error: {
+              name: genericError.name,
+              cause: genericError.cause,
+              message: genericError.message,
+              stack: genericError.stack,
+            },
+          });
+        }
+      });
+    });
+
+    describe('LobeTogetherAI chat with callback and headers', () => {
+      it('should handle callback and headers correctly', async () => {
+        // 模拟 chat.completions.create 方法返回一个可读流
+        const mockCreateMethod = vi
+          .spyOn(instance['client'].chat.completions, 'create')
+          .mockResolvedValue(
+            new ReadableStream({
+              start(controller) {
+                controller.enqueue({
+                  id: 'chatcmpl-8xDx5AETP8mESQN7UB30GxTN2H1SO',
+                  object: 'chat.completion.chunk',
+                  created: 1709125675,
+                  model: 'mistralai/mistral-7b-instruct:free',
+                  system_fingerprint: 'fp_86156a94a0',
+                  choices: [
+                    { index: 0, delta: { content: 'hello' }, logprobs: null, finish_reason: null },
+                  ],
+                });
+                controller.close();
+              },
+            }) as any,
+          );
+
+        // 准备 callback 和 headers
+        const mockCallback: ChatStreamCallbacks = {
+          onStart: vi.fn(),
+          onToken: vi.fn(),
+        };
+        const mockHeaders = { 'Custom-Header': 'TestValue' };
+
+        // 执行测试
+        const result = await instance.chat(
+          {
+            messages: [{ content: 'Hello', role: 'user' }],
+            model: 'mistralai/mistral-7b-instruct:free',
+            temperature: 0,
+          },
+          { callback: mockCallback, headers: mockHeaders },
+        );
+
+        // 验证 callback 被调用
+        await result.text(); // 确保流被消费
+        expect(mockCallback.onStart).toHaveBeenCalled();
+        expect(mockCallback.onToken).toHaveBeenCalledWith('hello');
+
+        // 验证 headers 被正确传递
+        expect(result.headers.get('Custom-Header')).toEqual('TestValue');
+
+        // 清理
+        mockCreateMethod.mockRestore();
+      });
+    });
+
+    describe('DEBUG', () => {
+      it('should call debugStream and return StreamingTextResponse when DEBUG_TOGETHERAI_CHAT_COMPLETION is 1', async () => {
+        // Arrange
+        const mockProdStream = new ReadableStream() as any; // 模拟的 prod 流
+        const mockDebugStream = new ReadableStream({
+          start(controller) {
+            controller.enqueue('Debug stream content');
+            controller.close();
+          },
+        }) as any;
+        mockDebugStream.toReadableStream = () => mockDebugStream; // 添加 toReadableStream 方法
+
+        // 模拟 chat.completions.create 返回值，包括模拟的 tee 方法
+        (instance['client'].chat.completions.create as Mock).mockResolvedValue({
+          tee: () => [mockProdStream, { toReadableStream: () => mockDebugStream }],
+        });
+
+        // 保存原始环境变量值
+        const originalDebugValue = process.env.DEBUG_TOGETHERAI_CHAT_COMPLETION;
+
+        // 模拟环境变量
+        process.env.DEBUG_TOGETHERAI_CHAT_COMPLETION = '1';
+        vi.spyOn(debugStreamModule, 'debugStream').mockImplementation(() => Promise.resolve());
+
+        // 执行测试
+        // 运行你的测试函数，确保它会在条件满足时调用 debugStream
+        // 假设的测试函数调用，你可能需要根据实际情况调整
+        await instance.chat({
+          messages: [{ content: 'Hello', role: 'user' }],
+          model: 'mistralai/mistral-7b-instruct:free',
+          temperature: 0,
+        });
+
+        // 验证 debugStream 被调用
+        expect(debugStreamModule.debugStream).toHaveBeenCalled();
+
+        // 恢复原始环境变量值
+        process.env.DEBUG_TOGETHERAI_CHAT_COMPLETION = originalDebugValue;
+      });
+    });
+  });
+});

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

@@ -0,0 +1,86 @@
+import { OpenAIStream, StreamingTextResponse } from 'ai';
+import OpenAI, { ClientOptions } from 'openai';
+
+import { LobeRuntimeAI } from '../BaseAI';
+import { AgentRuntimeErrorType } from '../error';
+import { ChatCompetitionOptions, ChatStreamPayload, ModelProvider } from '../types';
+import { AgentRuntimeError } from '../utils/createError';
+import { debugStream } from '../utils/debugStream';
+import { desensitizeUrl } from '../utils/desensitizeUrl';
+import { handleOpenAIError } from '../utils/handleOpenAIError';
+
+const DEFAULT_BASE_URL = 'https://api.together.xyz/v1';
+
+export class LobeTogetherAI implements LobeRuntimeAI {
+  private client: OpenAI;
+
+  baseURL: string;
+
+  constructor({ apiKey, baseURL = DEFAULT_BASE_URL, ...res }: ClientOptions) {
+    if (!apiKey) throw AgentRuntimeError.createError(AgentRuntimeErrorType.InvalidTogetherAIAPIKey);
+
+    this.client = new OpenAI({
+      apiKey,
+      baseURL,
+      defaultHeaders: {
+        'HTTP-Referer': 'https://chat-preview.lobehub.com',
+        'X-Title': 'Lobe Chat',
+      },
+      ...res,
+    });
+    this.baseURL = this.client.baseURL;
+  }
+
+  async chat(payload: ChatStreamPayload, options?: ChatCompetitionOptions) {
+    try {
+      const response = await this.client.chat.completions.create(
+        payload as unknown as OpenAI.ChatCompletionCreateParamsStreaming,
+      );
+      const [prod, debug] = response.tee();
+
+      if (process.env.DEBUG_TOGETHERAI_CHAT_COMPLETION === '1') {
+        debugStream(debug.toReadableStream()).catch(console.error);
+      }
+
+      return new StreamingTextResponse(OpenAIStream(prod, options?.callback), {
+        headers: options?.headers,
+      });
+    } 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.InvalidTogetherAIAPIKey,
+              provider: ModelProvider.TogetherAI,
+            });
+          }
+
+          default: {
+            break;
+          }
+        }
+      }
+
+      const { errorResult, RuntimeError } = handleOpenAIError(error);
+
+      const errorType = RuntimeError || AgentRuntimeErrorType.TogetherAIBizError;
+
+      throw AgentRuntimeError.chat({
+        endpoint: desensitizedEndpoint,
+        error: errorResult,
+        errorType,
+        provider: ModelProvider.TogetherAI,
+      });
+    }
+  }
+}
+
+export default LobeTogetherAI;

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

@@ -127,3 +127,10 @@ export interface ChatCompletionTool {
 }
 
 export type ChatStreamCallbacks = OpenAIStreamCallbacks;
+
+export interface OllamaChatMessage extends OpenAIChatMessage {
+  /**
+   * @description images for ollama vision models (https://ollama.com/blog/vision-models)
+   */
+  images?: string[];
+}

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

@@ -34,6 +34,7 @@ export enum ModelProvider {
   OpenAI = 'openai',
   OpenRouter = 'openrouter',
   Perplexity = 'perplexity',
+  TogetherAI = 'togetherai',
   Tongyi = 'tongyi',
   ZeroOne = 'zeroone',
   ZhiPu = 'zhipu',

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

@@ -26,7 +26,7 @@ export class LobeZeroOneAI implements LobeRuntimeAI {
   async chat(payload: ChatStreamPayload, options?: ChatCompetitionOptions) {
     try {
       const response = await this.client.chat.completions.create(
-        payload as unknown as OpenAI.ChatCompletionCreateParamsStreaming
+        payload as unknown as OpenAI.ChatCompletionCreateParamsStreaming,
       );
       const [prod, debug] = response.tee();
 

package/src/locales/default/common.ts

@@ -110,6 +110,7 @@ export default {
     openai: 'OpenAI',
     openrouter: 'OpenRouter',
     perplexity: 'Perplexity',
+    togetherai: 'TogetherAI',
     zeroone: '01.AI 零一万物',
     zhipu: '智谱AI',
   },

package/src/locales/default/error.ts

@@ -88,6 +88,9 @@ export default {
     InvalidOpenRouterAPIKey: 'OpenRouter API Key 不正确或为空，请检查 OpenRouter API Key 后重试',
     OpenRouterBizError: '请求 OpenRouter AI 服务出错，请根据以下信息排查或重试',
 
+    InvalidTogetherAIAPIKey: 'TogetherAI API Key 不正确或为空，请检查 TogetherAI API Key 后重试',
+    TogetherAIBizError: '请求 TogetherAI AI 服务出错，请根据以下信息排查或重试',
+
     ZeroOneBizError: '请求零一万物服务出错，请根据以下信息排查或重试',
     InvalidZeroOneAPIKey: '零一万物 API Key 不正确或为空，请检查零一万物 API Key 后重试',
 
@@ -145,6 +148,10 @@ export default {
         description: '输入你的 Perplexity API Key 即可开始会话。应用不会记录你的 API Key',
         title: '使用自定义 Perplexity API Key',
       },
+      TogetherAI: {
+        description: '输入你的 TogetherAI API Key 即可开始会话。应用不会记录你的 API Key',
+        title: '使用自定义 TogetherAI API Key',
+      },
       ZeroOne: {
         description: '输入你的零一万物 API Key 即可开始会话。应用不会记录你的 API Key',
         title: '使用自定义零一万物 API Key',

package/src/locales/default/setting.ts

@@ -202,13 +202,29 @@ export default {
         title: 'API Key',
       },
     },
+    TogetherAI: {
+      checker: {
+        desc: '测试代理地址是否正确填写',
+      },
+      customModelName: {
+        desc: '增加自定义模型，多个模型使用逗号（,）隔开',
+        placeholder: 'togethercomputer/Llama-2-7B-32K-Instruct,teknium/OpenHermes-2-Mistral-7B',
+        title: '自定义模型名称',
+      },
+      title: 'TogetherAI',
+      token: {
+        desc: '填入来自 TogetherAI AI 的 API Key',
+        placeholder: 'TogetherAI AI API Key',
+        title: 'API Key',
+      },
+    },
     ZeroOne: {
       title: '01.AI 零一万物',
       token: {
         desc: '填入来自 01.AI 零一万物的 API Key',
         placeholder: '01.AI 零一万物 API Key',
         title: 'API Key',
-      }
+      },
     },
     Zhipu: {
       title: '智谱',

package/src/services/_auth.test.ts

@@ -14,6 +14,7 @@ const mockGoogleAPIKey = 'google-api-key';
 const mockAnthropicAPIKey = 'anthropic-api-key';
 const mockMistralAPIKey = 'mistral-api-key';
 const mockOpenRouterAPIKey = 'openrouter-api-key';
+const mockTogetherAIAPIKey = 'togetherai-api-key';
 
 // mock the traditional zustand
 vi.mock('zustand/traditional');
@@ -73,6 +74,15 @@ describe('getProviderAuthPayload', () => {
     expect(payload).toEqual({ apiKey: mockOpenRouterAPIKey });
   });
 
+  it('should return correct payload for TogetherAI provider', () => {
+    act(() => {
+      setModelProviderConfig('togetherai', { apiKey: mockTogetherAIAPIKey });
+    });
+
+    const payload = getProviderAuthPayload(ModelProvider.TogetherAI);
+    expect(payload).toEqual({ apiKey: mockTogetherAIAPIKey });
+  });
+
   it('should return correct payload for Google provider', () => {
     act(() => {
       setModelProviderConfig('google', { apiKey: mockGoogleAPIKey });

package/src/services/_auth.ts

@@ -70,6 +70,10 @@ export const getProviderAuthPayload = (provider: string) => {
       return { apiKey: modelProviderSelectors.openrouterAPIKey(useGlobalStore.getState()) };
     }
 
+    case ModelProvider.TogetherAI: {
+      return { apiKey: modelProviderSelectors.togetheraiAPIKey(useGlobalStore.getState()) };
+    }
+
     case ModelProvider.ZeroOne: {
       return { apiKey: modelProviderSelectors.zerooneAPIKey(useGlobalStore.getState()) };
     }

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

@@ -139,7 +139,7 @@ describe('modelProviderSelectors', () => {
             openrouter: {
               enabled: true,
             },
-          }
+          },
         },
         serverConfig: {
           languageModel: {
@@ -148,7 +148,7 @@ describe('modelProviderSelectors', () => {
               customModelName:
                 '-all,+google/gemma-7b-it,+mistralai/mistral-7b-instruct=Mistral-7B-Instruct',
             },
-          }
+          },
         },
       }) as unknown as GlobalStore;
 

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

@@ -12,6 +12,7 @@ import {
   OpenAIProvider,
   OpenRouterProvider,
   PerplexityProvider,
+  TogetherAIProvider,
   ZeroOneProvider,
   ZhiPuProvider,
 } from '@/config/modelProviders';
@@ -67,6 +68,9 @@ const groqAPIKey = (s: GlobalStore) => modelProvider(s).groq.apiKey;
 const enableOpenrouter = (s: GlobalStore) => modelProvider(s).openrouter.enabled;
 const openrouterAPIKey = (s: GlobalStore) => modelProvider(s).openrouter.apiKey;
 
+const enableTogetherAI = (s: GlobalStore) => modelProvider(s).togetherai.enabled;
+const togetheraiAPIKey = (s: GlobalStore) => modelProvider(s).togetherai.apiKey;
+
 const enableZeroone = (s: GlobalStore) => modelProvider(s).zeroone.enabled;
 const zerooneAPIKey = (s: GlobalStore) => modelProvider(s).zeroone.apiKey;
 
@@ -151,14 +155,25 @@ const modelSelectList = (s: GlobalStore): ModelProviderCard[] => {
 
   const openrouterModelString = [
     s.serverConfig.languageModel?.openrouter?.customModelName,
-    currentSettings(s).languageModel.openrouter.customModelName
+    currentSettings(s).languageModel.openrouter.customModelName,
   ]
     .filter(Boolean)
     .join(',');
-  
+
   const openrouterModelConfig = parseModelString(openrouterModelString);
-  
-  const openrouterChatModels = processChatModels(openrouterModelConfig, OpenRouterProvider.chatModels);
+
+  const openrouterChatModels = processChatModels(
+    openrouterModelConfig,
+    OpenRouterProvider.chatModels,
+  );
+
+  const togetheraiModelConfig = parseModelString(
+    currentSettings(s).languageModel.togetherai.customModelName,
+  );
+  const togetheraiChatModels = processChatModels(
+    togetheraiModelConfig,
+    TogetherAIProvider.chatModels,
+  );
 
   return [
     {
@@ -175,8 +190,9 @@ const modelSelectList = (s: GlobalStore): ModelProviderCard[] => {
     { ...GroqProvider, enabled: enableGroq(s) },
     { ...ZhiPuProvider, enabled: enableZhipu(s) },
     { ...MoonshotProvider, enabled: enableMoonshot(s) },
-    { ...OpenRouterProvider, chatModels: openrouterChatModels, enabled: enableOpenrouter(s)},
-    { ...ZeroOneProvider, enabled: enableZeroone(s) }
+    { ...OpenRouterProvider, chatModels: openrouterChatModels, enabled: enableOpenrouter(s) },
+    { ...ZeroOneProvider, enabled: enableZeroone(s) },
+    { ...TogetherAIProvider, chatModels: togetheraiChatModels, enabled: enableTogetherAI(s) },
   ];
 };
 
@@ -271,4 +287,8 @@ export const modelProviderSelectors = {
   // ZeroOne 零一万物
   enableZeroone,
   zerooneAPIKey,
+
+  // TogetherAI
+  enableTogetherAI,
+  togetheraiAPIKey,
 };

package/src/types/settings/modelProvider.ts

@@ -88,6 +88,12 @@ export interface ZeroOneConfig {
   enabled?: boolean;
 }
 
+export interface TogetherAIConfig {
+  apiKey?: string;
+  customModelName?: string;
+  enabled?: boolean;
+}
+
 export interface GlobalLLMConfig {
   anthropic: AnthropicConfig;
   azure: AzureOpenAIConfig;
@@ -100,6 +106,7 @@ export interface GlobalLLMConfig {
   openAI: OpenAIConfig;
   openrouter: OpenRouterConfig;
   perplexity: PerplexityConfig;
+  togetherai: TogetherAIConfig;
   zeroone: ZeroOneConfig;
   zhipu: ZhiPuConfig;
 }