Coverage for tests/test_extract_usage.py: 98.81%

1import re (empty)

2from collections.abc import Mapping (empty)

3from decimal import Decimal (empty)

4from typing import Any (empty)

5 

6import pytest (empty)

7from inline_snapshot import snapshot (empty)

8 

9from voice_prices import Usage, calc_price, extract_usage (empty)

10from voice_prices.data import providers (empty)

11from voice_prices.types import ( (empty)

12 ArrayMatch, 

13 ClauseEquals, 

14 ExtractedUsage, 

15 ModelPrice, 

16 Provider, 

17 UsageExtractor, 

18 UsageExtractorMapping, 

19) 

20 

21 

22class MyMapping(Mapping[str, Any]): (empty)

23 def __init__(self, **data: Any): (empty)

24 self._data = data (empty)

25 

26 def __getitem__(self, key: str) -> Any: (empty)

27 return self._data[key] (empty)

28 

29 def __iter__(self) -> Any: (empty)

30 return iter(self._data) 

31 

32 def __len__(self) -> int: (empty)

33 return len(self._data) 

34 

35 

36@pytest.mark.parametrize( (empty)

37 'response_data,expected_model,expected_usage,expected_price', 

38 [ 

39 ( 

40 { 

41 'id': 'msg_0152tnC3YpjyASTB9qxqDJXu', 

42 'type': 'message', 

43 'role': 'assistant', 

44 'model': 'claude-sonnet-4-20250514', 

45 'stop_reason': 'tool_use', 

46 'stop_sequence': None, 

47 'usage': { 

48 'input_tokens': 504, 

49 'cache_creation_input_tokens': 123, 

50 'cache_read_input_tokens': 0, 

51 'output_tokens': 97, 

52 'service_tier': 'standard', 

53 }, 

54 }, 

55 snapshot('claude-sonnet-4-20250514'), 

56 snapshot(Usage(input_tokens=627, cache_write_tokens=123, cache_read_tokens=0, output_tokens=97)), 

57 snapshot(Decimal('0.00342825')), 

58 ), 

59 ( 

60 { 

61 'model': 'claude-3-5-haiku', 

62 'usage': {'input_tokens': 504, 'output_tokens': 97, 'service_tier': 'standard'}, 

63 }, 

64 snapshot('claude-3-5-haiku'), 

65 snapshot(Usage(input_tokens=504, output_tokens=97)), 

66 snapshot(Decimal('0.0007912')), 

67 ), 

68 ( 

69 MyMapping( 

70 model='claude-3.5-haiku', usage=MyMapping(input_tokens=504, output_tokens=97, service_tier='standard') 

71 ), 

72 snapshot('claude-3.5-haiku'), 

73 snapshot(Usage(input_tokens=504, output_tokens=97)), 

74 snapshot(Decimal('0.0007912')), 

75 ), 

76 ], 

77) 

78def test_extract_usage_ok(response_data: Any, expected_model: str, expected_usage: Usage, expected_price: Decimal): (empty)

79 provider = next(provider for provider in providers if provider.id == 'anthropic') (empty)

80 assert provider.name == 'Anthropic' (empty)

81 assert provider.extractors is not None (empty)

82 model, usage = provider.extract_usage(response_data) (empty)

83 assert model == expected_model (empty)

84 assert usage == expected_usage (empty)

85 

86 # also test the public simple API 

87 extracted_usage = extract_usage(response_data, provider_id='anthropic') (empty)

88 assert extracted_usage.usage == expected_usage (empty)

89 assert extracted_usage.provider.name == 'Anthropic' (empty)

90 

91 assert extracted_usage.calc_price().total_price == expected_price (empty)

92 

93 

94def test_openai(): (empty)

95 provider = next(provider for provider in providers if provider.id == 'openai') (empty)

96 assert provider.name == 'OpenAI' (empty)

97 assert provider.extractors is not None (empty)

98 response_data = { (empty)

99 'model': 'gpt-4.1', 

100 'usage': { 

101 'prompt_tokens': 100, 

102 'completion_tokens': 200, 

103 'prompt_tokens_details': {'cached_tokens': None}, 

104 'completion_tokens_details': None, 

105 }, 

106 } 

107 usage = provider.extract_usage(response_data, api_flavor='chat') (empty)

108 assert usage == snapshot(('gpt-4.1', Usage(input_tokens=100, output_tokens=200))) (empty)

109 

110 extracted_usage = extract_usage(response_data, provider_id='openai', api_flavor='chat') (empty)

111 assert extracted_usage.usage == snapshot(Usage(input_tokens=100, output_tokens=200)) (empty)

