FastAPI router for the Cohere Chat v2 API.
Exposes two routes:
POST /cohere/v2/chat - the chat endpoint itself.
POST /cohere/v2/chat/render - tokenize a request without running
generation, mirroring POST /v1/chat/completions/render.
The Cohere v2 protocol models are sourced from the official cohere
Python SDK (pip install cohere). To keep that an optional
dependency for vLLM, the SDK-dependent imports - and the route handlers
themselves - are gated on a one-shot probe at module load. If the SDK
isn't installed, attach_router becomes a no-op (with an info
log) and vLLM continues to boot normally.
Even when the SDK is installed, attach_router also requires
VLLM_ENABLE_COHERE_API=1 in the environment before it will expose
the chat route. The render route additionally requires scale-out
endpoints to be enabled, i.e. --enable-scale-out or
--tokens-only. These gates keep unrelated deployments from
accidentally exposing the APIs.
Note: the handlers must live at module scope (not inside
attach_router) so that FastAPI's typing.get_type_hints resolves
the CohereChatV2Request body annotation against the module's
globals. Defining them locally inside attach_router would hide the
type from get_type_hints, causing FastAPI to silently degrade the
body parameter into a query parameter and reject every request with
422.
Classes:
Functions:
-
attach_router
–
Register the enabled /cohere/v2/chat routes on app.
-
chat_v2_render
–
Tokenize a Cohere v2 chat request without running generation.
CohereErrorEnvelopeMiddleware
Bases: BaseHTTPMiddleware
Rewrite vLLM error bodies into the Cohere {message, id} shape.
The endpoint handler above already returns CohereError for
errors it owns, but globally-registered exception handlers (e.g.
validation_exception_handler for pydantic body errors,
http_exception_handler, engine error handlers) fire before
the handler runs and produce vLLM's internal
ErrorResponse shape ({"error": {"message": ...}}). That
shape doesn't match the CohereError schema advertised on the
route's OpenAPI responses, so clients (and schema-conformance
tests like test_openai_schema.py) would see a mismatch on
those paths. This middleware normalises any error body on
/cohere/* responses to CohereError.
Source code in vllm/entrypoints/cohere/api_router.py
| class CohereErrorEnvelopeMiddleware(BaseHTTPMiddleware):
"""Rewrite vLLM error bodies into the Cohere `{message, id}` shape.
The endpoint handler above already returns `CohereError` for
errors it owns, but globally-registered exception handlers (e.g.
`validation_exception_handler` for pydantic body errors,
`http_exception_handler`, engine error handlers) fire *before*
the handler runs and produce vLLM's internal
`ErrorResponse` shape (`{"error": {"message": ...}}`). That
shape doesn't match the `CohereError` schema advertised on the
route's OpenAPI `responses`, so clients (and schema-conformance
tests like `test_openai_schema.py`) would see a mismatch on
those paths. This middleware normalises any error body on
`/cohere/*` responses to `CohereError`.
"""
async def dispatch(self, request: Request, call_next):
response = await call_next(request)
if not request.url.path.startswith(_COHERE_PATH_PREFIX):
return response
if response.status_code < 400:
return response
content_type = response.headers.get("content-type", "")
if not content_type.startswith("application/json"):
return response
body = b"".join([chunk async for chunk in response.body_iterator])
translated = _translate_vllm_error_body(body, request)
if translated is not None:
return translated
passthrough_headers = {
k: v
for k, v in response.headers.items()
if k.lower() != "content-length"
}
return Response(
content=body,
status_code=response.status_code,
headers=passthrough_headers,
media_type=content_type,
)
|
_error_response(error, raw_request, *, fallback_status=HTTPStatus.BAD_REQUEST)
Translate vLLM's internal error envelope into Cohere's shape.
Source code in vllm/entrypoints/cohere/api_router.py
| def _error_response(
error: ErrorResponse,
raw_request: Request | None,
*,
fallback_status: int = HTTPStatus.BAD_REQUEST,
) -> JSONResponse:
"""Translate vLLM's internal error envelope into Cohere's shape."""
info = error.error
status = info.code or fallback_status
return JSONResponse(
status_code=status,
content=CohereError(
message=sanitize_message(info.message),
id=_request_id(raw_request),
).model_dump(exclude_none=True),
)
|
_request_id(raw_request)
Best-effort lookup of the active request id.
Prefers the id the underlying chat handler stamped onto
raw_request.state.request_metadata (if it got that far before
failing), falling back to the X-Request-Id HTTP header. May
return None if neither is available, in which case the field
is omitted from the response.
Source code in vllm/entrypoints/cohere/api_router.py
| def _request_id(raw_request: Request | None) -> str | None:
"""Best-effort lookup of the active request id.
Prefers the id the underlying chat handler stamped onto
`raw_request.state.request_metadata` (if it got that far before
failing), falling back to the `X-Request-Id` HTTP header. May
return `None` if neither is available, in which case the field
is omitted from the response.
"""
if raw_request is None:
return None
meta = getattr(raw_request.state, "request_metadata", None)
if meta is not None and getattr(meta, "request_id", None):
return meta.request_id
return raw_request.headers.get("X-Request-Id")
|
_translate_vllm_error_body(raw, request)
Translate a vLLM ErrorResponse body to a CohereError body.
Returns None if raw does not match the vLLM error envelope
(which signals the middleware to pass the body through unchanged).
Source code in vllm/entrypoints/cohere/api_router.py
| def _translate_vllm_error_body(raw: bytes, request: Request) -> JSONResponse | None:
"""Translate a vLLM `ErrorResponse` body to a `CohereError` body.
Returns `None` if `raw` does not match the vLLM error envelope
(which signals the middleware to pass the body through unchanged).
"""
try:
data = json.loads(raw)
except (json.JSONDecodeError, TypeError, ValueError):
return None
if not (
isinstance(data, dict)
and isinstance(data.get("error"), dict)
and "message" in data["error"]
):
return None
try:
err = ErrorResponse.model_validate(data)
except Exception: # noqa: BLE001 - malformed envelope; pass through
return None
return _error_response(err, request)
|
attach_router(app)
Register the enabled /cohere/v2/chat routes on app.
No-op when either:
- the
VLLM_ENABLE_COHERE_API env var isn't set to 1. The
Cohere v2 endpoints are opt-in because they carry Cohere-specific
request/response semantics (grounding citations, tool_plan,
PLAN/THINKING_CONTENT blocks) that are only meaningful when
serving a Cohere Command-family model.
- the optional
cohere SDK isn't installed (the v2 protocol
models live there)
The two skip paths log at different levels: an operator who set
VLLM_ENABLE_COHERE_API=1 but forgot to install cohere sees
a WARNING (they explicitly asked for the endpoints and they're
silently absent), whereas the default-off skip logs at debug.
Source code in vllm/entrypoints/cohere/api_router.py
| def attach_router(app: FastAPI) -> None:
"""Register the enabled `/cohere/v2/chat` routes on `app`.
No-op when either:
* the `VLLM_ENABLE_COHERE_API` env var isn't set to `1`. The
Cohere v2 endpoints are opt-in because they carry Cohere-specific
request/response semantics (grounding citations, tool_plan,
PLAN/THINKING_CONTENT blocks) that are only meaningful when
serving a Cohere Command-family model.
* the optional `cohere` SDK isn't installed (the v2 protocol
models live there)
The two skip paths log at different levels: an operator who set
`VLLM_ENABLE_COHERE_API=1` but forgot to install `cohere` sees
a WARNING (they explicitly asked for the endpoints and they're
silently absent), whereas the default-off skip logs at debug.
"""
enabled = envs.VLLM_ENABLE_COHERE_API
if not enabled:
logger.debug(
"VLLM_ENABLE_COHERE_API is not set; /cohere/v2/chat endpoints "
"disabled. Set VLLM_ENABLE_COHERE_API=1 to enable them."
)
return
if not _SDK_AVAILABLE:
logger.warning(
"VLLM_ENABLE_COHERE_API=1 but the `cohere` SDK is not "
"installed; /cohere/v2/chat will not be exposed. Install "
"with `pip install cohere` to enable the endpoints."
)
return
app.include_router(router)
args = getattr(app.state, "args", None)
if getattr(args, "enable_scale_out", False) or getattr(args, "tokens_only", False):
app.include_router(render_router)
app.add_middleware(CohereErrorEnvelopeMiddleware)
|
chat_v2_render(request, raw_request)
async
Tokenize a Cohere v2 chat request without running generation.
The Cohere counterpart to POST /v1/chat/completions/render.
The v2 body goes through the same conversion
/cohere/v2/chat uses and is then handed to the shared
ServingRender, so the returned GenerateRequest
carries the prompt tokens and sampling params the chat endpoint
would have sent to the engine.
Source code in vllm/entrypoints/cohere/api_router.py
| @render_router.post(
"/cohere/v2/chat/render",
dependencies=[Depends(validate_json_request)],
response_model=GenerateRequest,
responses={
HTTPStatus.BAD_REQUEST.value: {"model": CohereError},
HTTPStatus.NOT_FOUND.value: {"model": CohereError},
HTTPStatus.NOT_IMPLEMENTED.value: {"model": CohereError},
HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": CohereError},
},
)
@with_cancellation
async def chat_v2_render(request: CohereChatV2Request, raw_request: Request):
"""Tokenize a Cohere v2 chat request without running generation.
The Cohere counterpart to `POST /v1/chat/completions/render`.
The v2 body goes through the same conversion
`/cohere/v2/chat` uses and is then handed to the shared
`ServingRender`, so the returned `GenerateRequest`
carries the prompt tokens and sampling params the chat endpoint
would have sent to the engine.
"""
handler = _serving(raw_request)
render_handler = _serving_render(raw_request)
if handler is None or render_handler is None:
return JSONResponse(
status_code=HTTPStatus.NOT_IMPLEMENTED.value,
content=CohereError(
message=(
"The model does not support the Cohere v2 chat render API."
),
id=_request_id(raw_request),
).model_dump(exclude_none=True),
)
try:
chat_request = handler.to_chat_completion_request(request)
# Derender mode returns no citations: the position-to-source
# map can't go over JSON.
result = await render_handler.render_chat_request(
chat_request,
engine_chat_template_kwargs=handler._engine_chat_template_kwargs,
)
except Exception as e:
logger.exception("Error in /cohere/v2/chat/render: %s", e)
return _error_response(create_error_response(e), raw_request)
if isinstance(result, ErrorResponse):
return _error_response(result, raw_request)
if (kwargs := result.reasoning_parser_kwargs) is not None:
kwargs.chat_template_kwargs = handler._engine_chat_template_kwargs(
kwargs.chat_template_kwargs
)
return JSONResponse(content=result.model_dump())
|