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

pydantic_ai/agent.py

@@ -26,6 +26,7 @@ from .result import ResultData
 from .settings import ModelSettings, merge_model_settings
 from .tools import (
     AgentDeps,
+    DocstringFormat,
     RunContext,
     Tool,
     ToolDefinition,
@@ -242,9 +243,10 @@ class Agent(Generic[AgentDeps, ResultData]):
 
         agent = Agent('openai:gpt-4o')
 
-        result_sync = agent.run_sync('What is the capital of Italy?')
-        print(result_sync.data)
-        #> Rome
+        async def main():
+            result = await agent.run('What is the capital of France?')
+            print(result.data)
+            #> Paris
         ```
 
         Args:
@@ -382,10 +384,9 @@ class Agent(Generic[AgentDeps, ResultData]):
 
         agent = Agent('openai:gpt-4o')
 
-        async def main():
-            result = await agent.run('What is the capital of France?')
-            print(result.data)
-            #> Paris
+        result_sync = agent.run_sync('What is the capital of Italy?')
+        print(result_sync.data)
+        #> Rome
         ```
 
         Args:
@@ -535,7 +536,7 @@ class Agent(Generic[AgentDeps, ResultData]):
                         model_req_span.__exit__(None, None, None)
 
                         with _logfire.span('handle model response') as handle_span:
-                            maybe_final_result = await self._handle_streamed_model_response(
+                            maybe_final_result = await self._handle_streamed_response(
                                 model_response, run_context, result_schema
                             )
 
@@ -774,6 +775,8 @@ class Agent(Generic[AgentDeps, ResultData]):
         *,
         retries: int | None = None,
         prepare: ToolPrepareFunc[AgentDeps] | None = None,
+        docstring_format: DocstringFormat = 'auto',
+        require_parameter_descriptions: bool = False,
     ) -> Callable[[ToolFuncContext[AgentDeps, ToolParams]], ToolFuncContext[AgentDeps, ToolParams]]: ...
 
     def tool(
@@ -783,6 +786,8 @@ class Agent(Generic[AgentDeps, ResultData]):
         *,
         retries: int | None = None,
         prepare: ToolPrepareFunc[AgentDeps] | None = None,
+        docstring_format: DocstringFormat = 'auto',
+        require_parameter_descriptions: bool = False,
     ) -> Any:
         """Decorator to register a tool function which takes [`RunContext`][pydantic_ai.tools.RunContext] as its first argument.
 
@@ -820,6 +825,9 @@ class Agent(Generic[AgentDeps, ResultData]):
             prepare: custom method to prepare the tool definition for each step, return `None` to omit this
                 tool from a given step. This is useful if you want to customise a tool at call time,
                 or omit it completely from a step. See [`ToolPrepareFunc`][pydantic_ai.tools.ToolPrepareFunc].
+            docstring_format: The format of the docstring, see [`DocstringFormat`][pydantic_ai.tools.DocstringFormat].
+                Defaults to `'auto'`, such that the format is inferred from the structure of the docstring.
+            require_parameter_descriptions: If True, raise an error if a parameter description is missing. Defaults to False.
         """
         if func is None:
 
@@ -827,13 +835,13 @@ class Agent(Generic[AgentDeps, ResultData]):
                 func_: ToolFuncContext[AgentDeps, ToolParams],
             ) -> ToolFuncContext[AgentDeps, ToolParams]:
                 # noinspection PyTypeChecker
-                self._register_function(func_, True, retries, prepare)
+                self._register_function(func_, True, retries, prepare, docstring_format, require_parameter_descriptions)
                 return func_
 
             return tool_decorator
         else:
             # noinspection PyTypeChecker
-            self._register_function(func, True, retries, prepare)
+            self._register_function(func, True, retries, prepare, docstring_format, require_parameter_descriptions)
             return func
 
     @overload
@@ -846,6 +854,8 @@ class Agent(Generic[AgentDeps, ResultData]):
         *,
         retries: int | None = None,
         prepare: ToolPrepareFunc[AgentDeps] | None = None,
