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

pydantic_ai/_griffe.py

@@ -1,16 +1,24 @@
 from __future__ import annotations as _annotations
 
+import logging
 import re
+from contextlib import contextmanager
 from inspect import Signature
-from typing import Any, Callable, Literal, cast
+from typing import TYPE_CHECKING, Any, Callable, Literal, cast
 
 from griffe import Docstring, DocstringSectionKind, Object as GriffeObject
 
+if TYPE_CHECKING:
+    from .tools import DocstringFormat
+
 DocstringStyle = Literal['google', 'numpy', 'sphinx']
 
 
 def doc_descriptions(
-    func: Callable[..., Any], sig: Signature, *, style: DocstringStyle | None = None
+    func: Callable[..., Any],
+    sig: Signature,
+    *,
+    docstring_format: DocstringFormat,
 ) -> tuple[str, dict[str, str]]:
     """Extract the function description and parameter descriptions from a function's docstring.
 
@@ -24,8 +32,10 @@ def doc_descriptions(
     # see https://github.com/mkdocstrings/griffe/issues/293
     parent = cast(GriffeObject, sig)
 
-    docstring = Docstring(doc, lineno=1, parser=style or _infer_docstring_style(doc), parent=parent)
-    sections = docstring.parse()
+    docstring_style = _infer_docstring_style(doc) if docstring_format == 'auto' else docstring_format
+    docstring = Docstring(doc, lineno=1, parser=docstring_style, parent=parent)
+    with _disable_griffe_logging():
+        sections = docstring.parse()
 
     params = {}
     if parameters := next((p for p in sections if p.kind == DocstringSectionKind.parameters), None):
@@ -125,3 +135,12 @@ _docstring_style_patterns: list[tuple[str, list[str], DocstringStyle]] = [
         'numpy',
     ),
 ]
+
+
+@contextmanager
+def _disable_griffe_logging():
+    # Hacky, but suggested here: https://github.com/mkdocstrings/griffe/issues/293#issuecomment-2167668117
+    old_level = logging.root.getEffectiveLevel()
+    logging.root.setLevel(logging.ERROR)
+    yield
+    logging.root.setLevel(old_level)

pydantic_ai/_parts_manager.py

@@ -0,0 +1,239 @@
+"""This module provides functionality to manage and update parts of a model's streamed response.
+
+The manager tracks which parts (in particular, text and tool calls) correspond to which
+vendor-specific identifiers (e.g., `index`, `tool_call_id`, etc., as appropriate for a given model),
+and produces PydanticAI-format events as appropriate for consumers of the streaming APIs.
+
+The "vendor-specific identifiers" to use depend on the semantics of the responses of the responses from the vendor,
+and are tightly coupled to the specific model being used, and the PydanticAI Model subclass implementation.
+
+This `ModelResponsePartsManager` is used in each of the subclasses of `StreamedResponse` as a way to consolidate
+event-emitting logic.
+"""
+
+from __future__ import annotations as _annotations
+
+from collections.abc import Hashable
+from dataclasses import dataclass, field
+from typing import Any, Union
+
+from pydantic_ai.exceptions import UnexpectedModelBehavior
+from pydantic_ai.messages import (
+    ModelResponsePart,
+    ModelResponseStreamEvent,
+    PartDeltaEvent,
+    PartStartEvent,
+    TextPart,
+    TextPartDelta,
+    ToolCallPart,
+    ToolCallPartDelta,
+)
+
+VendorId = Hashable
+"""
+Type alias for a vendor identifier, which can be any hashable type (e.g., a string, UUID, etc.)
+"""
+
+ManagedPart = Union[ModelResponsePart, ToolCallPartDelta]
+"""
+A union of types that are managed by the ModelResponsePartsManager.
+Because many vendors have streaming APIs that may produce not-fully-formed tool calls,
+this includes ToolCallPartDelta's in addition to the more fully-formed ModelResponsePart's.
+"""
+
+
+@dataclass
+class ModelResponsePartsManager:
+    """Manages a sequence of parts that make up a model's streamed response.
+
+    Parts are generally added and/or updated by providing deltas, which are tracked by vendor-specific IDs.
+    """
+
+    _parts: list[ManagedPart] = field(default_factory=list, init=False)
+    """A list of parts (text or tool calls) that make up the current state of the model's response."""
+    _vendor_id_to_part_index: dict[VendorId, int] = field(default_factory=dict, init=False)
+    """Maps a vendor's "part" ID (if provided) to the index in `_parts` where that part resides."""
+
+    def get_parts(self) -> list[ModelResponsePart]:
+        """Return only model response parts that are complete (i.e., not ToolCallPartDelta's).
+
+        Returns:
+            A list of ModelResponsePart objects. ToolCallPartDelta objects are excluded.
+        """
+        return [p for p in self._parts if not isinstance(p, ToolCallPartDelta)]
+
+    def handle_text_delta(
+        self,
+        *,
+        vendor_part_id: Hashable | None,
+        content: str,
+    ) -> ModelResponseStreamEvent:
+        """Handle incoming text content, creating or updating a TextPart in the manager as appropriate.
+
+        When `vendor_part_id` is None, the latest part is updated if it exists and is a TextPart;
+        otherwise, a new TextPart is created. When a non-None ID is specified, the TextPart corresponding
+        to that vendor ID is either created or updated.
+
+        Args:
+            vendor_part_id: The ID the vendor uses to identify this piece
+                of text. If None, a new part will be created unless the latest part is already
+                a TextPart.
+            content: The text content to append to the appropriate TextPart.
+
+        Returns:
+            A `PartStartEvent` if a new part was created, or a `PartDeltaEvent` if an existing part was updated.
+
+        Raises:
+            UnexpectedModelBehavior: If attempting to apply text content to a part that is
+                not a TextPart.
+        """
+        existing_text_part_and_index: tuple[TextPart, int] | None = None
+
+        if vendor_part_id is None:
+            # If the vendor_part_id is None, check if the latest part is a TextPart to update
+            if self._parts:
+                part_index = len(self._parts) - 1
+                latest_part = self._parts[part_index]
+                if isinstance(latest_part, TextPart):
+                    existing_text_part_and_index = latest_part, part_index
+        else:
+            # Otherwise, attempt to look up an existing TextPart by vendor_part_id
+            part_index = self._vendor_id_to_part_index.get(vendor_part_id)
+            if part_index is not None:
+                existing_part = self._parts[part_index]
+                if not isinstance(existing_part, TextPart):
+                    raise UnexpectedModelBehavior(f'Cannot apply a text delta to {existing_part=}')
+                existing_text_part_and_index = existing_part, part_index
+
+        if existing_text_part_and_index is None:
+            # There is no existing text part that should be updated, so create a new one
+            new_part_index = len(self._parts)
+            part = TextPart(content=content)
+            if vendor_part_id is not None:
+                self._vendor_id_to_part_index[vendor_part_id] = new_part_index
+            self._parts.append(part)
+            return PartStartEvent(index=new_part_index, part=part)
+        else:
+            # Update the existing TextPart with the new content delta
+            existing_text_part, part_index = existing_text_part_and_index
+            part_delta = TextPartDelta(content_delta=content)
+            self._parts[part_index] = part_delta.apply(existing_text_part)
+            return PartDeltaEvent(index=part_index, delta=part_delta)
+
+    def handle_tool_call_delta(
+        self,
+        *,
+        vendor_part_id: Hashable | None,
+        tool_name: str | None,
+        args: str | dict[str, Any] | None,
+        tool_call_id: str | None,
+    ) -> ModelResponseStreamEvent | None:
+        """Handle or update a tool call, creating or updating a `ToolCallPart` or `ToolCallPartDelta`.
+
+        Managed items remain as `ToolCallPartDelta`s until they have both a tool_name and arguments, at which
+        point they are upgraded to `ToolCallPart`s.
+
+        If `vendor_part_id` is None, updates the latest matching ToolCallPart (or ToolCallPartDelta)
+        if any. Otherwise, a new part (or delta) may be created.
+
+        Args:
+            vendor_part_id: The ID the vendor uses for this tool call.
+                If None, the latest matching tool call may be updated.
+            tool_name: The name of the tool. If None, the manager does not enforce
+                a name match when `vendor_part_id` is None.
+            args: Arguments for the tool call, either as a string or a dictionary of key-value pairs.
+            tool_call_id: An optional string representing an identifier for this tool call.
+
+        Returns:
+            - A `PartStartEvent` if a new (fully realized) ToolCallPart is created.
+            - A `PartDeltaEvent` if an existing part is updated.
+            - `None` if no new event is emitted (e.g., the part is still incomplete).
+
+        Raises:
+            UnexpectedModelBehavior: If attempting to apply a tool call delta to a part that is not
+                a ToolCallPart or ToolCallPartDelta.
+        """
+        existing_matching_part_and_index: tuple[ToolCallPartDelta | ToolCallPart, int] | None = None
+
+        if vendor_part_id is None:
+            # vendor_part_id is None, so check if the latest part is a matching tool call or delta to update
+            # When the vendor_part_id is None, if the tool_name is _not_ None, assume this should be a new part rather
+            # than a delta on an existing one. We can change this behavior in the future if necessary for some model.
+            if tool_name is None and self._parts:
+                part_index = len(self._parts) - 1
+                latest_part = self._parts[part_index]
+                if isinstance(latest_part, (ToolCallPart, ToolCallPartDelta)):
+                    existing_matching_part_and_index = latest_part, part_index
+        else:
+            # vendor_part_id is provided, so look up the corresponding part or delta
+            part_index = self._vendor_id_to_part_index.get(vendor_part_id)
+            if part_index is not None:
+                existing_part = self._parts[part_index]
+                if not isinstance(existing_part, (ToolCallPartDelta, ToolCallPart)):
+                    raise UnexpectedModelBehavior(f'Cannot apply a tool call delta to {existing_part=}')
+                existing_matching_part_and_index = existing_part, part_index
+
+        if existing_matching_part_and_index is None:
+            # No matching part/delta was found, so create a new ToolCallPartDelta (or ToolCallPart if fully formed)
+            delta = ToolCallPartDelta(tool_name_delta=tool_name, args_delta=args, tool_call_id=tool_call_id)
+            part = delta.as_part() or delta
+            if vendor_part_id is not None:
+                self._vendor_id_to_part_index[vendor_part_id] = len(self._parts)
+            new_part_index = len(self._parts)
+            self._parts.append(part)
+            # Only emit a PartStartEvent if we have enough information to produce a full ToolCallPart
+            if isinstance(part, ToolCallPart):
+                return PartStartEvent(index=new_part_index, part=part)
+        else:
+            # Update the existing part or delta with the new information
+            existing_part, part_index = existing_matching_part_and_index
+            delta = ToolCallPartDelta(tool_name_delta=tool_name, args_delta=args, tool_call_id=tool_call_id)
+            updated_part = delta.apply(existing_part)
+            self._parts[part_index] = updated_part
+            if isinstance(updated_part, ToolCallPart):
+                if isinstance(existing_part, ToolCallPartDelta):
+                    # We just upgraded a delta to a full part, so emit a PartStartEvent
+                    return PartStartEvent(index=part_index, part=updated_part)
+                else:
+                    # We updated an existing part, so emit a PartDeltaEvent
+                    return PartDeltaEvent(index=part_index, delta=delta)
+
+    def handle_tool_call_part(
+        self,
+        *,
+        vendor_part_id: Hashable | None,
+        tool_name: str,
+        args: str | dict[str, Any],
+        tool_call_id: str | None = None,
+    ) -> ModelResponseStreamEvent:
+        """Immediately create or fully-overwrite a ToolCallPart with the given information.
+
+        This does not apply a delta; it directly sets the tool call part contents.
+
+        Args:
+            vendor_part_id: The vendor's ID for this tool call part. If not
+                None and an existing part is found, that part is overwritten.
+            tool_name: The name of the tool being invoked.
+            args: The arguments for the tool call, either as a string or a dictionary.
+            tool_call_id: An optional string identifier for this tool call.
+
+        Returns:
+            ModelResponseStreamEvent: A `PartStartEvent` indicating that a new tool call part
+            has been added to the manager, or replaced an existing part.
+        """
+        new_part = ToolCallPart.from_raw_args(tool_name=tool_name, args=args, tool_call_id=tool_call_id)
+        if vendor_part_id is None:
+            # vendor_part_id is None, so we unconditionally append a new ToolCallPart to the end of the list
+            new_part_index = len(self._parts)
+            self._parts.append(new_part)
+        else:
+            # vendor_part_id is provided, so find and overwrite or create a new ToolCallPart.
+            maybe_part_index = self._vendor_id_to_part_index.get(vendor_part_id)
+            if maybe_part_index is not None:
+                new_part_index = maybe_part_index
+                self._parts[new_part_index] = new_part
+            else:
+                new_part_index = len(self._parts)
+                self._parts.append(new_part)
+            self._vendor_id_to_part_index[vendor_part_id] = new_part_index
+        return PartStartEvent(index=new_part_index, part=new_part)

