pydantic-ai-slim 0.4.2__py3-none-any.whl → 0.4.4__py3-none-any.whl

pydantic_ai/_parts_manager.py

@@ -2,10 +2,10 @@
 
 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.
+and produces Pydantic AI-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.
+and are tightly coupled to the specific model being used, and the Pydantic AI Model subclass implementation.
 
 This `ModelResponsePartsManager` is used in each of the subclasses of `StreamedResponse` as a way to consolidate
 event-emitting logic.

pydantic_ai/_run_context.py

@@ -5,6 +5,7 @@ from collections.abc import Sequence
 from dataclasses import field
 from typing import TYPE_CHECKING, Generic
 
+from opentelemetry.trace import NoOpTracer, Tracer
 from typing_extensions import TypeVar
 
 from . import _utils, messages as _messages
@@ -27,10 +28,16 @@ class RunContext(Generic[AgentDepsT]):
     """The model used in this run."""
     usage: Usage
     """LLM usage associated with the run."""
-    prompt: str | Sequence[_messages.UserContent] | None
+    prompt: str | Sequence[_messages.UserContent] | None = None
     """The original user prompt passed to the run."""
     messages: list[_messages.ModelMessage] = field(default_factory=list)
     """Messages exchanged in the conversation so far."""
+    tracer: Tracer = field(default_factory=NoOpTracer)
+    """The tracer to use for tracing the run."""
+    trace_include_content: bool = False
+    """Whether to include the content of the messages in the trace."""
+    retries: dict[str, int] = field(default_factory=dict)
+    """Number of retries for each tool so far."""
     tool_call_id: str | None = None
     """The ID of the tool call."""
     tool_name: str | None = None
@@ -40,17 +47,4 @@ class RunContext(Generic[AgentDepsT]):
     run_step: int = 0
     """The current step in the run."""
 
-    def replace_with(
-        self,
-        retry: int | None = None,
-        tool_name: str | None | _utils.Unset = _utils.UNSET,
-    ) -> RunContext[AgentDepsT]:
-        # Create a new `RunContext` a new `retry` value and `tool_name`.
-        kwargs = {}
-        if retry is not None:
-            kwargs['retry'] = retry
-        if tool_name is not _utils.UNSET:  # pragma: no branch
-            kwargs['tool_name'] = tool_name
-        return dataclasses.replace(self, **kwargs)
-
     __repr__ = _utils.dataclasses_no_defaults_repr

pydantic_ai/_tool_manager.py

