[Bugfix][Frontend] Expose completion response schemas in OpenAPI (#60926)

Signed-off-by: Guan Luo <gluo@nvidia.com>
Signed-off-by: sfeng33 <4florafeng@gmail.com>
Co-authored-by: sfeng33 <4florafeng@gmail.com>

tests/entrypoints/openai/test_openai_schema.py

@@ -1,11 +1,14 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import json
+from argparse import Namespace
 from typing import Final
 
 import pytest
 import schemathesis
+from fastapi.testclient import TestClient
 from hypothesis import HealthCheck, settings
+from jsonschema import Draft202012Validator
 from schemathesis import GenerationMode
 from schemathesis.checks import not_a_server_error
 from schemathesis.config import (
@@ -30,6 +33,76 @@ DEFAULT_TIMEOUT_SECONDS: Final[int] = 10 * _ROCM_TIMEOUT_MULTIPLIER
 LONG_TIMEOUT_SECONDS: Final[int] = 60 * _ROCM_TIMEOUT_MULTIPLIER
 
 
+def test_openapi_exposes_completion_response_schemas() -> None:
+    """The published 200 contract of both completion APIs must reference the
+    real response models, for the JSON body and for each SSE data chunk.
+
+    Uses the production app: FastAPI renames models used as both input and
+    output (-Input/-Output), which a router-only app cannot show.
+    """
+    from vllm.entrypoints.launchers.app import build_app
+    from vllm.entrypoints.openai.chat_completion.protocol import (
+        ChatCompletionResponse,
+        ChatCompletionStreamResponse,
+    )
+    from vllm.entrypoints.openai.completion.protocol import (
+        CompletionResponse,
+        CompletionStreamResponse,
+    )
+
+    args = Namespace(
+        disable_fastapi_docs=False,
+        enable_offline_docs=False,
+        root_path="",
+        allowed_origins=["*"],
+        allow_credentials=False,
+        allowed_methods=["*"],
+        allowed_headers=["*"],
+        api_key=None,
+        enable_request_id_headers=False,
+        enable_fault_tolerance=False,
+        middleware=[],
+        log_error_stack=False,
+    )
+    # No lifespan: schema generation does not need an inference engine.
+    client = TestClient(build_app(args, supported_tasks=("generate",)))
+    document = client.get("/openapi.json").json()
+    message = {"role": "assistant", "content": "hi"}
+    for path, unary, unary_choice, chunk, chunk_choice in (
+        (
+            "/v1/chat/completions",
+            ChatCompletionResponse,
+            {"index": 0, "message": message},
+            ChatCompletionStreamResponse,
+            {"index": 0, "delta": {"content": "hi"}},
+        ),
+        (
+            "/v1/completions",
+            CompletionResponse,
+            {"index": 0, "text": "hi"},
+            CompletionStreamResponse,
+            {"index": 0, "text": "hi"},
+        ),
+    ):
+        content = document["paths"][path]["post"]["responses"]["200"]["content"]
+        data = content["text/event-stream"]["schema"]["properties"]["data"]
+        assert data["anyOf"][0] == {"const": "[DONE]"}
+        for schema, model, choice in (
+            (content["application/json"]["schema"], unary, unary_choice),
+            (data["anyOf"][1]["contentSchema"], chunk, chunk_choice),
+        ):
+            assert schema == {"$ref": f"#/components/schemas/{model.__name__}"}
+            validator = Draft202012Validator(
+                {**schema, "components": document["components"]}
+            )
+            payload = model(model="m", choices=[choice], usage={}).model_dump(
+                mode="json"
+            )
+            validator.validate(payload)
+            payload["choices"] = "not an array"
+            assert not validator.is_valid(payload)
+
+
 @pytest.fixture(scope="module")
 def server():
     args = [

tests/entrypoints/serve/utils/test_api_utils.py

@@ -1,9 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import asyncio
 from argparse import Namespace
 
 import pytest
+from fastapi import FastAPI, Request
 
 from vllm.entrypoints.generate.base.protocol import StreamOptions
 from vllm.entrypoints.serve.utils import api_utils
@@ -14,6 +16,55 @@ from vllm.entrypoints.serve.utils.api_utils import (
 )
 
 
+@pytest.mark.asyncio
+async def test_with_cancellation_disconnect_skips_response_validation():
+    """A disconnect must cancel the handler without a response_model 500."""
+    app = FastAPI()
+    started = asyncio.Event()
+    cancelled = asyncio.Event()
+
+    @app.post("/", response_model=str)
+    @api_utils.with_cancellation
+    async def handler(body: dict, raw_request: Request):
+        started.set()
+        try:
+            await asyncio.Future()
+        finally:
+            cancelled.set()
+
+    messages = []
+    body_sent = False
+
+    async def receive():
+        nonlocal body_sent
+        if not body_sent:
+            body_sent = True
+            return {"type": "http.request", "body": b"{}", "more_body": False}
+        await started.wait()
+        return {"type": "http.disconnect"}
+
+    async def send(message):
+        messages.append(message)
+
+    scope = {
+        "type": "http",
+        "asgi": {"version": "3.0"},
+        "http_version": "1.1",
+        "method": "POST",
+        "scheme": "http",
+        "path": "/",
+        "raw_path": b"/",
+        "query_string": b"",
+        "headers": [(b"content-type", b"application/json")],
+        "client": ("127.0.0.1", 1234),
+        "server": ("test", 80),
+    }
+    await asyncio.wait_for(app(scope, receive, send), timeout=5)
+    assert messages[0]["status"] == 200
+    assert messages[1]["body"] == b""
+    await asyncio.wait_for(cancelled.wait(), timeout=5)
+
+
 @pytest.mark.parametrize(
     ("stream_options", "expected"),
     [

vllm/entrypoints/openai/chat_completion/api_router.py

@@ -12,10 +12,12 @@ from vllm.entrypoints.openai.chat_completion.protocol import (
     BatchChatCompletionRequest,
     ChatCompletionRequest,
     ChatCompletionResponse,
+    ChatCompletionStreamResponse,
 )
 from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat
 from vllm.entrypoints.serve.engine.protocol import ErrorResponse
 from vllm.entrypoints.serve.utils.api_utils import (
+    json_or_sse_responses,
     load_aware_call,
     validate_json_request,
     with_cancellation,
@@ -41,13 +43,16 @@ def batch_chat(request: Request) -> OpenAIServingChatBatch | None:
 @router.post(
     "/v1/chat/completions",
     dependencies=[Depends(validate_json_request)],
-    responses={
-        HTTPStatus.OK.value: {"content": {"text/event-stream": {}}},
-        HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
-        HTTPStatus.NOT_FOUND.value: {"model": ErrorResponse},
-        HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
-        HTTPStatus.NOT_IMPLEMENTED.value: {"model": ErrorResponse},
-    },
+    **json_or_sse_responses(
+        ChatCompletionResponse,
+        ChatCompletionStreamResponse,
+        {
+            HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
+            HTTPStatus.NOT_FOUND.value: {"model": ErrorResponse},
+            HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
+            HTTPStatus.NOT_IMPLEMENTED.value: {"model": ErrorResponse},
+        },
+    ),
 )
 @with_cancellation
 @load_aware_call

vllm/entrypoints/openai/completion/api_router.py

@@ -10,10 +10,12 @@ from fastapi.responses import JSONResponse, StreamingResponse
 from vllm.entrypoints.openai.completion.protocol import (
     CompletionRequest,
     CompletionResponse,
+    CompletionStreamResponse,
 )
 from vllm.entrypoints.openai.completion.serving import OpenAIServingCompletion
 from vllm.entrypoints.serve.engine.protocol import ErrorResponse
 from vllm.entrypoints.serve.utils.api_utils import (
+    json_or_sse_responses,
     load_aware_call,
     validate_json_request,
     with_cancellation,
@@ -35,12 +37,15 @@ def completion(request: Request) -> OpenAIServingCompletion | None:
 @router.post(
     "/v1/completions",
     dependencies=[Depends(validate_json_request)],
-    responses={
-        HTTPStatus.OK.value: {"content": {"text/event-stream": {}}},
-        HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
-        HTTPStatus.NOT_FOUND.value: {"model": ErrorResponse},
-        HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
-    },
+    **json_or_sse_responses(
+        CompletionResponse,
+        CompletionStreamResponse,
+        {
+            HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
+            HTTPStatus.NOT_FOUND.value: {"model": ErrorResponse},
+            HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
+        },
+    ),
 )
 @with_cancellation
 @load_aware_call

vllm/entrypoints/serve/utils/api_utils.py

@@ -6,13 +6,15 @@ import dataclasses
 import functools
 import os
 from argparse import Namespace
+from http import HTTPStatus
 from logging import Logger
 from string import Template
 from typing import Any
 
 from fastapi import Request
 from fastapi.exceptions import RequestValidationError
-from fastapi.responses import JSONResponse, StreamingResponse
+from fastapi.responses import JSONResponse, Response, StreamingResponse
+from pydantic import BaseModel
 from starlette.background import BackgroundTask, BackgroundTasks
 
 from vllm import envs
@@ -88,11 +90,80 @@ def with_cancellation(handler_func):
 
         if handler_task in done:
             return handler_task.result()
-        return None
+        # The client is gone; an empty response skips response_model validation.
+        return Response()
 
     return wrapper
 
 
+_ERROR_REF = {"$ref": "#/components/schemas/ErrorResponse"}
+
+
+def _sse_event_schema(chunk_ref: dict[str, str]) -> dict[str, Any]:
+    # One SSE event: `data` is a JSON chunk, a JSON error, or "[DONE]". This is
+    # the envelope FastAPI emits for EventSourceResponse routes, but published
+    # under `schema` because OpenAPI 3.1 has no `itemSchema`.
+    return {
+        "type": "object",
+        "required": ["data"],
+        "properties": {
+            "data": {
+                "type": "string",
+                "anyOf": [
+                    {"const": "[DONE]"},
+                    {
+                        "contentMediaType": "application/json",
+                        "contentSchema": chunk_ref,
+                    },
+                    {
+                        "contentMediaType": "application/json",
+                        "contentSchema": _ERROR_REF,
+                    },
+                ],
+            },
+            "event": {"type": "string"},
+            "id": {"type": "string"},
+            "retry": {"type": "integer", "minimum": 0},
+        },
+    }
+
+
+def json_or_sse_responses(
+    unary_model: type[BaseModel],
+    chunk_model: type[BaseModel],
+    error_responses: dict[int, dict[str, Any]],
+) -> dict[str, Any]:
+    """Route kwargs for an endpoint answering with ``unary_model`` as JSON, or
+    with SSE ``data`` events carrying ``chunk_model`` when ``stream=true``.
+
+    FastAPI only registers models reachable from ``response_model`` or an
+    additional response's ``model``. The chunk model rides on the 200 entry,
+    which also overwrites the JSON schema, so ``openapi_extra`` restores it.
+    """
+    unary_ref = {"$ref": f"#/components/schemas/{unary_model.__name__}"}
+    chunk_ref = {"$ref": f"#/components/schemas/{chunk_model.__name__}"}
+    return {
+        "response_model": unary_model,
+        "responses": {
+            HTTPStatus.OK.value: {
+                "model": chunk_model,
+                "description": (
+                    "JSON body when stream=false; SSE events when stream=true."
+                ),
+                "content": {
+                    "text/event-stream": {"schema": _sse_event_schema(chunk_ref)}
+                },
+            },
+            **error_responses,
+        },
+        "openapi_extra": {
+            "responses": {
+                "200": {"content": {"application/json": {"schema": unary_ref}}}
+            }
+        },
+    }
+
+
 def decrement_server_load(request: Request):
     request.app.state.server_load_metrics -= 1