@lobehub/chat 0.138.2 → 0.139.1

package/src/libs/agent-runtime/openrouter/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 { LobeOpenRouterAI } from './index';
+
+const provider = 'openrouter';
+const defaultBaseURL = 'https://openrouter.ai/api/v1';
+const bizErrorType = 'OpenRouterBizError';
+const invalidErrorType = 'InvalidOpenRouterAPIKey';
+
+// Mock the console.error to avoid polluting test output
+vi.spyOn(console, 'error').mockImplementation(() => {});
+
+let instance: LobeOpenRouterAI;
+
+beforeEach(() => {
+  instance = new LobeOpenRouterAI({ apiKey: 'test' });
+
+  // 使用 vi.spyOn 来模拟 chat.completions.create 方法
+  vi.spyOn(instance['client'].chat.completions, 'create').mockResolvedValue(
+    new ReadableStream() as any,
+  );
+});
+
+afterEach(() => {
+  vi.clearAllMocks();
+});
+
+describe('LobeOpenRouterAI', () => {
+  describe('init', () => {
+    it('should correctly initialize with an API key', async () => {
+      const instance = new LobeOpenRouterAI({ apiKey: 'test_api_key' });
+      expect(instance).toBeInstanceOf(LobeOpenRouterAI);
+      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 OpenRouter 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 OpenRouterBizError 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 InvalidOpenRouterAPIKey if no apiKey is provided', async () => {
+        try {
+          new LobeOpenRouterAI({});
+        } catch (e) {
+          expect(e).toEqual({ errorType: invalidErrorType });
+        }
+      });
+
+      it('should return OpenRouterBizError 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 OpenRouterBizError 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 LobeOpenRouterAI({
+          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 InvalidOpenRouterAPIKey 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 InvalidMoonshotAPIKey
+          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('LobeOpenRouterAI 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_OPENROUTER_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_OPENROUTER_CHAT_COMPLETION;
+
+        // 模拟环境变量
+        process.env.DEBUG_OPENROUTER_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_OPENROUTER_CHAT_COMPLETION = originalDebugValue;
+      });
+    });
+  });
+});

package/src/libs/agent-runtime/openrouter/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://openrouter.ai/api/v1';
+
+export class LobeOpenRouterAI implements LobeRuntimeAI {
+  private client: OpenAI;
+
+  baseURL: string;
+
+  constructor({ apiKey, baseURL = DEFAULT_BASE_URL, ...res }: ClientOptions) {
+    if (!apiKey) throw AgentRuntimeError.createError(AgentRuntimeErrorType.InvalidOpenRouterAPIKey);
+    
+    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_OPENROUTER_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.InvalidOpenRouterAPIKey,
+              provider: ModelProvider.OpenRouter,
+            });
+          }
+
+          default: {
+            break;
+          }
+        }
+      }
+
+      const { errorResult, RuntimeError } = handleOpenAIError(error);
+
+      const errorType = RuntimeError || AgentRuntimeErrorType.OpenRouterBizError;
+
+      throw AgentRuntimeError.chat({
+        endpoint: desensitizedEndpoint,
+        error: errorResult,
+        errorType,
+        provider: ModelProvider.OpenRouter,
+      });
+    }
+  }
+}
+
+export default LobeOpenRouterAI;

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

