tatlock/scripts/benchmark_tool_calling.py

"""
Benchmark tool calling across different Ollama models via Tatlock API.

Sends test prompts through the full Tatlock pipeline (Steward -> Orchestration
-> Synthesis) and records tool selection accuracy, latency, and response quality.

Between models, swaps OLLAMA_DEFAULT_MODEL in .env and waits for uvicorn
auto-reload. Requires the server to be running via ./wakeup.sh.

Usage:
    .venv/bin/python scripts/benchmark_tool_calling.py
    .venv/bin/python scripts/benchmark_tool_calling.py --models "gemma4:e4b,gemma4:e2b"
    .venv/bin/python scripts/benchmark_tool_calling.py --iterations 3
"""
import argparse
import asyncio
import json
import re
import statistics
import sys
import time
from dataclasses import dataclass, field
from pathlib import Path

import httpx

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from scripts.ollama_residency import install_sigterm_handler, residency_guard

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

API_BASE = "http://localhost:8777"
CHAT_URL = f"{API_BASE}/v1/chat/completions"
HEALTH_URL = f"{API_BASE}/health"
OLLAMA_URL = "http://localhost:11434"
ENV_PATH = Path(__file__).parent.parent / ".env"

DEFAULT_MODELS = ["mistral-nemo-large:latest", "gemma4:e4b", "gemma4:e2b"]

# ---------------------------------------------------------------------------
# Test scenarios
# ---------------------------------------------------------------------------


@dataclass
class Scenario:
    name: str
    prompt: str
    expected_tool: str | None  # None = no tool expected
    # Patterns to check in the response text for indirect tool-use evidence
    success_patterns: list[str] = field(default_factory=list)
    category: str = "basic"


SCENARIOS = [
    # --- Should call calculate_math ---
    Scenario(
        name="Simple arithmetic",
        prompt="What is 144 divided by 12?",
        expected_tool="calculate_math",
        success_patterns=["12"],
        category="calculator",
    ),
    Scenario(
        name="Square root",
        prompt="What's the square root of 256?",
        expected_tool="calculate_math",
        success_patterns=["16"],
        category="calculator",
    ),
    Scenario(
        name="Complex math",
        prompt="Calculate pi times the square of 5",
        expected_tool="calculate_math",
        success_patterns=["78.5"],  # pi * 25 ≈ 78.54
        category="calculator",
    ),
    Scenario(
        name="Word problem",
        prompt="If I have 3 bags with 17 apples each and I eat 4, how many apples do I have?",
        expected_tool="calculate_math",
        success_patterns=["47"],
        category="calculator",
    ),

    # --- Should call get_current_time ---
    Scenario(
        name="Current date",
        prompt="What's today's date?",
        expected_tool="get_current_time",
        success_patterns=["2026"],  # Should contain current year
        category="datetime",
    ),
    Scenario(
        name="Current time",
        prompt="What time is it right now?",
        expected_tool="get_current_time",
        success_patterns=[":"],  # Time format contains colons
        category="datetime",
    ),

    # --- Should call calculate_date_offset ---
    Scenario(
        name="Relative date past",
        prompt="What was the date 2 weeks ago?",
        expected_tool="calculate_date_offset",
        success_patterns=["2026"],
        category="datetime",
    ),

    # --- Should call calculate_time_difference ---
    Scenario(
        name="Date difference",
        prompt="How many days between January 1st 2025 and March 15th 2025?",
        expected_tool="calculate_time_difference",
        success_patterns=["73", "74"],  # 73 or 74 days
        category="datetime",
    ),

    # --- Should NOT call any tool ---
    Scenario(
        name="Greeting",
        prompt="Hello! How are you?",
        expected_tool=None,
        success_patterns=["sir"],  # Butler personality
        category="no_tool",
    ),
    Scenario(
        name="Knowledge question",
        prompt="What is the capital of France?",
        expected_tool=None,
        success_patterns=["Paris"],
        category="no_tool",
    ),
    Scenario(
        name="Opinion request",
        prompt="What do you think about rainy days?",
        expected_tool=None,
        category="no_tool",
    ),
]

