pydantic-ai-slim 0.0.18__py3-none-any.whl → 0.0.20__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
@@ -72,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]
@@ -124,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:
@@ -159,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:
@@ -197,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):
@@ -226,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:
@@ -243,6 +252,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.
 
@@ -252,19 +264,209 @@ 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:
-        return cls([TextPart(content)], timestamp=timestamp or _now_utc())
 
-    @classmethod
-    def from_tool_call(cls, tool_call: ToolCallPart) -> Self:
-        return cls([tool_call])
+ModelMessage = Annotated[Union[ModelRequest, ModelResponse], pydantic.Discriminator('kind')]
+"""Any message sent to or returned by a model."""
 
+ModelMessagesTypeAdapter = pydantic.TypeAdapter(list[ModelMessage], config=pydantic.ConfigDict(defer_build=True))
+"""Pydantic [`TypeAdapter`][pydantic.type_adapter.TypeAdapter] for (de)serializing messages."""
 
-ModelMessage = Union[ModelRequest, ModelResponse]
-"""Any message send to or returned by a model."""
 
-ModelMessagesTypeAdapter = pydantic.TypeAdapter(
-    list[Annotated[ModelMessage, pydantic.Discriminator('kind')]], 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
 
 
@@ -59,6 +61,7 @@ KnownModelName = Literal[
     'mistral:codestral-latest',
     'mistral:mistral-moderation-latest',
     'ollama:codellama',
+    'ollama:deepseek-r1',
     'ollama:gemma',
     'ollama:gemma2',
     'ollama:llama3',
@@ -70,6 +73,7 @@ KnownModelName = Literal[
     'ollama:mistral-nemo',
     'ollama:mixtral',
     'ollama:phi3',
+    'ollama:phi4',
     'ollama:qwq',
     'ollama:qwen',
     'ollama:qwen2',
@@ -78,6 +82,22 @@ KnownModelName = Literal[
     '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',
     'test',
 ]
 """Known model names that can be used with the `model` parameter of [`Agent`][pydantic_ai.Agent].
@@ -129,88 +149,54 @@ 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()
+@dataclass
+class StreamedResponse(ABC):
+    """Streamed response from an LLM when calling a tool."""
 
-    @abstractmethod
-    def usage(self) -> Usage:
-        """Return the usage of the request.
+    _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)
 
-        NOTE: this won't return the full usage until the stream is finished.
-        """
-        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 timestamp(self) -> datetime:
-        """Get the timestamp of the response."""
-        raise NotImplementedError()
+    async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
+        """Return an async iterator of [`ModelResponseStreamEvent`][pydantic_ai.messages.ModelResponseStreamEvent]s.
 
-
-class StreamStructuredResponse(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.
+        This method should be implemented by subclasses to translate the vendor-specific stream of events into
+        pydantic_ai-format events.
         """
-        return self
-
-    @abstractmethod
-    async def __anext__(self) -> None:
-        """Process the next chunk of the response, see above for why this returns `None`."""
         raise NotImplementedError()
+        # noinspection PyUnreachableCode
+        yield
 
-    @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.
+    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(), model_name=self._model_name, timestamp=self.timestamp()
+        )
 
-        Args:
-            final: If True, this is the final call, after iteration is complete, the response should be fully validated.
-        """
-        raise NotImplementedError()
+    def model_name(self) -> str:
+        """Get the model name of the response."""
+        return self._model_name
 
-    @abstractmethod
     def usage(self) -> Usage:
-        """Get the usage of the request.
-
-        NOTE: this won't return the full usage until the stream is finished.
-        """
-        raise NotImplementedError()
+        """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:
@@ -218,9 +204,6 @@ class StreamStructuredResponse(ABC):
         raise NotImplementedError()
 
 
-EitherStreamedResponse = Union[StreamTextResponse, StreamStructuredResponse]
-
-
 ALLOW_MODEL_REQUESTS = True
 """Whether to allow requests to models.
 
@@ -269,6 +252,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