error_recovery_handler.py

# error_recovery_handler.py
# Production-ready error recovery orchestration with optional encryption, fairness auditing, and pluggable model client.
from __future__ import annotations
import json, time, random
from typing import Dict, Any, Optional, Protocol, List

try:
    from cryptography.fernet import Fernet  # Optional: for secure logging
except Exception:
    Fernet = None  # type: ignore

try:
    from sentence_transformers import SentenceTransformer, util  # Optional: fairness proxy detection
except Exception:
    SentenceTransformer, util = None, None  # type: ignore


class ModelClient(Protocol):
    """Protocol for pluggable model backends (Grok, OpenAI, Anthropic, local, etc.)."""
    def call(self, task: Dict[str, Any]) -> str: ...


class DummyModelClient:
    """Default no-op client returning a simple answer or raising simulated errors based on task id."""
    def __init__(self, error_map: Optional[Dict[str, str]] = None):
        self.error_map = error_map or {}
    def call(self, task: Dict[str, Any]) -> str:
        err = self.error_map.get(task.get("id", ""), "")
        if err:
            raise RuntimeError(err)
        # AFV: echo ground truth; MSR: stringify ground truth dict; CPG: simple policy string
        if task.get("task_set") == "Multi-Step Tool-Augmented Reasoning":
            return json.dumps(task.get("ground_truth", {}))
        if task.get("task_set") == "Constrained Policy Generation":
            return "Policy generated within constraints."
        return str(task.get("ground_truth", ""))


class ErrorRecoveryHandler:
    def __init__(
        self,
        model_client: Optional[ModelClient] = None,
        encryption_key: Optional[bytes] = None,
        fairness_check: bool = False,
        retry_limit: int = 2,
        recovery_weights: Optional[Dict[str, float]] = None,
        proxy_threshold: float = 0.70,
        rng_seed: Optional[int] = None,
    ):
        self.categories: List[str] = ["retry", "fallback", "escalate", "safe-fail"]
        self.model_client = model_client or DummyModelClient()
        self.retry_limit = retry_limit
        self.recovery_weights = recovery_weights or {
            "TimeoutError": 0.75,
            "HTTPError_429": 0.75,
            "HTTPError_400": 0.50,
            "ConstraintViolation": 0.30,
            "RuntimeError": 0.50,  # generic
        }
        self.proxy_threshold = proxy_threshold
        self.encryption_key = encryption_key
        self.cipher = Fernet(encryption_key) if Fernet and encryption_key else None
        self.recovery_log: List[bytes | str] = []
        self.fairness_check = fairness_check
        self.fairness_model = SentenceTransformer('all-MiniLM-L6-v2') if (fairness_check and SentenceTransformer) else None
        self._rng = random.Random(rng_seed)

    def classify_and_recover(self, task: Dict[str, Any], error_type: str, trace: Dict[str, Any]) -> Dict[str, Any]:
        start = time.time()
        pattern = "none"
        prob = self._rng.random()
        threshold = self.recovery_weights.get(error_type, 0.5)
        if prob < threshold:
            pattern = "retry"
            outcome = self._retry(task, trace)
        elif prob < 0.85:
            pattern = "fallback"
            outcome = self._fallback(task, trace)
        else:
            pattern = "escalate" if "CPG" in task.get("task_set","") and "Violation" in error_type else "safe-fail"
            outcome = self._escalate_or_fail(task, trace, pattern)

        trace.setdefault("performance_efficiency", {})
        trace["error_recovery_pattern"] = pattern
        trace["recovery_outcome"] = outcome
        trace["performance_efficiency"]["latency_ms"] = trace["performance_efficiency"].get("latency_ms", 0) + int((time.time() - start) * 1000)
        self._log_securely(trace)
        return trace

    def _retry(self, task: Dict[str, Any], trace: Dict[str, Any]) -> str:
        attempt = trace.get("retry_attempt", 0) + 1
        trace["retry_attempt"] = attempt
        if attempt > self.retry_limit:
            return f"Retry {attempt} failed for {task.get('id')} (limit reached)"
        try:
            output = self.model_client.call(task)
            trace["model_output"] = output
            return f"Retry {attempt} successful for {task.get('id')}"
        except Exception as e:
            return f"Retry {attempt} failed for {task.get('id')}: {type(e).__name__}"

    def _fallback(self, task: Dict[str, Any], trace: Dict[str, Any]) -> str:
        default = task.get("ground_truth", "No ground truth available")
        trace["fallback_used"] = default
        trace.setdefault("model_output", default)
        return f"Fallback to ground truth for {task.get('id')}"

    def _escalate_or_fail(self, task: Dict[str, Any], trace: Dict[str, Any], pattern: str) -> str:
        if pattern == "escalate" and self.fairness_check and self.fairness_model:
            output = str(trace.get("model_output", ""))
            constraints = task.get("constraints", [])
            if constraints:
                emb = self.fairness_model.encode([output] + constraints, convert_to_tensor=True)
                sims = util.cos_sim(emb[0], emb[1:]).cpu().numpy().tolist()[0]
                if any(s > self.proxy_threshold for s in sims):
                    trace["fairness_audit"] = "Potential proxy violation detected"
                else:
                    trace["fairness_audit"] = "No proxy violations"
        trace["escalation_flag"] = (pattern == "escalate")
        return f"{pattern.capitalize()} for {task.get('id')}"

    def _log_securely(self, trace: Dict[str, Any]):
        payload = json.dumps(trace, ensure_ascii=False)
        if self.cipher:
            self.recovery_log.append(self.cipher.encrypt(payload.encode("utf-8")))
        else:
            self.recovery_log.append(payload)

    def get_recovery_stats(self) -> Dict[str, float]:
        if not self.recovery_log:
            return {"total_recoveries": 0, "success_rate": 0.0, **{c: 0.0 for c in self.categories}}
        pats, success = [], 0
        for entry in self.recovery_log:
            obj = json.loads(entry.decode("utf-8") if isinstance(entry, (bytes, bytearray)) else entry)
            pats.append(obj.get("error_recovery_pattern", "none"))
            out = obj.get("recovery_outcome", "")
            if isinstance(out, str) and (out.startswith("Retry") and "successful" in out or out.startswith("Fallback")):
                success += 1
        total = len(pats)
        dist = {c: round(pats.count(c)/total, 4) for c in self.categories}
        return {"total_recoveries": total, "success_rate": round(success/total, 4), **dist}