| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Cohere Chat v2 API protocol. |
| 4 | |
| 5 | The bulk of the wire types come straight from the official ``cohere`` |
| 6 | Python SDK so we stay in sync with the upstream specification and |
| 7 | avoid re-declaring the schema. We only own three things locally: |
| 8 | |
| 9 | 1. The top-level request body model (the SDK doesn't ship one — its |
| 10 | ``ClientV2.chat`` takes the body as kwargs), with vLLM-specific |
| 11 | extensions (``kv_transfer_params`` / ``chat_template_kwargs``). |
| 12 | 2. The non-streaming response envelope (the SDK exposes the message |
| 13 | shape via :class:`AssistantMessageResponse` but no full response |
| 14 | wrapper). |
| 15 | 3. The streaming discriminated union (the SDK exports each event type |
| 16 | individually but not as a combined ``Annotated[Union[...], |
| 17 | discriminator]``). |
| 18 | |
| 19 | Importing this module pulls in the ``cohere`` package. The router that |
| 20 | mounts ``POST /cohere/v2/chat`` guards on that import succeeding so vLLM still |
| 21 | boots without the SDK installed. |
| 22 | |
| 23 | See https://docs.cohere.com/reference/chat for the upstream spec. |
| 24 | """ |
| 25 | |
| 26 | from __future__ import annotations |
| 27 | |
| 28 | from typing import Annotated, Any, Literal |
| 29 | |
| 30 | from cohere import types as _sdk |
| 31 | from cohere.types import ( |
| 32 | AssistantChatMessageV2, |
| 33 | AssistantMessageResponse, |
| 34 | ChatMessageV2, |
| 35 | ChatRequestSafetyMode, |
| 36 | Citation, |
| 37 | CitationOptions, |
| 38 | Document, |
| 39 | ResponseFormatV2, |
| 40 | SystemChatMessageV2, |
| 41 | Thinking, |
| 42 | ToolCallV2, |
| 43 | ToolChatMessageV2, |
| 44 | ToolV2, |
| 45 | UserChatMessageV2, |
| 46 | ) |
| 47 | from pydantic import BaseModel, Field, field_validator |
| 48 | |
| 49 | import vllm.envs as envs |
| 50 | from vllm.exceptions import VLLMValidationError |
| 51 | |
| 52 | # Re-export the SDK wire-format types alongside our local extensions so |
| 53 | # ``vllm.entrypoints.cohere.serving`` and friends can import everything |
| 54 | # they need from this module. |
| 55 | __all__ = [ |
| 56 | "AssistantChatMessageV2", |
| 57 | "AssistantMessageResponse", |
| 58 | "ChatMessageV2", |
| 59 | "Citation", |
| 60 | "CitationEndEvent", |
| 61 | "CitationOptions", |
| 62 | "CitationStartEvent", |
| 63 | "CohereChatV2Request", |
| 64 | "CohereChatV2Response", |
| 65 | "CohereError", |
| 66 | "CohereFinishReason", |
| 67 | "CohereLogprobItem", |
| 68 | "CohereUsage", |
| 69 | "CohereUsageBilledUnits", |
| 70 | "CohereUsageTokens", |
| 71 | "ContentDeltaEvent", |
| 72 | "ContentEndEvent", |
| 73 | "ContentStartEvent", |
| 74 | "Document", |
| 75 | "MessageEndEvent", |
| 76 | "MessageStartEvent", |
| 77 | "ResponseFormatV2", |
| 78 | "SystemChatMessageV2", |
| 79 | "Thinking", |
| 80 | "ToolCallDeltaEvent", |
| 81 | "ToolCallEndEvent", |
| 82 | "ToolCallStartEvent", |
| 83 | "ToolCallV2", |
| 84 | "ToolChatMessageV2", |
| 85 | "ToolPlanDeltaEvent", |
| 86 | "ToolV2", |
| 87 | "UserChatMessageV2", |
| 88 | ] |
| 89 | |
| 90 | |
| 91 | # --------------------------------------------------------------------------- |
| 92 | # Errors |
| 93 | # --------------------------------------------------------------------------- |
| 94 | |
| 95 | |
| 96 | class CohereError(BaseModel): |
| 97 | """Top-level error body returned by ``/cohere/v2/chat`` error responses. |
| 98 | |
| 99 | Cohere's documented error schemas are uniform: ``{message, id}``. |
| 100 | """ |
| 101 | |
| 102 | message: str |
| 103 | id: str | None = None |
| 104 | |
| 105 | |
| 106 | # --------------------------------------------------------------------------- |
| 107 | # Tool choice / finish reasons |
| 108 | # --------------------------------------------------------------------------- |
| 109 | # |
| 110 | # These literals aren't first-class enums in the SDK but are documented at |
| 111 | # https://docs.cohere.com/reference/chat. We declare them here so the |
| 112 | # request/response models can validate them. |
| 113 | |
| 114 | CohereToolChoice = Literal["REQUIRED", "NONE"] |
| 115 | |
| 116 | CohereFinishReason = Literal[ |
| 117 | "COMPLETE", |
| 118 | "STOP_SEQUENCE", |
| 119 | "MAX_TOKENS", |
| 120 | "TOOL_CALL", |
| 121 | "ERROR", |
| 122 | "TIMEOUT", |
| 123 | ] |
| 124 | |
| 125 | |
| 126 | # --------------------------------------------------------------------------- |
| 127 | # Request |
| 128 | # --------------------------------------------------------------------------- |
| 129 | |
| 130 | |
| 131 | class CohereChatV2Request(BaseModel): |
| 132 | """Cohere Chat v2 request body. |
| 133 | |
| 134 | Mirrors the schema documented at https://docs.cohere.com/reference/chat. |
| 135 | All structured fields delegate to the official SDK types so the body |
| 136 | schema stays in sync with the upstream spec. |
| 137 | """ |
| 138 | |
| 139 | model: str |
| 140 | messages: list[ChatMessageV2] |
| 141 | stream: bool | None = False |
| 142 | |
| 143 | # Tooling |
| 144 | tools: list[ToolV2] | None = None |
| 145 | strict_tools: bool | None = None |
| 146 | tool_choice: CohereToolChoice | None = None |
| 147 | |
| 148 | # Grounding |
| 149 | documents: list[str | Document] | None = None |
| 150 | citation_options: CitationOptions | None = None |
| 151 | |
| 152 | # Output |
| 153 | response_format: ResponseFormatV2 | None = None |
| 154 | safety_mode: ChatRequestSafetyMode | None = None |
| 155 | max_tokens: int | None = None |
| 156 | stop_sequences: ( |
| 157 | Annotated[list[str], Field(max_length=envs.VLLM_MAX_STOP_STRINGS)] | None |
| 158 | ) = None |
| 159 | |
| 160 | # Sampling |
| 161 | temperature: float | None = None |
| 162 | seed: int | None = None |
| 163 | frequency_penalty: float | None = None |
| 164 | presence_penalty: float | None = None |
| 165 | k: int | None = None |
| 166 | p: float | None = None |
| 167 | logprobs: bool | None = None |
| 168 | |
| 169 | # Reasoning |
| 170 | thinking: Thinking | None = None |
| 171 | |
| 172 | # Scheduling |
| 173 | priority: int | None = Field(default=None, ge=-(2**63), le=2**63 - 1) |
| 174 | |
| 175 | # vLLM-specific extensions (not in Cohere spec). These mirror what the |
| 176 | # Anthropic and OpenAI surfaces already expose so V2 callers can reach |
| 177 | # the same engine knobs when needed. |
| 178 | kv_transfer_params: dict[str, Any] | None = Field( |
| 179 | default=None, |
| 180 | description="KVTransfer parameters used for disaggregated serving.", |
| 181 | ) |
| 182 | chat_template_kwargs: dict[str, Any] | None = Field( |
| 183 | default=None, |
| 184 | description=( |
| 185 | "Additional keyword args to pass to the chat template renderer. " |
| 186 | "Will be accessible by the template." |
| 187 | ), |
| 188 | ) |
| 189 | |
| 190 | @field_validator("model") |
| 191 | @classmethod |
| 192 | def _validate_model(cls, v: str) -> str: |
| 193 | if not v: |
| 194 | raise VLLMValidationError("model is required", parameter="model") |
| 195 | return v |
| 196 | |
| 197 | @field_validator("max_tokens") |
| 198 | @classmethod |
| 199 | def _validate_max_tokens(cls, v: int | None) -> int | None: |
| 200 | if v is not None and v < 0: |
| 201 | raise VLLMValidationError( |
| 202 | "max_tokens must be non-negative", |
| 203 | parameter="max_tokens", |
| 204 | value=v, |
| 205 | ) |
| 206 | return v |
| 207 | |
| 208 | @field_validator("messages", mode="before") |
| 209 | @classmethod |
| 210 | def _normalize_message_roles(cls, v: Any) -> Any: |
| 211 | """Rewrite OpenAI-style ``developer`` roles to ``system``. |
| 212 | |
| 213 | Cohere's v2 ``ChatMessageV2`` discriminated union only admits |
| 214 | the four literal roles ``user`` / ``assistant`` / ``system`` / |
| 215 | ``tool``. OpenAI's ``developer`` role is documented as |
| 216 | high-priority system instructions, so we alias it onto |
| 217 | ``system`` *before* the SDK discriminator runs (otherwise it |
| 218 | rejects the message with a ``literal_error`` against each union |
| 219 | member). Mirrors ``_role_to_melody`` in the renderer so the |
| 220 | same alias is honoured no matter which surface the message |
| 221 | arrives through. |
| 222 | |
| 223 | ``mode="before"`` is required so the rewrite happens before the |
| 224 | ``list[ChatMessageV2]`` coercion runs the SDK's discriminated |
| 225 | union; a default-mode validator would never see ``developer`` |
| 226 | because validation would have already failed. On any |
| 227 | structural malformation (non-iterable input, items without a |
| 228 | dict-shaped ``role`` field, etc.) we hand ``v`` back unchanged |
| 229 | and let Pydantic's normal coercion surface a precise error. |
| 230 | """ |
| 231 | try: |
| 232 | return [ |
| 233 | {**msg, "role": "system"} |
| 234 | if msg.get("role", "").lower() == "developer" |
| 235 | else msg |
| 236 | for msg in v |
| 237 | ] |
| 238 | except (AttributeError, TypeError): |
| 239 | return v |
| 240 | |
| 241 | @field_validator("messages") |
| 242 | @classmethod |
| 243 | def _validate_messages(cls, v: list[ChatMessageV2]) -> list[ChatMessageV2]: |
| 244 | if not v: |
| 245 | raise VLLMValidationError( |
| 246 | "messages must contain at least one message", parameter="messages" |
| 247 | ) |
| 248 | return v |
| 249 | |
| 250 | |
| 251 | # --------------------------------------------------------------------------- |
| 252 | # Usage / Logprobs |
| 253 | # --------------------------------------------------------------------------- |
| 254 | # |
| 255 | # The Cohere SDK only exposes a v1 ``ApiMetaBilledUnits``. The v2 usage |
| 256 | # envelope is documented separately in the OpenAPI spec and we declare it |
| 257 | # here. |
| 258 | |
| 259 | |
| 260 | class CohereUsageBilledUnits(BaseModel): |
| 261 | input_tokens: float | None = None |
| 262 | output_tokens: float | None = None |
| 263 | search_units: float | None = None |
| 264 | classifications: float | None = None |
| 265 | |
| 266 | |
| 267 | class CohereUsageTokens(BaseModel): |
| 268 | input_tokens: float | None = None |
| 269 | output_tokens: float | None = None |
| 270 | |
| 271 | |
| 272 | class CohereUsage(BaseModel): |
| 273 | billed_units: CohereUsageBilledUnits | None = None |
| 274 | tokens: CohereUsageTokens | None = None |
| 275 | cached_tokens: float | None = None |
| 276 | |
| 277 | |
| 278 | class CohereLogprobItem(BaseModel): |
| 279 | text: str | None = None |
| 280 | token_ids: list[int] |
| 281 | logprobs: list[float] | None = None |
| 282 | |
| 283 | |
| 284 | # --------------------------------------------------------------------------- |
| 285 | # Non-streaming response |
| 286 | # --------------------------------------------------------------------------- |
| 287 | |
| 288 | |
| 289 | class CohereChatV2Response(BaseModel): |
| 290 | """Cohere Chat v2 non-streaming response body. |
| 291 | |
| 292 | Wraps the SDK :class:`AssistantMessageResponse` (the message shape) in |
| 293 | the documented v2 response envelope (``id``, ``finish_reason``, |
| 294 | ``usage``, ``logprobs``). The single constructor in |
| 295 | :class:`CohereServingChatV2._chat_completion_to_v2` is responsible for |
| 296 | supplying a non-empty ``id`` (falling back to a synthesized one if |
| 297 | the upstream response is missing it) to this model. |
| 298 | """ |
| 299 | |
| 300 | id: str |
| 301 | finish_reason: CohereFinishReason |
| 302 | message: AssistantMessageResponse |
| 303 | usage: CohereUsage | None = None |
| 304 | logprobs: list[CohereLogprobItem] | None = None |
| 305 | |
| 306 | # vLLM-specific extension. |
| 307 | kv_transfer_params: dict[str, Any] | None = Field( |
| 308 | default=None, description="KVTransfer parameters." |
| 309 | ) |
| 310 | |
| 311 | |
| 312 | # --------------------------------------------------------------------------- |
| 313 | # Streaming events |
| 314 | # --------------------------------------------------------------------------- |
| 315 | # |
| 316 | # Cohere V2 streams a sequence of typed JSON events delivered as Server- |
| 317 | # Sent Events; each event's ``type`` field carries the discriminator. The |
| 318 | # SDK exposes a Pydantic model per event but none of them declare ``type`` |
| 319 | # as a field (the SDK relies on its own deserializer for discrimination), |
| 320 | # so a naive ``model_dump_json()`` would silently drop the discriminator |
| 321 | # and break clients that demux on ``type``. |
| 322 | # |
| 323 | # We therefore subclass each SDK event and bake the wire-format ``type`` |
| 324 | # string in as a ``Literal`` field with a default. ``model_dump()`` now |
| 325 | # emits ``type`` for free, and surfaces can simply construct the event |
| 326 | # class and serialize it -- no manual ``type`` parameter needed. |
| 327 | # |
| 328 | # See https://docs.cohere.com/v2/docs/streaming and the OpenAPI |
| 329 | # ``StreamedChatResponseV2`` schema for the wire-format reference. |
| 330 | |
| 331 | |
| 332 | class MessageStartEvent(_sdk.ChatMessageStartEvent): |
| 333 | type: Literal["message-start"] = "message-start" |
| 334 | |
| 335 | |
| 336 | class ContentStartEvent(_sdk.ChatContentStartEvent): |
| 337 | type: Literal["content-start"] = "content-start" |
| 338 | |
| 339 | |
| 340 | class ContentDeltaEvent(_sdk.ChatContentDeltaEvent): |
| 341 | type: Literal["content-delta"] = "content-delta" |
| 342 | |
| 343 | |
| 344 | class ContentEndEvent(_sdk.ChatContentEndEvent): |
| 345 | type: Literal["content-end"] = "content-end" |
| 346 | |
| 347 | |
| 348 | class ToolPlanDeltaEvent(_sdk.ChatToolPlanDeltaEvent): |
| 349 | type: Literal["tool-plan-delta"] = "tool-plan-delta" |
| 350 | |
| 351 | |
| 352 | class ToolCallStartEvent(_sdk.ChatToolCallStartEvent): |
| 353 | type: Literal["tool-call-start"] = "tool-call-start" |
| 354 | |
| 355 | |
| 356 | class ToolCallDeltaEvent(_sdk.ChatToolCallDeltaEvent): |
| 357 | type: Literal["tool-call-delta"] = "tool-call-delta" |
| 358 | |
| 359 | |
| 360 | class ToolCallEndEvent(_sdk.ChatToolCallEndEvent): |
| 361 | type: Literal["tool-call-end"] = "tool-call-end" |
| 362 | |
| 363 | |
| 364 | class CitationStartEvent(_sdk.CitationStartEvent): |
| 365 | type: Literal["citation-start"] = "citation-start" |
| 366 | |
| 367 | |
| 368 | class CitationEndEvent(_sdk.CitationEndEvent): |
| 369 | type: Literal["citation-end"] = "citation-end" |
| 370 | |
| 371 | |
| 372 | class MessageEndEvent(_sdk.ChatMessageEndEvent): |
| 373 | type: Literal["message-end"] = "message-end" |