@@ -0,0 +1,190 @@
+from __future__ import annotations
+
+import json
+from collections.abc import Iterable
+from dataclasses import dataclass, replace
+from typing import Any, Generic
+
+from pydantic import ValidationError
+from typing_extensions import assert_never
+
+from pydantic_ai.output import DeferredToolCalls
+
+from . import messages as _messages
+from ._run_context import AgentDepsT, RunContext
+from .exceptions import ModelRetry, ToolRetryError, UnexpectedModelBehavior
+from .messages import ToolCallPart
+from .tools import ToolDefinition
+from .toolsets.abstract import AbstractToolset, ToolsetTool
+
+
+@dataclass
+class ToolManager(Generic[AgentDepsT]):
+    """Manages tools for an agent run step. It caches the agent run's toolset's tool definitions and handles calling tools and retries."""
+
+    ctx: RunContext[AgentDepsT]
+    """The agent run context for a specific run step."""
+    toolset: AbstractToolset[AgentDepsT]
+    """The toolset that provides the tools for this run step."""
+    tools: dict[str, ToolsetTool[AgentDepsT]]
+    """The cached tools for this run step."""
+
+    @classmethod
+    async def build(cls, toolset: AbstractToolset[AgentDepsT], ctx: RunContext[AgentDepsT]) -> ToolManager[AgentDepsT]:
+        """Build a new tool manager for a specific run step."""
+        return cls(
+            ctx=ctx,
+            toolset=toolset,
+            tools=await toolset.get_tools(ctx),
+        )
+
+    async def for_run_step(self, ctx: RunContext[AgentDepsT]) -> ToolManager[AgentDepsT]:
+        """Build a new tool manager for the next run step, carrying over the retries from the current run step."""
+        return await self.__class__.build(self.toolset, replace(ctx, retries=self.ctx.retries))
+
+    @property
+    def tool_defs(self) -> list[ToolDefinition]:
+        """The tool definitions for the tools in this tool manager."""
+        return [tool.tool_def for tool in self.tools.values()]
+
+    def get_tool_def(self, name: str) -> ToolDefinition | None:
+        """Get the tool definition for a given tool name, or `None` if the tool is unknown."""
+        try:
+            return self.tools[name].tool_def
+        except KeyError:
+            return None
+
+    async def handle_call(self, call: ToolCallPart, allow_partial: bool = False) -> Any:
+        """Handle a tool call by validating the arguments, calling the tool, and handling retries.
+
+        Args:
+            call: The tool call part to handle.
+            allow_partial: Whether to allow partial validation of the tool arguments.
+        """
+        if (tool := self.tools.get(call.tool_name)) and tool.tool_def.kind == 'output':
+            # Output tool calls are not traced
+            return await self._call_tool(call, allow_partial)
+        else:
+            return await self._call_tool_traced(call, allow_partial)
+
+    async def _call_tool(self, call: ToolCallPart, allow_partial: bool = False) -> Any:
+        name = call.tool_name
+        tool = self.tools.get(name)
+        try:
+            if tool is None:
+                if self.tools:
+                    msg = f'Available tools: {", ".join(f"{name!r}" for name in self.tools.keys())}'
+                else:
+                    msg = 'No tools available.'
+                raise ModelRetry(f'Unknown tool name: {name!r}. {msg}')
+
+            ctx = replace(
+                self.ctx,
+                tool_name=name,
+                tool_call_id=call.tool_call_id,
+                retry=self.ctx.retries.get(name, 0),
+            )
+
+            pyd_allow_partial = 'trailing-strings' if allow_partial else 'off'
+            validator = tool.args_validator
+            if isinstance(call.args, str):
+                args_dict = validator.validate_json(call.args or '{}', allow_partial=pyd_allow_partial)
+            else:
+                args_dict = validator.validate_python(call.args or {}, allow_partial=pyd_allow_partial)
+
+            output = await self.toolset.call_tool(name, args_dict, ctx, tool)
+        except (ValidationError, ModelRetry) as e:
+            max_retries = tool.max_retries if tool is not None else 1
+            current_retry = self.ctx.retries.get(name, 0)
+
+            if current_retry == max_retries:
+                raise UnexpectedModelBehavior(f'Tool {name!r} exceeded max retries count of {max_retries}') from e
+            else:
+                if isinstance(e, ValidationError):
+                    m = _messages.RetryPromptPart(
+                        tool_name=name,
+                        content=e.errors(include_url=False, include_context=False),
+                        tool_call_id=call.tool_call_id,
+                    )
+                    e = ToolRetryError(m)
+                elif isinstance(e, ModelRetry):
+                    m = _messages.RetryPromptPart(
+                        tool_name=name,
+                        content=e.message,
+                        tool_call_id=call.tool_call_id,
+                    )
+                    e = ToolRetryError(m)
+                else:
+                    assert_never(e)
+
+                self.ctx.retries[name] = current_retry + 1
+                raise e
+        else:
+            self.ctx.retries.pop(name, None)
+            return output
+
+    async def _call_tool_traced(self, call: ToolCallPart, allow_partial: bool = False) -> Any:
+        """See <https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-spans/#execute-tool-span>."""
+        span_attributes = {
+            'gen_ai.tool.name': call.tool_name,
+            # NOTE: this means `gen_ai.tool.call.id` will be included even if it was generated by pydantic-ai
+            'gen_ai.tool.call.id': call.tool_call_id,
+            **({'tool_arguments': call.args_as_json_str()} if self.ctx.trace_include_content else {}),
+            'logfire.msg': f'running tool: {call.tool_name}',
+            # add the JSON schema so these attributes are formatted nicely in Logfire
+            'logfire.json_schema': json.dumps(
+                {
+                    'type': 'object',
+                    'properties': {
+                        **(
+                            {
+                                'tool_arguments': {'type': 'object'},
+                                'tool_response': {'type': 'object'},
+                            }
+                            if self.ctx.trace_include_content
+                            else {}
+                        ),
+                        'gen_ai.tool.name': {},
+                        'gen_ai.tool.call.id': {},
+                    },
+                }
+            ),
+        }
+        with self.ctx.tracer.start_as_current_span('running tool', attributes=span_attributes) as span:
+            try:
+                tool_result = await self._call_tool(call, allow_partial)
+            except ToolRetryError as e:
+                part = e.tool_retry
+                if self.ctx.trace_include_content and span.is_recording():
+                    span.set_attribute('tool_response', part.model_response())
+                raise e
+
+            if self.ctx.trace_include_content and span.is_recording():
+                span.set_attribute(
+                    'tool_response',
+                    tool_result
+                    if isinstance(tool_result, str)
+                    else _messages.tool_return_ta.dump_json(tool_result).decode(),
+                )
+
+        return tool_result
+
+    def get_deferred_tool_calls(self, parts: Iterable[_messages.ModelResponsePart]) -> DeferredToolCalls | None:
+        """Get the deferred tool calls from the model response parts."""
+        deferred_calls_and_defs = [
+            (part, tool_def)
+            for part in parts
+            if isinstance(part, _messages.ToolCallPart)
+            and (tool_def := self.get_tool_def(part.tool_name))
+            and tool_def.kind == 'deferred'
+        ]
+        if not deferred_calls_and_defs:
+            return None
+
+        deferred_calls: list[_messages.ToolCallPart] = []
+        deferred_tool_defs: dict[str, ToolDefinition] = {}
+        for part, tool_def in deferred_calls_and_defs:
+            deferred_calls.append(part)
+            deferred_tool_defs[part.tool_name] = tool_def
+
+        return DeferredToolCalls(deferred_calls, deferred_tool_defs)

pydantic_ai/_utils.py

@@ -4,8 +4,10 @@ import asyncio
 import functools
 import inspect
 import re
+import sys
 import time
 import uuid
+import warnings
 from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Iterator
 from contextlib import asynccontextmanager, suppress
 from dataclasses import dataclass, fields, is_dataclass
@@ -29,7 +31,7 @@ from typing_extensions import (
 from typing_inspection import typing_objects
 from typing_inspection.introspection import is_union_origin
 
-from pydantic_graph._utils import AbstractSpan
+from pydantic_graph._utils import AbstractSpan, get_event_loop
 
 from . import exceptions
 
@@ -461,3 +463,18 @@ def get_union_args(tp: Any) -> tuple[Any, ...]:
         return get_args(tp)
     else:
         return ()
+
+
+# The `asyncio.Lock` `loop` argument was deprecated in 3.8 and removed in 3.10,
+# but 3.9 still needs it to have the intended behavior.
+
+if sys.version_info < (3, 10):
+
+    def get_async_lock() -> asyncio.Lock:  # pragma: lax no cover
+        with warnings.catch_warnings():
+            warnings.simplefilter('ignore', DeprecationWarning)
+            return asyncio.Lock(loop=get_event_loop())
+else:
+
+    def get_async_lock() -> asyncio.Lock:  # pragma: lax no cover
+        return asyncio.Lock()