Minor clean up in preparation of graph agent (#779)

Author: dmontagu

pydantic_ai_slim/pydantic_ai/agent.py

@@ -60,7 +60,7 @@
 - `'early'`: Stop processing other tool calls once a final result is found
 - `'exhaustive'`: Process all tool calls even after finding a final result
 """
-RunResultData = TypeVar('RunResultData')
+RunResultDataT = TypeVar('RunResultDataT')
 """Type variable for the result data of a run where `result_type` was customized on the run call."""
 
 
@@ -214,15 +214,15 @@ async def run(
         self,
         user_prompt: str,
         *,
-        result_type: type[RunResultData],
+        result_type: type[RunResultDataT],
         message_history: list[_messages.ModelMessage] | None = None,
         model: models.Model | models.KnownModelName | None = None,
         deps: AgentDepsT = None,
         model_settings: ModelSettings | None = None,
         usage_limits: _usage.UsageLimits | None = None,
         usage: _usage.Usage | None = None,
         infer_name: bool = True,
-    ) -> result.RunResult[RunResultData]: ...
+    ) -> result.RunResult[RunResultDataT]: ...
 
     async def run(
         self,
@@ -234,7 +234,7 @@ async def run(
         model_settings: ModelSettings | None = None,
         usage_limits: _usage.UsageLimits | None = None,
         usage: _usage.Usage | None = None,
-        result_type: type[RunResultData] | None = None,
+        result_type: type[RunResultDataT] | None = None,
         infer_name: bool = True,
     ) -> result.RunResult[Any]:
         """Run the agent with a user prompt in async mode.
@@ -352,21 +352,21 @@ def run_sync(
         self,
         user_prompt: str,
         *,
-        result_type: type[RunResultData] | None,
+        result_type: type[RunResultDataT] | None,
         message_history: list[_messages.ModelMessage] | None = None,
         model: models.Model | models.KnownModelName | None = None,
         deps: AgentDepsT = None,
         model_settings: ModelSettings | None = None,
         usage_limits: _usage.UsageLimits | None = None,
         usage: _usage.Usage | None = None,
         infer_name: bool = True,
-    ) -> result.RunResult[RunResultData]: ...
+    ) -> result.RunResult[RunResultDataT]: ...
 
     def run_sync(
         self,
         user_prompt: str,
         *,
-        result_type: type[RunResultData] | None = None,
+        result_type: type[RunResultDataT] | None = None,
         message_history: list[_messages.ModelMessage] | None = None,
         model: models.Model | models.KnownModelName | None = None,
         deps: AgentDepsT = None,
@@ -442,22 +442,22 @@ def run_stream(
         self,
         user_prompt: str,
         *,
-        result_type: type[RunResultData],
+        result_type: type[RunResultDataT],
         message_history: list[_messages.ModelMessage] | None = None,
         model: models.Model | models.KnownModelName | None = None,
         deps: AgentDepsT = None,
         model_settings: ModelSettings | None = None,
         usage_limits: _usage.UsageLimits | None = None,
         usage: _usage.Usage | None = None,
         infer_name: bool = True,
-    ) -> AbstractAsyncContextManager[result.StreamedRunResult[AgentDepsT, RunResultData]]: ...
+    ) -> AbstractAsyncContextManager[result.StreamedRunResult[AgentDepsT, RunResultDataT]]: ...
 
     @asynccontextmanager
     async def run_stream(
         self,
         user_prompt: str,
         *,
-        result_type: type[RunResultData] | None = None,
+        result_type: type[RunResultDataT] | None = None,
         message_history: list[_messages.ModelMessage] | None = None,
         model: models.Model | models.KnownModelName | None = None,
         deps: AgentDepsT = None,
@@ -572,7 +572,7 @@ async def on_complete():
                                 # there are result validators that might convert the result data from an overridden
                                 # `result_type` to a type that is not valid as such.
                                 result_validators = cast(
-                                    list[_result.ResultValidator[AgentDepsT, RunResultData]], self._result_validators
+                                    list[_result.ResultValidator[AgentDepsT, RunResultDataT]], self._result_validators
                                 )
 
                                 yield result.StreamedRunResult(
@@ -999,7 +999,7 @@ async def _get_model(self, model: models.Model | models.KnownModelName | None) -
         return model_
 
     async def _prepare_model(
-        self, run_context: RunContext[AgentDepsT], result_schema: _result.ResultSchema[RunResultData] | None
+        self, run_context: RunContext[AgentDepsT], result_schema: _result.ResultSchema[RunResultDataT] | None
     ) -> models.AgentModel:
         """Build tools and create an agent model."""
         function_tools: list[ToolDefinition] = []
@@ -1035,8 +1035,8 @@ async def _reevaluate_dynamic_prompts(
                                 )
 
     def _prepare_result_schema(
-        self, result_type: type[RunResultData] | None
-    ) -> _result.ResultSchema[RunResultData] | None:
+        self, result_type: type[RunResultDataT] | None
+    ) -> _result.ResultSchema[RunResultDataT] | None:
         if result_type is not None:
             if self._result_validators:
                 raise exceptions.UserError('Cannot set a custom run `result_type` when the agent has result validators')
@@ -1053,7 +1053,7 @@ async def _prepare_messages(
         run_context: RunContext[AgentDepsT],
     ) -> list[_messages.ModelMessage]:
         try:
-            ctx_messages = _messages_ctx_var.get()
+            ctx_messages = get_captured_run_messages()
         except LookupError:
             messages: list[_messages.ModelMessage] = []
         else:
@@ -1080,8 +1080,8 @@ async def _handle_model_response(
         self,
         model_response: _messages.ModelResponse,
         run_context: RunContext[AgentDepsT],
-        result_schema: _result.ResultSchema[RunResultData] | None,
-    ) -> tuple[_MarkFinalResult[RunResultData] | None, list[_messages.ModelRequestPart]]:
+        result_schema: _result.ResultSchema[RunResultDataT] | None,
+    ) -> tuple[_MarkFinalResult[RunResultDataT] | None, list[_messages.ModelRequestPart]]:
         """Process a non-streamed response from the model.
 
         Returns:
@@ -1110,11 +1110,11 @@ async def _handle_model_response(
             raise exceptions.UnexpectedModelBehavior('Received empty model response')
 
     async def _handle_text_response(
-        self, text: str, run_context: RunContext[AgentDepsT], result_schema: _result.ResultSchema[RunResultData] | None
-    ) -> tuple[_MarkFinalResult[RunResultData] | None, list[_messages.ModelRequestPart]]:
+        self, text: str, run_context: RunContext[AgentDepsT], result_schema: _result.ResultSchema[RunResultDataT] | None
+    ) -> tuple[_MarkFinalResult[RunResultDataT] | None, list[_messages.ModelRequestPart]]:
         """Handle a plain text response from the model for non-streaming responses."""
         if self._allow_text_result(result_schema):
-            result_data_input = cast(RunResultData, text)
+            result_data_input = cast(RunResultDataT, text)
             try:
                 result_data = await self._validate_result(result_data_input, run_context, None)
             except _result.ToolRetryError as e:
@@ -1133,13 +1133,13 @@ async def _handle_structured_response(
         self,
         tool_calls: list[_messages.ToolCallPart],
         run_context: RunContext[AgentDepsT],
-        result_schema: _result.ResultSchema[RunResultData] | None,
-    ) -> tuple[_MarkFinalResult[RunResultData] | None, list[_messages.ModelRequestPart]]:
+        result_schema: _result.ResultSchema[RunResultDataT] | None,
+    ) -> tuple[_MarkFinalResult[RunResultDataT] | None, list[_messages.ModelRequestPart]]:
         """Handle a structured response containing tool calls from the model for non-streaming responses."""
         assert tool_calls, 'Expected at least one tool call'
 
         # first look for the result tool call
-        final_result: _MarkFinalResult[RunResultData] | None = None
+        final_result: _MarkFinalResult[RunResultDataT] | None = None
 
         parts: list[_messages.ModelRequestPart] = []
         if result_schema is not None:
@@ -1168,7 +1168,7 @@ async def _process_function_tools(
         tool_calls: list[_messages.ToolCallPart],
         result_tool_name: str | None,
         run_context: RunContext[AgentDepsT],
-        result_schema: _result.ResultSchema[RunResultData] | None,
+        result_schema: _result.ResultSchema[RunResultDataT] | None,
     ) -> list[_messages.ModelRequestPart]:
         """Process function (non-result) tool calls in parallel.
 
