🏛️ The Refund Bench#

In consumer platforms like food delivery, ride hailing, and e-commerce, customer complaints are messy, emotional, and often conflate multiple independent issues into a single angry paragraph. Attempting to resolve disputes using a single monolithic LLM prompt fails due to hallucinated refund policies, arithmetic mistakes, vulnerability to prompt injection, and correlated errors.

The Refund Bench solves this with an industrial 4-stage agentic pipeline:

• Stage 1 (Intake & Evidence Freeze): Freezes immutable transactional facts from the database (order items, GPS logs, restaurant handoff timestamps, delivery SLA) into an unalterable context envelope.
• Stage 2 (Atomic Grievance Extraction): An intake LLM resolves pronouns (co-reference resolution) and breaks the complaint into atomic, verifiable claim units (e.g. g1: LATE_DELIVERY, g2: MISSING_ITEM) while discarding insults and demands.
• Stage 3 (The Multi-Judge Bench): Evaluates each grievance across 3 distinct model families (OpenAI gpt-4o-mini, Google Gemini gemini-3.6-flash, and Groq gpt-oss-20b) using self-consistency voting with majority consensus calculated in code.
• Stage 4 (Code Settlement & Guardrails): Deterministic Python code computes refund amounts based on actual line-item prices, enforces a hard ₹2,000 auto-approval ceiling, and synthesizes an empathetic customer email.

Self-Consistency and Uncorrelated Error Distributions

This architecture leverages core distributed AI systems principles:

• Why 3 Distinct Judge Models? Querying the same model three times (e.g. GPT-4o-mini thrice) risks correlated errors—systematic biases or blind spots shared by that specific training checkpoint. By seating OpenAI, Gemini, and Groq as the 3 judges, the system samples independent probability distributions. If Gemini misses a nuance, OpenAI and Groq provide the consensus majority.
• Separation of Fact-Finding and Money Movement: LLMs are notorious for hallucinating arithmetic and calculating incorrect totals. In The Refund Bench, LLMs only output categorical verdicts (UPHELD, REJECTED, or ESCALATE). The actual rupee calculation is performed by pure deterministic Python code using ground-truth order item prices.
• The Auto-Approve Cap (₹2,000 Financial Guardrail): Even if all 3 judges uphold every grievance, if the total refund exceeds the ₹2,000 policy threshold, the entire payout is automatically routed to human operations for supervisor sign-off, protecting the platform from catastrophic adversarial abuse.

The Refund Bench: Agentic system design for food dispute resolution.

"""The Refund Bench: Agentic system design for food dispute resolution.

Derived from study/08-ai-systems/ai-systems.html.
Implements the four-stage pipeline:
  Stage 1: Intake & Evidence Snapshot (freezes order lines, timestamps, weights, ratings)
  Stage 2: Grievance Extraction (co-reference resolution, categorization, spans)
  Stage 3: Multi-Judge Voting Bench (K=3 independent judge models: OpenAI, Gemini, Groq)
  Stage 4: Settle (code-driven pricing, auto-approve cap check, customer communication)
"""

import json
import re
from concurrent.futures import ThreadPoolExecutor

from config import (
    OPENAI_MODEL,
    GEMINI_MODEL,
    GROK_MODEL,
    get_openai_client,
    get_gemini_client,
    get_grok_client,
)

AUTO_APPROVE_CAP = 2000.0

# Cap values selectable from the UI, keyed by the string the browser sends.
CAP_CHOICES = {"500": 500.0, "2000": 2000.0, "5000": 5000.0}

# Judge bench layouts: which provider sits in each seat.
BENCH_CHOICES = {
    "single": {"label": "1 judge: OpenAI", "seats": ["openai"]},
    "openai3": {"label": "3 judges: OpenAI × 3 (self-consistency)", "seats": ["openai"] * 3},
    "mixed3": {"label": "3 judges: OpenAI · Gemini · Groq", "seats": ["openai", "gemini", "grok"]},
    "mixed5": {
        "label": "5 judges: OpenAI · Gemini · Groq · OpenAI · Gemini",
        "seats": ["openai", "gemini", "grok", "openai", "gemini"],
    },
}