112 assert extracted_usage.provider.name == snapshot('OpenAI') (empty)

113 assert extracted_usage.model is not None (empty)

114 assert extracted_usage.model.name == snapshot('gpt 4.1') (empty)

115 

116 assert extracted_usage.calc_price().total_price == snapshot(Decimal('0.0018')) (empty)

117 

118 response_data = { (empty)

119 'model': 'gpt-5', 

120 'usage': {'input_tokens': 100, 'output_tokens': 200}, 

121 } 

122 usage = provider.extract_usage(response_data, api_flavor='responses') (empty)

123 assert usage == snapshot(('gpt-5', Usage(input_tokens=100, output_tokens=200))) (empty)

124 

125 extracted_usage = extract_usage(response_data, provider_id='openai', api_flavor='responses') (empty)

126 assert extracted_usage.usage == snapshot(Usage(input_tokens=100, output_tokens=200)) (empty)

127 assert extracted_usage.provider.name == snapshot('OpenAI') (empty)

128 assert extracted_usage.model is not None (empty)

129 assert extracted_usage.model.name == snapshot('GPT-5') (empty)

130 

131 assert extracted_usage.calc_price().total_price == snapshot(Decimal('0.002125')) (empty)

132 

133 with pytest.raises(ValueError, match=re.escape("Unknown api_flavor 'default', allowed values: chat, responses")): (empty)

134 provider.extract_usage(response_data) (empty)

135 

136 

137def test_openrouter_chat_cache_write_tokens(): (empty)

138 provider = next(provider for provider in providers if provider.id == 'openrouter') (empty)

139 assert provider.name == 'OpenRouter' (empty)

140 assert provider.extractors is not None (empty)

141 response_data = { (empty)

142 'model': 'anthropic/claude-4.6-sonnet-20260217', 

143 'usage': { 

144 'prompt_tokens': 4819, 

145 'completion_tokens': 1906, 

146 'total_tokens': 6725, 

147 'prompt_tokens_details': { 

148 'cached_tokens': 0, 

149 'cache_write_tokens': 4800, 

150 'audio_tokens': 17, 

151 }, 

152 'completion_tokens_details': { 

153 'audio_tokens': 23, 

154 }, 

155 }, 

156 } 

157 usage = provider.extract_usage(response_data, api_flavor='chat') (empty)

158 assert usage == snapshot( (empty)

159 ( 

160 'anthropic/claude-4.6-sonnet-20260217', 

161 Usage( 

162 input_tokens=4819, 

163 cache_write_tokens=4800, 

164 cache_read_tokens=0, 

165 output_tokens=1906, 

166 input_audio_tokens=17, 

167 output_audio_tokens=23, 

168 ), 

169 ) 

170 ) 

171 

172 extracted_usage = extract_usage(response_data, provider_id='openrouter', api_flavor='chat') (empty)

173 assert extracted_usage.usage == snapshot( (empty)

174 Usage( 

175 input_tokens=4819, 

176 cache_write_tokens=4800, 

177 cache_read_tokens=0, 

178 output_tokens=1906, 

179 input_audio_tokens=17, 

180 output_audio_tokens=23, 

181 ) 

182 ) 

183 assert extracted_usage.provider.name == snapshot('OpenRouter') (empty)

184 assert extracted_usage.model is not None (empty)

185 assert extracted_usage.model.id == snapshot('anthropic/claude-sonnet-4.6') (empty)

186 

187 extracted_usage_by_url = extract_usage( (empty)

188 response_data, provider_api_url='https://openrouter.ai/api/v1', api_flavor='chat' 

189 ) 

190 assert extracted_usage_by_url.usage == extracted_usage.usage (empty)

191 

192 

193def test_extracted_usage_calc_price_requires_model(): (empty)

194 extracted_usage = ExtractedUsage( (empty)

195 usage=Usage(input_tokens=1), 

196 model=None, 

197 provider=Provider(id='test', name='Test', api_pattern='test'), 

198 auto_update_timestamp=None, 

199 ) 

200 

201 with pytest.raises(ValueError, match='No model reference found in response data and model not provided'): (empty)

202 extracted_usage.calc_price() (empty)

203 

204 

205@pytest.mark.parametrize( (empty)

206 'response_data,error', 

207 [ 

208 ({}, snapshot('Missing value at `usage`')), 

209 ({'model': None}, snapshot('Missing value at `usage`')), 

210 ({'model': 'x'}, snapshot('Missing value at `usage`')), 

211 ({'model': 'x', 'usage': {}}, snapshot('Missing value at `usage.input_tokens`')), 

212 ({'model': 'x', 'usage': 123}, snapshot('Expected `usage` value to be a Mapping, got int')), 

213 ( 

214 {'model': 'x', 'usage': {'input_tokens': 123.0}}, 

215 snapshot('Expected `usage.input_tokens` value to be a int, got float'), 

216 ), 

217 ( 

218 {'model': 'x', 'usage': {'input_tokens': []}}, 

219 snapshot('Expected `usage.input_tokens` value to be a int, got list'), 

220 ), 

221 ], 

222) 

