06_agent_workflow.py

#!/usr/bin/env python3
"""06 — End-to-end agent workflow with routed steps + fake tools.

Shows the harness pattern: plan → tools → verify, with per-step routing and
a structured run report (model, reason, cost).
"""

from __future__ import annotations

import json
import sys
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))

from lib.metrics import summarize_run
from lib.models import Step, StepType
from lib.router import StageRouter


def fake_tool(name: str, **kwargs) -> dict:
    time.sleep(0.02)
    if name == "billing.get_charges":
        return {"charges": [{"id": "c1", "amount": 49.0, "dup": True}, {"id": "c2", "amount": 49.0, "dup": True}]}
    if name == "policy.check_refund":
        return {"allowed": True, "max_days": 30, "reason": "duplicate_charge"}
    if name == "notify.cs":
        return {"ok": True, "ticket": "T-9912"}
    return {"ok": True, "echo": kwargs}


def run_agent() -> dict:
    router = StageRouter(efficient="lightning", capable="strong")
    transcript: list[dict] = []
    steps: list[Step] = []
    decisions = []

    plan = Step(
        StepType.PLAN,
        "Decompose refund for order 88421",
        tokens_in=1100,
        tokens_out=350,
    )
    d = router.route(plan)
    steps.append(plan)
    decisions.append(d)
    transcript.append({"role": "assistant", "model": d.model_id, "text": "1) fetch charges 2) policy 3) refund 4) notify"})

    # Execute: tools on efficient tier
    for content, tool, args, tin, tout in [
        ("Fetch charges", "billing.get_charges", {"order": "88421"}, 700, 120),
        ("Check policy", "policy.check_refund", {"order": "88421"}, 900, 100),
        ("Notify CS", "notify.cs", {"order": "88421"}, 400, 60),
    ]:
        signals = ["tool_result"]
        if tool == "policy.check_refund":
            signals = ["tool_result", "policy_check"]
        step = Step(StepType.EXECUTE if "policy" not in tool else StepType.VERIFY, content, signals=signals, tokens_in=tin, tokens_out=tout, risk="high" if "policy" in tool else "normal")
        d = router.route(step)
        result = fake_tool(tool, **args)
        steps.append(step)
        decisions.append(d)
        transcript.append({"role": "tool", "name": tool, "result": result, "model": d.model_id})

    report = summarize_run(steps, decisions)
    return {"transcript": transcript, "report": report, "decisions": report["decisions"]}


def main() -> int:
    out = run_agent()
    print("=== Transcript (abbrev) ===")
    for turn in out["transcript"]:
        if turn["role"] == "assistant":
            print(f"  [{turn['model']}] plan: {turn['text']}")
        else:
            print(f"  [{turn['model']}] tool {turn['name']} → {json.dumps(turn['result'])[:80]}")

    r = out["report"]
    print("\n=== Run report ===")
    print(json.dumps({k: r[k] for k in ("steps", "routed_cost_usd", "frontier_only_usd", "savings_pct", "tier_counts")}, indent=2))
    print("\nLog model_id + reason per step in production — that is your tokenomics dashboard.")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())