@@ -1227,7 +1227,7 @@ async def _handle_streamed_response(
         self,
         streamed_response: models.StreamedResponse,
         run_context: RunContext[AgentDepsT],
-        result_schema: _result.ResultSchema[RunResultData] | None,
+        result_schema: _result.ResultSchema[RunResultDataT] | None,
     ) -> _MarkFinalResult[models.StreamedResponse] | tuple[_messages.ModelResponse, list[_messages.ModelRequestPart]]:
         """Process a streamed response from the model.
 
@@ -1282,15 +1282,15 @@ async def _handle_streamed_response(
 
     async def _validate_result(
         self,
-        result_data: RunResultData,
+        result_data: RunResultDataT,
         run_context: RunContext[AgentDepsT],
         tool_call: _messages.ToolCallPart | None,
-    ) -> RunResultData:
+    ) -> RunResultDataT:
         if self._result_validators:
             agent_result_data = cast(ResultDataT, result_data)
             for validator in self._result_validators:
                 agent_result_data = await validator.validate(agent_result_data, tool_call, run_context)
-            return cast(RunResultData, agent_result_data)
+            return cast(RunResultDataT, agent_result_data)
         else:
             return result_data
 
@@ -1315,7 +1315,7 @@ async def _sys_parts(self, run_context: RunContext[AgentDepsT]) -> list[_message
     def _unknown_tool(
         self,
         tool_name: str,
-        result_schema: _result.ResultSchema[RunResultData] | None,
+        result_schema: _result.ResultSchema[RunResultDataT] | None,
     ) -> _messages.RetryPromptPart:
         names = list(self._function_tools.keys())
         if result_schema:
@@ -1358,7 +1358,7 @@ def _infer_name(self, function_frame: FrameType | None) -> None:
                             return
 
     @staticmethod
-    def _allow_text_result(result_schema: _result.ResultSchema[RunResultData] | None) -> bool:
+    def _allow_text_result(result_schema: _result.ResultSchema[RunResultDataT] | None) -> bool:
         return result_schema is None or result_schema.allow_text_result
 
     @property
@@ -1413,6 +1413,10 @@ def capture_run_messages() -> Iterator[list[_messages.ModelMessage]]:
             _messages_ctx_var.reset(token)
 
 
+def get_captured_run_messages() -> _RunMessages:
+    return _messages_ctx_var.get()
+
+
 @dataclasses.dataclass
 class _MarkFinalResult(Generic[ResultDataT]):
     """Marker class to indicate that the result is the final result.