pydantic_ai/_pydantic.py

@@ -20,7 +20,7 @@ from ._griffe import doc_descriptions
 from ._utils import check_object_json_schema, is_model_like
 
 if TYPE_CHECKING:
-    from .tools import ObjectJsonSchema
+    from .tools import DocstringFormat, ObjectJsonSchema
 
 
 __all__ = ('function_schema',)
@@ -38,12 +38,19 @@ class FunctionSchema(TypedDict):
     var_positional_field: str | None
 
 
-def function_schema(function: Callable[..., Any], takes_ctx: bool) -> FunctionSchema:  # noqa: C901
+def function_schema(  # noqa: C901
+    function: Callable[..., Any],
+    takes_ctx: bool,
+    docstring_format: DocstringFormat,
+    require_parameter_descriptions: bool,
+) -> FunctionSchema:
     """Build a Pydantic validator and JSON schema from a tool function.
 
     Args:
         function: The function to build a validator and JSON schema for.
         takes_ctx: Whether the function takes a `RunContext` first argument.
+        docstring_format: The docstring format to use.
+        require_parameter_descriptions: Whether to require descriptions for all tool function parameters.
 
     Returns:
         A `FunctionSchema` instance.
@@ -62,7 +69,13 @@ def function_schema(function: Callable[..., Any], takes_ctx: bool) -> FunctionSc
     var_positional_field: str | None = None
     errors: list[str] = []
     decorators = _decorators.DecoratorInfos()
-    description, field_descriptions = doc_descriptions(function, sig)
+
+    description, field_descriptions = doc_descriptions(function, sig, docstring_format=docstring_format)
+
+    if require_parameter_descriptions:
+        if len(field_descriptions) != len(sig.parameters):
+            missing_params = set(sig.parameters) - set(field_descriptions)
+            errors.append(f'Missing parameter descriptions for {", ".join(missing_params)}')
 
     for index, (name, p) in enumerate(sig.parameters.items()):
         if p.annotation is sig.empty:

pydantic_ai/_system_prompt.py

@@ -12,6 +12,7 @@ from .tools import AgentDeps, RunContext, SystemPromptFunc
 @dataclass
 class SystemPromptRunner(Generic[AgentDeps]):
     function: SystemPromptFunc[AgentDeps]
+    dynamic: bool = False
     _takes_ctx: bool = field(init=False)
     _is_async: bool = field(init=False)
 

pydantic_ai/_utils.py

@@ -15,7 +15,7 @@ from pydantic.json_schema import JsonSchemaValue
 from typing_extensions import ParamSpec, TypeAlias, TypeGuard, is_typeddict
 
 if TYPE_CHECKING:
-    from .messages import RetryPromptPart, ToolCallPart, ToolReturnPart
+    from . import messages as _messages
     from .tools import ObjectJsonSchema
 
 _P = ParamSpec('_P')
@@ -136,7 +136,7 @@ class Either(Generic[Left, Right]):
 
 @asynccontextmanager
 async def group_by_temporal(
-    aiter: AsyncIterator[T], soft_max_interval: float | None
+    aiterable: AsyncIterable[T], soft_max_interval: float | None
 ) -> AsyncIterator[AsyncIterable[list[T]]]:
     """Group items from an async iterable into lists based on time interval between them.
 
