pydantic-ai-slim 0.0.19__py3-none-any.whl → 0.0.21__py3-none-any.whl

pydantic_ai/messages.py

@@ -6,7 +6,6 @@ from typing import Annotated, Any, Literal, Union, cast, overload
 
 import pydantic
 import pydantic_core
-from typing_extensions import Self, assert_never
 
 from ._utils import now_utc as _now_utc
 from .exceptions import UnexpectedModelBehavior
@@ -168,22 +167,6 @@ class TextPart:
         return bool(self.content)
 
 
-@dataclass
-class ArgsJson:
-    """Tool arguments as a JSON string."""
-
-    args_json: str
-    """A JSON string of arguments."""
-
-
-@dataclass
-class ArgsDict:
-    """Tool arguments as a Python dictionary."""
-
-    args_dict: dict[str, Any]
-    """A python dictionary of arguments."""
-
-
 @dataclass
 class ToolCallPart:
     """A tool call from a model."""
@@ -191,10 +174,10 @@ class ToolCallPart:
     tool_name: str
     """The name of the tool to call."""
 
-    args: ArgsJson | ArgsDict
+    args: str | dict[str, Any]
     """The arguments to pass to the tool.
 
-    Either as JSON or a Python dictionary depending on how data was returned.
+    This is stored either as a JSON string or a Python dictionary depending on how data was received.
     """
 
     tool_call_id: str | None = None
@@ -203,24 +186,14 @@ class ToolCallPart:
     part_kind: Literal['tool-call'] = 'tool-call'
     """Part type identifier, this is available on all parts as a discriminator."""
 
-    @classmethod
-    def from_raw_args(cls, tool_name: str, args: str | dict[str, Any], tool_call_id: str | None = None) -> Self:
-        """Create a `ToolCallPart` from raw arguments, converting them to `ArgsJson` or `ArgsDict`."""
-        if isinstance(args, str):
-            return cls(tool_name, ArgsJson(args), tool_call_id)
-        elif isinstance(args, dict):
-            return cls(tool_name, ArgsDict(args), tool_call_id)
-        else:
-            assert_never(args)
-
     def args_as_dict(self) -> dict[str, Any]:
         """Return the arguments as a Python dictionary.
 
         This is just for convenience with models that require dicts as input.
         """
-        if isinstance(self.args, ArgsDict):
-            return self.args.args_dict
-        args = pydantic_core.from_json(self.args.args_json)
+        if isinstance(self.args, dict):
+            return self.args
+        args = pydantic_core.from_json(self.args)
         assert isinstance(args, dict), 'args should be a dict'
         return cast(dict[str, Any], args)
 
@@ -229,16 +202,18 @@ class ToolCallPart:
 
         This is just for convenience with models that require JSON strings as input.
         """
-        if isinstance(self.args, ArgsJson):
-            return self.args.args_json
-        return pydantic_core.to_json(self.args.args_dict).decode()
+        if isinstance(self.args, str):
+            return self.args
+        return pydantic_core.to_json(self.args).decode()
 
     def has_content(self) -> bool:
         """Return `True` if the arguments contain any data."""
-        if isinstance(self.args, ArgsDict):
-            return any(self.args.args_dict.values())
+        if isinstance(self.args, dict):
+            # TODO: This should probably return True if you have the value False, or 0, etc.
+            #   It makes sense to me to ignore empty strings, but not sure about empty lists or dicts
+            return any(self.args.values())
         else:
-            return bool(self.args.args_json)
+            return bool(self.args)
 
 
 ModelResponsePart = Annotated[Union[TextPart, ToolCallPart], pydantic.Discriminator('part_kind')]
@@ -252,6 +227,9 @@ class ModelResponse:
     parts: list[ModelResponsePart]
     """The parts of the model message."""
 
+    model_name: str | None = None
+    """The name of the model that generated the response."""
+
     timestamp: datetime = field(default_factory=_now_utc)
     """The timestamp of the response.
 
@@ -261,16 +239,6 @@ class ModelResponse:
     kind: Literal['response'] = 'response'
     """Message type identifier, this is available on all parts as a discriminator."""
 