pydantic_ai_slim/pydantic_ai/tools.py

@@ -158,6 +158,9 @@ class Tool(Generic[AgentDepsT]):
     _var_positional_field: str | None = field(init=False)
     _validator: SchemaValidator = field(init=False, repr=False)
     _parameters_json_schema: ObjectJsonSchema = field(init=False)
+
+    # TODO: Move this state off the Tool class, which is otherwise stateless.
+    #   This should be tracked inside a specific agent run, not the tool.
     current_retry: int = field(default=0, init=False)
 
     def __init__(
@@ -261,7 +264,7 @@ async def prepare_tool_def(self, ctx: RunContext[AgentDepsT]) -> ToolDefinition
 
     async def run(
         self, message: _messages.ToolCallPart, run_context: RunContext[AgentDepsT]
-    ) -> _messages.ModelRequestPart:
+    ) -> _messages.ToolReturnPart | _messages.RetryPromptPart:
         """Run the tool function asynchronously."""
         try:
             if isinstance(message.args, str):

pydantic_graph/pydantic_graph/graph.py

@@ -4,6 +4,7 @@
 import inspect
 import types
 from collections.abc import Sequence
+from contextlib import ExitStack
 from dataclasses import dataclass, field
 from functools import cached_property
 from pathlib import Path
