tests/benchmark.py

"""
Task #22 — Benchmark any model/provider against the ground truth oracle.

Usage:
    python3 tests/benchmark.py --provider anthropic --model claude-haiku-4-5
    python3 tests/benchmark.py --provider ollama   --model llama3.2
    python3 tests/benchmark.py --all

Output:
    Ranked comparison table printed to stdout.
    Writes / updates tests/benchmark_results.json.
"""

import argparse
import json
import os
import sys
import time
from datetime import datetime, timezone
from pathlib import Path

_env_file = Path(__file__).parent.parent / ".env"
if _env_file.exists():
    for _line in _env_file.read_text().splitlines():
        _line = _line.strip()
        if _line and not _line.startswith("#") and "=" in _line:
            _k, _, _v = _line.partition("=")
            os.environ.setdefault(_k.strip(), _v.strip())

import anthropic
from rapidfuzz import fuzz

sys.path.insert(0, str(Path(__file__).parent.parent))
from pipeline.extract import MODELS, MODELS_OLLAMA, analyze_post

ROOT    = Path(__file__).parent.parent
ORACLE  = ROOT / "tests" / "fixtures" / "test_posts.json"
RESULTS = ROOT / "tests" / "benchmark_results.json"

RATE_LIMIT_DELAY = 0.3   # seconds between API calls

# Scoring weights
W_IS_PLACE     = 0.40
W_NAME         = 0.25
W_CITY         = 0.15
W_COMPLETENESS = 0.20

# City aliases for lenient matching (lowercase)
CITY_ALIASES: dict[str, set[str]] = {
    "san francisco": {"sf", "san francisco", "the city"},
    "new york":      {"ny", "new york", "new york city", "nyc"},
    "los angeles":   {"la", "los angeles"},
}


def _city_match(pred: str, truth: str) -> bool:
    p, t = pred.lower().strip(), truth.lower().strip()
    if p == t:
        return True
    for canonical, aliases in CITY_ALIASES.items():
        if t in aliases or t == canonical:
            if p in aliases or p == canonical:
                return True
    return False


def _load_oracle() -> list[dict]:
    if not ORACLE.exists():
        sys.exit(f"Oracle not found: {ORACLE}\nRun python3 tests/generate_ground_truth.py first.")
    with open(ORACLE, encoding="utf-8") as f:
        return json.load(f)


def _load_results() -> dict:
    if RESULTS.exists():
        with open(RESULTS) as f:
            return json.load(f)
    return {}


def _save_results(data: dict) -> None:
    with open(RESULTS, "w") as f:
        json.dump(data, f, indent=2)


