Structured Output Unions

Letting the agent choose its own response shape

Lesson 3 gave an agent exactly one possible output shape. Real agents often need to pick between several: "answer the question" or "ask a clarifying question," "return search results" or "return an error explanation." output_type accepts a union of models for exactly this.

from dotenv import load_dotenv
from pydantic import BaseModel
from pydantic_ai import Agent

load_dotenv()


class WeatherAnswer(BaseModel):
    kind: str = "weather"
    forecast: str


class JokeAnswer(BaseModel):
    kind: str = "joke"
    text: str


agent = Agent(
    "google:gemini-3.5-flash-lite",
    output_type=WeatherAnswer | JokeAnswer,
)


def handle(prompt: str) -> None:
    result = agent.run_sync(prompt)
    match result.output:
        case JokeAnswer(text=text):
            print(f"[{prompt!r}] -> JokeAnswer: {text}")
        case WeatherAnswer(forecast=forecast):
            print(f"[{prompt!r}] -> WeatherAnswer: {forecast}")


def main() -> None:
    handle("Tell me a joke.")
    handle("Will it rain tomorrow in Paris?")

Pydantic AI turns the union into two candidate schemas the model can choose between, and result.output comes back as a real instance of whichever one the model picked, not a dict you have to inspect and branch on yourself.

Why this beats a single "kind" field on one model

You could instead define one BaseModel with an optional field for every possible answer type, but then every consumer has to remember which fields are actually populated for a given kind. A union makes "this is a JokeAnswer, so it has .text and nothing else" a fact your type checker verifies, not a convention you hope holds.