# ---------------------------------------------------------------------------
# Result tracking
# ---------------------------------------------------------------------------


@dataclass
class RunResult:
    scenario: str
    model: str
    iteration: int
    latency: float
    response_text: str
    has_correct_answer: bool
    error: str | None = None


@dataclass
class ModelStats:
    model: str
    results: list[RunResult] = field(default_factory=list)

    @property
    def total(self) -> int:
        return len(self.results)

    @property
    def errors(self) -> int:
        return sum(1 for r in self.results if r.error)

    @property
    def accuracy(self) -> float:
        valid = [r for r in self.results if not r.error]
        if not valid:
            return 0
        return sum(1 for r in valid if r.has_correct_answer) / len(valid) * 100

    @property
    def avg_latency(self) -> float:
        lats = [r.latency for r in self.results if not r.error]
        return statistics.mean(lats) if lats else 0

    @property
    def p95_latency(self) -> float:
        lats = sorted(r.latency for r in self.results if not r.error)
        if not lats:
            return 0
        return lats[min(int(len(lats) * 0.95), len(lats) - 1)]

    @property
    def max_latency(self) -> float:
        lats = [r.latency for r in self.results if not r.error]
        return max(lats) if lats else 0

    def category_accuracy(self, category: str) -> float:
        cat_scenarios = {s.name for s in SCENARIOS if s.category == category}
        valid = [r for r in self.results if not r.error and r.scenario in cat_scenarios]
        if not valid:
            return 0
        return sum(1 for r in valid if r.has_correct_answer) / len(valid) * 100


# ---------------------------------------------------------------------------
# .env manipulation
# ---------------------------------------------------------------------------


def swap_model_in_env(model_name: str):
    """Swap OLLAMA_DEFAULT_MODEL in .env file."""
    content = ENV_PATH.read_text()
    content = re.sub(
        r'^OLLAMA_DEFAULT_MODEL=.*$',
        f'OLLAMA_DEFAULT_MODEL={model_name}',
        content,
        flags=re.MULTILINE,
    )
    ENV_PATH.write_text(content)
    print(f"  .env updated: OLLAMA_DEFAULT_MODEL={model_name}")


async def wait_for_server_reload(client: httpx.AsyncClient, timeout: float = 30):
    """Wait for uvicorn to auto-reload after .env change."""
    # Give uvicorn a moment to detect the file change
    await asyncio.sleep(3)

    # Poll health endpoint
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        try:
            r = await client.get(HEALTH_URL, timeout=5)
            if r.status_code == 200:
                return
        except Exception:
            pass
        await asyncio.sleep(1)

    raise TimeoutError("Server did not come back after reload")


async def warm_up_ollama_model(client: httpx.AsyncClient, model_name: str):
    """Send a throwaway request to load the model into VRAM."""
    print(f"  Warming up {model_name} in Ollama...", end=" ", flush=True)
    try:
        r = await client.post(
            f"{OLLAMA_URL}/api/generate",
            json={"model": model_name, "prompt": "hi", "stream": False},
            timeout=120,
        )
        r.raise_for_status()
        duration = r.json().get("total_duration", 0) / 1e9
        print(f"OK ({duration:.1f}s)")
    except Exception as e:
        print(f"WARN: {e}")


# ---------------------------------------------------------------------------
# Core benchmark logic
# ---------------------------------------------------------------------------


async def run_scenario(
    client: httpx.AsyncClient,
    scenario: Scenario,
    model: str,
    iteration: int,
) -> RunResult:
    """Run a single scenario through the Tatlock API."""
    payload = {
        "model": "Tatlock",
        "messages": [{"role": "user", "content": scenario.prompt}],
    }

    start = time.monotonic()
    try:
        r = await client.post(CHAT_URL, json=payload, timeout=120)
        latency = time.monotonic() - start

        if r.status_code != 200:
            return RunResult(
                scenario=scenario.name,
                model=model,
                iteration=iteration,
                latency=latency,
                response_text="",
                has_correct_answer=False,
                error=f"HTTP {r.status_code}: {r.text[:100]}",
            )

        data = r.json()
        response_text = data["choices"][0]["message"]["content"]

        # Check if the response contains expected patterns
        has_correct = True
        if scenario.success_patterns:
            has_correct = any(
                p.lower() in response_text.lower()
                for p in scenario.success_patterns
            )

        return RunResult(
            scenario=scenario.name,
            model=model,
            iteration=iteration,
            latency=latency,
            response_text=response_text,
            has_correct_answer=has_correct,
        )

    except Exception as e:
        latency = time.monotonic() - start
        return RunResult(
            scenario=scenario.name,
            model=model,
            iteration=iteration,
            latency=latency,
            response_text="",
            has_correct_answer=False,
            error=str(e)[:200],
        )


