Allow raw tool call arguments to be pre-processed in case schema validation fails

[warmwind]

Question

When I set the output type of an agent to a TypedDict, sometimes the AI returns incorrect JSON (see attached screenshot).

I’d like to know if there’s a way to preprocess or adjust the agent’s response after the AI generates it but before it’s passed to the final_result tool.

• Does output_validator handle this kind of preprocessing?
• Or do I need to use stream_responses and validate_response_output manually to intercept and fix responses?

Additional Context

No response

[m7mdhka]

Yes, the output_validator runs immediately after the agent generates a response, and it can trigger a ModelRetry if a JSON validation error occurs.

For example:

import json
import re

class JSONValidationError(Exception):
    """Raised when JSON parsing fails from markdown output."""
    pass


def parse_json_from_markdown(text: str) -> dict:
    code_block_pattern = re.compile(r"```(?:json)?\s*(\{.*?\})\s*```", re.DOTALL)
    match = code_block_pattern.search(text)
    json_str = match.group(1) if match else text.strip()
    try:
        return json.loads(json_str)
    except json.JSONDecodeError as e:
        raise JSONValidationError(f"Invalid JSON format: {e}\nExtracted text: {json_str}") from e


@agent.output_validator
def extract_json(output: str) -> dict:
    try:
        parsed_output = parse_json_from_markdown(output)
        return parsed_output
    except JSONValidationError as e:
        # If the output isn't valid JSON, retry model generation
        raise ModelRetry(f"Failed to parse json configuration: {e}") from e

---

👉 This way:

• The validator ensures only valid JSON chart configurations are returned.
• If parsing fails, ModelRetry is raised so the agent regenerates the output.
• You end up with a clean, validated dictionary.

[DouweM]

@m7mdhka That's not quite right; the output_validator is called after the final_result output tool arguments are validated to match the schema. If the output_type is set to a specific type (like a StructuredDict in this case), the output validator will receive an object of that type. It can add additional validation, but not modify the raw JSON before it is parsed into that type.

---

@warmwind If the model is failing to generate JSON matching the schema, there's not currently a way to handle that by fixing the JSON yourself before validation. What you could do is define an output function that takes a str and outputs the desired type, and do the validation/fixing in there. You can then share the desired JSON schema with the model via the description. E.g.:

async def parse_output(json: str) -> MyDict:
    # ...
    return parsed

agent = Agent(output_type=ToolOutput(parse_output, description=f'Pass JSON matching this schema: {schema}'))
result = await agent.run(...)

# result.output will be a MyDict

The other option would be to use a model that's better at outputting JSON.

[warmwind]

Hi @DouweM , thanks for the clarification and update!

May I ask why there’s currently no preprocess hook before final_result is called?
It would be very useful to have a way to inspect or adjust the model’s raw output (for example, fixing malformed JSON or adding custom validation) before it’s passed to final_result.

Do you think adding such a preprocess step is feasible in the future?

[DouweM]

@warmwind Custom additional validation is already supported through https://ai.pydantic.dev/output/#output-functions, but I agree it'd be useful to also allow raw tool call arguments to be pre-processed before schema validation. Currently, you'd be able to implement that using a custom WrapperModel that overrides request (and request_stream if you want to support streaming as well) and modifies the ToolCallParts as it sees fit. Is that an option for you?

[warmwind]

For now, I’m using stream_responses together with validate_response_output. When a ValidationError occurs, I manually fix it and then process the chunk data.

It would be great to have the feature that allows pre-processing of raw tool call arguments before schema validation.

[DouweM]

@warmwind Do you have a suggestion for what the API should be?

[warmwind]

I think it could work well to add a new property to Agent, something like toolcall_preprocessor.
This function would receive the tool_name and the raw LLM response. For each tool call, the hook would run first, and its return value would be used as the tool call parameters. This way, users can modify or fix the response as needed based on the tool name.

I did consider using event_stream_handler, but I don’t think it’s the right place for this, since it should remain a read-only interface.

It would also be very important for this hook to support Streaming Structured Output, so users can preprocess tool call arguments in real time, especially for the final_result tool.

[DouweM]

@warmwind I think this feature is a little bit too niche and specific to warrant a new property on Agent.

I'd prefer a generic method of modifying model responses before they're handled by the framework. You could already create a custom WrapperModel that overrides request (and request_stream if you want to support streaming as well) and modifies the ToolCallParts as it sees fit. Have you tried that?

You can also modify tool execution by defining a custom WrapperToolset and overriding call_tool (see docs), but that currently happens after arguments have already been validated against the schema, so too late for your use case. We could add a new method like validate_tool_call that you could override.

You can already effectively accomplish that overriding the WrapperToolset.get_tools method, and replacing the args_validator field of each tool returned by super().get_tools() with a custom object matching this protocol:

pydantic-ai/pydantic_ai_slim/pydantic_ai/toolsets/abstract.py

Lines 22 to 35 in 3eaa11e

	class SchemaValidatorProt(Protocol):
	"""Protocol for a Pydantic Core `SchemaValidator` or `PluggableSchemaValidator` (which is private but API-compatible)."""
	def validate_json(
	self,
	input: str | bytes | bytearray,
	*,
	allow_partial: bool | Literal['off', 'on', 'trailing-strings'] = False,
	**kwargs: Any,
	) -> Any: ...
	def validate_python(
	self, input: Any, *, allow_partial: bool | Literal['off', 'on', 'trailing-strings'] = False, **kwargs: Any
	) -> Any: ...

validate_json will be called when the args are a (possibly incomplete) JSON string, and validate_python when args are already a dict.

Could you give either the WrapperModel or WrapperToolset approach a try?

[warmwind]

The reason I suggested adding it to the Agent is that an agent can have multiple tools, and it always has at least one final_result tool. Having the preprocessor defined at the agent level would provide a central place to handle preprocessing logic for all tools consistently.

I haven’t used WrapperModel or WrapperToolset before, but I’ll give them a try and let you know how it goes.

[warmwind]

@DouweM I started trying the WrapperModel approach and overriding request_stream, but I’m not sure it’s the right way to go.

I’m using OpenAIChatModel, so I need to create a wrapper for it. However, its request_stream implementation is quite complex — it seems I’d also need to override the handle_tool_call_delta call inside OpenAIStreamedResponse:

for dtc in choice.delta.tool_calls or []:
    maybe_event = self._parts_manager.handle_tool_call_delta(
        vendor_part_id=dtc.index,
        tool_name=dtc.function and dtc.function.name,
        args=dtc.function and dtc.function.arguments,
        tool_call_id=dtc.id,
    )
    if maybe_event is not None:
        yield maybe_event

Could you please confirm if this is indeed the correct approach?
If so, I’m a bit concerned about stability — since if the OpenAI code changes in the future, this override would likely break and need to be updated.

[DouweM]

@warmwind Hmm, you're right, to replace the events, you'd need to subclass OpenAIStreamedResponse, override __aiter__ and use _parts_manager, which is private. You could use handle_tool_call_part to replace an incomplete part with a new one, but that'd all be quite brittle.

If you just want to modify the final ModelResponse though, you could just override OpenAIStreamedResponse.get (and OpenAIChatModel.request), which is public. What would you expect the fixed-up events to look like anyway? A final tool call args delta after the real final one?