Upload aco/retry_optimizer.py with huggingface_hub

aco/retry_optimizer.py

@@ -1,224 +1,73 @@
-"""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
+"""Retry and Recovery Optimizer: Maps failure tags to specific recovery actions."""
+from typing import Dict, List, Optional
 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
+class RecoveryAction:
+    action: str  # "retry_same","retry_changed_prompt","repair_tool","retrieve_more",
+                 # "switch_model","ask_clarification","call_verifier","mark_blocked","terminate"
     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."""
+    new_tier: Optional[int] = None
+    additional_context: Optional[str] = None
 
-    # 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_RECOVERY_MAP = {
+    "tool_error": {"primary": "repair_tool", "fallback": "retry_changed_prompt"},
+    "tool_not_found": {"primary": "retry_changed_prompt", "fallback": "ask_clarification"},
+    "timeout": {"primary": "retry_same", "fallback": "switch_model"},
+    "context_too_large": {"primary": "retrieve_more", "fallback": "switch_model"},
+    "model_refused": {"primary": "retry_changed_prompt", "fallback": "switch_model"},
+    "hallucination": {"primary": "call_verifier", "fallback": "retrieve_more"},
+    "wrong_answer": {"primary": "switch_model", "fallback": "call_verifier"},
+    "incomplete": {"primary": "retrieve_more", "fallback": "retry_changed_prompt"},
+    "format_error": {"primary": "retry_changed_prompt", "fallback": "switch_model"},
+    "permission_denied": {"primary": "ask_clarification", "fallback": "mark_blocked"},
+    "rate_limit": {"primary": "retry_same", "fallback": "switch_model"},
+    "unknown": {"primary": "retry_changed_prompt", "fallback": "ask_clarification"},
+}
 
-    # 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,
-    }
+ESCALATION_TIERS = {
+    "switch_model": 1,  # upgrade by 1 tier
+    "call_verifier": 0,
+    "mark_blocked": 0,
+}
 
-    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] = {}
+class RetryOptimizer:
+    def __init__(self, max_retries: int = 3, max_total_retries: int = 5):
+        self.max_retries = max_retries
+        self.max_total_retries = max_total_retries
+        self.retry_counts: Dict[str, int] = {}
+        self.total_retries = 0
+        self.recovery_stats = {}
 
-    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 get_recovery(self, failure_tag: str, current_tier: int,
+                     retry_num: int, previous_actions: List[str] = None,
+                     run_cost_so_far: float = 0, max_run_cost: float = 5.0) -> RecoveryAction:
+        self.retry_counts[failure_tag] = self.retry_counts.get(failure_tag, 0) + 1
+        self.total_retries += 1
+        # Check if we should terminate
+        if self.total_retries >= self.max_total_retries:
+            return RecoveryAction("terminate", "max total retries reached")
+        if retry_num >= self.max_retries:
+            return RecoveryAction("mark_blocked", f"max retries ({self.max_retries}) for this failure")
+        if run_cost_so_far >= max_run_cost * 0.8:
+            return RecoveryAction("terminate", "approaching cost limit")
+        # Get recovery action
+        recovery = FAILURE_RECOVERY_MAP.get(failure_tag, FAILURE_RECOVERY_MAP["unknown"])
+        action_name = recovery["primary"]
+        # Check if primary was already tried
+        if previous_actions and action_name in previous_actions:
+            action_name = recovery["fallback"]
+        # Build action
+        new_tier = None
+        if action_name == "switch_model":
+            new_tier = min(current_tier + ESCALATION_TIERS["switch_model"], 5)
+        self.recovery_stats[action_name] = self.recovery_stats.get(action_name, 0) + 1
+        return RecoveryAction(
+            action=action_name,
+            reasoning=f"failure={failure_tag}, retry={retry_num}, action={action_name}",
+            new_tier=new_tier,
         )
 
-    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"]
+    def reset_run(self):
+        self.retry_counts = {}
+        self.total_retries = 0