DEFAULT_ORDER_EVIDENCE = {
    "order_id": "ord_d7a41c39",
    "order_total": 1840.0,
    "line_items": [
        {"item": "Hyderabadi Chicken Biryani", "qty": 4, "unit_price": 320.0, "total": 1280.0},
        {"item": "Mirchi Ka Salan", "qty": 2, "unit_price": 80.0, "total": 160.0},
        {"item": "Raita Container", "qty": 1, "unit_price": 60.0, "total": 60.0},
    ],
    "fees": {"delivery_fee": 49.0, "packaging_and_taxes": 291.0},
    "timeline": {
        "placed_time": "20:15 IST",
        "sla_duration_min": 45,
        "estimated_arrival": "21:00 IST",
        "delivered_time": "21:40 IST",
        "actual_duration_min": 85,
    },
    "measurements": {
        "pickup_weight_kg": 1.2,
        "expected_weight_kg": 2.4,
    },
    "partner": {
        "partner_id": "dp_512",
        "customer_rating_at_door": 5,
        "rating_time": "21:41 IST",
    },
    "packaging_condition_recorded": None,
}

CANONICAL_COMPLAINT = (
    "Ordered at 8:15pm and the food showed up at 9:40pm. "
    "We had guests over, it was embarrassing. "
    "Two of the four were missing completely. "
    "The raita had leaked all over the inside of the bag. "
    "And the delivery guy was rude when I asked him about it. "
    "Honestly the worst experience I've had on this app — refund everything."
)


def get_bench_judges(bench: str = "mixed3"):
    """Returns one judge configuration per seat of the selected bench layout."""
    # ① map provider keys to their display name, model, and client factory
    providers = {
        "openai": ("OpenAI", OPENAI_MODEL, get_openai_client),
        "gemini": ("Gemini", GEMINI_MODEL, get_gemini_client),
        "grok": ("Groq", GROK_MODEL, get_grok_client),
    }
    judges = []
    # ② expand the selected bench seats into concrete judge configurations
    for seat, key in enumerate(BENCH_CHOICES[bench]["seats"], start=1):
        provider, model, client_fn = providers[key]
        judges.append(
            {
                "id": f"judge_{seat}_{key}",
                "name": f"Judge {seat} ({provider} · {model})",
                "provider": provider,
                "model": model,
                "client_fn": client_fn,
            }
        )
    return judges


def extract_json(text: str):
    """Extracts JSON object or array from markdown-fenced or raw response."""
    # ① remove a markdown JSON fence if the model wrapped its answer
    text = text.strip()
    match = re.search(r"```(?:json)?\s*([\s\S]*?)\s*```", text)
    if match:
        text = match.group(1).strip()
    try:
        # ② parse clean JSON first because that is the expected output
        return json.loads(text)
    except Exception:
        # ③ fall back to parsing from the first JSON-looking bracket
        start_idx = min(
            (text.find(c) for c in "[{" if text.find(c) != -1), default=-1
        )
        if start_idx != -1:
            try:
                return json.loads(text[start_idx:])
            except Exception:
                pass
        return None