@@ -154,18 +154,18 @@ async def group_by_temporal(
     ```
 
     Args:
-        aiter: The async iterable to group.
+        aiterable: The async iterable to group.
         soft_max_interval: Maximum interval over which to group items, this should avoid a trickle of items causing
             a group to never be yielded. It's a soft max in the sense that once we're over this time, we yield items
             as soon as `aiter.__anext__()` returns. If `None`, no grouping/debouncing is performed
 
     Returns:
-        A context manager usable as an iterator async iterable of lists of items from the input async iterable.
+        A context manager usable as an async iterable of lists of items produced by the input async iterable.
     """
     if soft_max_interval is None:
 
         async def async_iter_groups_noop() -> AsyncIterator[list[T]]:
-            async for item in aiter:
+            async for item in aiterable:
                 yield [item]
 
         yield async_iter_groups_noop()
@@ -181,6 +181,7 @@ async def group_by_temporal(
         buffer: list[T] = []
         group_start_time = time.monotonic()
 
+        aiterator = aiterable.__aiter__()
         while True:
             if group_start_time is None:
                 # group hasn't started, we just wait for the maximum interval
@@ -193,7 +194,7 @@ async def group_by_temporal(
             if task is None:
                 # aiter.__anext__() returns an Awaitable[T], not a Coroutine which asyncio.create_task expects
                 # so far, this doesn't seem to be a problem
-                task = asyncio.create_task(aiter.__anext__())  # pyright: ignore[reportArgumentType]
+                task = asyncio.create_task(aiterator.__anext__())  # pyright: ignore[reportArgumentType]
 
             # we use asyncio.wait to avoid cancelling the coroutine if it's not done
             done, _ = await asyncio.wait((task,), timeout=wait_time)
@@ -232,16 +233,6 @@ async def group_by_temporal(
                 await task
 
 
-def add_optional(a: str | None, b: str | None) -> str | None:
-    """Add two optional strings."""
-    if a is None:
-        return b
-    elif b is None:
-        return a
-    else:
-        return a + b
-
-
 def sync_anext(iterator: Iterator[T]) -> T:
     """Get the next item from a sync iterator, raising `StopAsyncIteration` if it's exhausted.
 
@@ -257,7 +248,79 @@ def now_utc() -> datetime:
     return datetime.now(tz=timezone.utc)
 
 
-def guard_tool_call_id(t: ToolCallPart | ToolReturnPart | RetryPromptPart, model_source: str) -> str:
+def guard_tool_call_id(
+    t: _messages.ToolCallPart | _messages.ToolReturnPart | _messages.RetryPromptPart, model_source: str
+) -> str:
     """Type guard that checks a `tool_call_id` is not None both for static typing and runtime."""
     assert t.tool_call_id is not None, f'{model_source} requires `tool_call_id` to be set: {t}'
     return t.tool_call_id
+
+
+class PeekableAsyncStream(Generic[T]):
+    """Wraps an async iterable of type T and allows peeking at the *next* item without consuming it.
+
+    We only buffer one item at a time (the next item). Once that item is yielded, it is discarded.
+    This is a single-pass stream.
+    """
+
+    def __init__(self, source: AsyncIterable[T]):
+        self._source = source
+        self._source_iter: AsyncIterator[T] | None = None
+        self._buffer: T | Unset = UNSET
+        self._exhausted = False
+
+    async def peek(self) -> T | Unset:
+        """Returns the next item that would be yielded without consuming it.
+
+        Returns None if the stream is exhausted.
+        """
+        if self._exhausted:
+            return UNSET
+
+        # If we already have a buffered item, just return it.
+        if not isinstance(self._buffer, Unset):
+            return self._buffer
+
+        # Otherwise, we need to fetch the next item from the underlying iterator.
+        if self._source_iter is None:
+            self._source_iter = self._source.__aiter__()
+
+        try:
+            self._buffer = await self._source_iter.__anext__()
+        except StopAsyncIteration:
+            self._exhausted = True
+            return UNSET
+
+        return self._buffer
+
+    async def is_exhausted(self) -> bool:
+        """Returns True if the stream is exhausted, False otherwise."""
+        return isinstance(await self.peek(), Unset)
+
+    def __aiter__(self) -> AsyncIterator[T]:
+        # For a single-pass iteration, we can return self as the iterator.
+        return self
+
+    async def __anext__(self) -> T:
+        """Yields the buffered item if present, otherwise fetches the next item from the underlying source.
+
+        Raises StopAsyncIteration if the stream is exhausted.
+        """
+        if self._exhausted:
+            raise StopAsyncIteration
+
+        # If we have a buffered item, yield it.
+        if not isinstance(self._buffer, Unset):
+            item = self._buffer
+            self._buffer = UNSET
+            return item
+
+        # Otherwise, fetch the next item from the source.
+        if self._source_iter is None:
+            self._source_iter = self._source.__aiter__()
+
+        try:
+            return await self._source_iter.__anext__()
+        except StopAsyncIteration:
+            self._exhausted = True
+            raise