Home
Studio

Enforce structured output with a hook

Coming soon

An agent turn does not always end with a response that matches your response_format=. One turn can contain many model calls, and the model output is not deterministic.

This is a problem in programmatic usage, for example, a pipeline parses the response into a Pydantic model and stores the result. A response with a different shape stops the pipeline.

To prevent this, use pattern.hook.StructuredOutputHook. The hook validates the final response against your model. If the response is not valid, the hook sends the error to the agent, and the agent corrects its response.

The hook

The hook

import json
from typing import ClassVar, override

from pydantic import BaseModel, ValidationError

from mistralai.vibe.harness.core import protocol
from mistralai.vibe.harness.runtime.capabilities import (
    PostLlmCallHookAccept,
    PostLlmCallHookInvocation,
    PostLlmCallHookRetry,
    PreAgentTurnHookContinue,
    PreAgentTurnHookInvocation,
)
from mistralai.vibe_agents.sdk.authoring.contracts import Hook, HookContext

__all__ = ["StructuredOutputHook"]

_ERROR = "StructuredOutputHook subclasses must define output_model as a Pydantic BaseModel class"


class StructuredOutputHook(Hook):
    """Require every final model response to match a Pydantic model."""

    output_model: ClassVar[type[BaseModel]]
    _schema: ClassVar[str]

    def __init_subclass__(cls, **kwargs: object) -> None:
        super().__init_subclass__(**kwargs)
        output_model = cls.__dict__.get("output_model")
        if not isinstance(output_model, type) or not issubclass(output_model, BaseModel):
            raise TypeError(_ERROR)
        cls._schema = json.dumps(output_model.model_json_schema(), indent=2)

    def __init__(self) -> None:
        if type(self) is StructuredOutputHook:
            raise TypeError(_ERROR)
        super().__init__()

    @override
    async def pre_agent_turn(
        self,
        ctx: HookContext,
        request: PreAgentTurnHookInvocation,
    ) -> PreAgentTurnHookContinue:
        del ctx
        instruction = (
            "Return the final answer as one JSON object that matches the JSON Schema below. "
            "Do not wrap the JSON in Markdown.\n\n"
            f"JSON Schema:\n{type(self)._schema}"
        )
        return PreAgentTurnHookContinue(
            user_content=[*request.user_content, protocol.TextContentBlock(text=instruction)]
        )

    @override
    async def post_llm_call(
        self,
        ctx: HookContext,
        invocation: PostLlmCallHookInvocation,
    ) -> PostLlmCallHookAccept | PostLlmCallHookRetry:
        del ctx
        parts = invocation.candidate.message.content
        if any(isinstance(part, protocol.ModelToolCallPart) for part in parts):
            return PostLlmCallHookAccept()

        text = "".join(part.text for part in parts if isinstance(part, protocol.TextContentBlock))
        try:
            type(self).output_model.model_validate_json(text)
        except ValidationError as error:
            feedback = (
                "Your previous answer did not match the required JSON Schema. "
                "Return only a corrected JSON object.\n\n"
                f"Validation error:\n{error}\n\n"
                f"JSON Schema:\n{type(self)._schema}\n\n"
                f"Previous answer:\n{text}"
            )
            return PostLlmCallHookRetry(feedback=[protocol.TextContentBlock(text=feedback)])
        return PostLlmCallHookAccept()

pre_agent_turn adds the JSON Schema to the user message. post_llm_call runs after each model call. It accepts a tool call, because the turn is not finished. It validates a final response. If the validation fails, it returns PostLlmCallHookRetry with the error, the schema and the previous response. Each retry uses one iteration of max_iterations.

Usage

Usage

Subclass the hook and set output_model. Put the subclass last in hooks=[...].

import asyncio

from pydantic import BaseModel

from mistralai.agents import agents, pattern


class City(BaseModel):
    city: str
    country: str


class CityOutputHook(pattern.hook.StructuredOutputHook):
    output_model = City


class LandmarkArgs(BaseModel):
    name: str


@agents.tool(name="lookup_landmark", description="Find where a landmark is.", input_schema=LandmarkArgs)
async def lookup_landmark(args: LandmarkArgs) -> dict[str, str]:
    return {"landmark": args.name, "location": "Champ de Mars, 7th arrondissement, Paris, France"}


async def main() -> None:
    agent = agents.Agent(
        model="mistral-medium-3-5",
        name="structured-output-hook-example",
        instructions="Use lookup_landmark to find where a landmark is.",
        max_iterations=6,
        harness=agents.Harness(tools=[lookup_landmark], hooks=[CityOutputHook()]),
    )
    result = await agent.run("Which city and country contain the Eiffel Tower?", response_format=City)
    if not isinstance(result.output, City):
        raise RuntimeError("The run did not produce a parsed City output")
    print(result.output.model_dump())


if __name__ == "__main__":
    asyncio.run(main())
See also

See also