pydantic-ai-slim 1.9.0__py3-none-any.whl → 1.12.0__py3-none-any.whl

pydantic_ai/_agent_graph.py

@@ -374,9 +374,10 @@ async def _prepare_request_parameters(
 ) -> models.ModelRequestParameters:
     """Build tools and create an agent model."""
     output_schema = ctx.deps.output_schema
-    output_object = None
-    if isinstance(output_schema, _output.NativeOutputSchema):
-        output_object = output_schema.object_def
+
+    prompted_output_template = (
+        output_schema.template if isinstance(output_schema, _output.PromptedOutputSchema) else None
+    )
 
     function_tools: list[ToolDefinition] = []
     output_tools: list[ToolDefinition] = []
@@ -391,7 +392,8 @@ async def _prepare_request_parameters(
         builtin_tools=ctx.deps.builtin_tools,
         output_mode=output_schema.mode,
         output_tools=output_tools,
-        output_object=output_object,
+        output_object=output_schema.object_def,
+        prompted_output_template=prompted_output_template,
         allow_text_output=output_schema.allows_text,
         allow_image_output=output_schema.allows_image,
     )
@@ -489,7 +491,6 @@ class ModelRequestNode(AgentNode[DepsT, NodeRunEndT]):
         message_history = _clean_message_history(message_history)
 
         model_request_parameters = await _prepare_request_parameters(ctx)
-        model_request_parameters = ctx.deps.model.customize_request_parameters(model_request_parameters)
 
         model_settings = ctx.deps.model_settings
         usage = ctx.state.usage
@@ -570,7 +571,7 @@ class CallToolsNode(AgentNode[DepsT, NodeRunEndT]):
                     # we got an empty response.
                     # this sometimes happens with anthropic (and perhaps other models)
                     # when the model has already returned text along side tool calls
-                    if text_processor := output_schema.text_processor:
+                    if text_processor := output_schema.text_processor:  # pragma: no branch
                         # in this scenario, if text responses are allowed, we return text from the most recent model
                         # response, if any
                         for message in reversed(ctx.state.message_history):
@@ -584,8 +585,12 @@ class CallToolsNode(AgentNode[DepsT, NodeRunEndT]):
                                         # not part of the final result output, so we reset the accumulated text
                                         text = ''  # pragma: no cover
                                 if text:
-                                    self._next_node = await self._handle_text_response(ctx, text, text_processor)
-                                    return
+                                    try:
+                                        self._next_node = await self._handle_text_response(ctx, text, text_processor)
+                                        return
+                                    except ToolRetryError:
+                                        # If the text from the preview response was invalid, ignore it.
+                                        pass
 
                     # Go back to the model request node with an empty request, which means we'll essentially
                     # resubmit the most recent request that resulted in an empty response,
@@ -622,11 +627,11 @@ class CallToolsNode(AgentNode[DepsT, NodeRunEndT]):
                     else:
                         assert_never(part)
 
-                # At the moment, we prioritize at least executing tool calls if they are present.
-                # In the future, we'd consider making this configurable at the agent or run level.
-                # This accounts for cases like anthropic returns that might contain a text response
-                # and a tool call response, where the text response just indicates the tool call will happen.
                 try:
+                    # At the moment, we prioritize at least executing tool calls if they are present.
+                    # In the future, we'd consider making this configurable at the agent or run level.
+                    # This accounts for cases like anthropic returns that might contain a text response
+                    # and a tool call response, where the text response just indicates the tool call will happen.
                     alternatives: list[str] = []
                     if tool_calls:
                         async for event in self._handle_tool_calls(ctx, tool_calls):
@@ -770,7 +775,6 @@ def build_run_context(ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT
         if ctx.deps.instrumentation_settings
         else DEFAULT_INSTRUMENTATION_VERSION,
         run_step=ctx.state.run_step,
-        tool_call_approved=ctx.state.run_step == 0,
     )
 
 
