@lobehub/chat 1.108.1 → 1.109.0

package/src/libs/model-runtime/RouterRuntime/createRuntime.test.ts

@@ -0,0 +1,538 @@
+import { beforeEach, describe, expect, it, vi } from 'vitest';
+
+import { LobeRuntimeAI } from '../BaseAI';
+import { createRouterRuntime } from './createRuntime';
+
+describe('createRouterRuntime', () => {
+  beforeEach(() => {
+    vi.resetModules();
+  });
+
+  describe('initialization', () => {
+    it('should throw error when routers array is empty', () => {
+      expect(() => {
+        const Runtime = createRouterRuntime({
+          id: 'test-runtime',
+          routers: [],
+        });
+        new Runtime();
+      }).toThrow('empty providers');
+    });
+
+    it('should create UniformRuntime class with valid routers', () => {
+      class MockRuntime implements LobeRuntimeAI {
+        chat = vi.fn();
+        textToImage = vi.fn();
+        models = vi.fn();
+        embeddings = vi.fn();
+        textToSpeech = vi.fn();
+      }
+
+      const Runtime = createRouterRuntime({
+        id: 'test-runtime',
+        routers: [
+          {
+            apiType: 'openai',
+            options: { apiKey: 'test-key' },
+            runtime: MockRuntime as any,
+            models: ['gpt-4', 'gpt-3.5-turbo'],
+          },
+        ],
+      });
+
+      const runtime = new Runtime();
+      expect(runtime).toBeDefined();
+    });
+
+    it('should merge router options with constructor options', () => {
+      const mockConstructor = vi.fn();
+
+      class MockRuntime implements LobeRuntimeAI {
+        constructor(options: any) {
+          mockConstructor(options);
+        }
+        chat = vi.fn();
+      }
+
+      const Runtime = createRouterRuntime({
+        id: 'test-runtime',
+        routers: [
+          {
+            apiType: 'openai',
+            options: { baseURL: 'https://api.example.com' },
+            runtime: MockRuntime as any,
+          },
+        ],
+      });
+
+      new Runtime({ apiKey: 'constructor-key' });
+
+      expect(mockConstructor).toHaveBeenCalledWith(
+        expect.objectContaining({
+          baseURL: 'https://api.example.com',
+          apiKey: 'constructor-key',
+          id: 'test-runtime',
+        }),
+      );
+    });
+  });
+
+  describe('getModels', () => {
+    it('should return synchronous models array directly', async () => {
+      const mockRuntime = {
+        chat: vi.fn(),
+      } as unknown as LobeRuntimeAI;
+
+      const Runtime = createRouterRuntime({
+        id: 'test-runtime',
+        routers: [
+          {
+            apiType: 'openai',
+            options: {},
+            runtime: mockRuntime.constructor as any,
+            models: ['model-1', 'model-2'],
+          },
+        ],
+      });
+
+      const runtime = new Runtime();
+      const models = await runtime['getModels']({
+        id: 'test',
+        models: ['model-1', 'model-2'],
+        runtime: mockRuntime,
+      });
+
+      expect(models).toEqual(['model-1', 'model-2']);
+    });
+
+    it('should call and cache asynchronous models function', async () => {
+      const mockRuntime = {
+        chat: vi.fn(),
+      } as unknown as LobeRuntimeAI;
+
+      const mockModelsFunction = vi.fn().mockResolvedValue(['async-model-1', 'async-model-2']);
+
+      const Runtime = createRouterRuntime({
+        id: 'test-runtime',
+        routers: [
+          {
+            apiType: 'openai',
+            options: {},
+            runtime: mockRuntime.constructor as any,
+            models: mockModelsFunction,
+          },
+        ],
+      });
+
+      const runtime = new Runtime();
+      const runtimeItem = {
+        id: 'test',
+        models: mockModelsFunction,
+        runtime: mockRuntime,
+      };
+
+      // First call
+      const models1 = await runtime['getModels'](runtimeItem);
+      expect(models1).toEqual(['async-model-1', 'async-model-2']);
+      expect(mockModelsFunction).toHaveBeenCalledTimes(1);
+
+      // Second call should use cache
+      const models2 = await runtime['getModels'](runtimeItem);
+      expect(models2).toEqual(['async-model-1', 'async-model-2']);
+      expect(mockModelsFunction).toHaveBeenCalledTimes(1);
+    });
+
+    it('should return empty array when models is undefined', async () => {
+      const mockRuntime = {
+        chat: vi.fn(),
+      } as unknown as LobeRuntimeAI;
+
+      const Runtime = createRouterRuntime({
+        id: 'test-runtime',
+        routers: [
+          {
+            apiType: 'openai',
+            options: {},
+            runtime: mockRuntime.constructor as any,
+          },
+        ],
+      });
+
+      const runtime = new Runtime();
+      const models = await runtime['getModels']({
+        id: 'test',
+        runtime: mockRuntime,
+      });
+
+      expect(models).toEqual([]);
+    });
+  });
+
+  describe('getRuntimeByModel', () => {
+    it('should return runtime that supports the model', async () => {
+      class MockRuntime1 implements LobeRuntimeAI {
+        chat = vi.fn();
+      }
+
+      class MockRuntime2 implements LobeRuntimeAI {
+        chat = vi.fn();
+      }
+
+      const Runtime = createRouterRuntime({
+        id: 'test-runtime',
+        routers: [
+          {
+            apiType: 'openai',
+            options: {},
+            runtime: MockRuntime1 as any,
+            models: ['gpt-4'],
+          },
+          {
+            apiType: 'anthropic',
+            options: {},
+            runtime: MockRuntime2 as any,
+            models: ['claude-3'],
+          },
+        ],
+      });
+
+      const runtime = new Runtime();
+
+      const result1 = await runtime.getRuntimeByModel('gpt-4');
+      expect(result1).toBe(runtime['_runtimes'][0].runtime);
+
+      const result2 = await runtime.getRuntimeByModel('claude-3');
+      expect(result2).toBe(runtime['_runtimes'][1].runtime);
+    });
+
+    it('should return last runtime when no model matches', async () => {
+      class MockRuntime1 implements LobeRuntimeAI {
+        chat = vi.fn();
+      }
+
+      class MockRuntime2 implements LobeRuntimeAI {
+        chat = vi.fn();
+      }
+
+      const Runtime = createRouterRuntime({
+        id: 'test-runtime',
+        routers: [
+          {
+            apiType: 'openai',
+            options: {},
+            runtime: MockRuntime1 as any,
+            models: ['gpt-4'],
+          },
+          {
+            apiType: 'anthropic',
+            options: {},
+            runtime: MockRuntime2 as any,
+            models: ['claude-3'],
+          },
+        ],
+      });
+
+      const runtime = new Runtime();
+      const result = await runtime.getRuntimeByModel('unknown-model');
+
+      expect(result).toBe(runtime['_runtimes'][1].runtime);
+    });
+  });
+
+  describe('chat method', () => {
+    it('should call chat on the correct runtime based on model', async () => {
+      const mockChat = vi.fn().mockResolvedValue('chat-response');
+
+      class MockRuntime implements LobeRuntimeAI {
+        chat = mockChat;
+      }
+
+      const Runtime = createRouterRuntime({
+        id: 'test-runtime',
+        routers: [
+          {
+            apiType: 'openai',
+            options: {},
+            runtime: MockRuntime as any,
+            models: ['gpt-4'],
+          },
+        ],
+      });
+
+      const runtime = new Runtime();
+      const payload = { model: 'gpt-4', messages: [], temperature: 0.7 };
+
+      const result = await runtime.chat(payload);
+      expect(result).toBe('chat-response');
+      expect(mockChat).toHaveBeenCalledWith(payload, undefined);
+    });
+
+    it('should handle errors when provided with handleError', async () => {
+      const mockError = new Error('API Error');
+      const mockChat = vi.fn().mockRejectedValue(mockError);
+
+      class MockRuntime implements LobeRuntimeAI {
+        chat = mockChat;
+      }
+
+      const handleError = vi.fn().mockReturnValue({
+        errorType: 'APIError',
+        message: 'Handled error',
+      });
+
+      const Runtime = createRouterRuntime({
+        id: 'test-runtime',
+        routers: [
+          {
+            apiType: 'openai',
+            options: {},
+            runtime: MockRuntime as any,
+            models: ['gpt-4'],
+          },
+        ],
+      });
+
+      const runtime = new Runtime({
+        chat: {
+          handleError,
+        },
+      });
+
+      await expect(
+        runtime.chat({ model: 'gpt-4', messages: [], temperature: 0.7 }),
+      ).rejects.toEqual({
+        errorType: 'APIError',
+        message: 'Handled error',
+      });
+    });
+
+    it('should re-throw original error when handleError returns undefined', async () => {
+      const mockError = new Error('API Error');
+      const mockChat = vi.fn().mockRejectedValue(mockError);
+
+      class MockRuntime implements LobeRuntimeAI {
+        chat = mockChat;
+      }
+
+      const handleError = vi.fn().mockReturnValue(undefined);
+
+      const Runtime = createRouterRuntime({
+        id: 'test-runtime',
+        routers: [
+          {
+            apiType: 'openai',
+            options: {},
+            runtime: MockRuntime as any,
+            models: ['gpt-4'],
+          },
+        ],
+      });
+
+      const runtime = new Runtime({
+        chat: {
+          handleError,
+        },
+      });
+
+      await expect(runtime.chat({ model: 'gpt-4', messages: [], temperature: 0.7 })).rejects.toBe(
+        mockError,
+      );
+    });
+  });
+
+  describe('textToImage method', () => {
+    it('should call textToImage on the correct runtime based on model', async () => {
+      const mockTextToImage = vi.fn().mockResolvedValue('image-response');
+
+      class MockRuntime implements LobeRuntimeAI {
+        textToImage = mockTextToImage;
+      }
+
+      const Runtime = createRouterRuntime({
+        id: 'test-runtime',
+        routers: [
+          {
+            apiType: 'openai',
+            options: {},
+            runtime: MockRuntime as any,
+            models: ['dall-e-3'],
+          },
+        ],
+      });
+
+      const runtime = new Runtime();
+      const payload = { model: 'dall-e-3', prompt: 'test prompt' };
+
+      const result = await runtime.textToImage(payload);
+      expect(result).toBe('image-response');
+      expect(mockTextToImage).toHaveBeenCalledWith(payload);
+    });
+  });
+
+  describe('models method', () => {
+    it('should call models method on first runtime', async () => {
+      const mockModels = vi.fn().mockResolvedValue(['model-1', 'model-2']);
+
+      class MockRuntime implements LobeRuntimeAI {
+        models = mockModels;
+      }
+
+      const Runtime = createRouterRuntime({
+        id: 'test-runtime',
+        routers: [
+          {
+            apiType: 'openai',
+            options: {},
+            runtime: MockRuntime as any,
+          },
+        ],
+      });
+
+      const runtime = new Runtime();
+      const result = await runtime.models();
+
+      expect(result).toEqual(['model-1', 'model-2']);
+      expect(mockModels).toHaveBeenCalled();
+    });
+  });
+
+  describe('embeddings method', () => {
+    it('should call embeddings on the correct runtime based on model', async () => {
+      const mockEmbeddings = vi.fn().mockResolvedValue('embeddings-response');
+
+      class MockRuntime implements LobeRuntimeAI {
+        embeddings = mockEmbeddings;
+      }
+
+      const Runtime = createRouterRuntime({
+        id: 'test-runtime',
+        routers: [
+          {
+            apiType: 'openai',
+            options: {},
+            runtime: MockRuntime as any,
+            models: ['text-embedding-ada-002'],
+          },
+        ],
+      });
+
+      const runtime = new Runtime();
+      const payload = { model: 'text-embedding-ada-002', input: 'test input' };
+      const options = {} as any;
+
+      const result = await runtime.embeddings(payload, options);
+      expect(result).toBe('embeddings-response');
+      expect(mockEmbeddings).toHaveBeenCalledWith(payload, options);
+    });
+  });
+
+  describe('textToSpeech method', () => {
+    it('should call textToSpeech on the correct runtime based on model', async () => {
+      const mockTextToSpeech = vi.fn().mockResolvedValue('speech-response');
+
+      class MockRuntime implements LobeRuntimeAI {
+        textToSpeech = mockTextToSpeech;
+      }
+
+      const Runtime = createRouterRuntime({
+        id: 'test-runtime',
+        routers: [
+          {
+            apiType: 'openai',
+            options: {},
+            runtime: MockRuntime as any,
+            models: ['tts-1'],
+          },
+        ],
+      });
+
+      const runtime = new Runtime();
+      const payload = { model: 'tts-1', input: 'Hello world', voice: 'alloy' };
+      const options = {} as any;
+
+      const result = await runtime.textToSpeech(payload, options);
+      expect(result).toBe('speech-response');
+      expect(mockTextToSpeech).toHaveBeenCalledWith(payload, options);
+    });
+  });
+
+  describe('clearModelCache method', () => {
+    it('should clear specific runtime cache when runtimeId provided', async () => {
+      const mockModelsFunction = vi.fn().mockResolvedValue(['model-1']);
+
+      const Runtime = createRouterRuntime({
+        id: 'test-runtime',
+        routers: [
+          {
+            apiType: 'openai',
+            options: {},
+            runtime: vi.fn() as any,
+            models: mockModelsFunction,
+          },
+        ],
+      });
+
+      const runtime = new Runtime();
+      const runtimeItem = {
+        id: 'test-id',
+        models: mockModelsFunction,
+        runtime: {} as any,
+      };
+
+      // Build cache
+      await runtime['getModels'](runtimeItem);
+      expect(mockModelsFunction).toHaveBeenCalledTimes(1);
+
+      // Clear specific cache
+      runtime.clearModelCache('test-id');
+
+      // Should call function again
+      await runtime['getModels'](runtimeItem);
+      expect(mockModelsFunction).toHaveBeenCalledTimes(2);
+    });
+
+    it('should clear all cache when no runtimeId provided', async () => {
+      const mockModelsFunction1 = vi.fn().mockResolvedValue(['model-1']);
+      const mockModelsFunction2 = vi.fn().mockResolvedValue(['model-2']);
+
+      const Runtime = createRouterRuntime({
+        id: 'test-runtime',
+        routers: [
+          {
+            apiType: 'openai',
+            options: {},
+            runtime: vi.fn() as any,
+            models: mockModelsFunction1,
+          },
+        ],
+      });
+
+      const runtime = new Runtime();
+      const runtimeItem1 = {
+        id: 'test-id-1',
+        models: mockModelsFunction1,
+        runtime: {} as any,
+      };
+      const runtimeItem2 = {
+        id: 'test-id-2',
+        models: mockModelsFunction2,
+        runtime: {} as any,
+      };
+
+      // Build cache for both items
+      await runtime['getModels'](runtimeItem1);
+      await runtime['getModels'](runtimeItem2);
+      expect(mockModelsFunction1).toHaveBeenCalledTimes(1);
+      expect(mockModelsFunction2).toHaveBeenCalledTimes(1);
+
+      // Clear all cache
+      runtime.clearModelCache();
+
+      // Should call functions again
+      await runtime['getModels'](runtimeItem1);
+      await runtime['getModels'](runtimeItem2);
+      expect(mockModelsFunction1).toHaveBeenCalledTimes(2);
+      expect(mockModelsFunction2).toHaveBeenCalledTimes(2);
+    });
+  });
+});

