benchmark.py

"""
Prompt Technique Benchmark Suite
Compares: Zero-shot, Few-shot, Chain-of-Thought, Self-Consistency
Task: Customer Support Intent Classification
"""

import os
import json
import time
from datetime import datetime
from litellm import completion
from evals.scorer import score_output
from utils.langfuse_logger import log_run

# ── Config ────────────────────────────────────────────────────────────────────
MODEL = os.getenv("BENCHMARK_MODEL", "gemini/gemini-2.0-flash")
RESULTS_DIR = "results"
os.makedirs(RESULTS_DIR, exist_ok=True)

# ── Test Dataset ──────────────────────────────────────────────────────────────
TEST_CASES = [
    {"input": "My order hasn't arrived and it's been 2 weeks.", "label": "delivery_issue"},
    {"input": "I want to cancel my subscription immediately.", "label": "cancellation"},
    {"input": "How do I reset my password?", "label": "account_access"},
    {"input": "The item I received is broken.", "label": "damaged_item"},
    {"input": "Can I get a refund for my last purchase?", "label": "refund_request"},
    {"input": "I was charged twice for the same order.", "label": "billing_error"},
    {"input": "I'd like to upgrade my plan.", "label": "upgrade_request"},
    {"input": "The app keeps crashing on my phone.", "label": "technical_issue"},
    {"input": "Where can I find my invoice?", "label": "billing_inquiry"},
    {"input": "I never received a confirmation email.", "label": "missing_confirmation"},
]

VALID_LABELS = [
    "delivery_issue", "cancellation", "account_access", "damaged_item",
    "refund_request", "billing_error", "upgrade_request", "technical_issue",
    "billing_inquiry", "missing_confirmation"
]

# ── Prompt Templates ──────────────────────────────────────────────────────────
def zero_shot_prompt(text):
    return f"""Classify this customer support message into exactly one category.

Categories: {", ".join(VALID_LABELS)}

Message: "{text}"

Respond with only the category label, nothing else."""


def few_shot_prompt(text):
    return f"""Classify customer support messages into one of these categories:
{", ".join(VALID_LABELS)}

Examples:
Message: "My package never showed up after 10 days."
Category: delivery_issue

Message: "Please cancel my account, I don't want to be charged again."
Category: cancellation

Message: "I got billed $50 but my plan is only $25."
Category: billing_error

Message: "The screen on my device cracked when I opened the box."
Category: damaged_item

---
Now classify this message:
Message: "{text}"
Category:"""


def cot_prompt(text):
    return f"""Classify this customer support message into one of these categories:
{", ".join(VALID_LABELS)}

Message: "{text}"

Think step by step:
1. What is the customer's main problem or request?
2. What action or resolution are they seeking?
3. Which category best fits?

After reasoning, output your final answer on a new line in this exact format:
CATEGORY: <label>"""


def self_consistency_prompt(text):
    """Returns the base prompt — called 3x, majority vote taken."""
    return f"""Classify this customer support message. Be concise.

Categories: {", ".join(VALID_LABELS)}

Message: "{text}"

Respond with only the category label."""


# ── Rate limit config ─────────────────────────────────────────────────────────
# Free tier = 5 RPM. We wait 13s between calls to stay safely under the limit.
# Set to 0 if you have a paid API key.
RATE_LIMIT_DELAY = float(os.getenv("RATE_LIMIT_DELAY", "15"))
MAX_RETRIES = 3


# ── Runner ────────────────────────────────────────────────────────────────────
def call_model(prompt: str, technique: str, sample_idx: int = 0) -> dict:
    for attempt in range(MAX_RETRIES):
        try:
            time.sleep(RATE_LIMIT_DELAY)
            start = time.time()
            response = completion(
                model=MODEL,
                messages=[{"role": "user", "content": prompt}],
                temperature=0.0 if technique != "self_consistency" else 0.7,
                max_tokens=500,
            )
            latency = round(time.time() - start, 3)
            raw_output = response.choices[0].message.content.strip()

            return {
                "raw_output": raw_output,
                "latency_s": latency,
                "prompt_tokens": response.usage.prompt_tokens,
                "completion_tokens": response.usage.completion_tokens,
                "total_tokens": response.usage.total_tokens,
            }

        except Exception as e:
            if "429" in str(e) or "RateLimitError" in type(e).__name__:
                wait = 60 * (attempt + 1)
                print(f"  [rate limit] Waiting {wait}s before retry {attempt + 1}/{MAX_RETRIES}...")
                time.sleep(wait)
            else:
                raise  # non-rate-limit errors bubble up immediately

    raise RuntimeError(f"Failed after {MAX_RETRIES} retries due to rate limiting.")


