The Enterprise AI Masterclass

from pydantic import BaseModel, Field
from pydantic_ai import Agent

# 1. The Contract (Shared Model)
# This defines EXACTLY what the AI allows to output.
class SupportResult(BaseModel):
    summary: str = Field(description="A polite summary of the issue")
    risk_score: int = Field(description="Risk score 1-10 based on sentiment")
    escalate: bool = Field(description="True if human intervention is needed")

# 2. The Agent
agent = Agent(
    'openai:gpt-4o',
    # We enforce the return type here. 
    # If the LLM returns plain text, PydanticAI rejects it.
    result_type=SupportResult,  
    system_prompt='You are a helpful support assistant.'
)

# 3. Usage Example
# result = agent.run_sync('My account was hacked and I lost money!')
# print(result.data.risk_score) # -> 9 (Integer, not string)

from dataclasses import dataclass
from pydantic_ai import RunContext

@dataclass
class SupportDeps:
    user_id: int
    db_conn: str

# Tell the agent what dependencies to expect
agent = Agent('openai:gpt-4o', deps_type=SupportDeps)

@agent.tool
async def check_balance(ctx: RunContext[SupportDeps]) -> str:
    # ---------------------------------------------------------
    # SECURE ZONE: The LLM cannot control what happens here.
    # The 'ctx.deps' are injected by our Python code, not the Prompt.
    # ---------------------------------------------------------
    print(f"Checking DB: {ctx.deps.db_conn}")
    print(f"For Verified User: {ctx.deps.user_id}")
    return "$1,250.50"

from fastapi import FastAPI, Depends, HTTPException
from fastapi.security import HTTPBearer
from dataclasses import dataclass

app = FastAPI()
security = HTTPBearer()

@dataclass
class UserContext:
    user_id: int
    role: str

# 1. The Guard
async def get_current_user(creds=Depends(security)) -> UserContext:
    token = creds.credentials
    # In reality: Decode & Verify JWT here
    if token != "valid-secret-token":
        raise HTTPException(status_code=401, detail="Invalid Token")
    return UserContext(user_id=101, role="admin")

# 2. The Secured Endpoint
@app.post("/secure-chat")
async def secure_chat(
    query: str, 
    # FastAPI authenticates user BEFORE the function runs
    user: UserContext = Depends(get_current_user) 
):
    # We construct the deps using the VERIFIED user identity
    run_deps = SupportDeps(user_id=user.user_id, db_conn="pg://prod-db")
    
    # Run the agent. It now has the user's ID locked in its context.
    result = await agent.run(query, deps=run_deps)
    return result.data

# 1. The Pause (Agent returns Request object, not Result)
class ApprovalRequest(BaseModel):
    reason: str
    status: str = "WAITING"

# 2. Endpoint: Start Job
@app.post("/request-refund")
async def start_job(query: str):
    result = await agent.run(query)
    if isinstance(result.data, ApprovalRequest):
        # FREEZE STATE: Save entire conversation history to DB
        db.save_job(job_id="job_123", messages=result.all_messages())
        return {"status": "PAUSED", "id": "job_123"}
    return result.data

# 3. Endpoint: Resume Job (Manager Only)
@app.post("/manager-approve")
async def approve(job_id: str):
    # THAW STATE: Load history
    history = db.get_job(job_id)
    
    # WAKE UP AGENT: "Inject" the approval as a new user message
    final = await agent.run(
        "Manager approved the request. Proceed with transfer.", 
        message_history=history  # <--- Agent restores full context
    )
    return final.data