| 1 | import { describe, expect, it, vi } from 'vitest' |
| 2 | import { |
| 3 | ModelUsageCallbackHandler, |
| 4 | extractModelUsage |
| 5 | } from '../../src/main/agent-runtime/model' |
| 6 | |
| 7 | describe('model usage tracking', () => { |
| 8 | it('prefers provider-reported usage metadata', () => { |
| 9 | const usage = extractModelUsage({ |
| 10 | generations: [ |
| 11 | [ |
| 12 | { |
| 13 | text: 'done', |
| 14 | message: { |
| 15 | usage_metadata: { |
| 16 | input_tokens: 120, |
| 17 | output_tokens: 35, |
| 18 | total_tokens: 155 |
| 19 | } |
| 20 | } |
| 21 | } |
| 22 | ] |
| 23 | ] |
| 24 | } as never) |
| 25 | |
| 26 | expect(usage).toEqual({ |
| 27 | inputTokens: 120, |
| 28 | outputTokens: 35, |
| 29 | totalTokens: 155, |
| 30 | source: 'provider' |
| 31 | }) |
| 32 | }) |
| 33 | |
| 34 | it('uses heuristic estimates when the provider omits usage', async () => { |
| 35 | const recordModelUsage = vi.fn(async () => undefined) |
| 36 | const handler = new ModelUsageCallbackHandler({ |
| 37 | provider: 'openai', |
| 38 | model: 'compatible-model' |
| 39 | }, { record: recordModelUsage }) |
| 40 | |
| 41 | handler.handleLLMStart({} as never, ['Create a concise presentation outline.'], 'run-1') |
| 42 | await handler.handleLLMEnd( |
| 43 | { |
| 44 | generations: [[{ text: 'A short outline with three sections.' }]] |
| 45 | }, |
| 46 | 'run-1' |
| 47 | ) |
| 48 | |
| 49 | expect(recordModelUsage).toHaveBeenCalledWith( |
| 50 | expect.objectContaining({ |
| 51 | provider: 'openai', |
| 52 | model: 'compatible-model', |
| 53 | source: 'estimated', |
| 54 | inputTokens: expect.any(Number), |
| 55 | outputTokens: expect.any(Number), |
| 56 | totalTokens: expect.any(Number) |
| 57 | }) |
| 58 | ) |
| 59 | const recorded = recordModelUsage.mock.calls[0][0] |
| 60 | expect(recorded.inputTokens).toBeGreaterThan(0) |
| 61 | expect(recorded.outputTokens).toBeGreaterThan(0) |
| 62 | expect(recorded.totalTokens).toBe(recorded.inputTokens + recorded.outputTokens) |
| 63 | }) |
| 64 | }) |
| 65 |