Extract usage for unknown models

Author: alexmojaki

pydantic_ai_slim/pydantic_ai/models/anthropic.py

@@ -7,7 +7,6 @@
 from datetime import datetime
 from typing import Any, Literal, cast, overload
 
-from genai_prices import extract_usage
 from pydantic import TypeAdapter
 from typing_extensions import assert_never
 
@@ -638,12 +637,7 @@ def _map_usage(
         key: value for key, value in response_usage.model_dump().items() if isinstance(value, int)
     }
 
-    extracted_usage = extract_usage(dict(model=model, usage=details), provider_id=provider)
-
-    return usage.RequestUsage(
-        **{key: value for key, value in extracted_usage.usage.__dict__.items() if isinstance(value, int)},
-        details=details,
-    )
+    return usage.RequestUsage.extract(dict(model=model, usage=details), provider_id=provider, details=details)
 
 
 @dataclass

pydantic_ai_slim/pydantic_ai/usage.py

@@ -3,8 +3,9 @@
 import dataclasses
 from copy import copy
 from dataclasses import dataclass, fields
-from typing import Annotated
+from typing import Annotated, Any
 
+from genai_prices.data_snapshot import get_snapshot
 from pydantic import AliasChoices, BeforeValidator, Field
 from typing_extensions import deprecated, overload
 
@@ -120,6 +121,14 @@ def __add__(self, other: RequestUsage) -> RequestUsage:
         new_usage.incr(other)
         return new_usage
 
+    @classmethod
+    def extract(
+        cls, data: Any, *, provider_id: str, api_flavor: str | None = None, details: dict[str, Any] | None = None
+    ) -> RequestUsage:
+        provider = get_snapshot().find_provider(None, provider_id, None)
+        _model_ref, extracted_usage = provider.extract_usage(data, api_flavor=api_flavor)
+        return cls(**{k: v for k, v in extracted_usage.__dict__.items() if v is not None}, details=details or {})
+
 
 @dataclass(repr=False, kw_only=True)
 class RunUsage(UsageBase):

tests/models/test_anthropic.py

@@ -1730,14 +1730,14 @@ def anth_msg(usage: BetaUsage) -> BetaMessage:
 def test_usage(
     message_callback: Callable[[], BetaMessage | BetaRawMessageStartEvent | BetaRawMessageDeltaEvent], usage: RunUsage
 ):
-    assert _map_usage(message_callback(), 'anthropic', 'claude-sonnet-4-5') == usage
+    assert _map_usage(message_callback(), 'anthropic', 'unknown') == usage
 
 
 def test_streaming_usage():
     start = BetaRawMessageStartEvent(message=anth_msg(BetaUsage(input_tokens=1, output_tokens=1)), type='message_start')
-    initial_usage = _map_usage(start, 'anthropic', 'claude-sonnet-4-5')
+    initial_usage = _map_usage(start, 'anthropic', 'unknown')
     delta = BetaRawMessageDeltaEvent(delta=Delta(), usage=BetaMessageDeltaUsage(output_tokens=5), type='message_delta')
-    final_usage = _map_usage(delta, 'anthropic', 'claude-sonnet-4-5', existing_usage=initial_usage)
+    final_usage = _map_usage(delta, 'anthropic', 'unknown', existing_usage=initial_usage)
     assert final_usage == snapshot(
         RequestUsage(input_tokens=1, output_tokens=5, details={'input_tokens': 1, 'output_tokens': 5})
     )