223def test_extract_usage_error(response_data: Any, error: str): (empty)

224 provider = next(provider for provider in providers if provider.id == 'anthropic') (empty)

225 assert provider.name == 'Anthropic' (empty)

226 assert provider.extractors is not None (empty)

227 

228 with pytest.raises(ValueError) as exc_info: (empty)

229 provider.extract_usage(response_data) (empty)

230 

231 assert str(exc_info.value) == error (empty)

232 

233 

234def test_usage_extractor_errors_when_optional_mappings_find_no_usage_values(): (empty)

235 extractor = UsageExtractor( (empty)

236 root='usage', 

237 mappings=[ 

238 UsageExtractorMapping( 

239 path=[ArrayMatch(type='array-match', field='modality', match=ClauseEquals('AUDIO')), 'tokenCount'], 

240 dest='input_audio_tokens', 

241 required=False, 

242 ) 

243 ], 

244 ) 

245 

246 with pytest.raises(ValueError, match='No usage information found at usage'): (empty)

247 extractor.extract({'model': 'test-model', 'usage': {}}) (empty)

248 

249 

250def test_usage_extractor_errors_when_required_nested_path_has_wrong_type(): (empty)

251 extractor = UsageExtractor( (empty)

252 root='usage', 

253 mappings=[UsageExtractorMapping(path=['totals', 'input_tokens'], dest='input_tokens')], 

254 ) 

255 

256 with pytest.raises(ValueError, match='Expected `usage.totals` value to be a dict, got int'): (empty)

257 extractor.extract({'model': 'test-model', 'usage': {'totals': 1}}) (empty)

258 

259 

260def test_usage_extractor_skips_optional_nested_path_with_wrong_type(): (empty)

261 extractor = UsageExtractor( (empty)

262 root='usage', 

263 mappings=[ 

264 UsageExtractorMapping(path=['totals', 'input_tokens'], dest='input_tokens', required=False), 

265 UsageExtractorMapping(path='output_tokens', dest='output_tokens'), 

266 ], 

267 ) 

268 

269 assert extractor.extract({'model': 'test-model', 'usage': {'totals': 1, 'output_tokens': 2}}) == ( (empty)

270 'test-model', 

271 Usage(output_tokens=2), 

272 ) 

273 

274 

275def test_no_flavors(): (empty)

276 provider = Provider(id='test', name='Test', api_pattern='x') (empty)

277 

278 with pytest.raises(ValueError, match='No extraction logic defined for this provider'): (empty)

279 provider.extract_usage({}) (empty)

280 

281 

282gemini_response_data = { (empty)

283 'usageMetadata': { 

284 'promptTokenCount': 75, 

285 'candidatesTokenCount': 18, 

286 'totalTokenCount': 262, 

287 'trafficType': 'ON_DEMAND', 

288 'promptTokensDetails': [{'modality': 'TEXT', 'tokenCount': 75}], 

289 'candidatesTokensDetails': [{'modality': 'TEXT', 'tokenCount': 18}], 

290 'toolUsePromptTokenCount': 25, 

291 'thoughtsTokenCount': 144, 

292 }, 

293 'modelVersion': 'gemini-2.5-flash', 

294 'createTime': '2025-08-25T14:26:17.534704Z', 

295 'responseId': 'iXKsaLDRIPqsgLUPotqEyA0', 

296} 

297google_provider = next(provider for provider in providers if provider.id == 'google') (empty)

298assert google_provider.name == 'Google' (empty)

299assert google_provider.extractors is not None (empty)

300 

301 

302def test_google(): (empty)

303 usage = google_provider.extract_usage(gemini_response_data) (empty)

304 assert usage == snapshot(('gemini-2.5-flash', Usage(input_tokens=100, output_tokens=162))) (empty)

305 

306 

