scripts/benchmark_model.py

"""Measure whether the local model is usable on this machine.

The agent makes two model calls per conversational turn (intent parse, then
response generation), so per-turn latency is roughly the sum of the two. Run:

    uv run scripts/benchmark_model.py
"""

import time

from flight_agent.config import MODEL_NAME
from flight_agent.llm import LocalQwen

PROMPTS = [
    "Find flights from JFK to LAX on 2026-09-15",
    "Show my bookings",
    "Remember that my home airport is Boston",
]

CONTEXT = {
    "flight_results": [
        {"id": "F000101", "origin": "JFK", "destination": "LAX",
         "departure_date": "2026-09-15", "departure_time": "08:20",
         "arrival_time": "11:25", "base_price": 149.0},
    ],
    "preferences": {"home_airport": "BOS"},
    "intent": {"intent": "search_flights"},
}


def main() -> None:
    print(f"Model: {MODEL_NAME}")
    t0 = time.perf_counter()
    llm = LocalQwen()
    load_s = time.perf_counter() - t0
    print(f"Load: {load_s:.1f}s\n")

    tok = llm.tokenizer
    intent_times = []

    print("--- intent parsing ---")
    for prompt in PROMPTS:
        t0 = time.perf_counter()
        intent = llm.parse_intent(prompt)
        dt = time.perf_counter() - t0
        intent_times.append(dt)
        print(f"  {dt:6.1f}s  {prompt!r}")
        print(f"           -> {intent.intent} "
              f"{intent.origin or ''}->{intent.destination or ''} {intent.date or ''}")

    print("\n--- response generation ---")
    t0 = time.perf_counter()
    text = llm.respond(CONTEXT)
    respond_s = time.perf_counter() - t0
    n_tokens = len(tok(text)["input_ids"])
    print(f"  {respond_s:.1f}s for {n_tokens} tokens "
          f"({n_tokens / respond_s:.1f} tok/s)")
    print(f"  {text[:400]!r}")

    avg_intent = sum(intent_times) / len(intent_times)
    print("\n--- summary ---")
    print(f"  model load          {load_s:6.1f}s (once per process)")
    print(f"  intent parse (avg)  {avg_intent:6.1f}s")
    print(f"  response generate   {respond_s:6.1f}s")
    print(f"  => per chat turn    {avg_intent + respond_s:6.1f}s")


if __name__ == "__main__":
    main()