run_pipeline.py

"""LLM-backed pipeline + eval. Run after `analyze.py` to produce eval metrics.

Usage:
  python run_pipeline.py --provider anthropic                  # default
  python run_pipeline.py --provider openai
  python run_pipeline.py --provider anthropic --sample 20      # cheap dry-run
  python run_pipeline.py --provider anthropic --skip-eval      # skip eval pass

Cost note: at default settings the full 100-call run uses the cheap classifier
tier for tagging and the reasoning tier only for theme synthesis (one call). On
Anthropic with Haiku 4.5 + Sonnet 4.6 this is a few cents end-to-end. On
OpenAI with gpt-4o-mini + gpt-4o, similar.
"""
from __future__ import annotations

import argparse
import json
import os
from collections import defaultdict
from pathlib import Path

from dotenv import load_dotenv
from tqdm import tqdm

from pipeline.data_loader import load_dataset
from pipeline.call_classifier import classify
from pipeline.categorizer import (
    categorize_with_llm, categorize_from_ground_truth,
    aggregate_themes, synthesize_themes,
)
from pipeline.evals import categorization_eval


HERE = Path(__file__).parent
DATASET_DEFAULT = (HERE.parent / "interview-assignment" / "dataset").resolve()
OUT = HERE / "outputs"


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--provider", choices=["openai", "anthropic"], default="anthropic")
    ap.add_argument("--dataset", default=str(DATASET_DEFAULT))
    ap.add_argument("--sample", type=int, default=0,
                    help="If >0, run only first N calls (cheap dry-run).")
    ap.add_argument("--skip-eval", action="store_true")
    args = ap.parse_args()

    load_dotenv()
    from pipeline.llm_provider import LLMProvider
    llm = LLMProvider(provider=args.provider)
    print(f"Provider: {args.provider}  classifier={llm.classifier_model}  "
          f"reasoning={llm.reasoning_model}")

    transcripts = load_dataset(Path(args.dataset))
    if args.sample:
        transcripts = transcripts[:args.sample]
    print(f"Running on {len(transcripts)} transcripts.")

    # 1. LLM-driven categorization.
    predicted = []
    for t in tqdm(transcripts, desc="Categorize (LLM)"):
        try:
            c = categorize_with_llm(t, llm)
            predicted.append(c)
        except Exception as e:
            print(f"  [{t.meeting_id}] error: {e!r}; falling back to ground-truth projection")
            predicted.append(categorize_from_ground_truth(t))

    # 2. Aggregate + synthesize themes (one reasoning call).
    rollup = aggregate_themes(predicted)
    sample_titles = defaultdict(list)
    for t, c in zip(transcripts, predicted):
        for tag in c.tags:
            if len(sample_titles[tag]) < 3:
                sample_titles[tag].append(t.title)
    synth = synthesize_themes(rollup, dict(sample_titles), llm)

    # 3. Eval pass: compare LLM tags to lexical-projection of ground truth.
    eval_results = None
    if not args.skip_eval:
        reference = [categorize_from_ground_truth(t) for t in transcripts]
        eval_results = categorization_eval(
            predicted=[{"meeting_id": p.meeting_id, "tags": p.tags,
                        "primary_theme": p.primary_theme} for p in predicted],
            reference=[{"meeting_id": r.meeting_id, "tags": r.tags,
                        "primary_theme": r.primary_theme} for r in reference],
        )
        print("\n=== EVAL: LLM-categorizer vs ground-truth projection ===")
        print(f"  Tag F1            = {eval_results['tag_f1']:.3f}")
        print(f"  Tag precision     = {eval_results['tag_precision']:.3f}")
        print(f"  Tag recall        = {eval_results['tag_recall']:.3f}")
        print(f"  Primary-theme acc = {eval_results['primary_theme_accuracy']:.3f}")
        print(f"  N compared        = {eval_results['n_compared']}")

    # 4. Persist LLM outputs.
    OUT.mkdir(parents=True, exist_ok=True)
    (OUT / "llm_categorization.json").write_text(json.dumps({
        "provider": args.provider,
        "classifier_model": llm.classifier_model,
        "reasoning_model": llm.reasoning_model,
        "synthesis": synth,
        "eval": eval_results,
        "predicted": [
            {"meeting_id": p.meeting_id, "title": p.title,
             "tags": p.tags, "primary_theme": p.primary_theme,
             "rationale": p.rationale}
            for p in predicted
        ],
    }, indent=2), encoding="utf-8")
    print(f"\nWrote {OUT / 'llm_categorization.json'}")


if __name__ == "__main__":
    main()