p3_workflow_patterns.py

"""Pattern 3: the classic workflow patterns as code: routing, prompt chaining with a gate, and voting."""

import asyncio
import json

from claude_agent_sdk import ClaudeAgentOptions, ResultMessage, query
from demo_util import PROJECT


async def run(prompt, schema=None, **opts):
    """One node: a fresh-context agent. Returns typed JSON when a schema is given, else text."""
    options = ClaudeAgentOptions(cwd=PROJECT, setting_sources=[], max_turns=opts.pop("max_turns", 5),
                                 output_format={"type": "json_schema", "schema": schema} if schema else None, **opts)
    async for m in query(prompt=prompt, options=options):
        if isinstance(m, ResultMessage):
            return (m.structured_output if schema else m.result), (m.total_cost_usd or 0)


# ---------- 1. Routing: a cheap classifier picks the specialist -------------------------------
ROUTE_SCHEMA = {"type": "object", "properties": {"kind": {"type": "string", "enum": ["bug", "docs", "question"]},
                                                 "reason": {"type": "string"}}, "required": ["kind", "reason"]}
SPECIALISTS = {
    "bug":      dict(allowed_tools=["Read", "Grep", "Glob"], max_turns=8, system_prompt="You are a debugger. Locate the defect and describe the fix in two sentences."),
    "docs":     dict(allowed_tools=["Read", "Glob"],         system_prompt="You are a technical writer. Answer in two sentences."),
    "question": dict(tools=[],                               system_prompt="Answer from general knowledge in one sentence."),
}


async def routed(request):
    route, c1 = await run(f"Classify this request: {request!r}", ROUTE_SCHEMA, tools=[], model="claude-haiku-5-5", effort="low")
    answer, c2 = await run(request, **SPECIALISTS[route["kind"]])
    print(f"[route] {request!r}\n   -> kind={route['kind']} ({route['reason']})\n   -> {answer}\n   cost=${c1 + c2:.4f}")


# ---------- 2. Prompt chaining with a programmatic gate -----------------------------------------
OUTLINE_SCHEMA = {"type": "object", "properties": {"sections": {"type": "array", "items": {"type": "string"}, "minItems": 3}},
                  "required": ["sections"]}


async def chained():
    outline, c1 = await run("Outline a README for this project in 3-5 section titles.", OUTLINE_SCHEMA, allowed_tools=["Read", "Glob"])
    assert 3 <= len(outline["sections"]) <= 5, "gate failed"                       # deterministic gate between steps
    doc, c2 = await run("Write one sentence for each of these README sections, as a markdown list:\n"
                        + json.dumps(outline["sections"]), tools=[])
    print(f"[chain] outline={outline['sections']}\n{doc}\n   cost=${c1 + c2:.4f}")


# ---------- 3. Parallelization: voting -----------------------------------------------------------
VOTE_SCHEMA = {"type": "object", "properties": {"safe": {"type": "boolean"}, "reason": {"type": "string"}},
               "required": ["safe", "reason"]}


async def voted():
    prompt = "Is src/app.py safe to deploy to production as-is? Answer safe=true/false with a one-sentence reason."
    results = await asyncio.gather(*(run(prompt, VOTE_SCHEMA, allowed_tools=["Read"], effort="low") for _ in range(3)))
    votes = [r[0]["safe"] for r in results]
    print(f"[vote] votes={votes} -> majority safe={sum(votes) > 1}")
    for (v, c) in results:
        print(f"   - safe={v['safe']}: {v['reason']}")
    print(f"   cost=${sum(c for _, c in results):.4f}")


async def main():
    await routed("calculate_average crashes on an empty list, where is the bug?")
    await routed("What is the difference between a list and a tuple in Python?")
    await chained()
    await voted()


asyncio.run(main())