Skip to content

vllm.entrypoints.serve.utils.api_utils

Functions:

json_or_sse_responses(unary_model, chunk_model, error_responses)

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.

Source code in vllm/entrypoints/serve/utils/api_utils.py 
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
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}}}
            }
        },
    }

listen_for_disconnect(request) async

Returns if a disconnect message is received.

Source code in vllm/entrypoints/serve/utils/api_utils.py 
38
39
40
41
42
43
44
45
46
47
48
49
50
async def listen_for_disconnect(request: Request) -> None:
    """Returns if a disconnect message is received."""
    while True:
        message = await request.receive()
        if message["type"] == "http.disconnect":
            # If load tracking is enabled *and* the counter exists, decrement
            # it. Combines the previous nested checks into a single condition
            # to satisfy the linter rule.
            if getattr(
                request.app.state, "enable_server_load_tracking", False
            ) and hasattr(request.app.state, "server_load_metrics"):
                request.app.state.server_load_metrics -= 1
            break

redact_sensitive_args(args)

Return a copy of args with sensitive values redacted for logging.

Source code in vllm/entrypoints/serve/utils/api_utils.py 
345
346
347
348
349
350
351
352
def redact_sensitive_args(args: dict[str, Any]) -> dict[str, Any]:
    """Return a copy of `args` with sensitive values redacted for logging."""
    if not any(key in _SENSITIVE_ARG_FIELDS for key in args):
        return args
    return {
        key: ("***" if key in _SENSITIVE_ARG_FIELDS else value)
        for key, value in args.items()
    }

with_cancellation(handler_func)

Decorator that allows a route handler to be cancelled by client disconnections.

This does not use request.is_disconnected, which does not work with middleware. Instead this follows the pattern from starlette.StreamingResponse, which simultaneously awaits on two tasks- one to wait for an http disconnect message, and the other to do the work that we want done. When the first task finishes, the other is cancelled.

A core assumption of this method is that the body of the request has already been read. This is a safe assumption to make for fastapi handlers that have already parsed the body of the request into a pydantic model for us. This decorator is unsafe to use elsewhere, as it will consume and throw away all incoming messages for the request while it looks for a disconnect message.

In the case where a StreamingResponse is returned by the handler, this wrapper will stop listening for disconnects and instead the response object will start listening for disconnects.

Source code in vllm/entrypoints/serve/utils/api_utils.py 
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
def with_cancellation(handler_func):
    """Decorator that allows a route handler to be cancelled by client
    disconnections.

    This does _not_ use request.is_disconnected, which does not work with
    middleware. Instead this follows the pattern from
    starlette.StreamingResponse, which simultaneously awaits on two tasks- one
    to wait for an http disconnect message, and the other to do the work that we
    want done. When the first task finishes, the other is cancelled.

    A core assumption of this method is that the body of the request has already
    been read. This is a safe assumption to make for fastapi handlers that have
    already parsed the body of the request into a pydantic model for us.
    This decorator is unsafe to use elsewhere, as it will consume and throw away
    all incoming messages for the request while it looks for a disconnect
    message.

    In the case where a `StreamingResponse` is returned by the handler, this
    wrapper will stop listening for disconnects and instead the response object
    will start listening for disconnects.
    """

    # Functools.wraps is required for this wrapper to appear to fastapi as a
    # normal route handler, with the correct request type hinting.
    @functools.wraps(handler_func)
    async def wrapper(*args, **kwargs):
        # The request is either the second positional arg or `raw_request`
        request = args[1] if len(args) > 1 else kwargs["raw_request"]

        handler_task = asyncio.create_task(handler_func(*args, **kwargs))
        cancellation_task = asyncio.create_task(listen_for_disconnect(request))

        done, pending = await asyncio.wait(
            [handler_task, cancellation_task], return_when=asyncio.FIRST_COMPLETED
        )
        for task in pending:
            task.cancel()

        if handler_task in done:
            return handler_task.result()
        # The client is gone; an empty response skips response_model validation.
        return Response()

    return wrapper