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

pydantic_ai/format_as_xml.py

@@ -37,7 +37,8 @@ def format_as_xml(
         none_str: String to use for `None` values.
         indent: Indentation string to use for pretty printing.
 
-    Returns: XML representation of the object.
+    Returns:
+        XML representation of the object.
 
     Example:
     ```python {title="format_as_xml_example.py" lint="skip"}

pydantic_ai/messages.py

@@ -1,14 +1,15 @@
 from __future__ import annotations as _annotations
 
-from dataclasses import dataclass, field
+from dataclasses import dataclass, field, replace
 from datetime import datetime
-from typing import Annotated, Any, Literal, Union, cast
+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
 
 
 @dataclass
@@ -21,6 +22,12 @@ class SystemPromptPart:
     content: str
     """The content of the prompt."""
 
+    dynamic_ref: str | None = None
+    """The ref of the dynamic system prompt function that generated this part.
+
+    Only set if system prompt is dynamic, see [`system_prompt`][pydantic_ai.Agent.system_prompt] for more information.
+    """
+
     part_kind: Literal['system-prompt'] = 'system-prompt'
     """Part type identifier, this is available on all parts as a discriminator."""
 
@@ -66,12 +73,14 @@ class ToolReturnPart:
     """Part type identifier, this is available on all parts as a discriminator."""
 
     def model_response_str(self) -> str:
+        """Return a string representation of the content for the model."""
         if isinstance(self.content, str):
             return self.content
         else:
             return tool_return_ta.dump_json(self.content).decode()
 
     def model_response_object(self) -> dict[str, Any]:
+        """Return a dictionary representation of the content, wrapping non-dict types appropriately."""
         # gemini supports JSON dict return values, but no other JSON types, hence we wrap anything else in a dict
         if isinstance(self.content, dict):
             return tool_return_ta.dump_python(self.content, mode='json')  # pyright: ignore[reportUnknownMemberType]
@@ -118,6 +127,7 @@ class RetryPromptPart:
     """Part type identifier, this is available on all parts as a discriminator."""
 
     def model_response(self) -> str:
+        """Return a string message describing why the retry is requested."""
         if isinstance(self.content, str):
             description = self.content
         else:
@@ -153,6 +163,10 @@ class TextPart:
     part_kind: Literal['text'] = 'text'
     """Part type identifier, this is available on all parts as a discriminator."""
 
+    def has_content(self) -> bool:
+        """Return `True` if the text content is non-empty."""
+        return bool(self.content)
+
 
 @dataclass
 class ArgsJson:
@@ -191,7 +205,7 @@ class ToolCallPart:
 
     @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."""
+        """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):
@@ -220,6 +234,7 @@ class ToolCallPart:
         return pydantic_core.to_json(self.args.args_dict).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())
         else:
@@ -248,17 +263,217 @@ class ModelResponse:
 
     @classmethod
     def from_text(cls, content: str, timestamp: datetime | None = None) -> Self:
-        return cls([TextPart(content)], timestamp=timestamp or _now_utc())
+        """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 = Union[ModelRequest, ModelResponse]
-"""Any message send to or returned by a model."""
+ModelMessage = Annotated[Union[ModelRequest, ModelResponse], pydantic.Discriminator('kind')]
+"""Any message sent to or returned by a model."""
 
-ModelMessagesTypeAdapter = pydantic.TypeAdapter(
-    list[Annotated[ModelMessage, pydantic.Discriminator('kind')]], config=pydantic.ConfigDict(defer_build=True)
-)
+ModelMessagesTypeAdapter = pydantic.TypeAdapter(list[ModelMessage], config=pydantic.ConfigDict(defer_build=True))
 """Pydantic [`TypeAdapter`][pydantic.type_adapter.TypeAdapter] for (de)serializing messages."""
+
+
+@dataclass
+class TextPartDelta:
+    """A partial update (delta) for a `TextPart` to append new text content."""
+
+    content_delta: str
+    """The incremental text content to add to the existing `TextPart` content."""
+
+    part_delta_kind: Literal['text'] = 'text'
+    """Part delta type identifier, used as a discriminator."""
+
+    def apply(self, part: ModelResponsePart) -> TextPart:
+        """Apply this text delta to an existing `TextPart`.
+
+        Args:
+            part: The existing model response part, which must be a `TextPart`.
+
+        Returns:
+            A new `TextPart` with updated text content.
+
+        Raises:
+            ValueError: If `part` is not a `TextPart`.
+        """
+        if not isinstance(part, TextPart):
+            raise ValueError('Cannot apply TextPartDeltas to non-TextParts')
+        return replace(part, content=part.content + self.content_delta)
+
+
+@dataclass
+class ToolCallPartDelta:
+    """A partial update (delta) for a `ToolCallPart` to modify tool name, arguments, or tool call ID."""
+
+    tool_name_delta: str | None = None
+    """Incremental text to add to the existing tool name, if any."""
+
+    args_delta: str | dict[str, Any] | None = None
+    """Incremental data to add to the tool arguments.
+
+    If this is a string, it will be appended to existing JSON arguments.
+    If this is a dict, it will be merged with existing dict arguments.
+    """
+
+    tool_call_id: str | None = None
+    """Optional tool call identifier, this is used by some models including OpenAI.
+
+    Note this is never treated as a delta — it can replace None, but otherwise if a
+    non-matching value is provided an error will be raised."""
+
+    part_delta_kind: Literal['tool_call'] = 'tool_call'
+    """Part delta type identifier, used as a discriminator."""
+
+    def as_part(self) -> ToolCallPart | None:
+        """Convert this delta to a fully formed `ToolCallPart` if possible, otherwise return `None`.
+
+        Returns:
+            A `ToolCallPart` if both `tool_name_delta` and `args_delta` are set, otherwise `None`.
+        """
+        if self.tool_name_delta is None or self.args_delta is None:
+            return None
+
+        return ToolCallPart.from_raw_args(
+            self.tool_name_delta,
+            self.args_delta,
+            self.tool_call_id,
+        )
+
+    @overload
+    def apply(self, part: ModelResponsePart) -> ToolCallPart: ...
+
+    @overload
+    def apply(self, part: ModelResponsePart | ToolCallPartDelta) -> ToolCallPart | ToolCallPartDelta: ...
+
+    def apply(self, part: ModelResponsePart | ToolCallPartDelta) -> ToolCallPart | ToolCallPartDelta:
+        """Apply this delta to a part or delta, returning a new part or delta with the changes applied.
+
+        Args:
+            part: The existing model response part or delta to update.
+
+        Returns:
+            Either a new `ToolCallPart` or an updated `ToolCallPartDelta`.
+
+        Raises:
+            ValueError: If `part` is neither a `ToolCallPart` nor a `ToolCallPartDelta`.
+            UnexpectedModelBehavior: If applying JSON deltas to dict arguments or vice versa.
+        """
+        if isinstance(part, ToolCallPart):
+            return self._apply_to_part(part)
+
+        if isinstance(part, ToolCallPartDelta):
+            return self._apply_to_delta(part)
+
+        raise ValueError(f'Can only apply ToolCallPartDeltas to ToolCallParts or ToolCallPartDeltas, not {part}')
+
+    def _apply_to_delta(self, delta: ToolCallPartDelta) -> ToolCallPart | ToolCallPartDelta:
+        """Internal helper to apply this delta to another delta."""
+        if self.tool_name_delta:
+            # Append incremental text to the existing tool_name_delta
+            updated_tool_name_delta = (delta.tool_name_delta or '') + self.tool_name_delta
+            delta = replace(delta, tool_name_delta=updated_tool_name_delta)
+
+        if isinstance(self.args_delta, str):
+            if isinstance(delta.args_delta, dict):
+                raise UnexpectedModelBehavior(
+                    f'Cannot apply JSON deltas to non-JSON tool arguments ({delta=}, {self=})'
+                )
+            updated_args_delta = (delta.args_delta or '') + self.args_delta
+            delta = replace(delta, args_delta=updated_args_delta)
+        elif isinstance(self.args_delta, dict):
+            if isinstance(delta.args_delta, str):
+                raise UnexpectedModelBehavior(
+                    f'Cannot apply dict deltas to non-dict tool arguments ({delta=}, {self=})'
+                )
+            updated_args_delta = {**(delta.args_delta or {}), **self.args_delta}
+            delta = replace(delta, args_delta=updated_args_delta)
+
+        if self.tool_call_id:
+            # Set the tool_call_id if it wasn't present, otherwise error if it has changed
+            if delta.tool_call_id is not None and delta.tool_call_id != self.tool_call_id:
+                raise UnexpectedModelBehavior(
+                    f'Cannot apply a new tool_call_id to a ToolCallPartDelta that already has one ({delta=}, {self=})'
+                )
+            delta = replace(delta, tool_call_id=self.tool_call_id)
+
+        # 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(
+                delta.tool_name_delta,
+                delta.args_delta,
+                delta.tool_call_id,
+            )
+
+        return delta
+
+    def _apply_to_part(self, part: ToolCallPart) -> ToolCallPart:
+        """Internal helper to apply this delta directly to a `ToolCallPart`."""
+        if self.tool_name_delta:
+            # Append incremental text to the existing tool_name
+            tool_name = part.tool_name + self.tool_name_delta
+            part = replace(part, tool_name=tool_name)
+
+        if isinstance(self.args_delta, str):
+            if not isinstance(part.args, ArgsJson):
+                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))
+        elif isinstance(self.args_delta, dict):
+            if not isinstance(part.args, ArgsDict):
+                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))
+
+        if self.tool_call_id:
+            # Replace the tool_call_id entirely if given
+            if part.tool_call_id is not None and part.tool_call_id != self.tool_call_id:
+                raise UnexpectedModelBehavior(
+                    f'Cannot apply a new tool_call_id to a ToolCallPartDelta that already has one ({part=}, {self=})'
+                )
+            part = replace(part, tool_call_id=self.tool_call_id)
+        return part
+
+
+ModelResponsePartDelta = Annotated[Union[TextPartDelta, ToolCallPartDelta], pydantic.Discriminator('part_delta_kind')]
+"""A partial update (delta) for any model response part."""
+
+
+@dataclass
+class PartStartEvent:
+    """An event indicating that a new part has started.
+
+    If multiple `PartStartEvent`s are received with the same index,
+    the new one should fully replace the old one.
+    """
+
+    index: int
+    """The index of the part within the overall response parts list."""
+
+    part: ModelResponsePart
+    """The newly started `ModelResponsePart`."""
+
+    event_kind: Literal['part_start'] = 'part_start'
+    """Event type identifier, used as a discriminator."""
+
+
+@dataclass
+class PartDeltaEvent:
+    """An event indicating a delta update for an existing part."""
+
+    index: int
+    """The index of the part within the overall response parts list."""
+
+    delta: ModelResponsePartDelta
+    """The delta to apply to the specified part."""
+
+    event_kind: Literal['part_delta'] = 'part_delta'
+    """Event type identifier, used as a discriminator."""
+
+
+ModelResponseStreamEvent = Annotated[Union[PartStartEvent, PartDeltaEvent], pydantic.Discriminator('event_kind')]
+"""An event in the model response stream, either starting a new part or applying a delta to an existing one."""

pydantic_ai/models/__init__.py

@@ -7,20 +7,22 @@ specific LLM being used.
 from __future__ import annotations as _annotations
 
 from abc import ABC, abstractmethod
-from collections.abc import AsyncIterator, Iterable, Iterator
+from collections.abc import AsyncIterator, Iterator
 from contextlib import asynccontextmanager, contextmanager
+from dataclasses import dataclass, field
 from datetime import datetime
 from functools import cache
-from typing import TYPE_CHECKING, Literal, Union
+from typing import TYPE_CHECKING, Literal
 
 import httpx
 
+from .._parts_manager import ModelResponsePartsManager
 from ..exceptions import UserError
-from ..messages import ModelMessage, ModelResponse
+from ..messages import ModelMessage, ModelResponse, ModelResponseStreamEvent
 from ..settings import ModelSettings
+from ..usage import Usage
 
 if TYPE_CHECKING:
-    from ..result import Usage
     from ..tools import ToolDefinition
 
 
@@ -48,13 +50,12 @@ KnownModelName = Literal[
     'groq:mixtral-8x7b-32768',
     'groq:gemma2-9b-it',
     'groq:gemma-7b-it',
-    'gemini-1.5-flash',
-    'gemini-1.5-pro',
-    'gemini-2.0-flash-exp',
-    'vertexai:gemini-1.5-flash',
-    'vertexai:gemini-1.5-pro',
-    # since mistral models are supported by other providers (e.g. ollama), and some of their models (e.g. "codestral")
-    # don't start with "mistral", we add the "mistral:" prefix to all to be explicit
+    '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',
@@ -71,14 +72,15 @@ KnownModelName = Literal[
     'ollama:mistral-nemo',
     'ollama:mixtral',
     'ollama:phi3',
+    'ollama:phi4',
     'ollama:qwq',
     'ollama:qwen',
     'ollama:qwen2',
     'ollama:qwen2.5',
     'ollama:starcoder2',
-    'claude-3-5-haiku-latest',
-    'claude-3-5-sonnet-latest',
-    'claude-3-opus-latest',
+    'anthropic:claude-3-5-haiku-latest',
+    'anthropic:claude-3-5-sonnet-latest',
+    'anthropic:claude-3-opus-latest',
     'test',
 ]
 """Known model names that can be used with the `model` parameter of [`Agent`][pydantic_ai.Agent].
@@ -130,88 +132,47 @@ class AgentModel(ABC):
     @asynccontextmanager
     async def request_stream(
         self, messages: list[ModelMessage], model_settings: ModelSettings | None
-    ) -> AsyncIterator[EitherStreamedResponse]:
+    ) -> AsyncIterator[StreamedResponse]:
         """Make a request to the model and return a streaming response."""
+        # This method is not required, but you need to implement it if you want to support streamed responses
         raise NotImplementedError(f'Streamed requests not supported by this {self.__class__.__name__}')
         # yield is required to make this a generator for type checking
         # noinspection PyUnreachableCode
         yield  # pragma: no cover
 
 
-class StreamTextResponse(ABC):
-    """Streamed response from an LLM when returning text."""
-
-    def __aiter__(self) -> AsyncIterator[None]:
-        """Stream the response as an async iterable, building up the text as it goes.
-
-        This is an async iterator that yields `None` to avoid doing the work of validating the input and
-        extracting the text field when it will often be thrown away.
-        """
-        return self
-
-    @abstractmethod
-    async def __anext__(self) -> None:
-        """Process the next chunk of the response, see above for why this returns `None`."""
-        raise NotImplementedError()
-
-    @abstractmethod
-    def get(self, *, final: bool = False) -> Iterable[str]:
-        """Returns an iterable of text since the last call to `get()` — e.g. the text delta.
-
-        Args:
-            final: If True, this is the final call, after iteration is complete, the response should be fully validated
-                and all text extracted.
-        """
-        raise NotImplementedError()
-
-    @abstractmethod
-    def usage(self) -> Usage:
-        """Return the usage of the request.
-
-        NOTE: this won't return the full usage until the stream is finished.
-        """
-        raise NotImplementedError()
-
-    @abstractmethod
-    def timestamp(self) -> datetime:
-        """Get the timestamp of the response."""
-        raise NotImplementedError()
-
-
-class StreamStructuredResponse(ABC):
+@dataclass
+class StreamedResponse(ABC):
     """Streamed response from an LLM when calling a tool."""
 
-    def __aiter__(self) -> AsyncIterator[None]:
-        """Stream the response as an async iterable, building up the tool call as it goes.
-
-        This is an async iterator that yields `None` to avoid doing the work of building the final tool call when
-        it will often be thrown away.
-        """
-        return self
+    _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)
 
-    @abstractmethod
-    async def __anext__(self) -> None:
-        """Process the next chunk of the response, see above for why this returns `None`."""
-        raise NotImplementedError()
+    def __aiter__(self) -> AsyncIterator[ModelResponseStreamEvent]:
+        """Stream the response as an async iterable of [`ModelResponseStreamEvent`][pydantic_ai.messages.ModelResponseStreamEvent]s."""
+        if self._event_iterator is None:
+            self._event_iterator = self._get_event_iterator()
+        return self._event_iterator
 
     @abstractmethod
-    def get(self, *, final: bool = False) -> ModelResponse:
-        """Get the `ModelResponse` at this point.
-
-        The `ModelResponse` may or may not be complete, depending on whether the stream is finished.
+    async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
+        """Return an async iterator of [`ModelResponseStreamEvent`][pydantic_ai.messages.ModelResponseStreamEvent]s.
 
-        Args:
-            final: If True, this is the final call, after iteration is complete, the response should be fully validated.
+        This method should be implemented by subclasses to translate the vendor-specific stream of events into
+        pydantic_ai-format events.
         """
         raise NotImplementedError()
+        # noinspection PyUnreachableCode
+        yield
 
-    @abstractmethod
-    def usage(self) -> Usage:
-        """Get the usage of the request.
+    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())
 
-        NOTE: this won't return the full usage until the stream is finished.
-        """
-        raise NotImplementedError()
+    def usage(self) -> Usage:
+        """Get the usage of the response so far. This will not be the final usage until the stream is exhausted."""
+        return self._usage
 
     @abstractmethod
     def timestamp(self) -> datetime:
@@ -219,9 +180,6 @@ class StreamStructuredResponse(ABC):
         raise NotImplementedError()
 
 
-EitherStreamedResponse = Union[StreamTextResponse, StreamStructuredResponse]
-
-
 ALLOW_MODEL_REQUESTS = True
 """Whether to allow requests to models.
 
@@ -274,6 +232,15 @@ def infer_model(model: Model | KnownModelName) -> Model:
         from .openai import OpenAIModel
 
         return OpenAIModel(model[7:])
+    elif model.startswith(('gpt', 'o1')):
+        from .openai import OpenAIModel
+
+        return OpenAIModel(model)
+    elif model.startswith('google-gla'):
+        from .gemini import GeminiModel
+
+        return GeminiModel(model[11:])  # pyright: ignore[reportArgumentType]
+    # backwards compatibility with old model names (ex, gemini-1.5-flash -> google-gla:gemini-1.5-flash)
     elif model.startswith('gemini'):
         from .gemini import GeminiModel
 
@@ -283,6 +250,11 @@ def infer_model(model: Model | KnownModelName) -> Model:
         from .groq import GroqModel
 
         return GroqModel(model[5:])  # pyright: ignore[reportArgumentType]
+    elif model.startswith('google-vertex'):
+        from .vertexai import VertexAIModel
+
+        return VertexAIModel(model[14:])  # pyright: ignore[reportArgumentType]
+    # backwards compatibility with old model names (ex, vertexai:gemini-1.5-flash -> google-vertex:gemini-1.5-flash)
     elif model.startswith('vertexai:'):
         from .vertexai import VertexAIModel
 
@@ -295,6 +267,11 @@ def infer_model(model: Model | KnownModelName) -> Model:
         from .ollama import OllamaModel
 
         return OllamaModel(model[7:])
+    elif model.startswith('anthropic'):
+        from .anthropic import AnthropicModel
+
+        return AnthropicModel(model[10:])
+    # backwards compatibility with old model names (ex, claude-3-5-sonnet-latest -> anthropic:claude-3-5-sonnet-latest)
     elif model.startswith('claude'):
         from .anthropic import AnthropicModel
 

pydantic_ai/models/anthropic.py

@@ -8,7 +8,7 @@ from typing import Any, Literal, Union, cast, overload
 from httpx import AsyncClient as AsyncHTTPClient
 from typing_extensions import assert_never
 
-from .. import result
+from .. import usage
 from .._utils import guard_tool_call_id as _guard_tool_call_id
 from ..messages import (
     ArgsDict,
@@ -27,8 +27,8 @@ from ..settings import ModelSettings
 from ..tools import ToolDefinition
 from . import (
     AgentModel,
-    EitherStreamedResponse,
     Model,
+    StreamedResponse,
     cached_async_http_client,
     check_allow_model_requests,
 )
@@ -136,7 +136,7 @@ class AnthropicModel(Model):
         )
 
     def name(self) -> str:
-        return self.model_name
+        return f'anthropic:{self.model_name}'
 
     @staticmethod
     def _map_tool_definition(f: ToolDefinition) -> ToolParam:
@@ -158,14 +158,14 @@ class AnthropicAgentModel(AgentModel):
 
     async def request(
         self, messages: list[ModelMessage], model_settings: ModelSettings | None
-    ) -> tuple[ModelResponse, result.Usage]:
+    ) -> tuple[ModelResponse, usage.Usage]:
         response = await self._messages_create(messages, False, model_settings)
         return self._process_response(response), _map_usage(response)
 
     @asynccontextmanager
     async def request_stream(
         self, messages: list[ModelMessage], model_settings: ModelSettings | None
-    ) -> AsyncIterator[EitherStreamedResponse]:
+    ) -> AsyncIterator[StreamedResponse]:
         response = await self._messages_create(messages, True, model_settings)
         async with response:
             yield await self._process_streamed_response(response)
@@ -216,28 +216,28 @@ class AnthropicAgentModel(AgentModel):
         items: list[ModelResponsePart] = []
         for item in response.content:
             if isinstance(item, TextBlock):
-                items.append(TextPart(item.text))
+                items.append(TextPart(content=item.text))
             else:
                 assert isinstance(item, ToolUseBlock), 'unexpected item type'
                 items.append(
                     ToolCallPart.from_raw_args(
-                        item.name,
-                        cast(dict[str, Any], item.input),
-                        item.id,
+                        tool_name=item.name,
+                        args=cast(dict[str, Any], item.input),
+                        tool_call_id=item.id,
                     )
                 )
 
         return ModelResponse(items)
 
     @staticmethod
-    async def _process_streamed_response(response: AsyncStream[RawMessageStreamEvent]) -> EitherStreamedResponse:
+    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 AnthropicStreamStructuredResponse
+        # 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
@@ -315,30 +315,30 @@ def _map_tool_call(t: ToolCallPart) -> ToolUseBlockParam:
     )
 
 
-def _map_usage(message: AnthropicMessage | RawMessageStreamEvent) -> result.Usage:
+def _map_usage(message: AnthropicMessage | RawMessageStreamEvent) -> usage.Usage:
     if isinstance(message, AnthropicMessage):
-        usage = message.usage
+        response_usage = message.usage
     else:
         if isinstance(message, RawMessageStartEvent):
-            usage = message.message.usage
+            response_usage = message.message.usage
         elif isinstance(message, RawMessageDeltaEvent):
-            usage = message.usage
+            response_usage = message.usage
         else:
             # No usage information provided in:
             # - RawMessageStopEvent
             # - RawContentBlockStartEvent
             # - RawContentBlockDeltaEvent
             # - RawContentBlockStopEvent
-            usage = None
+            response_usage = None
 
-    if usage is None:
-        return result.Usage()
+    if response_usage is None:
+        return usage.Usage()
 
-    request_tokens = getattr(usage, 'input_tokens', None)
+    request_tokens = getattr(response_usage, 'input_tokens', None)
 
-    return result.Usage(
+    return usage.Usage(
         # Usage coming from the RawMessageDeltaEvent doesn't have input token data, hence this getattr
         request_tokens=request_tokens,
-        response_tokens=usage.output_tokens,
-        total_tokens=(request_tokens or 0) + usage.output_tokens,
+        response_tokens=response_usage.output_tokens,
+        total_tokens=(request_tokens or 0) + response_usage.output_tokens,
     )