| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """FastAPI router for the Cohere Chat v2 API. |
| 4 | |
| 5 | Exposes two routes: |
| 6 | |
| 7 | * ``POST /cohere/v2/chat`` - the chat endpoint itself. |
| 8 | * ``POST /cohere/v2/chat/render`` - tokenize a request without running |
| 9 | generation, mirroring ``POST /v1/chat/completions/render``. |
| 10 | |
| 11 | The Cohere v2 protocol models are sourced from the official ``cohere`` |
| 12 | Python SDK (``pip install cohere``). To keep that an *optional* |
| 13 | dependency for vLLM, the SDK-dependent imports - and the route handlers |
| 14 | themselves - are gated on a one-shot probe at module load. If the SDK |
| 15 | isn't installed, :func:`attach_router` becomes a no-op (with an info |
| 16 | log) and vLLM continues to boot normally. |
| 17 | |
| 18 | Even when the SDK is installed, :func:`attach_router` also requires |
| 19 | ``VLLM_ENABLE_COHERE_API=1`` in the environment before it will expose |
| 20 | the chat route. The render route additionally requires scale-out |
| 21 | endpoints to be enabled, i.e. ``--enable-scale-out`` or |
| 22 | ``--tokens-only``. These gates keep unrelated deployments from |
| 23 | accidentally exposing the APIs. |
| 24 | |
| 25 | Note: the handlers must live at module scope (not inside |
| 26 | ``attach_router``) so that FastAPI's ``typing.get_type_hints`` resolves |
| 27 | the ``CohereChatV2Request`` body annotation against the module's |
| 28 | globals. Defining them locally inside ``attach_router`` would hide the |
| 29 | type from ``get_type_hints``, causing FastAPI to silently degrade the |
| 30 | body parameter into a query parameter and reject every request with |
| 31 | 422. |
| 32 | """ |
| 33 | |
| 34 | import json |
| 35 | from http import HTTPStatus |
| 36 | |
| 37 | from fastapi import APIRouter, Depends, FastAPI, Request |
| 38 | from fastapi.responses import JSONResponse, Response, StreamingResponse |
| 39 | from starlette.middleware.base import BaseHTTPMiddleware |
| 40 | |
| 41 | import vllm.envs as envs |
| 42 | from vllm.entrypoints.serve.engine.protocol import ErrorResponse |
| 43 | from vllm.entrypoints.serve.exception_handling.utils import sanitize_message |
| 44 | from vllm.entrypoints.serve.utils.api_utils import ( |
| 45 | load_aware_call, |
| 46 | validate_json_request, |
| 47 | with_cancellation, |
| 48 | ) |
| 49 | from vllm.logger import init_logger |
| 50 | |
| 51 | _COHERE_PATH_PREFIX = "/cohere/" |
| 52 | |
| 53 | logger = init_logger(__name__) |
| 54 | |
| 55 | |
| 56 | try: |
| 57 | import cohere # noqa: F401 -- dependency probe |
| 58 | except ImportError: |
| 59 | _SDK_AVAILABLE = False |
| 60 | else: |
| 61 | _SDK_AVAILABLE = True |
| 62 | |
| 63 | |
| 64 | if _SDK_AVAILABLE: |
| 65 | from vllm.entrypoints.cohere.protocol import ( |
| 66 | CohereChatV2Request, |
| 67 | CohereChatV2Response, |
| 68 | CohereError, |
| 69 | ) |
| 70 | from vllm.entrypoints.cohere.serving import CohereServingChatV2 |
| 71 | from vllm.entrypoints.scale_out.render.serving import ServingRender |
| 72 | from vllm.entrypoints.scale_out.token_in_token_out.protocol import GenerateRequest |
| 73 | |
| 74 | router = APIRouter() |
| 75 | render_router = APIRouter() |
| 76 | |
| 77 | def _serving(request: Request) -> CohereServingChatV2 | None: |
| 78 | return getattr(request.app.state, "cohere_serving_chat_v2", None) |
| 79 | |
| 80 | def _serving_render(request: Request) -> ServingRender | None: |
| 81 | return getattr(request.app.state, "serving_render", None) |
| 82 | |
| 83 | def _request_id(raw_request: Request | None) -> str | None: |
| 84 | """Best-effort lookup of the active request id. |
| 85 | |
| 86 | Prefers the id the underlying chat handler stamped onto |
| 87 | ``raw_request.state.request_metadata`` (if it got that far before |
| 88 | failing), falling back to the ``X-Request-Id`` HTTP header. May |
| 89 | return ``None`` if neither is available, in which case the field |
| 90 | is omitted from the response. |
| 91 | """ |
| 92 | if raw_request is None: |
| 93 | return None |
| 94 | meta = getattr(raw_request.state, "request_metadata", None) |
| 95 | if meta is not None and getattr(meta, "request_id", None): |
| 96 | return meta.request_id |
| 97 | return raw_request.headers.get("X-Request-Id") |
| 98 | |
| 99 | def _error_response( |
| 100 | error: ErrorResponse, |
| 101 | raw_request: Request | None, |
| 102 | *, |
| 103 | fallback_status: int = HTTPStatus.BAD_REQUEST, |
| 104 | ) -> JSONResponse: |
| 105 | """Translate vLLM's internal error envelope into Cohere's shape.""" |
| 106 | info = error.error |
| 107 | status = info.code or fallback_status |
| 108 | return JSONResponse( |
| 109 | status_code=status, |
| 110 | content=CohereError( |
| 111 | message=sanitize_message(info.message), |
| 112 | id=_request_id(raw_request), |
| 113 | ).model_dump(exclude_none=True), |
| 114 | ) |
| 115 | |
| 116 | @router.post( |
| 117 | "/cohere/v2/chat", |
| 118 | dependencies=[Depends(validate_json_request)], |
| 119 | responses={ |
| 120 | HTTPStatus.OK.value: {"content": {"text/event-stream": {}}}, |
| 121 | HTTPStatus.BAD_REQUEST.value: {"model": CohereError}, |
| 122 | HTTPStatus.NOT_FOUND.value: {"model": CohereError}, |
| 123 | HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": CohereError}, |
| 124 | }, |
| 125 | ) |
| 126 | @with_cancellation |
| 127 | @load_aware_call |
| 128 | async def chat_v2(request: CohereChatV2Request, raw_request: Request): |
| 129 | handler = _serving(raw_request) |
| 130 | if handler is None: |
| 131 | return JSONResponse( |
| 132 | status_code=HTTPStatus.NOT_IMPLEMENTED.value, |
| 133 | content=CohereError( |
| 134 | message="The model does not support the Cohere v2 chat API.", |
| 135 | id=_request_id(raw_request), |
| 136 | ).model_dump(exclude_none=True), |
| 137 | ) |
| 138 | |
| 139 | try: |
| 140 | result = await handler.create_chat_v2(request, raw_request) |
| 141 | except Exception as e: # noqa: BLE001 - report as 500 for parity |
| 142 | logger.exception("Error in /cohere/v2/chat: %s", e) |
| 143 | return JSONResponse( |
| 144 | status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, |
| 145 | content=CohereError( |
| 146 | message=sanitize_message(str(e)), |
| 147 | id=_request_id(raw_request), |
| 148 | ).model_dump(exclude_none=True), |
| 149 | ) |
| 150 | |
| 151 | match result: |
| 152 | case ErrorResponse(): |
| 153 | return _error_response(result, raw_request) |
| 154 | case CohereChatV2Response(): |
| 155 | return JSONResponse(content=result.model_dump(exclude_none=True)) |
| 156 | case _: |
| 157 | return StreamingResponse(content=result, media_type="text/event-stream") |
| 158 | |
| 159 | @render_router.post( |
| 160 | "/cohere/v2/chat/render", |
| 161 | dependencies=[Depends(validate_json_request)], |
| 162 | response_model=GenerateRequest, |
| 163 | responses={ |
| 164 | HTTPStatus.BAD_REQUEST.value: {"model": CohereError}, |
| 165 | HTTPStatus.NOT_FOUND.value: {"model": CohereError}, |
| 166 | HTTPStatus.NOT_IMPLEMENTED.value: {"model": CohereError}, |
| 167 | HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": CohereError}, |
| 168 | }, |
| 169 | ) |
| 170 | @with_cancellation |
| 171 | async def chat_v2_render(request: CohereChatV2Request, raw_request: Request): |
| 172 | """Tokenize a Cohere v2 chat request without running generation. |
| 173 | |
| 174 | The Cohere counterpart to ``POST /v1/chat/completions/render``. |
| 175 | The v2 body goes through the same conversion |
| 176 | ``/cohere/v2/chat`` uses and is then handed to the shared |
| 177 | :class:`ServingRender`, so the returned ``GenerateRequest`` |
| 178 | carries the prompt tokens and sampling params the chat endpoint |
| 179 | would have sent to the engine. |
| 180 | """ |
| 181 | handler = _serving(raw_request) |
| 182 | render_handler = _serving_render(raw_request) |
| 183 | if handler is None or render_handler is None: |
| 184 | return JSONResponse( |
| 185 | status_code=HTTPStatus.NOT_IMPLEMENTED.value, |
| 186 | content=CohereError( |
| 187 | message=( |
| 188 | "The model does not support the Cohere v2 chat render API." |
| 189 | ), |
| 190 | id=_request_id(raw_request), |
| 191 | ).model_dump(exclude_none=True), |
| 192 | ) |
| 193 | |
| 194 | try: |
| 195 | chat_request = handler.to_chat_completion_request(request) |
| 196 | result = await render_handler.render_chat_request(chat_request) |
| 197 | except Exception as e: # noqa: BLE001 - report as 500 for parity |
| 198 | logger.exception("Error in /cohere/v2/chat/render: %s", e) |
| 199 | return JSONResponse( |
| 200 | status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, |
| 201 | content=CohereError( |
| 202 | message=sanitize_message(str(e)), |
| 203 | id=_request_id(raw_request), |
| 204 | ).model_dump(exclude_none=True), |
| 205 | ) |
| 206 | |
| 207 | if isinstance(result, ErrorResponse): |
| 208 | return _error_response(result, raw_request) |
| 209 | |
| 210 | return JSONResponse(content=result.model_dump()) |
| 211 | |
| 212 | class CohereErrorEnvelopeMiddleware(BaseHTTPMiddleware): |
| 213 | """Rewrite vLLM error bodies into the Cohere ``{message, id}`` shape. |
| 214 | |
| 215 | The endpoint handler above already returns :class:`CohereError` for |
| 216 | errors it owns, but globally-registered exception handlers (e.g. |
| 217 | :func:`validation_exception_handler` for pydantic body errors, |
| 218 | :func:`http_exception_handler`, engine error handlers) fire *before* |
| 219 | the handler runs and produce vLLM's internal |
| 220 | ``ErrorResponse`` shape (``{"error": {"message": ...}}``). That |
| 221 | shape doesn't match the ``CohereError`` schema advertised on the |
| 222 | route's OpenAPI ``responses``, so clients (and schema-conformance |
| 223 | tests like ``test_openai_schema.py``) would see a mismatch on |
| 224 | those paths. This middleware normalises any error body on |
| 225 | ``/cohere/*`` responses to :class:`CohereError`. |
| 226 | """ |
| 227 | |
| 228 | async def dispatch(self, request: Request, call_next): |
| 229 | response = await call_next(request) |
| 230 | if not request.url.path.startswith(_COHERE_PATH_PREFIX): |
| 231 | return response |
| 232 | if response.status_code < 400: |
| 233 | return response |
| 234 | content_type = response.headers.get("content-type", "") |
| 235 | if not content_type.startswith("application/json"): |
| 236 | return response |
| 237 | |
| 238 | body = b"".join([chunk async for chunk in response.body_iterator]) |
| 239 | translated = _translate_vllm_error_body(body, request) |
| 240 | if translated is not None: |
| 241 | return translated |
| 242 | passthrough_headers = { |
| 243 | k: v |
| 244 | for k, v in response.headers.items() |
| 245 | if k.lower() != "content-length" |
| 246 | } |
| 247 | return Response( |
| 248 | content=body, |
| 249 | status_code=response.status_code, |
| 250 | headers=passthrough_headers, |
| 251 | media_type=content_type, |
| 252 | ) |
| 253 | |
| 254 | def _translate_vllm_error_body(raw: bytes, request: Request) -> JSONResponse | None: |
| 255 | """Translate a vLLM ``ErrorResponse`` body to a ``CohereError`` body. |
| 256 | |
| 257 | Returns ``None`` if ``raw`` does not match the vLLM error envelope |
| 258 | (which signals the middleware to pass the body through unchanged). |
| 259 | """ |
| 260 | try: |
| 261 | data = json.loads(raw) |
| 262 | except (json.JSONDecodeError, TypeError, ValueError): |
| 263 | return None |
| 264 | if not ( |
| 265 | isinstance(data, dict) |
| 266 | and isinstance(data.get("error"), dict) |
| 267 | and "message" in data["error"] |
| 268 | ): |
| 269 | return None |
| 270 | try: |
| 271 | err = ErrorResponse.model_validate(data) |
| 272 | except Exception: # noqa: BLE001 - malformed envelope; pass through |
| 273 | return None |
| 274 | return _error_response(err, request) |
| 275 | |
| 276 | |
| 277 | def attach_router(app: FastAPI) -> None: |
| 278 | """Register the enabled ``/cohere/v2/chat`` routes on ``app``. |
| 279 | |
| 280 | No-op when either: |
| 281 | |
| 282 | * the ``VLLM_ENABLE_COHERE_API`` env var isn't set to ``1``. The |
| 283 | Cohere v2 endpoints are opt-in because they carry Cohere-specific |
| 284 | request/response semantics (grounding citations, tool_plan, |
| 285 | PLAN/THINKING_CONTENT blocks) that are only meaningful when |
| 286 | serving a Cohere Command-family model. |
| 287 | * the optional ``cohere`` SDK isn't installed (the v2 protocol |
| 288 | models live there) |
| 289 | |
| 290 | The two skip paths log at different levels: an operator who set |
| 291 | ``VLLM_ENABLE_COHERE_API=1`` but forgot to install ``cohere`` sees |
| 292 | a WARNING (they explicitly asked for the endpoints and they're |
| 293 | silently absent), whereas the default-off skip logs at debug. |
| 294 | """ |
| 295 | enabled = envs.VLLM_ENABLE_COHERE_API |
| 296 | if not enabled: |
| 297 | logger.debug( |
| 298 | "VLLM_ENABLE_COHERE_API is not set; /cohere/v2/chat endpoints " |
| 299 | "disabled. Set VLLM_ENABLE_COHERE_API=1 to enable them." |
| 300 | ) |
| 301 | return |
| 302 | if not _SDK_AVAILABLE: |
| 303 | logger.warning( |
| 304 | "VLLM_ENABLE_COHERE_API=1 but the `cohere` SDK is not " |
| 305 | "installed; /cohere/v2/chat will not be exposed. Install " |
| 306 | "with `pip install cohere` to enable the endpoints." |
| 307 | ) |
| 308 | return |
| 309 | app.include_router(router) |
| 310 | args = getattr(app.state, "args", None) |
| 311 | if getattr(args, "enable_scale_out", False) or getattr(args, "tokens_only", False): |
| 312 | app.include_router(render_router) |
| 313 | app.add_middleware(CohereErrorEnvelopeMiddleware) |