Reproduce the Lydia/Jake theory-of-mind translation experiment with OpenRouter and LiteLLM

run_experiment.py

#!/usr/bin/env python3
"""Reproduce the Lydia/Jake theory-of-mind translation experiment.

Usage:
    export OPENROUTER_API_KEY="..."
    python theory-of-mind-openrouter/run_experiment.py

Install dependency:
    python -m pip install litellm

The script calls OpenRouter through LiteLLM, writes raw outputs to JSONL, and
creates CSV and Markdown summaries of the pronoun used in Jake's final line.
"""

from __future__ import annotations

import argparse
import csv
import json
import os
import re
import sys
import time
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any


STORY = '''Lydia fölnézett a csillagos égre. Imádta, ha a hosszú haja lobog utána a szélben, de ezen a vidéken még nem járt, és nem akart ujjat húzni a helyi rendőrökkel. Minden tincsét gondosan a sisakja alá gyűrte. Megnyugtatta a Harleyjának dörmögése. Nem gépként, hanem érző lényként gondolt rá, akkor is, amikor lehúzodott az útról és begurult a kúthoz, hogy tankoljon. 

Jake a motorzúgásra ébredt. Gyakran elaludt a pénztár mögött. Fáradt volt már ehhez a munkához. Végigsimította a szakállát és ránézett a biztonsági kamera képére. Semmit nem látott a motor fényétől, ami pont a kamerába világított. "Istenem" - mormogott - "legalább a motorját leállíthatta volna."'''

PROMPT = f'''Translate the following story into English. No questions or commentary, just give me the text in English.

"""{STORY}
"""'''


@dataclass(frozen=True)
class ModelSpec:
    label: str
    slug: str

    @property
    def litellm_model(self) -> str:
        return f"openrouter/{self.slug}"


MODELS: tuple[ModelSpec, ...] = (
    # Dynamic OpenRouter aliases use a leading '~'. These are listed by
    # models.dev under the openrouter provider and may resolve to newer backing
    # model versions over time; the raw response metadata records what ran.
    ModelSpec("Claude Opus Latest", "~anthropic/claude-opus-latest"),
    ModelSpec("DeepSeek R1 0528", "deepseek/deepseek-r1-0528"),
    ModelSpec("DeepSeek V3.2", "deepseek/deepseek-v3.2"),
    ModelSpec("DeepSeek V4 Pro", "deepseek/deepseek-v4-pro"),
    ModelSpec("Gemini Flash Latest", "~google/gemini-flash-latest"),
    ModelSpec("Gemini Pro Latest", "~google/gemini-pro-latest"),
    ModelSpec("GLM 5.1", "z-ai/glm-5.1"),
    ModelSpec("Grok 4.20", "x-ai/grok-4.20"),
    ModelSpec("Kimi Latest", "~moonshotai/kimi-latest"),
    ModelSpec("MiniMax M3", "minimax/minimax-m3"),
    ModelSpec("MiMo-V2.5-Pro", "xiaomi/mimo-v2.5-pro"),
    ModelSpec("Mistral Medium 3.5", "mistralai/mistral-medium-3-5"),
    ModelSpec("Mistral Nemo", "mistralai/mistral-nemo"),
    ModelSpec("OpenAI GPT Latest", "~openai/gpt-latest"),
    ModelSpec("OpenAI o4 Mini High", "openai/o4-mini-high"),
    ModelSpec("Qwen3.7 Plus", "qwen/qwen3.7-plus"),
    ModelSpec("Step 3.7 Flash", "stepfun/step-3.7-flash"),
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Run the Lydia/Jake translation experiment via OpenRouter."
    )
    parser.add_argument(
        "--output-dir",
        default="",
        help="Directory for JSONL, CSV, and Markdown outputs. Default: ./results next to this script.",
    )
    parser.add_argument(
        "--temperature",
        type=float,
        default=0.0,
        help="Sampling temperature. Default: 0 for reproducibility.",
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=None,
        help="Optional seed. OpenRouter/provider support varies by model.",
    )
    parser.add_argument(
        "--timeout",
        type=float,
        default=120.0,
        help="Per-model timeout in seconds.",
    )
    parser.add_argument(
        "--only",
        default="",
        help="Comma-separated labels or OpenRouter slugs to run instead of all models.",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Print the prompt and model list without calling the API.",
    )
    return parser.parse_args()


