test_endpoint.py

"""Diagnostic: confirm the LLM endpoint works and measure latency.

Run it with the same environment variables as the app:

    export MODEL_BASE_URL="https://....modal.run/v1"
    export MODEL_API_KEY="sk-sokrates-demo-123"
    export MODEL_NAME="Qwen/Qwen3-14B"
    export SOKRATES_GUIDED_JSON=1
    export SOKRATES_NO_THINK=1
    python test_endpoint.py

It prints, step by step:
  - the resolved config,
  - whether the server lists the model,
  - the extraction call + how long it took,
  - the question-generation call + how long it took.

If the first call is slow (tens of seconds to minutes) but the second is fast,
that's a cold start — the model is loading on the GPU. Keep one replica warm:
    SOKRATES_KEEP_WARM=1 modal deploy modal_vllm.py
"""

import time

from sokrates.gaps import missing_mandatory_labels
from sokrates.llm import generate_questions, get_client, get_config, extract_form

TRANSCRIPT = (
    "Doctor: Good morning, how old are you?\n"
    "Patient: I'm 62, and I'm a man.\n"
    "Doctor: What brings you in?\n"
    "Patient: I've had a cough for three weeks and lost some weight. "
    "I had a chest CT that showed a mass in my right lung, but no biopsy yet."
)


def main() -> None:
    base_url, _key, model = get_config()
    print(f"• Endpoint : {base_url}")
    print(f"• Model    : {model}\n")

    print("→ Checking the server lists the model …")
    t0 = time.time()
    models = get_client().models.list()
    served = [m.id for m in models.data]
    print(f"  OK in {time.time() - t0:.1f}s — served models: {served}\n")

    print("→ Extraction call (fills the form) …")
    t0 = time.time()
    form = extract_form(TRANSCRIPT)
    print(f"  Done in {time.time() - t0:.1f}s")
    print(f"  Extracted: {form.model_dump(exclude_none=True)}\n")

    missing = missing_mandatory_labels(form)
    print(f"  Still-missing mandatory fields: {missing}\n")

    print("→ Question-generation call …")
    t0 = time.time()
    questions = generate_questions(TRANSCRIPT, missing)
    print(f"  Done in {time.time() - t0:.1f}s")
    for i, q in enumerate(questions, 1):
        print(f"   {i}. {q}")

    print("\n✅ Endpoint works end to end.")


if __name__ == "__main__":
    main()