async def benchmark_model(
    client: httpx.AsyncClient,
    model_name: str,
    iterations: int,
) -> ModelStats:
    """Run all scenarios for a single model."""
    stats = ModelStats(model=model_name)

    print(f"\n{'=' * 70}")
    print(f"  Model: {model_name}")
    print(f"{'=' * 70}")

    # Swap model in .env
    swap_model_in_env(model_name)

    # Warm up model in Ollama BEFORE server reload picks it up
    await warm_up_ollama_model(client, model_name)

    # Wait for server to reload with new model
    print("  Waiting for server reload...", end=" ", flush=True)
    await wait_for_server_reload(client)
    print("OK")

    # Run a throwaway request through the full pipeline to warm up
    print("  Warming up pipeline...", end=" ", flush=True)
    try:
        await client.post(
            CHAT_URL,
            json={"model": "Tatlock", "messages": [{"role": "user", "content": "hi"}]},
            timeout=120,
        )
        print("OK")
    except Exception as e:
        print(f"WARN: {e}")

    for iteration in range(iterations):
        if iterations > 1:
            print(f"\n  --- Iteration {iteration + 1}/{iterations} ---")

        for scenario in SCENARIOS:
            result = await run_scenario(client, scenario, model_name, iteration)
            stats.results.append(result)

            # Display
            if result.error:
                print(
                    f"  [ERR ] {scenario.name:30s} {result.latency:5.1f}s  "
                    f"{result.error[:60]}"
                )
            elif result.has_correct_answer:
                preview = result.response_text[:60].replace("\n", " ")
                print(f"  [OK  ] {scenario.name:30s} {result.latency:5.1f}s  {preview}")
            else:
                preview = result.response_text[:60].replace("\n", " ")
                print(f"  [MISS] {scenario.name:30s} {result.latency:5.1f}s  {preview}")

    return stats


def print_comparison(all_stats: list[ModelStats]):
    """Print side-by-side comparison table."""
    print("\n" + "=" * 80)
    print("  COMPARISON SUMMARY")
    print("=" * 80)

    col_width = max(len(s.model) for s in all_stats) + 2
    label_width = 32

    header = f"{'Metric':<{label_width}}"
    for s in all_stats:
        header += f"  {s.model:>{col_width}}"
    print(f"\n{header}")
    print("-" * (label_width + (col_width + 2) * len(all_stats)))

    # Answer accuracy
    row = f"{'Correct answer rate':<{label_width}}"
    for s in all_stats:
        row += f"  {s.accuracy:>{col_width - 1}.1f}%"
    print(row)

    # Latency
    row = f"{'Avg latency':<{label_width}}"
    for s in all_stats:
        row += f"  {s.avg_latency:>{col_width - 1}.1f}s"
    print(row)

    row = f"{'P95 latency':<{label_width}}"
    for s in all_stats:
        row += f"  {s.p95_latency:>{col_width - 1}.1f}s"
    print(row)

    row = f"{'Max latency':<{label_width}}"
    for s in all_stats:
        row += f"  {s.max_latency:>{col_width - 1}.1f}s"
    print(row)

    # Errors
    row = f"{'Errors':<{label_width}}"
    for s in all_stats:
        row += f"  {s.errors:>{col_width}}"
    print(row)

    # Per-category
    categories = sorted(set(sc.category for sc in SCENARIOS))
    print(f"\n{'Per-category accuracy':<{label_width}}")
    print("-" * (label_width + (col_width + 2) * len(all_stats)))
    for cat in categories:
        row = f"  {cat:<{label_width - 2}}"
        for s in all_stats:
            row += f"  {s.category_accuracy(cat):>{col_width - 1}.1f}%"
        print(row)

    # Mismatches
    print(f"\n{'Missed answers':<50}")
    print("-" * 80)
    any_miss = False
    for scenario in SCENARIOS:
        misses = []
        for s in all_stats:
            sc_results = [r for r in s.results if r.scenario == scenario.name]
            fails = [r for r in sc_results if not r.has_correct_answer and not r.error]
            if fails:
                preview = fails[0].response_text[:50].replace("\n", " ")
                misses.append(f"{s.model}: \"{preview}\"")
        if misses:
            any_miss = True
            print(f"  {scenario.name}")
            for m in misses:
                print(f"    {m}")

    if not any_miss:
        print("  (none)")

    print("\n" + "=" * 80)