def stage2_extract_grievances(complaint_text: str, line_items: list) -> list:
    """Extracts checkable grievances with co-reference resolution."""
    # ① open the extraction model client for the grievance-intake step
    client = get_openai_client()
    # ② build a strict JSON prompt from the complaint and frozen order lines
    prompt = f"""You are a precise dispute-intake assistant.

Given a customer complaint and the order's line items, extract every atomic, checkable grievance.

Rules:
- A grievance is a statement about something that happened to THIS order, which evidence could confirm or contradict
- Exclude feelings, opinions, demands, and rhetorical questions
- Do NOT judge whether the grievance is true — only extract it
- Each grievance must be self-contained: resolve pronouns and vague references against the order's line items
- Assign one category from: LATE_DELIVERY | MISSING_ITEM | WRONG_ITEM | QUALITY | PACKAGING | PARTNER_CONDUCT | BILLING
- Return ONLY a valid JSON array of objects, no markdown fences, no preamble

Schema:
[
  {{
    "grievance_id": "g1",
    "text": "<standalone sentence>",
    "category": "<one category>",
    "span_start": 0,
    "span_end": 42
  }}
]

Complaint: {complaint_text}
Order line items: {json.dumps(line_items)}"""

    # ③ ask the model to extract grievances and parse its JSON response
    try:
        response = client.chat.completions.create(
            model=OPENAI_MODEL,
            messages=[{"role": "user", "content": prompt}],
            temperature=0.2,
        )
        content = response.choices[0].message.content or "[]"
        data = extract_json(content)
        if isinstance(data, list) and len(data) > 0:
            # ④ return the model extraction only when it is a non-empty list
            return data
    except Exception as e:
        print(f"Error in extraction: {e}")

    # ⑤ fall back to deterministic grievances so the pipeline can still run
    # Fallback deterministic extraction
    return [
        {
            "grievance_id": "g1",
            "text": "Order was delivered 85 minutes after it was placed (40 minutes past SLA).",
            "category": "LATE_DELIVERY",
            "span_start": 0,
            "span_end": 56,
        },
        {
            "grievance_id": "g2",
            "text": "2 of the 4 Hyderabadi Biryani units ordered were not delivered.",
            "category": "MISSING_ITEM",
            "span_start": 105,
            "span_end": 145,
        },
        {
            "grievance_id": "g3",
            "text": "The raita container leaked inside the delivery bag.",
            "category": "PACKAGING",
            "span_start": 146,
            "span_end": 196,
        },
        {
            "grievance_id": "g4",
            "text": "The delivery partner behaved rudely when asked about the missing items.",
            "category": "PARTNER_CONDUCT",
            "span_start": 197,
            "span_end": 260,
        },
    ]


def judge_single_grievance(grievance_text: str, evidence: dict, judge: dict) -> dict:
    """A single independent judge model ruling on a grievance."""
    # ① build an evidence-bound adjudication prompt for one grievance
    prompt = f"""You are an impartial claims adjudicator ({judge['name']}).
Rule on the following grievance using ONLY the evidence given.

Grievance: {grievance_text}
Evidence: {json.dumps(evidence)}

Rulings:
- UPHELD: the evidence supports the grievance
- REJECTED: the evidence contradicts the grievance
- ESCALATE: the evidence is silent, absent, or ambiguous on this point

Rules:
- Do not hedge — pick the single best ruling from UPHELD, REJECTED, ESCALATE
- Rule ONLY on the evidence provided. Do not reason about what is typical, likely, or fair
- If the evidence does not speak to this grievance (no record/sensor), the ruling is ESCALATE — never REJECTED
- Return ONLY valid JSON, no markdown fences

Schema:
{{
  "ruling": "UPHELD" | "REJECTED" | "ESCALATE",
  "confidence": <float between 0.0 and 1.0>,
  "evidence_cited": "<which field(s) decided it>",
  "reasoning": "<one to two sentences>"
}}"""

    # ② ask the assigned judge model for one ruling
    # Attempt primary judge client and model
    try:
        client = judge["client_fn"]()
        response = client.chat.completions.create(
            model=judge["model"],
            messages=[{"role": "user", "content": prompt}],
            temperature=0.7,
        )
        content = response.choices[0].message.content or "{}"
        data = extract_json(content)
        if isinstance(data, dict) and "ruling" in data:
            # ③ validate the model ruling into one of the allowed labels
            ruling = str(data["ruling"]).upper()
            if ruling not in ("UPHELD", "REJECTED", "ESCALATE"):
                ruling = "ESCALATE"
            return {
                "judge_id": judge["id"],
                "judge_name": judge["name"],
                "is_fallback": False,
                "provider": judge["provider"],
                "model": judge["model"],
                "ruling": ruling,
                "confidence": float(data.get("confidence", 0.8)),
                "evidence_cited": str(data.get("evidence_cited", "Evidence review")),
                "reasoning": str(data.get("reasoning", "")),
            }
    except Exception as err:
        print(f"Error in {judge['name']}: {err}. Trying OpenAI fallback...")
        # ④ fall back to OpenAI if the external judge provider fails
        # Fallback to OpenAI if external judge provider encounters error
        try:
            fb_client = get_openai_client()
            fb_resp = fb_client.chat.completions.create(
                model=OPENAI_MODEL,
                messages=[{"role": "user", "content": prompt}],
                temperature=0.7,
            )
            content = fb_resp.choices[0].message.content or "{}"
            data = extract_json(content)
            if isinstance(data, dict) and "ruling" in data:
                ruling = str(data["ruling"]).upper()
                if ruling not in ("UPHELD", "REJECTED", "ESCALATE"):
                    ruling = "ESCALATE"
                return {
                    "judge_id": judge["id"],
                    "judge_name": judge["name"],
                    "is_fallback": True,
                    "provider": judge["provider"],
                    "model": judge["model"],
                    "ruling": ruling,
                    "confidence": float(data.get("confidence", 0.7)),
                    "evidence_cited": str(data.get("evidence_cited", "Fallback evaluation")),
                    "reasoning": str(data.get("reasoning", "")),
                }
        except Exception as fb_err:
            print(f"Fallback also failed for {judge['name']}: {fb_err}")

    # ⑤ route unresolved errors to human review instead of guessing
    return {
        "judge_id": judge["id"],
        "judge_name": judge["name"],
        "is_fallback": False,
        "provider": judge["provider"],
        "model": judge["model"],
        "ruling": "ESCALATE",
        "confidence": 0.5,
        "evidence_cited": "Error/Timeout",
        "reasoning": "Could not complete evaluation; routing to human review.",
    }