@@ -32,6 +32,7 @@ export enum ModelProvider {
   Moonshot = 'moonshot',
   Ollama = 'ollama',
   OpenAI = 'openai',
+  OpenRouter = 'openrouter',
   Perplexity = 'perplexity',
   Tongyi = 'tongyi',
   ZhiPu = 'zhipu',

package/src/locales/default/common.ts

@@ -110,6 +110,7 @@ export default {
     ollama: 'Ollama',
     oneapi: 'One API',
     openai: 'OpenAI',
+    openrouter: 'OpenRouter',
     perplexity: 'Perplexity',
     zhipu: '智谱AI',
   },

package/src/locales/default/error.ts

@@ -84,6 +84,9 @@ export default {
 
     InvalidGroqAPIKey: 'Groq API Key 不正确或为空，请检查 Groq API Key 后重试',
     GroqBizError: '请求 Groq 服务出错，请根据以下信息排查或重试',
+    
+    InvalidOpenRouterAPIKey: 'OpenRouter API Key 不正确或为空，请检查 OpenRouter API Key 后重试',
+    OpenRouterBizError: '请求 OpenRouter AI 服务出错，请根据以下信息排查或重试',
 
     InvalidOllamaArgs: 'Ollama 配置不正确，请检查 Ollama 配置后重试',
     OllamaBizError: '请求 Ollama 服务出错，请根据以下信息排查或重试',
@@ -131,6 +134,10 @@ export default {
         description: '输入你的 OpenAI API Key 即可开始会话。应用不会记录你的 API Key',
         title: '使用自定义 OpenAI API Key',
       },
+      OpenRouter: {
+        description: '输入你的 OpenRouter API Key 即可开始会话。应用不会记录你的 API Key',
+        title: '使用自定义 OpenRouter API Key',
+      },
       Perplexity: {
         description: '输入你的 Perplexity API Key 即可开始会话。应用不会记录你的 API Key',
         title: '使用自定义 Perplexity API Key',

package/src/locales/default/setting.ts

@@ -178,6 +178,22 @@ export default {
         title: 'Azure OpenAI',
       },
     },
+    OpenRouter: {
+      checker: {
+        desc: '测试代理地址是否正确填写',
+      },
+      customModelName: {
+        desc: '增加自定义模型，多个模型使用逗号（,）隔开',
+        placeholder: 'mistralai/mistral-7b-instruct:free,openchat/openchat-7b:free',
+        title: '自定义模型名称',
+      },
+      title: 'OpenRouter',
+      token: {
+        desc: '填入来自 OpenRouter AI 的 API Key',
+        placeholder: 'OpenRouter AI API Key',
+        title: 'API Key',
+      },
+    },
     Perplexity: {
       title: 'Perplexity',
       token: {

package/src/services/_auth.test.ts

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

package/src/services/_auth.ts

@@ -65,6 +65,10 @@ export const getProviderAuthPayload = (provider: string) => {
     case ModelProvider.Groq: {
       return { apiKey: modelProviderSelectors.groqAPIKey(useGlobalStore.getState()) };
     }
+    
+    case ModelProvider.OpenRouter: {
+      return { apiKey: modelProviderSelectors.openrouterAPIKey(useGlobalStore.getState()) };
+    }
 
     default:
     case ModelProvider.OpenAI: {

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

@@ -10,6 +10,7 @@ import {
   MoonshotProvider,
   OllamaProvider,
   OpenAIProvider,
+  OpenRouterProvider,
   PerplexityProvider,
   ZhiPuProvider,
 } from '@/config/modelProviders';
@@ -62,6 +63,9 @@ const anthropicProxyUrl = (s: GlobalStore) => modelProvider(s).anthropic.endpoin
 const enableGroq = (s: GlobalStore) => modelProvider(s).groq.enabled;
 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 azureModelList = (s: GlobalStore): ModelProviderCard => {
 //   const azure = azureConfig(s);
 //   return {
@@ -141,6 +145,11 @@ const modelSelectList = (s: GlobalStore): ModelProviderCard[] => {
 
   const ollamaChatModels = processChatModels(ollamaModelConfig, OllamaProvider.chatModels);
 
+  const openrouterModelConfig = parseModelString(
+    currentSettings(s).languageModel.openrouter.customModelName,
+  )
+  const openrouterChatModels = processChatModels(openrouterModelConfig, OpenRouterProvider.chatModels);
+
   return [
     {
       ...OpenAIProvider,
@@ -156,6 +165,7 @@ const modelSelectList = (s: GlobalStore): ModelProviderCard[] => {
     { ...GroqProvider, enabled: enableGroq(s) },
     { ...ZhiPuProvider, enabled: enableZhipu(s) },
     { ...MoonshotProvider, enabled: enableMoonshot(s) },
+    { ...OpenRouterProvider, chatModels: openrouterChatModels, enabled: enableOpenrouter(s)}
   ];
 };
 
@@ -242,4 +252,8 @@ export const modelProviderSelectors = {
   // Groq
   enableGroq,
   groqAPIKey,
+  
+  // OpenRouter
+  enableOpenrouter,
+  openrouterAPIKey,
 };

package/src/types/settings/modelProvider.ts

@@ -76,6 +76,12 @@ export interface GroqConfig {
   enabled: boolean;
 }
 
+export interface OpenRouterConfig {
+  apiKey?: string;
+  customModelName?: string;
+  enabled?: boolean;
+}
+
 export interface GlobalLLMConfig {
   anthropic: AnthropicConfig;
   azure: AzureOpenAIConfig;
@@ -86,6 +92,7 @@ export interface GlobalLLMConfig {
   moonshot: MoonshotConfig;
   ollama: OllamaConfig;
   openAI: OpenAIConfig;
+  openrouter: OpenRouterConfig;
   perplexity: PerplexityConfig;
   zhipu: ZhiPuConfig;
 }