measure_cache.py

"""Measure system-prompt prefix caching benefit on TTFT.

Sends the same long system prompt + varying user message, measures:
1. Cold first call (no cache)
2. Warm subsequent calls (SP prefix cached)
3. Control: short prompt with no SP

Reports per-call wall time + prompt-eval time from the server's timings field.
"""

import json
import time
import urllib.request

URL = "http://127.0.0.1:53658/v1/chat/completions"
SP = (
    "You are Alex, the front-desk assistant for Mike's Plumbing Solutions in South Florida "
    "(Miami-Dade, Broward, Palm Beach). GOAL: Handle only plumbing and plumbing-business questions. "
    "Collect details, schedule follow-up, route emergencies to Mike. STYLE: Warm, brief, confident, "
    "no corporate talk. Never say you are an AI or bot. Keep replies under 3 sentences. "
) * 6  # ~550 tokens


def call(label, messages, n=5):
    times, prompt_ms = [], []
    for i in range(n):
        payload = json.dumps(
            {"model": "x", "stream": False, "messages": messages, "max_tokens": 8, "temperature": 0}
        ).encode()
        t0 = time.time()
        r = urllib.request.urlopen(
            urllib.request.Request(URL, data=payload, headers={"Content-Type": "application/json"}), timeout=60
        )
        d = json.loads(r.read())
        wall = time.time() - t0
        times.append(wall)
        pm = d.get("timings", {}).get("prompt_ms", 0)
        prompt_ms.append(pm)
    avg_wall = sum(times) / len(times)
    avg_pm = sum(prompt_ms) / len(prompt_ms)
    print(f"  {label:35s} wall={avg_wall * 1000:6.0f}ms  prompt_eval={avg_pm:6.1f}ms  (n={n})")
    return avg_wall, avg_pm


print(f"System prompt: ~{len(SP) // 4} tokens\n")
print("Same SP + varying user (measures cache warm-up):")
warm_msgs = [[{"role": "system", "content": SP}, {"role": "user", "content": f"q{i} leak"}] for i in range(1, 7)]
for i, msgs in enumerate(warm_msgs):
    call(f"  call {i + 1} ({'cold' if i == 0 else 'warm'})", msgs, n=1)
print()
print("Repeated identical request (best-case cache):")
call("  same payload x5", [{"role": "system", "content": SP}, {"role": "user", "content": "toilet running"}], n=5)
print()
print("Control: no system prompt:")
call("  short prompt x5", [{"role": "user", "content": "hi"}], n=5)