pydantic-ai-slim 0.2.11__tar.gz → 0.2.12__tar.gz

{pydantic_ai_slim-0.2.11 → pydantic_ai_slim-0.2.12}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydantic-ai-slim
-Version: 0.2.11
+Version: 0.2.12
 Summary: Agent Framework / shim to use Pydantic with LLMs, slim package
 Author-email: Samuel Colvin <samuel@pydantic.dev>, Marcelo Trylesinski <marcelotryle@gmail.com>, David Montague <david@pydantic.dev>, Alex Hall <alex@pydantic.dev>
 License-Expression: MIT
@@ -30,11 +30,11 @@ Requires-Dist: exceptiongroup; python_version < '3.11'
 Requires-Dist: griffe>=1.3.2
 Requires-Dist: httpx>=0.27
 Requires-Dist: opentelemetry-api>=1.28.0
-Requires-Dist: pydantic-graph==0.2.11
+Requires-Dist: pydantic-graph==0.2.12
 Requires-Dist: pydantic>=2.10
 Requires-Dist: typing-inspection>=0.4.0
 Provides-Extra: a2a
-Requires-Dist: fasta2a==0.2.11; extra == 'a2a'
+Requires-Dist: fasta2a==0.2.12; extra == 'a2a'
 Provides-Extra: anthropic
 Requires-Dist: anthropic>=0.52.0; extra == 'anthropic'
 Provides-Extra: bedrock
@@ -48,7 +48,7 @@ Requires-Dist: cohere>=5.13.11; (platform_system != 'Emscripten') and extra == '
 Provides-Extra: duckduckgo
 Requires-Dist: duckduckgo-search>=7.0.0; extra == 'duckduckgo'
 Provides-Extra: evals
-Requires-Dist: pydantic-evals==0.2.11; extra == 'evals'
+Requires-Dist: pydantic-evals==0.2.12; extra == 'evals'
 Provides-Extra: google
 Requires-Dist: google-genai>=1.15.0; extra == 'google'
 Provides-Extra: groq

{pydantic_ai_slim-0.2.11 → pydantic_ai_slim-0.2.12}/pydantic_ai/_agent_graph.py

@@ -24,7 +24,7 @@ from . import (
     result,
     usage as _usage,
 )
-from .result import OutputDataT, ToolOutput
+from .result import OutputDataT
 from .settings import ModelSettings, merge_model_settings
 from .tools import RunContext, Tool, ToolDefinition, ToolsPrepareFunc
 
@@ -64,12 +64,14 @@ class GraphAgentState:
     retries: int
     run_step: int
 
-    def increment_retries(self, max_result_retries: int) -> None:
+    def increment_retries(self, max_result_retries: int, error: Exception | None = None) -> None:
         self.retries += 1
         if self.retries > max_result_retries:
-            raise exceptions.UnexpectedModelBehavior(
-                f'Exceeded maximum retries ({max_result_retries}) for result validation'
-            )
+            message = f'Exceeded maximum retries ({max_result_retries}) for result validation'
+            if error:
+                raise exceptions.UnexpectedModelBehavior(message) from error
+            else:
+                raise exceptions.UnexpectedModelBehavior(message)
 
 
 @dataclasses.dataclass
@@ -264,7 +266,7 @@ async def _prepare_request_parameters(
     output_schema = ctx.deps.output_schema
     return models.ModelRequestParameters(
         function_tools=function_tool_defs,
-        allow_text_output=allow_text_output(output_schema),
+        allow_text_output=_output.allow_text_output(output_schema),
         output_tools=output_schema.tool_defs() if output_schema is not None else [],
     )
 
@@ -450,7 +452,7 @@ class CallToolsNode(AgentNode[DepsT, NodeRunEndT]):
                     # when the model has already returned text along side tool calls
                     # in this scenario, if text responses are allowed, we return text from the most recent model
                     # response, if any
-                    if allow_text_output(ctx.deps.output_schema):
+                    if _output.allow_text_output(ctx.deps.output_schema):
                         for message in reversed(ctx.state.message_history):
                             if isinstance(message, _messages.ModelResponse):
                                 last_texts = [p.content for p in message.parts if isinstance(p, _messages.TextPart)]
@@ -471,6 +473,7 @@ class CallToolsNode(AgentNode[DepsT, NodeRunEndT]):
         tool_calls: list[_messages.ToolCallPart],
     ) -> AsyncIterator[_messages.HandleResponseEvent]:
         output_schema = ctx.deps.output_schema
