diff --git a/packages/services/src/__tests__/ChatService.test.ts b/packages/services/src/__tests__/ChatService.test.ts index 3cdc67f..5724049 100644 --- a/packages/services/src/__tests__/ChatService.test.ts +++ b/packages/services/src/__tests__/ChatService.test.ts @@ -37,6 +37,18 @@ vi.mock('../../lib/handleStreamData', () => ({ default: vi.fn().mockReturnValue(() => {}), })); +// Mock ProviderRepository +vi.mock('@open-gsio/ai/providers/_ProviderRepository.ts', () => { + return { + ProviderRepository: class { + constructor() {} + getProviders() { + return [{ name: 'openai', key: 'test-key', endpoint: 'https://api.openai.com/v1' }]; + } + }, + }; +}); + describe('ChatService', () => { let chatService: any; let mockEnv: any; @@ -221,6 +233,105 @@ describe('ChatService', () => { Response.json = originalResponseJson; localService.getSupportedModels = originalGetSupportedModels; }); + + it('should test the cache refresh mechanism when providers change', async () => { + // This test verifies that the cache is refreshed when providers change + // and that the cache is used when providers haven't changed. + + // Mock data for the first scenario (cache hit) + const cachedModels = [ + { id: 'model-1', provider: 'openai' }, + { id: 'model-2', provider: 'openai' }, + ]; + const providersSignature = JSON.stringify(['openai']); + + // Mock KV_STORAGE for the first scenario (cache hit) + const mockKVStorage = { + get: vi.fn().mockImplementation(key => { + if (key === 'supportedModels') return Promise.resolve(JSON.stringify(cachedModels)); + if (key === 'providersSignature') return Promise.resolve(providersSignature); + return Promise.resolve(null); + }), + put: vi.fn().mockResolvedValue(undefined), + }; + + // The ProviderRepository is already mocked at the top of the file + + // Create a service instance with the mocked environment + const service = ChatService.create({ + maxTokens: 2000, + systemPrompt: 'You are a helpful assistant.', + }); + + // Set up the environment with the mocked KV_STORAGE + service.setEnv({ + ...mockEnv, + KV_STORAGE: mockKVStorage, + }); + + // Scenario 1: Cache hit - providers haven't changed + const response1 = await service.getSupportedModels(); + const data1 = await response1.json(); + + // Verify the cache was used + expect(mockKVStorage.get).toHaveBeenCalledWith('supportedModels'); + expect(mockKVStorage.get).toHaveBeenCalledWith('providersSignature'); + expect(data1).toEqual(cachedModels); + expect(mockKVStorage.put).not.toHaveBeenCalled(); + + // Reset the mock calls for the next scenario + vi.clearAllMocks(); + + // Scenario 2: Cache miss - providers have changed + // Update the mock to return a different providers signature + mockKVStorage.get.mockImplementation(key => { + if (key === 'supportedModels') { + return Promise.resolve(JSON.stringify(cachedModels)); + } + if (key === 'providersSignature') { + // Different signature + return Promise.resolve(JSON.stringify(['openai', 'anthropic'])); + } + return Promise.resolve(null); + }); + + // Mock the provider models fetching to avoid actual API calls + const mockModels = [ + { id: 'new-model-1', provider: 'openai' }, + { id: 'new-model-2', provider: 'openai' }, + ]; + + // Mock OpenAI instance for the second scenario + const mockOpenAIInstance = { + models: { + list: vi.fn().mockResolvedValue({ + data: mockModels, + }), + retrieve: vi.fn().mockImplementation(id => { + return Promise.resolve({ id, provider: 'openai' }); + }), + }, + }; + + // Update the OpenAI mock + vi.mocked(OpenAI).mockImplementation(() => mockOpenAIInstance as any); + + // Call getSupportedModels again + const response2 = await service.getSupportedModels(); + + // Verify the cache was refreshed + expect(mockKVStorage.get).toHaveBeenCalledWith('supportedModels'); + expect(mockKVStorage.get).toHaveBeenCalledWith('providersSignature'); + expect(mockKVStorage.put).toHaveBeenCalledTimes(2); // Called twice: once for models, once for signature + expect(mockKVStorage.put).toHaveBeenCalledWith('supportedModels', expect.any(String), { + expirationTtl: 60 * 60 * 24, + }); + expect(mockKVStorage.put).toHaveBeenCalledWith('providersSignature', expect.any(String), { + expirationTtl: 60 * 60 * 24, + }); + + // No need to restore mocks as we're using vi.mock at the module level + }); }); // TODO: Fix this test suite diff --git a/packages/services/src/chat-service/ChatService.ts b/packages/services/src/chat-service/ChatService.ts index 98885d8..d210911 100644 --- a/packages/services/src/chat-service/ChatService.ts +++ b/packages/services/src/chat-service/ChatService.ts @@ -118,11 +118,19 @@ const ChatService = types const useCache = true; + // Create a signature of the current providers + const providerRepo = new ProviderRepository(self.env); + const providers = providerRepo.getProviders(); + const currentProvidersSignature = JSON.stringify(providers.map(p => p.name).sort()); + if (useCache) { // ----- 1. Try cached value --------------------------------------------- try { const cached = yield self.env.KV_STORAGE.get('supportedModels'); - if (cached) { + const cachedSignature = yield self.env.KV_STORAGE.get('providersSignature'); + + // Check if cache exists and providers haven't changed + if (cached && cachedSignature && cachedSignature === currentProvidersSignature) { const parsed = JSON.parse(cached as string); if (Array.isArray(parsed) && parsed.length > 0) { logger.info('Cache hit – returning supportedModels from KV'); @@ -130,6 +138,11 @@ const ChatService = types } logger.warn('Cache entry malformed – refreshing'); throw new Error('Malformed cache entry'); + } else if ( + cached && + (!cachedSignature || cachedSignature !== currentProvidersSignature) + ) { + logger.info('Providers changed – refreshing cache'); } } catch (err) { logger.warn('Error reading/parsing supportedModels cache', err); @@ -137,8 +150,6 @@ const ChatService = types } // ----- 2. Build fresh list --------------------------------------------- - const providerRepo = new ProviderRepository(self.env); - const providers = providerRepo.getProviders(); const providerModels = new Map(); const modelMeta = new Map(); @@ -195,11 +206,20 @@ const ChatService = types // ----- 4. Cache fresh list --------------------------------------------- try { + // Store the models yield self.env.KV_STORAGE.put( 'supportedModels', JSON.stringify(resultArr), - { expirationTtl: 60 * 60 * 24 }, // 24 + { expirationTtl: 60 * 60 * 24 }, // 24 hours ); + + // Store the providers signature + yield self.env.KV_STORAGE.put( + 'providersSignature', + currentProvidersSignature, + { expirationTtl: 60 * 60 * 24 }, // 24 hours + ); + logger.info('supportedModels cache refreshed'); } catch (err) { logger.error('KV put failed for supportedModels', err);