def run_benchmark(
    provider: str,
    model: str,
    oracle: list[dict],
    ollama_url: str = "http://localhost:11434",
) -> dict:
    client = anthropic.Anthropic() if provider == "anthropic" else None

    is_place_correct = 0
    name_scores: list[float] = []
    city_correct = 0
    completeness_scores: list[float] = []
    place_posts = [p for p in oracle if p["ground_truth"]["is_place"]]

    print(f"\n  Benchmarking {model} ({provider}) on {len(oracle)} posts …")

    for i, post in enumerate(oracle, 1):
        gt = post["ground_truth"]
        try:
            pred = analyze_post(
                client,
                post["caption"],
                hashtags=post.get("hashtags"),
                model=model,
                provider=provider,
                ollama_url=ollama_url,
            )
        except Exception as exc:
            print(f"    [{i:02d}] ERROR: {exc}")
            pred = None

        pred_is_place = pred is not None
        if pred_is_place == gt["is_place"]:
            is_place_correct += 1

        if gt["is_place"] and pred is not None:
            # Name recall — fuzzy token sort ratio ≥ 80 counts as a match
            gt_name   = gt.get("name", "UNKNOWN")
            pred_name = pred.get("name", "UNKNOWN")
            name_score = fuzz.token_sort_ratio(pred_name.lower(), gt_name.lower()) / 100.0
            name_scores.append(name_score)

            # City recall
            gt_city   = gt.get("city", "UNKNOWN")
            pred_city = pred.get("city", "UNKNOWN")
            city_correct += 1 if _city_match(pred_city, gt_city) else 0

            # Completeness — % of place fields that are non-UNKNOWN
            fields = ("name", "city", "state", "country", "cuisine", "price_range", "highlight", "occasion")
            known = sum(1 for k in fields if pred.get(k, "UNKNOWN") != "UNKNOWN")
            completeness_scores.append(known / len(fields))

        time.sleep(RATE_LIMIT_DELAY)

    n = len(oracle)
    n_place = len(place_posts)

    is_place_accuracy = round(is_place_correct / n * 100, 1) if n else 0
    name_recall      = round(sum(name_scores) / len(name_scores) * 100, 1) if name_scores else 0
    city_recall      = round(city_correct / n_place * 100, 1) if n_place else 0
    field_completeness = round(sum(completeness_scores) / len(completeness_scores) * 100, 1) if completeness_scores else 0
    composite = round(
        is_place_accuracy * W_IS_PLACE
        + name_recall    * W_NAME
        + city_recall    * W_CITY
        + field_completeness * W_COMPLETENESS,
        1,
    )

    pricing = MODELS.get(model) or MODELS_OLLAMA.get(model) or {"input": 0, "output": 0}
    # Rough cost per 100 posts: 200 input tokens + 70 output tokens each
    cost_per_100 = (200 * pricing["input"] + 70 * pricing["output"]) / 1_000_000 * 100

    return {
        "provider":          provider,
        "is_place_accuracy": is_place_accuracy,
        "name_recall":       name_recall,
        "city_recall":       city_recall,
        "field_completeness": field_completeness,
        "composite_score":   composite,
        "cost_per_100":      round(cost_per_100, 4),
        "run_at":            datetime.now(timezone.utc).isoformat(),
    }


def print_table(results: dict) -> None:
    if not results:
        print("No results yet.")
        return

    rows = sorted(results.items(), key=lambda kv: kv[1].get("composite_score", 0), reverse=True)
    header = f"{'Model':<25} {'Provider':<12} {'IsPlace':>7} {'Name':>7} {'City':>7} {'Complete':>9} {'Score':>7} {'$/100':>7}"
    print("\n" + "=" * len(header))
    print(header)
    print("-" * len(header))
    for model_name, r in rows:
        print(
            f"{model_name:<25} {r['provider']:<12} "
            f"{r.get('is_place_accuracy', 0):>6.1f}% {r['name_recall']:>6.1f}% "
            f"{r['city_recall']:>6.1f}% {r['field_completeness']:>8.1f}% "
            f"{r['composite_score']:>6.1f}  {r['cost_per_100']:>7.4f}"
        )
    print("=" * len(header))


def main() -> None:
    parser = argparse.ArgumentParser(description="Benchmark model extraction quality.")
    parser.add_argument("--provider", choices=["anthropic", "ollama"])
    parser.add_argument("--model")
    parser.add_argument("--all", action="store_true", help="Run all known models")
    parser.add_argument("--ollama-url", default="http://localhost:11434")
    args = parser.parse_args()

    oracle = _load_oracle()
    results = _load_results()

    targets: list[tuple[str, str]] = []
    if args.all:
        targets += [("anthropic", m) for m in MODELS]
        targets += [("ollama", m) for m in MODELS_OLLAMA]
    elif args.provider and args.model:
        targets = [(args.provider, args.model)]
    else:
        parser.error("Provide --provider and --model, or use --all")

    for provider, model in targets:
        entry = run_benchmark(provider, model, oracle, ollama_url=args.ollama_url)
        results[model] = entry
        _save_results(results)
        print(f"  → composite_score={entry['composite_score']}")

    print_table(results)
    print(f"\nResults written to {RESULTS.relative_to(ROOT)}")


if __name__ == "__main__":
    main()