Upload aco/retry_optimizer.py

aco/retry_optimizer.py

@@ -0,0 +1,224 @@
+"""Retry and Recovery Optimizer - Module 8.
+
+Avoids blind retry loops. For failures, decides:
+- retry same approach
+- retry with changed prompt
+- repair tool call
+- retrieve more context
+- switch model
+- ask clarification
+- call verifier
+- mark BLOCKED
+- terminate
+
+Uses trace-based recovery policies.
+"""
+
+from typing import Dict, List, Optional, Any
+from dataclasses import dataclass
+from enum import Enum
+
+from .trace_schema import Outcome, FailureTag, TraceStep, TaskType
+from .config import ACOConfig
+
+
+class RecoveryAction(Enum):
+    RETRY_SAME = "retry_same"
+    RETRY_CHANGED_PROMPT = "retry_changed_prompt"
+    REPAIR_TOOL = "repair_tool"
+    RETRIEVE_MORE_CONTEXT = "retrieve_more_context"
+    SWITCH_MODEL = "switch_model"
+    ASK_CLARIFICATION = "ask_clarification"
+    CALL_VERIFIER = "call_verifier"
+    MARK_BLOCKED = "mark_blocked"
+    TERMINATE = "terminate"
+    SKIP_AND_CONTINUE = "skip_and_continue"
+
+
+@dataclass
+class RecoveryDecision:
+    action: RecoveryAction
+    reasoning: str
+    confidence: float
+    new_model_tier: Optional[int] = None
+    context_additions: Optional[List[str]] = None
+    prompt_changes: Optional[Dict[str, str]] = None
+
+
+class RetryRecoveryOptimizer:
+    """Intelligently decides how to recover from failures."""
+
+    # Max retries per recovery type
+    MAX_RETRY_SAME = 1
+    MAX_RETRY_CHANGED = 2
+    MAX_REPAIR_TOOL = 2
+    MAX_RETRIEVE_CONTEXT = 1
+    MAX_SWITCH_MODEL = 2
+
+    # Failure pattern -> preferred recovery action
+    FAILURE_RECOVERY_MAP = {
+        FailureTag.MODEL_TOO_WEAK: RecoveryAction.SWITCH_MODEL,
+        FailureTag.CONTEXT_TOO_SMALL: RecoveryAction.RETRIEVE_MORE_CONTEXT,
+        FailureTag.TOOL_FAILED: RecoveryAction.REPAIR_TOOL,
+        FailureTag.TOOL_UNNECESSARY: RecoveryAction.SKIP_AND_CONTINUE,
+        FailureTag.TOOL_MISSED: RecoveryAction.RETRY_CHANGED_PROMPT,
+        FailureTag.RETRY_LOOP: RecoveryAction.MARK_BLOCKED,
+        FailureTag.CACHE_BREAK: RecoveryAction.RETRY_SAME,
+        FailureTag.HALLUCINATION: RecoveryAction.CALL_VERIFIER,
+        FailureTag.TIMEOUT: RecoveryAction.SWITCH_MODEL,
+        FailureTag.COST_EXCEEDED: RecoveryAction.TERMINATE,
+        FailureTag.UNSAFE_CHEAP_MODEL: RecoveryAction.SWITCH_MODEL,
+        FailureTag.MISSED_ESCALATION: RecoveryAction.SWITCH_MODEL,
+        FailureTag.VERIFIER_FALSE_PASS: RecoveryAction.RETRY_CHANGED_PROMPT,
+        FailureTag.VERIFIER_FALSE_REJECT: RecoveryAction.RETRY_SAME,
+    }
+
+    def __init__(self, config: Optional[ACOConfig] = None):
+        self.config = config or ACOConfig()
+        self.retry_counts: Dict[str, int] = {}  # failure_tag -> count
+        self.recovery_stats: Dict[str, Dict] = {}
+
+    def decide_recovery(
+        self,
+        task_type: TaskType,
+        current_step: TraceStep,
+        failure_tags: List[FailureTag],
+        total_cost_so_far: float,
+        predicted_cost: float,
+        current_tier: int,
+        step_number: int,
+        trace_history: Optional[List[TraceStep]] = None,
+    ) -> RecoveryDecision:
+        """Decide recovery action based on failure analysis."""
+        
+        history = trace_history or []
+        
+        # Count retries in trace
+        recent_retries = sum(1 for s in history[-5:] if s.retry_count > 0)
+        total_retries = sum(s.retry_count for s in history)
+        
+        # Detect retry loops
+        if recent_retries >= 3:
+            return RecoveryDecision(
+                action=RecoveryAction.MARK_BLOCKED,
+                reasoning=f"Retry loop detected: {recent_retries} retries in last 5 steps",
+                confidence=0.9,
+            )
+        
+        # Cost escalation check
+        cost_ratio = total_cost_so_far / max(predicted_cost, 0.001)
+        if cost_ratio > self.config.doom_max_cost_ratio * 1.5:
+            return RecoveryDecision(
+                action=RecoveryAction.TERMINATE,
+                reasoning=f"Cost exceeded {self.config.doom_max_cost_ratio * 1.5}x predicted cost ({total_cost_so_far:.4f} vs {predicted_cost:.4f})",
+                confidence=0.85,
+            )
+        
+        # Analyze primary failure tag
+        primary_failure = failure_tags[0] if failure_tags else FailureTag.MODEL_TOO_WEAK
+        preferred_action = self.FAILURE_RECOVERY_MAP.get(primary_failure, RecoveryAction.RETRY_CHANGED_PROMPT)
+        
+        # Check if we've exhausted this recovery path
+        failure_key = f"{primary_failure.value}_{preferred_action.value}"
+        current_count = self.retry_counts.get(failure_key, 0)
+        
+        max_map = {
+            RecoveryAction.RETRY_SAME: self.MAX_RETRY_SAME,
+            RecoveryAction.RETRY_CHANGED_PROMPT: self.MAX_RETRY_CHANGED,
+            RecoveryAction.REPAIR_TOOL: self.MAX_REPAIR_TOOL,
+            RecoveryAction.RETRIEVE_MORE_CONTEXT: self.MAX_RETRIEVE_CONTEXT,
+            RecoveryAction.SWITCH_MODEL: self.MAX_SWITCH_MODEL,
+        }
+        max_allowed = max_map.get(preferred_action, 1)
+        
+        if current_count >= max_allowed:
+            # Escalate to next recovery action
+            escalation_chain = [
+                RecoveryAction.RETRY_SAME,
+                RecoveryAction.RETRY_CHANGED_PROMPT,
+                RecoveryAction.REPAIR_TOOL,
+                RecoveryAction.RETRIEVE_MORE_CONTEXT,
+                RecoveryAction.SWITCH_MODEL,
+                RecoveryAction.ASK_CLARIFICATION,
+                RecoveryAction.MARK_BLOCKED,
+            ]
+            
+            try:
+                idx = escalation_chain.index(preferred_action)
+                preferred_action = escalation_chain[min(idx + 1, len(escalation_chain) - 1)]
+            except ValueError:
+                preferred_action = RecoveryAction.MARK_BLOCKED
+        
+        self.retry_counts[failure_key] = current_count + 1
+        
+        # Build decision
+        if preferred_action == RecoveryAction.SWITCH_MODEL:
+            new_tier = min(current_tier + 1, 5)
+            return RecoveryDecision(
+                action=preferred_action,
+                reasoning=f"Failure: {primary_failure.value}. Escalating from tier {current_tier} to tier {new_tier}",
+                confidence=0.8,
+                new_model_tier=new_tier,
+            )
+        
+        if preferred_action == RecoveryAction.RETRIEVE_MORE_CONTEXT:
+            return RecoveryDecision(
+                action=preferred_action,
+                reasoning=f"Failure: {primary_failure.value}. Adding retrieved context and retrying.",
+                confidence=0.75,
+                context_additions=["retrieved_docs", "tool_error_logs", "prior_attempt_summary"],
+            )
+        
+        if preferred_action == RecoveryAction.REPAIR_TOOL:
+            return RecoveryDecision(
+                action=preferred_action,
+                reasoning=f"Failure: {primary_failure.value}. Repairing tool call parameters.",
+                confidence=0.7,
+                prompt_changes={"tool_repair": "true", "validate_params": "true"},
+            )
+        
+        if preferred_action == RecoveryAction.RETRY_CHANGED_PROMPT:
+            return RecoveryDecision(
+                action=preferred_action,
+                reasoning=f"Failure: {primary_failure.value}. Retrying with modified prompt strategy.",
+                confidence=0.6,
+                prompt_changes={"add_examples": "true", "increase_temperature": "0.3"},
+            )
+        
+        if preferred_action == RecoveryAction.TERMINATE:
+            return RecoveryDecision(
+                action=preferred_action,
+                reasoning=f"Failure: {primary_failure.value}. Cost ratio {cost_ratio:.1f}x. Terminating.",
+                confidence=0.9,
+            )
+        
+        if preferred_action == RecoveryAction.MARK_BLOCKED:
+            return RecoveryDecision(
+                action=preferred_action,
+                reasoning=f"Failure: {primary_failure.value}. Exhausted recovery options. Marking BLOCKED.",
+                confidence=0.85,
+            )
+        
+        return RecoveryDecision(
+            action=preferred_action,
+            reasoning=f"Failure: {primary_failure.value}. Attempting recovery via {preferred_action.value}.",
+            confidence=0.6,
+        )
+
+    def record_recovery_outcome(
+        self,
+        failure_tag: FailureTag,
+        action: RecoveryAction,
+        succeeded: bool,
+        cost_delta: float,
+    ) -> None:
+        """Record outcome for policy improvement."""
+        key = f"{failure_tag.value}_{action.value}"
+        stats = self.recovery_stats.setdefault(key, {
+            "attempts": 0, "successes": 0, "total_cost_delta": 0.0,
+        })
+        stats["attempts"] += 1
+        if succeeded:
+            stats["successes"] += 1
+        stats["total_cost_delta"] += cost_delta
+        stats["success_rate"] = stats["successes"] / stats["attempts"]