+        run_context = build_run_context(ctx)
 
         # first, look for the output tool call
         final_result: result.FinalResult[NodeRunEndT] | None = None
@@ -478,12 +481,12 @@ class CallToolsNode(AgentNode[DepsT, NodeRunEndT]):
         if output_schema is not None:
             for call, output_tool in output_schema.find_tool(tool_calls):
                 try:
-                    result_data = output_tool.validate(call)
+                    result_data = await output_tool.process(call, run_context)
                     result_data = await _validate_output(result_data, ctx, call)
                 except _output.ToolRetryError as e:
                     # TODO: Should only increment retry stuff once per node execution, not for each tool call
                     #   Also, should increment the tool-specific retry count rather than the run retry count
-                    ctx.state.increment_retries(ctx.deps.max_result_retries)
+                    ctx.state.increment_retries(ctx.deps.max_result_retries, e)
                     parts.append(e.tool_retry)
                 else:
                     final_result = result.FinalResult(result_data, call.tool_name, call.tool_call_id)
@@ -505,7 +508,6 @@ class CallToolsNode(AgentNode[DepsT, NodeRunEndT]):
         else:
             if tool_responses:
                 parts.extend(tool_responses)
-            run_context = build_run_context(ctx)
             instructions = await ctx.deps.get_instructions(run_context)
             self._next_node = ModelRequestNode[DepsT, NodeRunEndT](
                 _messages.ModelRequest(parts=parts, instructions=instructions)
@@ -533,27 +535,22 @@ class CallToolsNode(AgentNode[DepsT, NodeRunEndT]):
         output_schema = ctx.deps.output_schema
 
         text = '\n\n'.join(texts)
-        if allow_text_output(output_schema):
-            # The following cast is safe because we know `str` is an allowed result type
-            result_data_input = cast(NodeRunEndT, text)
-            try:
-                result_data = await _validate_output(result_data_input, ctx, None)
-            except _output.ToolRetryError as e:
-                ctx.state.increment_retries(ctx.deps.max_result_retries)
-                return ModelRequestNode[DepsT, NodeRunEndT](_messages.ModelRequest(parts=[e.tool_retry]))
+        try:
+            if _output.allow_text_output(output_schema):
+                # The following cast is safe because we know `str` is an allowed result type
+                result_data = cast(NodeRunEndT, text)
             else:
-                return self._handle_final_result(ctx, result.FinalResult(result_data, None, None), [])
-        else:
-            ctx.state.increment_retries(ctx.deps.max_result_retries)
-            return ModelRequestNode[DepsT, NodeRunEndT](
-                _messages.ModelRequest(
-                    parts=[
-                        _messages.RetryPromptPart(
-                            content='Plain text responses are not permitted, please include your response in a tool call',
-                        )
-                    ]
+                m = _messages.RetryPromptPart(
+                    content='Plain text responses are not permitted, please include your response in a tool call',
                 )
-            )
+                raise _output.ToolRetryError(m)
+
+            result_data = await _validate_output(result_data, ctx, None)
+        except _output.ToolRetryError as e:
+            ctx.state.increment_retries(ctx.deps.max_result_retries, e)
+            return ModelRequestNode[DepsT, NodeRunEndT](_messages.ModelRequest(parts=[e.tool_retry]))
+        else:
+            return self._handle_final_result(ctx, result.FinalResult(result_data, None, None), [])
 
 
 def build_run_context(ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, Any]]) -> RunContext[DepsT]:
@@ -795,11 +792,6 @@ async def _validate_output(
     return result_data
 
 
-def allow_text_output(output_schema: _output.OutputSchema[Any] | None) -> bool:
-    """Check if the result schema allows text results."""
-    return output_schema is None or output_schema.allow_text_output
-
-
 @dataclasses.dataclass
 class _RunMessages:
     messages: list[_messages.ModelMessage]
@@ -849,7 +841,9 @@ def get_captured_run_messages() -> _RunMessages:
 
 
 def build_agent_graph(
-    name: str | None, deps_type: type[DepsT], output_type: type[OutputT] | ToolOutput[OutputT]
+    name: str | None,
+    deps_type: type[DepsT],
+    output_type: _output.OutputType[OutputT],
 ) -> Graph[GraphAgentState, GraphAgentDeps[DepsT, result.FinalResult[OutputT]], result.FinalResult[OutputT]]:
     """Build the execution [Graph][pydantic_graph.Graph] for a given agent."""
     nodes = (

pydantic_ai_slim-0.2.11/pydantic_ai/_pydantic.py → pydantic_ai_slim-0.2.12/pydantic_ai/_function_schema.py

@@ -5,6 +5,9 @@ This module has to use numerous internal Pydantic APIs and is therefore brittle
 
 from __future__ import annotations as _annotations
 
+import inspect
+from collections.abc import Awaitable
+from dataclasses import dataclass
 from inspect import Parameter, signature
 from typing import TYPE_CHECKING, Any, Callable, cast
 
@@ -15,10 +18,12 @@ from pydantic.fields import FieldInfo
 from pydantic.json_schema import GenerateJsonSchema
 from pydantic.plugin._schema_validator import create_schema_validator
 from pydantic_core import SchemaValidator, core_schema
-from typing_extensions import TypedDict, get_origin
+from typing_extensions import get_origin
+
+from pydantic_ai.tools import RunContext
 
 from ._griffe import doc_descriptions
-from ._utils import check_object_json_schema, is_model_like
+from ._utils import check_object_json_schema, is_model_like, run_in_executor
 
 if TYPE_CHECKING:
     from .tools import DocstringFormat, ObjectJsonSchema
@@ -27,24 +32,53 @@ if TYPE_CHECKING:
 __all__ = ('function_schema',)
 
 
-class FunctionSchema(TypedDict):
+@dataclass
+class FunctionSchema:
     """Internal information about a function schema."""
 
+    function: Callable[..., Any]
     description: str
     validator: SchemaValidator
     json_schema: ObjectJsonSchema
     # if not None, the function takes a single by that name (besides potentially `info`)
+    takes_ctx: bool
+    is_async: bool
     single_arg_name: str | None
     positional_fields: list[str]
     var_positional_field: str | None
 
+    async def call(self, args_dict: dict[str, Any], ctx: RunContext[Any]) -> Any:
+        args, kwargs = self._call_args(args_dict, ctx)
+        if self.is_async:
+            function = cast(Callable[[Any], Awaitable[str]], self.function)
+            return await function(*args, **kwargs)
+        else:
+            function = cast(Callable[[Any], str], self.function)
+            return await run_in_executor(function, *args, **kwargs)
+
+    def _call_args(
+        self,
+        args_dict: dict[str, Any],
+        ctx: RunContext[Any],
+    ) -> tuple[list[Any], dict[str, Any]]:
+        if self.single_arg_name:
+            args_dict = {self.single_arg_name: args_dict}
+
+        args = [ctx] if self.takes_ctx else []
+        for positional_field in self.positional_fields:
+            args.append(args_dict.pop(positional_field))  # pragma: no cover
+        if self.var_positional_field:
+            args.extend(args_dict.pop(self.var_positional_field))
+
+        return args, args_dict
+
 
 def function_schema(  # noqa: C901
     function: Callable[..., Any],
-    takes_ctx: bool,
-    docstring_format: DocstringFormat,
-    require_parameter_descriptions: bool,
     schema_generator: type[GenerateJsonSchema],
+    takes_ctx: bool | None = None,
+    docstring_format: DocstringFormat = 'auto',
+    require_parameter_descriptions: bool = False,
 ) -> FunctionSchema:
     """Build a Pydantic validator and JSON schema from a tool function.
 
@@ -58,6 +92,9 @@ def function_schema(  # noqa: C901
     Returns:
         A `FunctionSchema` instance.
     """
+    if takes_ctx is None:
+        takes_ctx = _takes_ctx(function)
+
     config = ConfigDict(title=function.__name__, use_attribute_docstrings=True)
     config_wrapper = ConfigWrapper(config)
     gen_schema = _generate_schema.GenerateSchema(config_wrapper)
@@ -176,10 +213,13 @@ def function_schema(  # noqa: C901
         single_arg_name=single_arg_name,
         positional_fields=positional_fields,
         var_positional_field=var_positional_field,
+        takes_ctx=takes_ctx,
+        is_async=inspect.iscoroutinefunction(function),
+        function=function,
     )
 
 
-def takes_ctx(function: Callable[..., Any]) -> bool:
+def _takes_ctx(function: Callable[..., Any]) -> bool:
     """Check if a function takes a `RunContext` first argument.
 
     Args:
@@ -196,7 +236,7 @@ def takes_ctx(function: Callable[..., Any]) -> bool:
     else:
         type_hints = _typing_extra.get_function_type_hints(function)
         annotation = type_hints[first_param_name]
-        return annotation is not sig.empty and _is_call_ctx(annotation)
+        return True is not sig.empty and _is_call_ctx(annotation)
 
 
 def _build_schema(