Upload aco/cli.py with huggingface_hub

aco/cli.py

@@ -1,123 +1,113 @@
-"""CLI for Agent Cost Optimizer."""
-
-import argparse
-import json
-import sys
-from pathlib import Path
-
-from aco.optimizer import AgentCostOptimizer
+"""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.benchmarks.benchmark_suite import BenchmarkSuite
-
+from aco.trace_schema import AgentTrace
 
 def main():
-    parser = argparse.ArgumentParser(description="Agent Cost Optimizer")
-    subparsers = parser.add_subparsers(dest="command", help="Command to run")
-    
-    # Optimize command
-    opt_parser = subparsers.add_parser("optimize", help="Optimize an agent request")
-    opt_parser.add_argument("--config", "-c", default="config.yaml", help="Config file path")
-    opt_parser.add_argument("--request", "-r", required=True, help="User request text")
-    opt_parser.add_argument("--output", "-o", default="-", help="Output file (default: stdout)")
-    
-    # Benchmark command
-    bench_parser = subparsers.add_parser("benchmark", help="Run benchmark suite")
-    bench_parser.add_argument("--config", "-c", default="config.yaml", help="Config file path")
-    bench_parser.add_argument("--tasks", "-n", type=int, default=1000, help="Number of tasks")
-    bench_parser.add_argument("--output", "-o", default="benchmark_results.json", help="Output path")
-    bench_parser.add_argument("--ablations", action="store_true", help="Run ablation study")
-    
-    # Report command
-    report_parser = subparsers.add_parser("report", help="Generate report from benchmark results")
-    report_parser.add_argument("--input", "-i", required=True, help="Benchmark results JSON")
-    report_parser.add_argument("--output", "-o", default="-", help="Output file")
-    
+    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 == "optimize":
-        _cmd_optimize(args)
-    elif args.command == "benchmark":
-        _cmd_benchmark(args)
-    elif args.command == "report":
-        _cmd_report(args)
-    else:
+    if args.command is None:
         parser.print_help()
-        sys.exit(1)
+        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))
 
-def _cmd_optimize(args):
-    config = ACOConfig.from_yaml(args.config) if Path(args.config).exists() else ACOConfig()
-    optimizer = AgentCostOptimizer(config)
-    result = optimizer.optimize(args.request)
-    
-    output = {
-        "trace_id": result.trace_id,
-        "model": result.routing_decision.model_id,
-        "tier": result.routing_decision.tier,
-        "estimated_cost": result.estimated_cost,
-        "estimated_latency_ms": result.estimated_latency_ms,
-        "confidence": result.confidence,
-        "reasoning": result.reasoning,
-        "tool_decisions": [
-            {"tool": d.tool_name, "decision": d.decision.value, "cost": d.estimated_cost}
-            for d in result.tool_decisions
-        ],
-        "verifier": result.verifier_decision.decision.value if result.verifier_decision else None,
-        "doom_score": result.doom_assessment.confidence if result.doom_assessment else None,
-        "meta_tool_match": result.meta_tool_match is not None,
-    }
-    
-    json_str = json.dumps(output, indent=2)
-    if args.output == "-":
-        print(json_str)
-    else:
-        with open(args.output, "w") as f:
-            f.write(json_str)
+    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))
 
-def _cmd_benchmark(args):
-    config = ACOConfig.from_yaml(args.config) if Path(args.config).exists() else ACOConfig()
-    suite = BenchmarkSuite(config)
-    
-    print(f"Generating {args.tasks} synthetic traces...")
-    traces = suite.generate_benchmark_data(args.tasks)
-    
-    print("Running baselines...")
-    results = suite.run_all_baselines(traces)
-    
-    if args.ablations:
-        print("Running ablations...")
-        ablation_results = suite.run_ablations(traces)
-        results.update(ablation_results)
-    
-    suite.export(results, args.output)
-    
-    # Print report
-    report = suite.report(results)
-    print(report)
-    
-    print(f"\nResults saved to {args.output}")
+    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))
 
-def _cmd_report(args):
-    with open(args.input, "r") as f:
-        data = json.load(f)
-    
-    # Reconstruct BenchmarkResults for reporting
-    from aco.benchmarks.benchmark_suite import BenchmarkResult
-    results = {}
-    for name, d in data.items():
-        results[name] = BenchmarkResult(**d)
-    
-    suite = BenchmarkSuite()
-    report = suite.report(results)
-    
-    if args.output == "-":
-        print(report)
-    else:
-        with open(args.output, "w") as f:
-            f.write(report)
+    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()