def parse_output(raw: str, technique: str) -> str:
    """Extract just the label from model output."""
    cleaned = raw.lower().strip()

    if technique == "cot":
        # Look for "CATEGORY: <label>"
        for line in cleaned.split("\n"):
            if "category:" in line:
                candidate = line.split("category:")[-1].strip()
                if candidate in VALID_LABELS:
                    return candidate
                for label in VALID_LABELS:
                    if label in candidate:
                        return label

    # Exact match on first line
    first_line = cleaned.split("\n")[0].strip().rstrip(".")
    if first_line in VALID_LABELS:
        return first_line

    # Longest match wins (avoids billing_ matching both billing_error and billing_inquiry)
    matches = [label for label in VALID_LABELS if label in cleaned]
    if matches:
        return max(matches, key=len)

    return first_line  # fallback


def run_self_consistency(text: str, n: int = 3) -> tuple[str, list]:
    """Run prompt n times, return majority vote."""
    votes = []
    raw_outputs = []
    for _ in range(n):
        result = call_model(self_consistency_prompt(text), "self_consistency")
        parsed = parse_output(result["raw_output"], "self_consistency")
        votes.append(parsed)
        raw_outputs.append(result["raw_output"])

    # Majority vote
    from collections import Counter
    winner = Counter(votes).most_common(1)[0][0]
    return winner, raw_outputs


# ── Main Benchmark Loop ───────────────────────────────────────────────────────
def run_benchmark():
    techniques = ["zero_shot", "few_shot", "cot", "self_consistency"]
    prompt_fns = {
        "zero_shot": zero_shot_prompt,
        "few_shot": few_shot_prompt,
        "cot": cot_prompt,
    }

    all_results = {}
    summary = {}

    print(f"\n{'='*60}")
    print(f"  Prompt Benchmark Suite — {MODEL}")
    print(f"  {len(TEST_CASES)} test cases × {len(techniques)} techniques")
    print(f"{'='*60}\n")

    for technique in techniques:
        print(f"▶ Running: {technique.upper().replace('_', '-')}")
        results = []
        correct = 0
        total_latency = 0
        total_tokens = 0

        for i, case in enumerate(TEST_CASES):
            try:
                if technique == "self_consistency":
                    predicted, raw_outputs = run_self_consistency(case["input"])
                    raw_output = " | ".join(raw_outputs)
                    # Approximate cost — 3 calls
                    result = call_model(self_consistency_prompt(case["input"]), technique)
                    latency = result["latency_s"]
                    tokens = result["total_tokens"] * 3
                else:
                    prompt = prompt_fns[technique](case["input"])
                    result = call_model(prompt, technique, i)
                    raw_output = result["raw_output"]
                    predicted = parse_output(raw_output, technique)
                    latency = result["latency_s"]
                    tokens = result["total_tokens"]

                is_correct = score_output(predicted, case["label"])
                correct += int(is_correct)
                total_latency += latency
                total_tokens += tokens

                run_data = {
                    "input": case["input"],
                    "expected": case["label"],
                    "predicted": predicted,
                    "correct": is_correct,
                    "raw_output": raw_output,
                    "latency_s": latency,
                    "tokens": tokens,
                }
                results.append(run_data)

                # Log to Langfuse
                log_run(
                    technique=technique,
                    input_text=case["input"],
                    predicted=predicted,
                    expected=case["label"],
                    correct=is_correct,
                    latency=latency,
                    tokens=tokens,
                    model=MODEL,
                )

                status = "✓" if is_correct else "✗"
                print(f"  [{status}] {case['input'][:45]:<45} → {predicted}")

            except Exception as e:
                import traceback
                print(f"  [!] Error on case {i}: {e}")
                traceback.print_exc()

        accuracy = round(correct / len(TEST_CASES) * 100, 1)
        avg_latency = round(total_latency / len(TEST_CASES), 3)
        avg_tokens = round(total_tokens / len(TEST_CASES))

        summary[technique] = {
            "accuracy": accuracy,
            "correct": correct,
            "total": len(TEST_CASES),
            "avg_latency_s": avg_latency,
            "avg_tokens": avg_tokens,
        }
        all_results[technique] = results

        print(f"  → Accuracy: {accuracy}% | Avg latency: {avg_latency}s | Avg tokens: {avg_tokens}\n")

    # ── Save results ──────────────────────────────────────────────────────────
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    output_path = f"{RESULTS_DIR}/run_{timestamp}.json"
    with open(output_path, "w") as f:
        json.dump({"summary": summary, "details": all_results}, f, indent=2)

    # ── Print summary table ───────────────────────────────────────────────────
    print(f"\n{'='*60}")
    print(f"  RESULTS SUMMARY")
    print(f"{'='*60}")
    print(f"  {'Technique':<20} {'Accuracy':>10} {'Avg Latency':>14} {'Avg Tokens':>12}")
    print(f"  {'-'*56}")
    for t, s in sorted(summary.items(), key=lambda x: -x[1]["accuracy"]):
        print(f"  {t:<20} {str(s['accuracy'])+'%':>10} {str(s['avg_latency_s'])+'s':>14} {s['avg_tokens']:>12}")
    print(f"{'='*60}")
    print(f"\n  Full results saved to: {output_path}\n")

    return summary


if __name__ == "__main__":
    run_benchmark()