From 64cb47874dee33e78ca3622f204f423a7f51148e Mon Sep 17 00:00:00 2001 From: Andrew Berneshawi Date: Wed, 3 Jun 2026 05:30:13 -0400 Subject: [PATCH 01/30] Initial chat v2 support --- vllm/entrypoints/cohere/__init__.py | 2 + vllm/entrypoints/cohere/api_router.py | 82 ++ vllm/entrypoints/cohere/protocol.py | 472 +++++++++ vllm/entrypoints/cohere/serving.py | 960 ++++++++++++++++++ vllm/entrypoints/generate/api_router.py | 27 + .../openai/chat_completion/protocol.py | 6 + vllm/entrypoints/openai/engine/protocol.py | 46 + vllm/renderers/cohere.py | 452 +++++++++ vllm/renderers/registry.py | 1 + vllm/tokenizers/registry.py | 3 + 10 files changed, 2051 insertions(+) create mode 100644 vllm/entrypoints/cohere/__init__.py create mode 100644 vllm/entrypoints/cohere/api_router.py create mode 100644 vllm/entrypoints/cohere/protocol.py create mode 100644 vllm/entrypoints/cohere/serving.py create mode 100644 vllm/renderers/cohere.py diff --git a/vllm/entrypoints/cohere/__init__.py b/vllm/entrypoints/cohere/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/vllm/entrypoints/cohere/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/entrypoints/cohere/api_router.py b/vllm/entrypoints/cohere/api_router.py new file mode 100644 index 000000000000..86ababaac0bc --- /dev/null +++ b/vllm/entrypoints/cohere/api_router.py @@ -0,0 +1,82 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from http import HTTPStatus + +from fastapi import APIRouter, Depends, FastAPI, Request +from fastapi.responses import JSONResponse, StreamingResponse + +from vllm.entrypoints.cohere.protocol import ( + CohereChatV2Request, + CohereChatV2Response, + CohereError, +) +from vllm.entrypoints.cohere.serving import CohereServingChatV2 +from vllm.entrypoints.openai.engine.protocol import ErrorResponse +from vllm.entrypoints.openai.utils import validate_json_request +from vllm.entrypoints.utils import load_aware_call, with_cancellation +from vllm.logger import init_logger + +logger = init_logger(__name__) + +router = APIRouter() + + +def _serving(request: Request) -> CohereServingChatV2 | None: + return getattr(request.app.state, "cohere_serving_chat_v2", None) + + +def _error_response( + error: ErrorResponse, *, fallback_status: int = HTTPStatus.BAD_REQUEST +) -> JSONResponse: + """Translate vLLM's internal error envelope into Cohere's error shape.""" + info = error.error + status = info.code or fallback_status + return JSONResponse( + status_code=status, + content=CohereError(message=info.message).model_dump(exclude_none=True), + ) + + +@router.post( + "/v2/chat", + dependencies=[Depends(validate_json_request)], + responses={ + HTTPStatus.OK.value: {"content": {"text/event-stream": {}}}, + HTTPStatus.BAD_REQUEST.value: {"model": CohereError}, + HTTPStatus.NOT_FOUND.value: {"model": CohereError}, + HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": CohereError}, + }, +) +@with_cancellation +@load_aware_call +async def chat_v2(request: CohereChatV2Request, raw_request: Request): + handler = _serving(raw_request) + if handler is None: + return JSONResponse( + status_code=HTTPStatus.NOT_IMPLEMENTED.value, + content=CohereError( + message="The model does not support the Cohere v2 chat API." + ).model_dump(exclude_none=True), + ) + + try: + result = await handler.create_chat_v2(request, raw_request) + except Exception as e: # noqa: BLE001 - report as 500 for parity w/ siblings + logger.exception("Error in /v2/chat: %s", e) + return JSONResponse( + status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, + content=CohereError(message=str(e)).model_dump(exclude_none=True), + ) + + if isinstance(result, ErrorResponse): + return _error_response(result) + + if isinstance(result, CohereChatV2Response): + return JSONResponse(content=result.model_dump(exclude_none=True)) + + return StreamingResponse(content=result, media_type="text/event-stream") + + +def attach_router(app: FastAPI) -> None: + app.include_router(router) diff --git a/vllm/entrypoints/cohere/protocol.py b/vllm/entrypoints/cohere/protocol.py new file mode 100644 index 000000000000..5adc3440c788 --- /dev/null +++ b/vllm/entrypoints/cohere/protocol.py @@ -0,0 +1,472 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Pydantic models for the Cohere Chat v2 API protocol. + +See https://docs.cohere.com/reference/chat for the upstream specification. + +This module mirrors the Cohere V2 wire format using Pydantic discriminated +unions for messages, content blocks, and stream events. +""" + +import time +from typing import Any, Literal + +from pydantic import BaseModel, Field, RootModel, field_validator + +# --------------------------------------------------------------------------- +# Errors +# --------------------------------------------------------------------------- + + +class CohereError(BaseModel): + """Top-level error body returned by /v2/chat error responses. + + Cohere's documented error schemas are uniform: ``{message, id}``. + """ + + message: str + id: str | None = None + + +# --------------------------------------------------------------------------- +# Content blocks +# --------------------------------------------------------------------------- + + +class CohereTextContent(BaseModel): + type: Literal["text"] = "text" + text: str + + +class CohereThinkingContent(BaseModel): + type: Literal["thinking"] = "thinking" + thinking: str + + +class CohereImageUrl(BaseModel): + url: str + detail: Literal["auto", "low", "high"] | None = None + + +class CohereImageContent(BaseModel): + type: Literal["image_url"] = "image_url" + image_url: CohereImageUrl + + +CohereUserContentBlock = CohereTextContent | CohereImageContent +CohereAssistantContentBlock = CohereTextContent | CohereThinkingContent + + +# --------------------------------------------------------------------------- +# Documents (request side) and Document content (tool message side) +# --------------------------------------------------------------------------- + + +class CohereDocument(BaseModel): + """A document supplied for grounded generation.""" + + data: dict[str, Any] | str + id: str | None = None + + +class CohereDocumentContent(BaseModel): + type: Literal["document"] = "document" + document: CohereDocument + + +CohereToolContent = CohereTextContent | CohereDocumentContent + + +# --------------------------------------------------------------------------- +# Tool calls +# --------------------------------------------------------------------------- + + +class CohereToolCallFunction(BaseModel): + name: str | None = None + # Cohere serializes arguments as a JSON-encoded string (same as OpenAI). + arguments: str | None = None + + +class CohereToolCallV2(BaseModel): + id: str + type: Literal["function"] = "function" + function: CohereToolCallFunction | None = None + + +# --------------------------------------------------------------------------- +# Messages +# --------------------------------------------------------------------------- + + +class CohereUserMessageV2(BaseModel): + role: Literal["user"] = "user" + content: str | list[CohereUserContentBlock] + + +class CohereSystemMessageV2(BaseModel): + role: Literal["system"] = "system" + content: str | list[CohereTextContent] + + +class CohereAssistantMessageV2(BaseModel): + role: Literal["assistant"] = "assistant" + content: str | list[CohereAssistantContentBlock] | None = None + tool_calls: list[CohereToolCallV2] | None = None + tool_plan: str | None = Field( + default=None, + description=( + "Chain-of-thought plan emitted by the model alongside tool calls." + ), + ) + citations: list["CohereCitation"] | None = None + + +class CohereToolMessageV2(BaseModel): + role: Literal["tool"] = "tool" + tool_call_id: str + content: str | list[CohereToolContent] + + +CohereChatMessageV2 = ( + CohereUserMessageV2 + | CohereAssistantMessageV2 + | CohereSystemMessageV2 + | CohereToolMessageV2 +) + + +# --------------------------------------------------------------------------- +# Tools (request side) and tool choice +# --------------------------------------------------------------------------- + + +class CohereToolFunction(BaseModel): + name: str + description: str | None = None + parameters: dict[str, Any] = Field(default_factory=dict) + + +class CohereToolV2(BaseModel): + type: Literal["function"] = "function" + function: CohereToolFunction + + +CohereToolChoice = Literal["REQUIRED", "NONE"] + + +# --------------------------------------------------------------------------- +# Citations / sources +# --------------------------------------------------------------------------- + + +class CohereChatToolSource(BaseModel): + type: Literal["tool"] = "tool" + id: str | None = None + tool_output: dict[str, Any] | None = None + + +class CohereChatDocumentSource(BaseModel): + type: Literal["document"] = "document" + id: str | None = None + document: dict[str, Any] | None = None + + +CohereSource = CohereChatToolSource | CohereChatDocumentSource + + +class CohereCitation(BaseModel): + start: int | None = None + end: int | None = None + text: str | None = None + sources: list[CohereSource] | None = None + content_index: int | None = None + type: Literal["TEXT_CONTENT", "PLAN"] | None = None + + +# --------------------------------------------------------------------------- +# Response format / output config +# --------------------------------------------------------------------------- + + +class CohereTextResponseFormat(BaseModel): + type: Literal["text"] = "text" + + +class CohereJsonResponseFormat(BaseModel): + type: Literal["json_object"] = "json_object" + json_schema: dict[str, Any] | None = None + + +CohereResponseFormatV2 = CohereTextResponseFormat | CohereJsonResponseFormat + + +# --------------------------------------------------------------------------- +# Misc options +# --------------------------------------------------------------------------- + + +CohereSafetyMode = Literal["CONTEXTUAL", "STRICT", "OFF"] +CohereCitationMode = Literal["ENABLED", "DISABLED", "FAST", "ACCURATE", "OFF"] + + +class CohereCitationOptions(BaseModel): + mode: CohereCitationMode | None = None + + +class CohereThinking(BaseModel): + type: Literal["enabled", "disabled"] + token_budget: int | None = None + + +# --------------------------------------------------------------------------- +# Request +# --------------------------------------------------------------------------- + + +class CohereChatV2Request(BaseModel): + """Cohere Chat v2 request body. + + Mirrors the schema documented at + https://docs.cohere.com/reference/chat. + """ + + model: str + messages: list[CohereChatMessageV2] + stream: bool | None = False + + # Tooling + tools: list[CohereToolV2] | None = None + strict_tools: bool | None = None + tool_choice: CohereToolChoice | None = None + + # Grounding + documents: list[str | CohereDocument] | None = None + citation_options: CohereCitationOptions | None = None + + # Output + response_format: CohereResponseFormatV2 | None = None + safety_mode: CohereSafetyMode | None = None + max_tokens: int | None = None + stop_sequences: list[str] | None = None + + # Sampling + temperature: float | None = None + seed: int | None = None + frequency_penalty: float | None = None + presence_penalty: float | None = None + k: int | None = None + p: float | None = None + logprobs: bool | None = None + + # Reasoning + thinking: CohereThinking | None = None + + # Scheduling + priority: int | None = None + + # vLLM-specific extensions (not in Cohere spec). These mirror what the + # Anthropic and OpenAI surfaces already expose so that V2 callers can + # reach the same engine knobs when needed. + kv_transfer_params: dict[str, Any] | None = Field( + default=None, + description="KVTransfer parameters used for disaggregated serving.", + ) + chat_template_kwargs: dict[str, Any] | None = Field( + default=None, + description=( + "Additional keyword args to pass to the chat template renderer. " + "Will be accessible by the template." + ), + ) + + @field_validator("model") + @classmethod + def _validate_model(cls, v: str) -> str: + if not v: + raise ValueError("model is required") + return v + + @field_validator("max_tokens") + @classmethod + def _validate_max_tokens(cls, v: int | None) -> int | None: + if v is not None and v <= 0: + raise ValueError("max_tokens must be positive") + return v + + +# --------------------------------------------------------------------------- +# Usage +# --------------------------------------------------------------------------- + + +class CohereUsageBilledUnits(BaseModel): + input_tokens: float | None = None + output_tokens: float | None = None + search_units: float | None = None + classifications: float | None = None + + +class CohereUsageTokens(BaseModel): + input_tokens: float | None = None + output_tokens: float | None = None + + +class CohereUsage(BaseModel): + billed_units: CohereUsageBilledUnits | None = None + tokens: CohereUsageTokens | None = None + cached_tokens: float | None = None + + +# --------------------------------------------------------------------------- +# Logprobs +# --------------------------------------------------------------------------- + + +class CohereLogprobItem(BaseModel): + text: str | None = None + token_ids: list[int] + logprobs: list[float] | None = None + + +# --------------------------------------------------------------------------- +# Non-streaming response +# --------------------------------------------------------------------------- + + +CohereFinishReason = Literal[ + "COMPLETE", + "STOP_SEQUENCE", + "MAX_TOKENS", + "TOOL_CALL", + "ERROR", + "TIMEOUT", +] + + +class CohereAssistantMessageResponse(BaseModel): + role: Literal["assistant"] = "assistant" + content: list[CohereAssistantContentBlock] | None = None + tool_calls: list[CohereToolCallV2] | None = None + tool_plan: str | None = None + citations: list[CohereCitation] | None = None + + +class CohereChatV2Response(BaseModel): + id: str + finish_reason: CohereFinishReason + message: CohereAssistantMessageResponse + usage: CohereUsage | None = None + logprobs: list[CohereLogprobItem] | None = None + + # vLLM-specific extension. + kv_transfer_params: dict[str, Any] | None = Field( + default=None, description="KVTransfer parameters." + ) + + def model_post_init(self, __context: Any) -> None: + if not self.id: + self.id = f"chat_{int(time.time() * 1000)}" + + +# --------------------------------------------------------------------------- +# Streaming events +# --------------------------------------------------------------------------- +# +# Cohere V2 streams a sequence of typed JSON events delivered as +# Server-Sent Events. Each event is a discriminated union over the ``type`` +# field. See: +# https://docs.cohere.com/v2/docs/streaming +# and the OpenAPI ``StreamedChatResponseV2`` schema. + + +class _CohereStreamEventBase(BaseModel): + """Common base – every event has a ``type`` discriminator and the + nested ``delta`` shapes are loose dicts so we don't have to redefine the + full Cohere schema for each event variant.""" + + +class CohereChatMessageStartEvent(_CohereStreamEventBase): + type: Literal["message-start"] = "message-start" + id: str | None = None + delta: dict[str, Any] | None = None + + +class CohereChatContentStartEvent(_CohereStreamEventBase): + type: Literal["content-start"] = "content-start" + index: int | None = None + delta: dict[str, Any] | None = None + + +class CohereChatContentDeltaEvent(_CohereStreamEventBase): + type: Literal["content-delta"] = "content-delta" + index: int | None = None + delta: dict[str, Any] | None = None + logprobs: CohereLogprobItem | None = None + + +class CohereChatContentEndEvent(_CohereStreamEventBase): + type: Literal["content-end"] = "content-end" + index: int | None = None + + +class CohereChatToolPlanDeltaEvent(_CohereStreamEventBase): + type: Literal["tool-plan-delta"] = "tool-plan-delta" + delta: dict[str, Any] | None = None + + +class CohereChatToolCallStartEvent(_CohereStreamEventBase): + type: Literal["tool-call-start"] = "tool-call-start" + index: int | None = None + delta: dict[str, Any] | None = None + + +class CohereChatToolCallDeltaEvent(_CohereStreamEventBase): + type: Literal["tool-call-delta"] = "tool-call-delta" + index: int | None = None + delta: dict[str, Any] | None = None + + +class CohereChatToolCallEndEvent(_CohereStreamEventBase): + type: Literal["tool-call-end"] = "tool-call-end" + index: int | None = None + + +class CohereCitationStartEvent(_CohereStreamEventBase): + type: Literal["citation-start"] = "citation-start" + index: int | None = None + delta: dict[str, Any] | None = None + + +class CohereCitationEndEvent(_CohereStreamEventBase): + type: Literal["citation-end"] = "citation-end" + index: int | None = None + + +class CohereChatMessageEndEvent(_CohereStreamEventBase): + type: Literal["message-end"] = "message-end" + id: str | None = None + delta: dict[str, Any] | None = None + + +CohereStreamedChatResponseV2 = ( + CohereChatMessageStartEvent + | CohereChatContentStartEvent + | CohereChatContentDeltaEvent + | CohereChatContentEndEvent + | CohereChatToolPlanDeltaEvent + | CohereChatToolCallStartEvent + | CohereChatToolCallDeltaEvent + | CohereChatToolCallEndEvent + | CohereCitationStartEvent + | CohereCitationEndEvent + | CohereChatMessageEndEvent +) + + +class CohereStreamedChatResponseV2Envelope(RootModel[CohereStreamedChatResponseV2]): + """Discriminated-union wrapper for stream events.""" + + +# Resolve forward references created by ``CohereAssistantMessageV2.citations``. +CohereAssistantMessageV2.model_rebuild() diff --git a/vllm/entrypoints/cohere/serving.py b/vllm/entrypoints/cohere/serving.py new file mode 100644 index 000000000000..e6e925f18517 --- /dev/null +++ b/vllm/entrypoints/cohere/serving.py @@ -0,0 +1,960 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from +# https://github.com/vllm-project/vllm/blob/main/vllm/entrypoints/anthropic/serving.py +"""Cohere Chat v2 API serving handler. + +Implements ``POST /v2/chat`` by translating the incoming Cohere v2 request +into a standard :class:`ChatCompletionRequest` and delegating to +:class:`OpenAIServingChat`. The actual prompt rendering is handled by +vLLM's renderer pipeline (``vllm.renderers``): + +- For Cohere Command-family models, set ``--tokenizer-mode cohere`` and the + :class:`vllm.renderers.cohere.CohereRenderer` will template the request + via the ``cohere_melody`` library (``render_cmd3`` / ``render_cmd4``) + and surface citations through the standard ``ChatMessage.citations`` / + ``DeltaMessage.citations`` fields. +- For any other model the default Jinja-based :class:`HfRenderer` is used, + and the endpoint behaves as a plain v2-shaped wrapper around chat + completions. +""" + +from __future__ import annotations + +import json +import logging +import time +from collections.abc import AsyncGenerator +from typing import TYPE_CHECKING, Any + +from fastapi import Request + +from vllm.engine.protocol import EngineClient +from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption +from vllm.entrypoints.cohere.protocol import ( + CohereAssistantMessageResponse, + CohereAssistantMessageV2, + CohereChatContentDeltaEvent, + CohereChatContentEndEvent, + CohereChatContentStartEvent, + CohereChatMessageEndEvent, + CohereChatMessageStartEvent, + CohereChatToolCallDeltaEvent, + CohereChatToolCallEndEvent, + CohereChatToolCallStartEvent, + CohereChatV2Request, + CohereChatV2Response, + CohereCitation, + CohereCitationEndEvent, + CohereCitationStartEvent, + CohereFinishReason, + CohereSystemMessageV2, + CohereTextContent, + CohereThinkingContent, + CohereToolCallFunction, + CohereToolCallV2, + CohereToolMessageV2, + CohereUsage, + CohereUsageBilledUnits, + CohereUsageTokens, + CohereUserMessageV2, +) +from vllm.entrypoints.logger import RequestLogger +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionRequest, + ChatCompletionResponse, + ChatCompletionStreamResponse, + ChatCompletionToolsParam, +) +from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat +from vllm.entrypoints.openai.engine.protocol import ( + ErrorResponse, + JsonSchemaResponseFormat, + ResponseFormat, + StreamOptions, +) +from vllm.entrypoints.openai.models.serving import OpenAIServingModels + +if TYPE_CHECKING: + from vllm.entrypoints.serve.render.serving import OpenAIServingRender + +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# SSE helpers +# --------------------------------------------------------------------------- + + +def _sse(data: str) -> str: + """Wrap a JSON payload in a Server-Sent Event frame. + + Cohere's stream uses bare ``data:`` lines (no ``event:`` prefix); each + JSON object's ``type`` field carries the event discriminator. + """ + return f"data: {data}\n\n" + + +# Mapping of vLLM/OpenAI finish reasons to Cohere's enum. +_FINISH_REASON_MAP: dict[str | None, CohereFinishReason] = { + "stop": "COMPLETE", + "length": "MAX_TOKENS", + "tool_calls": "TOOL_CALL", + "stop_sequence": "STOP_SEQUENCE", + "error": "ERROR", + None: "COMPLETE", +} + + +def _map_finish_reason(reason: str | None) -> CohereFinishReason: + return _FINISH_REASON_MAP.get(reason, "COMPLETE") + + +# --------------------------------------------------------------------------- +# Serving class +# --------------------------------------------------------------------------- + + +class CohereServingChatV2(OpenAIServingChat): + """Handler for the Cohere Chat v2 API (``POST /v2/chat``). + + The handler is intentionally thin: it converts the v2 request into a + :class:`ChatCompletionRequest` (preserving Cohere-specific fields such + as ``documents``, ``safety_mode`` and ``citation_options`` via + ``chat_template_kwargs``) and delegates to the underlying chat + completion machinery. All Cohere-specific templating happens in the + renderer (:class:`vllm.renderers.cohere.CohereRenderer`) when the + engine is started with ``--tokenizer-mode cohere``; citations are read + natively from the resulting :class:`ChatMessage` / :class:`DeltaMessage`. + """ + + def __init__( + self, + engine_client: EngineClient, + models: OpenAIServingModels, + response_role: str, + *, + openai_serving_render: OpenAIServingRender, + request_logger: RequestLogger | None, + chat_template: str | None, + chat_template_content_format: ChatTemplateContentFormatOption, + return_tokens_as_token_ids: bool = False, + reasoning_parser: str = "", + enable_auto_tools: bool = False, + tool_parser: str | None = None, + enable_prompt_tokens_details: bool = False, + enable_force_include_usage: bool = False, + default_chat_template_kwargs: dict[str, Any] | None = None, + ) -> None: + super().__init__( + engine_client=engine_client, + models=models, + response_role=response_role, + openai_serving_render=openai_serving_render, + request_logger=request_logger, + chat_template=chat_template, + chat_template_content_format=chat_template_content_format, + return_tokens_as_token_ids=return_tokens_as_token_ids, + reasoning_parser=reasoning_parser, + enable_auto_tools=enable_auto_tools, + tool_parser=tool_parser, + enable_prompt_tokens_details=enable_prompt_tokens_details, + enable_force_include_usage=enable_force_include_usage, + default_chat_template_kwargs=default_chat_template_kwargs, + ) + + # ------------------------------------------------------------------ + # Public entry point + # ------------------------------------------------------------------ + + async def create_chat_v2( + self, + request: CohereChatV2Request, + raw_request: Request | None = None, + ) -> AsyncGenerator[str, None] | CohereChatV2Response | ErrorResponse: + """Implements ``POST /v2/chat``.""" + if logger.isEnabledFor(logging.DEBUG): + logger.debug( + "Received Cohere v2 chat request %s", request.model_dump_json() + ) + + chat_req = self._convert_v2_to_chat_completion(request) + if logger.isEnabledFor(logging.DEBUG): + logger.debug( + "Converted Cohere v2 -> ChatCompletion: %s", + chat_req.model_dump_json(), + ) + + generator = await self.create_chat_completion(chat_req, raw_request) + + if isinstance(generator, ErrorResponse): + return generator + if isinstance(generator, ChatCompletionResponse): + return self._chat_completion_to_v2(generator, request) + return self._chat_completion_stream_to_v2(generator, request) + + # ================================================================== + # Request conversion: Cohere V2 -> ChatCompletionRequest + # ================================================================== + + @classmethod + def _convert_v2_to_chat_completion( + cls, request: CohereChatV2Request + ) -> ChatCompletionRequest: + openai_messages: list[dict[str, Any]] = [] + cls._convert_messages(request.messages, openai_messages) + + chat_req = cls._build_base_chat_completion(request, openai_messages) + cls._apply_streaming_options(chat_req, request) + cls._apply_response_format(chat_req, request) + cls._apply_tools(chat_req, request) + cls._apply_tool_choice(chat_req, request) + cls._apply_cohere_template_kwargs(chat_req, request) + return chat_req + + @classmethod + def _convert_messages( + cls, + messages: list, + openai_messages: list[dict[str, Any]], + ) -> None: + for msg in messages: + if isinstance(msg, CohereSystemMessageV2): + openai_messages.append( + { + "role": "system", + "content": cls._coerce_text_content(msg.content), + } + ) + elif isinstance(msg, CohereUserMessageV2): + openai_messages.append(cls._convert_user_message(msg)) + elif isinstance(msg, CohereAssistantMessageV2): + openai_messages.append(cls._convert_assistant_message(msg)) + elif isinstance(msg, CohereToolMessageV2): + openai_messages.append(cls._convert_tool_message(msg)) + else: # pragma: no cover - guarded by Pydantic discriminator + raise ValueError(f"Unsupported Cohere v2 message: {msg!r}") + + @staticmethod + def _coerce_text_content(content: str | list[Any]) -> str: + if isinstance(content, str): + return content + parts: list[str] = [] + for block in content: + text = getattr(block, "text", None) + if text: + parts.append(text) + return "".join(parts) + + @classmethod + def _convert_user_message( + cls, msg: CohereUserMessageV2 + ) -> dict[str, Any]: + if isinstance(msg.content, str): + return {"role": "user", "content": msg.content} + + content_parts: list[dict[str, Any]] = [] + for block in msg.content: + if isinstance(block, CohereTextContent): + content_parts.append({"type": "text", "text": block.text}) + else: # CohereImageContent + content_parts.append( + { + "type": "image_url", + "image_url": {"url": block.image_url.url}, + } + ) + if ( + len(content_parts) == 1 + and content_parts[0]["type"] == "text" + ): + return {"role": "user", "content": content_parts[0]["text"]} + return {"role": "user", "content": content_parts} + + @classmethod + def _convert_assistant_message( + cls, msg: CohereAssistantMessageV2 + ) -> dict[str, Any]: + out: dict[str, Any] = {"role": "assistant"} + + # Cohere splits reasoning out into ``thinking`` content blocks. We + # collapse them back into the OpenAI ``reasoning`` field for the + # downstream chat template, while text blocks become ``content``. + text_parts: list[str] = [] + thinking_parts: list[str] = [] + if isinstance(msg.content, str): + text_parts.append(msg.content) + elif msg.content is not None: + for block in msg.content: + if isinstance(block, CohereTextContent): + text_parts.append(block.text) + elif isinstance(block, CohereThinkingContent): + thinking_parts.append(block.thinking) + + # ``tool_plan`` is Cohere's chain-of-thought emitted alongside tool + # calls; preserve it as reasoning so templates that expect a + # planning block still see it. + if msg.tool_plan: + thinking_parts.append(msg.tool_plan) + + if text_parts: + out["content"] = "".join(text_parts) + if thinking_parts: + out["reasoning"] = "".join(thinking_parts) + + if msg.tool_calls: + out["tool_calls"] = [ + { + "id": tc.id, + "type": "function", + "function": { + "name": (tc.function.name if tc.function else "") or "", + "arguments": ( + tc.function.arguments if tc.function else None + ) + or "{}", + }, + } + for tc in msg.tool_calls + ] + return out + + @classmethod + def _convert_tool_message( + cls, msg: CohereToolMessageV2 + ) -> dict[str, Any]: + if isinstance(msg.content, str): + content: Any = msg.content + else: + text_parts: list[str] = [] + for block in msg.content: + if isinstance(block, CohereTextContent): + text_parts.append(block.text) + else: # CohereDocumentContent + # The OpenAI chat-completion shape only knows about + # text tool results. Serialize the document so its data + # is preserved in the prompt. + text_parts.append( + json.dumps(block.document.model_dump(exclude_none=True)) + ) + content = "\n".join(text_parts) + return { + "role": "tool", + "tool_call_id": msg.tool_call_id, + "content": content, + } + + @classmethod + def _build_base_chat_completion( + cls, + request: CohereChatV2Request, + openai_messages: list[dict[str, Any]], + ) -> ChatCompletionRequest: + return ChatCompletionRequest( + model=request.model, + messages=openai_messages, + max_tokens=request.max_tokens, + max_completion_tokens=request.max_tokens, + stop=request.stop_sequences, + temperature=request.temperature, + top_p=request.p, + top_k=request.k, + seed=request.seed, + frequency_penalty=request.frequency_penalty, + presence_penalty=request.presence_penalty, + logprobs=request.logprobs, + priority=request.priority or 0, + kv_transfer_params=request.kv_transfer_params, + chat_template_kwargs=request.chat_template_kwargs, + ) + + @classmethod + def _apply_streaming_options( + cls, + chat_req: ChatCompletionRequest, + request: CohereChatV2Request, + ) -> None: + if request.stream: + chat_req.stream = True + chat_req.stream_options = StreamOptions.model_validate( + {"include_usage": True, "continuous_usage_stats": True} + ) + + @classmethod + def _apply_response_format( + cls, + chat_req: ChatCompletionRequest, + request: CohereChatV2Request, + ) -> None: + rf = request.response_format + if rf is None or rf.type == "text": + return + chat_req.response_format = ResponseFormat( + type="json_schema" if rf.json_schema else "json_object", + json_schema=( + JsonSchemaResponseFormat( + name="cohere_v2_json_schema", + json_schema=rf.json_schema, + ) + if rf.json_schema + else None + ), + ) + + @classmethod + def _apply_tools( + cls, + chat_req: ChatCompletionRequest, + request: CohereChatV2Request, + ) -> None: + if not request.tools: + return + chat_req.tools = [ + ChatCompletionToolsParam.model_validate( + { + "type": "function", + "function": { + "name": tool.function.name, + "description": tool.function.description, + "parameters": tool.function.parameters, + }, + } + ) + for tool in request.tools + ] + + @classmethod + def _apply_tool_choice( + cls, + chat_req: ChatCompletionRequest, + request: CohereChatV2Request, + ) -> None: + if request.tool_choice == "REQUIRED": + chat_req.tool_choice = "required" + elif request.tool_choice == "NONE": + chat_req.tool_choice = "none" + elif chat_req.tools: + # Mirrors Cohere's "free choice" default when tools are present. + chat_req.tool_choice = "auto" + + @classmethod + def _apply_cohere_template_kwargs( + cls, + chat_req: ChatCompletionRequest, + request: CohereChatV2Request, + ) -> None: + """Forward Cohere-specific request fields into ``chat_template_kwargs``. + + The :class:`vllm.renderers.cohere.CohereRenderer` consumes these + kwargs to drive ``cohere_melody.render_cmd3`` / ``render_cmd4``. + Other renderers ignore unknown kwargs, so this is also a no-op for + non-cohere ``--tokenizer-mode`` settings. + """ + kwargs = dict(chat_req.chat_template_kwargs or {}) + + if request.documents: + documents: list[dict[str, Any]] = [] + for idx, doc in enumerate(request.documents): + if isinstance(doc, str): + documents.append({"id": f"doc_{idx}", "data": {"text": doc}}) + else: + documents.append( + { + "id": doc.id or f"doc_{idx}", + "data": ( + doc.data + if isinstance(doc.data, dict) + else {"text": doc.data} + ), + } + ) + kwargs.setdefault("documents", documents) + + if request.safety_mode is not None: + kwargs.setdefault("safety_mode", str(request.safety_mode).lower()) + if request.citation_options is not None: + kwargs.setdefault( + "citation_options", + request.citation_options.model_dump(exclude_none=True), + ) + if request.thinking is not None: + kwargs.setdefault( + "thinking", + request.thinking.model_dump(exclude_none=True), + ) + + if kwargs: + chat_req.chat_template_kwargs = kwargs + + # ================================================================== + # Response conversion: ChatCompletion -> Cohere V2 + # ================================================================== + + def _chat_completion_to_v2( + self, + response: ChatCompletionResponse, + request: CohereChatV2Request, + ) -> CohereChatV2Response: + choice = response.choices[0] + msg = choice.message + + content_blocks: list[Any] = [] + if msg.reasoning: + content_blocks.append(CohereThinkingContent(thinking=msg.reasoning)) + if msg.content: + content_blocks.append(CohereTextContent(text=msg.content)) + + tool_calls: list[CohereToolCallV2] | None = None + if msg.tool_calls: + tool_calls = [ + CohereToolCallV2( + id=tc.id, + function=CohereToolCallFunction( + name=tc.function.name, + arguments=tc.function.arguments, + ), + ) + for tc in msg.tool_calls + ] + + # Cohere's ``tool_plan`` is the planning text emitted before tool + # calls. The closest analogue in chat completions is the reasoning + # field; if a tool call is present, surface reasoning as tool_plan + # rather than as a thinking content block to match Cohere clients. + tool_plan: str | None = None + if tool_calls and msg.reasoning: + tool_plan = msg.reasoning + content_blocks = [ + blk + for blk in content_blocks + if not isinstance(blk, CohereThinkingContent) + ] + + assistant_msg = CohereAssistantMessageResponse( + content=content_blocks or None, + tool_calls=tool_calls, + tool_plan=tool_plan, + citations=self._extract_citations_if_any(msg), + ) + + usage = self._build_usage(response) + + return CohereChatV2Response( + id=response.id or f"chat_{int(time.time() * 1000)}", + finish_reason=_map_finish_reason(choice.finish_reason), + message=assistant_msg, + usage=usage, + kv_transfer_params=response.kv_transfer_params, + ) + + @staticmethod + def _extract_citations_if_any(msg: Any) -> list[CohereCitation] | None: + """Coerce ``ChatMessage.citations`` into the Cohere v2 wire shape. + + ``ChatMessage`` natively carries a ``citations: list[Citation] | + None`` field (see :mod:`vllm.entrypoints.openai.engine.protocol`). + Renderers/parsers (e.g. the ``cohere2`` reasoning parser) populate + it. We map each :class:`vllm...Citation` into the + :class:`CohereCitation` wire model, preserving sources and span. + """ + raw = getattr(msg, "citations", None) + if not raw: + return None + out: list[CohereCitation] = [] + for c in raw: + try: + if isinstance(c, CohereCitation): + out.append(c) + continue + if hasattr(c, "model_dump"): + c = c.model_dump(exclude_none=True) + if isinstance(c, dict): + out.append(CohereCitation.model_validate(c)) + except Exception: # pragma: no cover - defensive + logger.debug("Skipping malformed citation: %r", c, exc_info=True) + return out or None + + @staticmethod + def _build_usage(response: ChatCompletionResponse) -> CohereUsage | None: + if response.usage is None: + return None + prompt = response.usage.prompt_tokens + completion = response.usage.completion_tokens or 0 + return CohereUsage( + billed_units=CohereUsageBilledUnits( + input_tokens=prompt, + output_tokens=completion, + ), + tokens=CohereUsageTokens( + input_tokens=prompt, + output_tokens=completion, + ), + ) + + # ================================================================== + # Stream conversion: chat completion stream -> Cohere V2 SSE events + # ================================================================== + + async def _chat_completion_stream_to_v2( + self, + generator: AsyncGenerator[str, None], + request: CohereChatV2Request, + ) -> AsyncGenerator[str, None]: + """Translate an OpenAI-style chat completion SSE stream into Cohere's + v2 stream-event format. + + Cohere's v2 stream lifecycle is: + + message-start + [content-start, content-delta..., content-end]* + [tool-plan-delta]* + [tool-call-start, tool-call-delta..., tool-call-end]* + message-end + """ + state = _StreamState() + + try: + async for item in generator: + if not item.startswith("data:"): + continue + data_str = item[len("data:"):].strip().rstrip("\n") + if not data_str: + continue + if data_str == "[DONE]": + # Cohere closes the stream with message-end (already + # emitted on the final usage chunk); nothing to do. + return + + chunk = ChatCompletionStreamResponse.model_validate_json(data_str) + + if not state.started: + yield _sse( + CohereChatMessageStartEvent( + id=chunk.id, + delta={"message": {"role": "assistant"}}, + ).model_dump_json(exclude_none=True) + ) + state.started = True + + # The final OpenAI chunk has no choices and only carries usage. + if not chunk.choices: + yield from self._close_open_blocks(state) + yield _sse( + self._build_message_end_event( + chunk_id=chunk.id, + finish_reason=state.finish_reason, + usage_chunk=chunk, + ) + ) + continue + + choice = chunk.choices[0] + if choice.finish_reason is not None: + state.finish_reason = choice.finish_reason + + delta = choice.delta + + # Reasoning -> thinking content block + reasoning = getattr(delta, "reasoning", None) or getattr( + delta, "reasoning_content", None + ) + if reasoning: + for ev in self._handle_thinking_delta(state, reasoning): + yield ev + + if delta.content: + for ev in self._handle_text_delta(state, delta.content): + yield ev + + if delta.tool_calls: + for ev in self._handle_tool_call_deltas( + state, delta.tool_calls + ): + yield ev + + # Citations: a Cohere-specific extension on DeltaMessage that + # the cohere renderer/parsers may populate. + if getattr(delta, "citations", None): + for ev in self._handle_citation_deltas( + state, delta.citations + ): + yield ev + + except Exception as exc: + logger.exception("Error converting chat completion stream to v2") + yield _sse( + json.dumps( + { + "type": "message-end", + "delta": {"error": str(exc), "finish_reason": "ERROR"}, + } + ) + ) + + # -- per-delta helpers -------------------------------------------- + + def _handle_thinking_delta( + self, state: _StreamState, delta_text: str + ) -> list[str]: + events: list[str] = [] + if state.active_block != "thinking": + events.extend(self._close_open_blocks_inner(state)) + idx = state.next_content_index() + state.active_block = "thinking" + state.active_block_index = idx + events.append( + _sse( + CohereChatContentStartEvent( + index=idx, + delta={ + "message": { + "content": {"type": "thinking", "thinking": ""} + } + }, + ).model_dump_json(exclude_none=True) + ) + ) + events.append( + _sse( + CohereChatContentDeltaEvent( + index=state.active_block_index, + delta={ + "message": {"content": {"thinking": delta_text}} + }, + ).model_dump_json(exclude_none=True) + ) + ) + return events + + def _handle_text_delta( + self, state: _StreamState, delta_text: str + ) -> list[str]: + events: list[str] = [] + if state.active_block != "text": + events.extend(self._close_open_blocks_inner(state)) + idx = state.next_content_index() + state.active_block = "text" + state.active_block_index = idx + events.append( + _sse( + CohereChatContentStartEvent( + index=idx, + delta={ + "message": { + "content": {"type": "text", "text": ""} + } + }, + ).model_dump_json(exclude_none=True) + ) + ) + events.append( + _sse( + CohereChatContentDeltaEvent( + index=state.active_block_index, + delta={"message": {"content": {"text": delta_text}}}, + ).model_dump_json(exclude_none=True) + ) + ) + return events + + def _handle_tool_call_deltas( + self, state: _StreamState, deltas: list + ) -> list[str]: + events: list[str] = [] + for tc in deltas: + tc_index = tc.index + tool_id = tc.id + fn = tc.function + + if tool_id is not None and tc_index not in state.tool_calls_seen: + # New tool call. Close any open content/tool block first. + events.extend(self._close_open_blocks_inner(state)) + state.tool_calls_seen.add(tc_index) + state.active_tool_index = tc_index + state.active_block = "tool_call" + events.append( + _sse( + CohereChatToolCallStartEvent( + index=tc_index, + delta={ + "message": { + "tool_calls": { + "id": tool_id, + "type": "function", + "function": { + "name": (fn.name if fn else "") + or "", + "arguments": ( + fn.arguments if fn else None + ) + or "", + }, + } + } + }, + ).model_dump_json(exclude_none=True) + ) + ) + continue + + if fn and fn.arguments: + events.append( + _sse( + CohereChatToolCallDeltaEvent( + index=tc_index, + delta={ + "message": { + "tool_calls": { + "function": { + "arguments": fn.arguments, + } + } + } + }, + ).model_dump_json(exclude_none=True) + ) + ) + return events + + def _handle_citation_deltas( + self, state: _StreamState, citations: list + ) -> list[str]: + """Emit ``citation-start`` / ``citation-end`` events for a delta. + + The cohere renderer/parsers populate ``DeltaMessage.citations`` with + :class:`vllm.entrypoints.openai.engine.protocol.Citation` instances + once a citation has fully resolved (start/end indices known). We + emit a complete start+end pair for each citation in the delta so + Cohere clients can attach the annotation to the surrounding text. + """ + events: list[str] = [] + for c in citations: + payload = ( + c.model_dump(exclude_none=True) if hasattr(c, "model_dump") else c + ) + if not isinstance(payload, dict): + continue + try: + citation = CohereCitation.model_validate(payload) + except Exception: # pragma: no cover - defensive + logger.debug("Skipping malformed streamed citation: %r", payload) + continue + idx = state.next_citation_index() + events.append( + _sse( + CohereCitationStartEvent( + index=idx, + delta={"message": {"citations": citation.model_dump( + exclude_none=True + )}}, + ).model_dump_json(exclude_none=True) + ) + ) + events.append( + _sse( + CohereCitationEndEvent(index=idx).model_dump_json( + exclude_none=True + ) + ) + ) + return events + + # -- block lifecycle helpers -------------------------------------- + + def _close_open_blocks(self, state: _StreamState) -> list[str]: + return self._close_open_blocks_inner(state) + + def _close_open_blocks_inner(self, state: _StreamState) -> list[str]: + events: list[str] = [] + if state.active_block in ("text", "thinking"): + events.append( + _sse( + CohereChatContentEndEvent( + index=state.active_block_index, + ).model_dump_json(exclude_none=True) + ) + ) + elif state.active_block == "tool_call": + events.append( + _sse( + CohereChatToolCallEndEvent( + index=state.active_tool_index, + ).model_dump_json(exclude_none=True) + ) + ) + state.active_block = None + state.active_block_index = None + state.active_tool_index = None + return events + + def _build_message_end_event( + self, + chunk_id: str, + finish_reason: str | None, + usage_chunk: ChatCompletionStreamResponse, + ) -> str: + delta: dict[str, Any] = { + "finish_reason": _map_finish_reason(finish_reason), + } + if usage_chunk.usage is not None: + prompt = usage_chunk.usage.prompt_tokens + completion = usage_chunk.usage.completion_tokens or 0 + delta["usage"] = { + "billed_units": { + "input_tokens": prompt, + "output_tokens": completion, + }, + "tokens": { + "input_tokens": prompt, + "output_tokens": completion, + }, + } + return CohereChatMessageEndEvent( + id=chunk_id, delta=delta + ).model_dump_json(exclude_none=True) + + # ================================================================== + # Helpers for the router + # ================================================================== + + def create_error_response(self, message: str) -> ErrorResponse: + # Reuse the OpenAI engine's error response so that the router can + # translate it into Cohere's error envelope uniformly. + from vllm.entrypoints.openai.engine.protocol import ErrorInfo + + return ErrorResponse( + error=ErrorInfo( + message=message, + type="bad_request", + code=400, + ) + ) + + +# --------------------------------------------------------------------------- +# Stream state +# --------------------------------------------------------------------------- + + +class _StreamState: + """Tracks which Cohere v2 stream block (if any) is currently open.""" + + def __init__(self) -> None: + self.started: bool = False + self.finish_reason: str | None = None + self.active_block: str | None = None # "text" | "thinking" | "tool_call" + self.active_block_index: int | None = None + self.active_tool_index: int | None = None + self._next_index: int = 0 + self._next_citation_index: int = 0 + self.tool_calls_seen: set[int] = set() + + def next_content_index(self) -> int: + idx = self._next_index + self._next_index += 1 + return idx + + def next_citation_index(self) -> int: + idx = self._next_citation_index + self._next_citation_index += 1 + return idx diff --git a/vllm/entrypoints/generate/api_router.py b/vllm/entrypoints/generate/api_router.py index 713e2566bc56..5fdc56bd8264 100644 --- a/vllm/entrypoints/generate/api_router.py +++ b/vllm/entrypoints/generate/api_router.py @@ -41,6 +41,12 @@ def register_generate_api_routers(app: FastAPI): register_anthropic_api_router(app) + from vllm.entrypoints.cohere.api_router import ( + attach_router as register_cohere_api_router, + ) + + register_cohere_api_router(app) + from .generative_scoring.api_router import register_generative_scoring_api_router register_generative_scoring_api_router(app) @@ -55,6 +61,7 @@ async def init_generate_state( ): from vllm.entrypoints.anthropic.serving import AnthropicServingMessages from vllm.entrypoints.chat_utils import load_chat_template + from vllm.entrypoints.cohere.serving import CohereServingChatV2 from vllm.entrypoints.mcp.tool_server import ( DemoToolServer, MCPToolServer, @@ -175,6 +182,26 @@ async def init_generate_state( if "generate" in supported_tasks else None ) + state.cohere_serving_chat_v2 = ( + CohereServingChatV2( + engine_client, + state.openai_serving_models, + args.response_role, + openai_serving_render=state.openai_serving_render, + request_logger=request_logger, + chat_template=resolved_chat_template, + chat_template_content_format=args.chat_template_content_format, + return_tokens_as_token_ids=args.return_tokens_as_token_ids, + enable_auto_tools=args.enable_auto_tool_choice, + tool_parser=args.tool_call_parser, + reasoning_parser=args.structured_outputs_config.reasoning_parser, + enable_prompt_tokens_details=args.enable_prompt_tokens_details, + enable_force_include_usage=args.enable_force_include_usage, + default_chat_template_kwargs=args.default_chat_template_kwargs, + ) + if "generate" in supported_tasks + else None + ) state.serving_tokens = ( ServingTokens( engine_client, diff --git a/vllm/entrypoints/openai/chat_completion/protocol.py b/vllm/entrypoints/openai/chat_completion/protocol.py index 184ace56805c..8f919e2fb978 100644 --- a/vllm/entrypoints/openai/chat_completion/protocol.py +++ b/vllm/entrypoints/openai/chat_completion/protocol.py @@ -21,6 +21,7 @@ ) from vllm.entrypoints.openai.engine.protocol import ( AnyResponseFormat, + Citation, DeltaMessage, FunctionCall, FunctionDefinition, @@ -65,6 +66,11 @@ class ChatMessage(OpenAIBaseModel): # vLLM-specific fields that are not in OpenAI spec reasoning: str | None = None + # Citations grounding the message content in source material. Populated + # by parsers/renderers for grounded models (e.g. Cohere Command). Left + # unset (and therefore omitted from JSON) for ungrounded models so + # OpenAI-compatible clients see the standard shape. + citations: list[Citation] | None = None class ChatCompletionLogProb(OpenAIBaseModel): diff --git a/vllm/entrypoints/openai/engine/protocol.py b/vllm/entrypoints/openai/engine/protocol.py index 434888df9efa..c8d3f34998cc 100644 --- a/vllm/entrypoints/openai/engine/protocol.py +++ b/vllm/entrypoints/openai/engine/protocol.py @@ -339,11 +339,57 @@ class ExtractedToolCallInformation(BaseModel): content: str | None = None +class CitationSource(OpenAIBaseModel): + """Source attribution for a :class:`Citation`. + + Mirrors the shape used by Cohere's Chat v2 API. ``type`` is the source + discriminator (``document`` or ``tool``); ``id`` is the citing + document/tool-output identifier; ``document`` and ``tool_output`` carry + the original payload that produced the citation. + """ + + type: Literal["document", "tool"] + id: str | None = None + document: dict[str, Any] | None = None + tool_output: dict[str, Any] | None = None + + +class Citation(OpenAIBaseModel): + """A citation grounding a span of generated text in source material. + + This is a vLLM-specific extension to the OpenAI chat completion shape + so that grounded models (e.g. Cohere Command-family) can surface + citation metadata through the standard chat completion responses. It is + consumed by surfaces that expose citations (such as Cohere's + ``/v2/chat`` endpoint). + """ + + start: int | None = None + """Start character offset in the surrounding text content.""" + end: int | None = None + """End character offset (exclusive) in the surrounding text content.""" + text: str | None = None + """The cited text snippet.""" + sources: list[CitationSource] = Field(default_factory=list) + """Source documents / tool outputs that ground this citation.""" + content_index: int | None = None + """Index of the content block this citation refers to (when the message + has multiple content blocks).""" + type: Literal["TEXT_CONTENT", "PLAN"] | None = None + """Whether this citation grounds the user-visible text or a planning + block. ``None`` means unspecified.""" + + class DeltaMessage(OpenAIBaseModel): role: str | None = None content: str | None = None reasoning: str | None = None tool_calls: list[DeltaToolCall] = Field(default_factory=list) + # vLLM-specific extension: citations grounding the streamed content. + # Surfaces (e.g. Cohere Chat v2) consume this when present; OpenAI + # clients ignore it. Default is ``None`` so it is omitted from the wire + # for non-grounded models. + citations: list[Citation] | None = None class GenerationError(Exception): diff --git a/vllm/renderers/cohere.py b/vllm/renderers/cohere.py new file mode 100644 index 000000000000..7893def825a6 --- /dev/null +++ b/vllm/renderers/cohere.py @@ -0,0 +1,452 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Cohere prompt renderer. + +Templates the Cohere Command-family prompt formats (cmd3 / cmd4) using the +``cohere_melody`` Rust bindings instead of Jinja. Selecting this renderer is +a matter of setting ``--tokenizer-mode cohere`` on the engine; tokenization +itself still flows through the cached HuggingFace tokenizer. + +This renderer intentionally accepts the same ``chat_template_kwargs`` shape +used by the standard chat completions endpoint, plus a few Cohere-specific +fields that grounding/citation features require: + +* ``documents``: list of document dicts to expose to the model +* ``available_tools``: list of tool dicts (overrides ``tools``) +* ``safety_mode``: cmd3-only safety mode (``contextual`` / ``strict`` / ``none``) +* ``citation_quality``: cmd3 citation toggle (``on`` / ``off``) +* ``citation_options.mode``: cmd4 grounding (``fast`` / ``accurate`` / ``off``) +* ``reasoning_type``: ``enabled`` / ``disabled`` +* ``response_prefix``: optional response prefix +* ``json_schema`` / ``json_mode`` / ``response_format``: structured outputs +* ``template_id`` / ``template`` / ``template_jinja`` / ``use_jinja``: + template overrides +* ``cohere_format``: ``cmd3`` (default) or ``cmd4`` + +Citations produced by Cohere models are surfaced through the standard +``ChatMessage.citations`` / ``DeltaMessage.citations`` fields (populated by +the ``cohere2`` reasoning parser). +""" +from __future__ import annotations + +import copy +import json +from typing import TYPE_CHECKING, Any + +from vllm.config import VllmConfig +from vllm.entrypoints.chat_utils import ( + ChatCompletionMessageParam, + ConversationMessage, + parse_chat_messages, + parse_chat_messages_async, +) +from vllm.logger import init_logger +from vllm.tokenizers.hf import HfTokenizer +from vllm.utils.async_utils import make_async + +from .base import BaseRenderer +from .inputs import DictPrompt +from .inputs.preprocess import parse_dec_only_prompt +from .params import ChatParams + +if TYPE_CHECKING: + pass + +logger = init_logger(__name__) + + +_DEFAULT_FORMAT = "cmd3" +_VALID_FORMATS = ("cmd3", "cmd4") + +# Keys consumed by this renderer that should not also be forwarded to the +# `apply_chat_template` path of the underlying tokenizer. +_RENDERER_CONSUMED_KEYS = frozenset( + { + "cohere_format", + "documents", + "available_tools", + "tools", + "safety_mode", + "citation_quality", + "citation_options", + "grounding", + "reasoning_type", + "response_prefix", + "json_schema", + "json_mode", + "response_format", + "skip_preamble", + "dev_instruction", + "platform_instruction", + "template", + "template_id", + "template_jinja", + "use_jinja", + "additional_template_fields", + "escaped_special_tokens", + } +) + + +def _try_import_melody(): + try: + import cohere_melody # type: ignore + + return cohere_melody + except ImportError as e: # pragma: no cover - exercised at runtime + raise ImportError( + "The `cohere` tokenizer/renderer mode requires the " + "`cohere_melody` package. Install it via " + "`pip install cohere-melody` or build from " + "https://github.com/cohere-ai/melody." + ) from e + + +def _role_to_melody(role: str) -> str: + """Map an OpenAI role to the role string melody expects.""" + if role == "assistant": + return "chatbot" + return role + + +def _normalize_tool_call(tc: dict[str, Any] | Any) -> dict[str, Any]: + """Normalize an OpenAI tool call dict into melody's tool call shape. + + melody expects ``{id, name, parameters: }`` whereas OpenAI + delivers ``{id, type, function: {name, arguments: }}``. + """ + # Pydantic objects + if hasattr(tc, "model_dump"): + tc = tc.model_dump() + if not isinstance(tc, dict): + raise TypeError(f"Unexpected tool_call value: {tc!r}") + + fn = tc.get("function") or {} + name = fn.get("name") or tc.get("name") or "" + args = fn.get("arguments") + if args is None: + args = tc.get("arguments", {}) + # melody expects a JSON-encoded string + if not isinstance(args, str): + args = json.dumps(args, ensure_ascii=False) + return { + "id": tc.get("id") or "", + "name": name, + "parameters": args, + } + + +def _content_blocks(content: Any) -> list[dict[str, Any]]: + """Convert OpenAI ``ConversationMessage.content`` into melody content blocks. + + The chat_utils ``content_format="openai"`` produces a list of + ``{type: text, text: ...}`` dicts already; we pass those through with + minimal coercion. Plain string content is wrapped in a single text block. + Image / multimodal placeholder dicts (``{"type": "image"}``) are + forwarded as-is so that templates can reference image placeholders that + the upstream tokenizer expands separately. + """ + if content is None: + return [] + if isinstance(content, str): + return [{"type": "text", "text": content}] + blocks: list[dict[str, Any]] = [] + for part in content: + if isinstance(part, str): + blocks.append({"type": "text", "text": part}) + continue + if not isinstance(part, dict): + raise TypeError(f"Unexpected content part: {part!r}") + part_type = part.get("type", "text") + if part_type in ("text", "input_text", "output_text", "refusal"): + blocks.append({"type": "text", "text": part.get("text", "")}) + elif part_type == "thinking": + blocks.append( + {"type": "thinking", "thinking": part.get("thinking", "")} + ) + elif part_type == "image": + blocks.append( + { + "type": "image", + "image": { + "template_placeholder": part.get( + "template_placeholder", "" + ), + }, + } + ) + elif part_type == "document": + doc = part.get("document") + if isinstance(doc, dict): + blocks.append({"type": "document", "document": doc}) + else: + # Fall back to wrapping arbitrary string as text. + blocks.append({"type": "text", "text": json.dumps(doc)}) + elif part_type == "tool_reference": + # Tool references are rendered by name; emit as text so the + # renderer downstream is content-format agnostic. + blocks.append( + {"type": "text", "text": part.get("name") or part.get("text", "")} + ) + else: + # Unknown block type: render as text fallback. + text = part.get("text") or part.get(part_type) or "" + text_str = text if isinstance(text, str) else json.dumps(text) + blocks.append({"type": "text", "text": text_str}) + return blocks + + +def _document_to_melody(doc: Any) -> dict[str, Any]: + """Coerce a Cohere v2 document into melody's ``Document`` (dict) shape.""" + if isinstance(doc, str): + return {"text": doc} + if isinstance(doc, dict): + # Cohere v2 wraps documents in {id, data: {...}} or pure dicts. + if "data" in doc and isinstance(doc["data"], dict): + payload = dict(doc["data"]) + if "id" in doc and "id" not in payload: + payload["id"] = doc["id"] + return payload + return dict(doc) + raise TypeError(f"Unsupported document type: {type(doc).__name__}") + + +def _tool_to_melody(tool: Any) -> dict[str, Any]: + """Coerce a chat completions tool definition into melody's ``Tool`` shape. + + Accepts either the raw OpenAI tool wrapper ``{type:"function", function: + {name, description, parameters}}`` or a flat ``{name, description, + parameters}`` dict (which is what melody itself expects). + """ + if hasattr(tool, "model_dump"): + tool = tool.model_dump() + if not isinstance(tool, dict): + raise TypeError(f"Unsupported tool type: {type(tool).__name__}") + if "function" in tool and isinstance(tool["function"], dict): + fn = tool["function"] + else: + fn = tool + return { + "name": fn.get("name", ""), + "description": fn.get("description", "") or "", + "parameters": fn.get("parameters") or {}, + } + + +def _conversation_to_melody_messages( + conversation: list[ConversationMessage], +) -> list[dict[str, Any]]: + out: list[dict[str, Any]] = [] + for msg in conversation: + role = _role_to_melody(msg.get("role", "user")) + content_blocks = _content_blocks(msg.get("content")) + + # Treat reasoning content as a thinking block on assistant turns + # so multi-turn thinking is preserved in the rendered prompt. + reasoning = msg.get("reasoning") or msg.get("reasoning_content") + if role == "chatbot" and reasoning: + content_blocks.insert(0, {"type": "thinking", "thinking": reasoning}) + + tool_calls = [ + _normalize_tool_call(tc) for tc in (msg.get("tool_calls") or []) + ] + + out_msg: dict[str, Any] = { + "role": role, + "content": content_blocks, + "tool_calls": tool_calls, + } + tool_call_id = msg.get("tool_call_id") + if tool_call_id: + out_msg["tool_call_id"] = tool_call_id + out.append(out_msg) + return out + + +def _build_render_config( + conversation: list[ConversationMessage], + chat_template_kwargs: dict[str, Any], +) -> tuple[str, dict[str, Any]]: + """Build the ``render_cmd3`` / ``render_cmd4`` config dict. + + Returns ``(format, config_dict)`` where ``format`` is either ``"cmd3"`` + or ``"cmd4"``. + """ + fmt = chat_template_kwargs.get("cohere_format", _DEFAULT_FORMAT) + if fmt not in _VALID_FORMATS: + raise ValueError( + f"Invalid cohere_format={fmt!r}; expected one of {_VALID_FORMATS}" + ) + + config: dict[str, Any] = { + "messages": _conversation_to_melody_messages(conversation), + } + + # Optional template overrides + for k in ("template_id", "template", "template_jinja"): + if v := chat_template_kwargs.get(k): + config[k] = v + if (use_jinja := chat_template_kwargs.get("use_jinja")) is not None: + config["use_jinja"] = bool(use_jinja) + + # Documents + documents = chat_template_kwargs.get("documents") or [] + if documents: + config["documents"] = [_document_to_melody(d) for d in documents] + + # Tools - prefer explicit ``available_tools``, fall back to OpenAI ``tools`` + tools = ( + chat_template_kwargs.get("available_tools") + or chat_template_kwargs.get("tools") + or [] + ) + if tools: + config["available_tools"] = [_tool_to_melody(t) for t in tools] + + # Reasoning toggle (cmd3 + cmd4) + if (rt := chat_template_kwargs.get("reasoning_type")) is not None: + config["reasoning_type"] = str(rt) + elif "thinking" in chat_template_kwargs: + # Cohere v2 ``thinking: {type: enabled|disabled}`` shorthand + thinking = chat_template_kwargs["thinking"] + t = thinking.get("type") if isinstance(thinking, dict) else thinking + if t in ("enabled", "disabled"): + config["reasoning_type"] = t + + if (rp := chat_template_kwargs.get("response_prefix")) is not None: + config["response_prefix"] = str(rp) + if (di := chat_template_kwargs.get("dev_instruction")) is not None: + config["dev_instruction"] = str(di) + if (extra := chat_template_kwargs.get("additional_template_fields")): + config["additional_template_fields"] = dict(extra) + if (esc := chat_template_kwargs.get("escaped_special_tokens")): + config["escaped_special_tokens"] = dict(esc) + + # JSON / structured outputs + if (rf := chat_template_kwargs.get("response_format")) is not None: + rf = rf.model_dump() if hasattr(rf, "model_dump") else dict(rf) + rf_type = rf.get("type") + if rf_type == "json_object": + config["json_mode"] = True + elif rf_type in ("json_schema", "json"): + schema = rf.get("schema") or rf.get("json_schema") + if isinstance(schema, dict) and "schema" in schema: + schema = schema["schema"] + if schema is not None: + config["json_schema"] = ( + schema if isinstance(schema, str) else json.dumps(schema) + ) + if (js := chat_template_kwargs.get("json_schema")) is not None: + config["json_schema"] = js if isinstance(js, str) else json.dumps(js) + if "json_mode" in chat_template_kwargs: + config["json_mode"] = bool(chat_template_kwargs["json_mode"]) + + if fmt == "cmd3": + if (sm := chat_template_kwargs.get("safety_mode")) is not None: + config["safety_mode"] = str(sm).lower() + # citation_quality: ``on`` / ``off`` + cq = chat_template_kwargs.get("citation_quality") + if cq is None and (co := chat_template_kwargs.get("citation_options")): + mode = co.get("mode") if isinstance(co, dict) else None + if mode is not None: + cq = "on" if str(mode).lower() != "off" else "off" + if cq is not None: + config["citation_quality"] = str(cq).lower() + if "skip_preamble" in chat_template_kwargs: + config["skip_preamble"] = bool(chat_template_kwargs["skip_preamble"]) + else: # cmd4 + # cmd4 uses ``grounding`` rather than safety_mode/citation_quality + if (gr := chat_template_kwargs.get("grounding")) is not None: + config["grounding"] = str(gr).lower() + elif (co := chat_template_kwargs.get("citation_options")): + mode = co.get("mode") if isinstance(co, dict) else None + if mode is not None: + config["grounding"] = str(mode).lower() + if (pi := chat_template_kwargs.get("platform_instruction")) is not None: + config["platform_instruction"] = str(pi) + + return fmt, config + + +class CohereRenderer(BaseRenderer[HfTokenizer]): + """Renderer that templates Cohere prompts via the melody Rust bindings. + + Tokenization is delegated to the standard HF tokenizer; only the + chat-template step is replaced with ``cohere_melody.render_cmd3`` / + ``render_cmd4``. Picking this renderer is opt-in via + ``--tokenizer-mode cohere``. + """ + + def __init__( + self, + config: VllmConfig, + tokenizer: HfTokenizer | None, + ) -> None: + # Match HfRenderer in not mutating the cached tokenizer instance + tokenizer = copy.copy(tokenizer) + super().__init__(config, tokenizer) + + # Lazy import to keep `cohere_melody` an optional dependency + self._melody = _try_import_melody() + self._render_messages_async = make_async( + self.render_messages, executor=self._executor + ) + + def _render(self, fmt: str, config_dict: dict[str, Any]) -> str: + if fmt == "cmd3": + return self._melody.render_cmd3(config_dict) + return self._melody.render_cmd4(config_dict) + + def render_messages( + self, + messages: list[ChatCompletionMessageParam], + params: ChatParams, + ) -> tuple[list[ConversationMessage], DictPrompt]: + conversation, mm_data, mm_uuids = parse_chat_messages( + messages, + self.model_config, + content_format="openai", + media_io_kwargs=params.media_io_kwargs, + mm_processor_kwargs=params.mm_processor_kwargs, + ) + + chat_template_kwargs = dict(params.chat_template_kwargs) + fmt, config_dict = _build_render_config(conversation, chat_template_kwargs) + prompt_text = self._render(fmt, config_dict) + prompt = parse_dec_only_prompt(prompt_text) + + if mm_data is not None: + prompt["multi_modal_data"] = mm_data + if mm_uuids is not None: + prompt["multi_modal_uuids"] = mm_uuids + + return conversation, prompt + + async def render_messages_async( + self, + messages: list[ChatCompletionMessageParam], + params: ChatParams, + ) -> tuple[list[ConversationMessage], DictPrompt]: + conversation, mm_data, mm_uuids = await parse_chat_messages_async( + messages, + self.model_config, + content_format="openai", + media_io_kwargs=params.media_io_kwargs, + mm_processor_kwargs=params.mm_processor_kwargs, + ) + + chat_template_kwargs = dict(params.chat_template_kwargs) + fmt, config_dict = _build_render_config(conversation, chat_template_kwargs) + # render_cmd3/4 are pure CPU work; offload to thread pool for parity + # with HfRenderer's async path. + prompt_text = await make_async( + self._render, executor=self._executor + )(fmt, config_dict) + prompt = parse_dec_only_prompt(prompt_text) + + if mm_data is not None: + prompt["multi_modal_data"] = mm_data + if mm_uuids is not None: + prompt["multi_modal_uuids"] = mm_uuids + + return conversation, prompt diff --git a/vllm/renderers/registry.py b/vllm/renderers/registry.py index 8263dd713a49..8f1c4e1748d4 100644 --- a/vllm/renderers/registry.py +++ b/vllm/renderers/registry.py @@ -20,6 +20,7 @@ _VLLM_RENDERERS = { + "cohere": ("cohere", "CohereRenderer"), "deepseek_v32": ("deepseek_v32", "DeepseekV32Renderer"), "deepseek_v4": ("deepseek_v4", "DeepseekV4Renderer"), "grok2": ("grok2", "Grok2Renderer"), diff --git a/vllm/tokenizers/registry.py b/vllm/tokenizers/registry.py index 7578d3b43ab2..2d62755cbe65 100644 --- a/vllm/tokenizers/registry.py +++ b/vllm/tokenizers/registry.py @@ -41,6 +41,9 @@ _MODEL_TYPES_WITH_INCORRECT_TOKENIZER_CLASS: set[str] = {"step3_vl", "step3p7"} _VLLM_TOKENIZERS = { + # ``cohere`` mode uses the standard cached HF tokenizer; only the + # renderer (template stage) is replaced with a melody-based one. + "cohere": ("hf", "CachedHfTokenizer"), "deepseek_v32": ("deepseek_v32", "DeepseekV32Tokenizer"), "deepseek_v4": ("deepseek_v4", "DeepseekV4Tokenizer"), "grok2": ("grok2", "Grok2Tokenizer"), From 1c20e8921782a3eda072c566f80abd6991933c7a Mon Sep 17 00:00:00 2001 From: Andrew Berneshawi Date: Thu, 4 Jun 2026 10:25:50 -0400 Subject: [PATCH 02/30] fix yield --- vllm/config/model.py | 7 ++++++- vllm/entrypoints/cohere/serving.py | 3 ++- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/vllm/config/model.py b/vllm/config/model.py index b41ab189c108..19873597b972 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -93,7 +93,9 @@ def is_cumem_allocator_available() -> bool: RunnerOption = Literal["auto", RunnerType] ConvertType = Literal["none", "embed", "classify"] ConvertOption = Literal["auto", ConvertType] -TokenizerMode = Literal["auto", "hf", "slow", "mistral", "deepseek_v32", "deepseek_v4"] +TokenizerMode = Literal[ + "auto", "hf", "slow", "mistral", "deepseek_v32", "deepseek_v4", "cohere" +] ModelDType = Literal["auto", "half", "float16", "bfloat16", "float", "float32"] LogprobsMode = Literal[ "raw_logits", "raw_logprobs", "processed_logits", "processed_logprobs" @@ -146,6 +148,9 @@ class ModelConfig: - "deepseek_v32" will always use the tokenizer from `deepseek_v32`. - "deepseek_v4" will always use the tokenizer from `deepseek_v4`. - "qwen_vl" will always use the tokenizer from `qwen_vl`. + - "cohere" uses the standard HF tokenizer but renders the chat template + via the `cohere_melody` library (cmd3 / cmd4 templates) instead of + Jinja, and surfaces grounded-citation metadata on responses. - Other custom values can be supported via plugins. To swap the Rust BPE backend that powers HF fast tokenizers for the diff --git a/vllm/entrypoints/cohere/serving.py b/vllm/entrypoints/cohere/serving.py index e6e925f18517..14e15b5fdb0e 100644 --- a/vllm/entrypoints/cohere/serving.py +++ b/vllm/entrypoints/cohere/serving.py @@ -638,7 +638,8 @@ async def _chat_completion_stream_to_v2( # The final OpenAI chunk has no choices and only carries usage. if not chunk.choices: - yield from self._close_open_blocks(state) + for ev in self._close_open_blocks(state): + yield ev yield _sse( self._build_message_end_event( chunk_id=chunk.id, From f1f807ac95da04d48d86ea62d80662cec05f8bcc Mon Sep 17 00:00:00 2001 From: Andrew Berneshawi Date: Wed, 10 Jun 2026 11:08:33 -0400 Subject: [PATCH 03/30] fix image detail missing --- vllm/entrypoints/cohere/serving.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/vllm/entrypoints/cohere/serving.py b/vllm/entrypoints/cohere/serving.py index 14e15b5fdb0e..4b2c06a86383 100644 --- a/vllm/entrypoints/cohere/serving.py +++ b/vllm/entrypoints/cohere/serving.py @@ -258,10 +258,13 @@ def _convert_user_message( if isinstance(block, CohereTextContent): content_parts.append({"type": "text", "text": block.text}) else: # CohereImageContent + image_url: dict[str, Any] = {"url": block.image_url.url} + if block.image_url.detail is not None: + image_url["detail"] = block.image_url.detail content_parts.append( { "type": "image_url", - "image_url": {"url": block.image_url.url}, + "image_url": image_url, } ) if ( From b9c9bbe193ee8b7375390ad88bb7dc3c8eb41742 Mon Sep 17 00:00:00 2001 From: Andrew Berneshawi Date: Wed, 10 Jun 2026 13:47:58 -0400 Subject: [PATCH 04/30] fix strict tools conversion handling --- vllm/entrypoints/cohere/serving.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/vllm/entrypoints/cohere/serving.py b/vllm/entrypoints/cohere/serving.py index 4b2c06a86383..d8843eed3dfb 100644 --- a/vllm/entrypoints/cohere/serving.py +++ b/vllm/entrypoints/cohere/serving.py @@ -412,6 +412,13 @@ def _apply_tools( ) -> None: if not request.tools: return + # Cohere's ``strict_tools`` is the spec-equivalent of OpenAI's + # per-function ``strict`` flag: when true the API guarantees tool + # call arguments match the declared JSON schema. We surface it in + # both places so OpenAI-shaped consumers see it on the function + # definition and the cohere renderer can read it back from + # ``chat_template_kwargs`` for cmd3/cmd4 preamble selection. + strict = bool(request.strict_tools) if request.strict_tools else False chat_req.tools = [ ChatCompletionToolsParam.model_validate( { @@ -420,6 +427,7 @@ def _apply_tools( "name": tool.function.name, "description": tool.function.description, "parameters": tool.function.parameters, + **({"strict": True} if strict else {}), }, } ) @@ -485,6 +493,8 @@ def _apply_cohere_template_kwargs( "thinking", request.thinking.model_dump(exclude_none=True), ) + if request.strict_tools is not None: + kwargs.setdefault("strict_tools", bool(request.strict_tools)) if kwargs: chat_req.chat_template_kwargs = kwargs From 228eddebd145614f1be23c1892e780def95d1684 Mon Sep 17 00:00:00 2001 From: Andrew Berneshawi Date: Wed, 17 Jun 2026 08:50:28 -0400 Subject: [PATCH 05/30] fix document handling --- vllm/entrypoints/cohere/serving.py | 57 ++++++++++++++++++++++-------- 1 file changed, 42 insertions(+), 15 deletions(-) diff --git a/vllm/entrypoints/cohere/serving.py b/vllm/entrypoints/cohere/serving.py index d8843eed3dfb..3c39c85e2b2e 100644 --- a/vllm/entrypoints/cohere/serving.py +++ b/vllm/entrypoints/cohere/serving.py @@ -47,6 +47,7 @@ CohereCitation, CohereCitationEndEvent, CohereCitationStartEvent, + CohereDocumentContent, CohereFinishReason, CohereSystemMessageV2, CohereTextContent, @@ -327,24 +328,50 @@ def _convert_tool_message( cls, msg: CohereToolMessageV2 ) -> dict[str, Any]: if isinstance(msg.content, str): - content: Any = msg.content - else: - text_parts: list[str] = [] - for block in msg.content: - if isinstance(block, CohereTextContent): - text_parts.append(block.text) - else: # CohereDocumentContent - # The OpenAI chat-completion shape only knows about - # text tool results. Serialize the document so its data - # is preserved in the prompt. - text_parts.append( - json.dumps(block.document.model_dump(exclude_none=True)) - ) - content = "\n".join(text_parts) + return { + "role": "tool", + "tool_call_id": msg.tool_call_id, + "content": msg.content, + } + + # When the tool result is text-only, flatten to a string for maximum + # compatibility with vanilla chat templates. When it includes + # documents, preserve them as structured content parts so the cohere + # renderer can surface them as grounding sources (the + # :class:`CohereRenderer` understands ``{type: document, document: + # {...}}`` blocks). Non-cohere renderers may not honor document + # blocks, but that matches the broader "documents are no-op for OSS + # models" contract documented on this endpoint. + has_documents = any( + isinstance(block, CohereDocumentContent) for block in msg.content + ) + if not has_documents: + text = "\n".join( + block.text + for block in msg.content + if isinstance(block, CohereTextContent) + ) + return { + "role": "tool", + "tool_call_id": msg.tool_call_id, + "content": text, + } + + parts: list[dict[str, Any]] = [] + for block in msg.content: + if isinstance(block, CohereTextContent): + parts.append({"type": "text", "text": block.text}) + else: # CohereDocumentContent + parts.append( + { + "type": "document", + "document": block.document.model_dump(exclude_none=True), + } + ) return { "role": "tool", "tool_call_id": msg.tool_call_id, - "content": content, + "content": parts, } @classmethod From c7d09ccaff1c6cfe0d6bc38ab7c4dc58fe96f80e Mon Sep 17 00:00:00 2001 From: Andrew Berneshawi Date: Wed, 17 Jun 2026 09:43:11 -0400 Subject: [PATCH 06/30] cache renderer thread --- vllm/renderers/cohere.py | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/vllm/renderers/cohere.py b/vllm/renderers/cohere.py index 7893def825a6..c3df3cf7a6c8 100644 --- a/vllm/renderers/cohere.py +++ b/vllm/renderers/cohere.py @@ -388,9 +388,10 @@ def __init__( # Lazy import to keep `cohere_melody` an optional dependency self._melody = _try_import_melody() - self._render_messages_async = make_async( - self.render_messages, executor=self._executor - ) + # ``render_cmd3`` / ``render_cmd4`` are pure CPU work; cache the + # thread-pool wrapper once so the async path doesn't allocate a new + # adapter on every request. + self._render_async = make_async(self._render, executor=self._executor) def _render(self, fmt: str, config_dict: dict[str, Any]) -> str: if fmt == "cmd3": @@ -437,11 +438,7 @@ async def render_messages_async( chat_template_kwargs = dict(params.chat_template_kwargs) fmt, config_dict = _build_render_config(conversation, chat_template_kwargs) - # render_cmd3/4 are pure CPU work; offload to thread pool for parity - # with HfRenderer's async path. - prompt_text = await make_async( - self._render, executor=self._executor - )(fmt, config_dict) + prompt_text = await self._render_async(fmt, config_dict) prompt = parse_dec_only_prompt(prompt_text) if mm_data is not None: From 307eeff52db68394e7b97bfd23c5476475d01c1e Mon Sep 17 00:00:00 2001 From: Andrew Berneshawi Date: Mon, 22 Jun 2026 23:08:29 -0400 Subject: [PATCH 07/30] Initial refactor to reuse sdk types instead of redefining them --- vllm/entrypoints/cohere/api_router.py | 158 ++++---- vllm/entrypoints/cohere/protocol.py | 461 +++++++----------------- vllm/entrypoints/cohere/serving.py | 237 ++++++------ vllm/entrypoints/generate/api_router.py | 13 +- 4 files changed, 375 insertions(+), 494 deletions(-) diff --git a/vllm/entrypoints/cohere/api_router.py b/vllm/entrypoints/cohere/api_router.py index 86ababaac0bc..3e14d07c8344 100644 --- a/vllm/entrypoints/cohere/api_router.py +++ b/vllm/entrypoints/cohere/api_router.py @@ -1,82 +1,114 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""FastAPI router for the Cohere Chat v2 API (``POST /v2/chat``). -from http import HTTPStatus +The Cohere v2 protocol models are sourced from the official ``cohere`` +Python SDK (``pip install cohere``). To keep that an *optional* +dependency for vLLM, we don't import the SDK or any module that pulls +it in at top-level - everything that touches the SDK is imported lazily +inside :func:`attach_router`. If the SDK isn't installed, attach_router +logs a one-line notice and returns without registering the route, so +vLLM continues to boot normally. +""" +from __future__ import annotations -from fastapi import APIRouter, Depends, FastAPI, Request -from fastapi.responses import JSONResponse, StreamingResponse +from fastapi import FastAPI -from vllm.entrypoints.cohere.protocol import ( - CohereChatV2Request, - CohereChatV2Response, - CohereError, -) -from vllm.entrypoints.cohere.serving import CohereServingChatV2 -from vllm.entrypoints.openai.engine.protocol import ErrorResponse -from vllm.entrypoints.openai.utils import validate_json_request -from vllm.entrypoints.utils import load_aware_call, with_cancellation from vllm.logger import init_logger logger = init_logger(__name__) -router = APIRouter() +def attach_router(app: FastAPI) -> None: + """Register ``POST /v2/chat`` on ``app``. -def _serving(request: Request) -> CohereServingChatV2 | None: - return getattr(request.app.state, "cohere_serving_chat_v2", None) - - -def _error_response( - error: ErrorResponse, *, fallback_status: int = HTTPStatus.BAD_REQUEST -) -> JSONResponse: - """Translate vLLM's internal error envelope into Cohere's error shape.""" - info = error.error - status = info.code or fallback_status - return JSONResponse( - status_code=status, - content=CohereError(message=info.message).model_dump(exclude_none=True), - ) - - -@router.post( - "/v2/chat", - dependencies=[Depends(validate_json_request)], - responses={ - HTTPStatus.OK.value: {"content": {"text/event-stream": {}}}, - HTTPStatus.BAD_REQUEST.value: {"model": CohereError}, - HTTPStatus.NOT_FOUND.value: {"model": CohereError}, - HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": CohereError}, - }, -) -@with_cancellation -@load_aware_call -async def chat_v2(request: CohereChatV2Request, raw_request: Request): - handler = _serving(raw_request) - if handler is None: - return JSONResponse( - status_code=HTTPStatus.NOT_IMPLEMENTED.value, - content=CohereError( - message="The model does not support the Cohere v2 chat API." - ).model_dump(exclude_none=True), - ) - + No-op (with an info log) when the optional ``cohere`` SDK isn't + installed, since the v2 protocol models live there. + """ try: - result = await handler.create_chat_v2(request, raw_request) - except Exception as e: # noqa: BLE001 - report as 500 for parity w/ siblings - logger.exception("Error in /v2/chat: %s", e) - return JSONResponse( - status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, - content=CohereError(message=str(e)).model_dump(exclude_none=True), + import cohere # noqa: F401 -- dependency probe + except ImportError: + logger.info( + "cohere SDK not installed; /v2/chat endpoint disabled. " + "Install with `pip install cohere` to enable it." ) + return - if isinstance(result, ErrorResponse): - return _error_response(result) + # Defer the SDK-dependent imports until we know cohere is available. + # This keeps `import vllm.entrypoints.cohere.api_router` safe when the + # SDK isn't installed. + from http import HTTPStatus - if isinstance(result, CohereChatV2Response): - return JSONResponse(content=result.model_dump(exclude_none=True)) + from fastapi import APIRouter, Depends, Request + from fastapi.responses import JSONResponse, StreamingResponse - return StreamingResponse(content=result, media_type="text/event-stream") + from vllm.entrypoints.cohere.protocol import ( + CohereChatV2Request, + CohereChatV2Response, + CohereError, + ) + from vllm.entrypoints.cohere.serving import CohereServingChatV2 + from vllm.entrypoints.openai.engine.protocol import ErrorResponse + from vllm.entrypoints.openai.utils import validate_json_request + from vllm.entrypoints.utils import load_aware_call, with_cancellation + + router = APIRouter() + + def _serving(request: Request) -> CohereServingChatV2 | None: + return getattr(request.app.state, "cohere_serving_chat_v2", None) + + def _error_response( + error: ErrorResponse, + *, + 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=info.message).model_dump( + exclude_none=True + ), + ) + @router.post( + "/v2/chat", + dependencies=[Depends(validate_json_request)], + responses={ + HTTPStatus.OK.value: {"content": {"text/event-stream": {}}}, + HTTPStatus.BAD_REQUEST.value: {"model": CohereError}, + HTTPStatus.NOT_FOUND.value: {"model": CohereError}, + HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": CohereError}, + }, + ) + @with_cancellation + @load_aware_call + async def chat_v2(request: CohereChatV2Request, raw_request: Request): + handler = _serving(raw_request) + if handler is None: + return JSONResponse( + status_code=HTTPStatus.NOT_IMPLEMENTED.value, + content=CohereError( + message="The model does not support the Cohere v2 chat API." + ).model_dump(exclude_none=True), + ) + + try: + result = await handler.create_chat_v2(request, raw_request) + except Exception as e: # noqa: BLE001 - report as 500 for parity + logger.exception("Error in /v2/chat: %s", e) + return JSONResponse( + status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, + content=CohereError(message=str(e)).model_dump(exclude_none=True), + ) + + if isinstance(result, ErrorResponse): + return _error_response(result) + + if isinstance(result, CohereChatV2Response): + return JSONResponse(content=result.model_dump(exclude_none=True)) + + return StreamingResponse(content=result, media_type="text/event-stream") -def attach_router(app: FastAPI) -> None: app.include_router(router) diff --git a/vllm/entrypoints/cohere/protocol.py b/vllm/entrypoints/cohere/protocol.py index 5adc3440c788..b4eb679360b0 100644 --- a/vllm/entrypoints/cohere/protocol.py +++ b/vllm/entrypoints/cohere/protocol.py @@ -1,17 +1,98 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Pydantic models for the Cohere Chat v2 API protocol. - -See https://docs.cohere.com/reference/chat for the upstream specification. - -This module mirrors the Cohere V2 wire format using Pydantic discriminated -unions for messages, content blocks, and stream events. +"""Cohere Chat v2 API protocol. + +The bulk of the wire types come straight from the official ``cohere`` +Python SDK so we stay in lockstep with the upstream specification and +avoid hand-mirroring the schema. We only own three things locally: + +1. The top-level request body model (the SDK doesn't ship one — its + ``ClientV2.chat`` takes the body as kwargs), with vLLM-specific + extensions (``kv_transfer_params`` / ``chat_template_kwargs``). +2. The non-streaming response envelope (the SDK exposes the message + shape via :class:`AssistantMessageResponse` but no full response + wrapper). +3. The streaming discriminated union (the SDK exports each event type + individually but not as a combined ``Annotated[Union[...], + discriminator]``). + +Importing this module pulls in the ``cohere`` package. The router that +mounts ``POST /v2/chat`` guards on that import succeeding so vLLM still +boots without the SDK installed. + +See https://docs.cohere.com/reference/chat for the upstream spec. """ +from __future__ import annotations import time from typing import Any, Literal -from pydantic import BaseModel, Field, RootModel, field_validator +from cohere.types import ( + AssistantChatMessageV2, + AssistantMessageResponse, + ChatContentDeltaEvent, + ChatContentEndEvent, + ChatContentStartEvent, + ChatMessageEndEvent, + ChatMessageStartEvent, + ChatMessageV2, + ChatRequestSafetyMode, + ChatToolCallDeltaEvent, + ChatToolCallEndEvent, + ChatToolCallStartEvent, + ChatToolPlanDeltaEvent, + Citation, + CitationEndEvent, + CitationOptions, + CitationStartEvent, + Document, + ResponseFormatV2, + SystemChatMessageV2, + Thinking, + ToolCallV2, + ToolChatMessageV2, + ToolV2, + UserChatMessageV2, +) +from pydantic import BaseModel, Field, field_validator + +# Re-export the SDK types so that ``vllm.entrypoints.cohere.serving`` and +# friends can import everything they need from this module. +__all__ = [ + "AssistantChatMessageV2", + "AssistantMessageResponse", + "ChatContentDeltaEvent", + "ChatContentEndEvent", + "ChatContentStartEvent", + "ChatMessageEndEvent", + "ChatMessageStartEvent", + "ChatMessageV2", + "ChatToolCallDeltaEvent", + "ChatToolCallEndEvent", + "ChatToolCallStartEvent", + "ChatToolPlanDeltaEvent", + "Citation", + "CitationEndEvent", + "CitationOptions", + "CitationStartEvent", + "CohereChatV2Request", + "CohereChatV2Response", + "CohereError", + "CohereFinishReason", + "CohereLogprobItem", + "CohereUsage", + "CohereUsageBilledUnits", + "CohereUsageTokens", + "Document", + "ResponseFormatV2", + "SystemChatMessageV2", + "Thinking", + "ToolCallV2", + "ToolChatMessageV2", + "ToolV2", + "UserChatMessageV2", +] + # --------------------------------------------------------------------------- # Errors @@ -19,7 +100,7 @@ class CohereError(BaseModel): - """Top-level error body returned by /v2/chat error responses. + """Top-level error body returned by ``/v2/chat`` error responses. Cohere's documented error schemas are uniform: ``{message, id}``. """ @@ -29,194 +110,23 @@ class CohereError(BaseModel): # --------------------------------------------------------------------------- -# Content blocks -# --------------------------------------------------------------------------- - - -class CohereTextContent(BaseModel): - type: Literal["text"] = "text" - text: str - - -class CohereThinkingContent(BaseModel): - type: Literal["thinking"] = "thinking" - thinking: str - - -class CohereImageUrl(BaseModel): - url: str - detail: Literal["auto", "low", "high"] | None = None - - -class CohereImageContent(BaseModel): - type: Literal["image_url"] = "image_url" - image_url: CohereImageUrl - - -CohereUserContentBlock = CohereTextContent | CohereImageContent -CohereAssistantContentBlock = CohereTextContent | CohereThinkingContent - - -# --------------------------------------------------------------------------- -# Documents (request side) and Document content (tool message side) -# --------------------------------------------------------------------------- - - -class CohereDocument(BaseModel): - """A document supplied for grounded generation.""" - - data: dict[str, Any] | str - id: str | None = None - - -class CohereDocumentContent(BaseModel): - type: Literal["document"] = "document" - document: CohereDocument - - -CohereToolContent = CohereTextContent | CohereDocumentContent - - -# --------------------------------------------------------------------------- -# Tool calls -# --------------------------------------------------------------------------- - - -class CohereToolCallFunction(BaseModel): - name: str | None = None - # Cohere serializes arguments as a JSON-encoded string (same as OpenAI). - arguments: str | None = None - - -class CohereToolCallV2(BaseModel): - id: str - type: Literal["function"] = "function" - function: CohereToolCallFunction | None = None - - -# --------------------------------------------------------------------------- -# Messages +# Tool choice / finish reasons # --------------------------------------------------------------------------- - - -class CohereUserMessageV2(BaseModel): - role: Literal["user"] = "user" - content: str | list[CohereUserContentBlock] - - -class CohereSystemMessageV2(BaseModel): - role: Literal["system"] = "system" - content: str | list[CohereTextContent] - - -class CohereAssistantMessageV2(BaseModel): - role: Literal["assistant"] = "assistant" - content: str | list[CohereAssistantContentBlock] | None = None - tool_calls: list[CohereToolCallV2] | None = None - tool_plan: str | None = Field( - default=None, - description=( - "Chain-of-thought plan emitted by the model alongside tool calls." - ), - ) - citations: list["CohereCitation"] | None = None - - -class CohereToolMessageV2(BaseModel): - role: Literal["tool"] = "tool" - tool_call_id: str - content: str | list[CohereToolContent] - - -CohereChatMessageV2 = ( - CohereUserMessageV2 - | CohereAssistantMessageV2 - | CohereSystemMessageV2 - | CohereToolMessageV2 -) - - -# --------------------------------------------------------------------------- -# Tools (request side) and tool choice -# --------------------------------------------------------------------------- - - -class CohereToolFunction(BaseModel): - name: str - description: str | None = None - parameters: dict[str, Any] = Field(default_factory=dict) - - -class CohereToolV2(BaseModel): - type: Literal["function"] = "function" - function: CohereToolFunction - +# +# These literals aren't first-class enums in the SDK but are documented at +# https://docs.cohere.com/reference/chat. We declare them here so the +# request/response models can validate them. CohereToolChoice = Literal["REQUIRED", "NONE"] - -# --------------------------------------------------------------------------- -# Citations / sources -# --------------------------------------------------------------------------- - - -class CohereChatToolSource(BaseModel): - type: Literal["tool"] = "tool" - id: str | None = None - tool_output: dict[str, Any] | None = None - - -class CohereChatDocumentSource(BaseModel): - type: Literal["document"] = "document" - id: str | None = None - document: dict[str, Any] | None = None - - -CohereSource = CohereChatToolSource | CohereChatDocumentSource - - -class CohereCitation(BaseModel): - start: int | None = None - end: int | None = None - text: str | None = None - sources: list[CohereSource] | None = None - content_index: int | None = None - type: Literal["TEXT_CONTENT", "PLAN"] | None = None - - -# --------------------------------------------------------------------------- -# Response format / output config -# --------------------------------------------------------------------------- - - -class CohereTextResponseFormat(BaseModel): - type: Literal["text"] = "text" - - -class CohereJsonResponseFormat(BaseModel): - type: Literal["json_object"] = "json_object" - json_schema: dict[str, Any] | None = None - - -CohereResponseFormatV2 = CohereTextResponseFormat | CohereJsonResponseFormat - - -# --------------------------------------------------------------------------- -# Misc options -# --------------------------------------------------------------------------- - - -CohereSafetyMode = Literal["CONTEXTUAL", "STRICT", "OFF"] -CohereCitationMode = Literal["ENABLED", "DISABLED", "FAST", "ACCURATE", "OFF"] - - -class CohereCitationOptions(BaseModel): - mode: CohereCitationMode | None = None - - -class CohereThinking(BaseModel): - type: Literal["enabled", "disabled"] - token_budget: int | None = None +CohereFinishReason = Literal[ + "COMPLETE", + "STOP_SEQUENCE", + "MAX_TOKENS", + "TOOL_CALL", + "ERROR", + "TIMEOUT", +] # --------------------------------------------------------------------------- @@ -227,26 +137,27 @@ class CohereThinking(BaseModel): class CohereChatV2Request(BaseModel): """Cohere Chat v2 request body. - Mirrors the schema documented at - https://docs.cohere.com/reference/chat. + Mirrors the schema documented at https://docs.cohere.com/reference/chat. + All structured fields delegate to the official SDK types so the body + schema stays in sync with the upstream spec. """ model: str - messages: list[CohereChatMessageV2] + messages: list[ChatMessageV2] stream: bool | None = False # Tooling - tools: list[CohereToolV2] | None = None + tools: list[ToolV2] | None = None strict_tools: bool | None = None tool_choice: CohereToolChoice | None = None # Grounding - documents: list[str | CohereDocument] | None = None - citation_options: CohereCitationOptions | None = None + documents: list[str | Document] | None = None + citation_options: CitationOptions | None = None # Output - response_format: CohereResponseFormatV2 | None = None - safety_mode: CohereSafetyMode | None = None + response_format: ResponseFormatV2 | None = None + safety_mode: ChatRequestSafetyMode | None = None max_tokens: int | None = None stop_sequences: list[str] | None = None @@ -260,14 +171,14 @@ class CohereChatV2Request(BaseModel): logprobs: bool | None = None # Reasoning - thinking: CohereThinking | None = None + thinking: Thinking | None = None # Scheduling priority: int | None = None # vLLM-specific extensions (not in Cohere spec). These mirror what the - # Anthropic and OpenAI surfaces already expose so that V2 callers can - # reach the same engine knobs when needed. + # Anthropic and OpenAI surfaces already expose so V2 callers can reach + # the same engine knobs when needed. kv_transfer_params: dict[str, Any] | None = Field( default=None, description="KVTransfer parameters used for disaggregated serving.", @@ -296,8 +207,12 @@ def _validate_max_tokens(cls, v: int | None) -> int | None: # --------------------------------------------------------------------------- -# Usage +# Usage / Logprobs # --------------------------------------------------------------------------- +# +# The Cohere SDK only exposes a v1 ``ApiMetaBilledUnits``. The v2 usage +# envelope is documented separately in the OpenAPI spec and we declare it +# here. class CohereUsageBilledUnits(BaseModel): @@ -318,11 +233,6 @@ class CohereUsage(BaseModel): cached_tokens: float | None = None -# --------------------------------------------------------------------------- -# Logprobs -# --------------------------------------------------------------------------- - - class CohereLogprobItem(BaseModel): text: str | None = None token_ids: list[int] @@ -334,28 +244,17 @@ class CohereLogprobItem(BaseModel): # --------------------------------------------------------------------------- -CohereFinishReason = Literal[ - "COMPLETE", - "STOP_SEQUENCE", - "MAX_TOKENS", - "TOOL_CALL", - "ERROR", - "TIMEOUT", -] - - -class CohereAssistantMessageResponse(BaseModel): - role: Literal["assistant"] = "assistant" - content: list[CohereAssistantContentBlock] | None = None - tool_calls: list[CohereToolCallV2] | None = None - tool_plan: str | None = None - citations: list[CohereCitation] | None = None +class CohereChatV2Response(BaseModel): + """Cohere Chat v2 non-streaming response body. + Wraps the SDK :class:`AssistantMessageResponse` (the message shape) in + the documented v2 response envelope (``id``, ``finish_reason``, + ``usage``, ``logprobs``). + """ -class CohereChatV2Response(BaseModel): id: str finish_reason: CohereFinishReason - message: CohereAssistantMessageResponse + message: AssistantMessageResponse usage: CohereUsage | None = None logprobs: list[CohereLogprobItem] | None = None @@ -373,100 +272,14 @@ def model_post_init(self, __context: Any) -> None: # Streaming events # --------------------------------------------------------------------------- # -# Cohere V2 streams a sequence of typed JSON events delivered as -# Server-Sent Events. Each event is a discriminated union over the ``type`` -# field. See: -# https://docs.cohere.com/v2/docs/streaming -# and the OpenAPI ``StreamedChatResponseV2`` schema. - - -class _CohereStreamEventBase(BaseModel): - """Common base – every event has a ``type`` discriminator and the - nested ``delta`` shapes are loose dicts so we don't have to redefine the - full Cohere schema for each event variant.""" - - -class CohereChatMessageStartEvent(_CohereStreamEventBase): - type: Literal["message-start"] = "message-start" - id: str | None = None - delta: dict[str, Any] | None = None - - -class CohereChatContentStartEvent(_CohereStreamEventBase): - type: Literal["content-start"] = "content-start" - index: int | None = None - delta: dict[str, Any] | None = None - - -class CohereChatContentDeltaEvent(_CohereStreamEventBase): - type: Literal["content-delta"] = "content-delta" - index: int | None = None - delta: dict[str, Any] | None = None - logprobs: CohereLogprobItem | None = None - - -class CohereChatContentEndEvent(_CohereStreamEventBase): - type: Literal["content-end"] = "content-end" - index: int | None = None - - -class CohereChatToolPlanDeltaEvent(_CohereStreamEventBase): - type: Literal["tool-plan-delta"] = "tool-plan-delta" - delta: dict[str, Any] | None = None - - -class CohereChatToolCallStartEvent(_CohereStreamEventBase): - type: Literal["tool-call-start"] = "tool-call-start" - index: int | None = None - delta: dict[str, Any] | None = None - - -class CohereChatToolCallDeltaEvent(_CohereStreamEventBase): - type: Literal["tool-call-delta"] = "tool-call-delta" - index: int | None = None - delta: dict[str, Any] | None = None - - -class CohereChatToolCallEndEvent(_CohereStreamEventBase): - type: Literal["tool-call-end"] = "tool-call-end" - index: int | None = None - - -class CohereCitationStartEvent(_CohereStreamEventBase): - type: Literal["citation-start"] = "citation-start" - index: int | None = None - delta: dict[str, Any] | None = None - - -class CohereCitationEndEvent(_CohereStreamEventBase): - type: Literal["citation-end"] = "citation-end" - index: int | None = None - - -class CohereChatMessageEndEvent(_CohereStreamEventBase): - type: Literal["message-end"] = "message-end" - id: str | None = None - delta: dict[str, Any] | None = None - - -CohereStreamedChatResponseV2 = ( - CohereChatMessageStartEvent - | CohereChatContentStartEvent - | CohereChatContentDeltaEvent - | CohereChatContentEndEvent - | CohereChatToolPlanDeltaEvent - | CohereChatToolCallStartEvent - | CohereChatToolCallDeltaEvent - | CohereChatToolCallEndEvent - | CohereCitationStartEvent - | CohereCitationEndEvent - | CohereChatMessageEndEvent -) - - -class CohereStreamedChatResponseV2Envelope(RootModel[CohereStreamedChatResponseV2]): - """Discriminated-union wrapper for stream events.""" - - -# Resolve forward references created by ``CohereAssistantMessageV2.citations``. -CohereAssistantMessageV2.model_rebuild() +# Cohere V2 streams a sequence of typed JSON events delivered as Server- +# Sent Events. Each event class is exported above (``ChatMessageStartEvent``, +# ``ChatContentDeltaEvent``, ``CitationStartEvent``, etc.). We don't build +# a discriminated union here because the SDK's event classes don't carry +# the ``type`` literal as a Pydantic field (it's handled by the SDK's own +# deserializer), so pydantic ``Field(discriminator="type")`` would reject +# it. Surfaces emit each event with ``model_dump_json()`` and clients +# parse them by reading ``type`` directly. +# +# See https://docs.cohere.com/v2/docs/streaming and the OpenAPI +# ``StreamedChatResponseV2`` schema for the wire-format reference. diff --git a/vllm/entrypoints/cohere/serving.py b/vllm/entrypoints/cohere/serving.py index 3c39c85e2b2e..eeadab9dd657 100644 --- a/vllm/entrypoints/cohere/serving.py +++ b/vllm/entrypoints/cohere/serving.py @@ -27,38 +27,37 @@ from collections.abc import AsyncGenerator from typing import TYPE_CHECKING, Any +from cohere.types import ( + AssistantChatMessageV2, + AssistantMessageResponse, + ChatContentDeltaEvent, + ChatContentEndEvent, + ChatContentStartEvent, + ChatMessageEndEvent, + ChatMessageStartEvent, + ChatToolCallDeltaEvent, + ChatToolCallEndEvent, + ChatToolCallStartEvent, + Citation, + CitationEndEvent, + CitationStartEvent, + SystemChatMessageV2, + ToolCallV2, + ToolCallV2Function, + ToolChatMessageV2, + UserChatMessageV2, +) from fastapi import Request from vllm.engine.protocol import EngineClient from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption from vllm.entrypoints.cohere.protocol import ( - CohereAssistantMessageResponse, - CohereAssistantMessageV2, - CohereChatContentDeltaEvent, - CohereChatContentEndEvent, - CohereChatContentStartEvent, - CohereChatMessageEndEvent, - CohereChatMessageStartEvent, - CohereChatToolCallDeltaEvent, - CohereChatToolCallEndEvent, - CohereChatToolCallStartEvent, CohereChatV2Request, CohereChatV2Response, - CohereCitation, - CohereCitationEndEvent, - CohereCitationStartEvent, - CohereDocumentContent, CohereFinishReason, - CohereSystemMessageV2, - CohereTextContent, - CohereThinkingContent, - CohereToolCallFunction, - CohereToolCallV2, - CohereToolMessageV2, CohereUsage, CohereUsageBilledUnits, CohereUsageTokens, - CohereUserMessageV2, ) from vllm.entrypoints.logger import RequestLogger from vllm.entrypoints.openai.chat_completion.protocol import ( @@ -96,6 +95,21 @@ def _sse(data: str) -> str: return f"data: {data}\n\n" +def _emit(event: Any, type_str: str) -> str: + """Serialize an SDK stream event with the wire-format ``type`` field. + + The cohere SDK's stream-event Pydantic models (``ChatMessageStartEvent`` + etc.) don't declare ``type`` as a field -- the SDK handles type + discrimination in its own deserializer -- so a vanilla + ``model_dump_json()`` would drop the discriminator and break clients + that demux on ``type``. We dump the event, splice ``type`` back in, + and re-serialize as JSON before wrapping in an SSE frame. + """ + payload = event.model_dump(exclude_none=True) + payload["type"] = type_str + return _sse(json.dumps(payload)) + + # Mapping of vLLM/OpenAI finish reasons to Cohere's enum. _FINISH_REASON_MAP: dict[str | None, CohereFinishReason] = { "stop": "COMPLETE", @@ -220,18 +234,18 @@ def _convert_messages( openai_messages: list[dict[str, Any]], ) -> None: for msg in messages: - if isinstance(msg, CohereSystemMessageV2): + if isinstance(msg, SystemChatMessageV2): openai_messages.append( { "role": "system", "content": cls._coerce_text_content(msg.content), } ) - elif isinstance(msg, CohereUserMessageV2): + elif isinstance(msg, UserChatMessageV2): openai_messages.append(cls._convert_user_message(msg)) - elif isinstance(msg, CohereAssistantMessageV2): + elif isinstance(msg, AssistantChatMessageV2): openai_messages.append(cls._convert_assistant_message(msg)) - elif isinstance(msg, CohereToolMessageV2): + elif isinstance(msg, ToolChatMessageV2): openai_messages.append(cls._convert_tool_message(msg)) else: # pragma: no cover - guarded by Pydantic discriminator raise ValueError(f"Unsupported Cohere v2 message: {msg!r}") @@ -249,18 +263,22 @@ def _coerce_text_content(content: str | list[Any]) -> str: @classmethod def _convert_user_message( - cls, msg: CohereUserMessageV2 + cls, msg: UserChatMessageV2 ) -> dict[str, Any]: if isinstance(msg.content, str): return {"role": "user", "content": msg.content} + # Discriminate by ``type`` rather than isinstance so we don't have + # to import every individual ``*Content`` variant from the SDK - + # the union (``UserMessageV2Content``) covers both ``TextContent`` + # and ``ImageUrlContent`` and both expose ``type`` as a Literal. content_parts: list[dict[str, Any]] = [] for block in msg.content: - if isinstance(block, CohereTextContent): + if block.type == "text": content_parts.append({"type": "text", "text": block.text}) - else: # CohereImageContent + elif block.type == "image_url": image_url: dict[str, Any] = {"url": block.image_url.url} - if block.image_url.detail is not None: + if getattr(block.image_url, "detail", None) is not None: image_url["detail"] = block.image_url.detail content_parts.append( { @@ -277,7 +295,7 @@ def _convert_user_message( @classmethod def _convert_assistant_message( - cls, msg: CohereAssistantMessageV2 + cls, msg: AssistantChatMessageV2 ) -> dict[str, Any]: out: dict[str, Any] = {"role": "assistant"} @@ -290,9 +308,9 @@ def _convert_assistant_message( text_parts.append(msg.content) elif msg.content is not None: for block in msg.content: - if isinstance(block, CohereTextContent): + if block.type == "text": text_parts.append(block.text) - elif isinstance(block, CohereThinkingContent): + elif block.type == "thinking": thinking_parts.append(block.thinking) # ``tool_plan`` is Cohere's chain-of-thought emitted alongside tool @@ -325,7 +343,7 @@ def _convert_assistant_message( @classmethod def _convert_tool_message( - cls, msg: CohereToolMessageV2 + cls, msg: ToolChatMessageV2 ) -> dict[str, Any]: if isinstance(msg.content, str): return { @@ -342,14 +360,15 @@ def _convert_tool_message( # {...}}`` blocks). Non-cohere renderers may not honor document # blocks, but that matches the broader "documents are no-op for OSS # models" contract documented on this endpoint. - has_documents = any( - isinstance(block, CohereDocumentContent) for block in msg.content - ) + # + # Tool message content uses ``ToolMessageV2Content``, which is a + # union of ``TextToolContent`` and ``DocumentToolContent`` - + # distinct from the user-message text/image union. We discriminate + # on the ``type`` literal so we don't have to import each variant. + has_documents = any(block.type == "document" for block in msg.content) if not has_documents: text = "\n".join( - block.text - for block in msg.content - if isinstance(block, CohereTextContent) + block.text for block in msg.content if block.type == "text" ) return { "role": "tool", @@ -359,9 +378,9 @@ def _convert_tool_message( parts: list[dict[str, Any]] = [] for block in msg.content: - if isinstance(block, CohereTextContent): + if block.type == "text": parts.append({"type": "text", "text": block.text}) - else: # CohereDocumentContent + elif block.type == "document": parts.append( { "type": "document", @@ -538,18 +557,21 @@ def _chat_completion_to_v2( choice = response.choices[0] msg = choice.message - content_blocks: list[Any] = [] + # Build content blocks as dicts; ``AssistantMessageResponse`` + # validates them into the proper discriminated union variants + # (``Text/ThinkingAssistantMessageResponseContentItem``). + content_blocks: list[dict[str, Any]] = [] if msg.reasoning: - content_blocks.append(CohereThinkingContent(thinking=msg.reasoning)) + content_blocks.append({"type": "thinking", "thinking": msg.reasoning}) if msg.content: - content_blocks.append(CohereTextContent(text=msg.content)) + content_blocks.append({"type": "text", "text": msg.content}) - tool_calls: list[CohereToolCallV2] | None = None + tool_calls: list[ToolCallV2] | None = None if msg.tool_calls: tool_calls = [ - CohereToolCallV2( + ToolCallV2( id=tc.id, - function=CohereToolCallFunction( + function=ToolCallV2Function( name=tc.function.name, arguments=tc.function.arguments, ), @@ -565,12 +587,10 @@ def _chat_completion_to_v2( if tool_calls and msg.reasoning: tool_plan = msg.reasoning content_blocks = [ - blk - for blk in content_blocks - if not isinstance(blk, CohereThinkingContent) + blk for blk in content_blocks if blk.get("type") != "thinking" ] - assistant_msg = CohereAssistantMessageResponse( + assistant_msg = AssistantMessageResponse( content=content_blocks or None, tool_calls=tool_calls, tool_plan=tool_plan, @@ -588,28 +608,31 @@ def _chat_completion_to_v2( ) @staticmethod - def _extract_citations_if_any(msg: Any) -> list[CohereCitation] | None: + def _extract_citations_if_any(msg: Any) -> list[Citation] | None: """Coerce ``ChatMessage.citations`` into the Cohere v2 wire shape. ``ChatMessage`` natively carries a ``citations: list[Citation] | None`` field (see :mod:`vllm.entrypoints.openai.engine.protocol`). Renderers/parsers (e.g. the ``cohere2`` reasoning parser) populate - it. We map each :class:`vllm...Citation` into the - :class:`CohereCitation` wire model, preserving sources and span. + it. We map each :class:`vllm...Citation` into the SDK + :class:`cohere.types.Citation` wire model, preserving sources and + span. """ raw = getattr(msg, "citations", None) if not raw: return None - out: list[CohereCitation] = [] + out: list[Citation] = [] for c in raw: try: - if isinstance(c, CohereCitation): + if isinstance(c, Citation): out.append(c) continue if hasattr(c, "model_dump"): - c = c.model_dump(exclude_none=True) - if isinstance(c, dict): - out.append(CohereCitation.model_validate(c)) + payload = c.model_dump(exclude_none=True) + else: + payload = c + if isinstance(payload, dict): + out.append(Citation.model_validate(payload)) except Exception: # pragma: no cover - defensive logger.debug("Skipping malformed citation: %r", c, exc_info=True) return out or None @@ -668,11 +691,12 @@ async def _chat_completion_stream_to_v2( chunk = ChatCompletionStreamResponse.model_validate_json(data_str) if not state.started: - yield _sse( - CohereChatMessageStartEvent( + yield _emit( + ChatMessageStartEvent( id=chunk.id, delta={"message": {"role": "assistant"}}, - ).model_dump_json(exclude_none=True) + ), + "message-start", ) state.started = True @@ -680,12 +704,10 @@ async def _chat_completion_stream_to_v2( if not chunk.choices: for ev in self._close_open_blocks(state): yield ev - yield _sse( - self._build_message_end_event( - chunk_id=chunk.id, - finish_reason=state.finish_reason, - usage_chunk=chunk, - ) + yield self._build_message_end_event( + chunk_id=chunk.id, + finish_reason=state.finish_reason, + usage_chunk=chunk, ) continue @@ -744,25 +766,27 @@ def _handle_thinking_delta( state.active_block = "thinking" state.active_block_index = idx events.append( - _sse( - CohereChatContentStartEvent( + _emit( + ChatContentStartEvent( index=idx, delta={ "message": { "content": {"type": "thinking", "thinking": ""} } }, - ).model_dump_json(exclude_none=True) + ), + "content-start", ) ) events.append( - _sse( - CohereChatContentDeltaEvent( + _emit( + ChatContentDeltaEvent( index=state.active_block_index, delta={ "message": {"content": {"thinking": delta_text}} }, - ).model_dump_json(exclude_none=True) + ), + "content-delta", ) ) return events @@ -777,23 +801,25 @@ def _handle_text_delta( state.active_block = "text" state.active_block_index = idx events.append( - _sse( - CohereChatContentStartEvent( + _emit( + ChatContentStartEvent( index=idx, delta={ "message": { "content": {"type": "text", "text": ""} } }, - ).model_dump_json(exclude_none=True) + ), + "content-start", ) ) events.append( - _sse( - CohereChatContentDeltaEvent( + _emit( + ChatContentDeltaEvent( index=state.active_block_index, delta={"message": {"content": {"text": delta_text}}}, - ).model_dump_json(exclude_none=True) + ), + "content-delta", ) ) return events @@ -814,8 +840,8 @@ def _handle_tool_call_deltas( state.active_tool_index = tc_index state.active_block = "tool_call" events.append( - _sse( - CohereChatToolCallStartEvent( + _emit( + ChatToolCallStartEvent( index=tc_index, delta={ "message": { @@ -833,15 +859,16 @@ def _handle_tool_call_deltas( } } }, - ).model_dump_json(exclude_none=True) + ), + "tool-call-start", ) ) continue if fn and fn.arguments: events.append( - _sse( - CohereChatToolCallDeltaEvent( + _emit( + ChatToolCallDeltaEvent( index=tc_index, delta={ "message": { @@ -852,7 +879,8 @@ def _handle_tool_call_deltas( } } }, - ).model_dump_json(exclude_none=True) + ), + "tool-call-delta", ) ) return events @@ -876,26 +904,26 @@ def _handle_citation_deltas( if not isinstance(payload, dict): continue try: - citation = CohereCitation.model_validate(payload) + citation = Citation.model_validate(payload) except Exception: # pragma: no cover - defensive logger.debug("Skipping malformed streamed citation: %r", payload) continue idx = state.next_citation_index() events.append( - _sse( - CohereCitationStartEvent( + _emit( + CitationStartEvent( index=idx, delta={"message": {"citations": citation.model_dump( exclude_none=True )}}, - ).model_dump_json(exclude_none=True) + ), + "citation-start", ) ) events.append( - _sse( - CohereCitationEndEvent(index=idx).model_dump_json( - exclude_none=True - ) + _emit( + CitationEndEvent(index=idx), + "citation-end", ) ) return events @@ -909,18 +937,16 @@ def _close_open_blocks_inner(self, state: _StreamState) -> list[str]: events: list[str] = [] if state.active_block in ("text", "thinking"): events.append( - _sse( - CohereChatContentEndEvent( - index=state.active_block_index, - ).model_dump_json(exclude_none=True) + _emit( + ChatContentEndEvent(index=state.active_block_index), + "content-end", ) ) elif state.active_block == "tool_call": events.append( - _sse( - CohereChatToolCallEndEvent( - index=state.active_tool_index, - ).model_dump_json(exclude_none=True) + _emit( + ChatToolCallEndEvent(index=state.active_tool_index), + "tool-call-end", ) ) state.active_block = None @@ -950,9 +976,10 @@ def _build_message_end_event( "output_tokens": completion, }, } - return CohereChatMessageEndEvent( - id=chunk_id, delta=delta - ).model_dump_json(exclude_none=True) + return _emit( + ChatMessageEndEvent(id=chunk_id, delta=delta), + "message-end", + ) # ================================================================== # Helpers for the router diff --git a/vllm/entrypoints/generate/api_router.py b/vllm/entrypoints/generate/api_router.py index 5fdc56bd8264..2c9d696a5153 100644 --- a/vllm/entrypoints/generate/api_router.py +++ b/vllm/entrypoints/generate/api_router.py @@ -61,7 +61,16 @@ async def init_generate_state( ): from vllm.entrypoints.anthropic.serving import AnthropicServingMessages from vllm.entrypoints.chat_utils import load_chat_template - from vllm.entrypoints.cohere.serving import CohereServingChatV2 + + try: + from vllm.entrypoints.cohere.serving import CohereServingChatV2 + except ImportError: + # The Cohere serving handler depends on the optional `cohere` SDK + # for its wire-format protocol models. When it isn't installed, + # `register_cohere_api_router` already skips registering the route, + # so we simply leave `cohere_serving_chat_v2` unset. + CohereServingChatV2 = None # type: ignore[assignment,misc] + from vllm.entrypoints.mcp.tool_server import ( DemoToolServer, MCPToolServer, @@ -199,7 +208,7 @@ async def init_generate_state( enable_force_include_usage=args.enable_force_include_usage, default_chat_template_kwargs=args.default_chat_template_kwargs, ) - if "generate" in supported_tasks + if CohereServingChatV2 is not None and "generate" in supported_tasks else None ) state.serving_tokens = ( From 68ab8b927f87409d1d7a2a934674c1dc2a5c9c14 Mon Sep 17 00:00:00 2001 From: Andrew Berneshawi Date: Mon, 22 Jun 2026 23:25:25 -0400 Subject: [PATCH 08/30] improve event handling --- vllm/entrypoints/cohere/protocol.py | 96 ++++++++++++------ vllm/entrypoints/cohere/serving.py | 152 ++++++++++++++-------------- 2 files changed, 144 insertions(+), 104 deletions(-) diff --git a/vllm/entrypoints/cohere/protocol.py b/vllm/entrypoints/cohere/protocol.py index b4eb679360b0..9cb47cc62481 100644 --- a/vllm/entrypoints/cohere/protocol.py +++ b/vllm/entrypoints/cohere/protocol.py @@ -27,24 +27,14 @@ import time from typing import Any, Literal +from cohere import types as _sdk from cohere.types import ( AssistantChatMessageV2, AssistantMessageResponse, - ChatContentDeltaEvent, - ChatContentEndEvent, - ChatContentStartEvent, - ChatMessageEndEvent, - ChatMessageStartEvent, ChatMessageV2, ChatRequestSafetyMode, - ChatToolCallDeltaEvent, - ChatToolCallEndEvent, - ChatToolCallStartEvent, - ChatToolPlanDeltaEvent, Citation, - CitationEndEvent, CitationOptions, - CitationStartEvent, Document, ResponseFormatV2, SystemChatMessageV2, @@ -56,21 +46,13 @@ ) from pydantic import BaseModel, Field, field_validator -# Re-export the SDK types so that ``vllm.entrypoints.cohere.serving`` and -# friends can import everything they need from this module. +# Re-export the SDK wire-format types alongside our local extensions so +# ``vllm.entrypoints.cohere.serving`` and friends can import everything +# they need from this module. __all__ = [ "AssistantChatMessageV2", "AssistantMessageResponse", - "ChatContentDeltaEvent", - "ChatContentEndEvent", - "ChatContentStartEvent", - "ChatMessageEndEvent", - "ChatMessageStartEvent", "ChatMessageV2", - "ChatToolCallDeltaEvent", - "ChatToolCallEndEvent", - "ChatToolCallStartEvent", - "ChatToolPlanDeltaEvent", "Citation", "CitationEndEvent", "CitationOptions", @@ -83,12 +65,21 @@ "CohereUsage", "CohereUsageBilledUnits", "CohereUsageTokens", + "ContentDeltaEvent", + "ContentEndEvent", + "ContentStartEvent", "Document", + "MessageEndEvent", + "MessageStartEvent", "ResponseFormatV2", "SystemChatMessageV2", "Thinking", + "ToolCallDeltaEvent", + "ToolCallEndEvent", + "ToolCallStartEvent", "ToolCallV2", "ToolChatMessageV2", + "ToolPlanDeltaEvent", "ToolV2", "UserChatMessageV2", ] @@ -273,13 +264,60 @@ def model_post_init(self, __context: Any) -> None: # --------------------------------------------------------------------------- # # Cohere V2 streams a sequence of typed JSON events delivered as Server- -# Sent Events. Each event class is exported above (``ChatMessageStartEvent``, -# ``ChatContentDeltaEvent``, ``CitationStartEvent``, etc.). We don't build -# a discriminated union here because the SDK's event classes don't carry -# the ``type`` literal as a Pydantic field (it's handled by the SDK's own -# deserializer), so pydantic ``Field(discriminator="type")`` would reject -# it. Surfaces emit each event with ``model_dump_json()`` and clients -# parse them by reading ``type`` directly. +# Sent Events; each event's ``type`` field carries the discriminator. The +# SDK exposes a Pydantic model per event but none of them declare ``type`` +# as a field (the SDK relies on its own deserializer for discrimination), +# so a naive ``model_dump_json()`` would silently drop the discriminator +# and break clients that demux on ``type``. +# +# We therefore subclass each SDK event and bake the wire-format ``type`` +# string in as a ``Literal`` field with a default. ``model_dump()`` now +# emits ``type`` for free, and surfaces can simply construct the event +# class and serialize it -- no manual ``type`` parameter needed. # # See https://docs.cohere.com/v2/docs/streaming and the OpenAPI # ``StreamedChatResponseV2`` schema for the wire-format reference. + + +class MessageStartEvent(_sdk.ChatMessageStartEvent): + type: Literal["message-start"] = "message-start" + + +class ContentStartEvent(_sdk.ChatContentStartEvent): + type: Literal["content-start"] = "content-start" + + +class ContentDeltaEvent(_sdk.ChatContentDeltaEvent): + type: Literal["content-delta"] = "content-delta" + + +class ContentEndEvent(_sdk.ChatContentEndEvent): + type: Literal["content-end"] = "content-end" + + +class ToolPlanDeltaEvent(_sdk.ChatToolPlanDeltaEvent): + type: Literal["tool-plan-delta"] = "tool-plan-delta" + + +class ToolCallStartEvent(_sdk.ChatToolCallStartEvent): + type: Literal["tool-call-start"] = "tool-call-start" + + +class ToolCallDeltaEvent(_sdk.ChatToolCallDeltaEvent): + type: Literal["tool-call-delta"] = "tool-call-delta" + + +class ToolCallEndEvent(_sdk.ChatToolCallEndEvent): + type: Literal["tool-call-end"] = "tool-call-end" + + +class CitationStartEvent(_sdk.CitationStartEvent): + type: Literal["citation-start"] = "citation-start" + + +class CitationEndEvent(_sdk.CitationEndEvent): + type: Literal["citation-end"] = "citation-end" + + +class MessageEndEvent(_sdk.ChatMessageEndEvent): + type: Literal["message-end"] = "message-end" diff --git a/vllm/entrypoints/cohere/serving.py b/vllm/entrypoints/cohere/serving.py index eeadab9dd657..2f52fb2881fe 100644 --- a/vllm/entrypoints/cohere/serving.py +++ b/vllm/entrypoints/cohere/serving.py @@ -25,22 +25,13 @@ import logging import time from collections.abc import AsyncGenerator +from enum import StrEnum from typing import TYPE_CHECKING, Any from cohere.types import ( AssistantChatMessageV2, AssistantMessageResponse, - ChatContentDeltaEvent, - ChatContentEndEvent, - ChatContentStartEvent, - ChatMessageEndEvent, - ChatMessageStartEvent, - ChatToolCallDeltaEvent, - ChatToolCallEndEvent, - ChatToolCallStartEvent, Citation, - CitationEndEvent, - CitationStartEvent, SystemChatMessageV2, ToolCallV2, ToolCallV2Function, @@ -48,16 +39,27 @@ UserChatMessageV2, ) from fastapi import Request +from pydantic import BaseModel from vllm.engine.protocol import EngineClient from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption from vllm.entrypoints.cohere.protocol import ( + CitationEndEvent, + CitationStartEvent, CohereChatV2Request, CohereChatV2Response, CohereFinishReason, CohereUsage, CohereUsageBilledUnits, CohereUsageTokens, + ContentDeltaEvent, + ContentEndEvent, + ContentStartEvent, + MessageEndEvent, + MessageStartEvent, + ToolCallDeltaEvent, + ToolCallEndEvent, + ToolCallStartEvent, ) from vllm.entrypoints.logger import RequestLogger from vllm.entrypoints.openai.chat_completion.protocol import ( @@ -95,19 +97,29 @@ def _sse(data: str) -> str: return f"data: {data}\n\n" -def _emit(event: Any, type_str: str) -> str: - """Serialize an SDK stream event with the wire-format ``type`` field. +def _emit(event: BaseModel) -> str: + """Serialize a typed stream event into an SSE frame. - The cohere SDK's stream-event Pydantic models (``ChatMessageStartEvent`` - etc.) don't declare ``type`` as a field -- the SDK handles type - discrimination in its own deserializer -- so a vanilla - ``model_dump_json()`` would drop the discriminator and break clients - that demux on ``type``. We dump the event, splice ``type`` back in, - and re-serialize as JSON before wrapping in an SSE frame. + The typed event classes in ``vllm.entrypoints.cohere.protocol`` + (``MessageStartEvent``, ``ContentStartEvent``, ``CitationStartEvent``, + ...) bake the wire-format ``type`` field into the model definition, + so a plain ``model_dump_json()`` already carries the discriminator. """ - payload = event.model_dump(exclude_none=True) - payload["type"] = type_str - return _sse(json.dumps(payload)) + return _sse(event.model_dump_json(exclude_none=True)) + + +class ContentBlockType(StrEnum): + """Wire-format / internal discriminator for chat content blocks. + + ``THINKING`` and ``TEXT`` are the documented Cohere v2 content-block + type discriminators on the wire. ``TOOL_CALL`` is reserved for the + internal stream state machine (see :class:`_StreamState`) when a + tool call is the currently open block; it is never serialized. + """ + + THINKING = "thinking" + TEXT = "text" + TOOL_CALL = "tool_call" # Mapping of vLLM/OpenAI finish reasons to Cohere's enum. @@ -562,9 +574,13 @@ def _chat_completion_to_v2( # (``Text/ThinkingAssistantMessageResponseContentItem``). content_blocks: list[dict[str, Any]] = [] if msg.reasoning: - content_blocks.append({"type": "thinking", "thinking": msg.reasoning}) + content_blocks.append( + {"type": ContentBlockType.THINKING, "thinking": msg.reasoning} + ) if msg.content: - content_blocks.append({"type": "text", "text": msg.content}) + content_blocks.append( + {"type": ContentBlockType.TEXT, "text": msg.content} + ) tool_calls: list[ToolCallV2] | None = None if msg.tool_calls: @@ -587,7 +603,9 @@ def _chat_completion_to_v2( if tool_calls and msg.reasoning: tool_plan = msg.reasoning content_blocks = [ - blk for blk in content_blocks if blk.get("type") != "thinking" + blk + for blk in content_blocks + if blk.get("type") != ContentBlockType.THINKING ] assistant_msg = AssistantMessageResponse( @@ -692,11 +710,10 @@ async def _chat_completion_stream_to_v2( if not state.started: yield _emit( - ChatMessageStartEvent( + MessageStartEvent( id=chunk.id, delta={"message": {"role": "assistant"}}, - ), - "message-start", + ) ) state.started = True @@ -760,33 +777,34 @@ def _handle_thinking_delta( self, state: _StreamState, delta_text: str ) -> list[str]: events: list[str] = [] - if state.active_block != "thinking": + if state.active_block != ContentBlockType.THINKING: events.extend(self._close_open_blocks_inner(state)) idx = state.next_content_index() - state.active_block = "thinking" + state.active_block = ContentBlockType.THINKING state.active_block_index = idx events.append( _emit( - ChatContentStartEvent( + ContentStartEvent( index=idx, delta={ "message": { - "content": {"type": "thinking", "thinking": ""} + "content": { + "type": ContentBlockType.THINKING, + "thinking": "", + } } }, - ), - "content-start", + ) ) ) events.append( _emit( - ChatContentDeltaEvent( + ContentDeltaEvent( index=state.active_block_index, delta={ "message": {"content": {"thinking": delta_text}} }, - ), - "content-delta", + ) ) ) return events @@ -795,31 +813,32 @@ def _handle_text_delta( self, state: _StreamState, delta_text: str ) -> list[str]: events: list[str] = [] - if state.active_block != "text": + if state.active_block != ContentBlockType.TEXT: events.extend(self._close_open_blocks_inner(state)) idx = state.next_content_index() - state.active_block = "text" + state.active_block = ContentBlockType.TEXT state.active_block_index = idx events.append( _emit( - ChatContentStartEvent( + ContentStartEvent( index=idx, delta={ "message": { - "content": {"type": "text", "text": ""} + "content": { + "type": ContentBlockType.TEXT, + "text": "", + } } }, - ), - "content-start", + ) ) ) events.append( _emit( - ChatContentDeltaEvent( + ContentDeltaEvent( index=state.active_block_index, delta={"message": {"content": {"text": delta_text}}}, - ), - "content-delta", + ) ) ) return events @@ -838,10 +857,10 @@ def _handle_tool_call_deltas( events.extend(self._close_open_blocks_inner(state)) state.tool_calls_seen.add(tc_index) state.active_tool_index = tc_index - state.active_block = "tool_call" + state.active_block = ContentBlockType.TOOL_CALL events.append( _emit( - ChatToolCallStartEvent( + ToolCallStartEvent( index=tc_index, delta={ "message": { @@ -859,8 +878,7 @@ def _handle_tool_call_deltas( } } }, - ), - "tool-call-start", + ) ) ) continue @@ -868,7 +886,7 @@ def _handle_tool_call_deltas( if fn and fn.arguments: events.append( _emit( - ChatToolCallDeltaEvent( + ToolCallDeltaEvent( index=tc_index, delta={ "message": { @@ -879,8 +897,7 @@ def _handle_tool_call_deltas( } } }, - ), - "tool-call-delta", + ) ) ) return events @@ -916,16 +933,10 @@ def _handle_citation_deltas( delta={"message": {"citations": citation.model_dump( exclude_none=True )}}, - ), - "citation-start", - ) - ) - events.append( - _emit( - CitationEndEvent(index=idx), - "citation-end", + ) ) ) + events.append(_emit(CitationEndEvent(index=idx))) return events # -- block lifecycle helpers -------------------------------------- @@ -935,19 +946,13 @@ def _close_open_blocks(self, state: _StreamState) -> list[str]: def _close_open_blocks_inner(self, state: _StreamState) -> list[str]: events: list[str] = [] - if state.active_block in ("text", "thinking"): + if state.active_block in (ContentBlockType.TEXT, ContentBlockType.THINKING): events.append( - _emit( - ChatContentEndEvent(index=state.active_block_index), - "content-end", - ) + _emit(ContentEndEvent(index=state.active_block_index)) ) - elif state.active_block == "tool_call": + elif state.active_block == ContentBlockType.TOOL_CALL: events.append( - _emit( - ChatToolCallEndEvent(index=state.active_tool_index), - "tool-call-end", - ) + _emit(ToolCallEndEvent(index=state.active_tool_index)) ) state.active_block = None state.active_block_index = None @@ -976,10 +981,7 @@ def _build_message_end_event( "output_tokens": completion, }, } - return _emit( - ChatMessageEndEvent(id=chunk_id, delta=delta), - "message-end", - ) + return _emit(MessageEndEvent(id=chunk_id, delta=delta)) # ================================================================== # Helpers for the router @@ -1010,7 +1012,7 @@ class _StreamState: def __init__(self) -> None: self.started: bool = False self.finish_reason: str | None = None - self.active_block: str | None = None # "text" | "thinking" | "tool_call" + self.active_block: ContentBlockType | None = None self.active_block_index: int | None = None self.active_tool_index: int | None = None self._next_index: int = 0 From 4a703e8f2c9b3fcefc8d1c70a42b766a72ee4356 Mon Sep 17 00:00:00 2001 From: Andrew Berneshawi Date: Tue, 23 Jun 2026 00:15:27 -0400 Subject: [PATCH 09/30] rename api to include /cohere --- vllm/entrypoints/cohere/api_router.py | 10 +++++----- vllm/entrypoints/cohere/protocol.py | 4 ++-- vllm/entrypoints/cohere/serving.py | 6 +++--- vllm/entrypoints/openai/engine/protocol.py | 2 +- 4 files changed, 11 insertions(+), 11 deletions(-) diff --git a/vllm/entrypoints/cohere/api_router.py b/vllm/entrypoints/cohere/api_router.py index 3e14d07c8344..3cc019a952ee 100644 --- a/vllm/entrypoints/cohere/api_router.py +++ b/vllm/entrypoints/cohere/api_router.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""FastAPI router for the Cohere Chat v2 API (``POST /v2/chat``). +"""FastAPI router for the Cohere Chat v2 API (``POST /cohere/v2/chat``). The Cohere v2 protocol models are sourced from the official ``cohere`` Python SDK (``pip install cohere``). To keep that an *optional* @@ -20,7 +20,7 @@ def attach_router(app: FastAPI) -> None: - """Register ``POST /v2/chat`` on ``app``. + """Register ``POST /cohere/v2/chat`` on ``app``. No-op (with an info log) when the optional ``cohere`` SDK isn't installed, since the v2 protocol models live there. @@ -29,7 +29,7 @@ def attach_router(app: FastAPI) -> None: import cohere # noqa: F401 -- dependency probe except ImportError: logger.info( - "cohere SDK not installed; /v2/chat endpoint disabled. " + "cohere SDK not installed; /cohere/v2/chat endpoint disabled. " "Install with `pip install cohere` to enable it." ) return @@ -73,7 +73,7 @@ def _error_response( ) @router.post( - "/v2/chat", + "/cohere/v2/chat", dependencies=[Depends(validate_json_request)], responses={ HTTPStatus.OK.value: {"content": {"text/event-stream": {}}}, @@ -97,7 +97,7 @@ async def chat_v2(request: CohereChatV2Request, raw_request: Request): try: result = await handler.create_chat_v2(request, raw_request) except Exception as e: # noqa: BLE001 - report as 500 for parity - logger.exception("Error in /v2/chat: %s", e) + logger.exception("Error in /cohere/v2/chat: %s", e) return JSONResponse( status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, content=CohereError(message=str(e)).model_dump(exclude_none=True), diff --git a/vllm/entrypoints/cohere/protocol.py b/vllm/entrypoints/cohere/protocol.py index 9cb47cc62481..cf03bac16e01 100644 --- a/vllm/entrypoints/cohere/protocol.py +++ b/vllm/entrypoints/cohere/protocol.py @@ -17,7 +17,7 @@ discriminator]``). Importing this module pulls in the ``cohere`` package. The router that -mounts ``POST /v2/chat`` guards on that import succeeding so vLLM still +mounts ``POST /cohere/v2/chat`` guards on that import succeeding so vLLM still boots without the SDK installed. See https://docs.cohere.com/reference/chat for the upstream spec. @@ -91,7 +91,7 @@ class CohereError(BaseModel): - """Top-level error body returned by ``/v2/chat`` error responses. + """Top-level error body returned by ``/cohere/v2/chat`` error responses. Cohere's documented error schemas are uniform: ``{message, id}``. """ diff --git a/vllm/entrypoints/cohere/serving.py b/vllm/entrypoints/cohere/serving.py index 2f52fb2881fe..ba32a8e601bd 100644 --- a/vllm/entrypoints/cohere/serving.py +++ b/vllm/entrypoints/cohere/serving.py @@ -4,7 +4,7 @@ # https://github.com/vllm-project/vllm/blob/main/vllm/entrypoints/anthropic/serving.py """Cohere Chat v2 API serving handler. -Implements ``POST /v2/chat`` by translating the incoming Cohere v2 request +Implements ``POST /cohere/v2/chat`` by translating the incoming Cohere v2 request into a standard :class:`ChatCompletionRequest` and delegating to :class:`OpenAIServingChat`. The actual prompt rendering is handled by vLLM's renderer pipeline (``vllm.renderers``): @@ -143,7 +143,7 @@ def _map_finish_reason(reason: str | None) -> CohereFinishReason: class CohereServingChatV2(OpenAIServingChat): - """Handler for the Cohere Chat v2 API (``POST /v2/chat``). + """Handler for the Cohere Chat v2 API (``POST /cohere/v2/chat``). The handler is intentionally thin: it converts the v2 request into a :class:`ChatCompletionRequest` (preserving Cohere-specific fields such @@ -199,7 +199,7 @@ async def create_chat_v2( request: CohereChatV2Request, raw_request: Request | None = None, ) -> AsyncGenerator[str, None] | CohereChatV2Response | ErrorResponse: - """Implements ``POST /v2/chat``.""" + """Implements ``POST /cohere/v2/chat``.""" if logger.isEnabledFor(logging.DEBUG): logger.debug( "Received Cohere v2 chat request %s", request.model_dump_json() diff --git a/vllm/entrypoints/openai/engine/protocol.py b/vllm/entrypoints/openai/engine/protocol.py index c8d3f34998cc..92d95b185128 100644 --- a/vllm/entrypoints/openai/engine/protocol.py +++ b/vllm/entrypoints/openai/engine/protocol.py @@ -361,7 +361,7 @@ class Citation(OpenAIBaseModel): so that grounded models (e.g. Cohere Command-family) can surface citation metadata through the standard chat completion responses. It is consumed by surfaces that expose citations (such as Cohere's - ``/v2/chat`` endpoint). + ``/cohere/v2/chat`` endpoint). """ start: int | None = None From 702b85eb2dca4891f4a422c4b3aa298e7419aa23 Mon Sep 17 00:00:00 2001 From: Andrew Berneshawi Date: Tue, 23 Jun 2026 02:00:33 -0400 Subject: [PATCH 10/30] address some PR comments --- vllm/entrypoints/openai/engine/protocol.py | 7 ++++--- vllm/renderers/cohere.py | 10 +++------- 2 files changed, 7 insertions(+), 10 deletions(-) diff --git a/vllm/entrypoints/openai/engine/protocol.py b/vllm/entrypoints/openai/engine/protocol.py index 92d95b185128..b4e932b3a469 100644 --- a/vllm/entrypoints/openai/engine/protocol.py +++ b/vllm/entrypoints/openai/engine/protocol.py @@ -375,9 +375,10 @@ class Citation(OpenAIBaseModel): content_index: int | None = None """Index of the content block this citation refers to (when the message has multiple content blocks).""" - type: Literal["TEXT_CONTENT", "PLAN"] | None = None - """Whether this citation grounds the user-visible text or a planning - block. ``None`` means unspecified.""" + type: Literal["TEXT_CONTENT", "THINKING_CONTENT", "PLAN"] | None = None + """Which kind of content block this citation grounds: the user-visible + text (``TEXT_CONTENT``), a thinking block (``THINKING_CONTENT``), or a + tool-plan block (``PLAN``). ``None`` means unspecified.""" class DeltaMessage(OpenAIBaseModel): diff --git a/vllm/renderers/cohere.py b/vllm/renderers/cohere.py index c3df3cf7a6c8..fe89a5c3330c 100644 --- a/vllm/renderers/cohere.py +++ b/vllm/renderers/cohere.py @@ -283,11 +283,11 @@ def _build_render_config( } # Optional template overrides - for k in ("template_id", "template", "template_jinja"): + for k in ("template_id", "template_jinja"): if v := chat_template_kwargs.get(k): config[k] = v - if (use_jinja := chat_template_kwargs.get("use_jinja")) is not None: - config["use_jinja"] = bool(use_jinja) + # Only support Jinja with vllm + config["use_jinja"] = True # Documents documents = chat_template_kwargs.get("documents") or [] @@ -313,14 +313,10 @@ def _build_render_config( if t in ("enabled", "disabled"): config["reasoning_type"] = t - if (rp := chat_template_kwargs.get("response_prefix")) is not None: - config["response_prefix"] = str(rp) if (di := chat_template_kwargs.get("dev_instruction")) is not None: config["dev_instruction"] = str(di) if (extra := chat_template_kwargs.get("additional_template_fields")): config["additional_template_fields"] = dict(extra) - if (esc := chat_template_kwargs.get("escaped_special_tokens")): - config["escaped_special_tokens"] = dict(esc) # JSON / structured outputs if (rf := chat_template_kwargs.get("response_format")) is not None: From 3600591cace0cda7f82dbcc4fecdd35bed6c756d Mon Sep 17 00:00:00 2001 From: Andrew Berneshawi Date: Tue, 23 Jun 2026 05:48:51 -0400 Subject: [PATCH 11/30] fix tool plan output handling --- vllm/entrypoints/cohere/serving.py | 46 ++++++++++++++++++++++++++---- 1 file changed, 41 insertions(+), 5 deletions(-) diff --git a/vllm/entrypoints/cohere/serving.py b/vllm/entrypoints/cohere/serving.py index ba32a8e601bd..942243142d61 100644 --- a/vllm/entrypoints/cohere/serving.py +++ b/vllm/entrypoints/cohere/serving.py @@ -60,6 +60,7 @@ ToolCallDeltaEvent, ToolCallEndEvent, ToolCallStartEvent, + ToolPlanDeltaEvent, ) from vllm.entrypoints.logger import RequestLogger from vllm.entrypoints.openai.chat_completion.protocol import ( @@ -172,6 +173,7 @@ def __init__( enable_prompt_tokens_details: bool = False, enable_force_include_usage: bool = False, default_chat_template_kwargs: dict[str, Any] | None = None, + surface_reasoning_as_tool_plan: bool = False, ) -> None: super().__init__( engine_client=engine_client, @@ -189,6 +191,19 @@ def __init__( enable_force_include_usage=enable_force_include_usage, default_chat_template_kwargs=default_chat_template_kwargs, ) + # When True, the assistant's reasoning text is surfaced as Cohere's + # ``tool_plan`` (non-streaming) or as ``tool-plan-delta`` events + # (streaming) whenever the model emits tool calls. This matches the + # behavior of older non-reasoning Command models, which produce a + # tool plan in place of a thinking block. + # + # When False (default), reasoning is always surfaced as a thinking + # content block; this matches reasoning Command models that + # emit both a thinking block and tool calls. + # + # TODO: replace this manual flag with automatic detection based on + # the model's capabilities once we have that config + self._surface_reasoning_as_tool_plan = surface_reasoning_as_tool_plan # ------------------------------------------------------------------ # Public entry point @@ -596,11 +611,17 @@ def _chat_completion_to_v2( ] # Cohere's ``tool_plan`` is the planning text emitted before tool - # calls. The closest analogue in chat completions is the reasoning - # field; if a tool call is present, surface reasoning as tool_plan - # rather than as a thinking content block to match Cohere clients. + # calls on older, non-reasoning Command models. For those models we + # surface ``reasoning`` as ``tool_plan`` and drop the thinking + # block. Reasoning Command models emit a regular thinking + # block alongside tool calls, so the default leaves the thinking + # block in place and never sets ``tool_plan``. tool_plan: str | None = None - if tool_calls and msg.reasoning: + if ( + self._surface_reasoning_as_tool_plan + and tool_calls + and msg.reasoning + ): tool_plan = msg.reasoning content_blocks = [ blk @@ -776,7 +797,22 @@ async def _chat_completion_stream_to_v2( def _handle_thinking_delta( self, state: _StreamState, delta_text: str ) -> list[str]: - events: list[str] = [] + # Non-reasoning Command models: emit ``tool-plan-delta`` events + # directly instead of opening a thinking content block. The + # ``tool-plan-delta`` event has no start/end pair around it. + if self._surface_reasoning_as_tool_plan: + events: list[str] = list(self._close_open_blocks_inner(state)) + events.append( + _emit( + ToolPlanDeltaEvent( + delta={"message": {"tool_plan": delta_text}}, + ) + ) + ) + return events + + # Reasoning model (default): open / continue a thinking block. + events = [] if state.active_block != ContentBlockType.THINKING: events.extend(self._close_open_blocks_inner(state)) idx = state.next_content_index() From 35ec84deec8e2c30d10a559ac5746462a0ca6ef8 Mon Sep 17 00:00:00 2001 From: Andrew Berneshawi Date: Tue, 23 Jun 2026 23:45:53 -0400 Subject: [PATCH 12/30] cleanup --- vllm/entrypoints/cohere/serving.py | 14 +++++++------- vllm/renderers/cohere.py | 29 ----------------------------- 2 files changed, 7 insertions(+), 36 deletions(-) diff --git a/vllm/entrypoints/cohere/serving.py b/vllm/entrypoints/cohere/serving.py index 942243142d61..c3b5770d384f 100644 --- a/vllm/entrypoints/cohere/serving.py +++ b/vllm/entrypoints/cohere/serving.py @@ -801,7 +801,7 @@ def _handle_thinking_delta( # directly instead of opening a thinking content block. The # ``tool-plan-delta`` event has no start/end pair around it. if self._surface_reasoning_as_tool_plan: - events: list[str] = list(self._close_open_blocks_inner(state)) + events: list[str] = list(self._close_open_blocks(state)) events.append( _emit( ToolPlanDeltaEvent( @@ -814,7 +814,7 @@ def _handle_thinking_delta( # Reasoning model (default): open / continue a thinking block. events = [] if state.active_block != ContentBlockType.THINKING: - events.extend(self._close_open_blocks_inner(state)) + events.extend(self._close_open_blocks(state)) idx = state.next_content_index() state.active_block = ContentBlockType.THINKING state.active_block_index = idx @@ -850,7 +850,7 @@ def _handle_text_delta( ) -> list[str]: events: list[str] = [] if state.active_block != ContentBlockType.TEXT: - events.extend(self._close_open_blocks_inner(state)) + events.extend(self._close_open_blocks(state)) idx = state.next_content_index() state.active_block = ContentBlockType.TEXT state.active_block_index = idx @@ -890,7 +890,7 @@ def _handle_tool_call_deltas( if tool_id is not None and tc_index not in state.tool_calls_seen: # New tool call. Close any open content/tool block first. - events.extend(self._close_open_blocks_inner(state)) + events.extend(self._close_open_blocks(state)) state.tool_calls_seen.add(tc_index) state.active_tool_index = tc_index state.active_block = ContentBlockType.TOOL_CALL @@ -978,9 +978,9 @@ def _handle_citation_deltas( # -- block lifecycle helpers -------------------------------------- def _close_open_blocks(self, state: _StreamState) -> list[str]: - return self._close_open_blocks_inner(state) - - def _close_open_blocks_inner(self, state: _StreamState) -> list[str]: + """Emit ``content-end`` / ``tool-call-end`` for the currently open + block (if any) and reset the corresponding ``_StreamState`` slots. + """ events: list[str] = [] if state.active_block in (ContentBlockType.TEXT, ContentBlockType.THINKING): events.append( diff --git a/vllm/renderers/cohere.py b/vllm/renderers/cohere.py index fe89a5c3330c..ef3eb1bd3832 100644 --- a/vllm/renderers/cohere.py +++ b/vllm/renderers/cohere.py @@ -58,35 +58,6 @@ _DEFAULT_FORMAT = "cmd3" _VALID_FORMATS = ("cmd3", "cmd4") -# Keys consumed by this renderer that should not also be forwarded to the -# `apply_chat_template` path of the underlying tokenizer. -_RENDERER_CONSUMED_KEYS = frozenset( - { - "cohere_format", - "documents", - "available_tools", - "tools", - "safety_mode", - "citation_quality", - "citation_options", - "grounding", - "reasoning_type", - "response_prefix", - "json_schema", - "json_mode", - "response_format", - "skip_preamble", - "dev_instruction", - "platform_instruction", - "template", - "template_id", - "template_jinja", - "use_jinja", - "additional_template_fields", - "escaped_special_tokens", - } -) - def _try_import_melody(): try: From e7dc63e2d6727402da96ab1a51c17d67c2710d26 Mon Sep 17 00:00:00 2001 From: Andrew Berneshawi Date: Tue, 23 Jun 2026 23:51:22 -0400 Subject: [PATCH 13/30] tool call logic fix --- vllm/entrypoints/cohere/serving.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/vllm/entrypoints/cohere/serving.py b/vllm/entrypoints/cohere/serving.py index c3b5770d384f..cf99fd76003d 100644 --- a/vllm/entrypoints/cohere/serving.py +++ b/vllm/entrypoints/cohere/serving.py @@ -885,10 +885,9 @@ def _handle_tool_call_deltas( events: list[str] = [] for tc in deltas: tc_index = tc.index - tool_id = tc.id fn = tc.function - if tool_id is not None and tc_index not in state.tool_calls_seen: + if tc_index not in state.tool_calls_seen: # New tool call. Close any open content/tool block first. events.extend(self._close_open_blocks(state)) state.tool_calls_seen.add(tc_index) @@ -901,7 +900,7 @@ def _handle_tool_call_deltas( delta={ "message": { "tool_calls": { - "id": tool_id, + "id": tc.id or "", "type": "function", "function": { "name": (fn.name if fn else "") From 5e411632df3e95b99b1be2e5fe5e6bfa14d0c92b Mon Sep 17 00:00:00 2001 From: Andrew Berneshawi Date: Wed, 24 Jun 2026 00:42:45 -0400 Subject: [PATCH 14/30] add cached tokens to usage --- vllm/entrypoints/cohere/serving.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/vllm/entrypoints/cohere/serving.py b/vllm/entrypoints/cohere/serving.py index cf99fd76003d..0b80b2e7599a 100644 --- a/vllm/entrypoints/cohere/serving.py +++ b/vllm/entrypoints/cohere/serving.py @@ -682,6 +682,9 @@ def _build_usage(response: ChatCompletionResponse) -> CohereUsage | None: return None prompt = response.usage.prompt_tokens completion = response.usage.completion_tokens or 0 + cached: int | None = None + if response.usage.prompt_tokens_details is not None: + cached = response.usage.prompt_tokens_details.cached_tokens return CohereUsage( billed_units=CohereUsageBilledUnits( input_tokens=prompt, @@ -691,6 +694,7 @@ def _build_usage(response: ChatCompletionResponse) -> CohereUsage | None: input_tokens=prompt, output_tokens=completion, ), + cached_tokens=cached, ) # ================================================================== From 4eb7620bd6fe6db220154390e3c426c8a02f11b6 Mon Sep 17 00:00:00 2001 From: Andrew Berneshawi Date: Wed, 24 Jun 2026 00:50:30 -0400 Subject: [PATCH 15/30] fixing streaming done behavior --- vllm/entrypoints/cohere/serving.py | 70 +++++++++++++++++++++++++----- 1 file changed, 58 insertions(+), 12 deletions(-) diff --git a/vllm/entrypoints/cohere/serving.py b/vllm/entrypoints/cohere/serving.py index 0b80b2e7599a..c890c989a1b0 100644 --- a/vllm/entrypoints/cohere/serving.py +++ b/vllm/entrypoints/cohere/serving.py @@ -98,6 +98,14 @@ def _sse(data: str) -> str: return f"data: {data}\n\n" +# Cohere v2 SSE stream terminator. Declared by the upstream OpenAPI spec +# (``x-fern-streaming.terminator: "[DONE]"`` on the ``/v2/chat`` stream +# operation) and observed by the cohere-python / Fern-generated clients +# (the Python SDK breaks its read loop on ``_sse.data == "[DONE]"``). +# Must be emitted after the closing ``message-end`` event. +_DONE_FRAME = _sse("[DONE]") + + def _emit(event: BaseModel) -> str: """Serialize a typed stream event into an SSE frame. @@ -727,11 +735,13 @@ async def _chat_completion_stream_to_v2( if not data_str: continue if data_str == "[DONE]": - # Cohere closes the stream with message-end (already - # emitted on the final usage chunk); nothing to do. - return + # OpenAI's stream terminator. Fall through to the + # post-loop cleanup so we always emit ``message-end`` + # even if the usage-only chunk was skipped. + break chunk = ChatCompletionStreamResponse.model_validate_json(data_str) + state.last_chunk_id = chunk.id if not state.started: yield _emit( @@ -751,6 +761,7 @@ async def _chat_completion_stream_to_v2( finish_reason=state.finish_reason, usage_chunk=chunk, ) + state.ended = True continue choice = chunk.choices[0] @@ -787,14 +798,42 @@ async def _chat_completion_stream_to_v2( except Exception as exc: logger.exception("Error converting chat completion stream to v2") - yield _sse( - json.dumps( - { - "type": "message-end", - "delta": {"error": str(exc), "finish_reason": "ERROR"}, - } + if state.started and not state.ended: + yield _sse( + json.dumps( + { + "type": "message-end", + "delta": { + "error": str(exc), + "finish_reason": "ERROR", + }, + } + ) ) + state.ended = True + yield _DONE_FRAME + return + + # Normal completion or ``[DONE]``: ensure ``message-end`` is always + # emitted. Upstream may close the stream without sending the final + # usage-only chunk (e.g. on shutdown, or when ``[DONE]`` is the only + # terminator); without this fallback Cohere clients would hang + # waiting for the closing event. + if state.started and not state.ended: + for ev in self._close_open_blocks(state): + yield ev + yield self._build_message_end_event( + chunk_id=state.last_chunk_id, + finish_reason=state.finish_reason, + usage_chunk=None, ) + state.ended = True + + # Stream terminator. Cohere's v2 SSE protocol ends every stream + # with ``data: [DONE]\n\n`` after ``message-end``; Fern-generated + # clients (Go/Java) and cohere-python all key their read loop off + # this sentinel. + yield _DONE_FRAME # -- per-delta helpers -------------------------------------------- @@ -1002,15 +1041,15 @@ def _build_message_end_event( self, chunk_id: str, finish_reason: str | None, - usage_chunk: ChatCompletionStreamResponse, + usage_chunk: ChatCompletionStreamResponse | None = None, ) -> str: delta: dict[str, Any] = { "finish_reason": _map_finish_reason(finish_reason), } - if usage_chunk.usage is not None: + if usage_chunk is not None and usage_chunk.usage is not None: prompt = usage_chunk.usage.prompt_tokens completion = usage_chunk.usage.completion_tokens or 0 - delta["usage"] = { + usage_block: dict[str, Any] = { "billed_units": { "input_tokens": prompt, "output_tokens": completion, @@ -1020,6 +1059,11 @@ def _build_message_end_event( "output_tokens": completion, }, } + if usage_chunk.usage.prompt_tokens_details is not None: + cached = usage_chunk.usage.prompt_tokens_details.cached_tokens + if cached is not None: + usage_block["cached_tokens"] = cached + delta["usage"] = usage_block return _emit(MessageEndEvent(id=chunk_id, delta=delta)) # ================================================================== @@ -1050,7 +1094,9 @@ class _StreamState: def __init__(self) -> None: self.started: bool = False + self.ended: bool = False self.finish_reason: str | None = None + self.last_chunk_id: str = "" self.active_block: ContentBlockType | None = None self.active_block_index: int | None = None self.active_tool_index: int | None = None From bd3ff9b858b21f39c86804d52faa4a000512a04a Mon Sep 17 00:00:00 2001 From: Andrew Berneshawi Date: Wed, 24 Jun 2026 01:27:36 -0400 Subject: [PATCH 16/30] max tokens fix and simplification --- vllm/entrypoints/cohere/protocol.py | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/vllm/entrypoints/cohere/protocol.py b/vllm/entrypoints/cohere/protocol.py index cf03bac16e01..2248dd248c51 100644 --- a/vllm/entrypoints/cohere/protocol.py +++ b/vllm/entrypoints/cohere/protocol.py @@ -24,7 +24,6 @@ """ from __future__ import annotations -import time from typing import Any, Literal from cohere import types as _sdk @@ -192,8 +191,8 @@ def _validate_model(cls, v: str) -> str: @field_validator("max_tokens") @classmethod def _validate_max_tokens(cls, v: int | None) -> int | None: - if v is not None and v <= 0: - raise ValueError("max_tokens must be positive") + if v is not None and v < 0: + raise ValueError("max_tokens must be non-negative") return v @@ -240,7 +239,10 @@ class CohereChatV2Response(BaseModel): Wraps the SDK :class:`AssistantMessageResponse` (the message shape) in the documented v2 response envelope (``id``, ``finish_reason``, - ``usage``, ``logprobs``). + ``usage``, ``logprobs``). The single constructor in + :class:`CohereServingChatV2._chat_completion_to_v2` is responsible for + supplying a non-empty ``id`` (falling back to a synthesized one if + the upstream response is missing it) to this model. """ id: str @@ -254,10 +256,6 @@ class CohereChatV2Response(BaseModel): default=None, description="KVTransfer parameters." ) - def model_post_init(self, __context: Any) -> None: - if not self.id: - self.id = f"chat_{int(time.time() * 1000)}" - # --------------------------------------------------------------------------- # Streaming events From 1c3b0b5edc20a7a21794cd52966609441b7aa5b9 Mon Sep 17 00:00:00 2001 From: Andrew Berneshawi Date: Wed, 24 Jun 2026 02:19:43 -0400 Subject: [PATCH 17/30] fix additional_template_fields handling --- vllm/renderers/cohere.py | 55 ++++++++++++++++++++++++++++++++++++++-- 1 file changed, 53 insertions(+), 2 deletions(-) diff --git a/vllm/renderers/cohere.py b/vllm/renderers/cohere.py index ef3eb1bd3832..cf1fdd74eec2 100644 --- a/vllm/renderers/cohere.py +++ b/vllm/renderers/cohere.py @@ -23,6 +23,14 @@ template overrides * ``cohere_format``: ``cmd3`` (default) or ``cmd4`` +Any *other* keys in ``chat_template_kwargs`` are forwarded verbatim to +the melody render config as ``additional_template_fields`` -- i.e. they +become Jinja variables accessible inside the template. This mirrors +vLLM's documented contract for ``chat_template_kwargs`` ("kwargs +accessible by the template"), so e.g. +``chat_template_kwargs={"reasoning_effort": "low"}`` resolves +``{{ reasoning_effort }}`` inside cmd3 / cmd4 templates. + Citations produced by Cohere models are surfaced through the standard ``ChatMessage.citations`` / ``DeltaMessage.citations`` fields (populated by the ``cohere2`` reasoning parser). @@ -58,6 +66,38 @@ _DEFAULT_FORMAT = "cmd3" _VALID_FORMATS = ("cmd3", "cmd4") +# Keys this renderer interprets directly from ``chat_template_kwargs`` and +# maps onto typed melody render-config fields. Everything *not* in this +# set is forwarded verbatim to melody as ``additional_template_fields`` +# (i.e. as Jinja template variables), so callers can write +# ``chat_template_kwargs = {"my_var": "..."}`` and have ``{{ my_var }}`` +# resolve inside the template -- matching vLLM's documented contract for +# ``chat_template_kwargs`` and avoiding the older nested-namespace form +# (``chat_template_kwargs.additional_template_fields.my_var``). +_RENDERER_CONSUMED_KEYS = frozenset( + { + "cohere_format", + "template_id", + "template_jinja", + "use_jinja", + "documents", + "available_tools", + "tools", + "reasoning_type", + "thinking", + "dev_instruction", + "response_format", + "json_schema", + "json_mode", + "safety_mode", + "citation_quality", + "citation_options", + "skip_preamble", + "grounding", + "platform_instruction", + } +) + def _try_import_melody(): try: @@ -286,8 +326,6 @@ def _build_render_config( if (di := chat_template_kwargs.get("dev_instruction")) is not None: config["dev_instruction"] = str(di) - if (extra := chat_template_kwargs.get("additional_template_fields")): - config["additional_template_fields"] = dict(extra) # JSON / structured outputs if (rf := chat_template_kwargs.get("response_format")) is not None: @@ -332,6 +370,19 @@ def _build_render_config( if (pi := chat_template_kwargs.get("platform_instruction")) is not None: config["platform_instruction"] = str(pi) + # Anything we didn't explicitly interpret above is forwarded to melody + # as a Jinja template variable. This matches vLLM's documented contract + # for ``chat_template_kwargs`` ("kwargs accessible by the template") + # and lets callers write e.g. ``{"reasoning_effort": "low"}`` directly + # without a nested ``additional_template_fields`` wrapper. + extra = { + k: v + for k, v in chat_template_kwargs.items() + if k not in _RENDERER_CONSUMED_KEYS + } + if extra: + config["additional_template_fields"] = extra + return fmt, config From 26a71f5760ef4d933b57462fc9ebec01e3a2ab65 Mon Sep 17 00:00:00 2001 From: Andrew Berneshawi Date: Wed, 24 Jun 2026 02:39:01 -0400 Subject: [PATCH 18/30] update tool plan comment and add type enum to rendering logic --- vllm/renderers/cohere.py | 69 +++++++++++++++++++++++++++++++--------- 1 file changed, 54 insertions(+), 15 deletions(-) diff --git a/vllm/renderers/cohere.py b/vllm/renderers/cohere.py index cf1fdd74eec2..4db773aad5c2 100644 --- a/vllm/renderers/cohere.py +++ b/vllm/renderers/cohere.py @@ -39,6 +39,7 @@ import copy import json +from enum import StrEnum from typing import TYPE_CHECKING, Any from vllm.config import VllmConfig @@ -66,6 +67,21 @@ _DEFAULT_FORMAT = "cmd3" _VALID_FORMATS = ("cmd3", "cmd4") + +class MelodyContentType(StrEnum): + """Wire-format discriminator for melody content blocks. + + These strings are what the cmd3 / cmd4 Jinja templates check against + in ``message.content[0].type`` (e.g. the ``thinking`` branch in + ``cmd4-v1.jinja``). Keep new values in sync with melody's template + schema. + """ + + TEXT = "text" + THINKING = "thinking" + IMAGE = "image" + DOCUMENT = "document" + # Keys this renderer interprets directly from ``chat_template_kwargs`` and # maps onto typed melody render-config fields. Everything *not* in this # set is forwarded verbatim to melody as ``additional_template_fields`` @@ -160,25 +176,30 @@ def _content_blocks(content: Any) -> list[dict[str, Any]]: if content is None: return [] if isinstance(content, str): - return [{"type": "text", "text": content}] + return [{"type": MelodyContentType.TEXT, "text": content}] blocks: list[dict[str, Any]] = [] for part in content: if isinstance(part, str): - blocks.append({"type": "text", "text": part}) + blocks.append({"type": MelodyContentType.TEXT, "text": part}) continue if not isinstance(part, dict): raise TypeError(f"Unexpected content part: {part!r}") - part_type = part.get("type", "text") + part_type = part.get("type", MelodyContentType.TEXT) if part_type in ("text", "input_text", "output_text", "refusal"): - blocks.append({"type": "text", "text": part.get("text", "")}) - elif part_type == "thinking": blocks.append( - {"type": "thinking", "thinking": part.get("thinking", "")} + {"type": MelodyContentType.TEXT, "text": part.get("text", "")} + ) + elif part_type == MelodyContentType.THINKING: + blocks.append( + { + "type": MelodyContentType.THINKING, + "thinking": part.get("thinking", ""), + } ) - elif part_type == "image": + elif part_type == MelodyContentType.IMAGE: blocks.append( { - "type": "image", + "type": MelodyContentType.IMAGE, "image": { "template_placeholder": part.get( "template_placeholder", "" @@ -186,24 +207,31 @@ def _content_blocks(content: Any) -> list[dict[str, Any]]: }, } ) - elif part_type == "document": + elif part_type == MelodyContentType.DOCUMENT: doc = part.get("document") if isinstance(doc, dict): - blocks.append({"type": "document", "document": doc}) + blocks.append( + {"type": MelodyContentType.DOCUMENT, "document": doc} + ) else: # Fall back to wrapping arbitrary string as text. - blocks.append({"type": "text", "text": json.dumps(doc)}) + blocks.append( + {"type": MelodyContentType.TEXT, "text": json.dumps(doc)} + ) elif part_type == "tool_reference": # Tool references are rendered by name; emit as text so the # renderer downstream is content-format agnostic. blocks.append( - {"type": "text", "text": part.get("name") or part.get("text", "")} + { + "type": MelodyContentType.TEXT, + "text": part.get("name") or part.get("text", ""), + } ) else: # Unknown block type: render as text fallback. text = part.get("text") or part.get(part_type) or "" text_str = text if isinstance(text, str) else json.dumps(text) - blocks.append({"type": "text", "text": text_str}) + blocks.append({"type": MelodyContentType.TEXT, "text": text_str}) return blocks @@ -253,10 +281,21 @@ def _conversation_to_melody_messages( content_blocks = _content_blocks(msg.get("content")) # Treat reasoning content as a thinking block on assistant turns - # so multi-turn thinking is preserved in the rendered prompt. + # so multi-turn reasoning is preserved in the rendered prompt. + # + # Cohere models output thought as either a ``thinking`` + # block (reasoning models) or a ``tool_plan`` field (older non- + # reasoning Command models), but vLLM's ConversationMessage has a + # single unified ``reasoning`` field that drops that difference. + # The cmd3 / cmd4 jinja templates render ``thinking`` blocks in both + # the case of tool calls and regular thinking blocks, so this is fine + # from the renderer's perspective. reasoning = msg.get("reasoning") or msg.get("reasoning_content") if role == "chatbot" and reasoning: - content_blocks.insert(0, {"type": "thinking", "thinking": reasoning}) + content_blocks.insert( + 0, + {"type": MelodyContentType.THINKING, "thinking": reasoning}, + ) tool_calls = [ _normalize_tool_call(tc) for tc in (msg.get("tool_calls") or []) From 42e52f284ba15c0d77982df9da2e6f79bd6879fe Mon Sep 17 00:00:00 2001 From: Andrew Berneshawi Date: Wed, 24 Jun 2026 02:57:25 -0400 Subject: [PATCH 19/30] changes for tool_plan handling --- vllm/entrypoints/cohere/serving.py | 47 +++++++++++++------------ vllm/entrypoints/generate/api_router.py | 1 + vllm/entrypoints/openai/cli_args.py | 9 +++++ 3 files changed, 34 insertions(+), 23 deletions(-) diff --git a/vllm/entrypoints/cohere/serving.py b/vllm/entrypoints/cohere/serving.py index c890c989a1b0..e0577676fa25 100644 --- a/vllm/entrypoints/cohere/serving.py +++ b/vllm/entrypoints/cohere/serving.py @@ -181,7 +181,7 @@ def __init__( enable_prompt_tokens_details: bool = False, enable_force_include_usage: bool = False, default_chat_template_kwargs: dict[str, Any] | None = None, - surface_reasoning_as_tool_plan: bool = False, + is_reasoning_model: bool = True, ) -> None: super().__init__( engine_client=engine_client, @@ -199,19 +199,24 @@ def __init__( enable_force_include_usage=enable_force_include_usage, default_chat_template_kwargs=default_chat_template_kwargs, ) - # When True, the assistant's reasoning text is surfaced as Cohere's - # ``tool_plan`` (non-streaming) or as ``tool-plan-delta`` events - # (streaming) whenever the model emits tool calls. This matches the - # behavior of older non-reasoning Command models, which produce a - # tool plan in place of a thinking block. + # Controls how the assistant's chain-of-thought is surfaced on + # turns that also contain tool calls. # - # When False (default), reasoning is always surfaced as a thinking - # content block; this matches reasoning Command models that - # emit both a thinking block and tool calls. + # - ``True`` (default): the model is a reasoning Command-family + # model; reasoning is always surfaced as a ``thinking`` content + # block (non-streaming) or as ``content-start`` / ``content- + # delta`` events for a thinking block (streaming), regardless of + # whether tool calls also appear. + # - ``False``: the model is an older non-reasoning Command model + # that uses Cohere's ``tool_plan`` field for its chain-of- + # thought before tool calls; reasoning is surfaced as + # ``tool_plan`` (non-streaming) or as ``tool-plan-delta`` events + # (streaming) on tool-call turns, and the thinking content block + # is dropped. # - # TODO: replace this manual flag with automatic detection based on - # the model's capabilities once we have that config - self._surface_reasoning_as_tool_plan = surface_reasoning_as_tool_plan + # TODO: replace this manual flag with automatic detection from the + # model's capabilities config once that exists. + self._is_reasoning_model = is_reasoning_model # ------------------------------------------------------------------ # Public entry point @@ -619,17 +624,13 @@ def _chat_completion_to_v2( ] # Cohere's ``tool_plan`` is the planning text emitted before tool - # calls on older, non-reasoning Command models. For those models we - # surface ``reasoning`` as ``tool_plan`` and drop the thinking - # block. Reasoning Command models emit a regular thinking - # block alongside tool calls, so the default leaves the thinking - # block in place and never sets ``tool_plan``. + # calls on older, non-reasoning Command models. For those models + # we surface ``reasoning`` as ``tool_plan`` and drop the thinking + # block. Reasoning Command models emit a regular thinking block + # alongside tool calls, so for them we leave the thinking block + # in place and never set ``tool_plan``. tool_plan: str | None = None - if ( - self._surface_reasoning_as_tool_plan - and tool_calls - and msg.reasoning - ): + if not self._is_reasoning_model and tool_calls and msg.reasoning: tool_plan = msg.reasoning content_blocks = [ blk @@ -843,7 +844,7 @@ def _handle_thinking_delta( # Non-reasoning Command models: emit ``tool-plan-delta`` events # directly instead of opening a thinking content block. The # ``tool-plan-delta`` event has no start/end pair around it. - if self._surface_reasoning_as_tool_plan: + if not self._is_reasoning_model: events: list[str] = list(self._close_open_blocks(state)) events.append( _emit( diff --git a/vllm/entrypoints/generate/api_router.py b/vllm/entrypoints/generate/api_router.py index 2c9d696a5153..b8ef4c52299d 100644 --- a/vllm/entrypoints/generate/api_router.py +++ b/vllm/entrypoints/generate/api_router.py @@ -207,6 +207,7 @@ async def init_generate_state( enable_prompt_tokens_details=args.enable_prompt_tokens_details, enable_force_include_usage=args.enable_force_include_usage, default_chat_template_kwargs=args.default_chat_template_kwargs, + is_reasoning_model=args.cohere_is_reasoning_model, ) if CohereServingChatV2 is not None and "generate" in supported_tasks else None diff --git a/vllm/entrypoints/openai/cli_args.py b/vllm/entrypoints/openai/cli_args.py index d130e83422a9..95e20f16fcc1 100644 --- a/vllm/entrypoints/openai/cli_args.py +++ b/vllm/entrypoints/openai/cli_args.py @@ -149,6 +149,15 @@ class BaseFrontendArgs: """If set to False, output deltas will not be logged. Relevant only if --enable-log-outputs is set. """ + cohere_is_reasoning_model: bool = True + """Cohere ``/cohere/v2/chat`` only. Whether the served model is a + reasoning Command-family model. When True (default), the assistant's + chain-of-thought is surfaced as a ``thinking`` content block (or + ``content-*`` events on the stream). When False, reasoning is + surfaced as Cohere's ``tool_plan`` field (or ``tool-plan-delta`` + events) whenever the model emits tool calls, matching older non- + reasoning Command models. Has no effect on the non-Cohere + endpoints.""" log_error_stack: bool = envs.VLLM_SERVER_DEV_MODE """If set to True, log the stack trace of error responses""" tokens_only: bool = False From 1aea1545d9f17a2432a5bcbdd6fc3c8fe926e8b9 Mon Sep 17 00:00:00 2001 From: Andrew Berneshawi Date: Wed, 24 Jun 2026 05:16:59 -0400 Subject: [PATCH 20/30] add support note --- vllm/entrypoints/cohere/serving.py | 1 + 1 file changed, 1 insertion(+) diff --git a/vllm/entrypoints/cohere/serving.py b/vllm/entrypoints/cohere/serving.py index e0577676fa25..d37c221822ac 100644 --- a/vllm/entrypoints/cohere/serving.py +++ b/vllm/entrypoints/cohere/serving.py @@ -526,6 +526,7 @@ def _apply_tool_choice( chat_req: ChatCompletionRequest, request: CohereChatV2Request, ) -> None: + # TODO need to add support for this if request.tool_choice == "REQUIRED": chat_req.tool_choice = "required" elif request.tool_choice == "NONE": From 54cfaa8feb9c3eedb64aa6ae5c9af458715fb04b Mon Sep 17 00:00:00 2001 From: Andrew Berneshawi Date: Wed, 24 Jun 2026 05:57:43 -0400 Subject: [PATCH 21/30] log level fix --- vllm/entrypoints/cohere/serving.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/vllm/entrypoints/cohere/serving.py b/vllm/entrypoints/cohere/serving.py index d37c221822ac..4e9d311f1a52 100644 --- a/vllm/entrypoints/cohere/serving.py +++ b/vllm/entrypoints/cohere/serving.py @@ -683,7 +683,9 @@ def _extract_citations_if_any(msg: Any) -> list[Citation] | None: if isinstance(payload, dict): out.append(Citation.model_validate(payload)) except Exception: # pragma: no cover - defensive - logger.debug("Skipping malformed citation: %r", c, exc_info=True) + logger.warning( + "Skipping malformed citation: %r", c, exc_info=True + ) return out or None @staticmethod @@ -1003,7 +1005,11 @@ def _handle_citation_deltas( try: citation = Citation.model_validate(payload) except Exception: # pragma: no cover - defensive - logger.debug("Skipping malformed streamed citation: %r", payload) + logger.warning( + "Skipping malformed streamed citation: %r", + payload, + exc_info=True, + ) continue idx = state.next_citation_index() events.append( From 2e8a9c6bd5b4e350cde1405cd992f52c1340eace Mon Sep 17 00:00:00 2001 From: Andrew Berneshawi Date: Thu, 25 Jun 2026 03:00:13 -0400 Subject: [PATCH 22/30] fixes for package structure updates --- vllm/entrypoints/cohere/api_router.py | 12 ++-- vllm/entrypoints/cohere/serving.py | 84 ++++++++------------------- 2 files changed, 31 insertions(+), 65 deletions(-) diff --git a/vllm/entrypoints/cohere/api_router.py b/vllm/entrypoints/cohere/api_router.py index 3cc019a952ee..4ff2ade8fe8c 100644 --- a/vllm/entrypoints/cohere/api_router.py +++ b/vllm/entrypoints/cohere/api_router.py @@ -10,7 +10,6 @@ logs a one-line notice and returns without registering the route, so vLLM continues to boot normally. """ -from __future__ import annotations from fastapi import FastAPI @@ -49,8 +48,11 @@ def attach_router(app: FastAPI) -> None: ) from vllm.entrypoints.cohere.serving import CohereServingChatV2 from vllm.entrypoints.openai.engine.protocol import ErrorResponse - from vllm.entrypoints.openai.utils import validate_json_request - from vllm.entrypoints.utils import load_aware_call, with_cancellation + from vllm.entrypoints.serve.utils.api_utils import ( + load_aware_call, + validate_json_request, + with_cancellation, + ) router = APIRouter() @@ -67,9 +69,7 @@ def _error_response( status = info.code or fallback_status return JSONResponse( status_code=status, - content=CohereError(message=info.message).model_dump( - exclude_none=True - ), + content=CohereError(message=info.message).model_dump(exclude_none=True), ) @router.post( diff --git a/vllm/entrypoints/cohere/serving.py b/vllm/entrypoints/cohere/serving.py index 4e9d311f1a52..168b05fd1dfd 100644 --- a/vllm/entrypoints/cohere/serving.py +++ b/vllm/entrypoints/cohere/serving.py @@ -62,7 +62,6 @@ ToolCallStartEvent, ToolPlanDeltaEvent, ) -from vllm.entrypoints.logger import RequestLogger from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionRequest, ChatCompletionResponse, @@ -77,6 +76,7 @@ StreamOptions, ) from vllm.entrypoints.openai.models.serving import OpenAIServingModels +from vllm.entrypoints.serve.utils.request_logger import RequestLogger if TYPE_CHECKING: from vllm.entrypoints.serve.render.serving import OpenAIServingRender @@ -302,9 +302,7 @@ def _coerce_text_content(content: str | list[Any]) -> str: return "".join(parts) @classmethod - def _convert_user_message( - cls, msg: UserChatMessageV2 - ) -> dict[str, Any]: + def _convert_user_message(cls, msg: UserChatMessageV2) -> dict[str, Any]: if isinstance(msg.content, str): return {"role": "user", "content": msg.content} @@ -326,17 +324,12 @@ def _convert_user_message( "image_url": image_url, } ) - if ( - len(content_parts) == 1 - and content_parts[0]["type"] == "text" - ): + if len(content_parts) == 1 and content_parts[0]["type"] == "text": return {"role": "user", "content": content_parts[0]["text"]} return {"role": "user", "content": content_parts} @classmethod - def _convert_assistant_message( - cls, msg: AssistantChatMessageV2 - ) -> dict[str, Any]: + def _convert_assistant_message(cls, msg: AssistantChatMessageV2) -> dict[str, Any]: out: dict[str, Any] = {"role": "assistant"} # Cohere splits reasoning out into ``thinking`` content blocks. We @@ -371,9 +364,7 @@ def _convert_assistant_message( "type": "function", "function": { "name": (tc.function.name if tc.function else "") or "", - "arguments": ( - tc.function.arguments if tc.function else None - ) + "arguments": (tc.function.arguments if tc.function else None) or "{}", }, } @@ -382,9 +373,7 @@ def _convert_assistant_message( return out @classmethod - def _convert_tool_message( - cls, msg: ToolChatMessageV2 - ) -> dict[str, Any]: + def _convert_tool_message(cls, msg: ToolChatMessageV2) -> dict[str, Any]: if isinstance(msg.content, str): return { "role": "tool", @@ -607,9 +596,7 @@ def _chat_completion_to_v2( {"type": ContentBlockType.THINKING, "thinking": msg.reasoning} ) if msg.content: - content_blocks.append( - {"type": ContentBlockType.TEXT, "text": msg.content} - ) + content_blocks.append({"type": ContentBlockType.TEXT, "text": msg.content}) tool_calls: list[ToolCallV2] | None = None if msg.tool_calls: @@ -683,9 +670,7 @@ def _extract_citations_if_any(msg: Any) -> list[Citation] | None: if isinstance(payload, dict): out.append(Citation.model_validate(payload)) except Exception: # pragma: no cover - defensive - logger.warning( - "Skipping malformed citation: %r", c, exc_info=True - ) + logger.warning("Skipping malformed citation: %r", c, exc_info=True) return out or None @staticmethod @@ -735,7 +720,7 @@ async def _chat_completion_stream_to_v2( async for item in generator: if not item.startswith("data:"): continue - data_str = item[len("data:"):].strip().rstrip("\n") + data_str = item[len("data:") :].strip().rstrip("\n") if not data_str: continue if data_str == "[DONE]": @@ -787,17 +772,13 @@ async def _chat_completion_stream_to_v2( yield ev if delta.tool_calls: - for ev in self._handle_tool_call_deltas( - state, delta.tool_calls - ): + for ev in self._handle_tool_call_deltas(state, delta.tool_calls): yield ev # Citations: a Cohere-specific extension on DeltaMessage that # the cohere renderer/parsers may populate. if getattr(delta, "citations", None): - for ev in self._handle_citation_deltas( - state, delta.citations - ): + for ev in self._handle_citation_deltas(state, delta.citations): yield ev except Exception as exc: @@ -841,9 +822,7 @@ async def _chat_completion_stream_to_v2( # -- per-delta helpers -------------------------------------------- - def _handle_thinking_delta( - self, state: _StreamState, delta_text: str - ) -> list[str]: + def _handle_thinking_delta(self, state: _StreamState, delta_text: str) -> list[str]: # Non-reasoning Command models: emit ``tool-plan-delta`` events # directly instead of opening a thinking content block. The # ``tool-plan-delta`` event has no start/end pair around it. @@ -884,17 +863,13 @@ def _handle_thinking_delta( _emit( ContentDeltaEvent( index=state.active_block_index, - delta={ - "message": {"content": {"thinking": delta_text}} - }, + delta={"message": {"content": {"thinking": delta_text}}}, ) ) ) return events - def _handle_text_delta( - self, state: _StreamState, delta_text: str - ) -> list[str]: + def _handle_text_delta(self, state: _StreamState, delta_text: str) -> list[str]: events: list[str] = [] if state.active_block != ContentBlockType.TEXT: events.extend(self._close_open_blocks(state)) @@ -926,9 +901,7 @@ def _handle_text_delta( ) return events - def _handle_tool_call_deltas( - self, state: _StreamState, deltas: list - ) -> list[str]: + def _handle_tool_call_deltas(self, state: _StreamState, deltas: list) -> list[str]: events: list[str] = [] for tc in deltas: tc_index = tc.index @@ -950,11 +923,8 @@ def _handle_tool_call_deltas( "id": tc.id or "", "type": "function", "function": { - "name": (fn.name if fn else "") - or "", - "arguments": ( - fn.arguments if fn else None - ) + "name": (fn.name if fn else "") or "", + "arguments": (fn.arguments if fn else None) or "", }, } @@ -997,9 +967,7 @@ def _handle_citation_deltas( """ events: list[str] = [] for c in citations: - payload = ( - c.model_dump(exclude_none=True) if hasattr(c, "model_dump") else c - ) + payload = c.model_dump(exclude_none=True) if hasattr(c, "model_dump") else c if not isinstance(payload, dict): continue try: @@ -1016,9 +984,11 @@ def _handle_citation_deltas( _emit( CitationStartEvent( index=idx, - delta={"message": {"citations": citation.model_dump( - exclude_none=True - )}}, + delta={ + "message": { + "citations": citation.model_dump(exclude_none=True) + } + }, ) ) ) @@ -1033,13 +1003,9 @@ def _close_open_blocks(self, state: _StreamState) -> list[str]: """ events: list[str] = [] if state.active_block in (ContentBlockType.TEXT, ContentBlockType.THINKING): - events.append( - _emit(ContentEndEvent(index=state.active_block_index)) - ) + events.append(_emit(ContentEndEvent(index=state.active_block_index))) elif state.active_block == ContentBlockType.TOOL_CALL: - events.append( - _emit(ToolCallEndEvent(index=state.active_tool_index)) - ) + events.append(_emit(ToolCallEndEvent(index=state.active_tool_index))) state.active_block = None state.active_block_index = None state.active_tool_index = None From dfd0778646e44523010575d574c7e7e27e1f186d Mon Sep 17 00:00:00 2001 From: Andrew Berneshawi Date: Thu, 25 Jun 2026 03:03:38 -0400 Subject: [PATCH 23/30] improvements to role mapping, remove continuous usage setting --- vllm/entrypoints/cohere/serving.py | 2 +- vllm/renderers/cohere.py | 30 ++++++++++++++++++++++++++++-- 2 files changed, 29 insertions(+), 3 deletions(-) diff --git a/vllm/entrypoints/cohere/serving.py b/vllm/entrypoints/cohere/serving.py index 168b05fd1dfd..b2cf038eb6c3 100644 --- a/vllm/entrypoints/cohere/serving.py +++ b/vllm/entrypoints/cohere/serving.py @@ -455,7 +455,7 @@ def _apply_streaming_options( if request.stream: chat_req.stream = True chat_req.stream_options = StreamOptions.model_validate( - {"include_usage": True, "continuous_usage_stats": True} + {"include_usage": True} ) @classmethod diff --git a/vllm/renderers/cohere.py b/vllm/renderers/cohere.py index 4db773aad5c2..fd6c27de04ad 100644 --- a/vllm/renderers/cohere.py +++ b/vllm/renderers/cohere.py @@ -129,11 +129,37 @@ def _try_import_melody(): ) from e +_MELODY_ROLES = frozenset({"system", "user", "chatbot", "tool"}) + + def _role_to_melody(role: str) -> str: - """Map an OpenAI role to the role string melody expects.""" + """Map an OpenAI role to the role string melody expects. + + The cmd3 / cmd4 jinja templates only recognize ``system``, ``user``, + ``assistant``/``chatbot``, and ``tool``: any other role is silently + dropped by the template's role-dispatch chain (no fallback branch), + which produces a malformed prompt without any error. We therefore + refuse unknown roles up front rather than letting them disappear. + + Aliases: + + * ``assistant`` -> ``chatbot`` (Cohere's historical assistant role + name used by the templates). + * ``developer`` -> ``system`` (OpenAI's ``developer`` role is + documented as high-priority instructions, which maps onto the + ``system`` slot in Cohere's prompt format). + """ + role = role.lower() if role == "assistant": return "chatbot" - return role + if role == "developer": + return "system" + if role in _MELODY_ROLES: + return role + raise ValueError( + f"Unsupported message role for the cohere renderer: {role!r}. " + "Expected one of: system, developer, user, assistant, chatbot, tool." + ) def _normalize_tool_call(tc: dict[str, Any] | Any) -> dict[str, Any]: From 8f67de2a29907a499be5d80795b8dcb6fef79f6c Mon Sep 17 00:00:00 2001 From: Andrew Berneshawi Date: Tue, 30 Jun 2026 06:26:36 -0400 Subject: [PATCH 24/30] fix renderer import --- vllm/entrypoints/cohere/serving.py | 6 +++--- vllm/entrypoints/generate/api_router.py | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/vllm/entrypoints/cohere/serving.py b/vllm/entrypoints/cohere/serving.py index b2cf038eb6c3..399610436eef 100644 --- a/vllm/entrypoints/cohere/serving.py +++ b/vllm/entrypoints/cohere/serving.py @@ -79,7 +79,7 @@ from vllm.entrypoints.serve.utils.request_logger import RequestLogger if TYPE_CHECKING: - from vllm.entrypoints.serve.render.serving import OpenAIServingRender + from vllm.renderers.online_renderer import OnlineRenderer logger = logging.getLogger(__name__) @@ -170,7 +170,7 @@ def __init__( models: OpenAIServingModels, response_role: str, *, - openai_serving_render: OpenAIServingRender, + online_renderer: "OnlineRenderer", request_logger: RequestLogger | None, chat_template: str | None, chat_template_content_format: ChatTemplateContentFormatOption, @@ -187,7 +187,7 @@ def __init__( engine_client=engine_client, models=models, response_role=response_role, - openai_serving_render=openai_serving_render, + online_renderer=online_renderer, request_logger=request_logger, chat_template=chat_template, chat_template_content_format=chat_template_content_format, diff --git a/vllm/entrypoints/generate/api_router.py b/vllm/entrypoints/generate/api_router.py index 09f64a15cdf6..7c592c1a1188 100644 --- a/vllm/entrypoints/generate/api_router.py +++ b/vllm/entrypoints/generate/api_router.py @@ -196,7 +196,7 @@ async def init_generate_state( engine_client, state.openai_serving_models, args.response_role, - openai_serving_render=state.openai_serving_render, + online_renderer=state.online_renderer, request_logger=request_logger, chat_template=resolved_chat_template, chat_template_content_format=args.chat_template_content_format, From 47d3a303fcfd6fffc3c4ed01bfe507672d30640d Mon Sep 17 00:00:00 2001 From: Andrew Berneshawi Date: Tue, 30 Jun 2026 07:08:43 -0400 Subject: [PATCH 25/30] Fix developer validation check --- vllm/entrypoints/cohere/protocol.py | 42 +++++++++++++++++++++++++++++ 1 file changed, 42 insertions(+) diff --git a/vllm/entrypoints/cohere/protocol.py b/vllm/entrypoints/cohere/protocol.py index 2248dd248c51..7409dc1709db 100644 --- a/vllm/entrypoints/cohere/protocol.py +++ b/vllm/entrypoints/cohere/protocol.py @@ -195,6 +195,48 @@ def _validate_max_tokens(cls, v: int | None) -> int | None: raise ValueError("max_tokens must be non-negative") return v + @field_validator("messages", mode="before") + @classmethod + def _normalize_message_roles(cls, v: Any) -> Any: + """Rewrite OpenAI-style ``developer`` roles to ``system``. + + Cohere's v2 ``ChatMessageV2`` discriminated union only admits + the four literal roles ``user`` / ``assistant`` / ``system`` / + ``tool``. OpenAI's ``developer`` role is documented as + high-priority system instructions, so we alias it onto + ``system`` *before* the SDK discriminator runs (otherwise it + rejects the message with a ``literal_error`` against each union + member). Mirrors ``_role_to_melody`` in the renderer so the + same alias is honoured no matter which surface the message + arrives through. + + ``mode="before"`` is required so the rewrite happens before the + ``list[ChatMessageV2]`` coercion runs the SDK's discriminated + union; a default-mode validator would never see ``developer`` + because validation would have already failed. On any structural + surprise (non-iterable input, items without a dict-shaped + ``role`` field, etc.) we hand ``v`` back unchanged and let + Pydantic's normal coercion surface a precise error. + """ + try: + return [ + {**msg, "role": "system"} + if msg.get("role", "").lower() == "developer" + else msg + for msg in v + ] + except (AttributeError, TypeError): + return v + + @field_validator("messages") + @classmethod + def _validate_messages( + cls, v: list[ChatMessageV2] + ) -> list[ChatMessageV2]: + if not v: + raise ValueError("messages must contain at least one message") + return v + # --------------------------------------------------------------------------- # Usage / Logprobs From dbe5b05880b4464ae6d665f38b90fa4e2f0c5cc4 Mon Sep 17 00:00:00 2001 From: Andrew Berneshawi Date: Tue, 30 Jun 2026 07:57:19 -0400 Subject: [PATCH 26/30] add request id to router --- vllm/entrypoints/cohere/api_router.py | 80 +++++++++++++++++++-------- 1 file changed, 56 insertions(+), 24 deletions(-) diff --git a/vllm/entrypoints/cohere/api_router.py b/vllm/entrypoints/cohere/api_router.py index 4ff2ade8fe8c..f0b0ed586e9e 100644 --- a/vllm/entrypoints/cohere/api_router.py +++ b/vllm/entrypoints/cohere/api_router.py @@ -4,11 +4,18 @@ The Cohere v2 protocol models are sourced from the official ``cohere`` Python SDK (``pip install cohere``). To keep that an *optional* -dependency for vLLM, we don't import the SDK or any module that pulls -it in at top-level - everything that touches the SDK is imported lazily -inside :func:`attach_router`. If the SDK isn't installed, attach_router -logs a one-line notice and returns without registering the route, so +dependency for vLLM, the SDK-dependent imports - and the route handler +itself - are gated on a one-shot probe at module load. If the SDK isn't +installed, :func:`attach_router` becomes a no-op (with an info log) and vLLM continues to boot normally. + +Note: the handler 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 it 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. """ from fastapi import FastAPI @@ -18,24 +25,15 @@ logger = init_logger(__name__) -def attach_router(app: FastAPI) -> None: - """Register ``POST /cohere/v2/chat`` on ``app``. +try: + import cohere # noqa: F401 -- dependency probe +except ImportError: + _SDK_AVAILABLE = False +else: + _SDK_AVAILABLE = True - No-op (with an info log) when the optional ``cohere`` SDK isn't - installed, since the v2 protocol models live there. - """ - try: - import cohere # noqa: F401 -- dependency probe - except ImportError: - logger.info( - "cohere SDK not installed; /cohere/v2/chat endpoint disabled. " - "Install with `pip install cohere` to enable it." - ) - return - # Defer the SDK-dependent imports until we know cohere is available. - # This keeps `import vllm.entrypoints.cohere.api_router` safe when the - # SDK isn't installed. +if _SDK_AVAILABLE: from http import HTTPStatus from fastapi import APIRouter, Depends, Request @@ -59,8 +57,25 @@ def attach_router(app: FastAPI) -> None: def _serving(request: Request) -> CohereServingChatV2 | None: return getattr(request.app.state, "cohere_serving_chat_v2", None) + 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") + def _error_response( error: ErrorResponse, + raw_request: Request | None, *, fallback_status: int = HTTPStatus.BAD_REQUEST, ) -> JSONResponse: @@ -69,7 +84,7 @@ def _error_response( status = info.code or fallback_status return JSONResponse( status_code=status, - content=CohereError(message=info.message).model_dump(exclude_none=True), + content=CohereError(message=info.message, id=_request_id(raw_request)).model_dump(exclude_none=True), ) @router.post( @@ -90,7 +105,8 @@ async def chat_v2(request: CohereChatV2Request, raw_request: Request): return JSONResponse( status_code=HTTPStatus.NOT_IMPLEMENTED.value, content=CohereError( - message="The model does not support the Cohere v2 chat API." + message="The model does not support the Cohere v2 chat API.", + id=_request_id(raw_request), ).model_dump(exclude_none=True), ) @@ -100,15 +116,31 @@ async def chat_v2(request: CohereChatV2Request, raw_request: Request): logger.exception("Error in /cohere/v2/chat: %s", e) return JSONResponse( status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, - content=CohereError(message=str(e)).model_dump(exclude_none=True), + content=CohereError( + message=str(e), + id=_request_id(raw_request), + ).model_dump(exclude_none=True), ) if isinstance(result, ErrorResponse): - return _error_response(result) + return _error_response(result, raw_request) if isinstance(result, CohereChatV2Response): return JSONResponse(content=result.model_dump(exclude_none=True)) return StreamingResponse(content=result, media_type="text/event-stream") + +def attach_router(app: FastAPI) -> None: + """Register ``POST /cohere/v2/chat`` on ``app``. + + No-op (with an info log) when the optional ``cohere`` SDK isn't + installed, since the v2 protocol models live there. + """ + if not _SDK_AVAILABLE: + logger.info( + "cohere SDK not installed; /cohere/v2/chat endpoint disabled. " + "Install with `pip install cohere` to enable it." + ) + return app.include_router(router) From 4247e8dabc251c863af744bd7aab4631b868d65f Mon Sep 17 00:00:00 2001 From: Andrew Berneshawi Date: Tue, 30 Jun 2026 08:29:33 -0400 Subject: [PATCH 27/30] add cmd4 cli arg and grounding handling --- vllm/entrypoints/cohere/serving.py | 14 +++++--- vllm/entrypoints/generate/api_router.py | 21 +++++++++--- vllm/entrypoints/openai/cli_args.py | 10 ++++++ vllm/renderers/cohere.py | 45 +++++++++++++++++++++++-- 4 files changed, 79 insertions(+), 11 deletions(-) diff --git a/vllm/entrypoints/cohere/serving.py b/vllm/entrypoints/cohere/serving.py index 399610436eef..42df5c699334 100644 --- a/vllm/entrypoints/cohere/serving.py +++ b/vllm/entrypoints/cohere/serving.py @@ -569,8 +569,14 @@ def _apply_cohere_template_kwargs( "thinking", request.thinking.model_dump(exclude_none=True), ) - if request.strict_tools is not None: - kwargs.setdefault("strict_tools", bool(request.strict_tools)) + + # ``strict_tools`` is intentionally NOT forwarded here. It's a + # decoder-guidance flag that ``_apply_tools`` already maps onto + # per-function OpenAI ``strict: true`` on ``chat_req.tools``; + # ``cohere_melody.render_cmd3``/``render_cmd4`` have no + # ``strict_tools`` knob, and the renderer's residual-forward path + # would otherwise surface it as a Jinja variable ``{{ strict_tools }}`` + # that templates have no defined use for. if kwargs: chat_req.chat_template_kwargs = kwargs @@ -649,8 +655,8 @@ def _extract_citations_if_any(msg: Any) -> list[Citation] | None: ``ChatMessage`` natively carries a ``citations: list[Citation] | None`` field (see :mod:`vllm.entrypoints.openai.engine.protocol`). - Renderers/parsers (e.g. the ``cohere2`` reasoning parser) populate - it. We map each :class:`vllm...Citation` into the SDK + Renderers / parsers (the ``cohere_command3`` and + ``cohere_command4`` reasoning parsers) populate it. We map each :class:`vllm...Citation` into the SDK :class:`cohere.types.Citation` wire model, preserving sources and span. """ diff --git a/vllm/entrypoints/generate/api_router.py b/vllm/entrypoints/generate/api_router.py index 7c592c1a1188..a9c36f66089a 100644 --- a/vllm/entrypoints/generate/api_router.py +++ b/vllm/entrypoints/generate/api_router.py @@ -103,6 +103,19 @@ async def init_generate_state( tool_server = None resolved_chat_template = load_chat_template(args.chat_template) + # Fold the dedicated ``--cohere-format`` CLI flag into the renderer's + # default chat-template kwargs. The cohere renderer reads + # ``chat_template_kwargs["cohere_format"]`` to pick cmd3 vs cmd4 + # rendering; making this a first-class flag keeps the right format + # discoverable for ``vllm serve --tokenizer-mode cohere`` users + # without forcing them to hand-construct a JSON dict for + # ``--default-chat-template-kwargs``. Per-request overrides still + # take precedence (see ``merge_kwargs`` in + # ``ChatCompletionRequest.build_chat_params``). + default_chat_template_kwargs = dict(args.default_chat_template_kwargs or {}) + if getattr(args, "cohere_format", None): + default_chat_template_kwargs.setdefault("cohere_format", args.cohere_format) + # Render endpoints are always backed by OnlineRenderer so that # /v1/chat/completions/render and /v1/completions/render work on both # generate-mode and render-only servers. Created in init_app_state. @@ -123,7 +136,7 @@ async def init_generate_state( enable_prompt_tokens_details=args.enable_prompt_tokens_details, enable_force_include_usage=args.enable_force_include_usage, enable_log_outputs=args.enable_log_outputs, - default_chat_template_kwargs=args.default_chat_template_kwargs, + default_chat_template_kwargs=default_chat_template_kwargs, ) if "generate" in supported_tasks else None @@ -136,7 +149,7 @@ async def init_generate_state( request_logger=request_logger, chat_template=resolved_chat_template, chat_template_content_format=args.chat_template_content_format, - default_chat_template_kwargs=args.default_chat_template_kwargs, + default_chat_template_kwargs=default_chat_template_kwargs, trust_request_chat_template=args.trust_request_chat_template, return_tokens_as_token_ids=args.return_tokens_as_token_ids, enable_auto_tools=args.enable_auto_tool_choice, @@ -186,7 +199,7 @@ async def init_generate_state( reasoning_parser=args.structured_outputs_config.reasoning_parser, enable_prompt_tokens_details=args.enable_prompt_tokens_details, enable_force_include_usage=args.enable_force_include_usage, - default_chat_template_kwargs=args.default_chat_template_kwargs, + default_chat_template_kwargs=default_chat_template_kwargs, ) if "generate" in supported_tasks else None @@ -206,7 +219,7 @@ async def init_generate_state( reasoning_parser=args.structured_outputs_config.reasoning_parser, enable_prompt_tokens_details=args.enable_prompt_tokens_details, enable_force_include_usage=args.enable_force_include_usage, - default_chat_template_kwargs=args.default_chat_template_kwargs, + default_chat_template_kwargs=default_chat_template_kwargs, is_reasoning_model=args.cohere_is_reasoning_model, ) if CohereServingChatV2 is not None and "generate" in supported_tasks diff --git a/vllm/entrypoints/openai/cli_args.py b/vllm/entrypoints/openai/cli_args.py index 64746cdc011f..b1274ebfb7ae 100644 --- a/vllm/entrypoints/openai/cli_args.py +++ b/vllm/entrypoints/openai/cli_args.py @@ -158,6 +158,16 @@ class BaseFrontendArgs: events) whenever the model emits tool calls, matching older non- reasoning Command models. Has no effect on the non-Cohere endpoints.""" + cohere_format: str = "cmd4" + """Cohere ``--tokenizer-mode cohere`` only. Which Cohere prompt + format to render: ``cmd4`` (current Command A / R+ 2025+ models; + default) or ``cmd3`` (legacy Command R / R+ 2024 models). Selecting + the wrong format silently produces a prompt the model wasn't trained + on, which most commonly manifests as the model emitting text but no + citations / tool calls / thinking blocks. Equivalent to passing + ``--default-chat-template-kwargs '{"cohere_format": "..."}'`` -- any + explicit request-level ``chat_template_kwargs.cohere_format`` still + wins.""" log_error_stack: bool = envs.VLLM_SERVER_DEV_MODE """If set to True, log the stack trace of error responses""" tokens_only: bool = False diff --git a/vllm/renderers/cohere.py b/vllm/renderers/cohere.py index fd6c27de04ad..c6fc5b29608c 100644 --- a/vllm/renderers/cohere.py +++ b/vllm/renderers/cohere.py @@ -132,6 +132,39 @@ def _try_import_melody(): _MELODY_ROLES = frozenset({"system", "user", "chatbot", "tool"}) +# Cohere v2 ``citation_options.mode`` -> melody cmd4 ``grounding``. +# v2 surfaces three modes (``FAST`` / ``ACCURATE`` / ``OFF``) but cmd4's +# template doesn't differentiate fast vs accurate at the prompt layer -- +# both just turn grounding on. ``unknown`` / ``enabled`` / ``disabled`` +# are the values melody actually accepts. +_CMD4_GROUNDING_FROM_MODE = { + "fast": "enabled", + "accurate": "enabled", + "on": "enabled", + "enabled": "enabled", + "off": "disabled", + "disabled": "disabled", + "unknown": "unknown", +} + + +def _normalize_cmd4_grounding(value: Any) -> str: + """Coerce a user-facing grounding/citation mode to melody's vocab. + + Raises ``ValueError`` rather than letting an unrecognized value slip + through to ``render_cmd4`` and surface as a generic + ``Invalid config: grounding`` from melody. + """ + out = _CMD4_GROUNDING_FROM_MODE.get(value.lower()) + if out is None: + raise ValueError( + f"Unrecognized cmd4 grounding value: {value!r}. Expected one of " + f"FAST, ACCURATE, OFF (citation_options.mode), or " + f"enabled / disabled / unknown." + ) + return out + + def _role_to_melody(role: str) -> str: """Map an OpenAI role to the role string melody expects. @@ -425,13 +458,19 @@ def _build_render_config( if "skip_preamble" in chat_template_kwargs: config["skip_preamble"] = bool(chat_template_kwargs["skip_preamble"]) else: # cmd4 - # cmd4 uses ``grounding`` rather than safety_mode/citation_quality + # cmd4 uses ``grounding`` rather than safety_mode/citation_quality. + # melody's cmd4 only accepts ``unknown`` / ``enabled`` / ``disabled``, + # so the Cohere v2 ``citation_options.mode`` values + # (``FAST`` / ``ACCURATE`` / ``OFF``) have to be normalized -- a + # raw lowercased passthrough would raise from + # ``render_cmd4`` (cmd4 has no fast/accurate distinction at the + # prompt-template layer; both request grounding-on). if (gr := chat_template_kwargs.get("grounding")) is not None: - config["grounding"] = str(gr).lower() + config["grounding"] = _normalize_cmd4_grounding(gr) elif (co := chat_template_kwargs.get("citation_options")): mode = co.get("mode") if isinstance(co, dict) else None if mode is not None: - config["grounding"] = str(mode).lower() + config["grounding"] = _normalize_cmd4_grounding(mode) if (pi := chat_template_kwargs.get("platform_instruction")) is not None: config["platform_instruction"] = str(pi) From 2eb890bdc253172ce6dc96a534161cd605c4db3a Mon Sep 17 00:00:00 2001 From: Andrew Berneshawi Date: Tue, 30 Jun 2026 09:18:12 -0400 Subject: [PATCH 28/30] Add unit tests --- tests/entrypoints/cohere/__init__.py | 0 tests/entrypoints/cohere/test_api_router.py | 200 ++++ tests/entrypoints/cohere/test_chat_v2.py | 238 ++++ tests/entrypoints/cohere/test_protocol.py | 350 ++++++ .../cohere/test_registry_and_args.py | 91 ++ .../cohere/test_serving_conversion.py | 1036 +++++++++++++++++ .../cohere/test_serving_streaming.py | 815 +++++++++++++ tests/renderers/test_cohere.py | 743 ++++++++++++ 8 files changed, 3473 insertions(+) create mode 100644 tests/entrypoints/cohere/__init__.py create mode 100644 tests/entrypoints/cohere/test_api_router.py create mode 100644 tests/entrypoints/cohere/test_chat_v2.py create mode 100644 tests/entrypoints/cohere/test_protocol.py create mode 100644 tests/entrypoints/cohere/test_registry_and_args.py create mode 100644 tests/entrypoints/cohere/test_serving_conversion.py create mode 100644 tests/entrypoints/cohere/test_serving_streaming.py create mode 100644 tests/renderers/test_cohere.py diff --git a/tests/entrypoints/cohere/__init__.py b/tests/entrypoints/cohere/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/entrypoints/cohere/test_api_router.py b/tests/entrypoints/cohere/test_api_router.py new file mode 100644 index 000000000000..4c88d627c7aa --- /dev/null +++ b/tests/entrypoints/cohere/test_api_router.py @@ -0,0 +1,200 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for ``vllm/entrypoints/cohere/api_router.py``. + +Covers: + +* The optional-import guard: ``attach_router`` is a no-op when the + ``cohere`` SDK isn't installed. +* The router wiring: response shapes (JSON + SSE), error translation, + and the ``cohere_serving_chat_v2 is None`` fallback (501 Not + Implemented). +""" + +import json +import sys +from collections.abc import AsyncGenerator +from http import HTTPStatus + +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from vllm.entrypoints.cohere.api_router import attach_router +from vllm.entrypoints.cohere.protocol import ( + AssistantMessageResponse, + CohereChatV2Response, +) +from vllm.entrypoints.openai.engine.protocol import ErrorInfo, ErrorResponse + +# ---------------------------------------------------------------------- +# Fakes +# ---------------------------------------------------------------------- + + +class _Handler: + """Minimal stand-in for :class:`CohereServingChatV2` used by the + router. Each test sets ``self.result`` to either: + + * a :class:`CohereChatV2Response` (non-streaming JSON path); + * an async generator yielding SSE frames (streaming path); + * an :class:`ErrorResponse` (error envelope path); or + * an exception (router-level 500 path). + """ + + def __init__(self, result): + self.result = result + + async def create_chat_v2(self, request, raw_request): + if isinstance(self.result, Exception): + raise self.result + return self.result + + +def _build_app(handler: _Handler | None) -> FastAPI: + app = FastAPI() + attach_router(app) + app.state.cohere_serving_chat_v2 = handler + return app + + +def _minimal_request_body() -> dict: + return { + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + } + + +# ---------------------------------------------------------------------- +# Optional-import guard +# ---------------------------------------------------------------------- + + +class TestOptionalCohereImport: + def test_attach_router_noop_when_cohere_missing(self, monkeypatch, caplog): + # Simulate ``import cohere`` failing by injecting a sentinel that + # raises ImportError on first attribute access. The simplest + # approach: stash the real module, delete it from sys.modules, + # and shadow it with an entry that raises on next import. + real = sys.modules.pop("cohere", None) + + # Block re-import. + class _RaisingFinder: + def find_spec(self, name, path=None, target=None): + if name == "cohere": + raise ImportError("simulated missing cohere SDK") + return None + + finder = _RaisingFinder() + sys.meta_path.insert(0, finder) + try: + with caplog.at_level("INFO", logger="vllm.entrypoints.cohere.api_router"): + app = FastAPI() + attach_router(app) + # No routes should have been registered. + paths = [r.path for r in app.routes] + assert "/cohere/v2/chat" not in paths + assert any( + "cohere SDK not installed" in rec.message for rec in caplog.records + ) + finally: + sys.meta_path.remove(finder) + if real is not None: + sys.modules["cohere"] = real + + def test_attach_router_registers_route_when_cohere_present(self): + app = _build_app(handler=None) + paths = [getattr(r, "path", None) for r in app.routes] + assert "/cohere/v2/chat" in paths + + +# ---------------------------------------------------------------------- +# Endpoint behavior +# ---------------------------------------------------------------------- + + +class TestEndpoint: + def test_501_when_handler_missing(self): + app = _build_app(handler=None) + with TestClient(app) as client: + r = client.post("/cohere/v2/chat", json=_minimal_request_body()) + assert r.status_code == HTTPStatus.NOT_IMPLEMENTED + body = r.json() + assert "does not support" in body["message"] + assert "id" not in body # excluded by ``exclude_none=True`` + + def test_non_streaming_response_is_json(self): + msg = AssistantMessageResponse(content=[{"type": "text", "text": "hello"}]) + result = CohereChatV2Response(id="r1", finish_reason="COMPLETE", message=msg) + app = _build_app(handler=_Handler(result)) + with TestClient(app) as client: + r = client.post("/cohere/v2/chat", json=_minimal_request_body()) + assert r.status_code == HTTPStatus.OK + assert r.headers["content-type"].startswith("application/json") + body = r.json() + assert body["id"] == "r1" + assert body["finish_reason"] == "COMPLETE" + assert body["message"]["content"][0]["text"] == "hello" + + def test_streaming_response_is_sse(self): + async def _gen() -> AsyncGenerator[str, None]: + yield 'data: {"type":"message-start"}\n\n' + yield "data: [DONE]\n\n" + + app = _build_app(handler=_Handler(_gen())) + with TestClient(app) as client: + r = client.post( + "/cohere/v2/chat", + json={**_minimal_request_body(), "stream": True}, + ) + assert r.status_code == HTTPStatus.OK + assert r.headers["content-type"].startswith("text/event-stream") + body = r.text + assert "message-start" in body + assert body.rstrip().endswith("[DONE]") + + def test_error_response_translated_to_cohere_envelope(self): + err = ErrorResponse( + error=ErrorInfo( + message="bad request", + type="bad_request", + code=400, + ) + ) + app = _build_app(handler=_Handler(err)) + with TestClient(app) as client: + r = client.post("/cohere/v2/chat", json=_minimal_request_body()) + assert r.status_code == HTTPStatus.BAD_REQUEST + body = r.json() + assert body == {"message": "bad request"} + + def test_handler_exception_returns_500_envelope(self): + app = _build_app(handler=_Handler(RuntimeError("kaboom"))) + with TestClient(app) as client: + r = client.post("/cohere/v2/chat", json=_minimal_request_body()) + assert r.status_code == HTTPStatus.INTERNAL_SERVER_ERROR + body = r.json() + assert body == {"message": "kaboom"} + + def test_non_json_content_type_rejected(self): + """The ``validate_json_request`` dependency raises + ``RequestValidationError`` (HTTP 422) for non-JSON content + types, matching the behavior of the other vLLM API routers. + """ + app = _build_app(handler=None) + with TestClient(app) as client: + r = client.post( + "/cohere/v2/chat", + content=json.dumps(_minimal_request_body()), + headers={"content-type": "text/plain"}, + ) + assert r.status_code == HTTPStatus.UNPROCESSABLE_ENTITY + + def test_invalid_body_returns_422(self): + # ``model`` is required; omit it to trip Pydantic validation. + app = _build_app(handler=None) + with TestClient(app) as client: + r = client.post( + "/cohere/v2/chat", + json={"messages": [{"role": "user", "content": "hi"}]}, + ) + assert r.status_code == HTTPStatus.UNPROCESSABLE_ENTITY diff --git a/tests/entrypoints/cohere/test_chat_v2.py b/tests/entrypoints/cohere/test_chat_v2.py new file mode 100644 index 000000000000..505a2f93d836 --- /dev/null +++ b/tests/entrypoints/cohere/test_chat_v2.py @@ -0,0 +1,238 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""End-to-end integration tests for the ``POST /cohere/v2/chat`` endpoint. + +These tests spin up a real ``vllm serve`` process via +:class:`tests.utils.RemoteOpenAIServer` and exercise the Cohere Chat v2 +contract over HTTP. They mirror the pattern used by +:mod:`tests.entrypoints.anthropic.test_messages`. + +Two layers of integration are covered: + +1. Raw HTTP via :mod:`httpx` — always runs, verifies the wire contract. +2. Cohere SDK (``pip install cohere``) — auto-skipped when the optional + dependency isn't installed, verifies SDK-level interop. + +Like ``test_messages.py`` we use a small generic chat model +(``Qwen/Qwen3-0.6B``); the Cohere v2 endpoint is model-agnostic and +just performs the v2 ↔ OpenAI chat-completion translation, so the +choice of model only matters for response shape (not for the +translation logic itself, which is unit-tested elsewhere). +""" + +import json + +import httpx +import pytest +import pytest_asyncio + +from tests.utils import RemoteOpenAIServer + +MODEL_NAME = "Qwen/Qwen3-0.6B" +SERVED_MODEL_NAME = "command-r-plus-08-2024" + + +@pytest.fixture(scope="module") +def server(): + args = [ + "--max-model-len", + "2048", + "--enforce-eager", + "--enable-auto-tool-choice", + "--tool-call-parser", + "hermes", + # Advertise a Cohere model name so Cohere SDK ``model=`` calls + # round-trip without the ``model not found`` check. + "--served-model-name", + SERVED_MODEL_NAME, + # Disable the reasoning-model path so the test doesn't require + # the conversation to surface a thinking block (Qwen3-0.6B + # doesn't emit Cohere-style reasoning tokens out of the box). + "--no-cohere-is-reasoning-model", + ] + + with RemoteOpenAIServer(MODEL_NAME, args) as remote_server: + yield remote_server + + +# ---------------------------------------------------------------------- +# Layer 1: raw HTTP contract (no optional deps) +# ---------------------------------------------------------------------- + + +@pytest_asyncio.fixture +async def httpx_client(server): + async with httpx.AsyncClient( + base_url=server.url_root, timeout=httpx.Timeout(120.0) + ) as client: + yield client + + +@pytest.mark.asyncio +async def test_cohere_v2_chat_non_streaming(httpx_client: httpx.AsyncClient): + resp = await httpx_client.post( + "/cohere/v2/chat", + json={ + "model": SERVED_MODEL_NAME, + "messages": [{"role": "user", "content": "Say hi."}], + "max_tokens": 16, + "stream": False, + }, + ) + assert resp.status_code == 200, resp.text + payload = resp.json() + # The response envelope follows Cohere v2's schema. + assert "message" in payload + assert payload["message"]["role"] == "assistant" + # ``content`` is a list of content blocks; the synthesized + # ``CohereServingChatV2`` should always emit at least one ``text`` + # block (it falls back to an empty block when the model returned + # nothing). + content = payload["message"]["content"] + assert isinstance(content, list) and len(content) >= 1 + assert content[0]["type"] == "text" + # ``usage`` is always populated by the translator. + assert "usage" in payload + assert "billed_units" in payload["usage"] + assert "tokens" in payload["usage"] + # ``finish_reason`` is one of Cohere's enum values. + assert payload["finish_reason"] in { + "COMPLETE", + "MAX_TOKENS", + "STOP_SEQUENCE", + "TOOL_CALL", + "ERROR", + } + + +@pytest.mark.asyncio +async def test_cohere_v2_chat_streaming(httpx_client: httpx.AsyncClient): + """Streaming returns SSE frames in the v2 message-lifecycle shape.""" + events: list[dict] = [] + async with httpx_client.stream( + "POST", + "/cohere/v2/chat", + json={ + "model": SERVED_MODEL_NAME, + "messages": [{"role": "user", "content": "Hi"}], + "max_tokens": 8, + "stream": True, + }, + ) as resp: + assert resp.status_code == 200, await resp.aread() + assert resp.headers["content-type"].startswith("text/event-stream") + async for line in resp.aiter_lines(): + if not line.startswith("data: "): + continue + data = line[len("data: ") :] + if data == "[DONE]": + events.append({"type": "_DONE_"}) + continue + events.append(json.loads(data)) + + types = [ev["type"] for ev in events] + # The lifecycle always starts with message-start and ends with [DONE] + # preceded by message-end. + assert types[0] == "message-start" + assert types[-1] == "_DONE_" + assert types[-2] == "message-end" + # message-start must carry the chunk/message id. + assert events[0].get("id") + + +@pytest.mark.asyncio +async def test_cohere_v2_chat_validation_error_returns_422( + httpx_client: httpx.AsyncClient, +): + # Missing required ``model`` field → FastAPI/Pydantic returns 422. + resp = await httpx_client.post( + "/cohere/v2/chat", + json={"messages": []}, + ) + assert resp.status_code == 422 + + +@pytest.mark.asyncio +async def test_cohere_v2_chat_documents_field_accepted( + httpx_client: httpx.AsyncClient, +): + """The v2 endpoint forwards ``documents`` into chat_template_kwargs. + + We only assert the request is accepted and produces a 200 response — + the renderer-level effect is covered by ``tests/renderers/test_cohere.py``. + """ + resp = await httpx_client.post( + "/cohere/v2/chat", + json={ + "model": SERVED_MODEL_NAME, + "messages": [{"role": "user", "content": "Summarize."}], + "documents": [{"id": "d1", "data": {"title": "T", "snippet": "S"}}], + "max_tokens": 16, + "stream": False, + }, + ) + assert resp.status_code == 200, resp.text + + +# ---------------------------------------------------------------------- +# Layer 2: Cohere SDK round-trip (auto-skipped if SDK absent) +# ---------------------------------------------------------------------- + + +@pytest_asyncio.fixture +async def cohere_async_client(server): + cohere = pytest.importorskip("cohere") + # The vLLM endpoint is mounted at ``/cohere/v2/chat`` while the + # cohere SDK targets ``${base_url}/v2/chat``; point base_url at the + # ``/cohere`` prefix so paths line up. + client = cohere.AsyncClientV2( + api_key="dummy", + base_url=server.url_for("cohere"), + ) + try: + yield client + finally: + # ``AsyncClientV2`` exposes a sync close; if a future version + # adds aclose we still close cleanly. + close = getattr(client, "aclose", None) or getattr(client, "close", None) + if close is not None: + result = close() + if hasattr(result, "__await__"): + await result + + +@pytest.mark.asyncio +async def test_cohere_sdk_non_streaming(cohere_async_client): + resp = await cohere_async_client.chat( + model=SERVED_MODEL_NAME, + messages=[{"role": "user", "content": "Say hi."}], + max_tokens=16, + ) + # SDK parses our JSON into typed objects. + assert resp.message.role == "assistant" + assert resp.message.content is not None + assert len(resp.message.content) >= 1 + assert resp.message.content[0].type == "text" + assert resp.finish_reason in { + "COMPLETE", + "MAX_TOKENS", + "STOP_SEQUENCE", + "TOOL_CALL", + "ERROR", + } + + +@pytest.mark.asyncio +async def test_cohere_sdk_streaming(cohere_async_client): + events: list[str] = [] + stream = cohere_async_client.chat_stream( + model=SERVED_MODEL_NAME, + messages=[{"role": "user", "content": "Hi"}], + max_tokens=8, + ) + async for ev in stream: + events.append(ev.type) + + assert events, "SDK stream yielded no events" + assert events[0] == "message-start" + assert events[-1] == "message-end" diff --git a/tests/entrypoints/cohere/test_protocol.py b/tests/entrypoints/cohere/test_protocol.py new file mode 100644 index 000000000000..c2e56297cd0d --- /dev/null +++ b/tests/entrypoints/cohere/test_protocol.py @@ -0,0 +1,350 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for ``vllm/entrypoints/cohere/protocol.py``. + +The module is mostly a thin wrapper around the official ``cohere`` SDK +types plus a few local additions: + +* :class:`CohereError` envelope. +* :class:`CohereChatV2Request` (model required, ``max_tokens`` non-negative). +* :class:`CohereChatV2Response` plus the usage / logprob helpers. +* The streaming event subclasses that bake a wire-format ``type`` + discriminator into ``model_dump()`` so SSE consumers can demux on it. +""" + +import pytest +from pydantic import ValidationError + +from vllm.entrypoints.cohere.protocol import ( + AssistantChatMessageV2, + AssistantMessageResponse, + CitationEndEvent, + CitationStartEvent, + CohereChatV2Request, + CohereChatV2Response, + CohereError, + CohereLogprobItem, + CohereUsage, + CohereUsageBilledUnits, + CohereUsageTokens, + ContentDeltaEvent, + ContentEndEvent, + ContentStartEvent, + MessageEndEvent, + MessageStartEvent, + SystemChatMessageV2, + ToolCallDeltaEvent, + ToolCallEndEvent, + ToolCallStartEvent, + ToolPlanDeltaEvent, + UserChatMessageV2, +) + +# ====================================================================== +# CohereError +# ====================================================================== + + +class TestCohereError: + def test_message_only(self): + err = CohereError(message="boom") + assert err.message == "boom" + assert err.id is None + + def test_with_id(self): + err = CohereError(message="boom", id="req_123") + assert err.id == "req_123" + + def test_model_dump_excludes_none(self): + err = CohereError(message="boom") + assert err.model_dump(exclude_none=True) == {"message": "boom"} + + +# ====================================================================== +# CohereChatV2Request +# ====================================================================== + + +class TestCohereChatV2Request: + def test_minimal_required_fields(self): + req = CohereChatV2Request( + model="m", messages=[{"role": "user", "content": "hi"}] + ) + assert req.model == "m" + assert req.stream is False + assert req.max_tokens is None + assert req.tools is None + assert req.documents is None + assert req.kv_transfer_params is None + assert req.chat_template_kwargs is None + + def test_empty_model_rejected(self): + with pytest.raises(ValidationError, match="model is required"): + CohereChatV2Request(model="", messages=[{"role": "user", "content": "hi"}]) + + def test_negative_max_tokens_rejected(self): + with pytest.raises(ValidationError, match="non-negative"): + CohereChatV2Request( + model="m", + messages=[{"role": "user", "content": "hi"}], + max_tokens=-1, + ) + + def test_zero_max_tokens_allowed(self): + # Zero is allowed (the docs allow 0 -> return prompt only). + req = CohereChatV2Request( + model="m", + messages=[{"role": "user", "content": "hi"}], + max_tokens=0, + ) + assert req.max_tokens == 0 + + def test_full_schema(self): + """Sanity-check that every field listed in the model accepts a value.""" + req = CohereChatV2Request( + model="m", + stream=True, + messages=[ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": "hello", + "tool_plan": "p", + "tool_calls": [ + { + "id": "c1", + "type": "function", + "function": {"name": "f", "arguments": "{}"}, + } + ], + }, + ], + tools=[ + { + "type": "function", + "function": { + "name": "f", + "description": "d", + "parameters": {}, + }, + } + ], + strict_tools=True, + tool_choice="REQUIRED", + documents=["plain", {"id": "d1", "data": {"text": "t"}}], + citation_options={"mode": "accurate"}, + response_format={"type": "json_object"}, + safety_mode="CONTEXTUAL", + max_tokens=128, + stop_sequences=[""], + temperature=0.5, + seed=42, + frequency_penalty=0.1, + presence_penalty=0.2, + k=50, + p=0.95, + logprobs=True, + thinking={"type": "enabled", "token_budget": 1024}, + priority=1, + kv_transfer_params={"x": 1}, + chat_template_kwargs={"y": 2}, + ) + # Discriminated-union messages resolve to the SDK variants. + assert isinstance(req.messages[0], SystemChatMessageV2) + assert isinstance(req.messages[1], UserChatMessageV2) + assert isinstance(req.messages[2], AssistantChatMessageV2) + assert req.thinking.type == "enabled" + assert req.thinking.token_budget == 1024 + assert req.citation_options.mode == "accurate" + assert req.response_format.type == "json_object" + assert req.safety_mode == "CONTEXTUAL" + assert req.tool_choice == "REQUIRED" + assert req.tools[0].function.name == "f" + # Documents accept str OR Document. + assert isinstance(req.documents[0], str) + assert req.documents[1].id == "d1" + + def test_invalid_tool_choice_rejected(self): + with pytest.raises(ValidationError): + CohereChatV2Request( + model="m", + messages=[{"role": "user", "content": "hi"}], + tool_choice="ANY", # not REQUIRED/NONE + ) + + def test_unknown_field_ignored(self): + # Pydantic by default ignores extra fields. Make sure that contract + # holds so clients sending forward-compatible kwargs don't 422. + req = CohereChatV2Request( + model="m", + messages=[{"role": "user", "content": "hi"}], + some_future_field="x", # type: ignore[arg-type] + ) + assert req.model == "m" + + +# ====================================================================== +# Usage / Logprobs +# ====================================================================== + + +class TestUsage: + def test_billed_units_optional(self): + u = CohereUsageBilledUnits() + assert u.input_tokens is None + assert u.output_tokens is None + + def test_tokens_serialization(self): + u = CohereUsageTokens(input_tokens=10, output_tokens=5) + assert u.model_dump() == {"input_tokens": 10.0, "output_tokens": 5.0} + + def test_usage_envelope(self): + u = CohereUsage( + billed_units=CohereUsageBilledUnits(input_tokens=10, output_tokens=5), + tokens=CohereUsageTokens(input_tokens=10, output_tokens=5), + cached_tokens=3, + ) + assert u.cached_tokens == 3 + assert u.billed_units.input_tokens == 10 + + def test_logprob_item(self): + lp = CohereLogprobItem(text="hi", token_ids=[1, 2], logprobs=[-0.1, -0.2]) + assert lp.token_ids == [1, 2] + assert lp.logprobs == [-0.1, -0.2] + + +# ====================================================================== +# CohereChatV2Response +# ====================================================================== + + +class TestCohereChatV2Response: + def test_minimal(self): + msg = AssistantMessageResponse(content=[{"type": "text", "text": "hi"}]) + resp = CohereChatV2Response( + id="r1", + finish_reason="COMPLETE", + message=msg, + ) + assert resp.id == "r1" + assert resp.usage is None + assert resp.kv_transfer_params is None + + def test_invalid_finish_reason_rejected(self): + msg = AssistantMessageResponse(content=[{"type": "text", "text": "hi"}]) + with pytest.raises(ValidationError): + CohereChatV2Response( + id="r1", + finish_reason="NOT_A_REASON", # type: ignore[arg-type] + message=msg, + ) + + +# ====================================================================== +# Streaming event ``type`` discriminator baked into model_dump() +# ====================================================================== + + +class TestStreamingEventTypeField: + """Each event subclass adds a ``type: Literal[...]`` field with a + default so ``model_dump()`` always emits the wire-format discriminator + (the parent SDK classes don't declare ``type`` as a Pydantic field). + """ + + @pytest.mark.parametrize( + "cls, expected_type, kwargs", + [ + ( + MessageStartEvent, + "message-start", + {"id": "a", "delta": {"message": {"role": "assistant"}}}, + ), + ( + ContentStartEvent, + "content-start", + { + "index": 0, + "delta": {"message": {"content": {"type": "text", "text": ""}}}, + }, + ), + ( + ContentDeltaEvent, + "content-delta", + { + "index": 0, + "delta": {"message": {"content": {"text": "hi"}}}, + }, + ), + (ContentEndEvent, "content-end", {"index": 0}), + ( + ToolPlanDeltaEvent, + "tool-plan-delta", + {"delta": {"message": {"tool_plan": "thinking"}}}, + ), + ( + ToolCallStartEvent, + "tool-call-start", + { + "index": 0, + "delta": { + "message": { + "tool_calls": { + "id": "c1", + "type": "function", + "function": {"name": "f", "arguments": ""}, + } + } + }, + }, + ), + ( + ToolCallDeltaEvent, + "tool-call-delta", + { + "index": 0, + "delta": { + "message": {"tool_calls": {"function": {"arguments": "{}"}}} + }, + }, + ), + (ToolCallEndEvent, "tool-call-end", {"index": 0}), + ( + CitationStartEvent, + "citation-start", + { + "index": 0, + "delta": { + "message": { + "citations": {"start": 0, "end": 5, "text": "hello"} + } + }, + }, + ), + (CitationEndEvent, "citation-end", {"index": 0}), + ( + MessageEndEvent, + "message-end", + {"id": "a", "delta": {"finish_reason": "COMPLETE"}}, + ), + ], + ) + def test_type_field_default(self, cls, expected_type, kwargs): + ev = cls(**kwargs) + # type field is auto-populated from the Literal default. + assert ev.type == expected_type + # The discriminator must be present in the serialized payload so + # clients reading the stream can demux on it. + dumped = ev.model_dump(exclude_none=True) + assert dumped["type"] == expected_type + # Same in JSON form (what ``_emit`` serializes). + assert f'"type":"{expected_type}"' in ev.model_dump_json(exclude_none=True) + + def test_type_field_cannot_be_overridden_to_wrong_value(self): + # Literal types reject any value other than the bake-in default. + with pytest.raises(ValidationError): + MessageStartEvent( + id="a", + delta={"message": {"role": "assistant"}}, + type="other", # type: ignore[arg-type] + ) diff --git a/tests/entrypoints/cohere/test_registry_and_args.py b/tests/entrypoints/cohere/test_registry_and_args.py new file mode 100644 index 000000000000..5e907b0d9340 --- /dev/null +++ b/tests/entrypoints/cohere/test_registry_and_args.py @@ -0,0 +1,91 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Sanity tests for the small registry / CLI-arg / config additions made +by the Cohere v2 chat API: + +* ``vllm/renderers/registry.py``: new ``"cohere"`` renderer entry. +* ``vllm/tokenizers/registry.py``: new ``"cohere"`` tokenizer entry + (aliased to the cached HF tokenizer). +* ``vllm/entrypoints/openai/cli_args.py``: new + ``cohere_is_reasoning_model`` field. +* ``vllm/config/model.py``: ``"cohere"`` added to ``TokenizerMode`` + Literal. +""" + +import dataclasses +import typing + +import pytest + +from vllm.config.model import TokenizerMode +from vllm.entrypoints.openai.cli_args import BaseFrontendArgs, make_arg_parser +from vllm.renderers.registry import RENDERER_REGISTRY +from vllm.tokenizers.registry import TokenizerRegistry +from vllm.utils.argparse_utils import FlexibleArgumentParser + + +class TestRendererRegistry: + def test_cohere_renderer_registered(self): + # The registry resolves to the importable ``CohereRenderer`` class. + cls = RENDERER_REGISTRY.load_renderer_cls("cohere") + assert cls.__name__ == "CohereRenderer" + # Sanity: class lives in the cohere renderer module. + assert cls.__module__ == "vllm.renderers.cohere" + + +class TestTokenizerRegistry: + def test_cohere_aliased_to_cached_hf_tokenizer(self): + # ``cohere`` mode uses the standard HF tokenizer; only the + # renderer stage is replaced. This test guards against accidental + # divergence. + cls = TokenizerRegistry.load_tokenizer_cls("cohere") + assert cls.__name__ == "CachedHfTokenizer" + + +class TestTokenizerModeLiteral: + def test_cohere_is_a_valid_tokenizer_mode(self): + # The Literal must enumerate ``"cohere"`` so engine arg parsing + # accepts ``--tokenizer-mode cohere``. + modes = typing.get_args(TokenizerMode) + assert "cohere" in modes + + +# ---------------------------------------------------------------------- +# ``--cohere-is-reasoning-model`` CLI flag +# ---------------------------------------------------------------------- + + +class TestCohereCliArg: + """Verifies the new ``--cohere-is-reasoning-model`` flag end-to-end + through ``make_arg_parser`` — mirrors the pattern used in + :mod:`tests.entrypoints.openai.test_cli_args`. + """ + + def test_default_value_on_dataclass_is_true(self): + fields = {f.name: f for f in dataclasses.fields(BaseFrontendArgs)} + assert "cohere_is_reasoning_model" in fields + field = fields["cohere_is_reasoning_model"] + assert field.default is True + assert field.type is bool + + @pytest.fixture + def serve_parser(self) -> FlexibleArgumentParser: + parser = FlexibleArgumentParser() + return make_arg_parser(parser) + + def test_default_via_argparse_is_true(self, serve_parser: FlexibleArgumentParser): + # No flag supplied → dataclass default (True) wins. + args = serve_parser.parse_args(["--model", "m"]) + assert args.cohere_is_reasoning_model is True + + def test_explicit_false_via_argparse(self, serve_parser: FlexibleArgumentParser): + # Boolean dataclass fields are wired up as ``--flag value`` / + # ``--no-flag`` pairs by FlexibleArgumentParser. + args = serve_parser.parse_args( + ["--model", "m", "--no-cohere-is-reasoning-model"] + ) + assert args.cohere_is_reasoning_model is False + + def test_explicit_true_via_argparse(self, serve_parser: FlexibleArgumentParser): + args = serve_parser.parse_args(["--model", "m", "--cohere-is-reasoning-model"]) + assert args.cohere_is_reasoning_model is True diff --git a/tests/entrypoints/cohere/test_serving_conversion.py b/tests/entrypoints/cohere/test_serving_conversion.py new file mode 100644 index 000000000000..680ae775dd46 --- /dev/null +++ b/tests/entrypoints/cohere/test_serving_conversion.py @@ -0,0 +1,1036 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for the Cohere v2 -> OpenAI request / response conversion +implemented in ``vllm/entrypoints/cohere/serving.py``. + +These cover the pure-Python classmethods so we don't need an engine. +For the instance methods that read ``self._is_reasoning_model`` we +build a lightweight :class:`_FakeServing` subclass that skips the +heavy ``OpenAIServingChat.__init__`` chain (which would otherwise need +a real engine client, model registry, etc.) — the same pattern used in +``test_serving_streaming.py``. +""" + +from typing import Any + +import pytest + +from vllm.entrypoints.cohere.protocol import ( + CohereChatV2Request, + CohereChatV2Response, +) +from vllm.entrypoints.cohere.serving import ( + _FINISH_REASON_MAP, + CohereServingChatV2, + ContentBlockType, + _map_finish_reason, +) +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionRequest, + ChatCompletionResponse, +) +from vllm.entrypoints.openai.engine.protocol import ( + Citation as VLLMCitation, +) +from vllm.entrypoints.openai.engine.protocol import ( + CitationSource, +) + +# ---------------------------------------------------------------------- +# Helpers +# ---------------------------------------------------------------------- + + +def _make_request(**kwargs) -> CohereChatV2Request: + kwargs.setdefault("model", "m") + kwargs.setdefault("messages", [{"role": "user", "content": "hi"}]) + return CohereChatV2Request(**kwargs) + + +def _convert(request: CohereChatV2Request) -> ChatCompletionRequest: + return CohereServingChatV2._convert_v2_to_chat_completion(request) + + +class _FakeServing(CohereServingChatV2): + """Lightweight stand-in for :class:`CohereServingChatV2` that skips + the heavy ``OpenAIServingChat.__init__`` chain. + + Only ``_is_reasoning_model`` is read by the methods under test + (``_chat_completion_to_v2`` and friends); the rest of the parent + state is dead weight for unit testing. + """ + + def __init__(self, is_reasoning_model: bool = True) -> None: + # Intentionally skipping super().__init__ — see class docstring. + self._is_reasoning_model = is_reasoning_model + + +def _serving(is_reasoning_model: bool = True) -> CohereServingChatV2: + return _FakeServing(is_reasoning_model=is_reasoning_model) + + +def _build_chat_completion_response( + *, + response_id: str = "resp_1", + content: str | None = "hello", + reasoning: str | None = None, + tool_calls: list[dict[str, Any]] | None = None, + finish_reason: str | None = "stop", + citations: list[Any] | None = None, + usage: dict[str, Any] | None = None, + kv_transfer_params: dict[str, Any] | None = None, +) -> ChatCompletionResponse: + message: dict[str, Any] = {"role": "assistant"} + if content is not None: + message["content"] = content + if reasoning is not None: + message["reasoning"] = reasoning + if tool_calls is not None: + message["tool_calls"] = tool_calls + if citations is not None: + message["citations"] = citations + kwargs: dict[str, Any] = dict( + id=response_id, + object="chat.completion", + created=0, + model="m", + choices=[{"index": 0, "message": message, "finish_reason": finish_reason}], + # ``usage`` is a required field on ChatCompletionResponse, but the + # production code defensively handles ``None`` -> no usage block; + # we round-trip that behavior by post-setting the attribute below. + usage={"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}, + ) + if kv_transfer_params is not None: + kwargs["kv_transfer_params"] = kv_transfer_params + resp = ChatCompletionResponse(**kwargs) + if usage is None: + resp.usage = None + else: + # Replace the placeholder with the caller-provided usage. + resp = ChatCompletionResponse.model_validate( + {**resp.model_dump(), "usage": usage} + ) + if kv_transfer_params is not None: + resp.kv_transfer_params = kv_transfer_params + return resp + + +# ====================================================================== +# _map_finish_reason +# ====================================================================== + + +class TestMapFinishReason: + @pytest.mark.parametrize( + "openai, cohere", + [ + ("stop", "COMPLETE"), + ("length", "MAX_TOKENS"), + ("tool_calls", "TOOL_CALL"), + ("stop_sequence", "STOP_SEQUENCE"), + ("error", "ERROR"), + (None, "COMPLETE"), + ], + ) + def test_known_reasons(self, openai, cohere): + assert _map_finish_reason(openai) == cohere + + def test_unknown_reason_defaults_to_complete(self): + assert _map_finish_reason("not_a_real_reason") == "COMPLETE" + + def test_finish_reason_map_is_complete(self): + # Sanity check that the lookup table covers all documented states. + assert set(_FINISH_REASON_MAP) == { + "stop", + "length", + "tool_calls", + "stop_sequence", + "error", + None, + } + + +# ====================================================================== +# _coerce_text_content (system / tool string fallback) +# ====================================================================== + + +class TestCoerceTextContent: + def test_string_passthrough(self): + assert CohereServingChatV2._coerce_text_content("hi") == "hi" + + def test_concatenates_text_blocks(self): + from cohere.types import SystemChatMessageV2 + + sys_msg = SystemChatMessageV2( + content=[ + {"type": "text", "text": "a"}, + {"type": "text", "text": "b"}, + ] + ) + assert CohereServingChatV2._coerce_text_content(sys_msg.content) == "ab" + + +# ====================================================================== +# User message conversion +# ====================================================================== + + +class TestConvertUserMessage: + def test_string_content(self): + req = _make_request(messages=[{"role": "user", "content": "hi"}]) + result = _convert(req) + assert result.messages == [{"role": "user", "content": "hi"}] + + def test_text_only_list_flattened_to_string(self): + # Single-text-block list is flattened back to a string for maximum + # downstream-template compatibility. + req = _make_request( + messages=[ + { + "role": "user", + "content": [{"type": "text", "text": "hi"}], + } + ] + ) + result = _convert(req) + assert result.messages[0] == {"role": "user", "content": "hi"} + + def test_image_url_content_with_detail(self): + req = _make_request( + messages=[ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": { + "url": "data:image/png;base64,xxx", + "detail": "high", + }, + } + ], + } + ] + ) + result = _convert(req) + msg = result.messages[0] + assert msg["role"] == "user" + assert msg["content"] == [ + { + "type": "image_url", + "image_url": { + "url": "data:image/png;base64,xxx", + "detail": "high", + }, + } + ] + + def test_image_url_without_detail_omits_field(self): + req = _make_request( + messages=[ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": {"url": "https://x/i.png"}, + } + ], + } + ] + ) + result = _convert(req) + assert result.messages[0]["content"][0]["image_url"] == { + "url": "https://x/i.png" + } + + def test_text_plus_image_keeps_list(self): + req = _make_request( + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "describe"}, + { + "type": "image_url", + "image_url": {"url": "https://x/i.png"}, + }, + ], + } + ] + ) + result = _convert(req) + content = result.messages[0]["content"] + assert isinstance(content, list) + assert len(content) == 2 + assert content[0] == {"type": "text", "text": "describe"} + assert content[1]["type"] == "image_url" + + +# ====================================================================== +# Assistant message conversion +# ====================================================================== + + +class TestConvertAssistantMessage: + def test_string_content(self): + req = _make_request( + messages=[ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + ] + ) + result = _convert(req) + asst = result.messages[1] + assert asst == {"role": "assistant", "content": "hello"} + + def test_text_and_thinking_blocks(self): + # ``thinking`` blocks collapse back into the ``reasoning`` field on + # the OpenAI message; ``text`` blocks become ``content``. + req = _make_request( + messages=[ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "let me think"}, + {"type": "text", "text": "Hi!"}, + ], + }, + ] + ) + result = _convert(req) + asst = result.messages[1] + assert asst["role"] == "assistant" + assert asst["content"] == "Hi!" + assert asst["reasoning"] == "let me think" + + def test_thinking_only(self): + # Thinking-only assistant messages have no ``content`` set. + req = _make_request( + messages=[ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "ponder"}, + ], + }, + ] + ) + result = _convert(req) + asst = result.messages[1] + assert asst.get("reasoning") == "ponder" + assert "content" not in asst + + def test_multiple_thinking_blocks_concatenated(self): + req = _make_request( + messages=[ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "first."}, + {"type": "thinking", "thinking": "second."}, + {"type": "text", "text": "done."}, + ], + }, + ] + ) + result = _convert(req) + asst = result.messages[1] + assert asst["reasoning"] == "first.second." + assert asst["content"] == "done." + + def test_tool_plan_collapses_into_reasoning(self): + # Cohere's ``tool_plan`` is the older chain-of-thought field; it + # should be appended to ``reasoning`` so the rendered template + # preserves the planning context. + req = _make_request( + messages=[ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": [ + {"type": "text", "text": "I'll call a tool."}, + ], + "tool_plan": "plan: use calculator", + "tool_calls": [ + { + "id": "c1", + "type": "function", + "function": {"name": "calc", "arguments": "{}"}, + } + ], + }, + ] + ) + result = _convert(req) + asst = result.messages[1] + assert asst["content"] == "I'll call a tool." + assert asst["reasoning"] == "plan: use calculator" + assert asst["tool_calls"][0]["function"] == { + "name": "calc", + "arguments": "{}", + } + + def test_tool_calls_with_missing_function_pieces_get_defaults(self): + # The conversion defends against missing function name/arguments + # by emitting empty string / "{}" defaults so downstream + # validation never sees a None. + req = _make_request( + messages=[ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "tool_calls": [ + { + "id": "c1", + "type": "function", + "function": {"name": "", "arguments": ""}, + } + ], + }, + ] + ) + result = _convert(req) + tc = result.messages[1]["tool_calls"][0] + assert tc["id"] == "c1" + assert tc["type"] == "function" + assert tc["function"] == {"name": "", "arguments": "{}"} + + +# ====================================================================== +# Tool message conversion +# ====================================================================== + + +class TestConvertToolMessage: + def _request_with_tool_message(self, content: Any) -> CohereChatV2Request: + return _make_request( + messages=[ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "tool_calls": [ + { + "id": "c1", + "type": "function", + "function": {"name": "f", "arguments": "{}"}, + } + ], + }, + {"role": "tool", "tool_call_id": "c1", "content": content}, + ] + ) + + def test_string_content(self): + req = self._request_with_tool_message("result text") + result = _convert(req) + tool_msg = result.messages[-1] + assert tool_msg == { + "role": "tool", + "tool_call_id": "c1", + "content": "result text", + } + + def test_text_only_list_flattened_to_newline_string(self): + # Text-only tool results are flattened to a single newline-joined + # string for compatibility with vanilla chat templates. + req = self._request_with_tool_message( + [ + {"type": "text", "text": "line 1"}, + {"type": "text", "text": "line 2"}, + ] + ) + result = _convert(req) + tool_msg = result.messages[-1] + assert tool_msg["content"] == "line 1\nline 2" + + def test_with_document_preserves_structured_content(self): + # When documents appear in the tool result, we keep the list shape + # so the Cohere renderer can lift them into grounding sources. + req = self._request_with_tool_message( + [ + {"type": "text", "text": "see attachment"}, + { + "type": "document", + "document": {"data": {"text": "doc text"}, "id": "d1"}, + }, + ] + ) + result = _convert(req) + tool_msg = result.messages[-1] + assert isinstance(tool_msg["content"], list) + assert tool_msg["content"][0] == {"type": "text", "text": "see attachment"} + assert tool_msg["content"][1] == { + "type": "document", + "document": {"data": {"text": "doc text"}, "id": "d1"}, + } + + +# ====================================================================== +# System message +# ====================================================================== + + +class TestSystemMessage: + def test_system_string(self): + req = _make_request( + messages=[ + {"role": "system", "content": "be helpful"}, + {"role": "user", "content": "hi"}, + ] + ) + result = _convert(req) + assert result.messages[0] == { + "role": "system", + "content": "be helpful", + } + + def test_system_text_blocks_concatenated(self): + req = _make_request( + messages=[ + { + "role": "system", + "content": [ + {"type": "text", "text": "part1 "}, + {"type": "text", "text": "part2"}, + ], + }, + {"role": "user", "content": "hi"}, + ] + ) + result = _convert(req) + assert result.messages[0]["content"] == "part1 part2" + + +# ====================================================================== +# Base ChatCompletionRequest field mapping +# ====================================================================== + + +class TestBuildBaseChatCompletion: + def test_sampling_and_limits_mapped(self): + req = _make_request( + max_tokens=128, + stop_sequences=["", "STOP"], + temperature=0.5, + seed=42, + frequency_penalty=0.1, + presence_penalty=0.2, + k=50, + p=0.95, + logprobs=True, + priority=2, + kv_transfer_params={"x": 1}, + chat_template_kwargs={"y": 2}, + ) + result = _convert(req) + assert result.model == "m" + # ``max_tokens`` is deprecated in favor of ``max_completion_tokens`` + # but the serving code intentionally sets both for compatibility. + assert result.max_completion_tokens == 128 + assert result.stop == ["", "STOP"] + assert result.temperature == 0.5 + assert result.seed == 42 + assert result.frequency_penalty == 0.1 + assert result.presence_penalty == 0.2 + assert result.top_k == 50 + assert result.top_p == 0.95 + assert result.logprobs is True + assert result.priority == 2 + assert result.kv_transfer_params == {"x": 1} + # ``chat_template_kwargs`` may be expanded by _apply_cohere_*; the + # base build at least preserves what the caller passed. + assert (result.chat_template_kwargs or {}).get("y") == 2 + + def test_priority_defaults_to_zero(self): + # ChatCompletionRequest.priority defaults to 0; ``None`` Cohere + # priority must be coerced rather than passed through. + req = _make_request() + result = _convert(req) + assert result.priority == 0 + + +# ====================================================================== +# Streaming options +# ====================================================================== + + +class TestStreamingOptions: + def test_no_stream_leaves_defaults(self): + result = _convert(_make_request(stream=False)) + assert not result.stream + assert result.stream_options is None + + def test_stream_enables_usage_options(self): + # The v2 translator forces ``include_usage=True`` so the + # ``message-end`` event can surface ``billed_units`` / ``tokens``; + # ``continuous_usage_stats`` is intentionally left at its + # ``StreamOptions`` default (False) — Cohere v2 only reports + # usage on the terminal event. + result = _convert(_make_request(stream=True)) + assert result.stream is True + assert result.stream_options is not None + assert result.stream_options.include_usage is True + + +# ====================================================================== +# Response format +# ====================================================================== + + +class TestResponseFormat: + def test_text_is_passthrough(self): + result = _convert(_make_request(response_format={"type": "text"})) + assert result.response_format is None + + def test_json_object(self): + result = _convert(_make_request(response_format={"type": "json_object"})) + assert result.response_format is not None + assert result.response_format.type == "json_object" + assert result.response_format.json_schema is None + + def test_json_schema(self): + schema = {"type": "object", "properties": {"a": {"type": "string"}}} + result = _convert( + _make_request( + response_format={"type": "json_object", "json_schema": schema} + ) + ) + assert result.response_format is not None + assert result.response_format.type == "json_schema" + assert result.response_format.json_schema is not None + assert result.response_format.json_schema.name == "cohere_v2_json_schema" + # ``JsonSchemaResponseFormat.json_schema`` has alias=``schema`` on + # the Pydantic field, so we observe the value via the serialized + # payload (which is what downstream consumers actually read). + dumped = result.response_format.json_schema.model_dump(exclude_none=True) + assert dumped["json_schema"] == schema + + +# ====================================================================== +# Tools / tool_choice +# ====================================================================== + + +class TestApplyTools: + def test_no_tools(self): + result = _convert(_make_request()) + assert result.tools is None + + def test_basic_tool(self): + result = _convert( + _make_request( + tools=[ + { + "type": "function", + "function": { + "name": "calc", + "description": "calculator", + "parameters": {"type": "object"}, + }, + } + ] + ) + ) + assert result.tools is not None + assert len(result.tools) == 1 + tool = result.tools[0] + assert tool.type == "function" + assert tool.function.name == "calc" + assert tool.function.description == "calculator" + # ``strict`` is an extra attribute on FunctionDefinition (the + # field is only stamped onto the OpenAI tool when strict_tools is + # set on the request). The default path must not set it. + assert getattr(tool.function, "strict", None) is None + + def test_strict_tools_propagates_to_function(self): + result = _convert( + _make_request( + strict_tools=True, + tools=[ + { + "type": "function", + "function": { + "name": "calc", + "description": "", + "parameters": {}, + }, + } + ], + ) + ) + assert result.tools[0].function.strict is True + + +class TestApplyToolChoice: + def test_required(self): + result = _convert( + _make_request( + tool_choice="REQUIRED", + tools=[ + { + "type": "function", + "function": { + "name": "f", + "description": "", + "parameters": {}, + }, + } + ], + ) + ) + assert result.tool_choice == "required" + + def test_none(self): + result = _convert( + _make_request( + tool_choice="NONE", + tools=[ + { + "type": "function", + "function": { + "name": "f", + "description": "", + "parameters": {}, + }, + } + ], + ) + ) + assert result.tool_choice == "none" + + def test_default_to_auto_when_tools_present(self): + # No explicit ``tool_choice`` + tools present → auto, mirroring + # Cohere's documented "free choice" default. + result = _convert( + _make_request( + tools=[ + { + "type": "function", + "function": { + "name": "f", + "description": "", + "parameters": {}, + }, + } + ] + ) + ) + assert result.tool_choice == "auto" + + def test_no_tools_no_choice_left_unset(self): + # When there are no tools the underlying ChatCompletionRequest + # default applies; we must not stamp ``auto``. + result = _convert(_make_request()) + assert result.tool_choice != "auto" + + +# ====================================================================== +# Cohere-specific template kwargs forwarding +# ====================================================================== + + +class TestApplyCohereTemplateKwargs: + def test_string_documents_wrapped(self): + result = _convert(_make_request(documents=["doc 1", "doc 2"])) + docs = (result.chat_template_kwargs or {}).get("documents") + assert docs == [ + {"id": "doc_0", "data": {"text": "doc 1"}}, + {"id": "doc_1", "data": {"text": "doc 2"}}, + ] + + def test_document_with_explicit_id_preserved(self): + result = _convert( + _make_request( + documents=[ + {"id": "custom", "data": {"text": "t"}}, + {"data": {"text": "t2"}}, # no id -> synthesized + ] + ) + ) + docs = result.chat_template_kwargs["documents"] + assert docs[0] == {"id": "custom", "data": {"text": "t"}} + assert docs[1]["id"] == "doc_1" + + def test_safety_mode_normalized_to_lowercase(self): + result = _convert(_make_request(safety_mode="CONTEXTUAL")) + assert result.chat_template_kwargs["safety_mode"] == "contextual" + + def test_citation_options_forwarded_as_dict(self): + result = _convert(_make_request(citation_options={"mode": "accurate"})) + assert result.chat_template_kwargs["citation_options"] == {"mode": "accurate"} + + def test_thinking_forwarded_as_dict(self): + result = _convert( + _make_request(thinking={"type": "enabled", "token_budget": 16}) + ) + assert result.chat_template_kwargs["thinking"] == { + "type": "enabled", + "token_budget": 16, + } + + def test_strict_tools_forwarded(self): + result = _convert(_make_request(strict_tools=True)) + assert result.chat_template_kwargs["strict_tools"] is True + + def test_existing_chat_template_kwargs_preserved(self): + # User-supplied kwargs should not be clobbered by the v2 fields + # (setdefault semantics). + result = _convert( + _make_request( + chat_template_kwargs={ + "safety_mode": "user-explicit", + "extra": "x", + }, + safety_mode="CONTEXTUAL", + ) + ) + assert result.chat_template_kwargs["safety_mode"] == "user-explicit" + assert result.chat_template_kwargs["extra"] == "x" + + def test_no_template_kwargs_when_no_cohere_fields(self): + # Without any of the Cohere-specific fields and no caller-supplied + # kwargs, we must leave ``chat_template_kwargs`` as None so other + # renderers see a clean request. + result = _convert(_make_request()) + assert result.chat_template_kwargs is None + + +# ====================================================================== +# _chat_completion_to_v2 (non-streaming response builder) +# ====================================================================== + + +class TestChatCompletionToV2: + def test_text_only(self): + serving = _serving(is_reasoning_model=True) + resp = _build_chat_completion_response(content="hello") + v2 = serving._chat_completion_to_v2(resp, _make_request()) + assert isinstance(v2, CohereChatV2Response) + assert v2.id == "resp_1" + assert v2.finish_reason == "COMPLETE" + assert v2.message.role == "assistant" + assert len(v2.message.content) == 1 + assert v2.message.content[0].type == "text" + assert v2.message.content[0].text == "hello" + assert v2.message.tool_calls is None + assert v2.message.tool_plan is None + assert v2.usage is None + + def test_reasoning_model_keeps_thinking_with_tool_calls(self): + # Reasoning Command models: ``thinking`` block stays in + # ``message.content`` and ``tool_plan`` is left unset, even when + # tool calls are present. + serving = _serving(is_reasoning_model=True) + resp = _build_chat_completion_response( + content="resp text", + reasoning="thoughts", + tool_calls=[ + { + "id": "c1", + "type": "function", + "function": {"name": "f", "arguments": "{}"}, + } + ], + finish_reason="tool_calls", + ) + v2 = serving._chat_completion_to_v2(resp, _make_request()) + assert v2.finish_reason == "TOOL_CALL" + assert v2.message.tool_plan is None + types = [c.type for c in v2.message.content] + assert types == ["thinking", "text"] + assert v2.message.content[0].thinking == "thoughts" + assert v2.message.content[1].text == "resp text" + assert v2.message.tool_calls[0].id == "c1" + assert v2.message.tool_calls[0].function.name == "f" + + def test_non_reasoning_model_moves_reasoning_to_tool_plan(self): + # Older Command models surface reasoning as ``tool_plan`` on tool- + # call turns; the thinking block should be dropped from content. + serving = _serving(is_reasoning_model=False) + resp = _build_chat_completion_response( + content=None, + reasoning="plan", + tool_calls=[ + { + "id": "c1", + "type": "function", + "function": {"name": "f", "arguments": "{}"}, + } + ], + finish_reason="tool_calls", + ) + v2 = serving._chat_completion_to_v2(resp, _make_request()) + assert v2.message.tool_plan == "plan" + assert v2.message.content is None + assert v2.message.tool_calls[0].id == "c1" + + def test_non_reasoning_model_keeps_thinking_when_no_tool_calls(self): + # No tool calls => non-reasoning behavior is identical to + # reasoning behavior; the thinking block stays. + serving = _serving(is_reasoning_model=False) + resp = _build_chat_completion_response( + content="answer", reasoning="plan", tool_calls=None + ) + v2 = serving._chat_completion_to_v2(resp, _make_request()) + assert v2.message.tool_plan is None + types = [c.type for c in v2.message.content] + assert types == ["thinking", "text"] + + def test_id_synthesized_when_response_id_missing(self): + serving = _serving() + resp = _build_chat_completion_response(content="hi", response_id="") + v2 = serving._chat_completion_to_v2(resp, _make_request()) + assert v2.id.startswith("chat_") + + def test_kv_transfer_params_propagated(self): + serving = _serving() + resp = _build_chat_completion_response( + content="hi", kv_transfer_params={"k": 1} + ) + v2 = serving._chat_completion_to_v2(resp, _make_request()) + assert v2.kv_transfer_params == {"k": 1} + + +# ====================================================================== +# _build_usage +# ====================================================================== + + +class TestBuildUsage: + def test_none_passthrough(self): + resp = _build_chat_completion_response(content="hi") + # default usage is None + assert CohereServingChatV2._build_usage(resp) is None + + def test_basic_usage(self): + resp = _build_chat_completion_response( + content="hi", + usage={ + "prompt_tokens": 10, + "completion_tokens": 5, + "total_tokens": 15, + }, + ) + usage = CohereServingChatV2._build_usage(resp) + assert usage is not None + assert usage.billed_units.input_tokens == 10 + assert usage.billed_units.output_tokens == 5 + assert usage.tokens.input_tokens == 10 + assert usage.tokens.output_tokens == 5 + assert usage.cached_tokens is None + + def test_completion_tokens_default_to_zero_when_missing(self): + resp = _build_chat_completion_response( + content="hi", + usage={"prompt_tokens": 10, "completion_tokens": 0, "total_tokens": 10}, + ) + usage = CohereServingChatV2._build_usage(resp) + assert usage.billed_units.output_tokens == 0 + + def test_cached_tokens_propagated(self): + resp = _build_chat_completion_response( + content="hi", + usage={ + "prompt_tokens": 10, + "completion_tokens": 5, + "total_tokens": 15, + "prompt_tokens_details": {"cached_tokens": 3}, + }, + ) + usage = CohereServingChatV2._build_usage(resp) + assert usage.cached_tokens == 3 + + +# ====================================================================== +# _extract_citations_if_any +# ====================================================================== + + +class TestExtractCitations: + def test_none_or_empty_returns_none(self): + assert ( + CohereServingChatV2._extract_citations_if_any( + type("M", (), {"citations": None})() + ) + is None + ) + assert ( + CohereServingChatV2._extract_citations_if_any( + type("M", (), {"citations": []})() + ) + is None + ) + # Missing field entirely. + assert ( + CohereServingChatV2._extract_citations_if_any(type("M", (), {})()) is None + ) + + def test_vllm_citation_objects_normalized(self): + msg = type("M", (), {})() + msg.citations = [ + VLLMCitation( + start=0, + end=5, + text="hello", + sources=[CitationSource(type="document", id="d1")], + ) + ] + out = CohereServingChatV2._extract_citations_if_any(msg) + assert out is not None + assert len(out) == 1 + assert out[0].start == 0 + assert out[0].end == 5 + assert out[0].text == "hello" + + def test_dict_citation_payloads_accepted(self): + msg = type("M", (), {})() + msg.citations = [ + { + "start": 0, + "end": 3, + "text": "hi!", + "sources": [{"type": "document", "id": "d1"}], + } + ] + out = CohereServingChatV2._extract_citations_if_any(msg) + assert out is not None + assert out[0].start == 0 + assert out[0].text == "hi!" + + def test_malformed_citation_skipped_not_raised(self): + msg = type("M", (), {})() + msg.citations = [object()] # neither dict nor Pydantic + out = CohereServingChatV2._extract_citations_if_any(msg) + # All citations dropped → None + assert out is None + + +# ====================================================================== +# create_error_response sanity +# ====================================================================== + + +class TestCreateErrorResponse: + def test_envelope_uses_400(self): + serving = _serving() + err = serving.create_error_response("oops") + assert err.error.message == "oops" + assert err.error.code == 400 + assert err.error.type == "bad_request" + + +# ====================================================================== +# ContentBlockType enum +# ====================================================================== + + +class TestContentBlockType: + def test_values(self): + assert ContentBlockType.THINKING == "thinking" + assert ContentBlockType.TEXT == "text" + assert ContentBlockType.TOOL_CALL == "tool_call" diff --git a/tests/entrypoints/cohere/test_serving_streaming.py b/tests/entrypoints/cohere/test_serving_streaming.py new file mode 100644 index 000000000000..baf006cca627 --- /dev/null +++ b/tests/entrypoints/cohere/test_serving_streaming.py @@ -0,0 +1,815 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for the Cohere v2 SSE stream conversion in +``vllm/entrypoints/cohere/serving.py``. + +The stream-translation entry point is +:meth:`CohereServingChatV2._chat_completion_stream_to_v2`, which turns an +async iterable of OpenAI SSE chunks into Cohere's +``message-start → (content|tool-call|citation)* → message-end → [DONE]`` +event stream. + +We test the helpers (``_StreamState``, ``_handle_*_delta``, etc.) in +isolation, plus a handful of end-to-end scenarios that exercise the +state machine. +""" + +import json +from collections.abc import AsyncGenerator +from typing import Any + +import pytest + +from vllm.entrypoints.cohere.protocol import ( + CohereChatV2Request, + MessageStartEvent, +) +from vllm.entrypoints.cohere.serving import ( + _DONE_FRAME, + CohereServingChatV2, + ContentBlockType, + _emit, + _sse, + _StreamState, +) +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionStreamResponse, +) +from vllm.entrypoints.openai.engine.protocol import ( + Citation as VLLMCitation, +) +from vllm.entrypoints.openai.engine.protocol import ( + CitationSource, +) + +# ---------------------------------------------------------------------- +# Helpers +# ---------------------------------------------------------------------- + + +class _FakeServing(CohereServingChatV2): + """Lightweight stand-in for :class:`CohereServingChatV2` that skips + the heavy ``OpenAIServingChat.__init__`` chain (which would need a + real engine client, model registry, render service, etc.). + + Only ``_is_reasoning_model`` is read by the methods under test + (``_chat_completion_stream_to_v2`` and the per-delta handlers); the + rest is dead weight for unit testing. + """ + + def __init__(self, is_reasoning_model: bool = True) -> None: + # Intentionally skipping super().__init__ — see class docstring. + self._is_reasoning_model = is_reasoning_model + + +def _serving(is_reasoning_model: bool = True) -> CohereServingChatV2: + return _FakeServing(is_reasoning_model=is_reasoning_model) + + +def _parse_event(frame: str) -> dict[str, Any]: + """Strip the ``data: ... \\n\\n`` wrapper and parse the JSON payload.""" + assert frame.startswith("data: ") + assert frame.endswith("\n\n") + return json.loads(frame[len("data: ") : -2]) + + +def _make_chunk( + *, + chunk_id: str = "chunk_0", + role: str | None = None, + content: str | None = None, + reasoning: str | None = None, + tool_calls: list[dict[str, Any]] | None = None, + finish_reason: str | None = None, + usage: dict[str, Any] | None = None, + omit_choices: bool = False, + citations: list[Any] | None = None, +) -> str: + """Build the ``data: {...}\\n\\n`` SSE frame the production code + consumes.""" + delta: dict[str, Any] = {} + if role is not None: + delta["role"] = role + if content is not None: + delta["content"] = content + if reasoning is not None: + delta["reasoning"] = reasoning + if tool_calls is not None: + delta["tool_calls"] = tool_calls + if citations is not None: + delta["citations"] = citations + + payload: dict[str, Any] = { + "id": chunk_id, + "object": "chat.completion.chunk", + "created": 0, + "model": "m", + } + if not omit_choices: + payload["choices"] = [ + {"index": 0, "delta": delta, "finish_reason": finish_reason} + ] + else: + payload["choices"] = [] + if usage is not None: + payload["usage"] = usage + # Use ChatCompletionStreamResponse to normalize the payload shape. + chunk = ChatCompletionStreamResponse.model_validate(payload) + return f"data: {chunk.model_dump_json(exclude_none=False)}\n\n" + + +def _make_done() -> str: + return "data: [DONE]\n\n" + + +async def _stream_from(items: list[str]) -> AsyncGenerator[str, None]: + for item in items: + yield item + + +async def _drain(serving: CohereServingChatV2, items: list[str]) -> list[str]: + """Drive ``_chat_completion_stream_to_v2`` over ``items`` and collect + the emitted SSE frames.""" + request = CohereChatV2Request( + model="m", messages=[{"role": "user", "content": "hi"}], stream=True + ) + gen = serving._chat_completion_stream_to_v2(_stream_from(items), request) + return [frame async for frame in gen] + + +# ====================================================================== +# Low-level helpers: _sse, _emit, _DONE_FRAME +# ====================================================================== + + +class TestSSEHelpers: + def test_sse_wraps_payload(self): + assert _sse("hello") == "data: hello\n\n" + + def test_done_frame_constant(self): + # cohere-python and Fern-generated clients key off this exact + # sentinel; keep it byte-for-byte stable. + assert _DONE_FRAME == "data: [DONE]\n\n" + + def test_emit_serializes_event_with_type_discriminator(self): + frame = _emit( + MessageStartEvent(id="abc", delta={"message": {"role": "assistant"}}) + ) + payload = _parse_event(frame) + assert payload["type"] == "message-start" + assert payload["id"] == "abc" + assert payload["delta"] == {"message": {"role": "assistant"}} + + +# ====================================================================== +# _StreamState +# ====================================================================== + + +class TestStreamState: + def test_defaults(self): + st = _StreamState() + assert st.started is False + assert st.ended is False + assert st.finish_reason is None + assert st.last_chunk_id == "" + assert st.active_block is None + assert st.active_block_index is None + assert st.active_tool_index is None + assert st.tool_calls_seen == set() + + def test_content_index_monotonic(self): + st = _StreamState() + assert st.next_content_index() == 0 + assert st.next_content_index() == 1 + assert st.next_content_index() == 2 + + def test_citation_index_separate_from_content_index(self): + st = _StreamState() + st.next_content_index() # 0 + st.next_content_index() # 1 + # Citation indexing is independent of content indexing. + assert st.next_citation_index() == 0 + assert st.next_citation_index() == 1 + + +# ====================================================================== +# _close_open_blocks +# ====================================================================== + + +class TestCloseOpenBlocks: + def test_no_block_open_emits_nothing(self): + serving = _serving() + state = _StreamState() + assert serving._close_open_blocks(state) == [] + + def test_text_block_emits_content_end(self): + serving = _serving() + state = _StreamState() + state.active_block = ContentBlockType.TEXT + state.active_block_index = 1 + out = serving._close_open_blocks(state) + assert len(out) == 1 + payload = _parse_event(out[0]) + assert payload == {"type": "content-end", "index": 1} + assert state.active_block is None + assert state.active_block_index is None + + def test_thinking_block_emits_content_end(self): + serving = _serving() + state = _StreamState() + state.active_block = ContentBlockType.THINKING + state.active_block_index = 3 + out = serving._close_open_blocks(state) + payload = _parse_event(out[0]) + assert payload == {"type": "content-end", "index": 3} + + def test_tool_call_block_emits_tool_call_end(self): + serving = _serving() + state = _StreamState() + state.active_block = ContentBlockType.TOOL_CALL + state.active_tool_index = 7 + out = serving._close_open_blocks(state) + payload = _parse_event(out[0]) + assert payload == {"type": "tool-call-end", "index": 7} + assert state.active_tool_index is None + + +# ====================================================================== +# _handle_text_delta +# ====================================================================== + + +class TestHandleTextDelta: + def test_opens_block_first_time(self): + serving = _serving() + state = _StreamState() + events = serving._handle_text_delta(state, "Hi") + assert len(events) == 2 + start = _parse_event(events[0]) + delta = _parse_event(events[1]) + assert start["type"] == "content-start" + assert start["index"] == 0 + assert start["delta"]["message"]["content"]["type"] == "text" + assert delta["type"] == "content-delta" + assert delta["index"] == 0 + assert delta["delta"]["message"]["content"]["text"] == "Hi" + assert state.active_block == ContentBlockType.TEXT + assert state.active_block_index == 0 + + def test_continues_block_with_just_delta(self): + serving = _serving() + state = _StreamState() + serving._handle_text_delta(state, "Hi") + events = serving._handle_text_delta(state, " there") + # Only a delta event, no new content-start. + assert len(events) == 1 + delta = _parse_event(events[0]) + assert delta["type"] == "content-delta" + assert delta["delta"]["message"]["content"]["text"] == " there" + + def test_switches_from_thinking_block(self): + serving = _serving(is_reasoning_model=True) + state = _StreamState() + # Open a thinking block first, then switch to text. + serving._handle_thinking_delta(state, "ponder") + events = serving._handle_text_delta(state, "answer") + types = [_parse_event(ev)["type"] for ev in events] + assert types == ["content-end", "content-start", "content-delta"] + # The text block gets a new index (1), distinct from thinking's 0. + assert _parse_event(events[1])["index"] == 1 + assert state.active_block == ContentBlockType.TEXT + + +# ====================================================================== +# _handle_thinking_delta +# ====================================================================== + + +class TestHandleThinkingDelta: + def test_reasoning_model_opens_thinking_block(self): + serving = _serving(is_reasoning_model=True) + state = _StreamState() + events = serving._handle_thinking_delta(state, "thought") + assert len(events) == 2 + start = _parse_event(events[0]) + delta = _parse_event(events[1]) + assert start["type"] == "content-start" + assert start["delta"]["message"]["content"]["type"] == "thinking" + assert delta["type"] == "content-delta" + assert delta["delta"]["message"]["content"]["thinking"] == "thought" + assert state.active_block == ContentBlockType.THINKING + + def test_reasoning_model_continues_thinking_block(self): + serving = _serving(is_reasoning_model=True) + state = _StreamState() + serving._handle_thinking_delta(state, "first") + events = serving._handle_thinking_delta(state, " more") + assert len(events) == 1 + assert _parse_event(events[0])["type"] == "content-delta" + + def test_non_reasoning_model_emits_tool_plan_delta(self): + # Older Command models stream reasoning as ``tool_plan`` deltas; + # no content-start/end pair is emitted. + serving = _serving(is_reasoning_model=False) + state = _StreamState() + events = serving._handle_thinking_delta(state, "planning") + assert len(events) == 1 + payload = _parse_event(events[0]) + assert payload["type"] == "tool-plan-delta" + assert payload["delta"]["message"]["tool_plan"] == "planning" + # ``tool_plan`` deltas don't claim an active content block. + assert state.active_block is None + + def test_non_reasoning_model_closes_open_text_block(self): + serving = _serving(is_reasoning_model=False) + state = _StreamState() + # Open a text block first. + serving._handle_text_delta(state, "answer") + events = serving._handle_thinking_delta(state, "rethink") + types = [_parse_event(ev)["type"] for ev in events] + assert types == ["content-end", "tool-plan-delta"] + + +# ====================================================================== +# _handle_tool_call_deltas +# ====================================================================== + + +class TestHandleToolCallDeltas: + def test_new_tool_call_opens_tool_call_start(self): + serving = _serving() + state = _StreamState() + deltas = [ + type( + "Delta", + (), + { + "index": 0, + "id": "c1", + "function": type( + "Fn", (), {"name": "calc", "arguments": '{"x":'} + )(), + }, + )() + ] + events = serving._handle_tool_call_deltas(state, deltas) + assert len(events) == 1 + payload = _parse_event(events[0]) + assert payload["type"] == "tool-call-start" + assert payload["index"] == 0 + tc = payload["delta"]["message"]["tool_calls"] + assert tc["id"] == "c1" + assert tc["function"]["name"] == "calc" + assert tc["function"]["arguments"] == '{"x":' + assert state.active_block == ContentBlockType.TOOL_CALL + assert state.active_tool_index == 0 + assert 0 in state.tool_calls_seen + + def test_subsequent_arguments_emit_delta(self): + serving = _serving() + state = _StreamState() + # First call: start. + first = [ + type( + "Delta", + (), + { + "index": 0, + "id": "c1", + "function": type("Fn", (), {"name": "calc", "arguments": ""})(), + }, + )() + ] + serving._handle_tool_call_deltas(state, first) + # Second call: same index, additional arguments fragment. + more = [ + type( + "Delta", + (), + { + "index": 0, + "id": None, + "function": type("Fn", (), {"name": None, "arguments": "1}"})(), + }, + )() + ] + events = serving._handle_tool_call_deltas(state, more) + assert len(events) == 1 + payload = _parse_event(events[0]) + assert payload["type"] == "tool-call-delta" + assert payload["index"] == 0 + assert ( + payload["delta"]["message"]["tool_calls"]["function"]["arguments"] == "1}" + ) + + def test_new_tool_call_closes_existing_content_block(self): + serving = _serving() + state = _StreamState() + # Open a text block, then start a tool call. + serving._handle_text_delta(state, "I'll call:") + deltas = [ + type( + "Delta", + (), + { + "index": 0, + "id": "c1", + "function": type("Fn", (), {"name": "calc", "arguments": "{}"})(), + }, + )() + ] + events = serving._handle_tool_call_deltas(state, deltas) + types = [_parse_event(ev)["type"] for ev in events] + assert types == ["content-end", "tool-call-start"] + + +# ====================================================================== +# _handle_citation_deltas +# ====================================================================== + + +class TestHandleCitationDeltas: + def test_citation_objects_emit_start_and_end(self): + serving = _serving() + state = _StreamState() + citations = [ + VLLMCitation( + start=0, + end=5, + text="hello", + sources=[CitationSource(type="document", id="d1")], + ) + ] + events = serving._handle_citation_deltas(state, citations) + assert len(events) == 2 + start = _parse_event(events[0]) + end = _parse_event(events[1]) + assert start["type"] == "citation-start" + assert start["index"] == 0 + cit_payload = start["delta"]["message"]["citations"] + assert cit_payload["start"] == 0 + assert cit_payload["end"] == 5 + assert cit_payload["text"] == "hello" + assert end == {"type": "citation-end", "index": 0} + + def test_dict_citations_accepted(self): + serving = _serving() + state = _StreamState() + citations = [{"start": 0, "end": 3, "text": "Hi!", "sources": []}] + events = serving._handle_citation_deltas(state, citations) + assert len(events) == 2 + assert _parse_event(events[0])["type"] == "citation-start" + assert _parse_event(events[1])["type"] == "citation-end" + + def test_malformed_citation_skipped(self): + serving = _serving() + state = _StreamState() + # Neither a dict nor a Pydantic-model-like object. + events = serving._handle_citation_deltas(state, [object()]) + assert events == [] + + +# ====================================================================== +# _build_message_end_event +# ====================================================================== + + +class TestBuildMessageEndEvent: + def test_without_usage(self): + serving = _serving() + frame = serving._build_message_end_event(chunk_id="abc", finish_reason="stop") + payload = _parse_event(frame) + assert payload["type"] == "message-end" + assert payload["id"] == "abc" + assert payload["delta"]["finish_reason"] == "COMPLETE" + assert "usage" not in payload["delta"] + + def test_with_usage(self): + serving = _serving() + chunk = ChatCompletionStreamResponse.model_validate( + { + "id": "abc", + "object": "chat.completion.chunk", + "created": 0, + "model": "m", + "choices": [], + "usage": { + "prompt_tokens": 10, + "completion_tokens": 5, + "total_tokens": 15, + }, + } + ) + frame = serving._build_message_end_event( + chunk_id="abc", finish_reason="length", usage_chunk=chunk + ) + payload = _parse_event(frame) + assert payload["delta"]["finish_reason"] == "MAX_TOKENS" + usage = payload["delta"]["usage"] + assert usage["billed_units"] == {"input_tokens": 10, "output_tokens": 5} + assert usage["tokens"] == {"input_tokens": 10, "output_tokens": 5} + assert "cached_tokens" not in usage + + def test_with_cached_tokens(self): + serving = _serving() + chunk = ChatCompletionStreamResponse.model_validate( + { + "id": "abc", + "object": "chat.completion.chunk", + "created": 0, + "model": "m", + "choices": [], + "usage": { + "prompt_tokens": 10, + "completion_tokens": 5, + "total_tokens": 15, + "prompt_tokens_details": {"cached_tokens": 3}, + }, + } + ) + frame = serving._build_message_end_event( + chunk_id="abc", finish_reason="stop", usage_chunk=chunk + ) + payload = _parse_event(frame) + assert payload["delta"]["usage"]["cached_tokens"] == 3 + + +# ====================================================================== +# End-to-end: _chat_completion_stream_to_v2 +# ====================================================================== + + +class TestChatCompletionStreamToV2: + """End-to-end stream lifecycle tests.""" + + @pytest.mark.asyncio + async def test_text_only_happy_path_emits_full_lifecycle(self): + serving = _serving(is_reasoning_model=True) + items = [ + _make_chunk(role="assistant"), + _make_chunk(content="Hi"), + _make_chunk(content=" there"), + _make_chunk( + omit_choices=True, + usage={ + "prompt_tokens": 4, + "completion_tokens": 2, + "total_tokens": 6, + }, + ), + ] + frames = await _drain(serving, items) + types = [_parse_event(f)["type"] for f in frames[:-1]] + # message-start, content-start, content-delta, content-delta, + # content-end, message-end, then [DONE] as the last frame. + assert types == [ + "message-start", + "content-start", + "content-delta", + "content-delta", + "content-end", + "message-end", + ] + assert frames[-1] == _DONE_FRAME + # message-end should carry usage stats from the trailing chunk. + end_payload = _parse_event(frames[-2]) + assert end_payload["delta"]["finish_reason"] == "COMPLETE" + assert end_payload["delta"]["usage"]["billed_units"]["input_tokens"] == 4 + + @pytest.mark.asyncio + async def test_text_then_tool_call_closes_text_first(self): + serving = _serving(is_reasoning_model=True) + items = [ + _make_chunk(role="assistant"), + _make_chunk(content="planning..."), + _make_chunk( + tool_calls=[ + { + "index": 0, + "id": "c1", + "type": "function", + "function": {"name": "calc", "arguments": "{}"}, + } + ] + ), + _make_chunk(finish_reason="tool_calls"), + _make_chunk( + omit_choices=True, + usage={ + "prompt_tokens": 1, + "completion_tokens": 1, + "total_tokens": 2, + }, + ), + ] + frames = await _drain(serving, items) + types = [_parse_event(f)["type"] for f in frames[:-1]] + # Text block is opened+delta, then closed before the tool-call + # opens, and the final close happens on the usage chunk path. + assert types == [ + "message-start", + "content-start", + "content-delta", + "content-end", + "tool-call-start", + "tool-call-end", + "message-end", + ] + # finish_reason captured from the prior chunk. + end = _parse_event(frames[-2]) + assert end["delta"]["finish_reason"] == "TOOL_CALL" + + @pytest.mark.asyncio + async def test_thinking_then_text_reasoning_model(self): + serving = _serving(is_reasoning_model=True) + items = [ + _make_chunk(role="assistant"), + _make_chunk(reasoning="thinking..."), + _make_chunk(content="answer"), + _make_chunk(finish_reason="stop"), + _make_chunk( + omit_choices=True, + usage={ + "prompt_tokens": 1, + "completion_tokens": 1, + "total_tokens": 2, + }, + ), + ] + frames = await _drain(serving, items) + types = [_parse_event(f)["type"] for f in frames[:-1]] + # Thinking block opens with index 0; text block reopens with index 1. + assert types == [ + "message-start", + "content-start", + "content-delta", + "content-end", + "content-start", + "content-delta", + "content-end", + "message-end", + ] + thinking_start = _parse_event(frames[1]) + assert thinking_start["delta"]["message"]["content"]["type"] == "thinking" + text_start = _parse_event(frames[4]) + assert text_start["delta"]["message"]["content"]["type"] == "text" + + @pytest.mark.asyncio + async def test_reasoning_on_non_reasoning_model_emits_tool_plan_delta(self): + serving = _serving(is_reasoning_model=False) + items = [ + _make_chunk(role="assistant"), + _make_chunk(reasoning="planning"), + _make_chunk( + tool_calls=[ + { + "index": 0, + "id": "c1", + "type": "function", + "function": {"name": "f", "arguments": "{}"}, + } + ] + ), + _make_chunk(finish_reason="tool_calls"), + _make_chunk( + omit_choices=True, + usage={ + "prompt_tokens": 1, + "completion_tokens": 1, + "total_tokens": 2, + }, + ), + ] + frames = await _drain(serving, items) + types = [_parse_event(f)["type"] for f in frames[:-1]] + assert types == [ + "message-start", + "tool-plan-delta", + "tool-call-start", + "tool-call-end", + "message-end", + ] + # No thinking content blocks should be present. + assert "content-start" not in types + assert "content-end" not in types + + @pytest.mark.asyncio + async def test_done_marker_in_middle_closes_open_block(self): + # Some upstreams send [DONE] without a trailing usage-only chunk. + # The translator must still emit message-end before [DONE] so + # Cohere clients don't hang. + serving = _serving(is_reasoning_model=True) + items = [ + _make_chunk(role="assistant"), + _make_chunk(content="Hi"), + _make_done(), + ] + frames = await _drain(serving, items) + types = [_parse_event(f)["type"] for f in frames[:-1]] + assert types == [ + "message-start", + "content-start", + "content-delta", + "content-end", + "message-end", + ] + assert frames[-1] == _DONE_FRAME + + @pytest.mark.asyncio + async def test_skips_empty_and_non_data_lines(self): + serving = _serving() + items = [ + "\n", + "event: ping\n\n", + _make_chunk(role="assistant"), + "data: \n\n", # empty data + _make_chunk(content="Hi"), + _make_chunk( + omit_choices=True, + usage={ + "prompt_tokens": 1, + "completion_tokens": 1, + "total_tokens": 2, + }, + ), + ] + frames = await _drain(serving, items) + assert frames[-1] == _DONE_FRAME + # Should still produce a complete lifecycle. + types = [_parse_event(f)["type"] for f in frames[:-1]] + assert types[0] == "message-start" + assert types[-1] == "message-end" + + @pytest.mark.asyncio + async def test_exception_in_chunk_parsing_emits_error_message_end(self): + # An invalid JSON payload after the first chunk triggers the + # error path: a synthetic message-end with finish_reason=ERROR + # followed by [DONE]. + serving = _serving() + items = [ + _make_chunk(role="assistant"), + "data: {not valid json}\n\n", + ] + frames = await _drain(serving, items) + assert frames[-1] == _DONE_FRAME + # Find the error-shaped message-end. + error_end = _parse_event(frames[-2]) + assert error_end["type"] == "message-end" + assert error_end["delta"]["finish_reason"] == "ERROR" + assert "error" in error_end["delta"] + + @pytest.mark.asyncio + async def test_citations_in_delta_emit_citation_events(self): + serving = _serving() + # Pass citation dicts the way the cohere2 reasoning parser does. + items = [ + _make_chunk(role="assistant"), + _make_chunk(content="hello"), + _make_chunk( + citations=[ + { + "start": 0, + "end": 5, + "text": "hello", + "sources": [{"type": "document", "id": "d1"}], + } + ] + ), + _make_chunk(finish_reason="stop"), + _make_chunk( + omit_choices=True, + usage={ + "prompt_tokens": 1, + "completion_tokens": 1, + "total_tokens": 2, + }, + ), + ] + frames = await _drain(serving, items) + types = [_parse_event(f)["type"] for f in frames[:-1]] + assert "citation-start" in types + assert "citation-end" in types + + @pytest.mark.asyncio + async def test_first_chunk_emits_message_start_with_chunk_id(self): + serving = _serving() + items = [ + _make_chunk(chunk_id="my-id", role="assistant"), + _make_chunk(chunk_id="my-id", content="hi"), + _make_chunk( + chunk_id="my-id", + omit_choices=True, + usage={ + "prompt_tokens": 1, + "completion_tokens": 1, + "total_tokens": 2, + }, + ), + ] + frames = await _drain(serving, items) + ms = _parse_event(frames[0]) + assert ms["type"] == "message-start" + assert ms["id"] == "my-id" + assert ms["delta"]["message"]["role"] == "assistant" diff --git a/tests/renderers/test_cohere.py b/tests/renderers/test_cohere.py new file mode 100644 index 000000000000..4a7a23d2b737 --- /dev/null +++ b/tests/renderers/test_cohere.py @@ -0,0 +1,743 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for ``vllm/renderers/cohere.py``. + +The tests focus on the pure-Python helpers that produce the render-config +dicts passed to ``cohere_melody.render_cmd3`` / ``render_cmd4``. We also +include a class-level instantiation + async-non-blocking test that +mirrors the analogous ``test_mistral.py`` pattern, exercising the +:class:`CohereRenderer` end-to-end with mocked ``model_config`` / +tokenizer / melody bindings. +""" + +import asyncio +import json +import time +from dataclasses import dataclass +from typing import Any +from unittest.mock import Mock + +import pytest + +from vllm.renderers import ChatParams +from vllm.renderers.cohere import ( + CohereRenderer, + MelodyContentType, + _build_render_config, + _content_blocks, + _conversation_to_melody_messages, + _document_to_melody, + _normalize_tool_call, + _role_to_melody, + _tool_to_melody, +) +from vllm.tokenizers.hf import HfTokenizer + +# ====================================================================== +# _role_to_melody +# ====================================================================== + + +class TestRoleToMelody: + def test_assistant_maps_to_chatbot(self): + # melody's templates use the legacy Cohere ``chatbot`` role name. + assert _role_to_melody("assistant") == "chatbot" + + def test_developer_aliases_to_system(self): + # OpenAI's ``developer`` role is documented as high-priority + # instructions; map it onto the ``system`` slot rather than + # letting the templates drop it on the floor. + assert _role_to_melody("developer") == "system" + + @pytest.mark.parametrize("role", ["user", "system", "tool", "chatbot"]) + def test_recognized_roles_passthrough(self, role): + assert _role_to_melody(role) == role + + @pytest.mark.parametrize( + "role,expected", + [ + ("ASSISTANT", "chatbot"), + ("Developer", "system"), + ("User", "user"), + ("SYSTEM", "system"), + ], + ) + def test_role_normalization_is_case_insensitive(self, role, expected): + # cmd3 / cmd4 templates lowercase the role before matching, so + # accept any casing the caller provides. + assert _role_to_melody(role) == expected + + @pytest.mark.parametrize("role", ["function", "moderator", "", "anything"]) + def test_unknown_roles_raise(self, role): + # Silently dropping unknown roles produces malformed prompts + # (the templates' role chain has no else branch). + with pytest.raises(ValueError, match="Unsupported message role"): + _role_to_melody(role) + + def test_non_string_role_rejected(self): + # The function is typed ``role: str`` and the implementation + # relies on Python's attribute lookup (``role.lower()``) to + # reject non-strings — any exception type is acceptable as long + # as we don't silently produce a malformed prompt. + with pytest.raises((AttributeError, TypeError, ValueError)): + _role_to_melody(None) # type: ignore[arg-type] + + +# ====================================================================== +# _normalize_tool_call +# ====================================================================== + + +class TestNormalizeToolCall: + def test_openai_dict_with_dict_arguments_json_encoded(self): + # melody expects ``parameters`` as a JSON-encoded string even when + # OpenAI delivers an already-parsed dict. + out = _normalize_tool_call( + { + "id": "c1", + "type": "function", + "function": {"name": "f", "arguments": {"a": 1}}, + } + ) + assert out == {"id": "c1", "name": "f", "parameters": '{"a": 1}'} + + def test_openai_dict_with_string_arguments_preserved(self): + out = _normalize_tool_call( + { + "id": "c1", + "type": "function", + "function": {"name": "f", "arguments": '{"a":1}'}, + } + ) + assert out["parameters"] == '{"a":1}' + + def test_flat_dict_without_function_wrapper(self): + out = _normalize_tool_call({"id": "c1", "name": "f", "arguments": '{"k": 1}'}) + # Falls back to top-level ``name`` / ``arguments``. + assert out == {"id": "c1", "name": "f", "parameters": '{"k": 1}'} + + def test_missing_id_becomes_empty_string(self): + out = _normalize_tool_call({"function": {"name": "f", "arguments": "{}"}}) + assert out["id"] == "" + + def test_pydantic_model_dump_supported(self): + class _Fake: + def model_dump(self): + return { + "id": "c1", + "type": "function", + "function": {"name": "f", "arguments": "{}"}, + } + + out = _normalize_tool_call(_Fake()) + assert out == {"id": "c1", "name": "f", "parameters": "{}"} + + def test_invalid_type_rejected(self): + with pytest.raises(TypeError, match="Unexpected tool_call value"): + _normalize_tool_call(42) # type: ignore[arg-type] + + +# ====================================================================== +# _content_blocks +# ====================================================================== + + +class TestContentBlocks: + def test_none_returns_empty_list(self): + assert _content_blocks(None) == [] + + def test_string_wrapped_in_text_block(self): + out = _content_blocks("hi") + assert out == [{"type": MelodyContentType.TEXT, "text": "hi"}] + + def test_string_item_in_list_wrapped(self): + out = _content_blocks(["a", "b"]) + assert out == [ + {"type": MelodyContentType.TEXT, "text": "a"}, + {"type": MelodyContentType.TEXT, "text": "b"}, + ] + + @pytest.mark.parametrize( + "part_type", + ["text", "input_text", "output_text", "refusal"], + ) + def test_text_variants_normalized(self, part_type): + out = _content_blocks([{"type": part_type, "text": "hello"}]) + assert out == [{"type": MelodyContentType.TEXT, "text": "hello"}] + + def test_thinking_block(self): + out = _content_blocks([{"type": "thinking", "thinking": "thoughts"}]) + assert out == [{"type": MelodyContentType.THINKING, "thinking": "thoughts"}] + + def test_image_block_with_default_placeholder(self): + out = _content_blocks([{"type": "image"}]) + assert out == [ + { + "type": MelodyContentType.IMAGE, + "image": {"template_placeholder": ""}, + } + ] + + def test_image_block_custom_placeholder(self): + out = _content_blocks([{"type": "image", "template_placeholder": "[[IMG]]"}]) + assert out[0]["image"]["template_placeholder"] == "[[IMG]]" + + def test_document_block_dict_passthrough(self): + out = _content_blocks( + [{"type": "document", "document": {"data": {"text": "doc"}}}] + ) + assert out == [ + { + "type": MelodyContentType.DOCUMENT, + "document": {"data": {"text": "doc"}}, + } + ] + + def test_document_block_with_non_dict_falls_back_to_json_text(self): + out = _content_blocks([{"type": "document", "document": "raw string doc"}]) + assert out[0]["type"] == MelodyContentType.TEXT + # JSON-encoded for safety since melody expects a structured doc. + assert out[0]["text"] == json.dumps("raw string doc") + + def test_tool_reference_emitted_as_text(self): + out = _content_blocks([{"type": "tool_reference", "name": "calc"}]) + assert out == [{"type": MelodyContentType.TEXT, "text": "calc"}] + + def test_unknown_block_type_fallback_to_text(self): + # Unknown block type with a string value is wrapped in a text block. + out = _content_blocks([{"type": "custom", "custom": "value"}]) + assert out == [{"type": MelodyContentType.TEXT, "text": "value"}] + + def test_unknown_block_type_dict_value_json_encoded(self): + out = _content_blocks([{"type": "custom", "custom": {"k": 1}}]) + assert out == [{"type": MelodyContentType.TEXT, "text": json.dumps({"k": 1})}] + + def test_non_string_non_dict_part_rejected(self): + with pytest.raises(TypeError, match="Unexpected content part"): + _content_blocks([42]) # type: ignore[list-item] + + +# ====================================================================== +# _document_to_melody +# ====================================================================== + + +class TestDocumentToMelody: + def test_string_wrapped_in_text_dict(self): + assert _document_to_melody("hello") == {"text": "hello"} + + def test_pure_dict_passthrough(self): + out = _document_to_melody({"text": "x", "id": "d1"}) + assert out == {"text": "x", "id": "d1"} + # Must be a defensive copy (mutating output should not affect input). + out["new_key"] = "value" + + def test_data_wrapper_flattened(self): + # Cohere v2 documents use ``{id, data: {...}}``; melody expects + # the flat shape with ``id`` merged into the payload. + out = _document_to_melody({"id": "d1", "data": {"text": "hello", "title": "t"}}) + assert out == {"id": "d1", "text": "hello", "title": "t"} + + def test_data_wrapper_preserves_inner_id(self): + # If the inner ``data`` already has an ``id``, it wins. + out = _document_to_melody({"id": "outer", "data": {"id": "inner", "text": "x"}}) + assert out["id"] == "inner" + + def test_invalid_type_rejected(self): + with pytest.raises(TypeError, match="Unsupported document type"): + _document_to_melody(42) # type: ignore[arg-type] + + +# ====================================================================== +# _tool_to_melody +# ====================================================================== + + +class TestToolToMelody: + def test_openai_wrapper(self): + out = _tool_to_melody( + { + "type": "function", + "function": { + "name": "calc", + "description": "calculate", + "parameters": {"type": "object"}, + }, + } + ) + assert out == { + "name": "calc", + "description": "calculate", + "parameters": {"type": "object"}, + } + + def test_flat_dict(self): + out = _tool_to_melody({"name": "calc", "description": "d", "parameters": {}}) + assert out["name"] == "calc" + assert out["parameters"] == {} + + def test_pydantic_like_model_dump(self): + class _Fake: + def model_dump(self): + return { + "type": "function", + "function": { + "name": "calc", + "description": "x", + "parameters": {}, + }, + } + + out = _tool_to_melody(_Fake()) + assert out["name"] == "calc" + + def test_missing_description_becomes_empty(self): + out = _tool_to_melody({"name": "calc"}) + assert out["description"] == "" + assert out["parameters"] == {} + + def test_invalid_type_rejected(self): + with pytest.raises(TypeError, match="Unsupported tool type"): + _tool_to_melody(42) # type: ignore[arg-type] + + +# ====================================================================== +# _conversation_to_melody_messages +# ====================================================================== + + +class TestConversationToMelody: + def test_basic_user_assistant_pair(self): + conv = [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + ] + out = _conversation_to_melody_messages(conv) # type: ignore[arg-type] + assert out == [ + { + "role": "user", + "content": [{"type": MelodyContentType.TEXT, "text": "hi"}], + "tool_calls": [], + }, + { + "role": "chatbot", + "content": [{"type": MelodyContentType.TEXT, "text": "hello"}], + "tool_calls": [], + }, + ] + + def test_assistant_reasoning_prepended_as_thinking_block(self): + # ``reasoning`` (or ``reasoning_content``) is prepended as a + # ``thinking`` block on assistant turns, preserving multi-turn + # chain-of-thought across the rendered prompt. + conv = [ + { + "role": "assistant", + "content": "answer", + "reasoning": "thoughts", + } + ] + out = _conversation_to_melody_messages(conv) # type: ignore[arg-type] + assert out[0]["content"] == [ + {"type": MelodyContentType.THINKING, "thinking": "thoughts"}, + {"type": MelodyContentType.TEXT, "text": "answer"}, + ] + + def test_assistant_reasoning_content_alias_accepted(self): + conv = [ + { + "role": "assistant", + "content": "answer", + "reasoning_content": "thoughts", + } + ] + out = _conversation_to_melody_messages(conv) # type: ignore[arg-type] + assert out[0]["content"][0] == { + "type": MelodyContentType.THINKING, + "thinking": "thoughts", + } + + def test_user_reasoning_ignored(self): + # Only assistant turns get reasoning-as-thinking lifting; user + # turns with a ``reasoning`` key (which shouldn't happen in + # practice) must not produce a phantom thinking block. + conv = [ + { + "role": "user", + "content": "hi", + "reasoning": "should be ignored", + } + ] + out = _conversation_to_melody_messages(conv) # type: ignore[arg-type] + assert out[0]["content"] == [{"type": MelodyContentType.TEXT, "text": "hi"}] + + def test_tool_calls_normalized(self): + conv = [ + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "c1", + "type": "function", + "function": {"name": "f", "arguments": '{"a":1}'}, + } + ], + } + ] + out = _conversation_to_melody_messages(conv) # type: ignore[arg-type] + assert out[0]["tool_calls"] == [ + {"id": "c1", "name": "f", "parameters": '{"a":1}'} + ] + + def test_tool_call_id_preserved_on_tool_role(self): + conv = [ + { + "role": "tool", + "content": "result", + "tool_call_id": "c1", + } + ] + out = _conversation_to_melody_messages(conv) # type: ignore[arg-type] + assert out[0]["tool_call_id"] == "c1" + + +# ====================================================================== +# _build_render_config +# ====================================================================== + + +class TestBuildRenderConfig: + def _conv(self): + return [{"role": "user", "content": "hi"}] + + def test_default_format_is_cmd3(self): + fmt, cfg = _build_render_config(self._conv(), {}) # type: ignore[arg-type] + assert fmt == "cmd3" + assert cfg["use_jinja"] is True + assert isinstance(cfg["messages"], list) + # No additional_template_fields when no extra kwargs are set. + assert "additional_template_fields" not in cfg + + def test_explicit_cmd4(self): + fmt, cfg = _build_render_config(self._conv(), {"cohere_format": "cmd4"}) # type: ignore[arg-type] + assert fmt == "cmd4" + + def test_invalid_format_raises(self): + with pytest.raises(ValueError, match="Invalid cohere_format"): + _build_render_config(self._conv(), {"cohere_format": "cmd5"}) # type: ignore[arg-type] + + def test_documents_converted(self): + _, cfg = _build_render_config( + self._conv(), + { + "documents": [ + "doc text", + {"id": "d1", "data": {"text": "wrapped"}}, + ] + }, + ) # type: ignore[arg-type] + assert cfg["documents"] == [ + {"text": "doc text"}, + {"id": "d1", "text": "wrapped"}, + ] + + def test_available_tools_take_precedence_over_tools(self): + _, cfg = _build_render_config( + self._conv(), + { + "tools": [{"type": "function", "function": {"name": "from_tools"}}], + "available_tools": [ + {"type": "function", "function": {"name": "preferred"}} + ], + }, + ) # type: ignore[arg-type] + names = [t["name"] for t in cfg["available_tools"]] + assert names == ["preferred"] + + def test_tools_used_when_no_available_tools(self): + _, cfg = _build_render_config( + self._conv(), + {"tools": [{"type": "function", "function": {"name": "from_tools"}}]}, + ) # type: ignore[arg-type] + assert [t["name"] for t in cfg["available_tools"]] == ["from_tools"] + + @pytest.mark.parametrize("value", ["enabled", "disabled"]) + def test_reasoning_type_direct(self, value): + _, cfg = _build_render_config(self._conv(), {"reasoning_type": value}) # type: ignore[arg-type] + assert cfg["reasoning_type"] == value + + def test_thinking_dict_shorthand_resolves_reasoning_type(self): + _, cfg = _build_render_config(self._conv(), {"thinking": {"type": "enabled"}}) # type: ignore[arg-type] + assert cfg["reasoning_type"] == "enabled" + + def test_thinking_shorthand_ignores_unknown_type(self): + _, cfg = _build_render_config(self._conv(), {"thinking": {"type": "auto"}}) # type: ignore[arg-type] + assert "reasoning_type" not in cfg + + def test_dev_instruction_forwarded(self): + _, cfg = _build_render_config(self._conv(), {"dev_instruction": "be brief"}) # type: ignore[arg-type] + assert cfg["dev_instruction"] == "be brief" + + def test_response_format_json_object_sets_json_mode(self): + _, cfg = _build_render_config( + self._conv(), {"response_format": {"type": "json_object"}} + ) # type: ignore[arg-type] + assert cfg["json_mode"] is True + assert "json_schema" not in cfg + + def test_response_format_json_schema_sets_json_schema(self): + schema = {"type": "object"} + _, cfg = _build_render_config( + self._conv(), + {"response_format": {"type": "json_schema", "schema": schema}}, + ) # type: ignore[arg-type] + # JSON-encoded for melody (string-only schema field). + assert cfg["json_schema"] == json.dumps(schema) + + def test_response_format_nested_json_schema_unwrapped(self): + # When the SDK shape is ``{type: json_schema, schema: {schema: + # {...}}}``, the inner ``schema`` value is used. + inner = {"type": "object"} + _, cfg = _build_render_config( + self._conv(), + { + "response_format": { + "type": "json_schema", + "schema": {"schema": inner}, + } + }, + ) # type: ignore[arg-type] + assert cfg["json_schema"] == json.dumps(inner) + + def test_json_schema_kwarg_direct(self): + # Caller can also pass ``json_schema`` directly, both as dict and + # as a pre-stringified value. + _, cfg = _build_render_config(self._conv(), {"json_schema": {"a": 1}}) # type: ignore[arg-type] + assert cfg["json_schema"] == '{"a": 1}' + _, cfg = _build_render_config( + self._conv(), {"json_schema": "raw-string-schema"} + ) # type: ignore[arg-type] + assert cfg["json_schema"] == "raw-string-schema" + + def test_json_mode_kwarg_overrides(self): + _, cfg = _build_render_config(self._conv(), {"json_mode": True}) # type: ignore[arg-type] + assert cfg["json_mode"] is True + + def test_cmd3_safety_mode_lowercased(self): + _, cfg = _build_render_config(self._conv(), {"safety_mode": "CONTEXTUAL"}) # type: ignore[arg-type] + assert cfg["safety_mode"] == "contextual" + + def test_cmd3_citation_quality_direct(self): + _, cfg = _build_render_config(self._conv(), {"citation_quality": "ACCURATE"}) # type: ignore[arg-type] + assert cfg["citation_quality"] == "accurate" + + def test_cmd3_citation_quality_derived_from_citation_options(self): + # When ``citation_quality`` is unset, ``citation_options.mode`` is + # collapsed to on/off so cmd3's binary toggle has a value. + _, cfg = _build_render_config( + self._conv(), {"citation_options": {"mode": "accurate"}} + ) # type: ignore[arg-type] + assert cfg["citation_quality"] == "on" + + _, cfg = _build_render_config( + self._conv(), {"citation_options": {"mode": "off"}} + ) # type: ignore[arg-type] + assert cfg["citation_quality"] == "off" + + def test_cmd3_skip_preamble_forwarded(self): + _, cfg = _build_render_config(self._conv(), {"skip_preamble": True}) # type: ignore[arg-type] + assert cfg["skip_preamble"] is True + + def test_cmd3_no_grounding_field(self): + # cmd3 should never emit a cmd4-only ``grounding`` field. + _, cfg = _build_render_config( + self._conv(), + {"cohere_format": "cmd3", "grounding": "fast"}, + ) # type: ignore[arg-type] + assert "grounding" not in cfg + + @pytest.mark.parametrize( + "raw,expected", + [ + ("FAST", "enabled"), + ("ACCURATE", "enabled"), + ("OFF", "disabled"), + ("enabled", "enabled"), + ("disabled", "disabled"), + ("unknown", "unknown"), + ], + ) + def test_cmd4_grounding_direct(self, raw, expected): + # melody's cmd4 only accepts ``unknown``/``enabled``/``disabled``, + # so the renderer normalizes any of the v2-facing values into + # that vocab. + _, cfg = _build_render_config( + self._conv(), + {"cohere_format": "cmd4", "grounding": raw}, + ) # type: ignore[arg-type] + assert cfg["grounding"] == expected + + @pytest.mark.parametrize( + "mode,expected", + [ + ("ACCURATE", "enabled"), + ("FAST", "enabled"), + ("OFF", "disabled"), + ], + ) + def test_cmd4_grounding_from_citation_options_mode(self, mode, expected): + _, cfg = _build_render_config( + self._conv(), + { + "cohere_format": "cmd4", + "citation_options": {"mode": mode}, + }, + ) # type: ignore[arg-type] + assert cfg["grounding"] == expected + + def test_cmd4_grounding_rejects_unknown_value(self): + with pytest.raises(ValueError, match="Unrecognized cmd4 grounding"): + _build_render_config( + self._conv(), + {"cohere_format": "cmd4", "grounding": "foobar"}, + ) # type: ignore[arg-type] + + def test_cmd4_platform_instruction(self): + _, cfg = _build_render_config( + self._conv(), + { + "cohere_format": "cmd4", + "platform_instruction": "do this", + }, + ) # type: ignore[arg-type] + assert cfg["platform_instruction"] == "do this" + + def test_cmd4_no_safety_mode_field(self): + # cmd4 should never carry cmd3-only ``safety_mode``/``citation_quality``. + _, cfg = _build_render_config( + self._conv(), + { + "cohere_format": "cmd4", + "safety_mode": "contextual", + "citation_quality": "on", + }, + ) # type: ignore[arg-type] + assert "safety_mode" not in cfg + assert "citation_quality" not in cfg + + def test_extra_kwargs_become_additional_template_fields(self): + # Anything not in the renderer's consumed-keys set is forwarded + # verbatim under ``additional_template_fields`` so jinja templates + # can resolve ``{{ var }}`` directly. + _, cfg = _build_render_config( + self._conv(), + { + "reasoning_effort": "low", + "my_var": "x", + "documents": ["doc"], # consumed, must NOT leak through + }, + ) # type: ignore[arg-type] + extras = cfg["additional_template_fields"] + assert extras == {"reasoning_effort": "low", "my_var": "x"} + # Sanity: the consumed key still produced its dedicated config slot. + assert cfg["documents"] == [{"text": "doc"}] + + def test_template_id_and_template_jinja_passthrough(self): + _, cfg = _build_render_config( + self._conv(), + { + "template_id": "tpl1", + "template_jinja": "raw {{ jinja }}", + }, + ) # type: ignore[arg-type] + assert cfg["template_id"] == "tpl1" + assert cfg["template_jinja"] == "raw {{ jinja }}" + # use_jinja is always True, regardless of caller input. + assert cfg["use_jinja"] is True + + +# ====================================================================== +# End-to-end async rendering (mirrors ``test_mistral.py``) +# ====================================================================== +# +# Verifies that the synchronous melody bindings run on the renderer's +# thread pool so the asyncio event loop stays responsive under +# concurrent load. Mirrors +# ``test_async_mistral_tokenizer_does_not_block_event_loop`` so future +# regressions in either path are caught uniformly. + + +@dataclass +class _MockHFConfig: + model_type: str = "any" + + +@dataclass +class _MockModelConfig: + runner_type = "generate" + model: str = "cohere-test" + tokenizer: str = "cohere-test" + trust_remote_code: bool = False + max_model_len: int = 100 + tokenizer_revision = None + tokenizer_mode = "cohere" + hf_config = _MockHFConfig() + hf_text_config = _MockHFConfig() + encoder_config: dict[str, Any] | None = None + enable_prompt_embeds: bool = True + skip_tokenizer_init: bool = True + is_encoder_decoder: bool = False + is_multimodal_model: bool = False + renderer_num_workers: int = 1 + + +@dataclass +class _MockParallelConfig: + _api_process_rank: int = 0 + + +@dataclass +class _MockVllmConfig: + model_config: _MockModelConfig + parallel_config: _MockParallelConfig + + +@pytest.mark.asyncio +async def test_async_cohere_renderer_does_not_block_event_loop(): + expected_prompt = "MOCK_RENDERED_PROMPT" + + def slow_render(*_a, **_kw): + time.sleep(2) + return expected_prompt + + mock_tokenizer = Mock(spec=HfTokenizer) + renderer = CohereRenderer( + _MockVllmConfig(_MockModelConfig(), _MockParallelConfig()), + tokenizer=mock_tokenizer, + ) + + # Replace the (already-imported) ``cohere_melody`` bindings with a + # blocking mock. ``_render`` reads ``self._melody`` at call time, so + # this works even though ``_render_async`` was bound at __init__. + fake_melody = Mock() + fake_melody.render_cmd3 = slow_render + fake_melody.render_cmd4 = slow_render + renderer._melody = fake_melody + + task = renderer.render_messages_async([], ChatParams()) + + # Ensure the event loop is not blocked while the (blocking) render + # call is in flight on the thread pool. + blocked_count = 0 + for _ in range(20): # ~2 seconds at 0.1s slices + start = time.perf_counter() + await asyncio.sleep(0) + elapsed = time.perf_counter() - start + if elapsed >= 0.5: + blocked_count += 1 + await asyncio.sleep(0.1) + + _, prompt = await task + assert prompt["prompt"] == expected_prompt, "Mocked blocking render was not called" + assert blocked_count == 0, "Event loop blocked during rendering" From 2304ce4c7e7837d568043dab6ebd5e2ca2e7f4ca Mon Sep 17 00:00:00 2001 From: Andrew Berneshawi Date: Tue, 30 Jun 2026 10:50:34 -0400 Subject: [PATCH 29/30] Add citation output handling --- .../openai/chat_completion/serving.py | 17 ++++ .../cohere_command_reasoning_parser.py | 80 ++++++++++++++++++- 2 files changed, 96 insertions(+), 1 deletion(-) diff --git a/vllm/entrypoints/openai/chat_completion/serving.py b/vllm/entrypoints/openai/chat_completion/serving.py index 284b2511dbaa..f613f3e1ac28 100644 --- a/vllm/entrypoints/openai/chat_completion/serving.py +++ b/vllm/entrypoints/openai/chat_completion/serving.py @@ -863,10 +863,23 @@ async def chat_completion_full_generator( ) if not request.include_reasoning: reasoning = None + # Reasoning parsers that extract grounding citations + # (Cohere Command family) cache them on the parser + # instance after ``parse``. We surface them here so + # grounded surfaces (e.g. /cohere/v2/chat) can attach + # them to the response message. Non-citation parsers + # leave the attribute absent and ``citations`` stays + # ``None`` so it round-trips out of the OpenAI envelope. + citations = getattr( + getattr(parser, "reasoning_parser", None), + "last_unary_citations", + None, + ) else: reasoning = None content = output.text tool_calls = [] + citations = None auto_tools_called = False is_named_tool_choice = ( @@ -930,6 +943,10 @@ async def chat_completion_full_generator( "completion." ) message = ChatMessage(role=role, reasoning=reasoning, content=content) + + if citations: + message.citations = citations + # In OpenAI's API, when a tool is called, the finish_reason is: # "tool_calls" for "auto" or "required" tool calls, # and "stop" for named tool calls. diff --git a/vllm/reasoning/cohere_command_reasoning_parser.py b/vllm/reasoning/cohere_command_reasoning_parser.py index 34066ef2d922..2e1a60ba007d 100644 --- a/vllm/reasoning/cohere_command_reasoning_parser.py +++ b/vllm/reasoning/cohere_command_reasoning_parser.py @@ -25,6 +25,8 @@ ) from vllm.entrypoints.openai.engine.protocol import ( AnyResponseFormat, + Citation, + CitationSource, DeltaFunctionCall, DeltaMessage, DeltaToolCall, @@ -404,6 +406,59 @@ def _schema_dict_from_structured_outputs( ) +def _melody_sources_to_vllm(raw_sources: Any) -> list[CitationSource]: + """Convert melody's ``Source`` objects into :class:`CitationSource`. + + melody's ``Source`` shape is ``{tool_call_index, tool_result_indices, + document_ids}``. ``document_ids`` may not be set; if it is empty and + we have no resolvable identifier then we fall back to a generic + ``tool``-style source carrying the tool-call index for visibility. + """ + out: list[CitationSource] = [] + for s in raw_sources or []: + # TODO Verify the tool vs doc logic + doc_ids: list[str] = list(getattr(s, "document_ids", None) or []) + if doc_ids: + for did in doc_ids: + if did: + out.append(CitationSource(type="document", id=did)) + continue + tool_call_index = getattr(s, "tool_call_index", None) + out.append( + CitationSource( + type="tool", + id=( + str(tool_call_index) + if tool_call_index is not None + else None + ), + ) + ) + return out + + +def _melody_citations_to_vllm(raw_citations: Any) -> list[Citation] | None: + """Convert melody's ``FilterCitation`` objects into :class:`Citation`.""" + if not raw_citations: + return None + out: list[Citation] = [] + for c in raw_citations: + out.append( + Citation( + start=getattr(c, "start_index", None), + end=getattr(c, "end_index", None), + text=getattr(c, "text", None), + sources=_melody_sources_to_vllm(getattr(c, "sources", None)), + type=( + "THINKING_CONTENT" + if getattr(c, "is_thinking", False) + else "TEXT_CONTENT" + ), + ) + ) + return out + + class BaseCohereCommandReasoningParser(ReasoningParser): def __init__( self, @@ -418,6 +473,13 @@ def __init__( self.unary_opts = unary_opts self.melody_unary = PyFilter(unary_opts) self.melody_streaming = PyFilter(streaming_opts) + # Citations extracted by the most recent ``extract_reasoning`` call. + # The non-streaming chat-completion path reads this back from the + # parser instance (which is constructed per-request) and attaches + # the result to ``ChatMessage.citations`` so grounded surfaces + # (e.g. ``/cohere/v2/chat``) can surface them. ``None`` when the + # last parse produced no citations. + self.last_unary_citations: list[Citation] | None = None @property def reasoning_start_str(self) -> str | None: @@ -437,7 +499,13 @@ def extract_reasoning_streaming( delta_token_ids: Sequence[int], ) -> DeltaMessage | None: r = self.melody_streaming.write_decoded(delta_text) - if r.content is None and r.reasoning is None and not r.tool_calls: + citations = _melody_citations_to_vllm(getattr(r, "citations", None)) + if ( + r.content is None + and r.reasoning is None + and not r.tool_calls + and not citations + ): return None msg = DeltaMessage() if r.content is not None: @@ -454,12 +522,22 @@ def extract_reasoning_streaming( ) for tc in r.tool_calls ] + if citations: + msg.citations = citations return msg def extract_reasoning( self, model_output: str, request: ChatCompletionRequest | ResponsesRequest ) -> tuple[str | None, str | None]: result = self.melody_unary.process_full_text(model_output) + # Cache citations so the non-streaming chat-completion path can + # surface them on ``ChatMessage.citations`` (the ``parse`` return + # tuple is locked to ``(reasoning, content, tool_calls)`` across + # all parsers, so we ferry citations via parser-instance state -- + # safe because the parser is constructed per-request). + self.last_unary_citations = _melody_citations_to_vllm( + getattr(result, "citations", None) + ) return result.reasoning, result.content def extract_content_ids(self, input_ids: list[int]) -> list[int]: From 00fe9eff777b332aaf3eff5c283522926e1fa4fb Mon Sep 17 00:00:00 2001 From: Andrew Berneshawi Date: Tue, 30 Jun 2026 11:11:13 -0400 Subject: [PATCH 30/30] Attempt fix to streaming citations --- .../cohere_command_reasoning_parser.py | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/vllm/reasoning/cohere_command_reasoning_parser.py b/vllm/reasoning/cohere_command_reasoning_parser.py index 2e1a60ba007d..2dea67b3c9b2 100644 --- a/vllm/reasoning/cohere_command_reasoning_parser.py +++ b/vllm/reasoning/cohere_command_reasoning_parser.py @@ -611,12 +611,25 @@ def adjust_request( return request +# melody's streaming filter only buffers a partial ```` citation +# across ``write_decoded`` calls when ``stream_non_grounded_answer`` is +# set: otherwise, the moment an opening ```` in the same delta, the filter emits the partial +# marker bytes verbatim as plain content. In vLLM's streaming path the +# parser is fed one token (1-4 chars) per call, so an unbuffered filter +# will leak ````-style markers into ``delta.content`` and never +# emit a ``FilterCitation`` for them. Enabling the flag flips the +# partial-match branch in melody's ``parse_citations`` (see +# ``src/parsing/citations_filter.rs``) to ``return (None, 0)`` -- i.e. +# keep buffering -- which lets a full citation eventually resolve. +# Non-streaming (unary) parsing receives the whole output in one call so +# the flag is a no-op there and we leave ``unary_opts`` alone. class CohereCommand3ReasoningParser(BaseCohereCommandReasoningParser): def __init__(self, tokenizer: TokenizerLike, *args, **kwargs): super().__init__( tokenizer, *args, - streaming_opts=PyFilterOptions().cmd3(), + streaming_opts=PyFilterOptions().cmd3().stream_non_grounded_answer(), unary_opts=PyFilterOptions().cmd3().no_tools(), **kwargs, ) @@ -627,7 +640,7 @@ def __init__(self, tokenizer: TokenizerLike, *args, **kwargs): super().__init__( tokenizer, *args, - streaming_opts=PyFilterOptions().cmd4(), + streaming_opts=PyFilterOptions().cmd4().stream_non_grounded_answer(), unary_opts=PyFilterOptions().cmd4().no_tools(), **kwargs, )