aco/cli.py

"""ACO CLI: Command-line interface for the Agent Cost Optimizer."""
import argparse, json, sys
sys.path.insert(0, "/app")
from aco.optimizer import ACOOptimizer
from aco.config import ACOConfig
from aco.trace_schema import AgentTrace

def main():
    parser = argparse.ArgumentParser(description="ACO: Agent Cost Optimizer")
    sub = parser.add_subparsers(dest="command")
    # Route
    rp = sub.add_parser("route", help="Route a request to the optimal model")
    rp.add_argument("request", help="The request text")
    rp.add_argument("--task-type", default=None, help="Override task type")
    rp.add_argument("--difficulty", type=int, default=None, help="Override difficulty")
    # Budget
    bp = sub.add_parser("budget", help="Compute context budget for a task")
    bp.add_argument("request", help="The request text")
    # Gate
    gp = sub.add_parser("gate", help="Check if a tool call is worth it")
    gp.add_argument("tool", help="Tool name")
    gp.add_argument("--task-type", default="coding")
    # Verify
    vp = sub.add_parser("verify", help="Check if verification is needed")
    vp.add_argument("--risk", default="medium", choices=["low","medium","high","critical"])
    vp.add_argument("--confidence", type=float, default=0.8)
    # Doom
    dp = sub.add_parser("doom", help="Assess if a run is doomed")
    dp.add_argument("--cost", type=float, default=0.0)
    # Stats
    sp = sub.add_parser("stats", help="Show optimizer statistics")
    # Eval
    ep = sub.add_parser("eval", help="Run benchmark evaluation")
    ep.add_argument("--n", type=int, default=2000, help="Number of traces")
    # Version
    sub.add_parser("version", help="Show version")

    args = parser.parse_args()
    if args.command is None:
        parser.print_help()
        return

    opt = ACOOptimizer()

    if args.command == "route":
        prediction = opt.classifier.classify(args.request)
        tt = args.task_type or prediction["task_type"]
        diff = args.difficulty or prediction["difficulty"]
        routing = opt.router.route(args.request, tt, diff, prediction)
        print(json.dumps({
            "request": args.request[:80],
            "task_type": tt,
            "difficulty": diff,
            "model_id": routing.model_id,
            "tier": routing.tier,
            "confidence": round(routing.confidence, 3),
            "cost_estimate": routing.cost_estimate,
            "escalated": routing.escalated,
            "downgraded": routing.downgraded,
            "reasoning": routing.reasoning,
        }, indent=2))

    elif args.command == "budget":
        prediction = opt.classifier.classify(args.request)
        budget = opt.context_budgeter.budget(
            prediction["task_type"], prediction["difficulty"],
            prediction["needs_retrieval"], prediction["needs_tools"])
        print(json.dumps({
            "total_tokens": budget.total_tokens,
            "keep_exact": budget.keep_exact,
            "summarize": budget.summarize,
            "omit": budget.omit,
            "retrieve_on_demand": budget.retrieve_on_demand,
            "cache_prefix": budget.cache_prefix,
        }, indent=2))

    elif args.command == "gate":
        decision = opt.tool_gate.gate(args.tool, {}, args.task_type, 1, 5, 0.5)
        print(json.dumps({
            "tool": args.tool,
            "action": decision.action,
            "confidence": decision.confidence,
            "reasoning": decision.reasoning,
            "estimated_cost": decision.estimated_cost,
        }, indent=2))

    elif args.command == "verify":
        decision = opt.verifier_budgeter.should_verify(
            "coding", args.risk, args.confidence)
        print(json.dumps({
            "should_verify": decision.should_verify,
            "verifier_type": decision.verifier_type,
            "reasoning": decision.reasoning,
            "estimated_cost": decision.estimated_cost,
        }, indent=2))

    elif args.command == "doom":
        assessment = opt.check_doom(args.cost)
        print(json.dumps({
            "doomed": assessment.doomed,
            "severity": assessment.severity,
            "action": assessment.recommended_action,
            "reasoning": assessment.reasoning,
        }, indent=2))

    elif args.command == "stats":
        print(json.dumps(opt.get_stats(), indent=2))

    elif args.command == "version":
        print("ACO v8.0 - Agent Cost Optimizer")

if __name__ == "__main__":
    main()