pass actual model

Author: alexmojaki

pydantic_ai_slim/pydantic_ai/models/anthropic.py

@@ -352,7 +352,7 @@ def _process_response(self, response: BetaMessage) -> ModelResponse:
 
         return ModelResponse(
             parts=items,
-            usage=_map_usage(response, self._provider.name),
+            usage=_map_usage(response, self._provider.name, self._model_name),
             model_name=response.model,
             provider_response_id=response.id,
             provider_name=self._provider.name,
@@ -619,7 +619,8 @@ def _map_tool_definition(f: ToolDefinition) -> BetaToolParam:
 
 def _map_usage(
     message: BetaMessage | BetaRawMessageStartEvent | BetaRawMessageDeltaEvent,
-    provider: str = 'anthropic',
+    provider: str,
+    model: str,
     existing_usage: usage.RequestUsage | None = None,
 ) -> usage.RequestUsage:
     if isinstance(message, BetaMessage):
@@ -637,10 +638,7 @@ def _map_usage(
         key: value for key, value in response_usage.model_dump().items() if isinstance(value, int)
     }
 
-    # `extract_usage` expects a mapping with a `model` and `usage` key.
-    # Not all the actual types of messages here have a model so we just make a dummy message.
-    # We only care about the numbers of tokens etc.
-    extracted_usage = extract_usage(dict(model='claude-sonnet-4-5', usage=details), provider_id=provider)
+    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)},
@@ -662,7 +660,7 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
 
         async for event in self._response:
             if isinstance(event, BetaRawMessageStartEvent):
-                self._usage = _map_usage(event, self._provider_name)
+                self._usage = _map_usage(event, self._provider_name, self._model_name)
                 self.provider_response_id = event.message.id
 
             elif isinstance(event, BetaRawContentBlockStartEvent):
@@ -743,7 +741,7 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
                     pass
 
             elif isinstance(event, BetaRawMessageDeltaEvent):
-                self._usage = _map_usage(event, self._provider_name, self._usage)
+                self._usage = _map_usage(event, self._provider_name, self._model_name, self._usage)
                 if raw_finish_reason := event.delta.stop_reason:  # pragma: no branch
                     self.provider_details = {'finish_reason': raw_finish_reason}
                     self.finish_reason = _FINISH_REASON_MAP.get(raw_finish_reason)

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()) == usage
+    assert _map_usage(message_callback(), 'anthropic', 'claude-sonnet-4-5') == 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)
+    initial_usage = _map_usage(start, 'anthropic', 'claude-sonnet-4-5')
     delta = BetaRawMessageDeltaEvent(delta=Delta(), usage=BetaMessageDeltaUsage(output_tokens=5), type='message_delta')
-    final_usage = _map_usage(delta, existing_usage=initial_usage)
+    final_usage = _map_usage(delta, 'anthropic', 'claude-sonnet-4-5', existing_usage=initial_usage)
     assert final_usage == snapshot(
         RequestUsage(input_tokens=1, output_tokens=5, details={'input_tokens': 1, 'output_tokens': 5})
     )