307gemini_response_data_caching = { (empty)

308 'usageMetadata': { 

309 'promptTokenCount': 14152, 

310 'candidatesTokenCount': 60, 

311 'totalTokenCount': 14271, 

312 'cachedContentTokenCount': 12239, 

313 'trafficType': 'ON_DEMAND', 

314 'promptTokensDetails': [{'modality': 'TEXT', 'tokenCount': 14002}, {'modality': 'AUDIO', 'tokenCount': 150}], 

315 'cacheTokensDetails': [{'modality': 'AUDIO', 'tokenCount': 129}, {'modality': 'TEXT', 'tokenCount': 12110}], 

316 'candidatesTokensDetails': [{'modality': 'TEXT', 'tokenCount': 50}, {'modality': 'AUDIO', 'tokenCount': 10}], 

317 'thoughtsTokenCount': 69, 

318 }, 

319 'modelVersion': 'gemini-2.5-flash', 

320} 

321 

322 

323def test_google_caching(): (empty)

324 model, usage = google_provider.extract_usage(gemini_response_data_caching) (empty)

325 assert model == snapshot('gemini-2.5-flash') (empty)

326 assert usage == snapshot( (empty)

327 Usage( 

328 input_tokens=14152, 

329 cache_read_tokens=12239, 

330 output_tokens=129, 

331 input_audio_tokens=150, 

332 cache_audio_read_tokens=129, 

333 output_audio_tokens=10, 

334 ), 

335 ) 

336 assert model is not None (empty)

337 assert calc_price(usage, model).total_price == snapshot(Decimal('0.0012873')) (empty)

338 

339 

340def test_google_caching_public_extraction_parity(): (empty)

341 extracted_usage = extract_usage(gemini_response_data_caching, provider_id='google') (empty)

342 

343 assert extracted_usage.usage == snapshot( (empty)

344 Usage( 

345 input_tokens=14152, 

346 cache_read_tokens=12239, 

347 output_tokens=129, 

348 input_audio_tokens=150, 

349 cache_audio_read_tokens=129, 

350 output_audio_tokens=10, 

351 ) 

352 ) 

353 assert extracted_usage.model is not None (empty)

354 assert ( (empty)

355 extracted_usage.calc_price().total_price 

356 == calc_price( 

357 extracted_usage.usage, 

358 extracted_usage.model.id, 

359 provider_id='google', 

360 ).total_price 

361 ) 

362 

363 

364gemini_response_data_thoughtless = { (empty)

365 'usageMetadata': { 

366 'promptTokenCount': 75, 

367 'candidatesTokenCount': 18, 

368 'totalTokenCount': 237, 

369 'trafficType': 'ON_DEMAND', 

370 'promptTokensDetails': [{'modality': 'TEXT', 'tokenCount': 75}], 

371 'candidatesTokensDetails': [{'modality': 'TEXT', 'tokenCount': 18}], 

372 }, 

373 'modelVersion': 'gemini-2.5-flash', 

374 'createTime': '2025-08-25T14:26:17.534704Z', 

375 'responseId': 'iXKsaLDRIPqsgLUPotqEyA0', 

376} 

377 

378 

379def test_gemini_response_thoughtless(): (empty)

380 usage = google_provider.extract_usage(gemini_response_data_thoughtless) (empty)

381 assert usage == snapshot(('gemini-2.5-flash', Usage(input_tokens=75, output_tokens=18))) (empty)

382 

383 

384def test_bedrock(): (empty)

385 provider = next(provider for provider in providers if provider.id == 'aws') (empty)

386 response_data = {'usage': {'inputTokens': 406, 'outputTokens': 53}} (empty)

387 usage = provider.extract_usage(response_data) (empty)

388 assert usage == snapshot((None, Usage(input_tokens=406, output_tokens=53))) (empty)

389 

390 extracted_usage = extract_usage(response_data, provider_id='aws') (empty)

391 assert extracted_usage.usage == snapshot(Usage(input_tokens=406, output_tokens=53)) (empty)

392 assert extracted_usage.provider.name == snapshot('AWS') (empty)

393 assert extracted_usage.model == snapshot(None) (empty)

394 

395 

396anthropic_response_data = { (empty)

397 'model': 'claude-sonnet-4-20250514', 

398 'usage': { 

399 'input_tokens': 483, 

400 'cache_creation_input_tokens': 0, 

401 'cache_read_input_tokens': 0, 

402 'output_tokens': 78, 

403 }, 

404} 

405 

406 

407def test_google_anthropic(): (empty)

408 usage = google_provider.extract_usage(anthropic_response_data, api_flavor='anthropic') (empty)

409 assert usage == snapshot( (empty)

410 ( 

411 'claude-sonnet-4-20250514', 

412 Usage(input_tokens=483, cache_write_tokens=0, cache_read_tokens=0, output_tokens=78), 

413 ) 

414 ) 

415 

416 

417def test_extractor_accumulates_by_destination_string() -> None: (empty)

