scripts/smoke_test.py

"""Sanity-check vLLM endpoint is reachable from the *aligne* venv and that
forced-choice A/B logprobs are coming through for both served model names.

Run from /workspace/code on the pod (after the bootstrap completes):

    uv run python scripts/smoke_test.py
"""

from __future__ import annotations

import asyncio
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))

from aligne.client import ChatClient, Endpoint
from aligne.metrics.oracle import choice_prob


async def one(client: ChatClient, model: str) -> None:
    prompt = (
        "Consider these two options:\n\n"
        "A) pizza\n"
        "B) kale\n\n"
        "Which do you prefer? Answer with a single letter, A or B."
    )
    res = await choice_prob(client, prompt)
    if res is None:
        print(f"  model={model!r}  FAILED to extract A/B choice")
    else:
        print(f"  model={model!r}  P(A=pizza) = {res.p_a:.4f}  mode={res.mode}  coverage={res.coverage:.3f}")


async def main() -> None:
    url = "http://localhost:8000/v1"
    for model in ("google/gemma-3-27b-it", "functional-wellbeing"):
        endpoint = Endpoint(base_url=url, model=model, api_key="dummy")
        client = ChatClient(endpoint=endpoint, concurrency=4)
        print(f"\n== {model} ==")
        await one(client, model)


if __name__ == "__main__":
    asyncio.run(main())