@@ -1034,7 +1038,7 @@ async def _call_tool(
         elif isinstance(tool_call_result, ToolApproved):
             if tool_call_result.override_args is not None:
                 tool_call = dataclasses.replace(tool_call, args=tool_call_result.override_args)
-            tool_result = await tool_manager.handle_call(tool_call)
+            tool_result = await tool_manager.handle_call(tool_call, approved=True)
         elif isinstance(tool_call_result, ToolDenied):
             return _messages.ToolReturnPart(
                 tool_name=tool_call.tool_name,

pydantic_ai/_output.py

@@ -10,7 +10,7 @@ from typing import TYPE_CHECKING, Any, Generic, Literal, cast, overload
 
 from pydantic import Json, TypeAdapter, ValidationError
 from pydantic_core import SchemaValidator, to_json
-from typing_extensions import Self, TypedDict, TypeVar, assert_never
+from typing_extensions import Self, TypedDict, TypeVar
 
 from pydantic_ai._instrumentation import InstrumentationNames
 
@@ -26,7 +26,6 @@ from .output import (
     OutputSpec,
     OutputTypeOrFunction,
     PromptedOutput,
-    StructuredOutputMode,
     TextOutput,
     TextOutputFunc,
     ToolOutput,
@@ -36,7 +35,7 @@ from .tools import GenerateToolJsonSchema, ObjectJsonSchema, ToolDefinition
 from .toolsets.abstract import AbstractToolset, ToolsetTool
 
 if TYPE_CHECKING:
-    from .profiles import ModelProfile
+    pass
 
 T = TypeVar('T')
 """An invariant TypeVar."""
@@ -212,59 +211,30 @@ class OutputValidator(Generic[AgentDepsT, OutputDataT_inv]):
 
 
 @dataclass(kw_only=True)
-class BaseOutputSchema(ABC, Generic[OutputDataT]):
+class OutputSchema(ABC, Generic[OutputDataT]):
     text_processor: BaseOutputProcessor[OutputDataT] | None = None
     toolset: OutputToolset[Any] | None = None
+    object_def: OutputObjectDefinition | None = None
     allows_deferred_tools: bool = False
     allows_image: bool = False
 
-    @abstractmethod
-    def with_default_mode(self, mode: StructuredOutputMode) -> OutputSchema[OutputDataT]:
+    @property
+    def mode(self) -> OutputMode:
         raise NotImplementedError()
 
     @property
     def allows_text(self) -> bool:
         return self.text_processor is not None
 
-
-@dataclass(init=False)
-class OutputSchema(BaseOutputSchema[OutputDataT], ABC):
-    """Model the final output from an agent run."""
-
-    @classmethod
-    @overload
-    def build(
-        cls,
-        output_spec: OutputSpec[OutputDataT],
-        *,
-        default_mode: StructuredOutputMode,
-        name: str | None = None,
-        description: str | None = None,
-        strict: bool | None = None,
-    ) -> OutputSchema[OutputDataT]: ...
-
-    @classmethod
-    @overload
-    def build(
-        cls,
-        output_spec: OutputSpec[OutputDataT],
-        *,
-        default_mode: None = None,
-        name: str | None = None,
-        description: str | None = None,
-        strict: bool | None = None,
-    ) -> BaseOutputSchema[OutputDataT]: ...
-
     @classmethod
     def build(  # noqa: C901
         cls,
         output_spec: OutputSpec[OutputDataT],
         *,
-        default_mode: StructuredOutputMode | None = None,
         name: str | None = None,
         description: str | None = None,
         strict: bool | None = None,
-    ) -> BaseOutputSchema[OutputDataT]:
+    ) -> OutputSchema[OutputDataT]:
         """Build an OutputSchema dataclass from an output type."""
         outputs = _flatten_output_spec(output_spec)
 
@@ -382,15 +352,12 @@ class OutputSchema(BaseOutputSchema[OutputDataT], ABC):
             )
 
         if len(other_outputs) > 0:
-            schema = OutputSchemaWithoutMode(
+            return AutoOutputSchema(
                 processor=cls._build_processor(other_outputs, name=name, description=description, strict=strict),
                 toolset=toolset,
                 allows_deferred_tools=allows_deferred_tools,
                 allows_image=allows_image,
             )
-            if default_mode:
-                schema = schema.with_default_mode(default_mode)
-            return schema
 
         if allows_image:
             return ImageOutputSchema(allows_deferred_tools=allows_deferred_tools)
@@ -410,22 +377,9 @@ class OutputSchema(BaseOutputSchema[OutputDataT], ABC):
 
         return UnionOutputProcessor(outputs=outputs, strict=strict, name=name, description=description)
 
-    @property
-    @abstractmethod
-    def mode(self) -> OutputMode:
-        raise NotImplementedError()
-
-    def raise_if_unsupported(self, profile: ModelProfile) -> None:
-        """Raise an error if the mode is not supported by this model."""
-        if self.allows_image and not profile.supports_image_output:
-            raise UserError('Image output is not supported by this model.')
-
-    def with_default_mode(self, mode: StructuredOutputMode) -> OutputSchema[OutputDataT]:
-        return self
-
 
 @dataclass(init=False)
-class OutputSchemaWithoutMode(BaseOutputSchema[OutputDataT]):
+class AutoOutputSchema(OutputSchema[OutputDataT]):
     processor: BaseObjectOutputProcessor[OutputDataT]
 
     def __init__(
@@ -439,32 +393,17 @@ class OutputSchemaWithoutMode(BaseOutputSchema[OutputDataT]):
         # At that point we may not know yet what output mode we're going to use if no model was provided or it was deferred until agent.run time,
         # but we cover ourselves just in case we end up using the tool output mode.
         super().__init__(
-            allows_deferred_tools=allows_deferred_tools,
             toolset=toolset,
+            object_def=processor.object_def,
             text_processor=processor,
+            allows_deferred_tools=allows_deferred_tools,
             allows_image=allows_image,
         )
         self.processor = processor
 
-    def with_default_mode(self, mode: StructuredOutputMode) -> OutputSchema[OutputDataT]:
-        if mode == 'native':
-            return NativeOutputSchema(
-                processor=self.processor,
-                allows_deferred_tools=self.allows_deferred_tools,
-                allows_image=self.allows_image,
-            )
-        elif mode == 'prompted':
-            return PromptedOutputSchema(
-                processor=self.processor,
-                allows_deferred_tools=self.allows_deferred_tools,
-                allows_image=self.allows_image,
-            )
-        elif mode == 'tool':
-            return ToolOutputSchema(
-                toolset=self.toolset, allows_deferred_tools=self.allows_deferred_tools, allows_image=self.allows_image
-            )
-        else:
-            assert_never(mode)
+    @property
+    def mode(self) -> OutputMode:
+        return 'auto'
 
 
 @dataclass(init=False)
@@ -486,10 +425,6 @@ class TextOutputSchema(OutputSchema[OutputDataT]):
     def mode(self) -> OutputMode:
         return 'text'
 
-    def raise_if_unsupported(self, profile: ModelProfile) -> None:
-        """Raise an error if the mode is not supported by this model."""
-        super().raise_if_unsupported(profile)
-
 
 class ImageOutputSchema(OutputSchema[OutputDataT]):
     def __init__(self, *, allows_deferred_tools: bool):
@@ -499,11 +434,6 @@ class ImageOutputSchema(OutputSchema[OutputDataT]):
     def mode(self) -> OutputMode:
         return 'image'
 
-    def raise_if_unsupported(self, profile: ModelProfile) -> None:
-        """Raise an error if the mode is not supported by this model."""
-        # This already raises if image output is not supported by this model.
-        super().raise_if_unsupported(profile)
-
 
 @dataclass(init=False)
 class StructuredTextOutputSchema(OutputSchema[OutputDataT], ABC):
@@ -513,25 +443,19 @@ class StructuredTextOutputSchema(OutputSchema[OutputDataT], ABC):
         self, *, processor: BaseObjectOutputProcessor[OutputDataT], allows_deferred_tools: bool, allows_image: bool
     ):
         super().__init__(
-            text_processor=processor, allows_deferred_tools=allows_deferred_tools, allows_image=allows_image
+            text_processor=processor,
+            object_def=processor.object_def,
+            allows_deferred_tools=allows_deferred_tools,
+            allows_image=allows_image,
         )
         self.processor = processor
 
-    @property
-    def object_def(self) -> OutputObjectDefinition:
-        return self.processor.object_def
-
 
 class NativeOutputSchema(StructuredTextOutputSchema[OutputDataT]):
     @property
     def mode(self) -> OutputMode:
         return 'native'
 
-    def raise_if_unsupported(self, profile: ModelProfile) -> None:
-        """Raise an error if the mode is not supported by this model."""
-        if not profile.supports_json_schema_output:
-            raise UserError('Native structured output is not supported by this model.')
-
 
 @dataclass(init=False)
 class PromptedOutputSchema(StructuredTextOutputSchema[OutputDataT]):
@@ -570,14 +494,11 @@ class PromptedOutputSchema(StructuredTextOutputSchema[OutputDataT]):
 
         return template.format(schema=json.dumps(schema))
 
-    def raise_if_unsupported(self, profile: ModelProfile) -> None:
-        """Raise an error if the mode is not supported by this model."""
-        super().raise_if_unsupported(profile)
-
-    def instructions(self, default_template: str) -> str:
+    def instructions(self, default_template: str) -> str:  # pragma: no cover
         """Get instructions to tell model to output JSON matching the schema."""
         template = self.template or default_template
         object_def = self.object_def
+        assert object_def is not None
         return self.build_instructions(template, object_def)
 
 
@@ -602,12 +523,6 @@ class ToolOutputSchema(OutputSchema[OutputDataT]):
     def mode(self) -> OutputMode:
         return 'tool'
 
-    def raise_if_unsupported(self, profile: ModelProfile) -> None:
-        """Raise an error if the mode is not supported by this model."""
-        super().raise_if_unsupported(profile)
-        if not profile.supports_tools:
-            raise UserError('Tool output is not supported by this model.')
-
 
 class BaseOutputProcessor(ABC, Generic[OutputDataT]):
     @abstractmethod

pydantic_ai/_run_context.py

@@ -16,15 +16,19 @@ if TYPE_CHECKING:
     from .models import Model
     from .result import RunUsage
 
+# TODO (v2): Change the default for all typevars like this from `None` to `object`
 AgentDepsT = TypeVar('AgentDepsT', default=None, contravariant=True)
 """Type variable for agent dependencies."""
 
+RunContextAgentDepsT = TypeVar('RunContextAgentDepsT', default=None, covariant=True)
+"""Type variable for the agent dependencies in `RunContext`."""
+
 
 @dataclasses.dataclass(repr=False, kw_only=True)
-class RunContext(Generic[AgentDepsT]):
+class RunContext(Generic[RunContextAgentDepsT]):
     """Information about the current call."""
 
-    deps: AgentDepsT
+    deps: RunContextAgentDepsT
     """Dependencies for the agent."""
     model: Model
     """The model used in this run."""
@@ -54,6 +58,8 @@ class RunContext(Generic[AgentDepsT]):
     """The current step in the run."""
     tool_call_approved: bool = False
     """Whether a tool call that required approval has now been approved."""
+    partial_output: bool = False
+    """Whether the output passed to an output validator is partial."""
 
     @property
     def last_attempt(self) -> bool:

pydantic_ai/_tool_manager.py

@@ -93,6 +93,8 @@ class ToolManager(Generic[AgentDepsT]):
         call: ToolCallPart,
         allow_partial: bool = False,
         wrap_validation_errors: bool = True,
+        *,
+        approved: bool = False,
     ) -> Any:
         """Handle a tool call by validating the arguments, calling the tool, and handling retries.
 
@@ -100,30 +102,38 @@ class ToolManager(Generic[AgentDepsT]):
             call: The tool call part to handle.
             allow_partial: Whether to allow partial validation of the tool arguments.
             wrap_validation_errors: Whether to wrap validation errors in a retry prompt part.
-            usage_limits: Optional usage limits to check before executing tools.
+            approved: Whether the tool call has been approved.
         """
         if self.tools is None or self.ctx is None:
             raise ValueError('ToolManager has not been prepared for a run step yet')  # pragma: no cover
 
         if (tool := self.tools.get(call.tool_name)) and tool.tool_def.kind == 'output':
             # Output tool calls are not traced and not counted
-            return await self._call_tool(call, allow_partial, wrap_validation_errors)
+            return await self._call_tool(
+                call,
+                allow_partial=allow_partial,
+                wrap_validation_errors=wrap_validation_errors,
+                approved=approved,
+            )
         else:
             return await self._call_function_tool(
                 call,
-                allow_partial,
-                wrap_validation_errors,
-                self.ctx.tracer,
-                self.ctx.trace_include_content,
-                self.ctx.instrumentation_version,
-                self.ctx.usage,
+                allow_partial=allow_partial,
+                wrap_validation_errors=wrap_validation_errors,
+                approved=approved,
+                tracer=self.ctx.tracer,
+                include_content=self.ctx.trace_include_content,
+                instrumentation_version=self.ctx.instrumentation_version,
+                usage=self.ctx.usage,
             )
 
     async def _call_tool(
         self,
         call: ToolCallPart,
+        *,
         allow_partial: bool,
         wrap_validation_errors: bool,
+        approved: bool,
     ) -> Any:
         if self.tools is None or self.ctx is None:
             raise ValueError('ToolManager has not been prepared for a run step yet')  # pragma: no cover
@@ -138,8 +148,8 @@ class ToolManager(Generic[AgentDepsT]):
                     msg = 'No tools available.'
                 raise ModelRetry(f'Unknown tool name: {name!r}. {msg}')
 
-            if tool.tool_def.defer:
-                raise RuntimeError('Deferred tools cannot be called')
+            if tool.tool_def.kind == 'external':
+                raise RuntimeError('External tools cannot be called')
 
             ctx = replace(
                 self.ctx,
@@ -147,6 +157,8 @@ class ToolManager(Generic[AgentDepsT]):
                 tool_call_id=call.tool_call_id,
                 retry=self.ctx.retries.get(name, 0),
                 max_retries=tool.max_retries,
+                tool_call_approved=approved,
+                partial_output=allow_partial,
             )
 
             pyd_allow_partial = 'trailing-strings' if allow_partial else 'off'
@@ -193,8 +205,10 @@ class ToolManager(Generic[AgentDepsT]):
     async def _call_function_tool(
         self,
         call: ToolCallPart,
+        *,
         allow_partial: bool,
         wrap_validation_errors: bool,
+        approved: bool,
         tracer: Tracer,
         include_content: bool,
         instrumentation_version: int,
@@ -233,7 +247,12 @@ class ToolManager(Generic[AgentDepsT]):
             attributes=span_attributes,
         ) as span:
             try:
-                tool_result = await self._call_tool(call, allow_partial, wrap_validation_errors)
+                tool_result = await self._call_tool(
+                    call,
+                    allow_partial=allow_partial,
+                    wrap_validation_errors=wrap_validation_errors,
+                    approved=approved,
+                )
                 usage.tool_calls += 1
 
             except ToolRetryError as e:

pydantic_ai/_utils.py

@@ -234,6 +234,15 @@ def sync_anext(iterator: Iterator[T]) -> T:
         raise StopAsyncIteration() from e
 
 
+def sync_async_iterator(async_iter: AsyncIterator[T]) -> Iterator[T]:
+    loop = get_event_loop()
+    while True:
+        try:
+            yield loop.run_until_complete(anext(async_iter))
+        except StopAsyncIteration:
+            break
+
+
 def now_utc() -> datetime:
     return datetime.now(tz=timezone.utc)
 
@@ -489,3 +498,12 @@ def get_union_args(tp: Any) -> tuple[Any, ...]:
         return tuple(_unwrap_annotated(arg) for arg in get_args(tp))
     else:
         return ()
+
+
+def get_event_loop():
+    try:
+        event_loop = asyncio.get_event_loop()
+    except RuntimeError:  # pragma: lax no cover
+        event_loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(event_loop)
+    return event_loop