package/src/libs/model-runtime/RouterRuntime/createRuntime.ts

@@ -28,7 +28,7 @@ import { baseRuntimeMap } from './baseRuntimeMap';
 
 export interface RuntimeItem {
   id: string;
-  models?: string[];
+  models?: string[] | (() => Promise<string[]>);
   runtime: LobeRuntimeAI;
 }
 
@@ -44,11 +44,13 @@ interface ProviderIniOptions extends Record<string, any> {
   sessionToken?: string;
 }
 
+export type RuntimeClass = typeof LobeOpenAI;
+
 interface RouterInstance {
   apiType: keyof typeof baseRuntimeMap;
-  models?: string[];
+  models?: string[] | (() => Promise<string[]>);
   options: ProviderIniOptions;
-  runtime?: typeof LobeOpenAI;
+  runtime?: RuntimeClass;
 }
 
 type ConstructorOptions<T extends Record<string, any> = any> = ClientOptions & T;
@@ -117,6 +119,7 @@ export const createRouterRuntime = ({
   return class UniformRuntime implements LobeRuntimeAI {
     private _runtimes: RuntimeItem[];
     private _options: ClientOptions & Record<string, any>;
+    private _modelCache = new Map<string, string[]>();
 
     constructor(options: ClientOptions & Record<string, any> = {}) {
       const _options = {
@@ -142,18 +145,40 @@ export const createRouterRuntime = ({
       this._options = _options;
     }
 
-    // 检查下是否能匹配到特定模型，否则默认使用最后一个 runtime
-    getRuntimeByModel(model: string) {
-      const runtimeItem =
-        this._runtimes.find((runtime) => runtime.models && runtime.models.includes(model)) ||
-        this._runtimes.at(-1)!;
+    // 获取 runtime 的 models 列表，支持同步数组和异步函数，带缓存机制
+    private async getModels(runtimeItem: RuntimeItem): Promise<string[]> {
+      const cacheKey = runtimeItem.id;
+
+      // 如果是同步数组，直接返回不需要缓存
+      if (typeof runtimeItem.models !== 'function') {
+        return runtimeItem.models || [];
+      }
+
+      // 检查缓存
+      if (this._modelCache.has(cacheKey)) {
+        return this._modelCache.get(cacheKey)!;
+      }
+
+      // 获取模型列表并缓存结果
+      const models = await runtimeItem.models();
+      this._modelCache.set(cacheKey, models);
+      return models;
+    }
 
-      return runtimeItem.runtime;
+    // 检查下是否能匹配到特定模型，否则默认使用最后一个 runtime
+    async getRuntimeByModel(model: string) {
+      for (const runtimeItem of this._runtimes) {
+        const models = await this.getModels(runtimeItem);
+        if (models.includes(model)) {
+          return runtimeItem.runtime;
+        }
+      }
+      return this._runtimes.at(-1)!.runtime;
     }
 
     async chat(payload: ChatStreamPayload, options?: ChatMethodOptions) {
       try {
-        const runtime = this.getRuntimeByModel(payload.model);
+        const runtime = await this.getRuntimeByModel(payload.model);
 
         return await runtime.chat!(payload, options);
       } catch (e) {
@@ -170,7 +195,7 @@ export const createRouterRuntime = ({
     }
 
     async textToImage(payload: TextToImagePayload) {
-      const runtime = this.getRuntimeByModel(payload.model);
+      const runtime = await this.getRuntimeByModel(payload.model);
 
       return runtime.textToImage!(payload);
     }
@@ -180,15 +205,27 @@ export const createRouterRuntime = ({
     }
 
     async embeddings(payload: EmbeddingsPayload, options?: EmbeddingsOptions) {
-      const runtime = this.getRuntimeByModel(payload.model);
+      const runtime = await this.getRuntimeByModel(payload.model);
 
       return runtime.embeddings!(payload, options);
     }
 
     async textToSpeech(payload: TextToSpeechPayload, options?: EmbeddingsOptions) {
-      const runtime = this.getRuntimeByModel(payload.model);
+      const runtime = await this.getRuntimeByModel(payload.model);
 
       return runtime.textToSpeech!(payload, options);
     }
+
+    /**
+     * 清除模型列表缓存，强制下次获取时重新加载
+     * @param runtimeId - 可选，指定清除特定 runtime 的缓存，不传则清除所有缓存
+     */
+    clearModelCache(runtimeId?: string) {
+      if (runtimeId) {
+        this._modelCache.delete(runtimeId);
+      } else {
+        this._modelCache.clear();
+      }
+    }
   };
 };

package/src/libs/model-runtime/RouterRuntime/index.ts

@@ -2,7 +2,7 @@ import { LobeRuntimeAI } from '../BaseAI';
 
 export interface RuntimeItem {
   id: string;
-  models?: string[];
+  models?: string[] | (() => Promise<string[]>);
   runtime: LobeRuntimeAI;
 }
 

package/src/libs/model-runtime/aihubmix/index.ts

@@ -1,6 +1,5 @@
 import urlJoin from 'url-join';
 
-import { LOBE_DEFAULT_MODEL_LIST } from '@/config/aiModels';
 import AiHubMixModels from '@/config/aiModels/aihubmix';
 import type { ChatModelCard } from '@/types/llm';
 
@@ -98,14 +97,20 @@ export const LobeAiHubMixAI = createRouterRuntime({
   routers: [
     {
       apiType: 'anthropic',
-      models: LOBE_DEFAULT_MODEL_LIST.map((m) => m.id).filter(
-        (id) => id.startsWith('claude') || id.startsWith('kimi-k2'),
-      ),
+      models: async () => {
+        const { LOBE_DEFAULT_MODEL_LIST } = await import('@/config/aiModels');
+        return LOBE_DEFAULT_MODEL_LIST.map((m) => m.id).filter(
+          (id) => id.startsWith('claude') || id.startsWith('kimi-k2'),
+        );
+      },
       options: { baseURL },
     },
     {
       apiType: 'google',
-      models: LOBE_DEFAULT_MODEL_LIST.map((m) => m.id).filter((id) => id.startsWith('gemini')),
+      models: async () => {
+        const { LOBE_DEFAULT_MODEL_LIST } = await import('@/config/aiModels');
+        return LOBE_DEFAULT_MODEL_LIST.map((m) => m.id).filter((id) => id.startsWith('gemini'));
+      },
       options: { baseURL: urlJoin(baseURL, '/gemini') },
     },
     {