diff --git a/server/src/llm/__tests__/llmService.test.ts b/server/src/llm/__tests__/llmService.test.ts index 0e7048c..f35369e 100644 --- a/server/src/llm/__tests__/llmService.test.ts +++ b/server/src/llm/__tests__/llmService.test.ts @@ -6,68 +6,157 @@ describe('llmService', () => { beforeEach(() => { originalFetch = globalThis.fetch; + vi.useFakeTimers(); }); afterEach(() => { globalThis.fetch = originalFetch; delete process.env.OPENROUTER_API_KEY; + delete process.env.LLM_MODEL; + delete process.env.LLM_FALLBACK_MODEL; + vi.useRealTimers(); }); - it('returns null when LLM is disabled (no API key)', async () => { + function setupEnv(opts?: { fallback?: string }) { + process.env.OPENROUTER_API_KEY = 'test-key'; + process.env.LLM_MODEL = 'free/model:free'; + if (opts?.fallback) { + process.env.LLM_FALLBACK_MODEL = opts.fallback; + } + } + + function mockFetchOk(content: string) { + globalThis.fetch = vi.fn().mockResolvedValue({ + ok: true, + json: () => Promise.resolve({ + choices: [{ message: { content } }], + }), + }); + } + + function mockFetch429(resetAt?: number) { + const metadata = resetAt != null + ? { headers: { 'X-RateLimit-Reset': String(resetAt) } } + : undefined; + globalThis.fetch = vi.fn().mockResolvedValue({ + ok: false, + status: 429, + statusText: 'Too Many Requests', + json: () => Promise.resolve({ + error: { code: 429, message: 'Rate limit exceeded', metadata }, + }), + }); + } + + it('returns null when LLM is disabled', async () => { delete process.env.OPENROUTER_API_KEY; const service = createLlmService(); const result = await service.generate('backstory', { npcName: 'Test' }); expect(result).toBeNull(); }); - it('generates text using a named template', async () => { - process.env.OPENROUTER_API_KEY = 'test-key'; - globalThis.fetch = vi.fn().mockResolvedValue({ - ok: true, - json: () => Promise.resolve({ - choices: [{ message: { content: 'A brave warrior from the north.' } }], - }), - }); + it('generates text using primary model', async () => { + setupEnv(); + mockFetchOk('A brave warrior.'); const service = createLlmService(); const result = await service.generate('backstory', { npcName: 'Brynn', - stats: 'STR:15, DEX:12, CON:14, INT:8, PER:10, SOC:6, COU:16, CUR:11, EMP:7, TMP:13', + stats: 'STR:15', }); + await vi.advanceTimersByTimeAsync(100); - expect(result).toBe('A brave warrior from the north.'); + expect(result).toBe('A brave warrior.'); }); - it('returns false for isDailyLimitReached when disabled', () => { - delete process.env.OPENROUTER_API_KEY; + it('switches to fallback model on 429', async () => { + setupEnv({ fallback: 'paid/model' }); + const resetAt = Date.now() + 3600000; // 1 hour from now + + // First call hits 429 + mockFetch429(resetAt); const service = createLlmService(); - expect(service.isDailyLimitReached()).toBe(false); + const result1 = await service.generate('backstory', { + npcName: 'Brynn', + stats: 'STR:15', + }); + await vi.advanceTimersByTimeAsync(100); + expect(result1).toBeNull(); // The 429 request itself returns null + + // Second call should use fallback model + mockFetchOk('Fallback response.'); + const result2 = await service.generate('backstory', { + npcName: 'Brynn', + stats: 'STR:15', + }); + await vi.advanceTimersByTimeAsync(100); + expect(result2).toBe('Fallback response.'); + + // Verify the model sent was the fallback + const body = JSON.parse( + (globalThis.fetch as ReturnType).mock.calls[0][1].body, + ); + expect(body.model).toBe('paid/model'); }); - it('exposes isDailyLimitReached when enabled', () => { - process.env.OPENROUTER_API_KEY = 'test-key'; - globalThis.fetch = vi.fn().mockResolvedValue({ - ok: true, - json: () => Promise.resolve({ - choices: [{ message: { content: 'ok' } }], - }), - }); + it('switches back to primary model after reset time', async () => { + setupEnv({ fallback: 'paid/model' }); + const resetAt = Date.now() + 5000; // 5 seconds from now + + // Hit 429 + mockFetch429(resetAt); + const service = createLlmService(); + await service.generate('backstory', { npcName: 'X', stats: 'S:1' }); + await vi.advanceTimersByTimeAsync(100); + + expect(service.activeModel()).toBe('paid/model'); + + // Advance past reset time + await vi.advanceTimersByTimeAsync(6000); + + expect(service.activeModel()).toBe('free/model:free'); + }); + + it('returns null on 429 when no fallback model configured', async () => { + setupEnv(); // no fallback + mockFetch429(Date.now() + 3600000); const service = createLlmService(); - expect(service.isDailyLimitReached()).toBe(false); + const result = await service.generate('backstory', { + npcName: 'X', + stats: 'S:1', + }); + await vi.advanceTimersByTimeAsync(100); + + expect(result).toBeNull(); + }); + + it('tracks usage per model', async () => { + setupEnv({ fallback: 'paid/model' }); + mockFetchOk('ok'); + + const service = createLlmService(); + await service.generate('backstory', { npcName: 'X', stats: 'S:1' }); + await vi.advanceTimersByTimeAsync(100); + + const stats = service.usageStats(); + expect(stats['free/model:free']?.total).toBe(1); }); it('exposes queue depth and clear', async () => { - process.env.OPENROUTER_API_KEY = 'test-key'; - globalThis.fetch = vi.fn().mockResolvedValue({ - ok: true, - json: () => Promise.resolve({ - choices: [{ message: { content: 'ok' } }], - }), - }); + setupEnv(); + mockFetchOk('ok'); const service = createLlmService(); expect(service.queueDepth()).toBe(0); service.clear(); }); + + it('exposes activeModel', () => { + setupEnv({ fallback: 'paid/model' }); + mockFetchOk('ok'); + + const service = createLlmService(); + expect(service.activeModel()).toBe('free/model:free'); + }); }); diff --git a/server/src/llm/llmService.ts b/server/src/llm/llmService.ts index bb8dd92..5056036 100644 --- a/server/src/llm/llmService.ts +++ b/server/src/llm/llmService.ts @@ -1,6 +1,7 @@ import { getLlmConfig } from '../config/llmConfig.js'; -import { createOpenRouterClient } from './openRouterClient.js'; +import { createOpenRouterClient, isRateLimited } from './openRouterClient.js'; import { createGenerationQueue } from './generationQueue.js'; +import { createUsageCounters } from './usageCounters.js'; import { renderTemplate } from './promptTemplate.js'; import { templates } from './templates.js'; @@ -8,18 +9,22 @@ export interface LlmService { generate(templateName: string, variables: Record): Promise; queueDepth(): number; clear(): void; - isDailyLimitReached(): boolean; + activeModel(): string; + usageStats(): Record; + usageSummary(): string; } export function createLlmService(): LlmService { const config = getLlmConfig(); - if (!config.enabled) { + if (!config.enabled || !config.model) { return { generate: async () => null, queueDepth: () => 0, clear: () => {}, - isDailyLimitReached: () => false, + activeModel: () => '', + usageStats: () => ({}), + usageSummary: () => 'LLM disabled', }; } @@ -27,6 +32,35 @@ export function createLlmService(): LlmService { const queue = createGenerationQueue(client, { requestsPerMinute: config.requestsPerMinute, }); + const counters = createUsageCounters(); + + let currentModel = config.model; + let switchBackTimer: ReturnType | null = null; + + function getNextMidnightUTC(): number { + const now = new Date(); + return new Date(Date.UTC( + now.getUTCFullYear(), now.getUTCMonth(), now.getUTCDate() + 1, + 0, 0, 0, 0, + )).getTime(); + } + + function switchToFallback(resetAt: number | null): void { + if (!config.fallbackModel) return; + + currentModel = config.fallbackModel; + console.log(`[LLM] Switched to fallback model: ${currentModel}`); + + const resetTime = resetAt ?? getNextMidnightUTC(); + const delay = Math.max(0, resetTime - Date.now()); + + if (switchBackTimer) clearTimeout(switchBackTimer); + switchBackTimer = setTimeout(() => { + currentModel = config.model; + switchBackTimer = null; + console.log(`[LLM] Switched back to primary model: ${currentModel}`); + }, delay); + } return { async generate( @@ -39,7 +73,18 @@ export function createLlmService(): LlmService { return null; } const rendered = renderTemplate(template, variables); - return queue.enqueue(rendered); + const result = await queue.enqueue({ ...rendered, model: currentModel }); + + if (isRateLimited(result)) { + switchToFallback(result.resetAt); + return null; + } + + if (typeof result === 'string') { + counters.record(currentModel); + } + + return result; }, queueDepth(): number { @@ -50,8 +95,16 @@ export function createLlmService(): LlmService { queue.clear(); }, - isDailyLimitReached(): boolean { - return queue.isDailyLimitReached(); + activeModel(): string { + return currentModel; + }, + + usageStats(): Record { + return counters.getStats(); + }, + + usageSummary(): string { + return counters.getSummary(); }, }; }