Fix Anthropic streaming (#1686)

Author: samuelcolvin

pydantic_ai_slim/pydantic_ai/models/anthropic.py

@@ -393,36 +393,31 @@ def _map_tool_definition(f: ToolDefinition) -> ToolParam:
 def _map_usage(message: AnthropicMessage | RawMessageStreamEvent) -> usage.Usage:
     if isinstance(message, AnthropicMessage):
         response_usage = message.usage
+    elif isinstance(message, RawMessageStartEvent):
+        response_usage = message.message.usage
+    elif isinstance(message, RawMessageDeltaEvent):
+        response_usage = message.usage
     else:
-        if isinstance(message, RawMessageStartEvent):
-            response_usage = message.message.usage
-        elif isinstance(message, RawMessageDeltaEvent):
-            response_usage = message.usage
-        else:
-            # No usage information provided in:
-            # - RawMessageStopEvent
-            # - RawContentBlockStartEvent
-            # - RawContentBlockDeltaEvent
-            # - RawContentBlockStopEvent
-            response_usage = None
-
-    if response_usage is None:
+        # No usage information provided in:
+        # - RawMessageStopEvent
+        # - RawContentBlockStartEvent
+        # - RawContentBlockDeltaEvent
+        # - RawContentBlockStopEvent
         return usage.Usage()
 
-    # Store all integer-typed usage values in the details dict
-    response_usage_dict = response_usage.model_dump()
-    details: dict[str, int] = {}
-    for key, value in response_usage_dict.items():
-        if isinstance(value, int):
-            details[key] = value
+    # Store all integer-typed usage values in the details, except 'output_tokens' which is represented exactly by
+    # `response_tokens`
+    details: dict[str, int] = {
+        key: value for key, value in response_usage.model_dump().items() if isinstance(value, int)
+    }
 
-    # Usage coming from the RawMessageDeltaEvent doesn't have input token data, hence the getattr call
+    # Usage coming from the RawMessageDeltaEvent doesn't have input token data, hence using `get`
     # Tokens are only counted once between input_tokens, cache_creation_input_tokens, and cache_read_input_tokens
     # This approach maintains request_tokens as the count of all input tokens, with cached counts as details
     request_tokens = (
-        getattr(response_usage, 'input_tokens', 0)
-        + (getattr(response_usage, 'cache_creation_input_tokens', 0) or 0)  # These can be missing, None, or int
-        + (getattr(response_usage, 'cache_read_input_tokens', 0) or 0)
+        details.get('input_tokens', 0)
+        + details.get('cache_creation_input_tokens', 0)
+        + details.get('cache_read_input_tokens', 0)
     )
 
     return usage.Usage(

tests/models/test_anthropic.py

@@ -6,7 +6,7 @@
 from dataclasses import dataclass, field
 from datetime import timezone
 from functools import cached_property
-from typing import Any, TypeVar, Union, cast
+from typing import Any, Callable, TypeVar, Union, cast
 
 import httpx
 import pytest
@@ -52,7 +52,11 @@
     )
     from anthropic.types.raw_message_delta_event import Delta
 
-    from pydantic_ai.models.anthropic import AnthropicModel, AnthropicModelSettings
+    from pydantic_ai.models.anthropic import (
+        AnthropicModel,
+        AnthropicModelSettings,
+        _map_usage,  # pyright: ignore[reportPrivateUsage]
+    )
     from pydantic_ai.providers.anthropic import AnthropicProvider
 
     # note: we use Union here so that casting works with Python 3.9
@@ -921,3 +925,74 @@ def simple_instructions():
             ),
         ]
     )
+
+
+def anth_msg(usage: AnthropicUsage) -> AnthropicMessage:
+    return AnthropicMessage(
+        id='x',
+        content=[],
+        model='claude-3-7-sonnet-latest',
+        role='assistant',
+        type='message',
+        usage=usage,
+    )
+
+
+@pytest.mark.parametrize(
+    'message_callback,usage',
+    [
+        pytest.param(
+            lambda: anth_msg(AnthropicUsage(input_tokens=1, output_tokens=1)),
+            snapshot(
+                Usage(
+                    request_tokens=1, response_tokens=1, total_tokens=2, details={'input_tokens': 1, 'output_tokens': 1}
+                )
+            ),
+            id='AnthropicMessage',
+        ),
+        pytest.param(
+            lambda: anth_msg(
+                AnthropicUsage(
+                    input_tokens=1, output_tokens=1, cache_creation_input_tokens=2, cache_read_input_tokens=3
+                )
+            ),
+            snapshot(
+                Usage(
+                    request_tokens=6,
+                    response_tokens=1,
+                    total_tokens=7,
+                    details={
+                        'cache_creation_input_tokens': 2,
+                        'cache_read_input_tokens': 3,
+                        'input_tokens': 1,
+                        'output_tokens': 1,
+                    },
+                )
+            ),
+            id='AnthropicMessage-cached',
+        ),
+        pytest.param(
+            lambda: RawMessageStartEvent(
+                message=anth_msg(AnthropicUsage(input_tokens=1, output_tokens=1)), type='message_start'
+            ),
+            snapshot(
+                Usage(
+                    request_tokens=1, response_tokens=1, total_tokens=2, details={'input_tokens': 1, 'output_tokens': 1}
+                )
+            ),
+            id='RawMessageStartEvent',
+        ),
+        pytest.param(
+            lambda: RawMessageDeltaEvent(
+                delta=Delta(),
+                usage=MessageDeltaUsage(output_tokens=5),
+                type='message_delta',
+            ),
+            snapshot(Usage(response_tokens=5, total_tokens=5, details={'output_tokens': 5})),
+            id='RawMessageDeltaEvent',
+        ),
+        pytest.param(lambda: RawMessageStopEvent(type='message_stop'), snapshot(Usage()), id='RawMessageStopEvent'),
+    ],
+)
+def test_usage(message_callback: Callable[[], AnthropicMessage | RawMessageStreamEvent], usage: Usage):
+    assert _map_usage(message_callback()) == usage