pydantic-ai-slim 0.8.0__py3-none-any.whl → 1.0.0__py3-none-any.whl

pydantic_ai/ag_ui.py

@@ -8,12 +8,11 @@ from __future__ import annotations
 
 import json
 import uuid
-from collections.abc import AsyncIterator, Iterable, Mapping, Sequence
+from collections.abc import AsyncIterator, Callable, Iterable, Mapping, Sequence
 from dataclasses import Field, dataclass, replace
 from http import HTTPStatus
 from typing import (
     Any,
-    Callable,
     ClassVar,
     Final,
     Generic,
@@ -46,11 +45,11 @@ from .messages import (
     UserPromptPart,
 )
 from .models import KnownModelName, Model
-from .output import DeferredToolCalls, OutputDataT, OutputSpec
+from .output import OutputDataT, OutputSpec
 from .settings import ModelSettings
-from .tools import AgentDepsT, ToolDefinition
+from .tools import AgentDepsT, DeferredToolRequests, ToolDefinition
 from .toolsets import AbstractToolset
-from .toolsets.deferred import DeferredToolset
+from .toolsets.external import ExternalToolset
 from .usage import RunUsage, UsageLimits
 
 try:
@@ -69,6 +68,9 @@ try:
         TextMessageContentEvent,
         TextMessageEndEvent,
         TextMessageStartEvent,
+        # TODO: Enable once https://github.com/ag-ui-protocol/ag-ui/issues/289 is resolved.
+        # ThinkingEndEvent,
+        # ThinkingStartEvent,
         ThinkingTextMessageContentEvent,
         ThinkingTextMessageEndEvent,
         ThinkingTextMessageStartEvent,
@@ -343,7 +345,7 @@ async def run_ag_ui(
 
         async with agent.iter(
             user_prompt=None,
-            output_type=[output_type or agent.output_type, DeferredToolCalls],
+            output_type=[output_type or agent.output_type, DeferredToolRequests],
             message_history=messages,
             model=model,
             deps=deps,
@@ -393,6 +395,12 @@ async def _agent_stream(run: AgentRun[AgentDepsT, Any]) -> AsyncIterator[BaseEve
                 if stream_ctx.part_end:  # pragma: no branch
                     yield stream_ctx.part_end
                     stream_ctx.part_end = None
+                if stream_ctx.thinking:
+                    # TODO: Enable once https://github.com/ag-ui-protocol/ag-ui/issues/289 is resolved.
+                    # yield ThinkingEndEvent(
+                    #     type=EventType.THINKING_END,
+                    # )
+                    stream_ctx.thinking = False
         elif isinstance(node, CallToolsNode):
             async with node.stream(run.ctx) as handle_stream:
                 async for event in handle_stream:
@@ -401,7 +409,7 @@ async def _agent_stream(run: AgentRun[AgentDepsT, Any]) -> AsyncIterator[BaseEve
                             yield msg
 
 
-async def _handle_model_request_event(
+async def _handle_model_request_event(  # noqa: C901
     stream_ctx: _RequestStreamContext,
     agent_event: ModelResponseStreamEvent,
 ) -> AsyncIterator[BaseEvent]:
@@ -421,56 +429,70 @@ async def _handle_model_request_event(
             stream_ctx.part_end = None
 
         part = agent_event.part
-        if isinstance(part, TextPart):
-            message_id = stream_ctx.new_message_id()
-            yield TextMessageStartEvent(
-                message_id=message_id,
-            )
-            if part.content:  # pragma: no branch
-                yield TextMessageContentEvent(
-                    message_id=message_id,
+        if isinstance(part, ThinkingPart):  # pragma: no branch
+            if not stream_ctx.thinking:
+                # TODO: Enable once https://github.com/ag-ui-protocol/ag-ui/issues/289 is resolved.
+                # yield ThinkingStartEvent(
+                #     type=EventType.THINKING_START,
+                # )
+                stream_ctx.thinking = True
+
+            if part.content:
+                yield ThinkingTextMessageStartEvent(
+                    type=EventType.THINKING_TEXT_MESSAGE_START,
+                )
+                yield ThinkingTextMessageContentEvent(
+                    type=EventType.THINKING_TEXT_MESSAGE_CONTENT,
                     delta=part.content,
                 )
-            stream_ctx.part_end = TextMessageEndEvent(
-                message_id=message_id,
-            )
-        elif isinstance(part, ToolCallPart):  # pragma: no branch
-            message_id = stream_ctx.message_id or stream_ctx.new_message_id()
-            yield ToolCallStartEvent(
-                tool_call_id=part.tool_call_id,
-                tool_call_name=part.tool_name,
-                parent_message_id=message_id,
-            )
-            if part.args:
-                yield ToolCallArgsEvent(
+                stream_ctx.part_end = ThinkingTextMessageEndEvent(
+                    type=EventType.THINKING_TEXT_MESSAGE_END,
+                )
+        else:
+            if stream_ctx.thinking:
+                # TODO: Enable once https://github.com/ag-ui-protocol/ag-ui/issues/289 is resolved.
+                # yield ThinkingEndEvent(
+                #     type=EventType.THINKING_END,
+                # )
+                stream_ctx.thinking = False
+
+            if isinstance(part, TextPart):
+                message_id = stream_ctx.new_message_id()
+                yield TextMessageStartEvent(
+                    message_id=message_id,
+                )
+                if part.content:  # pragma: no branch
+                    yield TextMessageContentEvent(
+                        message_id=message_id,
+                        delta=part.content,
+                    )
+                stream_ctx.part_end = TextMessageEndEvent(
+                    message_id=message_id,
+                )
+            elif isinstance(part, ToolCallPart):  # pragma: no branch
+                message_id = stream_ctx.message_id or stream_ctx.new_message_id()
+                yield ToolCallStartEvent(
+                    tool_call_id=part.tool_call_id,
+                    tool_call_name=part.tool_name,
+                    parent_message_id=message_id,
+                )
+                if part.args:
+                    yield ToolCallArgsEvent(
+                        tool_call_id=part.tool_call_id,
+                        delta=part.args if isinstance(part.args, str) else json.dumps(part.args),
+                    )
+                stream_ctx.part_end = ToolCallEndEvent(
                     tool_call_id=part.tool_call_id,
-                    delta=part.args if isinstance(part.args, str) else json.dumps(part.args),
                 )
-            stream_ctx.part_end = ToolCallEndEvent(
-                tool_call_id=part.tool_call_id,
-            )
-
-        elif isinstance(part, ThinkingPart):  # pragma: no branch
-            yield ThinkingTextMessageStartEvent(
-                type=EventType.THINKING_TEXT_MESSAGE_START,
-            )
-            # Always send the content even if it's empty, as it may be
-            # used to indicate the start of thinking.
-            yield ThinkingTextMessageContentEvent(
-                type=EventType.THINKING_TEXT_MESSAGE_CONTENT,
-                delta=part.content,
-            )
-            stream_ctx.part_end = ThinkingTextMessageEndEvent(
-                type=EventType.THINKING_TEXT_MESSAGE_END,
-            )
 
     elif isinstance(agent_event, PartDeltaEvent):
         delta = agent_event.delta
         if isinstance(delta, TextPartDelta):
-            yield TextMessageContentEvent(
-                message_id=stream_ctx.message_id,
-                delta=delta.content_delta,
-            )
+            if delta.content_delta:  # pragma: no branch
+                yield TextMessageContentEvent(
+                    message_id=stream_ctx.message_id,
+                    delta=delta.content_delta,
+                )
         elif isinstance(delta, ToolCallPartDelta):  # pragma: no branch
             assert delta.tool_call_id, '`ToolCallPartDelta.tool_call_id` must be set'
             yield ToolCallArgsEvent(
@@ -479,6 +501,14 @@ async def _handle_model_request_event(
             )
         elif isinstance(delta, ThinkingPartDelta):  # pragma: no branch
             if delta.content_delta:  # pragma: no branch
+                if not isinstance(stream_ctx.part_end, ThinkingTextMessageEndEvent):
+                    yield ThinkingTextMessageStartEvent(
+                        type=EventType.THINKING_TEXT_MESSAGE_START,
+                    )
+                    stream_ctx.part_end = ThinkingTextMessageEndEvent(
+                        type=EventType.THINKING_TEXT_MESSAGE_END,
+                    )
+
                 yield ThinkingTextMessageContentEvent(
                     type=EventType.THINKING_TEXT_MESSAGE_CONTENT,
                     delta=delta.content_delta,
@@ -515,7 +545,7 @@ async def _handle_tool_result_event(
     content = result.content
     if isinstance(content, BaseEvent):
         yield content
-    elif isinstance(content, (str, bytes)):  # pragma: no branch
+    elif isinstance(content, str | bytes):  # pragma: no branch
         # Avoid iterable check for strings and bytes.
         pass
     elif isinstance(content, Iterable):  # pragma: no branch
@@ -630,6 +660,7 @@ class _RequestStreamContext:
 
     message_id: str = ''
     part_end: BaseEvent | None = None
+    thinking: bool = False
 
     def new_message_id(self) -> str:
         """Generate a new message ID for the request stream.
@@ -681,7 +712,7 @@ class _ToolCallNotFoundError(_RunError, ValueError):
         )
 
 
-class _AGUIFrontendToolset(DeferredToolset[AgentDepsT]):
+class _AGUIFrontendToolset(ExternalToolset[AgentDepsT]):
     def __init__(self, tools: list[AGUITool]):
         super().__init__(
             [

pydantic_ai/agent/__init__.py

@@ -5,14 +5,14 @@ import inspect
 import json
 import warnings
 from asyncio import Lock
-from collections.abc import AsyncIterator, Awaitable, Iterator, Sequence
+from collections.abc import AsyncIterator, Awaitable, Callable, Iterator, Sequence
 from contextlib import AbstractAsyncContextManager, AsyncExitStack, asynccontextmanager, contextmanager
 from contextvars import ContextVar
-from typing import TYPE_CHECKING, Any, Callable, ClassVar, cast, overload
+from typing import TYPE_CHECKING, Any, ClassVar, cast, overload
 
 from opentelemetry.trace import NoOpTracer, use_span
 from pydantic.json_schema import GenerateJsonSchema
-from typing_extensions import TypeVar, deprecated
+from typing_extensions import Self, TypeVar, deprecated
 
 from pydantic_graph import Graph
 
@@ -45,10 +45,15 @@ from ..run import AgentRun, AgentRunResult
 from ..settings import ModelSettings, merge_model_settings
 from ..tools import (
     AgentDepsT,
+    DeferredToolCallResult,
+    DeferredToolResult,
+    DeferredToolResults,
     DocstringFormat,
     GenerateToolJsonSchema,
     RunContext,
     Tool,
+    ToolApproved,
+    ToolDenied,
     ToolFuncContext,
     ToolFuncEither,
     ToolFuncPlain,
@@ -321,7 +326,7 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]):
 
         self._instructions = ''
         self._instructions_functions = []
-        if isinstance(instructions, (str, Callable)):
+        if isinstance(instructions, str | Callable):
             instructions = [instructions]
         for instruction in instructions or []:
             if isinstance(instruction, str):
@@ -346,7 +351,9 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]):
         if self._output_toolset:
             self._output_toolset.max_retries = self._max_result_retries
 
-        self._function_toolset = _AgentFunctionToolset(tools, max_retries=self._max_tool_retries)
+        self._function_toolset = _AgentFunctionToolset(
+            tools, max_retries=self._max_tool_retries, output_schema=self._output_schema
+        )
         self._dynamic_toolsets = [
             DynamicToolset[AgentDepsT](toolset_func=toolset)
             for toolset in toolsets or []
@@ -367,7 +374,7 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]):
             _utils.Option[Sequence[Tool[AgentDepsT] | ToolFuncEither[AgentDepsT, ...]]]
         ] = ContextVar('_override_tools', default=None)
 
-        self._enter_lock = _utils.get_async_lock()
+        self._enter_lock = Lock()
         self._entered_count = 0
         self._exit_stack = None
 
@@ -427,6 +434,7 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]):
         *,
         output_type: None = None,
         message_history: list[_messages.ModelMessage] | None = None,
+        deferred_tool_results: DeferredToolResults | None = None,
         model: models.Model | models.KnownModelName | str | None = None,
         deps: AgentDepsT = None,
         model_settings: ModelSettings | None = None,
@@ -443,6 +451,7 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]):
         *,
         output_type: OutputSpec[RunOutputDataT],
         message_history: list[_messages.ModelMessage] | None = None,
+        deferred_tool_results: DeferredToolResults | None = None,
         model: models.Model | models.KnownModelName | str | None = None,
         deps: AgentDepsT = None,
         model_settings: ModelSettings | None = None,
@@ -453,12 +462,13 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]):
     ) -> AbstractAsyncContextManager[AgentRun[AgentDepsT, RunOutputDataT]]: ...
 
     @asynccontextmanager
-    async def iter(
+    async def iter(  # noqa: C901
         self,
         user_prompt: str | Sequence[_messages.UserContent] | None = None,
         *,
         output_type: OutputSpec[RunOutputDataT] | None = None,
         message_history: list[_messages.ModelMessage] | None = None,
+        deferred_tool_results: DeferredToolResults | None = None,
         model: models.Model | models.KnownModelName | str | None = None,
         deps: AgentDepsT = None,
         model_settings: ModelSettings | None = None,
@@ -531,6 +541,7 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]):
             output_type: Custom output type to use for this run, `output_type` may only be used if the agent has no
                 output validators since output validators would expect an argument that matches the agent's output type.
             message_history: History of the conversation so far.
+            deferred_tool_results: Optional results for deferred tool calls in the message history.
             model: Optional model to use for this run, required if `model` was not set when creating the agent.
             deps: Optional dependencies to use for this run.
             model_settings: Optional settings to use for this model's request.
@@ -609,6 +620,23 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]):
             instrumentation_settings = None
             tracer = NoOpTracer()
 
+        tool_call_results: dict[str, DeferredToolResult] | None = None
+        if deferred_tool_results is not None:
+            tool_call_results = {}
+            for tool_call_id, approval in deferred_tool_results.approvals.items():
+                if approval is True:
+                    approval = ToolApproved()
+                elif approval is False:
+                    approval = ToolDenied()
+                tool_call_results[tool_call_id] = approval
+
+            if calls := deferred_tool_results.calls:
+                call_result_types = _utils.get_union_args(DeferredToolCallResult)
+                for tool_call_id, result in calls.items():
+                    if not isinstance(result, call_result_types):
+                        result = _messages.ToolReturn(result)
+                    tool_call_results[tool_call_id] = result
+
         graph_deps = _agent_graph.GraphAgentDeps[AgentDepsT, RunOutputDataT](
             user_deps=deps,
             prompt=user_prompt,
@@ -623,6 +651,7 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]):
             history_processors=self.history_processors,
             builtin_tools=list(self._builtin_tools),
             tool_manager=tool_manager,
+            tool_call_results=tool_call_results,
             tracer=tracer,
             get_instructions=get_instructions,
             instrumentation_settings=instrumentation_settings,
@@ -678,22 +707,28 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]):
         self, state: _agent_graph.GraphAgentState, usage: _usage.RunUsage, settings: InstrumentationSettings
     ):
         if settings.version == 1:
-            attr_name = 'all_messages_events'
-            value = [
-                InstrumentedModel.event_to_dict(e) for e in settings.messages_to_otel_events(state.message_history)
-            ]
+            attrs = {
+                'all_messages_events': json.dumps(
+                    [
+                        InstrumentedModel.event_to_dict(e)
+                        for e in settings.messages_to_otel_events(state.message_history)
+                    ]
+                )
+            }
         else:
-            attr_name = 'pydantic_ai.all_messages'
-            value = settings.messages_to_otel_messages(state.message_history)
+            attrs = {
+                'pydantic_ai.all_messages': json.dumps(settings.messages_to_otel_messages(state.message_history)),
+                **settings.system_instructions_attributes(self._instructions),
+            }
 
         return {
             **usage.opentelemetry_attributes(),
-            attr_name: json.dumps(value),
+            **attrs,
             'logfire.json_schema': json.dumps(
                 {
                     'type': 'object',
                     'properties': {
-                        attr_name: {'type': 'array'},
+                        **{attr: {'type': 'array'} for attr in attrs.keys()},
                         'final_result': {'type': 'object'},
                     },
                 }
@@ -970,6 +1005,7 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]):
         require_parameter_descriptions: bool = False,
         schema_generator: type[GenerateJsonSchema] = GenerateToolJsonSchema,
         strict: bool | None = None,
+        requires_approval: bool = False,
     ) -> Callable[[ToolFuncContext[AgentDepsT, ToolParams]], ToolFuncContext[AgentDepsT, ToolParams]]: ...
 
     def tool(
@@ -984,6 +1020,7 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]):
         require_parameter_descriptions: bool = False,
         schema_generator: type[GenerateJsonSchema] = GenerateToolJsonSchema,
         strict: bool | None = None,
+        requires_approval: bool = False,
     ) -> Any:
         """Decorator to register a tool function which takes [`RunContext`][pydantic_ai.tools.RunContext] as its first argument.
 
@@ -1028,6 +1065,8 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]):
             schema_generator: The JSON schema generator class to use for this tool. Defaults to `GenerateToolJsonSchema`.
             strict: Whether to enforce JSON schema compliance (only affects OpenAI).
                 See [`ToolDefinition`][pydantic_ai.tools.ToolDefinition] for more info.
+            requires_approval: Whether this tool requires human-in-the-loop approval. Defaults to False.
+                See the [tools documentation](../deferred-tools.md#human-in-the-loop-tool-approval) for more info.
         """
 
         def tool_decorator(
@@ -1044,6 +1083,7 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]):
                 require_parameter_descriptions,
                 schema_generator,
                 strict,
+                requires_approval,
             )
             return func_
 
@@ -1064,6 +1104,7 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]):
         require_parameter_descriptions: bool = False,
         schema_generator: type[GenerateJsonSchema] = GenerateToolJsonSchema,
         strict: bool | None = None,
+        requires_approval: bool = False,
     ) -> Callable[[ToolFuncPlain[ToolParams]], ToolFuncPlain[ToolParams]]: ...
 
     def tool_plain(
@@ -1078,6 +1119,7 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]):
         require_parameter_descriptions: bool = False,
         schema_generator: type[GenerateJsonSchema] = GenerateToolJsonSchema,
         strict: bool | None = None,
+        requires_approval: bool = False,
     ) -> Any:
         """Decorator to register a tool function which DOES NOT take `RunContext` as an argument.
 
@@ -1122,6 +1164,8 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]):
             schema_generator: The JSON schema generator class to use for this tool. Defaults to `GenerateToolJsonSchema`.
             strict: Whether to enforce JSON schema compliance (only affects OpenAI).
                 See [`ToolDefinition`][pydantic_ai.tools.ToolDefinition] for more info.
+            requires_approval: Whether this tool requires human-in-the-loop approval. Defaults to False.
+                See the [tools documentation](../deferred-tools.md#human-in-the-loop-tool-approval) for more info.
         """
 
         def tool_decorator(func_: ToolFuncPlain[ToolParams]) -> ToolFuncPlain[ToolParams]:
@@ -1136,6 +1180,7 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]):
                 require_parameter_descriptions,
                 schema_generator,
                 strict,
+                requires_approval,
             )
             return func_
 
@@ -1279,7 +1324,9 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]):
         toolsets: list[AbstractToolset[AgentDepsT]] = []
 
         if some_tools := self._override_tools.get():
-            function_toolset = _AgentFunctionToolset(some_tools.value, max_retries=self._max_tool_retries)
+            function_toolset = _AgentFunctionToolset(
+                some_tools.value, max_retries=self._max_tool_retries, output_schema=self._output_schema
+            )
         else:
             function_toolset = self._function_toolset
         toolsets.append(function_toolset)
@@ -1308,7 +1355,7 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]):
 
         return schema  # pyright: ignore[reportReturnType]
 
-    async def __aenter__(self) -> AbstractAgent[AgentDepsT, OutputDataT]:
+    async def __aenter__(self) -> Self:
         """Enter the agent context.
 
         This will start all [`MCPServerStdio`s][pydantic_ai.mcp.MCPServerStdio] registered as `toolsets` so they are ready to be used.
@@ -1376,6 +1423,19 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]):
 
 @dataclasses.dataclass(init=False)
 class _AgentFunctionToolset(FunctionToolset[AgentDepsT]):
+    output_schema: _output.BaseOutputSchema[Any]
+
+    def __init__(
+        self,
+        tools: Sequence[Tool[AgentDepsT] | ToolFuncEither[AgentDepsT, ...]] = [],
+        *,
+        max_retries: int = 1,
+        id: str | None = None,
+        output_schema: _output.BaseOutputSchema[Any],
+    ):
+        self.output_schema = output_schema
+        super().__init__(tools, max_retries=max_retries, id=id)
+
     @property
     def id(self) -> str:
         return '<agent>'
@@ -1383,3 +1443,10 @@ class _AgentFunctionToolset(FunctionToolset[AgentDepsT]):
     @property
     def label(self) -> str:
         return 'the agent'
+
+    def add_tool(self, tool: Tool[AgentDepsT]) -> None:
+        if tool.requires_approval and not self.output_schema.allows_deferred_tools:
+            raise exceptions.UserError(
+                'To use tools that require approval, add `DeferredToolRequests` to the list of output types for this agent.'
+            )
+        super().add_tool(tool)