aco/retry_optimizer.py

"""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"]