Skip to main content

Pydantic AI Integration

Enforce Control Zero policies on Pydantic AI agents, tool calls, and structured outputs.

First-party integration module available

As of Python SDK v1.4.0, Control Zero ships a first-party Pydantic AI integration at controlzero.integrations.pydantic_ai. It exposes governed_tool (a decorator that wraps @agent.tool functions with a guard() call) and GovernedAgent (a drop-in subclass that adds policy enforcement around model selection and tool calls). The manual examples below still work — use whichever style fits your codebase.

from controlzero.integrations.pydantic_ai import governed_tool, GovernedAgent

Overview​

Pydantic AI builds type-safe AI agents with structured inputs and outputs. Control Zero integrates at the tool level: every tool call your Pydantic AI agent makes is checked against your policies before execution. The integration is lightweight since Pydantic AI tools are plain Python functions that you can wrap with a single enforce() call.

Installation​

pip install controlzero pydantic-ai

Setup​

Tool-Level Enforcement​

Wrap your Pydantic AI tools with Control Zero policy checks:

from controlzero import Client
from pydantic_ai import Agent

cz = Client(api_key="cz_live_your_api_key_here")

agent = Agent(
    "openai:gpt-5.4",
    system_prompt="You are a helpful research assistant.",
)


@agent.tool_plain
def search_web(query: str) -> str:
    """Search the web for information."""
    cz.guard("search_web", args={"agent_id": "research-agent"})
    return f"Results for: {query}"


@agent.tool_plain
def read_file(path: str) -> str:
    """Read a file from the local filesystem."""
    cz.guard("read_file", args={"agent_id": "research-agent", "path": path})
    with open(path) as f:
        return f.read()

Dependency Injection Approach​

Pydantic AI's dependency injection system works well with Control Zero. Pass the client as a dependency:

from dataclasses import dataclass
from controlzero import Client
from pydantic_ai import Agent, RunContext

cz = Client(api_key="cz_live_your_api_key_here")


@dataclass
class AgentDeps:
    cz_client: Client
    agent_id: str


agent = Agent(
    "openai:gpt-5.4",
    deps_type=AgentDeps,
    system_prompt="You are a data analyst.",
)


@agent.tool
def query_database(ctx: RunContext[AgentDeps], sql: str) -> str:
    """Run a read-only SQL query."""
    ctx.deps.cz_client.guard("query_database", args={"agent_id": ctx.deps.agent_id, "sql": sql})
    # ... execute query ...
    return f"Query results for: {sql}"


@agent.tool
def write_database(ctx: RunContext[AgentDeps], table: str, data: str) -> str:
    """Insert data into a database table."""
    ctx.deps.cz_client.guard("write_database", args={"agent_id": ctx.deps.agent_id, "table": table})
    return f"Inserted into {table}"


# Run with dependencies
result = agent.run_sync(
    "What were last quarter's top products?",
    deps=AgentDeps(cz_client=cz, agent_id="analyst-agent"),
)

What Gets Enforced​

Pydantic AI EventPolicy ActionPolicy Resource
Tool call (@agent.tool)tool.calltool/{tool_name}
Plain tool (@agent.tool_plain)tool.calltool/{tool_name}
Model selectionllm.generatemodel/{model_name}

Full Example: Customer Data Agent​

from dataclasses import dataclass
from controlzero import Client
import controlzero
from pydantic_ai import Agent, RunContext
from pydantic import BaseModel

# --- Control Zero setup ---
cz = Client(api_key="cz_live_your_api_key_here")


# --- Structured output ---
class CustomerReport(BaseModel):
    customer_name: str
    total_orders: int
    lifetime_value: float
    risk_score: str
    summary: str


# --- Dependencies ---
@dataclass
class Deps:
    cz: Client
    agent_id: str = "customer-data-agent"


# --- Agent ---
agent = Agent(
    "openai:gpt-5.4",
    deps_type=Deps,
    output_type=CustomerReport,
    system_prompt=(
        "You are a customer data analyst. Use the available tools to gather "
        "customer information and produce a structured report."
    ),
)


