vllm-frontdoor/router/tests/test_model_resolution.py

"""Model resolution: JSON / multipart / defaults / embeddings / unknown."""

from __future__ import annotations

import json

TEXT = "Qwen3.6-35B-A3B-FP8"
OCR = "OvisOCR2"
EMBED = "Qwen3-Embedding-8B"


async def test_json_model_exact(pub, backend):
    r = await pub.post("/v1/chat/completions", json={"model": OCR, "messages": []})
    assert r.status_code == 200
    assert r.json()["service"] == "vllm-ocr"


async def test_json_model_case_insensitive(pub):
    r = await pub.post("/v1/chat/completions", json={"model": "qwen3.6-35b-a3b-fp8"})
    assert r.status_code == 200
    assert r.json()["service"] == "vllm-text"


async def test_alias_matches(pub):
    r = await pub.post("/v1/chat/completions", json={"model": "ocr"})
    assert r.status_code == 200
    assert r.json()["service"] == "vllm-ocr"
    r = await pub.post("/v1/chat/completions", json={"model": "Embedding"})
    assert r.json()["service"] == "vllm-embed"


async def test_missing_model_defaults_to_text_on_chat(pub):
    r = await pub.post("/v1/chat/completions", json={"messages": [{"role": "user", "content": "hi"}]})
    assert r.status_code == 200
    assert r.json()["service"] == "vllm-text"


async def test_missing_model_defaults_to_text_on_completions(pub):
    r = await pub.post("/v1/completions", json={"prompt": "hi"})
    assert r.status_code == 200
    assert r.json()["service"] == "vllm-text"


async def test_missing_model_is_400_elsewhere(pub, backend):
    r = await pub.post("/v1/rerank", json={"query": "hi"})
    assert r.status_code == 400
    assert r.json()["error"]["code"] == "missing_model"
    assert backend.count("/v1/rerank") == 0


async def test_invalid_json_is_400(pub, backend):
    r = await pub.post("/v1/chat/completions",
                       content=b"{not json",
                       headers={"content-type": "application/json"})
    assert r.status_code == 400
    assert backend.count("/v1/") == 0


async def test_embeddings_always_routes_to_embed(pub):
    r = await pub.post("/v1/embeddings", json={"input": "hello"})
    assert r.status_code == 200
    assert r.json()["service"] == "vllm-embed"
    # ... even when the body names a different model
    r = await pub.post("/v1/embeddings", json={"input": "hello", "model": TEXT})
    assert r.status_code == 200
    assert r.json()["service"] == "vllm-embed"


async def test_unknown_model_404_and_no_wake(pub, backend):
    backend.set_sleeping("ocr", True)
    r = await pub.post("/v1/chat/completions", json={"model": "gpt-4o"})
    assert r.status_code == 404
    body = r.json()["error"]
    assert body["code"] == "model_not_found"
    assert body["type"] == "invalid_request_error"
    assert body["param"] == "model"
    # no probe, no wake, no proxy
    assert backend.calls == []


async def test_multipart_with_model_field(pub):
    r = await pub.post(
        "/v1/chat/completions",
        data={"model": OCR},
        files={"image": ("page.png", b"PNGDATA" * 64, "image/png")},
    )
    assert r.status_code == 200
    echo = r.json()
    assert echo["service"] == "vllm-ocr"
    # raw body forwarded unchanged: the file bytes and the boundary survive
    assert "PNGDATA" * 64 in echo["echo_body"]
    assert echo["echo_content_type"].startswith("multipart/form-data; boundary=")


async def test_multipart_without_model_defaults_to_ocr(pub):
    r = await pub.post(
        "/v1/chat/completions",
        files={"image": ("page.png", b"PNGDATA", "image/png")},
    )
    assert r.status_code == 200
    assert r.json()["service"] == "vllm-ocr"


async def test_multipart_unknown_model_404(pub, backend):
    r = await pub.post(
        "/v1/chat/completions",
        data={"model": "nope"},
        files={"image": ("page.png", b"x", "image/png")},
    )
    assert r.status_code == 404
    assert r.json()["error"]["code"] == "model_not_found"
    assert backend.calls == []


async def test_multipart_model_extraction_ignores_file_parts(pub):
    # a file part literally named "model" must not be read as the model field
    r = await pub.post(
        "/v1/chat/completions",
        data={"prompt": "hi"},
        files={"model": ("fake.json", b"NOT-A-MODEL-NAME", "application/octet-stream")},
    )
    # no usable model field -> multipart default on the chat path is OCR
    assert r.status_code == 200
    assert r.json()["service"] == "vllm-ocr"
    assert "NOT-A-MODEL-NAME" in r.json()["echo_body"]


async def test_model_in_path_for_models_detail(pub):
    r = await pub.get(f"/v1/models/{OCR}")
    assert r.status_code == 200
    assert r.json()["service"] == "vllm-ocr"
    r = await pub.get("/v1/models/does-not-exist")
    assert r.status_code == 404
    assert r.json()["error"]["code"] == "model_not_found"


async def test_body_and_headers_forwarded(pub):
    payload = {"model": TEXT, "messages": [{"role": "user", "content": "hi"}], "stream": False}
    r = await pub.post("/v1/chat/completions", json=payload,
                       headers={"Authorization": "Bearer x", "x-custom": "1"})
    assert r.status_code == 200
    echo = r.json()
    assert echo["echo_path"] == "/v1/chat/completions"
    assert echo["echo_method"] == "POST"
    assert json.loads(echo["echo_body"]) == payload


async def test_query_string_forwarded(pub):
    r = await pub.post("/v1/chat/completions?foo=bar&baz=1", json={"model": TEXT})
    assert r.status_code == 200
    assert r.json()["echo_query"] == "foo=bar&baz=1"


async def test_streaming_passthrough(pub, backend):
    backend.stream_chunks = [b"data: {\"a\":1}\n\n", b"data: {\"a\":2}\n\n", b"data: [DONE]\n\n"]
    r = await pub.post("/v1/chat/completions", json={"model": TEXT, "stream": True})
    assert r.status_code == 200
    assert r.headers["content-type"].startswith("text/event-stream")
    assert r.text == "".join(chunk.decode() for chunk in backend.stream_chunks)
    backend.stream_chunks = []