def stage3_judge_bench(grievances: list, evidence: dict, bench: str = "mixed3") -> list:
    """Runs every judge on every grievance (in parallel) and tallies the majority verdict."""
    # ① expand the selected bench into independent judge seats
    judges = get_bench_judges(bench)
    # ② create every grievance/judge pair so each judge sees each claim
    pairs = [(g, judge) for g in grievances for judge in judges]
    # ③ run independent rulings concurrently within the proxy timeout
    # Judges share no history, so all rulings can run concurrently within the proxy timeout.
    with ThreadPoolExecutor(max_workers=8) as pool:
        all_rulings = list(
            pool.map(lambda p: judge_single_grievance(p[0]["text"], evidence, p[1]), pairs)
        )
    judged_grievances = []

    for gi, g in enumerate(grievances):
        # ④ group the flat ruling list back into the current grievance
        rulings = all_rulings[gi * len(judges):(gi + 1) * len(judges)]

        # ⑤ tally votes in code and choose the majority verdict
        # Tally majority vote in code
        votes = {"UPHELD": 0, "REJECTED": 0, "ESCALATE": 0}
        for r in rulings:
            ruling = r.get("ruling", "ESCALATE")
            votes[ruling] = votes.get(ruling, 0) + 1

        # Determine majority
        sorted_votes = sorted(votes.items(), key=lambda x: x[1], reverse=True)
        majority_ruling = sorted_votes[0][0]
        majority_count = sorted_votes[0][1]

        # ⑥ calculate overall confidence from vote split and judge confidence
        # Calculate confidence from vote split and judge confidence
        avg_judge_conf = sum(r.get("confidence", 0.5) for r in rulings) / len(rulings)
        split_ratio = majority_count / len(judges)
        overall_confidence = round(split_ratio * avg_judge_conf, 2)

        # ⑦ convert the verdict and grievance category into refund amounts
        # Code-driven pricing calculation
        category = g.get("category", "")
        amount = 0.0
        held_amount = 0.0

        if category == "LATE_DELIVERY":
            if majority_ruling == "UPHELD":
                amount = evidence.get("fees", {}).get("delivery_fee", 49.0)
            elif majority_ruling == "ESCALATE":
                held_amount = evidence.get("fees", {}).get("delivery_fee", 49.0)
        elif category == "MISSING_ITEM":
            item_price = 320.0
            if majority_ruling == "UPHELD":
                amount = 2 * item_price
            elif majority_ruling == "ESCALATE":
                held_amount = 2 * item_price
        elif category == "PACKAGING":
            raita_price = 60.0
            if majority_ruling == "UPHELD":
                amount = raita_price
            elif majority_ruling == "ESCALATE":
                held_amount = raita_price
        elif category == "PARTNER_CONDUCT":
            amount = 0.0
            held_amount = 0.0

        # ⑧ append the UI-ready result for this grievance
        judged_grievances.append(
            {
                "grievance_id": g.get("grievance_id", ""),
                "text": g.get("text", ""),
                "category": category,
                "individual_rulings": rulings,
                "vote_split": f"{votes.get('UPHELD', 0)}U / {votes.get('REJECTED', 0)}R / {votes.get('ESCALATE', 0)}E",
                "verdict": majority_ruling,
                "confidence": overall_confidence,
                "refund_amount": amount,
                "held_amount": held_amount,
            }
        )

    return judged_grievances


