diff --git a/server/src/llm/__tests__/llmService.test.ts b/server/src/llm/__tests__/llmService.test.ts index b16a7f8..984d3ec 100644 --- a/server/src/llm/__tests__/llmService.test.ts +++ b/server/src/llm/__tests__/llmService.test.ts @@ -161,6 +161,56 @@ describe('llmService', () => { expect(service.activeModel()).toBe('free/model:free'); }); + it('records successful call to stats service', async () => { + setupEnv(); + globalThis.fetch = vi.fn().mockResolvedValue({ + ok: true, + json: () => Promise.resolve({ + choices: [{ message: { content: 'ok' } }], + usage: { prompt_tokens: 100, completion_tokens: 50 }, + }), + }); + + const mockStatsService = { + record: vi.fn(), + getStats: vi.fn(), + isTrackingEnabled: vi.fn().mockReturnValue(true), + setTrackingEnabled: vi.fn(), + }; + + const service = createLlmService(undefined, mockStatsService); + await service.generate('backstory', { npcName: 'X', stats: 'S:1' }); + await vi.advanceTimersByTimeAsync(100); + + expect(mockStatsService.record).toHaveBeenCalledWith( + 'backstory', 100, 50, 'free/model:free', 0, false, 0, + ); + }); + + it('records failed call to stats service', async () => { + setupEnv(); + globalThis.fetch = vi.fn().mockResolvedValue({ + ok: false, + status: 400, + statusText: 'Bad Request', + }); + + const mockStatsService = { + record: vi.fn(), + getStats: vi.fn(), + isTrackingEnabled: vi.fn().mockReturnValue(true), + setTrackingEnabled: vi.fn(), + }; + + const service = createLlmService(undefined, mockStatsService); + await service.generate('backstory', { npcName: 'X', stats: 'S:1' }); + await vi.advanceTimersByTimeAsync(100); + + expect(mockStatsService.record).toHaveBeenCalledWith( + 'backstory', 0, 0, 'free/model:free', 0, true, 0, + ); + }); + it('destroy cleans up switch-back timer', async () => { setupEnv({ fallback: 'paid/model' }); const resetAt = Date.now() + 60000; diff --git a/server/src/llm/llmService.ts b/server/src/llm/llmService.ts index 72b0dd3..3d97f38 100644 --- a/server/src/llm/llmService.ts +++ b/server/src/llm/llmService.ts @@ -6,6 +6,7 @@ import { renderTemplate } from './promptTemplate.js'; import { templates } from './templates.js'; import { createTokenTracker, type TokenTracker } from './tokenTracker.js'; import type { LogService } from '../services/logService.js'; +import type { LlmStatsService } from './llmStatsService.js'; export interface LlmService { generate(templateName: string, variables: Record): Promise; @@ -18,7 +19,7 @@ export interface LlmService { tokenUsage(): TokenTracker; } -export function createLlmService(logService?: LogService): LlmService { +export function createLlmService(logService?: LogService, statsService?: LlmStatsService): LlmService { const config = getLlmConfig(); if (!config.enabled || !config.model) { @@ -87,15 +88,19 @@ export function createLlmService(logService?: LogService): LlmService { if (isRateLimited(result)) { logService?.log('warning', 'LLM', 'Rate limited, switching to fallback'); switchToFallback(result.resetAt); + statsService?.record(templateName, 0, 0, currentModel, 0, true, 0); return null; } if (isSuccess(result)) { counters.record(currentModel); + const inputTokens = result.usage?.promptTokens ?? 0; + const outputTokens = result.usage?.completionTokens ?? 0; if (result.usage) { - tokenTracker.record(templateName, result.usage.promptTokens, result.usage.completionTokens); + tokenTracker.record(templateName, inputTokens, outputTokens); console.log(tokenTracker.lastSummary()); } + statsService?.record(templateName, inputTokens, outputTokens, currentModel, result.retries, false, 0); const totalRequests = Object.values(counters.getStats()).reduce((sum, s) => sum + s.total, 0); if (totalRequests % 100 === 0) { console.log(`[LLM] ${counters.getSummary()}`); @@ -104,6 +109,8 @@ export function createLlmService(logService?: LogService): LlmService { return result.content; } + // null result = failure + statsService?.record(templateName, 0, 0, currentModel, 0, true, 0); return result; },