File size: 8,322 Bytes
f17c0fd c55ff1a f17c0fd c55ff1a f17c0fd c55ff1a f17c0fd c55ff1a f17c0fd c55ff1a f17c0fd c55ff1a f17c0fd c55ff1a f17c0fd c55ff1a f17c0fd c55ff1a f17c0fd c55ff1a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 | """ACO Optimizer: Main orchestrator that coordinates all modules."""
import json, time, uuid
from typing import Dict, List, Optional, Any
from .config import ACOConfig, RoutingPolicy
from .trace_schema import AgentTrace, TraceStep, ModelCall, ToolCall
from .classifier import TaskCostClassifier
from .router import ModelCascadeRouter, RoutingDecision
from .context_budgeter import ContextBudgeter, ContextBudget
from .cache_layout import CacheAwareLayout, PromptLayout
from .tool_gate import ToolCostGate, ToolDecision
from .verifier_budgeter import VerifierBudgeter, VerifierDecision
from .retry_optimizer import RetryOptimizer, RecoveryAction
from .meta_tool_miner import MetaToolMiner, MacroTool
from .doom_detector import DoomDetector, DoomAssessment
class ACOOptimizer:
def __init__(self, config: ACOConfig = None):
self.config = config or ACOConfig()
self.classifier = TaskCostClassifier()
self.router = ModelCascadeRouter(
model_path=self.config.router_model_path,
safety_threshold=self.config.routing_policy.safety_threshold,
downgrade_threshold=self.config.routing_policy.downgrade_threshold,
task_floor=self.config.task_floors,
tier_costs=self.config.tier_costs,
)
self.context_budgeter = ContextBudgeter()
self.cache_layout = CacheAwareLayout()
self.tool_gate = ToolCostGate()
self.verifier_budgeter = VerifierBudgeter()
self.retry_optimizer = RetryOptimizer(
max_retries=self.config.routing_policy.max_retries,
)
self.meta_tool_miner = MetaToolMiner()
self.doom_detector = DoomDetector()
self._current_trace: Optional[AgentTrace] = None
self._step_num = 0
self._traces: List[AgentTrace] = []
def start_run(self, request: str) -> Dict:
prediction = self.classifier.classify(request)
routing = self.router.route(request, prediction["task_type"], prediction["difficulty"], prediction)
context_budget = self.context_budgeter.budget(
prediction["task_type"], prediction["difficulty"],
prediction["needs_retrieval"], prediction["needs_tools"],
)
# Check for meta-tool match
macro = self.meta_tool_miner.match_macro(request, prediction["task_type"]) if self.config.enable_meta_tools else None
self._current_trace = AgentTrace(
request=request,
task_type=prediction["task_type"],
difficulty=prediction["difficulty"],
predicted_tier=routing.tier,
)
self._step_num = 0
self.retry_optimizer.reset_run()
self.verifier_budgeter.reset_run()
return {
"trace_id": self._current_trace.trace_id,
"prediction": prediction,
"routing": {
"model_id": routing.model_id,
"tier": routing.tier,
"confidence": routing.confidence,
"cost_estimate": routing.cost_estimate,
"dynamic_difficulty": routing.dynamic_difficulty,
"escalated": routing.escalated,
"downgraded": routing.downgraded,
"reasoning": routing.reasoning,
},
"context_budget": {
"total_tokens": context_budget.total_tokens,
"keep_exact": context_budget.keep_exact,
"summarize": context_budget.summarize,
"omit": context_budget.omit,
"retrieve_on_demand": context_budget.retrieve_on_demand,
"cache_prefix": context_budget.cache_prefix,
},
"macro_tool": macro.name if macro else None,
}
def record_step(self, model_call: Dict = None, tool_calls: List[Dict] = None,
context_size: int = 0, verifier_called: bool = False,
verifier_result: str = None, retry_num: int = 0,
recovery_action: str = None) -> None:
self._step_num += 1
mc = None
if model_call:
mc = ModelCall(**model_call)
tcs = [ToolCall(**tc) for tc in (tool_calls or [])]
step = TraceStep(
step_num=self._step_num,
model_call=mc,
tool_calls=tcs,
context_size=context_size,
verifier_called=verifier_called,
verifier_result=verifier_result,
retry_num=retry_num,
recovery_action=recovery_action,
)
if self._current_trace:
self._current_trace.steps.append(step)
def check_doom(self, current_cost: float = 0.0) -> DoomAssessment:
if not self._current_trace:
return DoomAssessment(False, 0.0, [], "continue", "no active trace")
return self.doom_detector.assess(
[s.__dict__ for s in self._current_trace.steps],
current_cost, self.config.routing_policy.max_cost_per_task, 4)
def should_verify(self, is_irreversible: bool = False,
has_prior_failures: bool = False) -> VerifierDecision:
if not self._current_trace:
return VerifierDecision(False, "skip", 0.0, "no active trace", 0.0)
return self.verifier_budgeter.should_verify(
self._current_trace.task_type, "medium", 0.8,
is_irreversible, has_prior_failures,
self._current_trace.predicted_tier)
def gate_tool(self, tool_name: str, args: Dict) -> ToolDecision:
if not self._current_trace:
return ToolDecision("skip", tool_name, 0.0, "no active trace", 0.0, 0.0)
return self.tool_gate.gate(tool_name, args, self._current_trace.task_type,
self._step_num, self._step_num + 1, 0.5)
def get_recovery(self, failure_tag: str, current_tier: int,
retry_num: int, previous_actions: List[str] = None,
run_cost: float = 0.0) -> RecoveryAction:
return self.retry_optimizer.get_recovery(
failure_tag, current_tier, retry_num,
previous_actions, run_cost,
self.config.routing_policy.max_cost_per_task)
def end_run(self, success: bool, outcome: str = "completed",
artifacts: List[str] = None, failure_tags: List[str] = None,
user_correction: bool = False) -> AgentTrace:
if self._current_trace:
self._current_trace.task_success = success
self._current_trace.final_outcome = outcome
self._current_trace.artifacts_created = artifacts or []
self._current_trace.failure_tags = failure_tags or []
self._current_trace.user_correction = user_correction
summary = self._current_trace.compute_summary()
self._current_trace.total_cost = summary["total_cost"]
self._current_trace.total_tokens = summary["total_tokens"]
self._current_trace.total_tool_calls = summary["total_tool_calls"]
self._current_trace.total_retries = summary["total_retries"]
self._current_trace.total_verifier_calls = summary["total_verifier_calls"]
self._current_trace.cache_hit_rate = summary["cache_hit_rate"]
self._traces.append(self._current_trace)
trace = self._current_trace
self._current_trace = None
return trace
def layout_prompt(self, sources: Dict[str, str]) -> PromptLayout:
if not self._current_trace:
budget = self.context_budgeter.budget("unknown_ambiguous", 3, False, False)
else:
budget = self.context_budgeter.budget(
self._current_trace.task_type,
self._current_trace.difficulty,
False, False)
return self.cache_layout.layout(sources, budget)
def get_stats(self) -> Dict:
return {
"total_runs": len(self._traces),
"successful_runs": sum(1 for t in self._traces if t.task_success),
"avg_cost": sum(t.total_cost for t in self._traces) / max(len(self._traces),1),
"cache_stats": self.cache_layout.stats(),
"tool_stats": self.tool_gate.call_stats,
"verifier_stats": self.verifier_budgeter.stats,
"retry_stats": self.retry_optimizer.recovery_stats,
}
|