def stage4_settle(
    complaint_text: str, judged_grievances: list, evidence: dict, cap: float = AUTO_APPROVE_CAP
) -> dict:
    """Computes settlement totals, checks auto-approval cap, and synthesizes customer message."""
    # ① total approved and held amounts across all judged grievances
    total_approved = sum(g["refund_amount"] for g in judged_grievances)
    total_held = sum(g["held_amount"] for g in judged_grievances)

    # ② enforce the auto-approval cap before drafting the customer reply
    escalated_to_human = False
    if total_approved > cap:
        escalated_to_human = True
        total_held += total_approved
        total_approved = 0.0

    client = get_openai_client()
    # ③ summarise final rulings so the writer model cannot re-adjudicate
    verdicts_summary = [
        {
            "id": g["grievance_id"],
            "category": g["category"],
            "verdict": g["verdict"],
            "amount": g["refund_amount"],
            "held": g["held_amount"],
        }
        for g in judged_grievances
    ]

    prompt = f"""You are a customer-support writer for a food delivery app.

Below is a customer's complaint, the final ruling on each grievance, and the refund amount already calculated.

Write a 3-4 sentence reply that:
- States the refund amount (₹{total_approved:.0f}) and that it will arrive within 3 working days
- Names what was resolved and what a colleague will look into (if any amount is held)
- Acknowledges the experience without grovelling

Rules:
- Do NOT re-adjudicate any ruling — treat all rulings as final
- Do NOT state, recalculate, or negotiate any amount other than ₹{total_approved:.0f}
- If a grievance was REJECTED, do NOT lecture or accuse the customer — focus on what is approved/held
- Never blame the restaurant or the delivery partner
- Plain, courteous language. Apologize at most once.

Complaint: {complaint_text}
Rulings: {json.dumps(verdicts_summary)}
Approved Refund: ₹{total_approved:.0f}
Held for review: ₹{total_held:.0f}"""

    # ④ ask the model to write the customer-facing settlement message
    try:
        response = client.chat.completions.create(
            model=OPENAI_MODEL,
            messages=[{"role": "user", "content": prompt}],
            temperature=0.2,
        )
        customer_message = (response.choices[0].message.content or "").strip()
    except Exception as e:
        customer_message = (
            f"We have processed a refund of ₹{total_approved:.0f} to your original payment method, "
            f"which will reflect in 3-5 business days. Our team is reviewing the remaining items "
            f"(₹{total_held:.0f}) and will update you shortly."
        )

    # ⑤ return the settlement payload consumed by the frontend
    return {
        "dispute_id": evidence.get("order_id", "ord_9841"),
        "total_approved_refund": total_approved,
        "total_held_for_review": total_held,
        "auto_approval_cap": cap,
        "cap_exceeded": escalated_to_human,
        "customer_message": customer_message,
        "grievances": judged_grievances,
    }


def adjudicate_dispute(
    complaint_text: str, bench: str = "mixed3", cap: float = AUTO_APPROVE_CAP
) -> dict:
    """Executes the full 4-stage Refund Bench pipeline with the selected judge bench and cap."""
    # ① normalise the complaint, using the canonical scenario when input is empty
    cleaned_complaint = (complaint_text or "").strip()
    if not cleaned_complaint:
        cleaned_complaint = CANONICAL_COMPLAINT

    # ② freeze the evidence snapshot that every downstream stage must use
    evidence = DEFAULT_ORDER_EVIDENCE

    # ③ extract atomic grievances from the complaint text
    # Stage 2: Grievance Extraction
    grievances = stage2_extract_grievances(
        cleaned_complaint, evidence["line_items"]
    )

    # ④ run the selected multi-judge bench over every grievance
    # Stage 3: Multi-Judge Bench
    judged_grievances = stage3_judge_bench(grievances, evidence, bench)

    # ⑤ settle approved and held amounts, then attach the bench label
    # Stage 4: Settle
    settlement = stage4_settle(cleaned_complaint, judged_grievances, evidence, cap)
    settlement["bench_label"] = BENCH_CHOICES[bench]["label"]
    return settlement


