tifkin/NextTokenPrediction

1import httpx2import pytest3 4import backend.main as main5import next_token_prediction.catalog as catalog6from next_token_prediction import ModelUnavailableError, UnknownModelError7 8 9@pytest.fixture10def anyio_backend():11    return "asyncio"12 13 14def fake_prediction(prompt, settings, model_id="test-model"):15    return {16        "model_id": model_id,17        "supports_reasoning": False,18        "prompt": prompt,19        "settings": {20            "temperature": settings.temperature,21            "top_k": settings.top_k,22            "top_p": settings.top_p,23            "filter_mode": settings.filter_mode,24        },25        "is_empty": prompt == "",26        "sampled_token": None,27        "token_spans": [],28        "reasoning_token_spans": [],29        "summary": {"kept_count": 0, "vocab_size": 0, "retained_mass": 0.0},30        "bars": [],31        "table_rows": [],32        "nucleus_rows": [],33    }34 35 36@pytest.mark.anyio37async def test_health_returns_model_id():38    async with httpx.AsyncClient(transport=httpx.ASGITransport(app=main.app), base_url="http://testserver") as client:39        response = await client.get("/api/health")40 41    assert response.status_code == 20042    assert response.json()["ok"] is True43    assert "model_id" in response.json()44 45 46@pytest.mark.anyio47async def test_models_returns_model_catalog(monkeypatch):48    monkeypatch.delenv(catalog.PHI4_GGUF_PREFETCH_ENV, raising=False)49 50    async with httpx.AsyncClient(transport=httpx.ASGITransport(app=main.app), base_url="http://testserver") as client:51        response = await client.get("/api/models")52 53    assert response.status_code == 20054    body = response.json()55    assert body["default_model_id"] == "HuggingFaceTB/SmolLM2-135M"56    assert any(model["model_id"] == "Phi-4-mini-reasoning-Q2_K.gguf" for model in body["models"])57    assert any(model["model_id"] == "Qwen3-0.6B-Q8_0.gguf" for model in body["models"])58 59 60@pytest.mark.anyio61async def test_models_keeps_phi_when_prefetch_is_disabled(monkeypatch):62    monkeypatch.setenv(catalog.PHI4_GGUF_PREFETCH_ENV, "0")63 64    async with httpx.AsyncClient(transport=httpx.ASGITransport(app=main.app), base_url="http://testserver") as client:65        response = await client.get("/api/models")66 67    assert response.status_code == 20068    assert [model["model_id"] for model in response.json()["models"]] == [69        catalog.MODEL_ID,70        catalog.PHI4_GGUF_MODEL_ID,71        catalog.QWEN3_GGUF_MODEL_ID,72    ]73 74 75@pytest.mark.anyio76async def test_predict_delegates_to_model_layer(monkeypatch):77    monkeypatch.setattr(main, "predict_prompt", fake_prediction)78 79    async with httpx.AsyncClient(transport=httpx.ASGITransport(app=main.app), base_url="http://testserver") as client:80        response = await client.post(81            "/api/predict",82            json={"prompt": "Hello", "temperature": 0.5, "top_k": 10, "top_p": 0.8},83        )84 85    assert response.status_code == 20086    body = response.json()87    assert body["prompt"] == "Hello"88    assert body["model_id"] == "HuggingFaceTB/SmolLM2-135M"89    assert body["settings"] == {"temperature": 0.5, "top_k": 0, "top_p": 0.8, "filter_mode": "top_p"}90 91 92@pytest.mark.anyio93async def test_predict_passes_selected_model_to_model_layer(monkeypatch):94    monkeypatch.setattr(main, "predict_prompt", fake_prediction)95 96    async with httpx.AsyncClient(transport=httpx.ASGITransport(app=main.app), base_url="http://testserver") as client:97        response = await client.post(98            "/api/predict",99            json={"prompt": "Hello", "model_id": "Phi-4-mini-reasoning-Q2_K.gguf"},100        )101 102    assert response.status_code == 200103    assert response.json()["model_id"] == "Phi-4-mini-reasoning-Q2_K.gguf"104 105 106@pytest.mark.anyio107async def test_predict_returns_not_found_for_unknown_models(monkeypatch):108    def unknown_model(prompt, settings, model_id):109        raise UnknownModelError(f"Unknown model '{model_id}'.")110 111    monkeypatch.setattr(main, "predict_prompt", unknown_model)112 113    async with httpx.AsyncClient(transport=httpx.ASGITransport(app=main.app), base_url="http://testserver") as client:114        response = await client.post("/api/predict", json={"prompt": "Hello", "model_id": "missing"})115 116    assert response.status_code == 404117 118 119@pytest.mark.anyio120async def test_predict_returns_unavailable_for_missing_local_models(monkeypatch):121    def unavailable_model(prompt, settings, model_id):122        raise ModelUnavailableError("Run uv sync --extra gguf.")123 124    monkeypatch.setattr(main, "predict_prompt", unavailable_model)125 126    async with httpx.AsyncClient(transport=httpx.ASGITransport(app=main.app), base_url="http://testserver") as client:127        response = await client.post("/api/predict", json={"prompt": "Hello", "model_id": "phi"})128 129    assert response.status_code == 503130 131 132@pytest.mark.anyio133async def test_predict_uses_only_selected_filter(monkeypatch):134    monkeypatch.setattr(main, "predict_prompt", fake_prediction)135 136    async with httpx.AsyncClient(transport=httpx.ASGITransport(app=main.app), base_url="http://testserver") as client:137        response = await client.post(138            "/api/predict",139            json={"prompt": "Hello", "temperature": 0.5, "top_k": 10, "top_p": 0.2, "filter_mode": "top_k"},140        )141 142    assert response.status_code == 200143    body = response.json()144    assert body["settings"] == {"temperature": 0.5, "top_k": 10, "top_p": 1.0, "filter_mode": "top_k"}145 146 147@pytest.mark.anyio148async def test_predict_allows_zero_top_k_to_disable_filter(monkeypatch):149    monkeypatch.setattr(main, "predict_prompt", fake_prediction)150 151    async with httpx.AsyncClient(transport=httpx.ASGITransport(app=main.app), base_url="http://testserver") as client:152        response = await client.post(153            "/api/predict",154            json={"prompt": "Hello", "temperature": 0.5, "top_k": 0, "top_p": 0.2, "filter_mode": "top_k"},155        )156 157    assert response.status_code == 200158    body = response.json()159    assert body["settings"] == {"temperature": 0.5, "top_k": 0, "top_p": 1.0, "filter_mode": "top_k"}160 161 162@pytest.mark.anyio163async def test_sample_delegates_to_model_layer(monkeypatch):164    def fake_sample(prompt, settings, model_id):165        result = fake_prediction(prompt + "!", settings, model_id)166        result["sampled_token"] = {"token_id": 1, "raw": "!", "text": "!"}167        return result168 169    monkeypatch.setattr(main, "sample_prompt", fake_sample)170 171    async with httpx.AsyncClient(transport=httpx.ASGITransport(app=main.app), base_url="http://testserver") as client:172        response = await client.post("/api/sample", json={"prompt": "Hello"})173 174    assert response.status_code == 200175    body = response.json()176    assert body["prompt"] == "Hello!"177    assert body["sampled_token"]["raw"] == "!"178 179 180@pytest.mark.anyio181async def test_invalid_sampling_settings_return_validation_error():182    async with httpx.AsyncClient(transport=httpx.ASGITransport(app=main.app), base_url="http://testserver") as client:183        response = await client.post("/api/predict", json={"prompt": "Hello", "top_k": 31})184 185    assert response.status_code == 422186