import time, json, httpx, statistics
from app.config import settings
from app.database import SessionLocal
from sqlalchemy import text as sa_text

db = SessionLocal()
stems = [r[0][:700] for r in db.execute(sa_text(
    "select question_text from questions where question_text is not null order by id limit 200")).fetchall()]
base = settings.LITELLM_API_BASE.rstrip('/').removesuffix('/v1')
key = settings.LITELLM_API_KEY

def call(model, docs, q="what is the workup for a febrile seizure in a toddler"):
    t = time.perf_counter()
    r = httpx.post(f"{base}/v1/rerank",
                   headers={"Authorization": f"Bearer {key}"},
                   json={"model": model, "query": q, "documents": docs, "top_n": len(docs)},
                   timeout=30)
    ms = (time.perf_counter()-t)*1000
    r.raise_for_status()
    return ms, r.json()

for model in ["jina-reranker-v2-base-multilingual", "cohere-rerank-v4.0-fast", "cohere-rerank-v4.0-pro"]:
    for n in (10, 30, 50, 100):
        docs = stems[:n]
        times = []
        for i in range(3):
            # vary the query so the proxy's own cache does not answer
            try:
                ms, _ = call(model, docs, f"what is the workup for a febrile seizure in a toddler #{i}{n}")
            except Exception as e:
                print(model, n, "ERR", str(e)[:120]); break
            times.append(ms)
        if times:
            print(f"{model:38s} n={n:3d}  {min(times):7.0f} / {statistics.median(times):7.0f} / {max(times):7.0f} ms")
avg = sum(len(s) for s in stems)/len(stems)
print("avg stem chars", round(avg))