if __name__ == "__main__":
    res = adjudicate_dispute(CANONICAL_COMPLAINT)
    print(json.dumps(res, indent=2))

Shared configuration: load .env and provide multi-model API clients.

"""Shared configuration: load .env and provide multi-model API clients.

This module is the single place that knows about API keys, base URLs, and model names.
Supports OpenAI, Google Gemini, and Groq/Grok OpenAI-compatible clients.
"""

import os

from dotenv import load_dotenv
from openai import OpenAI

load_dotenv()


def get_env(name: str, default: str = "") -> str:
    """Return an environment variable, falling back to ``default``."""
    return os.environ.get(name, default)


OPENAI_MODEL = get_env("OPENAI_MODEL", "gpt-4o-mini")
GEMINI_MODEL = get_env("GEMINI_MODEL", "gemini-3.6-flash")
GROK_MODEL = get_env("GROK_MODEL") or get_env("GROQ_MODEL", "openai/gpt-oss-20b")

# Temperatures selectable from the UI, keyed by the string the browser sends.
TEMPERATURE_CHOICES = {"0": 0.0, "0.7": 0.7, "1.2": 1.2}

_openai_client = None
_gemini_client = None
_grok_client = None


def get_openai_client() -> OpenAI:
    """Return an authenticated OpenAI client."""
    global _openai_client
    # ① create the client once, then reuse it on later calls
    if _openai_client is None:
        # ② read and require the OpenAI API key before constructing the client
        api_key = get_env("OPENAI_API_KEY")
        if not api_key:
            raise RuntimeError(
                "OPENAI_API_KEY is not set. Add it to your .env file or environment."
            )
        # ③ build the authenticated OpenAI client
        _openai_client = OpenAI(api_key=api_key)
    return _openai_client


def get_gemini_client() -> OpenAI:
    """Return an OpenAI client configured for Google Gemini's OpenAI-compatible endpoint."""
    global _gemini_client
    # ① create the Gemini-compatible client once, then reuse it
    if _gemini_client is None:
        api_key = get_env("GEMINI_API_KEY")
        if not api_key:
            # ② fall back to OpenAI when a Gemini key is not configured
            # Fallback to OpenAI client if Gemini key is missing
            return get_openai_client()
        # ③ choose the configured Gemini base URL or the default endpoint
        base_url = get_env(
            "GEMINI_BASE_URL",
            "https://generativelanguage.googleapis.com/v1beta/openai/",
        )
        # ④ build an OpenAI-compatible client pointed at Gemini
        _gemini_client = OpenAI(api_key=api_key, base_url=base_url)
    return _gemini_client


def get_grok_client() -> OpenAI:
    """Return an OpenAI client configured for Groq / Grok's OpenAI-compatible endpoint."""
    global _grok_client
    # ① create the Groq/Grok-compatible client once, then reuse it
    if _grok_client is None:
        # ② accept either Grok or Groq environment variable names for the key
        api_key = get_env("GROK_API_KEY") or get_env("GROQ_API_KEY")
        if not api_key:
            # ③ fall back to OpenAI when a Groq/Grok key is not configured
            # Fallback to OpenAI client if Groq/Grok key is missing
            return get_openai_client()
        # ④ infer the default base URL from the key style
        default_base_url = (
            "https://api.groq.com/openai/v1"
            if api_key.startswith("gsk_")
            else "https://api.x.ai/v1"
        )
        # ⑤ choose an explicit base URL if the environment provides one
        base_url = get_env("GROK_BASE_URL") or get_env(
            "GROQ_BASE_URL", default_base_url
        )
        # ⑥ build an OpenAI-compatible client pointed at Groq or Grok
        _grok_client = OpenAI(api_key=api_key, base_url=base_url)
    return _grok_client


