output_schema.py

import asyncio

from agno.agent import Agent
from agno.models.openai import OpenAIChat
from pydantic import BaseModel, Field
from rich.pretty import pprint


class PersonInfo(BaseModel):
    name: str = Field(..., description="Person's full name")
    age: int = Field(..., description="Person's age")
    occupation: str = Field(..., description="Person's occupation")


class BookInfo(BaseModel):
    title: str = Field(..., description="Book title")
    author: str = Field(..., description="Author name")
    year: int = Field(..., description="Publication year")


async def main():
    agent = Agent(
        model=OpenAIChat(id="gpt-4o-mini"),
        markdown=False,
    )

    # Test 1: arun with output_schema passed at runtime (no agent-level schema)
    print("Test 1: arun with output_schema passed at runtime")
    book_response = await agent.arun(
        "Tell me about '1984' by George Orwell", output_schema=BookInfo, stream=False
    )
    print(f"Response type: {type(book_response.content)}")
    print(f"Is BookInfo: {isinstance(book_response.content, BookInfo)}")
    pprint(book_response.content)

    # Test 2: arun streaming with output_schema passed at runtime
    print("\nTest 2: arun streaming with output_schema passed at runtime")
    final_response = None
    async for event in agent.arun(
        "Tell me about 'Pride and Prejudice'", output_schema=BookInfo, stream=True
    ):
        final_response = event
    print(f"Response type: {type(final_response.content)}")
    print(f"Is BookInfo: {isinstance(final_response.content, BookInfo)}")
    pprint(final_response.content)

    # Test 3: Agent with output_schema, then override with arun
    print("\nTest 3: Agent with output_schema, then override with arun")
    agent_with_schema = Agent(
        model=OpenAIChat(id="gpt-4o-mini"),
        output_schema=PersonInfo,
        markdown=False,
    )

    # First, use default schema
    person_response = await agent_with_schema.arun(
        "Tell me about Albert Einstein", stream=False
    )
    print(f"Default schema response type: {type(person_response.content)}")
    print(f"Is PersonInfo: {isinstance(person_response.content, PersonInfo)}")
    pprint(person_response.content)

    # Now override with BookInfo
    print("\nOverriding with BookInfo:")
    book_response = await agent_with_schema.arun(
        "Tell me about 'The Great Gatsby'", output_schema=BookInfo, stream=False
    )
    print(f"Override response type: {type(book_response.content)}")
    print(f"Is BookInfo: {isinstance(book_response.content, BookInfo)}")
    pprint(book_response.content)


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