def selected_models(only: str) -> list[ModelSpec]:
    if not only.strip():
        return list(MODELS)
    wanted = {item.strip().lower() for item in only.split(",") if item.strip()}
    models = [
        model
        for model in MODELS
        if model.label.lower() in wanted or model.slug.lower() in wanted
    ]
    missing = wanted - {model.label.lower() for model in models} - {
        model.slug.lower() for model in models
    }
    if missing:
        raise SystemExit(f"Unknown --only value(s): {', '.join(sorted(missing))}")
    return models


def classify_pronoun(text: str) -> tuple[str, str]:
    """Classify the pronoun in Jake's translated final utterance.

    Returns (classification, evidence). The heuristic intentionally focuses on
    the last quoted sentence, because earlier parts of the story correctly refer
    to Lydia as she/her.
    """
    normalized = (
        text.replace("“", '"')
        .replace("”", '"')
        .replace("‘", "'")
        .replace("’", "'")
    )
    quoted = re.findall(r'"([^"\n]+)"', normalized)
    evidence = quoted[-1] if quoted else normalized[-300:]
    evidence_lower = evidence.lower()

    patterns = (
        ("he", r"\b(he|him|his)\b"),
        ("she", r"\b(she|her|hers)\b"),
        ("they", r"\b(they|them|their|theirs)\b"),
    )
    matches = [label for label, pattern in patterns if re.search(pattern, evidence_lower)]
    if len(matches) == 1:
        return matches[0], evidence
    if len(matches) > 1:
        return "mixed", evidence
    return "omitted_or_other", evidence


def response_to_dict(response: Any) -> dict[str, Any]:
    if hasattr(response, "model_dump"):
        return response.model_dump()
    if hasattr(response, "dict"):
        return response.dict()
    return json.loads(json.dumps(response, default=str))


def extract_text(response_dict: dict[str, Any]) -> str:
    try:
        message = response_dict["choices"][0]["message"]
    except (KeyError, IndexError, TypeError) as exc:
        raise ValueError("Unexpected LiteLLM/OpenRouter response shape") from exc
    content = message.get("content", "")
    if isinstance(content, str):
        return content.strip()
    return json.dumps(content, ensure_ascii=False)


def call_model(
    completion: Any,
    model: ModelSpec,
    temperature: float,
    seed: int | None,
    timeout: float,
) -> dict[str, Any]:
    started = time.time()
    kwargs: dict[str, Any] = {
        "model": model.litellm_model,
        "messages": [{"role": "user", "content": PROMPT}],
        "temperature": temperature,
        "timeout": timeout,
    }
    if seed is not None:
        kwargs["seed"] = seed

    response = completion(**kwargs)
    response_dict = response_to_dict(response)
    text = extract_text(response_dict)
    pronoun, evidence = classify_pronoun(text)
    return {
        "ok": True,
        "label": model.label,
        "slug": model.slug,
        "litellm_model": model.litellm_model,
        "response_model": response_dict.get("model"),
        "pronoun": pronoun,
        "evidence": evidence,
        "output_text": text,
        "latency_seconds": round(time.time() - started, 3),
        "raw_response": response_dict,
    }


def write_jsonl(path: Path, records: list[dict[str, Any]]) -> None:
    with path.open("w", encoding="utf-8") as handle:
        for record in records:
            handle.write(json.dumps(record, ensure_ascii=False, default=str) + "\n")