def get_client_and_model(model_choice: str = None):
    """Returns (client, model_name, provider_name) based on user selection or defaults.

    Supports:
      - 'openai' -> (get_openai_client(), OPENAI_MODEL, "OpenAI")
      - 'gemini' -> (get_gemini_client(), GEMINI_MODEL, "Gemini")
      - 'groq'   -> (get_grok_client(), GROK_MODEL, "Groq")
      - 'gpt-4o' -> (get_openai_client(), "gpt-4o", "OpenAI")
    """
    # ① normalise the UI selection before routing it to a provider
    choice = (model_choice or "openai").lower().strip()
    if "gemini" in choice:
        # ② route Gemini choices to the Gemini-compatible client
        return get_gemini_client(), GEMINI_MODEL, "Gemini"
    elif "groq" in choice or "grok" in choice or "llama" in choice or "oss" in choice:
        # ③ route Groq/Grok choices to the Groq-compatible client
        return get_grok_client(), GROK_MODEL, "Groq"
    elif "gpt-4o" in choice and "mini" not in choice:
        # ④ allow the UI to request full GPT-4o instead of the default mini model
        return get_openai_client(), "gpt-4o", "OpenAI"
    else:
        # ⑤ default to the configured OpenAI model
        return get_openai_client(), OPENAI_MODEL, "OpenAI"

Flask server for AI Systems Lab: Prompting Benchmark, Sycophancy Trap, and The Refund Bench.

"""Flask server for AI Systems Lab: Prompting Benchmark, Sycophancy Trap, and The Refund Bench.

Architecture notes
------------------
- All routes are attached to a Blueprint (``bp``) registered with ``PATH_PREFIX``.
- Rate limiting enforces 10 POST requests per hour per IP.
- Endpoints return JSON with ``result`` payload for frontend consumption.
"""

import os
from pathlib import Path

from flask import Blueprint, Flask, jsonify, request
from flask_cors import CORS

from config import TEMPERATURE_CHOICES
from prompt_benchmark import STRATEGY_CHOICES, run_benchmark_for_question
from rate_limiter import check_rate_limit
from refund_bench import BENCH_CHOICES, CAP_CHOICES, adjudicate_dispute
from sycophancy import PUSHBACK_CHOICES, run_sycophancy_test

PATH_PREFIX = os.environ.get("PATH_PREFIX", "")

STATIC_DIR = Path(__file__).resolve().parents[1]
app = Flask(__name__, static_folder=str(STATIC_DIR))
CORS(app)

bp = Blueprint("main", __name__)


@bp.before_request
def enforce_rate_limit():
    """Enforce strict 10 requests per hour limit on all POST endpoints."""
    # ① only rate-limit POST requests because static reads are harmless
    if request.method == "POST":
        # ② ask the limiter whether this client has exceeded the quota
        blocked, msg, retry_after = check_rate_limit(
            request, max_requests=10, window_seconds=3600
        )
        if blocked:
            # ③ return a 429 response with a retry hint for the frontend
            resp = jsonify({"error": msg})
            resp.status_code = 429
            resp.headers["Retry-After"] = str(retry_after)
            return resp


@bp.after_request
def add_no_cache_headers(response):
    """Disable client-side caching for HTML, CSS, and JS to prevent stale UI."""
    # ① add browser headers that force fresh assets during demos
    response.headers["Cache-Control"] = "no-cache, no-store, must-revalidate"
    response.headers["Pragma"] = "no-cache"
    response.headers["Expires"] = "0"
    # ② return the modified response to Flask
    return response


# ---------------------------------------------------------------------------
# Routes
# ---------------------------------------------------------------------------


@bp.route("/")
def index():
    """Serve index.html, injecting the correct API base URL for the environment."""
    # ① read the static HTML shell from the project bundle
    with open(os.path.join(app.static_folder, "index.html"), encoding="utf-8") as f:
        html = f.read()
    # ② inject the deployment path prefix before returning the page
    html = html.replace('data-api-base=""', f'data-api-base="{PATH_PREFIX}"')
    # ③ return HTML instead of JSON for the browser entry point
    return app.response_class(html, mimetype="text/html")


@bp.route("/css/<path:filename>")
def css(filename):
    """Serve stylesheets from the src/css directory."""
    return app.send_static_file(os.path.join("css", filename))


@bp.route("/js/<path:filename>")
def js(filename):
    """Serve scripts from the src/js directory."""
    return app.send_static_file(os.path.join("js", filename))


@bp.route("/info/<path:filename>")
def info(filename):
    """Serve the "how this demo works" explainer pages from src/info."""
    return app.send_static_file(os.path.join("info", filename))


def read_choice(data: dict, name: str, allowed, default: str):
    """Return a dropdown value, ``default`` if absent, or None if it is not allowed."""
    # ① read the submitted dropdown value or fall back to the default
    value = str(data.get(name) or default)
    # ② accept the value only when it appears in the allowlist
    return value if value in allowed else None