418 extractor = UsageExtractor( (empty)

419 root='usage', 

420 mappings=[ 

421 UsageExtractorMapping(path='prompt_tokens', dest='input_tokens'), 

422 UsageExtractorMapping(path='cached_tokens', dest='input_tokens', required=False), 

423 UsageExtractorMapping(path='missing_tokens', dest='output_tokens', required=False), 

424 ], 

425 ) 

426 

427 assert extractor.extract({'model': 'test-model', 'usage': {'prompt_tokens': 100, 'cached_tokens': 25}}) == ( (empty)

428 'test-model', 

429 Usage(input_tokens=125), 

430 ) 

431 

432 

433def test_extractor_accumulates_repeated_destination_string_with_zero_values() -> None: (empty)

434 extractor = UsageExtractor( (empty)

435 root='usage', 

436 mappings=[ 

437 UsageExtractorMapping(path='prompt_tokens', dest='input_tokens'), 

438 UsageExtractorMapping(path='cached_tokens', dest='input_tokens'), 

439 ], 

440 ) 

441 

442 assert extractor.extract({'model': 'test-model', 'usage': {'prompt_tokens': 0, 'cached_tokens': 25}}) == ( (empty)

443 'test-model', 

444 Usage(input_tokens=25), 

445 ) 

446 

447 

448def test_extractor_ignores_unknown_response_extras() -> None: (empty)

449 extractor = UsageExtractor( (empty)

450 root='usage', 

451 mappings=[UsageExtractorMapping(path='prompt_tokens', dest='input_tokens')], 

452 ) 

453 

454 assert extractor.extract( (empty)

455 {'model': 'test-model', 'usage': {'prompt_tokens': 100, 'provider_specific_tokens': 999}} 

456 ) == ( 

457 'test-model', 

458 Usage(input_tokens=100), 

459 ) 

460 

461 

462def test_extractor_stores_registered_contradictions_until_pricing_interprets_them() -> None: (empty)

463 extractor = UsageExtractor( (empty)

464 root='usage', 

465 mappings=[ 

466 UsageExtractorMapping(path='prompt_tokens', dest='input_tokens'), 

467 UsageExtractorMapping(path='audio_tokens', dest='input_audio_tokens'), 

468 ], 

469 ) 

470 

471 _, usage = extractor.extract({'model': 'test-model', 'usage': {'prompt_tokens': 50, 'audio_tokens': 100}}) (empty)

472 

473 assert usage == Usage(input_tokens=50, input_audio_tokens=100) (empty)

474 assert usage.input_tokens == 50 (empty)

475 with pytest.raises(ValueError, match='negative|Impossible usage data'): (empty)

476 ModelPrice(input_mtok=Decimal('1'), input_audio_mtok=Decimal('2')).calc_price(usage) (empty)

477 

478 

479def test_accumulate_extracted_usage(): (empty)

480 extracted = extract_usage(gemini_response_data, provider_id='google') (empty)

481 assert extracted.usage == Usage(input_tokens=100, output_tokens=162) (empty)

482 with pytest.raises(TypeError): (empty)

483 _ = extracted + 1 (empty)

484 with pytest.raises(TypeError): (empty)

485 _ = None + extracted (empty)

486 with pytest.raises(TypeError): (empty)

487 _ = extracted + None (empty)

488 with pytest.raises(ValueError): (empty)

489 _ = extracted + extract_usage(anthropic_response_data, provider_id='anthropic') (empty)

490 double_extracted = extracted + extracted (empty)

491 assert double_extracted.usage == Usage(input_tokens=100 * 2, output_tokens=162 * 2) (empty)

492 assert Usage(input_tokens=10, output_tokens=10) + Usage(output_tokens=10) == Usage( (empty)

493 input_tokens=10, output_tokens=20 

494 ) 

495 assert Usage(input_audio_tokens=10) + Usage(input_tokens=10) == Usage(input_audio_tokens=10, input_tokens=10) (empty)

496 

497 

498def test_xai_native(): (empty)

499 provider = next(provider for provider in providers if provider.id == 'x-ai') (empty)

500 response_data = { (empty)

501 'model': 'grok-4-fast-non-reasoning', 

502 'usage': { 

503 'prompt_tokens': 181, 

504 'cached_prompt_text_tokens': 162, 

505 'completion_tokens': 27, 

506 'prompt_text_tokens': 181, 

507 'total_tokens': 208, 

508 }, 

509 } 

510 model, usage = provider.extract_usage(response_data) (empty)

511 assert model == 'grok-4-fast-non-reasoning' (empty)

512 assert usage == Usage(input_tokens=181, cache_read_tokens=162, output_tokens=27) (empty)