Stream tool calls and structured output from Anthropic as it's received instead of in one go (#1669)

Author: DouweM

pydantic_ai_slim/pydantic_ai/models/anthropic.py

@@ -5,7 +5,6 @@
 from contextlib import asynccontextmanager
 from dataclasses import dataclass, field
 from datetime import datetime, timezone
-from json import JSONDecodeError, loads as json_loads
 from typing import Any, Literal, Union, cast, overload
 
 from typing_extensions import assert_never
@@ -440,7 +439,6 @@ class AnthropicStreamedResponse(StreamedResponse):
 
     async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
         current_block: ContentBlock | None = None
-        current_json: str = ''
 
         async for event in self._response:
             self._usage += _map_usage(event)
@@ -455,7 +453,7 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
                     maybe_event = self._parts_manager.handle_tool_call_delta(
                         vendor_part_id=current_block.id,
                         tool_name=current_block.name,
-                        args=cast(dict[str, Any], current_block.input),
+                        args=cast(dict[str, Any], current_block.input) or None,
                         tool_call_id=current_block.id,
                     )
                     if maybe_event is not None:  # pragma: no branch
@@ -469,20 +467,10 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
                 elif (  # pragma: no branch
                     current_block and event.delta.type == 'input_json_delta' and isinstance(current_block, ToolUseBlock)
                 ):
-                    # Try to parse the JSON immediately, otherwise cache the value for later. This handles
-                    # cases where the JSON is not currently valid but will be valid once we stream more tokens.
-                    try:
-                        parsed_args = json_loads(current_json + event.delta.partial_json)
-                        current_json = ''
-                    except JSONDecodeError:
-                        current_json += event.delta.partial_json
-                        continue
-
-                    # For tool calls, we need to handle partial JSON updates
                     maybe_event = self._parts_manager.handle_tool_call_delta(
                         vendor_part_id=current_block.id,
                         tool_name='',
-                        args=parsed_args,
+                        args=event.delta.partial_json,
                         tool_call_id=current_block.id,
                     )
                     if maybe_event is not None:  # pragma: no branch

tests/models/test_anthropic.py

@@ -583,18 +583,23 @@ async def test_stream_structured(allow_model_requests: None):
         RawContentBlockStartEvent(
             type='content_block_start',
             index=0,
-            content_block=ToolUseBlock(type='tool_use', id='tool_1', name='my_tool', input={'first': 'One'}),
+            content_block=ToolUseBlock(type='tool_use', id='tool_1', name='my_tool', input={}),
         ),
         # Add more data through an incomplete JSON delta
         RawContentBlockDeltaEvent(
             type='content_block_delta',
             index=0,
-            delta=InputJSONDelta(type='input_json_delta', partial_json='{"second":'),
+            delta=InputJSONDelta(type='input_json_delta', partial_json='{"first": "One'),
         ),
         RawContentBlockDeltaEvent(
             type='content_block_delta',
             index=0,
-            delta=InputJSONDelta(type='input_json_delta', partial_json='"Two"}'),
+            delta=InputJSONDelta(type='input_json_delta', partial_json='", "second": "Two"'),
+        ),
+        RawContentBlockDeltaEvent(
+            type='content_block_delta',
+            index=0,
+            delta=InputJSONDelta(type='input_json_delta', partial_json='}'),
         ),
         # Mark tool block as complete
         RawContentBlockStopEvent(type='content_block_stop', index=0),