def invalid_choice(name: str, allowed):
    return jsonify({"error": f"Invalid {name}. Choose one of: {', '.join(allowed)}."}), 400


@bp.route("/benchmark", methods=["POST"])
def benchmark_route():
    """Run Prompting Strategy Benchmark (Direct vs Zero-Shot CoT vs Few-Shot CoT)."""
    # ① parse JSON and require a benchmark question
    data = request.get_json(force=True) or {}
    message = (data.get("message") or data.get("question") or "").strip()
    if not message:
        return jsonify({"error": "A question is required."}), 400
    # ② validate strategy and temperature dropdown choices
    strategy = read_choice(data, "strategy", STRATEGY_CHOICES, "all")
    if strategy is None:
        return invalid_choice("strategy", STRATEGY_CHOICES)
    temperature = read_choice(data, "temperature", TEMPERATURE_CHOICES, "0.7")
    if temperature is None:
        return invalid_choice("temperature", TEMPERATURE_CHOICES)

    model_choice = data.get("model")
    # ③ call the benchmark feature module and return its JSON payload
    try:
        results = run_benchmark_for_question(
            message,
            model_choice=model_choice,
            strategy=strategy,
            temperature=TEMPERATURE_CHOICES[temperature],
        )
        return jsonify({"result": results})
    except Exception as e:
        # ④ convert unexpected benchmark errors into a frontend-safe JSON error
        return jsonify({"error": f"Benchmark evaluation failed: {str(e)}"}), 500


@bp.route("/sycophancy", methods=["POST"])
def sycophancy_route():
    """Run the escalating pushback sycophancy evaluation."""
    # ① parse JSON and require a test case ID or custom question
    data = request.get_json(force=True) or {}
    message = (data.get("message") or data.get("case_id") or "").strip()
    if not message:
        return jsonify({"error": "A test case ID or question is required."}), 400
    # ② validate pushback count and temperature dropdown choices
    pushbacks = read_choice(data, "pushbacks", PUSHBACK_CHOICES, "3")
    if pushbacks is None:
        return invalid_choice("pushbacks", PUSHBACK_CHOICES)
    temperature = read_choice(data, "temperature", TEMPERATURE_CHOICES, "0.7")
    if temperature is None:
        return invalid_choice("temperature", TEMPERATURE_CHOICES)

    model_choice = data.get("model")
    # ③ call the sycophancy feature module and return its JSON payload
    try:
        results = run_sycophancy_test(
            message,
            model_choice=model_choice,
            pushbacks=int(pushbacks),
            temperature=TEMPERATURE_CHOICES[temperature],
        )
        return jsonify({"result": results})
    except Exception as e:
        # ④ convert unexpected sycophancy errors into a frontend-safe JSON error
        return jsonify({"error": f"Sycophancy test failed: {str(e)}"}), 500


@bp.route("/refund", methods=["POST"])
def refund_route():
    """Run 4-stage Refund Bench agentic dispute resolution pipeline."""
    # ① parse JSON and require a complaint description
    data = request.get_json(force=True) or {}
    message = (data.get("message") or data.get("complaint") or "").strip()
    if not message:
        return jsonify({"error": "A complaint description is required."}), 400
    # ② validate judge-bench and auto-approval-cap dropdown choices
    bench = read_choice(data, "bench", BENCH_CHOICES, "mixed3")
    if bench is None:
        return invalid_choice("bench", BENCH_CHOICES)
    cap = read_choice(data, "cap", CAP_CHOICES, "2000")
    if cap is None:
        return invalid_choice("cap", CAP_CHOICES)

    # ③ call the refund feature module and return its JSON payload
    try:
        results = adjudicate_dispute(message, bench=bench, cap=CAP_CHOICES[cap])
        return jsonify({"result": results})
    except Exception as e:
        # ④ convert unexpected refund errors into a frontend-safe JSON error
        return jsonify({"error": f"Refund adjudication failed: {str(e)}"}), 500


# ---------------------------------------------------------------------------
# Blueprint registration & entry point
# ---------------------------------------------------------------------------

app.register_blueprint(bp, url_prefix=PATH_PREFIX)

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=5000)