Passing Pydantic validation context for output and tool arguments validation

[NicolasPllr1]

The problem

When running an agent, I want to influence the validation of its output with a context object.

My use-case and an example

I'm using Pydantic AI's agents primarily for their structured-output capabilities.

Concretely, I have LLMs generating complex Pydantic models with custom validators. Some of these validators use the info.context property to gain access to dynamic validation data (relevant pydantic doc).

Example

`Image` example with a custom validator on its `source` field

The source is a string but that's not enough. I want to validate that it's within a dynamic list of allowed sources:

You can think of this list of available images as provided by the user at runtime, and changing over time (in particular its "value" can't be known at schema/agent-definition time).

from pydantic import BaseModel, ValidationInfo, field_validator

from pydantic_ai import Agent, UnexpectedModelBehavior
from pydantic_ai.models.test import TestModel


class Image(BaseModel):
    source: str

    @field_validator('source', mode='after')
    def check_allowed_src(cls, source: str, info: ValidationInfo) -> str:
        if isinstance(info.context, dict):
            allowed_src = info.context.get('allowed_src', [])
        else:
            allowed_src = []

        if source not in allowed_src:
            raise ValueError(f"Source '{source}' is not available. Choose from: {allowed_src}")
        return source


agent = Agent(
    model=TestModel(custom_output_args={'src': 'dog.png'}),
    output_type=Image,  # Image needs a validation context to validate successfully
    retries=1,
)

try:
    agent.run_sync(
        user_prompt='lorem'
    )  # the Image validation will fail as no context will be provided to the Pydantic model
except UnexpectedModelBehavior as e:
    print(e)
    'Exceeded maximum retries (1) for output validation'

Goal and motivation

Essentially, I'm looking for a way to provide the context: Any | None argument to the model_validate / model_validate_json calls that Pydantic AI ultimately perform in an agent run.

For comparison, with Instructor one can pass this context object to the 'patched' completion API. This context is then drilled down until it arrives to a Pydantic validation function call where it's used.

Personally, it's while porting my AI workflows code from Instructor to Pydantic AI that I ran into this context problem.

So far, I have not found a solution reading through Pydantic AI's docs and code, and I don't think there is a "simple" one if any. But I may be wrong here!

Options I considered:

• Output function / ToolOutput : execute after the model_type has been validated. I saw this comment for instance. But I need to influence the model_type validation itself.
  • One would need to define their models twice. One with only static validation constraints, and another with both static and "dynamic" validation using the context. I don't think this is practical.
• Dynamic tools with prepare : allow for tweaking models/schemas before validation but does not expose the context object which could be used during validation.

So I don't think those options are really adequate for my use case.

• Also, reading through this issue and as suggested in this particular comment, one could use a WrapperToolSet and override call_tool ?
  • But the comment continues with: "but that currently happens after arguments have already been validated against the schema" making me not very hopeful in this solution
  • Also, looking at the code in the _tool_manager module, the call to call_tool on the toolset happens after the validation against the output schema ... And in my local fork, it's precisely to those validation calls that I'm passing the validation context. I.e., before calling call_tool.

Contribution

I have a working fork where I added a new validation_ctx: Any | None = None argument to the relevant Agent run methods.

It's passed down to the GraphAgentDeps and RunContext, making it available when calling model_validate/ model_validate_json in:

• ToolManager._call_tool
• ObjectOutputProcessor.validate

Essentially, anytime there is the output_type arg, I added a new accompanying validation_ctx argument.

I'm not sure this is the design you guys would want for this feature or if you want this feature at all.
I was also surprised to not see this context exposed in Pydantic AI as it's a feature in Pydantic itself. Is that a conscious design choice on your part ?

Waiting on your feedback here before pushing any PR.

Thanks for the awesome library!

[DouweM]

@NicolasPllr1 Thanks for the report! We should definitely support this considering that we are Pydantic :D

Note that you can already get this to work using a combination of these features:

• https://ai.pydantic.dev/output/#structured-dict
• https://ai.pydantic.dev/output/#output-functions
• https://ai.pydantic.dev/dependencies/#dependencies

from __future__ import annotations

from dataclasses import dataclass
from typing import Any

import logfire
from pydantic import BaseModel, ValidationInfo, field_validator
from pydantic_ai import Agent, RunContext, StructuredDict, UnexpectedModelBehavior
from pydantic_ai.models.test import TestModel

# NOTE: Helpful during debugging https://ai.pydantic.dev/logfire/
logfire.configure()
logfire.instrument_pydantic_ai()


class Image(BaseModel):
    source: str

    @field_validator("source", mode="after")
    def check_allowed_src(cls, source: str, info: ValidationInfo) -> str:
        if isinstance(info.context, dict):
            allowed_src = info.context.get("allowed_src", [])
        else:
            allowed_src = []

        if source not in allowed_src:
            raise ValueError(
                f"Source '{source}' is not available. Choose from: {allowed_src}"
            )
        return source


# NOTE: This will use the same JSON schema as the `Image`, but will result in a `dict[str, Any]` without Pydantic validation
ImageDict = StructuredDict(Image.model_json_schema())


@dataclass
class ContextDeps:
    validation_ctx: dict[str, Any]


def validate_image(ctx: RunContext[ContextDeps], image_dict: ImageDict) -> Image:
	# NOTE: Here we validate the image dict with context from deps
    return Image.model_validate(image_dict, context=ctx.deps.validation_ctx)


agent = Agent(
    # NOTE: This said `src` before, which was incompatible with the `source` field anyway
    model=TestModel(custom_output_args={"source": "dog.png"}), 
    # NOTE: This means that the model will generate input args for the function (a dict matching the schema), 
    # then the function is called to create a real Image using validation context, 
    # and then that Image becomes the agent run output
    output_type=validate_image,   
    deps_type=ContextDeps,
    retries=1,
)

try:
    result = agent.run_sync(
        user_prompt="lorem",
        deps=ContextDeps(validation_ctx={"allowed_src": ["dog.png"]}),
    )  # the Image validation will fail as no context will be provided to the Pydantic model
except UnexpectedModelBehavior as e:
    print(e)
else:
    print(result.output)

As for built-in support, I'm hesitant to add a new arg to the run methods, and in Pydantic AI deps seem like the right place to put this type of contextual information that may be needed inside the run (like how RunContext with .deps is already being made available to all tool functions, output functions, etc).

We wouldn't want to always pass the deps as validation context, but it would be useful to be able to define for a given output type or tool args how deps or the entire run context can be converted into a validation context using a callable. In the example above that would just be lambda ctx: ctx.deps.validation_ctx.

That could be a validation_context argument on @agent.tool, Tool and ToolOutput that takes Any | Callable[[RunContext], Any]. In you case that would mean output_type=Image becomes output_type=ToolOutput(Image, validation_context=lambda ctx: ctx.deps.validation_ctx).

validation_context could also be an argument on the Agent constructor if you want to use it for all output (and Pydantic models used in tool args). Passing it at agent.run-time would then still be done through deps.

What do you think of these options?

[NicolasPllr1]

Thanks @DouweM for the rapid response!

Note that you can already get this to work

This is great, thanks a lot for the working code! I didn't have the StructuredDict in mind at all.

And it's much more practical/elegant than what I expected (I thought "one would need to define their models twice"). I will use that for now ;)

---

We should definitely support this considering that we are Pydantic :D

Aha I knew it! ;)

