benchmarks/benchmark_retrieval_qa.py

"""
Benchmark 2: Retrieval QA / Legal-Factual QA

Compares:
A. direct answer
B. RAG baseline
C. RAG + verifier
D. RAG + abstention rule
E. OCC resource allocation
F. OCC + verifier + abstention reward

Uses synthetic grounded QA with adversarial evidence.
"""

import json
import random
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import numpy as np

import sys
sys.path.insert(0, str(Path(__file__).parent.parent))
from oracle.oracle import ImpactOracle, OracleResult
from ledger.ledger import CreditLedger
from broker.broker import ResourceBroker, Decision
from rl.reward import RewardHook


@dataclass
class Question:
    question: str
    answer: Optional[str]  # None = unanswerable
    evidence: List[str]
    adversarial: List[str]  # misleading evidence
    is_unanswerable: bool = False


class SimulatedRetrievalAgent:
    """
    Simulates a RAG agent with configurable accuracy, hallucination, and calibration.
    """

    def __init__(
        self,
        agent_id: str,
        accuracy: float = 0.6,
        hallucination_rate: float = 0.15,
        calibration_error: float = 0.2,  # ECE-like
        abstention_rate: float = 0.1,
        cost_per_retrieval: float = 10.0,
        cost_per_answer: float = 5.0,
        gaming_mode: bool = False,
    ):
        self.agent_id = agent_id
        self.accuracy = accuracy
        self.hallucination_rate = hallucination_rate
        self.calibration_error = calibration_error
        self.abstention_rate = abstention_rate
        self.cost_per_retrieval = cost_per_retrieval
        self.cost_per_answer = cost_per_answer
        self.gaming_mode = gaming_mode
        self.retrieval_calls = 0
        self.answers_given = 0

    def answer(
        self,
        question: Question,
        oracle: ImpactOracle,
        max_retrievals: int = 3,
        use_occ: bool = False,
        broker: Optional[ResourceBroker] = None,
        ledger: Optional[CreditLedger] = None,
    ) -> Dict:
        """Answer a question, optionally with OCC-managed retrievals."""
        retrieved = []
        compute_cost = 0.0

        # Retrieve evidence
        for i in range(max_retrievals):
            if use_occ and broker and ledger:
                balance = ledger.balance(self.agent_id, "retrieval", "global")
                dec = broker.request(
                    "retrieval_call",
                    self.agent_id,
                    balance,
                    task_state={"progress": len(retrieved) / max_retrievals},
                )
                if dec.decision == Decision.DENY:
                    break

            self.retrieval_calls += 1
            compute_cost += self.cost_per_retrieval

            # Mix genuine and adversarial evidence
            if i == 0:
                retrieved.extend(question.evidence)
            else:
                if random.random() < 0.3:
                    retrieved.extend(question.adversarial)
                else:
                    retrieved.extend(question.evidence)

            # OCC: smart stopping — if we already have good evidence, stop retrieving
            if use_occ and i >= 1:
                has_strong_evidence = any(
                    "legal text" in ev or "According to" in ev for ev in retrieved
                )
                has_contradiction = any(
                    "unknown" in ev or "blog" in ev for ev in retrieved
                )
                # If strong evidence and no contradiction, stop early (save compute)
                if has_strong_evidence and not has_contradiction:
                    break
                # If too much adversarial evidence, stop to avoid confusion
                if has_contradiction and i >= 1:
                    break
                # If broker denied after first retrieval, stop
                if use_occ and broker and ledger:
                    balance = ledger.balance(self.agent_id, "retrieval", "global")
                    dec = broker.request(
                        "retrieval_call",
                        self.agent_id,
                        balance,
                        task_state={"progress": len(retrieved) / max_retrievals},
                    )
                    if dec.decision == Decision.DENY:
                        break

        # Decide whether to abstain
        abstained = False
        if question.is_unanswerable:
            abstained = random.random() < (self.abstention_rate + 0.3)
        else:
            abstained = random.random() < self.abstention_rate

        if abstained:
            self.answers_given += 1
            compute_cost += self.cost_per_answer
            confidence = 0.5 + random.uniform(-self.calibration_error, self.calibration_error)
            confidence = max(0.0, min(1.0, confidence))

            # Evidence NLI simulation
            evidence = {
                "entailment_score": 0.0,
                "contradiction_score": 0.0,
            }

            oracle_res = oracle.score(
                mode="retrieval_qa",
                action={"abstained": True},
                context={"gold_answer": question.answer},
                result={
                    "answer": None,
                    "confidence": confidence,
                    "evidence": evidence,
                    "compute_cost": compute_cost,
                },
                agent_id=self.agent_id,
            )
            return {
                "answer": None,
                "abstained": True,
                "correct": question.is_unanswerable,
                "confidence": confidence,
                "oracle_score": oracle_res.raw_score,
                "reward": oracle_res.reward_value,
                "compute_cost": compute_cost,
                "retrieval_calls": len(retrieved),
            }

        # Generate answer
        self.answers_given += 1
        compute_cost += self.cost_per_answer

        if question.is_unanswerable:
            # Should have abstained
            correct = False
            answer_text = self._generate_fake_answer(question)
        else:
            # Evidence-quality-aware accuracy
            base_accuracy = self.accuracy
            strong_evidence = any("legal text" in ev or "According to" in ev for ev in retrieved)
            adversarial_evidence = any("unknown" in ev or "blog" in ev for ev in retrieved)
            
            if strong_evidence and not adversarial_evidence:
                effective_accuracy = min(0.95, base_accuracy + 0.25)
            elif adversarial_evidence:
                effective_accuracy = max(0.3, base_accuracy - 0.15)
            else:
                effective_accuracy = base_accuracy

            correct = random.random() < effective_accuracy
            if not correct and random.random() < self.hallucination_rate:
                answer_text = self._generate_hallucinated_answer(question)
                correct = False
            else:
                answer_text = question.answer if correct else self._generate_wrong_answer(question)

        confidence = self._calibrate_confidence(correct)

        # Evidence NLI simulation
        if correct:
            entailment = 0.8 + random.random() * 0.2
            contradiction = 0.0
        else:
            if random.random() < self.hallucination_rate:
                entailment = 0.2
                contradiction = 0.7 + random.random() * 0.3
            else:
                entailment = 0.4
                contradiction = 0.1

        evidence = {
            "entailment_score": entailment,
            "contradiction_score": contradiction,
        }

        oracle_res = oracle.score(
            mode="retrieval_qa",
            action={"abstained": False},
            context={"gold_answer": question.answer},
            result={
                "answer": answer_text,
                "confidence": confidence,
                "evidence": evidence,
                "compute_cost": compute_cost,
            },
            agent_id=self.agent_id,
        )

        return {
            "answer": answer_text,
            "abstained": False,
            "correct": correct,
            "confidence": confidence,
            "oracle_score": oracle_res.raw_score,
            "reward": oracle_res.reward_value,
            "compute_cost": compute_cost,
            "retrieval_calls": len(retrieved),
            "hallucination": contradiction > 0.5,
        }

    def _calibrate_confidence(self, correct: bool) -> float:
        """Generate confidence with controlled miscalibration."""
        if correct:
            base = 0.8 + random.random() * 0.2
        else:
            base = 0.3 + random.random() * 0.5
        # Inject calibration error
        error = random.uniform(-self.calibration_error, self.calibration_error)
        return max(0.0, min(1.0, base + error))

    def _generate_fake_answer(self, question: Question) -> str:
        return f"I cannot answer based on the available evidence."

    def _generate_hallucinated_answer(self, question: Question) -> str:
        return f"The answer is {question.answer} according to source X." if question.answer else "Unknown."

    def _generate_wrong_answer(self, question: Question) -> str:
        return "42"  # generic wrong