-    @classmethod
-    def from_text(cls, content: str, timestamp: datetime | None = None) -> Self:
-        """Create a `ModelResponse` containing a single `TextPart`."""
-        return cls([TextPart(content=content)], timestamp=timestamp or _now_utc())
-
-    @classmethod
-    def from_tool_call(cls, tool_call: ToolCallPart) -> Self:
-        """Create a `ModelResponse` containing a single `ToolCallPart`."""
-        return cls([tool_call])
-
 
 ModelMessage = Annotated[Union[ModelRequest, ModelResponse], pydantic.Discriminator('kind')]
 """Any message sent to or returned by a model."""
@@ -338,7 +306,7 @@ class ToolCallPartDelta:
         if self.tool_name_delta is None or self.args_delta is None:
             return None
 
-        return ToolCallPart.from_raw_args(
+        return ToolCallPart(
             self.tool_name_delta,
             self.args_delta,
             self.tool_call_id,
@@ -403,7 +371,7 @@ class ToolCallPartDelta:
 
         # If we now have enough data to create a full ToolCallPart, do so
         if delta.tool_name_delta is not None and delta.args_delta is not None:
-            return ToolCallPart.from_raw_args(
+            return ToolCallPart(
                 delta.tool_name_delta,
                 delta.args_delta,
                 delta.tool_call_id,
@@ -419,15 +387,15 @@ class ToolCallPartDelta:
             part = replace(part, tool_name=tool_name)
 
         if isinstance(self.args_delta, str):
-            if not isinstance(part.args, ArgsJson):
+            if not isinstance(part.args, str):
                 raise UnexpectedModelBehavior(f'Cannot apply JSON deltas to non-JSON tool arguments ({part=}, {self=})')
-            updated_json = part.args.args_json + self.args_delta
-            part = replace(part, args=ArgsJson(updated_json))
+            updated_json = part.args + self.args_delta
+            part = replace(part, args=updated_json)
         elif isinstance(self.args_delta, dict):
-            if not isinstance(part.args, ArgsDict):
+            if not isinstance(part.args, dict):
                 raise UnexpectedModelBehavior(f'Cannot apply dict deltas to non-dict tool arguments ({part=}, {self=})')
-            updated_dict = {**(part.args.args_dict or {}), **self.args_delta}
-            part = replace(part, args=ArgsDict(updated_dict))
+            updated_dict = {**(part.args or {}), **self.args_delta}
+            part = replace(part, args=updated_dict)
 
         if self.tool_call_id:
             # Replace the tool_call_id entirely if given

pydantic_ai/models/__init__.py

@@ -12,9 +12,10 @@ from contextlib import asynccontextmanager, contextmanager
 from dataclasses import dataclass, field
 from datetime import datetime
 from functools import cache
-from typing import TYPE_CHECKING, Literal
+from typing import TYPE_CHECKING
 
 import httpx
+from typing_extensions import Literal
 
 from .._parts_manager import ModelResponsePartsManager
 from ..exceptions import UserError
@@ -27,60 +28,123 @@ if TYPE_CHECKING:
 
 
 KnownModelName = Literal[
-    'openai:gpt-4o',
-    'openai:gpt-4o-mini',
-    'openai:gpt-4-turbo',
-    'openai:gpt-4',
-    'openai:o1-preview',
-    'openai:o1-mini',
-    'openai:o1',
-    'openai:gpt-3.5-turbo',
-    'groq:llama-3.3-70b-versatile',
-    'groq:llama-3.1-70b-versatile',
-    'groq:llama3-groq-70b-8192-tool-use-preview',
-    'groq:llama3-groq-8b-8192-tool-use-preview',
-    'groq:llama-3.1-70b-specdec',
+    'anthropic:claude-3-5-haiku-latest',
+    'anthropic:claude-3-5-sonnet-latest',
+    'anthropic:claude-3-opus-latest',
+    'claude-3-5-haiku-latest',
+    'claude-3-5-sonnet-latest',
+    'claude-3-opus-latest',
+    'cohere:c4ai-aya-expanse-32b',
+    'cohere:c4ai-aya-expanse-8b',
+    'cohere:command',
+    'cohere:command-light',
+    'cohere:command-light-nightly',
+    'cohere:command-nightly',
+    'cohere:command-r',
+    'cohere:command-r-03-2024',
+    'cohere:command-r-08-2024',
+    'cohere:command-r-plus',
+    'cohere:command-r-plus-04-2024',
+    'cohere:command-r-plus-08-2024',
+    'cohere:command-r7b-12-2024',
+    'google-gla:gemini-1.0-pro',
+    'google-gla:gemini-1.5-flash',
+    'google-gla:gemini-1.5-flash-8b',
+    'google-gla:gemini-1.5-pro',
+    'google-gla:gemini-2.0-flash-exp',
+    'google-vertex:gemini-1.0-pro',
+    'google-vertex:gemini-1.5-flash',
+    'google-vertex:gemini-1.5-flash-8b',
+    'google-vertex:gemini-1.5-pro',
+    'google-vertex:gemini-2.0-flash-exp',
+    'gpt-3.5-turbo',
+    'gpt-3.5-turbo-0125',
+    'gpt-3.5-turbo-0301',
+    'gpt-3.5-turbo-0613',
+    'gpt-3.5-turbo-1106',
+    'gpt-3.5-turbo-16k',
+    'gpt-3.5-turbo-16k-0613',
+    'gpt-4',
+    'gpt-4-0125-preview',
+    'gpt-4-0314',
+    'gpt-4-0613',
+    'gpt-4-1106-preview',
+    'gpt-4-32k',
+    'gpt-4-32k-0314',
+    'gpt-4-32k-0613',
+    'gpt-4-turbo',
+    'gpt-4-turbo-2024-04-09',
+    'gpt-4-turbo-preview',
+    'gpt-4-vision-preview',
+    'gpt-4o',
+    'gpt-4o-2024-05-13',
+    'gpt-4o-2024-08-06',
+    'gpt-4o-2024-11-20',
+    'gpt-4o-audio-preview',
+    'gpt-4o-audio-preview-2024-10-01',
+    'gpt-4o-audio-preview-2024-12-17',
+    'gpt-4o-mini',
+    'gpt-4o-mini-2024-07-18',
+    'gpt-4o-mini-audio-preview',
+    'gpt-4o-mini-audio-preview-2024-12-17',
+    'groq:gemma2-9b-it',
     'groq:llama-3.1-8b-instant',
+    'groq:llama-3.2-11b-vision-preview',
     'groq:llama-3.2-1b-preview',
     'groq:llama-3.2-3b-preview',
-    'groq:llama-3.2-11b-vision-preview',
     'groq:llama-3.2-90b-vision-preview',
+    'groq:llama-3.3-70b-specdec',
+    'groq:llama-3.3-70b-versatile',
     'groq:llama3-70b-8192',
     'groq:llama3-8b-8192',
     'groq:mixtral-8x7b-32768',
-    'groq:gemma2-9b-it',
-    'groq:gemma-7b-it',
-    'google-gla:gemini-1.5-flash',
-    'google-gla:gemini-1.5-pro',
-    'google-gla:gemini-2.0-flash-exp',
-    'google-vertex:gemini-1.5-flash',
-    'google-vertex:gemini-1.5-pro',
-    'google-vertex:gemini-2.0-flash-exp',
-    'mistral:mistral-small-latest',
-    'mistral:mistral-large-latest',
     'mistral:codestral-latest',
+    'mistral:mistral-large-latest',
     'mistral:mistral-moderation-latest',
-    'ollama:codellama',
-    'ollama:gemma',
-    'ollama:gemma2',
-    'ollama:llama3',
-    'ollama:llama3.1',
-    'ollama:llama3.2',
-    'ollama:llama3.2-vision',
-    'ollama:llama3.3',
-    'ollama:mistral',
-    'ollama:mistral-nemo',
-    'ollama:mixtral',
-    'ollama:phi3',
-    'ollama:phi4',
-    'ollama:qwq',
-    'ollama:qwen',
-    'ollama:qwen2',
-    'ollama:qwen2.5',
-    'ollama:starcoder2',
-    'anthropic:claude-3-5-haiku-latest',
-    'anthropic:claude-3-5-sonnet-latest',
-    'anthropic:claude-3-opus-latest',
+    'mistral:mistral-small-latest',
+    'o1',
+    'o1-2024-12-17',
+    'o1-mini',
+    'o1-mini-2024-09-12',
+    'o1-preview',
+    'o1-preview-2024-09-12',
+    'openai:chatgpt-4o-latest',
+    'openai:gpt-3.5-turbo',
+    'openai:gpt-3.5-turbo-0125',
+    'openai:gpt-3.5-turbo-0301',
+    'openai:gpt-3.5-turbo-0613',
+    'openai:gpt-3.5-turbo-1106',
+    'openai:gpt-3.5-turbo-16k',
+    'openai:gpt-3.5-turbo-16k-0613',
+    'openai:gpt-4',
+    'openai:gpt-4-0125-preview',
+    'openai:gpt-4-0314',
+    'openai:gpt-4-0613',
+    'openai:gpt-4-1106-preview',
+    'openai:gpt-4-32k',
+    'openai:gpt-4-32k-0314',
+    'openai:gpt-4-32k-0613',
+    'openai:gpt-4-turbo',
+    'openai:gpt-4-turbo-2024-04-09',
+    'openai:gpt-4-turbo-preview',
+    'openai:gpt-4-vision-preview',
+    'openai:gpt-4o',
+    'openai:gpt-4o-2024-05-13',
+    'openai:gpt-4o-2024-08-06',
+    'openai:gpt-4o-2024-11-20',
+    'openai:gpt-4o-audio-preview',
+    'openai:gpt-4o-audio-preview-2024-10-01',
+    'openai:gpt-4o-audio-preview-2024-12-17',
+    'openai:gpt-4o-mini',
+    'openai:gpt-4o-mini-2024-07-18',
+    'openai:gpt-4o-mini-audio-preview',
+    'openai:gpt-4o-mini-audio-preview-2024-12-17',
+    'openai:o1',
+    'openai:o1-2024-12-17',
+    'openai:o1-mini',
+    'openai:o1-mini-2024-09-12',
+    'openai:o1-preview',
+    'openai:o1-preview-2024-09-12',
     'test',
 ]
 """Known model names that can be used with the `model` parameter of [`Agent`][pydantic_ai.Agent].
@@ -145,6 +209,7 @@ class AgentModel(ABC):
 class StreamedResponse(ABC):
     """Streamed response from an LLM when calling a tool."""
 
+    _model_name: str
     _usage: Usage = field(default_factory=Usage, init=False)
     _parts_manager: ModelResponsePartsManager = field(default_factory=ModelResponsePartsManager, init=False)
     _event_iterator: AsyncIterator[ModelResponseStreamEvent] | None = field(default=None, init=False)
@@ -168,7 +233,13 @@ class StreamedResponse(ABC):
 
     def get(self) -> ModelResponse:
         """Build a [`ModelResponse`][pydantic_ai.messages.ModelResponse] from the data received from the stream so far."""
-        return ModelResponse(parts=self._parts_manager.get_parts(), timestamp=self.timestamp())
+        return ModelResponse(
+            parts=self._parts_manager.get_parts(), model_name=self._model_name, timestamp=self.timestamp()
+        )
+
+    def model_name(self) -> str:
+        """Get the model name of the response."""
+        return self._model_name
 
     def usage(self) -> Usage:
         """Get the usage of the response so far. This will not be the final usage until the stream is exhausted."""
@@ -228,6 +299,10 @@ def infer_model(model: Model | KnownModelName) -> Model:
         from .test import TestModel
 
         return TestModel()
+    elif model.startswith('cohere:'):
+        from .cohere import CohereModel
+
+        return CohereModel(model[7:])
     elif model.startswith('openai:'):
         from .openai import OpenAIModel
 
@@ -263,10 +338,6 @@ def infer_model(model: Model | KnownModelName) -> Model:
         from .mistral import MistralModel
 
         return MistralModel(model[8:])
-    elif model.startswith('ollama:'):
-        from .ollama import OllamaModel
-
-        return OllamaModel(model[7:])
     elif model.startswith('anthropic'):
         from .anthropic import AnthropicModel
 

pydantic_ai/models/anthropic.py

@@ -1,21 +1,23 @@
 from __future__ import annotations as _annotations
 
-from collections.abc import AsyncIterator
+from collections.abc import AsyncIterable, AsyncIterator
 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 httpx import AsyncClient as AsyncHTTPClient
 from typing_extensions import assert_never
 
-from .. import usage
+from .. import UnexpectedModelBehavior, _utils, usage
 from .._utils import guard_tool_call_id as _guard_tool_call_id
 from ..messages import (
-    ArgsDict,
     ModelMessage,
     ModelRequest,
     ModelResponse,
     ModelResponsePart,
+    ModelResponseStreamEvent,
     RetryPromptPart,
     SystemPromptPart,
     TextPart,
@@ -38,11 +40,17 @@ try:
     from anthropic.types import (
         Message as AnthropicMessage,
         MessageParam,
+        MetadataParam,
+        RawContentBlockDeltaEvent,
+        RawContentBlockStartEvent,
+        RawContentBlockStopEvent,
         RawMessageDeltaEvent,
         RawMessageStartEvent,
+        RawMessageStopEvent,
         RawMessageStreamEvent,
         TextBlock,
         TextBlockParam,
+        TextDelta,
         ToolChoiceParam,
         ToolParam,
         ToolResultBlockParam,
@@ -71,6 +79,15 @@ Since [the Anthropic docs](https://docs.anthropic.com/en/docs/about-claude/model
 """
 
 
+class AnthropicModelSettings(ModelSettings):
+    """Settings used for an Anthropic model request."""
+
+    anthropic_metadata: MetadataParam
+    """An object describing metadata about the request.
+
+    Contains `user_id`, an external identifier for the user who is associated with the request."""
+
+
 @dataclass(init=False)
 class AnthropicModel(Model):
     """A model that uses the Anthropic API.
@@ -152,50 +169,54 @@ class AnthropicAgentModel(AgentModel):
     """Implementation of `AgentModel` for Anthropic models."""
 
     client: AsyncAnthropic
-    model_name: str
+    model_name: AnthropicModelName
     allow_text_result: bool
     tools: list[ToolParam]
 
     async def request(
         self, messages: list[ModelMessage], model_settings: ModelSettings | None
     ) -> tuple[ModelResponse, usage.Usage]:
-        response = await self._messages_create(messages, False, model_settings)
+        response = await self._messages_create(messages, False, cast(AnthropicModelSettings, model_settings or {}))
         return self._process_response(response), _map_usage(response)
 
     @asynccontextmanager
     async def request_stream(
         self, messages: list[ModelMessage], model_settings: ModelSettings | None
     ) -> AsyncIterator[StreamedResponse]:
-        response = await self._messages_create(messages, True, model_settings)
+        response = await self._messages_create(messages, True, cast(AnthropicModelSettings, model_settings or {}))
         async with response:
             yield await self._process_streamed_response(response)
 
     @overload
     async def _messages_create(
-        self, messages: list[ModelMessage], stream: Literal[True], model_settings: ModelSettings | None
+        self, messages: list[ModelMessage], stream: Literal[True], model_settings: AnthropicModelSettings
     ) -> AsyncStream[RawMessageStreamEvent]:
         pass
 
     @overload
     async def _messages_create(
-        self, messages: list[ModelMessage], stream: Literal[False], model_settings: ModelSettings | None
+        self, messages: list[ModelMessage], stream: Literal[False], model_settings: AnthropicModelSettings
     ) -> AnthropicMessage:
         pass
 
     async def _messages_create(
-        self, messages: list[ModelMessage], stream: bool, model_settings: ModelSettings | None
+        self, messages: list[ModelMessage], stream: bool, model_settings: AnthropicModelSettings
     ) -> AnthropicMessage | AsyncStream[RawMessageStreamEvent]:
         # standalone function to make it easier to override
+        tool_choice: ToolChoiceParam | None
+
         if not self.tools:
-            tool_choice: ToolChoiceParam | None = None
-        elif not self.allow_text_result:
-            tool_choice = {'type': 'any'}
+            tool_choice = None
         else:
-            tool_choice = {'type': 'auto'}
+            if not self.allow_text_result:
+                tool_choice = {'type': 'any'}
+            else:
+                tool_choice = {'type': 'auto'}
 
-        system_prompt, anthropic_messages = self._map_message(messages)
+            if (allow_parallel_tool_calls := model_settings.get('parallel_tool_calls')) is not None:
+                tool_choice['disable_parallel_tool_use'] = not allow_parallel_tool_calls
 
-        model_settings = model_settings or {}
+        system_prompt, anthropic_messages = self._map_message(messages)
 
         return await self.client.messages.create(
             max_tokens=model_settings.get('max_tokens', 1024),
@@ -208,10 +229,10 @@ class AnthropicAgentModel(AgentModel):
             temperature=model_settings.get('temperature', NOT_GIVEN),
             top_p=model_settings.get('top_p', NOT_GIVEN),
             timeout=model_settings.get('timeout', NOT_GIVEN),
+            metadata=model_settings.get('anthropic_metadata', NOT_GIVEN),
         )
 
-    @staticmethod
-    def _process_response(response: AnthropicMessage) -> ModelResponse:
+    def _process_response(self, response: AnthropicMessage) -> ModelResponse:
         """Process a non-streamed response, and prepare a message to return."""
         items: list[ModelResponsePart] = []
         for item in response.content:
@@ -220,33 +241,24 @@ class AnthropicAgentModel(AgentModel):
             else:
                 assert isinstance(item, ToolUseBlock), 'unexpected item type'
                 items.append(
-                    ToolCallPart.from_raw_args(
+                    ToolCallPart(
                         tool_name=item.name,
                         args=cast(dict[str, Any], item.input),
                         tool_call_id=item.id,
                     )
                 )
 
-        return ModelResponse(items)
+        return ModelResponse(items, model_name=self.model_name)
 
-    @staticmethod
-    async def _process_streamed_response(response: AsyncStream[RawMessageStreamEvent]) -> StreamedResponse:
-        """TODO: Process a streamed response, and prepare a streaming response to return."""
-        # We don't yet support streamed responses from Anthropic, so we raise an error here for now.
-        # Streamed responses will be supported in a future release.
-
-        raise RuntimeError('Streamed responses are not yet supported for Anthropic models.')
-
-        # Should be returning some sort of AnthropicStreamTextResponse or AnthropicStreamedResponse
-        # depending on the type of chunk we get, but we need to establish how we handle (and when we get) the following:
-        # RawMessageStartEvent
-        # RawMessageDeltaEvent
-        # RawMessageStopEvent
-        # RawContentBlockStartEvent
-        # RawContentBlockDeltaEvent
-        # RawContentBlockDeltaEvent
-        #
-        # We might refactor streaming internally before we implement this...
+    async def _process_streamed_response(self, response: AsyncStream[RawMessageStreamEvent]) -> StreamedResponse:
+        peekable_response = _utils.PeekableAsyncStream(response)
+        first_chunk = await peekable_response.peek()
+        if isinstance(first_chunk, _utils.Unset):
+            raise UnexpectedModelBehavior('Streamed response ended without content or tool calls')
+
+        # Since Anthropic doesn't provide a timestamp in the message, we'll use the current time
+        timestamp = datetime.now(tz=timezone.utc)
+        return AnthropicStreamedResponse(_model_name=self.model_name, _response=peekable_response, _timestamp=timestamp)
 
     @staticmethod
     def _map_message(messages: list[ModelMessage]) -> tuple[str, list[MessageParam]]:
@@ -306,7 +318,6 @@ class AnthropicAgentModel(AgentModel):
 
 
 def _map_tool_call(t: ToolCallPart) -> ToolUseBlockParam:
-    assert isinstance(t.args, ArgsDict), f'Expected ArgsDict, got {t.args}'
     return ToolUseBlockParam(
         id=_guard_tool_call_id(t=t, model_source='Anthropic'),
         type='tool_use',
@@ -342,3 +353,63 @@ def _map_usage(message: AnthropicMessage | RawMessageStreamEvent) -> usage.Usage
         response_tokens=response_usage.output_tokens,
         total_tokens=(request_tokens or 0) + response_usage.output_tokens,
     )
+
+
+@dataclass
+class AnthropicStreamedResponse(StreamedResponse):
+    """Implementation of `StreamedResponse` for Anthropic models."""
+
+    _response: AsyncIterable[RawMessageStreamEvent]
+    _timestamp: datetime
+
+    async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
+        current_block: TextBlock | ToolUseBlock | None = None
+        current_json: str = ''
+
+        async for event in self._response:
+            self._usage += _map_usage(event)
+
+            if isinstance(event, RawContentBlockStartEvent):
+                current_block = event.content_block
+                if isinstance(current_block, TextBlock) and current_block.text:
+                    yield self._parts_manager.handle_text_delta(vendor_part_id='content', content=current_block.text)
+                elif isinstance(current_block, ToolUseBlock):
+                    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),
+                        tool_call_id=current_block.id,
+                    )
+                    if maybe_event is not None:
+                        yield maybe_event
+
+            elif isinstance(event, RawContentBlockDeltaEvent):
+                if isinstance(event.delta, TextDelta):
+                    yield self._parts_manager.handle_text_delta(vendor_part_id='content', content=event.delta.text)
+                elif (
+                    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,
+                        tool_call_id=current_block.id,
+                    )
+                    if maybe_event is not None:
+                        yield maybe_event
+
+            elif isinstance(event, (RawContentBlockStopEvent, RawMessageStopEvent)):
+                current_block = None
+
+    def timestamp(self) -> datetime:
+        return self._timestamp