Pydantic AI Traces

pip install 'inference-catalyst-tracing[pydantic-ai]'

from inference_catalyst_tracing import setup
from pydantic import BaseModel, Field
from pydantic_ai import Agent, RunContext

class CityWeather(BaseModel):
    city: str
    temp_c: float = Field(description="Temperature in Celsius.")
    condition: str

class WeatherReport(BaseModel):
    cities: list[CityWeather]
    summary: str = Field(description="One sentence comparing the conditions.")

tracing = setup(service_name="weather-agent")

agent = Agent(
    "openai:gpt-4o-mini",
    output_type=WeatherReport,
    system_prompt="Use get_weather for every requested city.",
)

@agent.tool
def get_weather(_ctx: RunContext[None], city: str) -> str:
    """Look up current weather for a city."""
    weather = {
        "Paris": {"temp_c": 12, "condition": "overcast"},
        "Tokyo": {"temp_c": 18, "condition": "sunny"},
    }
    record = weather.get(city, {"temp_c": 0, "condition": "unknown"})
    return (
        f'{{"city": "{city}", "temp_c": {record["temp_c"]}, '
        f'"condition": "{record["condition"]}"}}'
    )

result = agent.run_sync("What's the weather in Paris and Tokyo?")
print(result.output.summary)
tracing.shutdown()