feedback: Shove the validation context inside the RunContext

Author: NicolasPllr1

pydantic_ai_slim/pydantic_ai/_agent_graph.py

@@ -18,7 +18,7 @@
 
 from pydantic_ai._function_schema import _takes_ctx as is_takes_ctx  # type: ignore
 from pydantic_ai._instrumentation import DEFAULT_INSTRUMENTATION_VERSION
-from pydantic_ai._tool_manager import ToolManager, build_validation_context
+from pydantic_ai._tool_manager import ToolManager
 from pydantic_ai._utils import dataclasses_no_defaults_repr, get_union_args, is_async_callable, run_in_executor
 from pydantic_ai.builtin_tools import AbstractBuiltinTool
 from pydantic_graph import BaseNode, GraphRunContext
@@ -144,6 +144,7 @@ class GraphAgentDeps(Generic[DepsT, OutputDataT]):
 
     output_schema: _output.OutputSchema[OutputDataT]
     output_validators: list[_output.OutputValidator[DepsT, OutputDataT]]
+    validation_context: Any | Callable[[RunContext[DepsT]], Any]
 
     history_processors: Sequence[HistoryProcessor[DepsT]]
 
@@ -477,6 +478,8 @@ async def _prepare_request(
         ctx.state.run_step += 1
 
         run_context = build_run_context(ctx)
+        validation_context = build_validation_context(ctx.deps.validation_context, run_context)
+        run_context = replace(run_context, validation_context=validation_context)
 
         # This will raise errors for any tool name conflicts
         ctx.deps.tool_manager = await ctx.deps.tool_manager.for_run_step(run_context)
@@ -720,9 +723,11 @@ async def _handle_text_response(
         text_processor: _output.BaseOutputProcessor[NodeRunEndT],
     ) -> ModelRequestNode[DepsT, NodeRunEndT] | End[result.FinalResult[NodeRunEndT]]:
         run_context = build_run_context(ctx)
-        validation_context = build_validation_context(ctx.deps.tool_manager.validation_ctx, run_context)
+        validation_context = build_validation_context(ctx.deps.validation_context, run_context)
 
-        result_data = await text_processor.process(text, run_context=run_context, validation_context=validation_context)
+        run_context = replace(run_context, validation_context=validation_context)
+
+        result_data = await text_processor.process(text, run_context=run_context)
 
         for validator in ctx.deps.output_validators:
             result_data = await validator.validate(result_data, run_context)
@@ -773,6 +778,7 @@ def build_run_context(ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT
         usage=ctx.state.usage,
         prompt=ctx.deps.prompt,
         messages=ctx.state.message_history,
+        validation_context=None,
         tracer=ctx.deps.tracer,
         trace_include_content=ctx.deps.instrumentation_settings is not None
         and ctx.deps.instrumentation_settings.include_content,
@@ -784,6 +790,18 @@ def build_run_context(ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT
     )
 
 
+def build_validation_context(
+    validation_ctx: Any | Callable[[RunContext[DepsT]], Any],
+    run_context: RunContext[DepsT],
+) -> Any:
+    """Build a Pydantic validation context, potentially from the current agent run context."""
+    if callable(validation_ctx):
+        fn = cast(Callable[[RunContext[DepsT]], Any], validation_ctx)
+        return fn(run_context)
+    else:
+        return validation_ctx
+
+
 async def process_tool_calls(  # noqa: C901
     tool_manager: ToolManager[DepsT],
     tool_calls: list[_messages.ToolCallPart],

pydantic_ai_slim/pydantic_ai/_output.py

@@ -531,7 +531,6 @@ async def process(
         data: str,
         *,
         run_context: RunContext[AgentDepsT],
-        validation_context: Any | None = None,
         allow_partial: bool = False,
         wrap_validation_errors: bool = True,
     ) -> OutputDataT:
@@ -566,7 +565,6 @@ async def process(
         return await self.wrapped.process(
             text,
             run_context=run_context,
-            validation_context=validation_context,
             allow_partial=allow_partial,
             wrap_validation_errors=wrap_validation_errors,
         )
@@ -648,7 +646,6 @@ async def process(
         data: str | dict[str, Any] | None,
         *,
         run_context: RunContext[AgentDepsT],
-        validation_context: Any | None = None,
         allow_partial: bool = False,
         wrap_validation_errors: bool = True,
     ) -> OutputDataT:
@@ -665,7 +662,7 @@ async def process(
             Either the validated output data (left) or a retry message (right).
         """
         try:
-            output = self.validate(data, allow_partial, validation_context)
+            output = self.validate(data, allow_partial, run_context.validation_context)
         except ValidationError as e:
             if wrap_validation_errors:
                 m = _messages.RetryPromptPart(
@@ -814,14 +811,12 @@ async def process(
         data: str,
         *,
         run_context: RunContext[AgentDepsT],
-        validation_context: Any | None = None,
         allow_partial: bool = False,
         wrap_validation_errors: bool = True,
     ) -> OutputDataT:
         union_object = await self._union_processor.process(
             data,
             run_context=run_context,
-            validation_context=validation_context,
             allow_partial=allow_partial,
             wrap_validation_errors=wrap_validation_errors,
         )
@@ -841,7 +836,6 @@ async def process(
         return await processor.process(
             inner_data,
             run_context=run_context,
-            validation_context=validation_context,
             allow_partial=allow_partial,
             wrap_validation_errors=wrap_validation_errors,
         )

pydantic_ai_slim/pydantic_ai/_run_context.py

@@ -3,7 +3,7 @@
 import dataclasses
 from collections.abc import Sequence
 from dataclasses import field
-from typing import TYPE_CHECKING, Generic
+from typing import TYPE_CHECKING, Any, Generic
 
 from opentelemetry.trace import NoOpTracer, Tracer
 from typing_extensions import TypeVar
@@ -38,6 +38,8 @@ class RunContext(Generic[RunContextAgentDepsT]):
     """The original user prompt passed to the run."""
     messages: list[_messages.ModelMessage] = field(default_factory=list)
     """Messages exchanged in the conversation so far."""
+    validation_context: Any = None
+    """Additional Pydantic validation context for the run outputs."""
     tracer: Tracer = field(default_factory=NoOpTracer)
     """The tracer to use for tracing the run."""
     trace_include_content: bool = False

pydantic_ai_slim/pydantic_ai/_tool_manager.py

@@ -1,11 +1,11 @@
 from __future__ import annotations
 
 import json
-from collections.abc import Callable, Iterator
+from collections.abc import Iterator
 from contextlib import contextmanager
 from contextvars import ContextVar
 from dataclasses import dataclass, field, replace
-from typing import Any, Generic, cast
+from typing import Any, Generic
 
 from opentelemetry.trace import Tracer
 from pydantic import ValidationError
@@ -31,8 +31,6 @@ class ToolManager(Generic[AgentDepsT]):
     """The toolset that provides the tools for this run step."""
     ctx: RunContext[AgentDepsT] | None = None
     """The agent run context for a specific run step."""
-    validation_ctx: Any | Callable[[RunContext[AgentDepsT]], Any] = None
-    """Additional Pydantic validation context for the run."""
     tools: dict[str, ToolsetTool[AgentDepsT]] | None = None
     """The cached tools for this run step."""
     failed_tools: set[str] = field(default_factory=set)
@@ -63,7 +61,6 @@ async def for_run_step(self, ctx: RunContext[AgentDepsT]) -> ToolManager[AgentDe
         return self.__class__(
             toolset=self.toolset,
             ctx=ctx,
-            validation_ctx=self.validation_ctx,
             tools=await self.toolset.get_tools(ctx),
         )
 
@@ -164,17 +161,15 @@ async def _call_tool(
                 partial_output=allow_partial,
             )
 
-            validation_ctx = build_validation_context(self.validation_ctx, self.ctx)
-
             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, context=validation_ctx
+                    call.args or '{}', allow_partial=pyd_allow_partial, context=ctx.validation_context
                 )
             else:
                 args_dict = validator.validate_python(
-                    call.args or {}, allow_partial=pyd_allow_partial, context=validation_ctx
+                    call.args or {}, allow_partial=pyd_allow_partial, context=ctx.validation_context
                 )
 
             result = await self.toolset.call_tool(name, args_dict, ctx, tool)
@@ -279,14 +274,3 @@ async def _call_function_tool(
                 )
 
         return tool_result
-
-
-def build_validation_context(
-    validation_ctx: Any | Callable[[RunContext[AgentDepsT]], Any], run_context: RunContext[AgentDepsT]
-) -> Any:
-    """Build a Pydantic validation context, potentially from the current agent run context."""
-    if callable(validation_ctx):
-        fn = cast(Callable[[RunContext[AgentDepsT]], Any], validation_ctx)
-        return fn(run_context)
-    else:
-        return validation_ctx

pydantic_ai_slim/pydantic_ai/agent/__init__.py

@@ -569,7 +569,7 @@ async def main():
                 output_toolset.max_retries = self._max_result_retries
                 output_toolset.output_validators = output_validators
         toolset = self._get_toolset(output_toolset=output_toolset, additional_toolsets=toolsets)
-        tool_manager = ToolManager[AgentDepsT](toolset, validation_ctx=self._validation_context)
+        tool_manager = ToolManager[AgentDepsT](toolset)
 
         # Build the graph
         graph = _agent_graph.build_agent_graph(self.name, self._deps_type, output_type_)
@@ -619,6 +619,7 @@ async def get_instructions(run_context: RunContext[AgentDepsT]) -> str | None:
             end_strategy=self.end_strategy,
             output_schema=output_schema,
             output_validators=output_validators,
+            validation_context=self._validation_context,
             history_processors=self.history_processors,
             builtin_tools=[*self._builtin_tools, *(builtin_tools or [])],
             tool_manager=tool_manager,

pydantic_ai_slim/pydantic_ai/result.py

@@ -18,7 +18,7 @@
     TextOutputSchema,
 )
 from ._run_context import AgentDepsT, RunContext
-from ._tool_manager import ToolManager, build_validation_context
+from ._tool_manager import ToolManager
 from .messages import ModelResponseStreamEvent
 from .output import (
     DeferredToolRequests,
@@ -197,12 +197,9 @@ async def validate_response_output(
                     # not part of the final result output, so we reset the accumulated text
                     text = ''
 
-            validation_context = build_validation_context(self._tool_manager.validation_ctx, self._run_ctx)
-
             result_data = await text_processor.process(
                 text,
                 run_context=self._run_ctx,
-                validation_context=validation_context,
                 allow_partial=allow_partial,
                 wrap_validation_errors=False,
             )