def save_results(all_stats: list[ModelStats], output_path: Path):
    """Save detailed results to JSON."""
    data = {}
    for stats in all_stats:
        data[stats.model] = {
            "summary": {
                "accuracy": stats.accuracy,
                "avg_latency": round(stats.avg_latency, 2),
                "p95_latency": round(stats.p95_latency, 2),
                "max_latency": round(stats.max_latency, 2),
                "errors": stats.errors,
                "total_runs": stats.total,
            },
            "runs": [
                {
                    "scenario": r.scenario,
                    "iteration": r.iteration,
                    "latency": round(r.latency, 3),
                    "has_correct_answer": r.has_correct_answer,
                    "response_text": r.response_text,
                    "error": r.error,
                }
                for r in stats.results
            ],
        }

    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text(json.dumps(data, indent=2))
    print(f"\nDetailed results saved to: {output_path}")


async def main():
    parser = argparse.ArgumentParser(description="Benchmark tool calling across Ollama models via Tatlock API")
    parser.add_argument(
        "--iterations", type=int, default=1,
        help="Iterations per model (default: 1)",
    )
    parser.add_argument(
        "--models", type=str, default=",".join(DEFAULT_MODELS),
        help=f"Comma-separated models (default: {','.join(DEFAULT_MODELS)})",
    )
    parser.add_argument(
        "--output", type=str, default="logs/benchmark_results.json",
        help="JSON output path (default: logs/benchmark_results.json)",
    )
    args = parser.parse_args()

    models = [m.strip() for m in args.models.split(",")]

    # Verify server is running
    async with httpx.AsyncClient() as client:
        try:
            r = await client.get(HEALTH_URL, timeout=5)
            r.raise_for_status()
            print("Server is running.")
        except Exception:
            print("ERROR: Server not running. Start it with ./wakeup.sh first.")
            return

    print("=" * 70)
    print("  Tool Calling Benchmark (via Tatlock API)")
    print("=" * 70)
    print(f"  Models:     {', '.join(models)}")
    print(f"  Scenarios:  {len(SCENARIOS)}")
    print(f"  Iterations: {args.iterations}")
    print(f"  Total runs: {len(SCENARIOS) * args.iterations * len(models)}")

    # Remember original model to restore after benchmark
    original_env = ENV_PATH.read_text()

    all_stats = []
    # Both restores must survive a crash or an interrupt. The .env one especially:
    # this script rewrites OLLAMA_DEFAULT_MODEL and lets uvicorn reload onto it,
    # so bailing out mid-run used to leave the *running server* pointed at the
    # benchmark model — and DEFAULT_MODELS starts at mistral-nemo-large, the 9.2G
    # model implicated in the 2026-08-07 VRAM outage.
    install_sigterm_handler()
    try:
        with residency_guard(models_used=models):
            async with httpx.AsyncClient() as client:
                for model in models:
                    stats = await benchmark_model(client, model, args.iterations)
                    all_stats.append(stats)
    finally:
        ENV_PATH.write_text(original_env)
        print("\n  .env restored to original")

    print_comparison(all_stats)
    save_results(all_stats, Path(args.output))


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        # .env and GPU residency are both restored by now; do not bury that
        # output under a traceback.
        print("\ninterrupted", file=sys.stderr)
        raise SystemExit(130) from None