PydanticAI for AI Engineering: Type Safety, Dependencies & Logfire

from dataclasses import dataclass
from pydantic import BaseModel, Field
from pydantic_ai import Agent, RunContext

# 1. Define runtime dependency container (e.g. database pool, API tokens)
@dataclass
class DatabaseDependencies:
  db_connection_url: str
  api_key: str

# 2. Define structured result output model
class FinancialAuditReport(BaseModel):
  account_name: str = Field(description="Name of the audited corporate entity")
  total_exposure_usd: float = Field(description="Total calculated financial exposure")
  risk_level: str = Field(description="Assessed risk level: LOW, MEDIUM, HIGH")

# 3. Initialize PydanticAI Agent with typed dependencies and result model
audit_agent = Agent[DatabaseDependencies, FinancialAuditReport](
  'anthropic:claude-3-5-sonnet-20241022',
  deps_type=DatabaseDependencies,
  result_type=FinancialAuditReport,
  system_prompt="You are an enterprise financial auditor. Query tools to fetch ledger data and output a structured audit report."
)

# 4. Define typed tool accessing context dependencies
@audit_agent.tool
def fetch_ledger_balance(ctx: RunContext[DatabaseDependencies], account_id: str) -> str:
  """Fetch real-time ledger balance using the injected DB connection."""
  # Context gives access to ctx.deps.db_connection_url safely
  return f"Account {account_id} Ledger Balance: $1,250,000.00 USD. Status: Active."

if __name__ == "__main__":
  deps = DatabaseDependencies(db_connection_url="postgresql://admin:pass@localhost:5432/finance", api_key="secret")
  result = audit_agent.run_sync("Audit corporate account ACC-9012.", deps=deps)
  
  # Strongly-typed Pydantic result object
  report: FinancialAuditReport = result.data
  print("Audit Report Account:", report.account_name)
  print("Total Exposure:", report.total_exposure_usd)
  print("Risk Level:", report.risk_level)