Pydantic AI Traces

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

from inference_catalyst_tracing import agent_span, 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"]}"}}'
    )

with agent_span(
    tracing.tracer,
    agent_id="weather-agent",
    agent_name="Weather Agent",
    span_name="weather-agent.run",
    session_id="conversation-weather-paris",
    agent_role="weather",
    system="pydantic-ai",
) as span:
    user_input = "What's the weather in Paris and Tokyo?"
    span.set_input(user_input)
    result = agent.run_sync(user_input)
    span.set_output(result.output.model_dump())

print(result.output.summary)
tracing.shutdown()