def write_csv(path: Path, records: list[dict[str, Any]]) -> None:
    fields = [
        "ok",
        "label",
        "slug",
        "response_model",
        "pronoun",
        "evidence",
        "latency_seconds",
        "error",
    ]
    with path.open("w", encoding="utf-8", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=fields, extrasaction="ignore")
        writer.writeheader()
        writer.writerows(records)


def write_markdown(path: Path, records: list[dict[str, Any]], metadata: dict[str, Any]) -> None:
    counts: dict[str, int] = {}
    for record in records:
        key = record.get("pronoun", "error") if record.get("ok") else "error"
        counts[key] = counts.get(key, 0) + 1

    lines = [
        "# Theory-of-Mind Translation Experiment Results",
        "",
        f"Run timestamp: `{metadata['timestamp_utc']}`",
        f"Temperature: `{metadata['temperature']}`",
        f"Seed: `{metadata['seed']}`",
        "",
        "## Prompt",
        "",
        "```",
        PROMPT,
        "```",
        "",
        "## Summary Counts",
        "",
    ]
    for key in sorted(counts):
        lines.append(f"- `{key}`: {counts[key]}")

    lines.extend(
        [
            "",
            "## Model Results",
            "",
            "| Model | Pronoun | Evidence |",
            "| --- | --- | --- |",
        ]
    )
    for record in records:
        if not record.get("ok"):
            pronoun = "error"
            evidence = str(record.get("error", ""))
        else:
            pronoun = str(record.get("pronoun", ""))
            evidence = str(record.get("evidence", "")).replace("|", "\\|")
        lines.append(f"| {record['label']} | {pronoun} | {evidence} |")

    path.write_text("\n".join(lines) + "\n", encoding="utf-8")


def main() -> int:
    args = parse_args()
    models = selected_models(args.only)
    timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
    output_dir = Path(args.output_dir) if args.output_dir else Path(__file__).parent / "results"
    output_dir.mkdir(parents=True, exist_ok=True)

    metadata = {
        "timestamp_utc": timestamp,
        "temperature": args.temperature,
        "seed": args.seed,
        "models": [model.__dict__ for model in models],
    }

    if args.dry_run:
        print(PROMPT)
        print("\nModels:")
        for model in models:
            print(f"- {model.label}: {model.litellm_model}")
        return 0

    if not os.environ.get("OPENROUTER_API_KEY"):
        print("OPENROUTER_API_KEY is not set.", file=sys.stderr)
        return 2

    try:
        from litellm import completion
    except ImportError:
        print("Missing dependency: install with `python -m pip install litellm`.", file=sys.stderr)
        return 2

    records: list[dict[str, Any]] = []
    for index, model in enumerate(models, start=1):
        print(f"[{index}/{len(models)}] {model.label} ({model.slug})", flush=True)
        try:
            record = call_model(
                completion=completion,
                model=model,
                temperature=args.temperature,
                seed=args.seed,
                timeout=args.timeout,
            )
        except Exception as exc:  # noqa: BLE001 - record provider failures and continue.
            record = {
                "ok": False,
                "label": model.label,
                "slug": model.slug,
                "litellm_model": model.litellm_model,
                "pronoun": "error",
                "evidence": "",
                "output_text": "",
                "latency_seconds": None,
                "error": repr(exc),
            }
            print(f"  ERROR: {exc}", file=sys.stderr, flush=True)
        else:
            print(f"  pronoun={record['pronoun']} evidence={record['evidence']!r}", flush=True)
        records.append(record)

    base = output_dir / f"run_{timestamp}"
    (base.with_suffix(".metadata.json")).write_text(
        json.dumps(metadata, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )
    write_jsonl(base.with_suffix(".jsonl"), records)
    write_csv(base.with_suffix(".csv"), records)
    write_markdown(base.with_suffix(".md"), records, metadata)

    print(f"\nWrote {base.with_suffix('.jsonl')}")
    print(f"Wrote {base.with_suffix('.csv')}")
    print(f"Wrote {base.with_suffix('.md')}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())