+        docstring_format: DocstringFormat = 'auto',
+        require_parameter_descriptions: bool = False,
     ) -> Callable[[ToolFuncPlain[ToolParams]], ToolFuncPlain[ToolParams]]: ...
 
     def tool_plain(
@@ -855,6 +865,8 @@ class Agent(Generic[AgentDeps, ResultData]):
         *,
         retries: int | None = None,
         prepare: ToolPrepareFunc[AgentDeps] | None = None,
+        docstring_format: DocstringFormat = 'auto',
+        require_parameter_descriptions: bool = False,
     ) -> Any:
         """Decorator to register a tool function which DOES NOT take `RunContext` as an argument.
 
@@ -892,17 +904,22 @@ class Agent(Generic[AgentDeps, ResultData]):
             prepare: custom method to prepare the tool definition for each step, return `None` to omit this
                 tool from a given step. This is useful if you want to customise a tool at call time,
                 or omit it completely from a step. See [`ToolPrepareFunc`][pydantic_ai.tools.ToolPrepareFunc].
+            docstring_format: The format of the docstring, see [`DocstringFormat`][pydantic_ai.tools.DocstringFormat].
+                Defaults to `'auto'`, such that the format is inferred from the structure of the docstring.
+            require_parameter_descriptions: If True, raise an error if a parameter description is missing. Defaults to False.
         """
         if func is None:
 
             def tool_decorator(func_: ToolFuncPlain[ToolParams]) -> ToolFuncPlain[ToolParams]:
                 # noinspection PyTypeChecker
-                self._register_function(func_, False, retries, prepare)
+                self._register_function(
+                    func_, False, retries, prepare, docstring_format, require_parameter_descriptions
+                )
                 return func_
 
             return tool_decorator
         else:
-            self._register_function(func, False, retries, prepare)
+            self._register_function(func, False, retries, prepare, docstring_format, require_parameter_descriptions)
             return func
 
     def _register_function(
@@ -911,10 +928,19 @@ class Agent(Generic[AgentDeps, ResultData]):
         takes_ctx: bool,
         retries: int | None,
         prepare: ToolPrepareFunc[AgentDeps] | None,
+        docstring_format: DocstringFormat,
+        require_parameter_descriptions: bool,
     ) -> None:
         """Private utility to register a function as a tool."""
         retries_ = retries if retries is not None else self._default_retries
-        tool = Tool(func, takes_ctx=takes_ctx, max_retries=retries_, prepare=prepare)
+        tool = Tool(
+            func,
+            takes_ctx=takes_ctx,
+            max_retries=retries_,
+            prepare=prepare,
+            docstring_format=docstring_format,
+            require_parameter_descriptions=require_parameter_descriptions,
+        )
         self._register_tool(tool)
 
     def _register_tool(self, tool: Tool[AgentDeps]) -> None:
@@ -1100,7 +1126,7 @@ class Agent(Generic[AgentDeps, ResultData]):
         final_result: _MarkFinalResult[RunResultData] | None = None
 
         parts: list[_messages.ModelRequestPart] = []
-        if result_schema := result_schema:
+        if result_schema is not None:
             if match := result_schema.find_tool(tool_calls):
                 call, result_tool = match
                 try:
@@ -1179,76 +1205,58 @@ class Agent(Generic[AgentDeps, ResultData]):
                 parts.extend(task_results)
         return parts
 
-    async def _handle_streamed_model_response(
+    async def _handle_streamed_response(
         self,
-        model_response: models.EitherStreamedResponse,
+        streamed_response: models.StreamedResponse,
         run_context: RunContext[AgentDeps],
         result_schema: _result.ResultSchema[RunResultData] | None,
-    ) -> (
-        _MarkFinalResult[models.EitherStreamedResponse]
-        | tuple[_messages.ModelResponse, list[_messages.ModelRequestPart]]
-    ):
+    ) -> _MarkFinalResult[models.StreamedResponse] | tuple[_messages.ModelResponse, list[_messages.ModelRequestPart]]:
         """Process a streamed response from the model.
 
         Returns:
             Either a final result or a tuple of the model response and the tool responses for the next request.
             If a final result is returned, the conversation should end.
         """