deps seem like the right place to put this type of contextual information that may be needed inside the run

I agree!
My first shot at this actually used deps, but I had hacked its use in downstream internals (reaching for .deps.validation_ctx or similar) which felt wrong.

Your approach with the Callable sounds awesome to me. It removes any weird hard-coded logic and is way more flexible than my proposal.

However on second thought, I believed that having to wrap the output type in what might seem like an arbitrary ToolOutput isn't ideal. In the sense that structured-outputs may not be associated with "tool calling" in the user's mind, even though that's the mechanism being used under the hood. From the user's perspective, they have their output type, they have their context object, why would they need to use this wrapper ToolOutput ?

But I think your next proposal:

validation_context could also be an argument on the Agent constructor if you want to use it for all output

makes it easier for users who just want an agent for 1-off structure-generation calls that include a validation context. No wrapping of the output type needed with this.
And It's also nice that the context would hit all outputs in the run.

--> I'm getting to work on this, will try to open a PR soon!

[DouweM]

However on second thought, I believed that having to wrap the output type in what might seem like an arbitrary ToolOutput isn't ideal.

@NicolasPllr1 Yeah I agree it'd be annoying to require the user to explicitly choose an output mode in this case instead of relying on the automatic output mode selection.

But I think your next proposal:

validation_context could also be an argument on the Agent constructor if you want to use it for all output

makes it easier for users who just want an agent for 1-off structure-generation calls that include a validation context. No wrapping of the output type needed with this. And It's also nice that the context would hit all outputs in the run.

That would require all Pydantic models to use the same validation context, though, which is a bit more limited, but I suppose that'd already be the case most of the time.

--> I'm getting to work on this, will try to open a PR soon!

Sounds good! Let me know if you have any questions.