@agent.tool
def lookup_customer(ctx: RunContext[Deps], customer_id: str) -> str:
    """Look up basic customer information."""
    ctx.deps.cz.guard("lookup_customer", args={"agent_id": ctx.deps.agent_id})
    return f"Customer {customer_id}: Acme Corp, Premium tier"


@agent.tool
def get_order_history(ctx: RunContext[Deps], customer_id: str) -> str:
    """Get the order history for a customer."""
    ctx.deps.cz.guard("get_order_history", args={"agent_id": ctx.deps.agent_id})
    return f"Customer {customer_id}: 47 orders, $125,000 lifetime value"


@agent.tool
def get_support_tickets(ctx: RunContext[Deps], customer_id: str) -> str:
    """Get support ticket history for a customer."""
    ctx.deps.cz.guard("get_support_tickets", args={"agent_id": ctx.deps.agent_id})
    return f"Customer {customer_id}: 3 open tickets, 12 resolved"


@agent.tool
def export_customer_data(ctx: RunContext[Deps], customer_id: str) -> str:
    """Export all customer data to an external file."""
    ctx.deps.cz.guard("export_customer_data", args={"agent_id": ctx.deps.agent_id})
    return f"Exported data for {customer_id}"


# --- Run ---
try:
    result = agent.run_sync(
        "Generate a report for customer CUST-789",
        deps=Deps(cz=cz),
    )
    print(result.output.model_dump_json(indent=2))
except controlzero.PolicyDeniedError as e:
    print(f"Blocked by policy: {e.decision.reason}")

Example Policy​

Allow read-only lookups but deny data exports:

{
  "name": "customer-data-policy",
  "rules": [
    {
      "effect": "allow",
      "action": "llm:generate",
      "resource": "model/gpt-4"
    },
    {
      "effect": "allow",
      "action": "tool:call",
      "resource": "tool/lookup_customer"
    },
    {
      "effect": "allow",
      "action": "tool:call",
      "resource": "tool/get_order_history"
    },
    {
      "effect": "allow",
      "action": "tool:call",
      "resource": "tool/get_support_tickets"
    },
    {
      "effect": "deny",
      "action": "tool:call",
      "resource": "tool/export_customer_data"
    }
  ]
}

What happens at runtime:

  • Customer lookup, order history, and support tickets: ALLOWED.
  • Data export: DENIED. The agent cannot exfiltrate customer data.
  • The structured CustomerReport output is still produced using only the allowed data sources.

Pydantic AI-Specific Patterns​

Structured Output Validation​

Pydantic AI enforces output types at the framework level. Control Zero complements this by enforcing which tools the agent can use to gather the data that populates the structured output. Together they form two layers of safety:

  1. Control Zero: Controls what data the agent can access (tool-level governance).
  2. Pydantic AI: Controls what shape the output takes (type-level validation).

Model Selection Enforcement​

Control which models your Pydantic AI agents can use via the gateway:

# Route through the Control Zero gateway to enforce model policies
agent = Agent(
    "openai:gpt-5.4",
    # The gateway will block if gpt-5.4 is not in the allow list
)

Or enforce model selection in code:

cz.guard("llm", method="generate", args={"model": "gpt-5.4", "agent_id": "analyst-agent"})

Retry with Policy Awareness​

Pydantic AI supports retries on validation failures. If a tool call is denied by policy, the agent can retry with a different tool:

agent = Agent(
    "openai:gpt-5.4",
    deps_type=Deps,
    output_type=CustomerReport,
    retries=3,  # agent retries if a tool call fails (including policy denials)
)

The agent sees the PolicyDeniedError as a tool error and adjusts its approach on the next attempt.

Gateway Alternative​

For zero-code integration, point your Pydantic AI agent's model to the Control Zero gateway:

export OPENAI_BASE_URL=https://gateway.controlzero.ai/v1

The gateway intercepts all LLM calls and tool call responses. See the Gateway Guide.

Next Steps​