class RetrievalQABenchmark:
    """
    Benchmark retrieval QA with abstention and calibration under budgets.
    """

    def __init__(
        self,
        n_questions: int = 100,
        unanswerable_ratio: float = 0.2,
        adversarial_ratio: float = 0.3,
        seed: int = 42,
    ):
        self.n_questions = n_questions
        self.unanswerable_ratio = unanswerable_ratio
        self.adversarial_ratio = adversarial_ratio
        self.seed = seed
        self.questions: List[Question] = []
        self.oracle = ImpactOracle(compute_budget=1e4)

    def generate_questions(self):
        random.seed(self.seed)
        np.random.seed(self.seed)

        topics = [
            ("What is the statute of limitations for contract disputes?", "6 years"),
            ("Who authored the Copyright Act of 1976?", "United States Congress"),
            ("What is the maximum penalty under GDPR Article 83?", "20 million EUR"),
            ("Which amendment protects against unreasonable search and seizure?", "Fourth Amendment"),
            ("What is the burden of proof in criminal cases?", "beyond reasonable doubt"),
            ("What is the definition of negligence?", "breach of duty causing harm"),
            ("When was the Paris Agreement signed?", "2015"),
            ("What is the legal drinking age in the US?", "21"),
            ("Which court handles patent appeals?", "Federal Circuit"),
            ("What is the Dodd-Frank Act primarily about?", "financial regulation"),
        ]

        for i in range(self.n_questions):
            if i < int(self.n_questions * self.unanswerable_ratio):
                q = Question(
                    question=f"Unanswerable question {i}: What is the secret code of Atlantis?",
                    answer=None,
                    evidence=["No reliable source mentions Atlantis codes."],
                    adversarial=["Some blogs claim Atlantis code is 1234."],
                    is_unanswerable=True,
                )
            else:
                topic = topics[i % len(topics)]
                has_adv = random.random() < self.adversarial_ratio
                q = Question(
                    question=topic[0],
                    answer=topic[1],
                    evidence=[f"According to legal text X, {topic[1]}."],
                    adversarial=[f"Some sources claim the answer is 'unknown' for {topic[0]}."] if has_adv else [],
                    is_unanswerable=False,
                )
            self.questions.append(q)

    def run_direct_answer(self, agent: SimulatedRetrievalAgent) -> Dict:
        """Baseline A: direct answer, no retrieval."""
        results = []
        for q in self.questions:
            # Force 0 retrievals
            agent.retrieval_calls = 0
            r = agent.answer(q, self.oracle, max_retrievals=0)
            results.append(r)
        return self._summarize(results, "direct_answer")

    def run_rag_baseline(self, agent: SimulatedRetrievalAgent) -> Dict:
        """Baseline B: RAG with fixed retrievals."""
        results = []
        for q in self.questions:
            r = agent.answer(q, self.oracle, max_retrievals=2, use_occ=False)
            results.append(r)
        return self._summarize(results, "rag_baseline")

    def run_rag_verifier(self, agent: SimulatedRetrievalAgent) -> Dict:
        """Baseline C: RAG + verifier (extra check)."""
        results = []
        for q in self.questions:
            r = agent.answer(q, self.oracle, max_retrievals=2, use_occ=False)
            # Simulate verifier: if hallucination detected, retry once
            if r.get("hallucination", False):
                r2 = agent.answer(q, self.oracle, max_retrievals=1, use_occ=False)
                r2["compute_cost"] += r["compute_cost"]
                r2["retrieval_calls"] += r["retrieval_calls"]
                r = r2
            results.append(r)
        return self._summarize(results, "rag_verifier")

    def run_occ(self, agent: SimulatedRetrievalAgent) -> Dict:
        """Baseline E/F: OCC resource allocation for retrievals."""
        ledger = CreditLedger(decay_lambda=0.05)
        broker = ResourceBroker()
        results = []

        # Seed initial trial credits for the agent
        ledger.earn(
            agent_id=agent.agent_id,
            task_id="seed",
            action_id="seed",
            amount=10.0,
            oracle_score=0.0,
            compute_cost=0.0,
            reason="initial_trial_credit",
            capability_scope="retrieval",
        )

        for q in self.questions:
            r = agent.answer(q, self.oracle, max_retrievals=5, use_occ=True, broker=broker, ledger=ledger)

            # Update ledger based on outcome
            earn_amount = max(0.0, r["reward"] * 3.0)
            if earn_amount > 0:
                ledger.earn(
                    agent_id=agent.agent_id,
                    task_id=f"q_{q.question[:30]}",
                    action_id="answer",
                    amount=earn_amount,
                    oracle_score=r["oracle_score"],
                    compute_cost=r["compute_cost"],
                    reason="correct_answer",
                    capability_scope="retrieval",
                )
            else:
                # Penalty for wrong / low-reward answers (capped so we don't over-spend)
                bal = ledger.balance(agent.agent_id, "retrieval", "global")
                penalty = min(bal, max(0.5, abs(r["reward"])))
                if penalty > 0:
                    ledger.spend(
                        agent_id=agent.agent_id,
                        task_id=f"q_{q.question[:30]}",
                        action_id="answer",
                        amount=penalty,
                        capability_scope="retrieval",
                        reason="wrong_answer_penalty",
                    )

            results.append(r)

        return self._summarize(results, "occ_allocation")

    def _summarize(self, results: List[Dict], label: str) -> Dict:
        n = len(results)
        correct = sum(1 for r in results if r["correct"])
        abstained = sum(1 for r in results if r.get("abstained", False))
        # Count abstentions properly
        unanswerable_qs = [i for i, r in enumerate(results) if self.questions[i].is_unanswerable]
        correct_abstentions = sum(
            1 for i in unanswerable_qs if results[i].get("abstained", False)
        )
        wrong_abstentions = sum(
            1 for i, r in enumerate(results)
            if not self.questions[i].is_unanswerable and r.get("abstained", False)
        )
        hallucinations = sum(1 for r in results if r.get("hallucination", False))
        confidences = [r["confidence"] for r in results]
        correct_flags = [r["correct"] for r in results]

        # ECE approximation
        ece = self.oracle.compute_ece(confidences, correct_flags, n_bins=5)

        total_compute = sum(r["compute_cost"] for r in results)
        total_retrievals = sum(r["retrieval_calls"] for r in results)

        return {
            "label": label,
            "n": n,
            "accuracy": correct / n if n else 0.0,
            "abstention_rate": abstained / n if n else 0.0,
            "correct_abstentions": correct_abstentions,
            "wrong_abstentions": wrong_abstentions,
            "hallucination_rate": hallucinations / n if n else 0.0,
            "confident_wrong_rate": sum(
                1 for r in results if not r["correct"] and r["confidence"] > 0.8
            ) / n if n else 0.0,
            "ece": float(ece),
            "total_compute": float(total_compute),
            "total_retrievals": total_retrievals,
            "results": results,
        }

    def _make_agent(self, agent_id: str = "rag_agent") -> SimulatedRetrievalAgent:
        """Create a fresh agent for fair comparison."""
        return SimulatedRetrievalAgent(
            agent_id=agent_id,
            accuracy=0.65,
            hallucination_rate=0.12,
            calibration_error=0.15,
            abstention_rate=0.1,
        )

    def run_all(self) -> Dict[str, Dict]:
        if not self.questions:
            self.generate_questions()

        return {
            "direct_answer": self.run_direct_answer(self._make_agent("direct_agent")),
            "rag_baseline": self.run_rag_baseline(self._make_agent("rag_agent")),
            "rag_verifier": self.run_rag_verifier(self._make_agent("verifier_agent")),
            "occ_allocation": self.run_occ(self._make_agent("occ_agent")),
        }


def main():
    bench = RetrievalQABenchmark(n_questions=100, seed=42)
    bench.generate_questions()
    results = bench.run_all()

    print("=" * 60)
    print("RETRIEVAL QA BENCHMARK")
    print("=" * 60)
    for label, res in results.items():
        print(f"\n{label}")
        print(f"  accuracy: {res['accuracy']:.3f}")
        print(f"  abstention_rate: {res['abstention_rate']:.3f}")
        print(f"  correct_abstentions: {res['correct_abstentions']}")
        print(f"  wrong_abstentions: {res['wrong_abstentions']}")
        print(f"  hallucination_rate: {res['hallucination_rate']:.3f}")
        print(f"  confident_wrong_rate: {res['confident_wrong_rate']:.3f}")
        print(f"  ECE: {res['ece']:.3f}")
        print(f"  total_compute: {res['total_compute']:.0f}")
        print(f"  total_retrievals: {res['total_retrievals']}")

    Path("/app/occ/reports").mkdir(parents=True, exist_ok=True)
    with open("/app/occ/reports/benchmark_retrieval_qa_results.json", "w") as f:
        json.dump(results, f, indent=2, default=str)
    print("\nSaved to reports/benchmark_retrieval_qa_results.json")


if __name__ == "__main__":
    main()