@@ -75,6 +76,7 @@ async def run(self, ctx: GraphRunContext) -> Increment | End[int]:
     snapshot_state: Callable[[StateT], StateT]
     _state_type: type[StateT] | _utils.Unset = field(repr=False)
     _run_end_type: type[RunEndT] | _utils.Unset = field(repr=False)
+    _auto_instrument: bool = field(repr=False)
 
     def __init__(
         self,
@@ -84,6 +86,7 @@ def __init__(
         state_type: type[StateT] | _utils.Unset = _utils.UNSET,
         run_end_type: type[RunEndT] | _utils.Unset = _utils.UNSET,
         snapshot_state: Callable[[StateT], StateT] = deep_copy_state,
+        auto_instrument: bool = True,
     ):
         """Create a graph from a sequence of nodes.
 
@@ -97,10 +100,12 @@ def __init__(
             snapshot_state: A function to snapshot the state of the graph, this is used in
                 [`NodeStep`][pydantic_graph.state.NodeStep] and [`EndStep`][pydantic_graph.state.EndStep] to record
                 the state before each step.
+            auto_instrument: Whether to create a span for the graph run and the execution of each node's run method.
         """
         self.name = name
         self._state_type = state_type
         self._run_end_type = run_end_type
+        self._auto_instrument = auto_instrument
         self.snapshot_state = snapshot_state
 
         parent_namespace = _utils.get_parent_namespace(inspect.currentframe())
@@ -155,26 +160,32 @@ async def main():
             self._infer_name(inspect.currentframe())
 
         history: list[HistoryStep[StateT, T]] = []
-        with _logfire.span(
-            '{graph_name} run {start=}',
-            graph_name=self.name or 'graph',
-            start=start_node,
-        ) as run_span:
-            while True:
-                next_node = await self.next(start_node, history, state=state, deps=deps, infer_name=False)
-                if isinstance(next_node, End):
-                    history.append(EndStep(result=next_node))
+        with ExitStack() as stack:
+            run_span: logfire_api.LogfireSpan | None = None
+            if self._auto_instrument:
+                run_span = stack.enter_context(
+                    _logfire.span(
+                        '{graph_name} run {start=}',
+                        graph_name=self.name or 'graph',
+                        start=start_node,
+                    )
+                )
+        while True:
+            next_node = await self.next(start_node, history, state=state, deps=deps, infer_name=False)
+            if isinstance(next_node, End):
+                history.append(EndStep(result=next_node))
+                if run_span is not None:
                     run_span.set_attribute('history', history)
-                    return next_node.data, history
-                elif isinstance(next_node, BaseNode):
-                    start_node = next_node
+                return next_node.data, history
+            elif isinstance(next_node, BaseNode):
+                start_node = next_node
+            else:
+                if TYPE_CHECKING:
+                    typing_extensions.assert_never(next_node)
                 else:
-                    if TYPE_CHECKING:
-                        typing_extensions.assert_never(next_node)
-                    else:
-                        raise exceptions.GraphRuntimeError(
-                            f'Invalid node return type: `{type(next_node).__name__}`. Expected `BaseNode` or `End`.'
-                        )
+                    raise exceptions.GraphRuntimeError(
+                        f'Invalid node return type: `{type(next_node).__name__}`. Expected `BaseNode` or `End`.'
+                    )
 
     def run_sync(
         self: Graph[StateT, DepsT, T],
@@ -232,8 +243,10 @@ async def next(
         if node_id not in self.node_defs:
             raise exceptions.GraphRuntimeError(f'Node `{node}` is not in the graph.')
 
-        ctx = GraphRunContext(state, deps)
-        with _logfire.span('run node {node_id}', node_id=node_id, node=node):
+        with ExitStack() as stack:
+            if self._auto_instrument:
+                stack.enter_context(_logfire.span('run node {node_id}', node_id=node_id, node=node))
+            ctx = GraphRunContext(state, deps)
             start_ts = _utils.now_utc()
             start = perf_counter()
             next_node = await node.run(ctx)