-        if isinstance(model_response, models.StreamTextResponse):
-            # plain string response
-            if self._allow_text_result(result_schema):
-                return _MarkFinalResult(model_response, None)
-            else:
-                self._incr_result_retry(run_context)
-                response = _messages.RetryPromptPart(
-                    content='Plain text responses are not permitted, please call one of the functions instead.',
-                )
-                # stream the response, so usage is correct
-                async for _ in model_response:
-                    pass
-
-                text = ''.join(model_response.get(final=True))
-                return _messages.ModelResponse([_messages.TextPart(text)]), [response]
-        elif isinstance(model_response, models.StreamStructuredResponse):
-            if result_schema is not None:
-                # if there's a result schema, iterate over the stream until we find at least one tool
-                # NOTE: this means we ignore any other tools called here
-                structured_msg = model_response.get()
-                while not structured_msg.parts:
-                    try:
-                        await model_response.__anext__()
-                    except StopAsyncIteration:
-                        break
-                    structured_msg = model_response.get()
-
-                if match := result_schema.find_tool(structured_msg.parts):
-                    call, _ = match
-                    return _MarkFinalResult(model_response, call.tool_name)
-
-            # the model is calling a tool function, consume the response to get the next message
-            async for _ in model_response:
-                pass
-            model_response_msg = model_response.get()
-            if not model_response_msg.parts:
-                raise exceptions.UnexpectedModelBehavior('Received empty tool call message')
-
-            # we now run all tool functions in parallel
-            tasks: list[asyncio.Task[_messages.ModelRequestPart]] = []
-            parts: list[_messages.ModelRequestPart] = []
-            for item in model_response_msg.parts:
-                if isinstance(item, _messages.ToolCallPart):
-                    call = item
-                    if tool := self._function_tools.get(call.tool_name):
-                        tasks.append(asyncio.create_task(tool.run(call, run_context), name=call.tool_name))
-                    else:
-                        parts.append(self._unknown_tool(call.tool_name, run_context, result_schema))
+        received_text = False
+
+        async for maybe_part_event in streamed_response:
+            if isinstance(maybe_part_event, _messages.PartStartEvent):
+                new_part = maybe_part_event.part
+                if isinstance(new_part, _messages.TextPart):
+                    received_text = True
+                    if self._allow_text_result(result_schema):
+                        return _MarkFinalResult(streamed_response, None)
+                elif isinstance(new_part, _messages.ToolCallPart):
+                    if result_schema is not None and (match := result_schema.find_tool([new_part])):
+                        call, _ = match
+                        return _MarkFinalResult(streamed_response, call.tool_name)
+                else:
+                    assert_never(new_part)
 
-            with _logfire.span('running {tools=}', tools=[t.get_name() for t in tasks]):
-                task_results: Sequence[_messages.ModelRequestPart] = await asyncio.gather(*tasks)
-                parts.extend(task_results)
-            return model_response_msg, parts
-        else:
-            assert_never(model_response)
+        tasks: list[asyncio.Task[_messages.ModelRequestPart]] = []
+        parts: list[_messages.ModelRequestPart] = []
+        model_response = streamed_response.get()
+        if not model_response.parts:
+            raise exceptions.UnexpectedModelBehavior('Received empty model response')
+        for p in model_response.parts:
+            if isinstance(p, _messages.ToolCallPart):
+                if tool := self._function_tools.get(p.tool_name):
+                    tasks.append(asyncio.create_task(tool.run(p, run_context), name=p.tool_name))
+                else:
+                    parts.append(self._unknown_tool(p.tool_name, run_context, result_schema))
+
+        if received_text and not tasks and not parts:
+            # Can only get here if self._allow_text_result returns `False` for the provided result_schema
+            self._incr_result_retry(run_context)
+            model_response = _messages.RetryPromptPart(
+                content='Plain text responses are not permitted, please call one of the functions instead.',
+            )
+            return streamed_response.get(), [model_response]
+
+        with _logfire.span('running {tools=}', tools=[t.get_name() for t in tasks]):
+            task_results: Sequence[_messages.ModelRequestPart] = await asyncio.gather(*tasks)
+            parts.extend(task_results)
+        return model_response, parts
 
     async def _validate_result(
         self,

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:
@@ -254,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."""