| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | # Adapted from |
| 4 | # https://github.com/vllm-project/vllm/blob/main/vllm/entrypoints/anthropic/serving.py |
| 5 | """Cohere Chat v2 API serving handler. |
| 6 | |
| 7 | Implements ``POST /cohere/v2/chat`` by translating the incoming Cohere v2 request |
| 8 | into a standard :class:`ChatCompletionRequest` and delegating to |
| 9 | :class:`OpenAIServingChat`. The actual prompt rendering is handled by |
| 10 | vLLM's renderer pipeline (``vllm.renderers``): |
| 11 | |
| 12 | - For Cohere Command-family models, set ``--tokenizer-mode cohere`` and the |
| 13 | :class:`vllm.renderers.cohere.CohereRenderer` will template the request |
| 14 | via the ``cohere_melody`` library (``render_cmd3`` / ``render_cmd4``) |
| 15 | and surface citations through the Cohere-scoped |
| 16 | :class:`CohereChatMessage` / :class:`CohereDeltaMessage` subclasses |
| 17 | (see :mod:`vllm.entrypoints.cohere.cohere_chat_message`). |
| 18 | - For any other model the default Jinja-based :class:`HfRenderer` is used, |
| 19 | and the endpoint behaves as a plain v2-shaped wrapper around chat |
| 20 | completions. |
| 21 | """ |
| 22 | |
| 23 | from __future__ import annotations |
| 24 | |
| 25 | import json |
| 26 | import logging |
| 27 | import time |
| 28 | from collections.abc import AsyncGenerator |
| 29 | from enum import Enum |
| 30 | from http import HTTPStatus |
| 31 | from typing import TYPE_CHECKING, Any, NamedTuple, cast |
| 32 | |
| 33 | from cohere.types import ( |
| 34 | AssistantChatMessageV2, |
| 35 | AssistantMessageResponse, |
| 36 | Citation, |
| 37 | Document, |
| 38 | SystemChatMessageV2, |
| 39 | ToolCallV2, |
| 40 | ToolCallV2Function, |
| 41 | ToolChatMessageV2, |
| 42 | UserChatMessageV2, |
| 43 | ) |
| 44 | from fastapi import Request |
| 45 | from pydantic import BaseModel |
| 46 | |
| 47 | from vllm.engine.protocol import EngineClient |
| 48 | from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption |
| 49 | from vllm.entrypoints.cohere.cohere_chat_message import ( |
| 50 | CitationSource as InternalCitationSource, |
| 51 | ) |
| 52 | from vllm.entrypoints.cohere.cohere_chat_message import ( |
| 53 | CohereChatMessage, |
| 54 | ) |
| 55 | from vllm.entrypoints.cohere.protocol import ( |
| 56 | CitationEndEvent, |
| 57 | CitationStartEvent, |
| 58 | CohereChatV2Request, |
| 59 | CohereChatV2Response, |
| 60 | CohereFinishReason, |
| 61 | CohereUsage, |
| 62 | CohereUsageBilledUnits, |
| 63 | CohereUsageTokens, |
| 64 | ContentDeltaEvent, |
| 65 | ContentEndEvent, |
| 66 | ContentStartEvent, |
| 67 | MessageEndEvent, |
| 68 | MessageStartEvent, |
| 69 | ToolCallDeltaEvent, |
| 70 | ToolCallEndEvent, |
| 71 | ToolCallStartEvent, |
| 72 | ToolPlanDeltaEvent, |
| 73 | ) |
| 74 | from vllm.entrypoints.generate.base.protocol import ( |
| 75 | JsonSchemaResponseFormat, |
| 76 | ResponseFormat, |
| 77 | StreamOptions, |
| 78 | ) |
| 79 | from vllm.entrypoints.openai.chat_completion.protocol import ( |
| 80 | ChatCompletionRequest, |
| 81 | ChatCompletionResponse, |
| 82 | ChatCompletionStreamResponse, |
| 83 | ChatCompletionToolsParam, |
| 84 | ChatMessage, |
| 85 | ) |
| 86 | from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat |
| 87 | from vllm.entrypoints.openai.models.serving import OpenAIServingModels |
| 88 | from vllm.entrypoints.serve.engine.protocol import ErrorInfo, ErrorResponse |
| 89 | from vllm.entrypoints.serve.exception_handling.utils import sanitize_message |
| 90 | from vllm.entrypoints.serve.utils.request_logger import RequestLogger |
| 91 | from vllm.parser.abstract_parser import Parser |
| 92 | from vllm.renderers.cohere import ( |
| 93 | MESSAGES_CITATIONS_KEY, |
| 94 | POSITION_TO_SOURCE_KEY, |
| 95 | TOOL_MESSAGE_V2_CONTENT_KEY, |
| 96 | ) |
| 97 | |
| 98 | if TYPE_CHECKING: |
| 99 | from vllm.renderers.online_renderer import OnlineRenderer |
| 100 | |
| 101 | logger = logging.getLogger(__name__) |
| 102 | |
| 103 | |
| 104 | # Reserved ``chat_template_kwargs`` entries this handler populates for |
| 105 | # the API-server-side parser and renderer only. Stripped before the dict |
| 106 | # is forwarded to the engine core -- see |
| 107 | # :meth:`CohereServingChatV2._engine_chat_template_kwargs`. |
| 108 | _API_SERVER_ONLY_TEMPLATE_KWARGS = frozenset( |
| 109 | { |
| 110 | MESSAGES_CITATIONS_KEY, |
| 111 | POSITION_TO_SOURCE_KEY, |
| 112 | TOOL_MESSAGE_V2_CONTENT_KEY, |
| 113 | } |
| 114 | ) |
| 115 | |
| 116 | |
| 117 | # --------------------------------------------------------------------------- |
| 118 | # SSE helpers |
| 119 | # --------------------------------------------------------------------------- |
| 120 | |
| 121 | |
| 122 | def _sse(data: str) -> str: |
| 123 | """Wrap a JSON payload in a Server-Sent Event frame. |
| 124 | |
| 125 | Cohere's stream uses bare ``data:`` lines (no ``event:`` prefix); each |
| 126 | JSON object's ``type`` field carries the event discriminator. |
| 127 | """ |
| 128 | return f"data: {data}\n\n" |
| 129 | |
| 130 | |
| 131 | # Cohere v2 SSE stream terminator. Declared by the upstream OpenAPI spec |
| 132 | # (``x-fern-streaming.terminator: "[DONE]"`` on the ``/v2/chat`` stream |
| 133 | # operation) and observed by the cohere-python / Fern-generated clients |
| 134 | # (the Python SDK breaks its read loop on ``_sse.data == "[DONE]"``). |
| 135 | # Must be emitted after the closing ``message-end`` event. |
| 136 | _DONE_FRAME = _sse("[DONE]") |
| 137 | |
| 138 | |
| 139 | def _emit(event: BaseModel) -> str: |
| 140 | """Serialize a typed stream event into an SSE frame. |
| 141 | |
| 142 | The typed event classes in ``vllm.entrypoints.cohere.protocol`` |
| 143 | (``MessageStartEvent``, ``ContentStartEvent``, ``CitationStartEvent``, |
| 144 | ...) bake the wire-format ``type`` field into the model definition, |
| 145 | so a plain ``model_dump_json()`` already carries the discriminator. |
| 146 | """ |
| 147 | return _sse(event.model_dump_json(exclude_none=True)) |
| 148 | |
| 149 | |
| 150 | class ContentBlockType(str, Enum): |
| 151 | """Wire-format / internal discriminator for chat content blocks. |
| 152 | |
| 153 | ``THINKING`` and ``TEXT`` are the documented Cohere v2 content-block |
| 154 | type discriminators on the wire. ``TOOL_CALL`` is reserved for the |
| 155 | internal stream state machine (see :class:`_StreamState`) when a |
| 156 | tool call is the currently open block; it is never serialized. |
| 157 | """ |
| 158 | |
| 159 | THINKING = "thinking" |
| 160 | TEXT = "text" |
| 161 | TOOL_CALL = "tool_call" |
| 162 | |
| 163 | |
| 164 | # Mapping of vLLM/OpenAI finish reasons to Cohere's enum. |
| 165 | _FINISH_REASON_MAP: dict[str | None, CohereFinishReason] = { |
| 166 | "stop": "COMPLETE", |
| 167 | "length": "MAX_TOKENS", |
| 168 | "tool_calls": "TOOL_CALL", |
| 169 | "stop_sequence": "STOP_SEQUENCE", |
| 170 | "error": "ERROR", |
| 171 | None: "COMPLETE", |
| 172 | } |
| 173 | |
| 174 | |
| 175 | def _map_finish_reason( |
| 176 | reason: str | None, stop_reason: int | str | None = None |
| 177 | ) -> CohereFinishReason: |
| 178 | # vLLM reports both EOS and matched stop strings as finish_reason="stop". |
| 179 | # The latter is distinguished by the matched string in stop_reason. |
| 180 | if reason == "stop" and isinstance(stop_reason, str): |
| 181 | return "STOP_SEQUENCE" |
| 182 | return _FINISH_REASON_MAP.get(reason, "COMPLETE") |
| 183 | |
| 184 | |
| 185 | class _CitablePosition(NamedTuple): |
| 186 | """One entry in melody's citation numbering scheme, resolved to wire shape. |
| 187 | |
| 188 | Populated by :meth:`CohereServingChatV2._walk_citable_positions` |
| 189 | (the single source of truth for melody's numbering rule) and |
| 190 | consumed by both the inbound |
| 191 | (:meth:`CohereServingChatV2._build_doc_id_to_prompt_position`, |
| 192 | which drives citation resolution on the request side) and outbound |
| 193 | (:meth:`CohereServingChatV2._build_position_to_source`, whose |
| 194 | result is forwarded to the reasoning parser via |
| 195 | ``chat_template_kwargs[POSITION_TO_SOURCE_KEY]``) helpers -- so |
| 196 | the numbering rule only lives in one place. |
| 197 | |
| 198 | ``ids`` lists every id under which this position is addressable: |
| 199 | top-level documents contribute both their explicit id (if set) and |
| 200 | the ``doc_{idx}`` synthetic fallback, in that order so |
| 201 | ``setdefault``-style resolution prefers the explicit id; |
| 202 | tool-message documents contribute their single ``document.id``. |
| 203 | ``source`` is the fully-resolved wire-shape source (``type`` / |
| 204 | ``id`` / ``document`` or ``tool_output`` payload) for the |
| 205 | outbound direction. |
| 206 | """ |
| 207 | |
| 208 | bucket: int |
| 209 | result_idx: int |
| 210 | ids: tuple[str, ...] |
| 211 | source: InternalCitationSource |
| 212 | |
| 213 | |
| 214 | # --------------------------------------------------------------------------- |
| 215 | # Serving class |
| 216 | # --------------------------------------------------------------------------- |
| 217 | |
| 218 | |
| 219 | class CohereServingChatV2(OpenAIServingChat): |
| 220 | """Handler for the Cohere Chat v2 API (``POST /cohere/v2/chat``). |
| 221 | |
| 222 | The handler is intentionally thin: it converts the v2 request into a |
| 223 | :class:`ChatCompletionRequest` (preserving Cohere-specific fields such |
| 224 | as ``documents``, ``safety_mode`` and ``citation_options`` via |
| 225 | ``chat_template_kwargs``) and delegates to the underlying chat |
| 226 | completion machinery. All Cohere-specific templating happens in the |
| 227 | renderer (:class:`vllm.renderers.cohere.CohereRenderer`) when the |
| 228 | engine is started with ``--tokenizer-mode cohere``. |
| 229 | |
| 230 | Citations flow through ``CohereChatMessage`` / ``CohereDeltaMessage`` |
| 231 | subclasses (see :mod:`vllm.entrypoints.cohere.cohere_chat_message`) so |
| 232 | the OpenAI-shared :class:`ChatMessage` / :class:`DeltaMessage` keep |
| 233 | their declared schemas unchanged. Non-streaming responses pick up |
| 234 | the subclass by overriding :meth:`_create_chat_message`; streaming |
| 235 | deltas are emitted directly by the Cohere reasoning parser. |
| 236 | """ |
| 237 | |
| 238 | def __init__( |
| 239 | self, |
| 240 | engine_client: EngineClient, |
| 241 | models: OpenAIServingModels, |
| 242 | response_role: str, |
| 243 | *, |
| 244 | online_renderer: OnlineRenderer, |
| 245 | request_logger: RequestLogger | None, |
| 246 | chat_template: str | None, |
| 247 | chat_template_content_format: ChatTemplateContentFormatOption, |
| 248 | return_tokens_as_token_ids: bool = False, |
| 249 | reasoning_parser: str = "", |
| 250 | enable_auto_tools: bool = False, |
| 251 | tool_parser: str | None = None, |
| 252 | tool_strict_level: str = "auto", |
| 253 | enable_prompt_tokens_details: bool = False, |
| 254 | enable_force_include_usage: bool = False, |
| 255 | default_chat_template_kwargs: dict[str, Any] | None = None, |
| 256 | is_reasoning_model: bool = True, |
| 257 | ) -> None: |
| 258 | super().__init__( |
| 259 | engine_client=engine_client, |
| 260 | models=models, |
| 261 | response_role=response_role, |
| 262 | online_renderer=online_renderer, |
| 263 | request_logger=request_logger, |
| 264 | chat_template=chat_template, |
| 265 | chat_template_content_format=chat_template_content_format, |
| 266 | return_tokens_as_token_ids=return_tokens_as_token_ids, |
| 267 | reasoning_parser=reasoning_parser, |
| 268 | enable_auto_tools=enable_auto_tools, |
| 269 | tool_strict_level=tool_strict_level, |
| 270 | tool_parser=tool_parser, |
| 271 | enable_prompt_tokens_details=enable_prompt_tokens_details, |
| 272 | enable_force_include_usage=enable_force_include_usage, |
| 273 | default_chat_template_kwargs=default_chat_template_kwargs, |
| 274 | ) |
| 275 | # Controls how the assistant's chain-of-thought is surfaced on |
| 276 | # turns that also contain tool calls. |
| 277 | # |
| 278 | # - ``True`` (default): the model is a reasoning Command-family |
| 279 | # model; reasoning is always surfaced as a ``thinking`` content |
| 280 | # block (non-streaming) or as ``content-start`` / ``content- |
| 281 | # delta`` events for a thinking block (streaming), regardless of |
| 282 | # whether tool calls also appear. |
| 283 | # - ``False``: the model is an older non-reasoning Command model |
| 284 | # that uses Cohere's ``tool_plan`` field for its chain-of- |
| 285 | # thought before tool calls; reasoning is surfaced as |
| 286 | # ``tool_plan`` (non-streaming) or as ``tool-plan-delta`` events |
| 287 | # (streaming) on tool-call turns, and the thinking content block |
| 288 | # is dropped. |
| 289 | # |
| 290 | # TODO: replace this manual flag with automatic detection from the |
| 291 | # model's capabilities config once that exists. |
| 292 | self._is_reasoning_model = is_reasoning_model |
| 293 | |
| 294 | # ------------------------------------------------------------------ |
| 295 | # Public entry point |
| 296 | # ------------------------------------------------------------------ |
| 297 | |
| 298 | async def create_chat_v2( |
| 299 | self, |
| 300 | request: CohereChatV2Request, |
| 301 | raw_request: Request | None = None, |
| 302 | ) -> AsyncGenerator[str, None] | CohereChatV2Response | ErrorResponse: |
| 303 | """Implements ``POST /cohere/v2/chat``.""" |
| 304 | if logger.isEnabledFor(logging.DEBUG): |
| 305 | logger.debug( |
| 306 | "Received Cohere v2 chat request %s", request.model_dump_json() |
| 307 | ) |
| 308 | |
| 309 | chat_req = self.to_chat_completion_request(request) |
| 310 | generator = await self.create_chat_completion(chat_req, raw_request) |
| 311 | |
| 312 | match generator: |
| 313 | case ErrorResponse(): |
| 314 | return generator |
| 315 | case ChatCompletionResponse(): |
| 316 | return self._chat_completion_to_v2(generator, request) |
| 317 | case _: |
| 318 | return self._chat_completion_stream_to_v2(generator, request) |
| 319 | |
| 320 | def to_chat_completion_request( |
| 321 | self, request: CohereChatV2Request |
| 322 | ) -> ChatCompletionRequest: |
| 323 | """Convert a Cohere v2 request into its ``ChatCompletionRequest``. |
| 324 | |
| 325 | Exposed for callers that need the converted request without |
| 326 | running generation - notably the render endpoint |
| 327 | (``POST /cohere/v2/chat/render``), which hands the result to |
| 328 | :meth:`vllm.entrypoints.scale_out.render.serving.ServingRender.render_chat_request` |
| 329 | so that Cohere requests tokenize through exactly the same path |
| 330 | as ``/v1/chat/completions/render``. |
| 331 | """ |
| 332 | chat_req = self._convert_v2_to_chat_completion(request) |
| 333 | if logger.isEnabledFor(logging.DEBUG): |
| 334 | logger.debug( |
| 335 | "Converted Cohere v2 -> ChatCompletion: %s", |
| 336 | chat_req.model_dump_json(), |
| 337 | ) |
| 338 | return chat_req |
| 339 | |
| 340 | def _engine_chat_template_kwargs( |
| 341 | self, chat_template_kwargs: dict[str, Any] |
| 342 | ) -> dict[str, Any]: |
| 343 | """Keep this handler's internal entries out of the engine hop. |
| 344 | |
| 345 | Overrides :meth:`OpenAIServingChat._engine_chat_template_kwargs`. |
| 346 | The reserved keys in :data:`_API_SERVER_ONLY_TEMPLATE_KWARGS` |
| 347 | exist only to hand request-side state to the API-server-side |
| 348 | parser and renderer; the engine-core-side reasoning parser (used |
| 349 | for structured-output reasoning-end gating) never reads them. They |
| 350 | also can't survive the trip: ``POSITION_TO_SOURCE_KEY`` holds |
| 351 | :class:`CitationSource` models, and msgpack -- which encodes |
| 352 | ``reasoning_parser_kwargs`` on the way to the engine core -- has |
| 353 | no native encoding for Pydantic models, so forwarding them fails |
| 354 | the request with a ``TypeError``. |
| 355 | """ |
| 356 | return { |
| 357 | k: v |
| 358 | for k, v in chat_template_kwargs.items() |
| 359 | if k not in _API_SERVER_ONLY_TEMPLATE_KWARGS |
| 360 | } |
| 361 | |
| 362 | # ================================================================== |
| 363 | # Request conversion: Cohere V2 -> ChatCompletionRequest |
| 364 | # ================================================================== |
| 365 | |
| 366 | @classmethod |
| 367 | def _convert_v2_to_chat_completion( |
| 368 | cls, request: CohereChatV2Request |
| 369 | ) -> ChatCompletionRequest: |
| 370 | openai_messages: list[dict[str, Any]] = [] |
| 371 | cls._convert_messages(request.messages, openai_messages) |
| 372 | |
| 373 | chat_req = cls._build_base_chat_completion(request, openai_messages) |
| 374 | cls._apply_streaming_options(chat_req, request) |
| 375 | cls._apply_response_format(chat_req, request) |
| 376 | cls._apply_tools(chat_req, request) |
| 377 | cls._apply_tool_choice(chat_req, request) |
| 378 | cls._apply_cohere_template_kwargs(chat_req, request) |
| 379 | return chat_req |
| 380 | |
| 381 | @classmethod |
| 382 | def _convert_messages( |
| 383 | cls, |
| 384 | messages: list, |
| 385 | openai_messages: list[dict[str, Any]], |
| 386 | ) -> None: |
| 387 | for msg in messages: |
| 388 | match msg: |
| 389 | case SystemChatMessageV2(): |
| 390 | openai_messages.append( |
| 391 | { |
| 392 | "role": "system", |
| 393 | "content": cls._coerce_text_content(msg.content), |
| 394 | } |
| 395 | ) |
| 396 | case UserChatMessageV2(): |
| 397 | openai_messages.append(cls._convert_user_message(msg)) |
| 398 | case AssistantChatMessageV2(): |
| 399 | openai_messages.append(cls._convert_assistant_message(msg)) |
| 400 | case ToolChatMessageV2(): |
| 401 | openai_messages.append(cls._convert_tool_message(msg)) |
| 402 | case _: # pragma: no cover - guarded by Pydantic discriminator |
| 403 | raise ValueError(f"Unsupported Cohere v2 message: {msg!r}") |
| 404 | |
| 405 | @staticmethod |
| 406 | def _coerce_text_content(content: str | list[Any]) -> str: |
| 407 | if isinstance(content, str): |
| 408 | return content |
| 409 | parts: list[str] = [] |
| 410 | for block in content: |
| 411 | text = getattr(block, "text", None) |
| 412 | if text: |
| 413 | parts.append(text) |
| 414 | return "".join(parts) |
| 415 | |
| 416 | @classmethod |
| 417 | def _convert_user_message(cls, msg: UserChatMessageV2) -> dict[str, Any]: |
| 418 | if isinstance(msg.content, str): |
| 419 | return {"role": "user", "content": msg.content} |
| 420 | |
| 421 | # Discriminate by ``type`` rather than isinstance so we don't have |
| 422 | # to import every individual ``*Content`` variant from the SDK - |
| 423 | # the union (``UserMessageV2Content``) covers both ``TextContent`` |
| 424 | # and ``ImageUrlContent`` and both expose ``type`` as a Literal. |
| 425 | content_parts: list[dict[str, Any]] = [] |
| 426 | for block in msg.content: |
| 427 | if block.type == "text": |
| 428 | content_parts.append({"type": "text", "text": block.text}) |
| 429 | elif block.type == "image_url": |
| 430 | image_url: dict[str, Any] = {"url": block.image_url.url} |
| 431 | if getattr(block.image_url, "detail", None) is not None: |
| 432 | image_url["detail"] = block.image_url.detail |
| 433 | content_parts.append( |
| 434 | { |
| 435 | "type": "image_url", |
| 436 | "image_url": image_url, |
| 437 | } |
| 438 | ) |
| 439 | if len(content_parts) == 1 and content_parts[0]["type"] == "text": |
| 440 | return {"role": "user", "content": content_parts[0]["text"]} |
| 441 | return {"role": "user", "content": content_parts} |
| 442 | |
| 443 | @classmethod |
| 444 | def _convert_assistant_message(cls, msg: AssistantChatMessageV2) -> dict[str, Any]: |
| 445 | out: dict[str, Any] = {"role": "assistant"} |
| 446 | |
| 447 | # Cohere splits reasoning out into ``thinking`` content blocks. We |
| 448 | # collapse them back into the OpenAI ``reasoning`` field for the |
| 449 | # downstream chat template, while text blocks become ``content``. |
| 450 | text_parts: list[str] = [] |
| 451 | thinking_parts: list[str] = [] |
| 452 | if isinstance(msg.content, str): |
| 453 | text_parts.append(msg.content) |
| 454 | elif msg.content is not None: |
| 455 | for block in msg.content: |
| 456 | if block.type == "text": |
| 457 | text_parts.append(block.text) |
| 458 | elif block.type == "thinking": |
| 459 | thinking_parts.append(block.thinking) |
| 460 | |
| 461 | # ``tool_plan`` is Cohere's chain-of-thought emitted alongside tool |
| 462 | # calls; preserve it as reasoning so templates that expect a |
| 463 | # planning block still see it. |
| 464 | if msg.tool_plan: |
| 465 | thinking_parts.append(msg.tool_plan) |
| 466 | |
| 467 | if text_parts: |
| 468 | out["content"] = "".join(text_parts) |
| 469 | if thinking_parts: |
| 470 | out["reasoning"] = "".join(thinking_parts) |
| 471 | |
| 472 | if msg.tool_calls: |
| 473 | out["tool_calls"] = [ |
| 474 | { |
| 475 | "id": tc.id, |
| 476 | "type": "function", |
| 477 | "function": { |
| 478 | "name": (tc.function.name if tc.function else "") or "", |
| 479 | "arguments": (tc.function.arguments if tc.function else None) |
| 480 | or "{}", |
| 481 | }, |
| 482 | } |
| 483 | for tc in msg.tool_calls |
| 484 | ] |
| 485 | return out |
| 486 | |
| 487 | @classmethod |
| 488 | def _convert_tool_message(cls, msg: ToolChatMessageV2) -> dict[str, Any]: |
| 489 | if isinstance(msg.content, str): |
| 490 | return { |
| 491 | "role": "tool", |
| 492 | "tool_call_id": msg.tool_call_id, |
| 493 | "content": msg.content, |
| 494 | } |
| 495 | |
| 496 | # Map every content block to a plain ``{"type": "text", "text": |
| 497 | # ...}`` part, mirroring the shape :meth:`_convert_user_message` |
| 498 | # emits for multi-part user content. vLLM's |
| 499 | # ``_parse_chat_message_content`` collapses a text-only tool |
| 500 | # part list back into ``"\n".join(...)`` on the way to the |
| 501 | # renderer, so simple chat templates that expect string tool |
| 502 | # content still work; keeping the list here just preserves the |
| 503 | # per-block structure at the OpenAI-shape request level for |
| 504 | # symmetry with the user-message conversion. |
| 505 | # |
| 506 | # Tool message content uses ``ToolMessageV2Content``, a union of |
| 507 | # ``TextToolContent`` and ``DocumentToolContent``. The latter |
| 508 | # produces ``{"type": "document", "document": {...}}`` blocks, |
| 509 | # which vLLM's shared ``_parse_chat_message_content_parts`` in |
| 510 | # ``vllm/entrypoints/chat_utils.py`` rejects because the OpenAI |
| 511 | # ``ChatCompletionContentPartParam`` union only knows text / |
| 512 | # image / audio / etc. Emitting them anyway triggers a Pydantic |
| 513 | # ``ValidatorIterator`` failure at chat-parsing time, so we |
| 514 | # JSON-serialize each document into a text part instead. |
| 515 | # |
| 516 | # The cohere renderer doesn't consume this flattened content: |
| 517 | # :meth:`_collect_v2_tool_message_content` forwards the |
| 518 | # original v2 content blocks under |
| 519 | # :data:`TOOL_MESSAGE_V2_CONTENT_KEY`, and the renderer applies |
| 520 | # :func:`vllm.renderers.cohere._v2_tool_content_to_melody_block` |
| 521 | # to produce the melody-shape ``document`` blocks cmd3/cmd4 |
| 522 | # tool-result rendering expects. Outbound citation binding is |
| 523 | # unaffected either way: sources resolve through the |
| 524 | # ``POSITION_TO_SOURCE_KEY`` map (built from the original v2 |
| 525 | # messages, not the converted OpenAI messages). |
| 526 | content_parts: list[dict[str, Any]] = [] |
| 527 | for block in msg.content: |
| 528 | if block.type == "text": |
| 529 | content_parts.append({"type": "text", "text": block.text}) |
| 530 | elif block.type == "document": |
| 531 | content_parts.append( |
| 532 | { |
| 533 | "type": "text", |
| 534 | "text": json.dumps( |
| 535 | block.document.model_dump(exclude_none=True) |
| 536 | ), |
| 537 | } |
| 538 | ) |
| 539 | return { |
| 540 | "role": "tool", |
| 541 | "tool_call_id": msg.tool_call_id, |
| 542 | "content": content_parts, |
| 543 | } |
| 544 | |
| 545 | @classmethod |
| 546 | def _build_base_chat_completion( |
| 547 | cls, |
| 548 | request: CohereChatV2Request, |
| 549 | openai_messages: list[dict[str, Any]], |
| 550 | ) -> ChatCompletionRequest: |
| 551 | return ChatCompletionRequest( |
| 552 | model=request.model, |
| 553 | messages=openai_messages, |
| 554 | max_tokens=request.max_tokens, |
| 555 | max_completion_tokens=request.max_tokens, |
| 556 | stop=request.stop_sequences, |
| 557 | temperature=request.temperature, |
| 558 | top_p=request.p, |
| 559 | top_k=request.k, |
| 560 | seed=request.seed, |
| 561 | frequency_penalty=request.frequency_penalty, |
| 562 | presence_penalty=request.presence_penalty, |
| 563 | logprobs=request.logprobs, |
| 564 | priority=request.priority or 0, |
| 565 | kv_transfer_params=request.kv_transfer_params, |
| 566 | chat_template_kwargs=request.chat_template_kwargs, |
| 567 | ) |
| 568 | |
| 569 | @classmethod |
| 570 | def _apply_streaming_options( |
| 571 | cls, |
| 572 | chat_req: ChatCompletionRequest, |
| 573 | request: CohereChatV2Request, |
| 574 | ) -> None: |
| 575 | if request.stream: |
| 576 | chat_req.stream = True |
| 577 | chat_req.stream_options = StreamOptions.model_validate( |
| 578 | {"include_usage": True} |
| 579 | ) |
| 580 | |
| 581 | @classmethod |
| 582 | def _apply_response_format( |
| 583 | cls, |
| 584 | chat_req: ChatCompletionRequest, |
| 585 | request: CohereChatV2Request, |
| 586 | ) -> None: |
| 587 | rf = request.response_format |
| 588 | if rf is None or rf.type == "text": |
| 589 | return |
| 590 | chat_req.response_format = ResponseFormat( |
| 591 | type="json_schema" if rf.json_schema else "json_object", |
| 592 | json_schema=( |
| 593 | JsonSchemaResponseFormat( |
| 594 | name="cohere_v2_json_schema", |
| 595 | json_schema=rf.json_schema, |
| 596 | ) |
| 597 | if rf.json_schema |
| 598 | else None |
| 599 | ), |
| 600 | ) |
| 601 | |
| 602 | @classmethod |
| 603 | def _apply_tools( |
| 604 | cls, |
| 605 | chat_req: ChatCompletionRequest, |
| 606 | request: CohereChatV2Request, |
| 607 | ) -> None: |
| 608 | if not request.tools: |
| 609 | return |
| 610 | # Cohere's ``strict_tools`` is the spec-equivalent of OpenAI's |
| 611 | # per-function ``strict`` flag: when true the API guarantees tool |
| 612 | # call arguments match the declared JSON schema. We surface it in |
| 613 | # both places so OpenAI-shaped consumers see it on the function |
| 614 | # definition and the cohere renderer can read it back from |
| 615 | # ``chat_template_kwargs`` for cmd3/cmd4 preamble selection. |
| 616 | strict = bool(request.strict_tools) |
| 617 | chat_req.tools = [ |
| 618 | ChatCompletionToolsParam.model_validate( |
| 619 | { |
| 620 | "type": "function", |
| 621 | "function": { |
| 622 | "name": tool.function.name, |
| 623 | "description": tool.function.description, |
| 624 | "parameters": tool.function.parameters, |
| 625 | **({"strict": True} if strict else {}), |
| 626 | }, |
| 627 | } |
| 628 | ) |
| 629 | for tool in request.tools |
| 630 | ] |
| 631 | |
| 632 | @classmethod |
| 633 | def _apply_tool_choice( |
| 634 | cls, |
| 635 | chat_req: ChatCompletionRequest, |
| 636 | request: CohereChatV2Request, |
| 637 | ) -> None: |
| 638 | # Cohere v2's ``tool_choice`` only admits ``REQUIRED`` / ``NONE`` |
| 639 | # (see :data:`CohereToolChoice`); named-tool selection is not |
| 640 | # part of the v2 spec. Map both onto OpenAI's equivalents. |
| 641 | if request.tool_choice == "REQUIRED": |
| 642 | chat_req.tool_choice = "required" |
| 643 | elif request.tool_choice == "NONE": |
| 644 | chat_req.tool_choice = "none" |
| 645 | elif chat_req.tools: |
| 646 | # Mirrors Cohere's "free choice" default when tools are present. |
| 647 | chat_req.tool_choice = "auto" |
| 648 | |
| 649 | @classmethod |
| 650 | def _apply_cohere_template_kwargs( |
| 651 | cls, |
| 652 | chat_req: ChatCompletionRequest, |
| 653 | request: CohereChatV2Request, |
| 654 | ) -> None: |
| 655 | """Forward Cohere-specific request fields into ``chat_template_kwargs``. |
| 656 | |
| 657 | The :class:`vllm.renderers.cohere.CohereRenderer` consumes these |
| 658 | kwargs to drive ``cohere_melody.render_cmd3`` / ``render_cmd4``. |
| 659 | Other renderers ignore unknown kwargs, so this is also a no-op for |
| 660 | non-cohere ``--tokenizer-mode`` settings. |
| 661 | """ |
| 662 | kwargs = dict(chat_req.chat_template_kwargs or {}) |
| 663 | |
| 664 | if request.documents: |
| 665 | documents: list[dict[str, Any]] = [] |
| 666 | for idx, doc in enumerate(request.documents): |
| 667 | if isinstance(doc, str): |
| 668 | documents.append({"id": f"doc_{idx}", "data": {"text": doc}}) |
| 669 | else: |
| 670 | documents.append( |
| 671 | { |
| 672 | "id": doc.id or f"doc_{idx}", |
| 673 | "data": ( |
| 674 | doc.data |
| 675 | if isinstance(doc.data, dict) |
| 676 | else {"text": doc.data} |
| 677 | ), |
| 678 | } |
| 679 | ) |
| 680 | kwargs.setdefault("documents", documents) |
| 681 | |
| 682 | if request.safety_mode is not None: |
| 683 | kwargs.setdefault("safety_mode", str(request.safety_mode).lower()) |
| 684 | if request.citation_options is not None: |
| 685 | kwargs.setdefault( |
| 686 | "citation_options", |
| 687 | request.citation_options.model_dump(exclude_none=True), |
| 688 | ) |
| 689 | if request.thinking is not None: |
| 690 | kwargs.setdefault( |
| 691 | "thinking", |
| 692 | request.thinking.model_dump(exclude_none=True), |
| 693 | ) |
| 694 | |
| 695 | # ``strict_tools`` is intentionally NOT forwarded here. It's a |
| 696 | # decoder-guidance flag that ``_apply_tools`` already maps onto |
| 697 | # per-function OpenAI ``strict: true`` on ``chat_req.tools``; |
| 698 | # ``cohere_melody.render_cmd3``/``render_cmd4`` have no |
| 699 | # ``strict_tools`` knob, and the renderer's residual-forward path |
| 700 | # would otherwise surface it as a Jinja variable ``{{ strict_tools }}`` |
| 701 | # that templates have no defined use for. |
| 702 | |
| 703 | # Citations on assistant messages in the request don't fit |
| 704 | # OpenAI's ``ChatMessage`` / ``ConversationMessage`` schemas, so |
| 705 | # we forward them under the reserved ``MESSAGES_CITATIONS_KEY`` |
| 706 | # chat_template_kwargs entry keyed by request-message index. |
| 707 | # See ``_sdk_citation_to_melody`` for the id-resolution rules. |
| 708 | citations_by_index = cls._collect_message_citations(request) |
| 709 | if citations_by_index: |
| 710 | kwargs.setdefault(MESSAGES_CITATIONS_KEY, citations_by_index) |
| 711 | |
| 712 | # Forward the ``(bucket, result_idx) -> resolved wire source`` |
| 713 | # map to the reasoning parser via chat_template_kwargs. The |
| 714 | # parser (constructed with the same kwargs) uses it to attach |
| 715 | # real document ids / types / payloads to every source it |
| 716 | # emits, so citation deltas already carry SDK-shape sources by |
| 717 | # the time they reach the streaming loop. See |
| 718 | # ``_build_position_to_source`` for the numbering rule and |
| 719 | # ``_melody_sources_to_vllm`` in ``vllm.parser.cohere_command`` for |
| 720 | # the consumer. |
| 721 | position_to_source = cls._build_position_to_source(request) |
| 722 | if position_to_source: |
| 723 | kwargs.setdefault(POSITION_TO_SOURCE_KEY, position_to_source) |
| 724 | |
| 725 | # Forward the original v2 tool-message content blocks so the |
| 726 | # renderer can apply its v2 → melody mapping over them (the |
| 727 | # renderer owns the melody-shape rules, this handler only |
| 728 | # ships the raw structural info that OpenAI validation would |
| 729 | # otherwise destroy -- see the flatten comment in |
| 730 | # :meth:`_convert_tool_message`). |
| 731 | v2_tool_content = cls._collect_v2_tool_message_content(request) |
| 732 | if v2_tool_content: |
| 733 | kwargs.setdefault(TOOL_MESSAGE_V2_CONTENT_KEY, v2_tool_content) |
| 734 | |
| 735 | if kwargs: |
| 736 | chat_req.chat_template_kwargs = kwargs |
| 737 | |
| 738 | @classmethod |
| 739 | def _collect_message_citations( |
| 740 | cls, request: CohereChatV2Request |
| 741 | ) -> dict[int, list[dict[str, Any]]]: |
| 742 | """Extract ``AssistantChatMessageV2.citations`` keyed by message index. |
| 743 | |
| 744 | Returns an empty dict when no assistant message in the request |
| 745 | carries any resolvable citations, so callers can use truthiness |
| 746 | as the "should I emit the ``MESSAGES_CITATIONS_KEY`` entry?" |
| 747 | test. Citations whose sources can't be resolved are dropped |
| 748 | individually; messages that end up with zero surviving |
| 749 | citations are omitted from the result. |
| 750 | """ |
| 751 | doc_positions = cls._build_doc_id_to_prompt_position( |
| 752 | request.messages, request.documents |
| 753 | ) |
| 754 | out: dict[int, list[dict[str, Any]]] = {} |
| 755 | for idx, msg in enumerate(request.messages): |
| 756 | if not isinstance(msg, AssistantChatMessageV2) or not msg.citations: |
| 757 | continue |
| 758 | resolved: list[dict[str, Any]] = [] |
| 759 | for c in msg.citations: |
| 760 | converted = cls._sdk_citation_to_melody(c, doc_positions) |
| 761 | if converted is not None: |
| 762 | resolved.append(converted) |
| 763 | if resolved: |
| 764 | out[idx] = resolved |
| 765 | return out |
| 766 | |
| 767 | @classmethod |
| 768 | def _walk_citable_positions( |
| 769 | cls, messages: list, documents: list | None |
| 770 | ) -> list[_CitablePosition]: |
| 771 | """Single-pass implementation of melody's citation numbering rule. |
| 772 | |
| 773 | Mirrors ``PromptRenderIds::from_messages`` in |
| 774 | ``melody/src/templating/util.rs``. Emits one |
| 775 | :class:`_CitablePosition` per document-bearing slot melody |
| 776 | assigns while rendering the prompt: |
| 777 | |
| 778 | * **Bucket numbering.** Bucket ``0`` is reserved for the |
| 779 | top-level ``documents`` array when present, and each unique |
| 780 | ``tool_call_id`` claims the next integer on first-seen basis |
| 781 | -- registered from whichever of the assistant's |
| 782 | ``tool_calls[].id`` or the tool message's ``tool_call_id`` |
| 783 | appears first in the history. |
| 784 | * **Per-bucket result-index slots.** Consumed by *every* |
| 785 | content item, not just documents -- a text block occupies a |
| 786 | slot too, and the source at that slot carries an id-less |
| 787 | ``tool_output`` payload synthesized via |
| 788 | :meth:`_text_to_tool_output_payload`, matching the cohere |
| 789 | api's behavior for text-only tool content. Strings and |
| 790 | empty contents each consume exactly one slot because the |
| 791 | cohere renderer wraps them into a single text block before |
| 792 | handing the message off to melody. Multiple tool messages |
| 793 | sharing a ``tool_call_id`` accumulate into the same bucket, |
| 794 | so later docs' slots are offset by the sum of previous |
| 795 | messages' content lengths. |
| 796 | * **Doc ids attached to each position.** Top-level docs |
| 797 | contribute both the explicit ``id`` (if set) and the |
| 798 | ``doc_{idx}`` fallback used for inbound addressing; tool |
| 799 | document blocks contribute their single ``document.id`` |
| 800 | when present. The ``ids`` tuple is empty for slots that |
| 801 | are only reachable positionally (text-only tool content, |
| 802 | docs without a client-provided id) -- those slots still |
| 803 | have a fully-populated :class:`InternalCitationSource`, |
| 804 | just no client-visible id to cite them by. |
| 805 | |
| 806 | Ids alone key the inbound helper |
| 807 | (:meth:`_build_doc_id_to_prompt_position`, which drives |
| 808 | request-side citation resolution); the resolved |
| 809 | :class:`InternalCitationSource` (with ``type`` / ``id`` / |
| 810 | payload) keys the outbound helper |
| 811 | (:meth:`_build_position_to_source`, whose result is forwarded |
| 812 | to the reasoning parser via |
| 813 | ``chat_template_kwargs[POSITION_TO_SOURCE_KEY]``). |
| 814 | |
| 815 | Cohere's on-wire semantics -- ``ToolSource.id`` and |
| 816 | ``DocumentSource.id`` both carry a document id, with ``type`` |
| 817 | acting only as a payload-shape hint -- are the invariant this |
| 818 | walk is mirroring. |
| 819 | """ |
| 820 | out: list[_CitablePosition] = [] |
| 821 | tool_call_id_to_bucket: dict[str, int] = {} |
| 822 | # Next unassigned ``tool_result_index`` per bucket. Advances |
| 823 | # by ``len(content)`` (or 1 for string / missing content) for |
| 824 | # every tool message, whether or not the content had any doc |
| 825 | # ids in it -- matching melody's per-slot accounting. |
| 826 | bucket_result_len: dict[int, int] = {} |
| 827 | next_bucket = 0 |
| 828 | |
| 829 | if documents: |
| 830 | for idx, doc in enumerate(documents): |
| 831 | # Every top-level doc is addressable inbound by a |
| 832 | # synthetic ``doc_{idx}`` fallback (so clients can cite |
| 833 | # positional docs). On the *outbound* wire we prefer |
| 834 | # the client's explicit id when present, and omit the |
| 835 | # ``id`` field entirely when it's not -- matching |
| 836 | # cohere's ``DocumentSource`` where ``id`` is optional |
| 837 | # and preserving "no id in, no id out" symmetry. |
| 838 | fallback_id = f"doc_{idx}" |
| 839 | if isinstance(doc, str): |
| 840 | payload: dict[str, Any] = {"text": doc} |
| 841 | explicit_id: str | None = None |
| 842 | else: |
| 843 | payload = ( |
| 844 | dict(doc.data) |
| 845 | if isinstance(doc.data, dict) |
| 846 | else {"text": doc.data} |
| 847 | ) |
| 848 | explicit_id = doc.id or None |
| 849 | if explicit_id: |
| 850 | payload.setdefault("id", explicit_id) |
| 851 | # Register explicit id first so it wins over the |
| 852 | # fallback for ``setdefault``-style inbound lookup. |
| 853 | ids: tuple[str, ...] = ( |
| 854 | (explicit_id, fallback_id) if explicit_id else (fallback_id,) |
| 855 | ) |
| 856 | out.append( |
| 857 | _CitablePosition( |
| 858 | bucket=0, |
| 859 | result_idx=idx, |
| 860 | ids=ids, |
| 861 | source=InternalCitationSource( |
| 862 | type="document", id=explicit_id, document=payload |
| 863 | ), |
| 864 | ) |
| 865 | ) |
| 866 | next_bucket = 1 |
| 867 | |
| 868 | def register_bucket(tool_call_id: str) -> int: |
| 869 | nonlocal next_bucket |
| 870 | bucket = tool_call_id_to_bucket.get(tool_call_id) |
| 871 | if bucket is None: |
| 872 | bucket = next_bucket |
| 873 | tool_call_id_to_bucket[tool_call_id] = bucket |
| 874 | next_bucket += 1 |
| 875 | return bucket |
| 876 | |
| 877 | for msg in messages: |
| 878 | if isinstance(msg, ToolChatMessageV2): |
| 879 | if not msg.tool_call_id: |
| 880 | continue |
| 881 | bucket = register_bucket(msg.tool_call_id) |
| 882 | base = bucket_result_len.get(bucket, 0) |
| 883 | if isinstance(msg.content, list): |
| 884 | for offset, block in enumerate(msg.content): |
| 885 | info = cls._tool_content_block_to_source(block) |
| 886 | if info is None: |
| 887 | continue |
| 888 | # Only real document ids get registered in the |
| 889 | # inbound-facing ``ids`` tuple; text-only tool |
| 890 | # content produces a source with ``id=None`` |
| 891 | # to match the cohere api -- the text is |
| 892 | # exposed as ``tool_output.content`` but |
| 893 | # there's no id for clients to cite by. |
| 894 | pos_ids: tuple[str, ...] = (info.id,) if info.id else () |
| 895 | out.append( |
| 896 | _CitablePosition( |
| 897 | bucket=bucket, |
| 898 | result_idx=base + offset, |
| 899 | ids=pos_ids, |
| 900 | source=info, |
| 901 | ) |
| 902 | ) |
| 903 | bucket_result_len[bucket] = base + len(msg.content) |
| 904 | else: |
| 905 | # A string or missing content is wrapped by the |
| 906 | # cohere renderer into a single text block before |
| 907 | # melody sees it, so it consumes one slot. For |
| 908 | # string content, emit a synthetic id-less tool |
| 909 | # source carrying the text as its payload so a |
| 910 | # model citation against this slot still has |
| 911 | # something to attach on the wire (matching the |
| 912 | # id-less tool_output behavior for text-only |
| 913 | # ``ToolChatMessageV2`` content blocks). Missing |
| 914 | # content still advances the slot cursor but |
| 915 | # produces no source -- a citation against it |
| 916 | # would be dropped by ``_to_wire_citation``. |
| 917 | if isinstance(msg.content, str): |
| 918 | payload = cls._text_to_tool_output_payload(msg.content) |
| 919 | out.append( |
| 920 | _CitablePosition( |
| 921 | bucket=bucket, |
| 922 | result_idx=base, |
| 923 | ids=(), |
| 924 | source=InternalCitationSource( |
| 925 | type="tool", id=None, tool_output=payload |
| 926 | ), |
| 927 | ) |
| 928 | ) |
| 929 | bucket_result_len[bucket] = base + 1 |
| 930 | elif isinstance(msg, AssistantChatMessageV2) and msg.tool_calls: |
| 931 | for tc in msg.tool_calls: |
| 932 | if tc.id: |
| 933 | register_bucket(tc.id) |
| 934 | |
| 935 | return out |
| 936 | |
| 937 | @classmethod |
| 938 | def _build_doc_id_to_prompt_position( |
| 939 | cls, messages: list, documents: list | None |
| 940 | ) -> dict[str, tuple[int, int]]: |
| 941 | """Map every citable document id to its ``(bucket, result_index)``. |
| 942 | |
| 943 | Inbound helper: converts request-side citation ``source.id`` |
| 944 | refs into melody's numeric ``FilterCitation`` addressing (see |
| 945 | :meth:`_sdk_citation_to_melody`). First-seen-id wins so an |
| 946 | explicit id on a top-level document takes precedence over the |
| 947 | ``doc_{idx}`` fallback that :meth:`_walk_citable_positions` |
| 948 | also emits for the same slot. |
| 949 | |
| 950 | Thin wrapper over :meth:`_walk_citable_positions`; the |
| 951 | numbering rule lives there. |
| 952 | """ |
| 953 | doc_positions: dict[str, tuple[int, int]] = {} |
| 954 | for pos in cls._walk_citable_positions(messages, documents): |
| 955 | for pos_id in pos.ids: |
| 956 | doc_positions.setdefault(pos_id, (pos.bucket, pos.result_idx)) |
| 957 | return doc_positions |
| 958 | |
| 959 | @staticmethod |
| 960 | def _sdk_citation_to_melody( |
| 961 | citation: Any, |
| 962 | doc_positions: dict[str, tuple[int, int]], |
| 963 | ) -> dict[str, Any] | None: |
| 964 | """Convert a Cohere SDK ``Citation`` to melody's ``FilterCitation``. |
| 965 | |
| 966 | Melody's ``Source`` uses a two-level numeric address: |
| 967 | ``tool_call_index`` picks a bucket (``0`` = top-level |
| 968 | ``documents`` when present, ``1..N`` = tool results in history |
| 969 | order) and ``tool_result_indices`` picks positions inside that |
| 970 | bucket. The Cohere ``Citation.sources`` schema is per-document: |
| 971 | each entry's ``id`` is the id of the specific document that |
| 972 | grounds the citation (``type`` distinguishes payload shape, not |
| 973 | id semantics). We resolve each source id against |
| 974 | ``doc_positions`` and aggregate hits sharing a bucket into one |
| 975 | melody ``Source`` so the renderer emits compact |
| 976 | ``<co>text</co: N:[i,j]>`` markers. |
| 977 | |
| 978 | Returns ``None`` when any source id fails to resolve or when |
| 979 | the input has no sources at all: emitting a citation with an |
| 980 | empty ``sources=[]`` would render a malformed |
| 981 | ``<co>text</co: :>`` marker, and silently attributing a |
| 982 | citation to the wrong bucket is worse than dropping the |
| 983 | ``<co>`` markup. This matches the cohere api's behavior. |
| 984 | |
| 985 | ``Source.document_ids`` is intentionally NOT populated: it is a |
| 986 | parser-*output* lookup table (see ``PromptRenderIds`` in |
| 987 | melody), and melody's deserializer ignores it on input. |
| 988 | |
| 989 | The Cohere ``type`` literal (``TEXT_CONTENT`` / ``THINKING_CONTENT`` |
| 990 | / ``PLAN``) collapses to melody's boolean ``is_thinking``. |
| 991 | ``PLAN`` is treated as a thinking-style citation because cmd3 / |
| 992 | cmd4 templates render both inline with the same ``<co>...</co>`` |
| 993 | markup. |
| 994 | """ |
| 995 | raw_sources = getattr(citation, "sources", None) or [] |
| 996 | if not raw_sources: |
| 997 | return None |
| 998 | |
| 999 | grouped: dict[int, list[int]] = {} |
| 1000 | for src in raw_sources: |
| 1001 | src_id = getattr(src, "id", None) |
| 1002 | if not src_id: |
| 1003 | return None |
| 1004 | position = doc_positions.get(src_id) |
| 1005 | if position is None: |
| 1006 | return None |
| 1007 | bucket, result_idx = position |
| 1008 | grouped.setdefault(bucket, []).append(result_idx) |
| 1009 | |
| 1010 | sources = [ |
| 1011 | {"tool_call_index": bucket, "tool_result_indices": indices} |
| 1012 | for bucket, indices in grouped.items() |
| 1013 | ] |
| 1014 | |
| 1015 | cite_type = str(getattr(citation, "type", "") or "").upper() |
| 1016 | is_thinking = cite_type in ("THINKING_CONTENT", "PLAN") |
| 1017 | return { |
| 1018 | "start_index": getattr(citation, "start", None) or 0, |
| 1019 | "end_index": getattr(citation, "end", None) or 0, |
| 1020 | "text": getattr(citation, "text", None) or "", |
| 1021 | "sources": sources, |
| 1022 | "is_thinking": is_thinking, |
| 1023 | } |
| 1024 | |
| 1025 | # ================================================================== |
| 1026 | # Response conversion: ChatCompletion -> Cohere V2 |
| 1027 | # ================================================================== |
| 1028 | |
| 1029 | def _chat_completion_to_v2( |
| 1030 | self, |
| 1031 | response: ChatCompletionResponse, |
| 1032 | request: CohereChatV2Request, |
| 1033 | ) -> CohereChatV2Response: |
| 1034 | choice = response.choices[0] |
| 1035 | msg = choice.message |
| 1036 | |
| 1037 | # Build content blocks as dicts; ``AssistantMessageResponse`` |
| 1038 | # validates them into the proper discriminated union variants |
| 1039 | # (``Text/ThinkingAssistantMessageResponseContentItem``). |
| 1040 | content_blocks: list[dict[str, Any]] = [] |
| 1041 | if msg.reasoning: |
| 1042 | content_blocks.append( |
| 1043 | {"type": ContentBlockType.THINKING, "thinking": msg.reasoning} |
| 1044 | ) |
| 1045 | if msg.content: |
| 1046 | content_blocks.append({"type": ContentBlockType.TEXT, "text": msg.content}) |
| 1047 | |
| 1048 | tool_calls: list[ToolCallV2] | None = None |
| 1049 | if msg.tool_calls: |
| 1050 | tool_calls = [ |
| 1051 | ToolCallV2( |
| 1052 | id=tc.id, |
| 1053 | function=ToolCallV2Function( |
| 1054 | name=tc.function.name, |
| 1055 | arguments=tc.function.arguments, |
| 1056 | ), |
| 1057 | ) |
| 1058 | for tc in msg.tool_calls |
| 1059 | ] |
| 1060 | |
| 1061 | # Cohere's ``tool_plan`` is the planning text emitted before tool |
| 1062 | # calls on older, non-reasoning Command models. For those models |
| 1063 | # we surface ``reasoning`` as ``tool_plan`` and drop the thinking |
| 1064 | # block. Reasoning Command models emit a regular thinking block |
| 1065 | # alongside tool calls, so for them we leave the thinking block |
| 1066 | # in place and never set ``tool_plan``. |
| 1067 | tool_plan: str | None = None |
| 1068 | if not self._is_reasoning_model and tool_calls and msg.reasoning: |
| 1069 | tool_plan = msg.reasoning |
| 1070 | content_blocks = [ |
| 1071 | blk |
| 1072 | for blk in content_blocks |
| 1073 | if blk.get("type") != ContentBlockType.THINKING |
| 1074 | ] |
| 1075 | |
| 1076 | assistant_msg = AssistantMessageResponse( |
| 1077 | content=content_blocks or None, |
| 1078 | tool_calls=tool_calls, |
| 1079 | tool_plan=tool_plan, |
| 1080 | citations=self._extract_citations_if_any(msg), |
| 1081 | ) |
| 1082 | |
| 1083 | usage = self._build_usage(response) |
| 1084 | |
| 1085 | return CohereChatV2Response( |
| 1086 | id=response.id or f"chat_{int(time.time() * 1000)}", |
| 1087 | finish_reason=_map_finish_reason( |
| 1088 | choice.finish_reason, |
| 1089 | choice.stop_reason, |
| 1090 | ), |
| 1091 | message=assistant_msg, |
| 1092 | usage=usage, |
| 1093 | kv_transfer_params=response.kv_transfer_params, |
| 1094 | ) |
| 1095 | |
| 1096 | def _create_chat_message(self, *args: Any, **kwargs: Any) -> ChatMessage: |
| 1097 | """Route response construction through the citation-carrying subclass. |
| 1098 | |
| 1099 | Overrides :meth:`OpenAIServingChat._create_chat_message` so every |
| 1100 | non-streaming construction site in the base full-generator |
| 1101 | produces a :class:`CohereChatMessage`. That lets |
| 1102 | :meth:`_finalize_response_message` below set citations on |
| 1103 | ``message.citations`` directly, without relying on |
| 1104 | ``extra="allow"`` attribute injection. |
| 1105 | """ |
| 1106 | return CohereChatMessage(*args, **kwargs) |
| 1107 | |
| 1108 | def _finalize_response_message( |
| 1109 | self, |
| 1110 | message: ChatMessage, |
| 1111 | *, |
| 1112 | parser: Parser | None, |
| 1113 | ) -> ChatMessage: |
| 1114 | """Copy grounding citations off the unified parser onto the message. |
| 1115 | |
| 1116 | :class:`vllm.parser.cohere_command.CohereCommandParser` caches the |
| 1117 | citations produced by its most recent unary ``parse`` call on |
| 1118 | ``last_unary_citations``. We surface them here on |
| 1119 | :class:`CohereChatMessage` so downstream response conversion can |
| 1120 | pick them up without the base :class:`OpenAIServingChat` having to |
| 1121 | know about citations. |
| 1122 | """ |
| 1123 | # ``_create_chat_message`` above guarantees the concrete type at |
| 1124 | # runtime. ``cast`` narrows the declared ``ChatMessage`` for |
| 1125 | # mypy without adding a runtime check we don't need: even in the |
| 1126 | # invariant-violated case (a subclass reset the factory back to |
| 1127 | # plain ``ChatMessage``), ``OpenAIBaseModel``'s ``extra="allow"`` |
| 1128 | # config lets the ``.citations`` write survive into |
| 1129 | # ``model_dump`` via the extras bucket, so the wire is correct |
| 1130 | # either way. |
| 1131 | message = cast(CohereChatMessage, message) |
| 1132 | citations = getattr(parser, "last_unary_citations", None) |
| 1133 | if citations: |
| 1134 | message.citations = citations |
| 1135 | return message |
| 1136 | |
| 1137 | def _extract_citations_if_any(self, msg: Any) -> list[Citation] | None: |
| 1138 | """Coerce ``CohereChatMessage.citations`` into the Cohere v2 wire shape. |
| 1139 | |
| 1140 | :class:`CohereChatMessage` carries a |
| 1141 | ``citations: list[vllm...Citation] | None`` field populated by |
| 1142 | :meth:`_finalize_response_message` (which reads it off the |
| 1143 | unified parser's ``last_unary_citations`` cache). Sources are |
| 1144 | already fully resolved by the parser (see |
| 1145 | :func:`_melody_sources_to_vllm` in |
| 1146 | :mod:`vllm.parser.cohere_command`) using |
| 1147 | the position map forwarded via ``chat_template_kwargs``. This |
| 1148 | method's remaining job is: |
| 1149 | |
| 1150 | * Drop citations whose sources all failed to resolve (empty |
| 1151 | ``sources`` list) -- fail-closed policy matching the inbound |
| 1152 | ``_sdk_citation_to_melody``. |
| 1153 | * Rewrite ``THINKING_CONTENT`` -> ``PLAN`` for non-reasoning |
| 1154 | models, since ``_chat_completion_to_v2`` surfaces |
| 1155 | ``msg.reasoning`` as ``tool_plan`` on the same |
| 1156 | ``is_reasoning_model`` flag. |
| 1157 | * Coerce to :class:`cohere.types.Citation` for the wire. |
| 1158 | """ |
| 1159 | raw = getattr(msg, "citations", None) |
| 1160 | if not raw: |
| 1161 | return None |
| 1162 | out: list[Citation] = [] |
| 1163 | for c in raw: |
| 1164 | resolved = self._to_wire_citation(c) |
| 1165 | if resolved is not None: |
| 1166 | out.append(resolved) |
| 1167 | return out or None |
| 1168 | |
| 1169 | def _to_wire_citation(self, citation: Any) -> Citation | None: |
| 1170 | """Coerce one parser-produced citation into the SDK wire shape. |
| 1171 | |
| 1172 | Accepts both :class:`InternalCitation` instances (what the |
| 1173 | unary path receives directly from the parser via |
| 1174 | ``last_unary_citations``) and plain ``dict`` payloads (what |
| 1175 | the streaming path sees, because ``DeltaMessage.citations`` is |
| 1176 | an untyped extras field on the OpenAI wire protocol and dict |
| 1177 | shapes survive ``model_validate_json`` unchanged). Both flows |
| 1178 | end up carrying the same key names, so a small ``_get`` helper |
| 1179 | handles both without a shape-specific code path. |
| 1180 | """ |
| 1181 | |
| 1182 | def _get(obj: Any, key: str, default: Any = None) -> Any: |
| 1183 | if isinstance(obj, dict): |
| 1184 | return obj.get(key, default) |
| 1185 | return getattr(obj, key, default) |
| 1186 | |
| 1187 | raw_sources = _get(citation, "sources") or [] |
| 1188 | wire_sources: list[dict[str, Any]] = [] |
| 1189 | for src in raw_sources: |
| 1190 | # Sources arrive pre-resolved from the parser (or as trusted |
| 1191 | # dicts on the streaming path). Match the cohere api's |
| 1192 | # behavior of preserving id-less sources (e.g. text-only |
| 1193 | # tool results): cohere's ``Source`` types treat ``id`` as |
| 1194 | # optional, so we simply drop the field via |
| 1195 | # ``exclude_none=True`` instead of dropping the source. |
| 1196 | if isinstance(src, dict): |
| 1197 | wire = {k: v for k, v in src.items() if v is not None} |
| 1198 | if wire: |
| 1199 | wire_sources.append(wire) |
| 1200 | elif isinstance(src, InternalCitationSource): |
| 1201 | wire_sources.append(src.model_dump(exclude_none=True)) |
| 1202 | if not wire_sources: |
| 1203 | return None |
| 1204 | |
| 1205 | cite_type = _get(citation, "type") |
| 1206 | if cite_type == "THINKING_CONTENT" and not self._is_reasoning_model: |
| 1207 | # Non-reasoning Command models surface the reasoning block |
| 1208 | # as ``tool_plan`` (see ``_chat_completion_to_v2``), so any |
| 1209 | # ``is_thinking`` citation from melody is actually a PLAN |
| 1210 | # citation on the wire. |
| 1211 | cite_type = "PLAN" |
| 1212 | |
| 1213 | payload = { |
| 1214 | "start": _get(citation, "start"), |
| 1215 | "end": _get(citation, "end"), |
| 1216 | "text": _get(citation, "text"), |
| 1217 | "sources": wire_sources, |
| 1218 | "type": cite_type, |
| 1219 | } |
| 1220 | try: |
| 1221 | return Citation.model_validate(payload) |
| 1222 | except Exception: # pragma: no cover - defensive |
| 1223 | logger.warning( |
| 1224 | "Skipping malformed citation payload: %r", payload, exc_info=True |
| 1225 | ) |
| 1226 | return None |
| 1227 | |
| 1228 | @classmethod |
| 1229 | def _build_position_to_source( |
| 1230 | cls, request: CohereChatV2Request |
| 1231 | ) -> dict[tuple[int, int], InternalCitationSource]: |
| 1232 | """Invert the request-side numbering into ``(bucket, idx) -> source``. |
| 1233 | |
| 1234 | Outbound helper: forwarded to the reasoning parser via |
| 1235 | ``chat_template_kwargs[POSITION_TO_SOURCE_KEY]`` so |
| 1236 | :func:`_melody_sources_to_vllm` can attach real document ids / |
| 1237 | types / payloads to each melody source it emits. Each entry |
| 1238 | carries the fully-resolved wire-shape info: ``type`` |
| 1239 | (``"document"`` for the reserved top-level bucket 0, ``"tool"`` |
| 1240 | for any tool-call bucket), ``id`` (explicit doc id, the |
| 1241 | ``doc_{idx}`` fallback, or ``None`` for id-less sources), and |
| 1242 | ``document`` / ``tool_output`` payload -- matching the cohere |
| 1243 | api's ``Source`` shape. |
| 1244 | |
| 1245 | Thin wrapper over :meth:`_walk_citable_positions`; the |
| 1246 | numbering rule lives there. |
| 1247 | """ |
| 1248 | return { |
| 1249 | (pos.bucket, pos.result_idx): pos.source |
| 1250 | for pos in cls._walk_citable_positions(request.messages, request.documents) |
| 1251 | } |
| 1252 | |
| 1253 | @staticmethod |
| 1254 | def _text_to_tool_output_payload(text: str) -> dict[str, Any]: |
| 1255 | """Wrap a bare tool-content text into a ``tool_output`` payload. |
| 1256 | |
| 1257 | Matches the cohere api's handling of text-only tool content: |
| 1258 | attempt to interpret ``text`` as a JSON object first; fall |
| 1259 | back to ``{"content": <text>}`` so a text-only tool result |
| 1260 | still has a structured payload we can attach to a citation |
| 1261 | source. |
| 1262 | """ |
| 1263 | if not text: |
| 1264 | return {} |
| 1265 | try: |
| 1266 | parsed = json.loads(text) |
| 1267 | except (json.JSONDecodeError, TypeError): |
| 1268 | return {"content": text} |
| 1269 | if isinstance(parsed, dict): |
| 1270 | return parsed |
| 1271 | return {"content": text} |
| 1272 | |
| 1273 | @classmethod |
| 1274 | def _tool_content_block_to_source( |
| 1275 | cls, |
| 1276 | block: Any, |
| 1277 | ) -> InternalCitationSource | None: |
| 1278 | """Extract the wire-shape source for one tool-message content block. |
| 1279 | |
| 1280 | Handles both cohere ``DocumentToolContent`` and |
| 1281 | ``TextToolContent`` blocks: |
| 1282 | |
| 1283 | * ``document`` blocks with an explicit id produce a |
| 1284 | ``tool`` source keyed by that id (payload is the document's |
| 1285 | ``data`` dict, or ``{"text": ...}`` for scalar data). |
| 1286 | * ``document`` blocks *without* an id and ``text`` blocks |
| 1287 | produce an id-less ``tool`` source: the cohere api emits |
| 1288 | the payload with an empty id in this case, and cohere's |
| 1289 | SDK treats ``ToolSource.id`` as optional -- so we simply |
| 1290 | omit ``id`` on the wire instead of inventing a synthetic |
| 1291 | one. |
| 1292 | |
| 1293 | Returns ``None`` only for shapes we don't recognise at all; |
| 1294 | that slot still gets counted (callers advance the bucket |
| 1295 | cursor over the block regardless) so numbering stays aligned |
| 1296 | with melody. |
| 1297 | """ |
| 1298 | block_type = getattr(block, "type", None) |
| 1299 | if block_type == "document": |
| 1300 | doc: Document | None = getattr(block, "document", None) |
| 1301 | if doc is None: |
| 1302 | return None |
| 1303 | doc_id = doc.id or None |
| 1304 | data = doc.data |
| 1305 | if isinstance(data, dict): |
| 1306 | payload: dict[str, Any] = dict(data) |
| 1307 | elif data is None: |
| 1308 | payload = {} |
| 1309 | else: |
| 1310 | # Match the top-level ``documents`` handling (see |
| 1311 | # ``_apply_cohere_template_kwargs``): non-dict payloads |
| 1312 | # get wrapped as ``{"text": ...}`` so clients get a |
| 1313 | # renderable body instead of a bare id. |
| 1314 | payload = {"text": data} |
| 1315 | if doc_id: |
| 1316 | payload.setdefault("id", doc_id) |
| 1317 | return InternalCitationSource( |
| 1318 | type="tool", |
| 1319 | id=doc_id, |
| 1320 | tool_output=payload, |
| 1321 | ) |
| 1322 | if block_type == "text": |
| 1323 | text = getattr(block, "text", None) or "" |
| 1324 | payload = cls._text_to_tool_output_payload(text) |
| 1325 | return InternalCitationSource( |
| 1326 | type="tool", |
| 1327 | id=None, |
| 1328 | tool_output=payload, |
| 1329 | ) |
| 1330 | return None |
| 1331 | |
| 1332 | @classmethod |
| 1333 | def _collect_v2_tool_message_content( |
| 1334 | cls, request: CohereChatV2Request |
| 1335 | ) -> dict[int, list[dict[str, Any]]]: |
| 1336 | """Collect the original v2 tool-message content blocks so the |
| 1337 | renderer can apply its v2 → melody mapping. |
| 1338 | |
| 1339 | Builds the map that |
| 1340 | :meth:`_apply_cohere_template_kwargs` stashes under |
| 1341 | :data:`TOOL_MESSAGE_V2_CONTENT_KEY` for the renderer to |
| 1342 | consume (see the constant definition in |
| 1343 | :mod:`vllm.renderers.cohere` for why the map exists). Walks |
| 1344 | ``request.messages`` -- not the OpenAI-shape messages built by |
| 1345 | :meth:`_convert_messages` -- because that flattening step |
| 1346 | already threw away the structural distinction between text and |
| 1347 | document blocks. Each surviving block is |
| 1348 | :meth:`~pydantic.BaseModel.model_dump`'d to a plain dict so the |
| 1349 | forwarded map carries no Pydantic instances (msgpack- / |
| 1350 | general-serialization-safe on principle even though internal |
| 1351 | underscore-prefixed keys are filtered before crossing the |
| 1352 | engine-core boundary). |
| 1353 | |
| 1354 | Only emits an entry for a tool message whose content is a |
| 1355 | non-empty list of blocks; a bare-string tool message renders |
| 1356 | fine through the default text-block path in |
| 1357 | :func:`vllm.renderers.cohere._content_blocks`, so no forwarded |
| 1358 | blocks are needed there. Returns an empty dict when there is |
| 1359 | nothing to forward, so the caller can use truthiness to decide |
| 1360 | whether to stash the key at all. |
| 1361 | """ |
| 1362 | blocks_by_index: dict[int, list[dict[str, Any]]] = {} |
| 1363 | for i, msg in enumerate(request.messages): |
| 1364 | if not isinstance(msg, ToolChatMessageV2): |
| 1365 | continue |
| 1366 | if isinstance(msg.content, str) or not msg.content: |
| 1367 | continue |
| 1368 | v2_blocks: list[dict[str, Any]] = [] |
| 1369 | for block in msg.content: |
| 1370 | v2_blocks.append(block.model_dump(exclude_none=True)) |
| 1371 | if v2_blocks: |
| 1372 | blocks_by_index[i] = v2_blocks |
| 1373 | return blocks_by_index |
| 1374 | |
| 1375 | @staticmethod |
| 1376 | def _build_usage(response: ChatCompletionResponse) -> CohereUsage | None: |
| 1377 | if response.usage is None: |
| 1378 | return None |
| 1379 | prompt = response.usage.prompt_tokens |
| 1380 | completion = response.usage.completion_tokens or 0 |
| 1381 | cached: int | None = None |
| 1382 | if response.usage.prompt_tokens_details is not None: |
| 1383 | cached = response.usage.prompt_tokens_details.cached_tokens |
| 1384 | return CohereUsage( |
| 1385 | billed_units=CohereUsageBilledUnits( |
| 1386 | input_tokens=prompt, |
| 1387 | output_tokens=completion, |
| 1388 | ), |
| 1389 | tokens=CohereUsageTokens( |
| 1390 | input_tokens=prompt, |
| 1391 | output_tokens=completion, |
| 1392 | ), |
| 1393 | cached_tokens=cached, |
| 1394 | ) |
| 1395 | |
| 1396 | # ================================================================== |
| 1397 | # Stream conversion: chat completion stream -> Cohere V2 SSE events |
| 1398 | # ================================================================== |
| 1399 | |
| 1400 | async def _chat_completion_stream_to_v2( |
| 1401 | self, |
| 1402 | generator: AsyncGenerator[str, None], |
| 1403 | request: CohereChatV2Request, |
| 1404 | ) -> AsyncGenerator[str, None]: |
| 1405 | """Translate an OpenAI-style chat completion SSE stream into Cohere's |
| 1406 | v2 stream-event format. |
| 1407 | |
| 1408 | Cohere's v2 stream lifecycle is: |
| 1409 | |
| 1410 | message-start |
| 1411 | [content-start, content-delta..., content-end]* |
| 1412 | [tool-plan-delta]* |
| 1413 | [tool-call-start, tool-call-delta..., tool-call-end]* |
| 1414 | message-end |
| 1415 | """ |
| 1416 | state = _StreamState() |
| 1417 | |
| 1418 | try: |
| 1419 | async for item in generator: |
| 1420 | if not item.startswith("data:"): |
| 1421 | continue |
| 1422 | data_str = item[len("data:") :].strip().rstrip("\n") |
| 1423 | if not data_str: |
| 1424 | continue |
| 1425 | if data_str == "[DONE]": |
| 1426 | # OpenAI's stream terminator. Fall through to the |
| 1427 | # post-loop cleanup so we always emit ``message-end`` |
| 1428 | # even if the usage-only chunk was skipped. |
| 1429 | break |
| 1430 | |
| 1431 | chunk = ChatCompletionStreamResponse.model_validate_json(data_str) |
| 1432 | state.last_chunk_id = chunk.id |
| 1433 | |
| 1434 | if not state.started: |
| 1435 | yield _emit( |
| 1436 | MessageStartEvent( |
| 1437 | id=chunk.id, |
| 1438 | delta={"message": {"role": "assistant"}}, |
| 1439 | ) |
| 1440 | ) |
| 1441 | state.started = True |
| 1442 | |
| 1443 | # The final OpenAI chunk has no choices and only carries usage. |
| 1444 | if not chunk.choices: |
| 1445 | for ev in self._close_open_blocks(state): |
| 1446 | yield ev |
| 1447 | yield self._build_message_end_event( |
| 1448 | chunk_id=chunk.id, |
| 1449 | finish_reason=state.finish_reason, |
| 1450 | stop_reason=state.stop_reason, |
| 1451 | usage_chunk=chunk, |
| 1452 | ) |
| 1453 | state.ended = True |
| 1454 | continue |
| 1455 | |
| 1456 | choice = chunk.choices[0] |
| 1457 | if choice.finish_reason is not None: |
| 1458 | state.finish_reason = choice.finish_reason |
| 1459 | state.stop_reason = choice.stop_reason |
| 1460 | |
| 1461 | delta = choice.delta |
| 1462 | |
| 1463 | # Reasoning -> thinking content block |
| 1464 | reasoning = getattr(delta, "reasoning", None) or getattr( |
| 1465 | delta, "reasoning_content", None |
| 1466 | ) |
| 1467 | if reasoning: |
| 1468 | for ev in self._handle_thinking_delta(state, reasoning): |
| 1469 | yield ev |
| 1470 | |
| 1471 | if delta.content: |
| 1472 | for ev in self._handle_text_delta(state, delta.content): |
| 1473 | yield ev |
| 1474 | |
| 1475 | if delta.tool_calls: |
| 1476 | for ev in self._handle_tool_call_deltas(state, delta.tool_calls): |
| 1477 | yield ev |
| 1478 | |
| 1479 | # Citations: a Cohere-specific extension on DeltaMessage that |
| 1480 | # the cohere renderer/parsers may populate. |
| 1481 | delta_citations = getattr(delta, "citations", None) |
| 1482 | if delta_citations: |
| 1483 | for ev in self._handle_citation_deltas(state, delta_citations): |
| 1484 | yield ev |
| 1485 | |
| 1486 | except Exception as exc: |
| 1487 | logger.exception("Error converting chat completion stream to v2") |
| 1488 | if state.started and not state.ended: |
| 1489 | yield _sse( |
| 1490 | json.dumps( |
| 1491 | { |
| 1492 | "type": "message-end", |
| 1493 | "delta": { |
| 1494 | "error": sanitize_message(str(exc)), |
| 1495 | "finish_reason": "ERROR", |
| 1496 | }, |
| 1497 | } |
| 1498 | ) |
| 1499 | ) |
| 1500 | state.ended = True |
| 1501 | yield _DONE_FRAME |
| 1502 | return |
| 1503 | |
| 1504 | # Normal completion or ``[DONE]``: ensure ``message-end`` is always |
| 1505 | # emitted. Upstream may close the stream without sending the final |
| 1506 | # usage-only chunk (e.g. on shutdown, or when ``[DONE]`` is the only |
| 1507 | # terminator); without this fallback Cohere clients would hang |
| 1508 | # waiting for the closing event. |
| 1509 | if state.started and not state.ended: |
| 1510 | for ev in self._close_open_blocks(state): |
| 1511 | yield ev |
| 1512 | yield self._build_message_end_event( |
| 1513 | chunk_id=state.last_chunk_id, |
| 1514 | finish_reason=state.finish_reason, |
| 1515 | stop_reason=state.stop_reason, |
| 1516 | usage_chunk=None, |
| 1517 | ) |
| 1518 | state.ended = True |
| 1519 | |
| 1520 | # Stream terminator. Cohere's v2 SSE protocol ends every stream |
| 1521 | # with ``data: [DONE]\n\n`` after ``message-end``; Fern-generated |
| 1522 | # clients (Go/Java) and cohere-python all key their read loop off |
| 1523 | # this sentinel. |
| 1524 | yield _DONE_FRAME |
| 1525 | |
| 1526 | # -- per-delta helpers -------------------------------------------- |
| 1527 | |
| 1528 | def _handle_thinking_delta(self, state: _StreamState, delta_text: str) -> list[str]: |
| 1529 | # Non-reasoning Command models: emit ``tool-plan-delta`` events |
| 1530 | # directly instead of opening a thinking content block. The |
| 1531 | # ``tool-plan-delta`` event has no start/end pair around it. |
| 1532 | if not self._is_reasoning_model: |
| 1533 | events: list[str] = list(self._close_open_blocks(state)) |
| 1534 | events.append( |
| 1535 | _emit( |
| 1536 | ToolPlanDeltaEvent( |
| 1537 | delta={"message": {"tool_plan": delta_text}}, |
| 1538 | ) |
| 1539 | ) |
| 1540 | ) |
| 1541 | return events |
| 1542 | |
| 1543 | # Reasoning model (default): open / continue a thinking block. |
| 1544 | events = [] |
| 1545 | if state.active_block != ContentBlockType.THINKING: |
| 1546 | events.extend(self._close_open_blocks(state)) |
| 1547 | idx = state.next_content_index() |
| 1548 | state.active_block = ContentBlockType.THINKING |
| 1549 | state.active_block_index = idx |
| 1550 | events.append( |
| 1551 | _emit( |
| 1552 | ContentStartEvent( |
| 1553 | index=idx, |
| 1554 | delta={ |
| 1555 | "message": { |
| 1556 | "content": { |
| 1557 | "type": ContentBlockType.THINKING, |
| 1558 | "thinking": "", |
| 1559 | } |
| 1560 | } |
| 1561 | }, |
| 1562 | ) |
| 1563 | ) |
| 1564 | ) |
| 1565 | events.append( |
| 1566 | _emit( |
| 1567 | ContentDeltaEvent( |
| 1568 | index=state.active_block_index, |
| 1569 | delta={"message": {"content": {"thinking": delta_text}}}, |
| 1570 | ) |
| 1571 | ) |
| 1572 | ) |
| 1573 | return events |
| 1574 | |
| 1575 | def _handle_text_delta(self, state: _StreamState, delta_text: str) -> list[str]: |
| 1576 | events: list[str] = [] |
| 1577 | if state.active_block != ContentBlockType.TEXT: |
| 1578 | events.extend(self._close_open_blocks(state)) |
| 1579 | idx = state.next_content_index() |
| 1580 | state.active_block = ContentBlockType.TEXT |
| 1581 | state.active_block_index = idx |
| 1582 | events.append( |
| 1583 | _emit( |
| 1584 | ContentStartEvent( |
| 1585 | index=idx, |
| 1586 | delta={ |
| 1587 | "message": { |
| 1588 | "content": { |
| 1589 | "type": ContentBlockType.TEXT, |
| 1590 | "text": "", |
| 1591 | } |
| 1592 | } |
| 1593 | }, |
| 1594 | ) |
| 1595 | ) |
| 1596 | ) |
| 1597 | events.append( |
| 1598 | _emit( |
| 1599 | ContentDeltaEvent( |
| 1600 | index=state.active_block_index, |
| 1601 | delta={"message": {"content": {"text": delta_text}}}, |
| 1602 | ) |
| 1603 | ) |
| 1604 | ) |
| 1605 | return events |
| 1606 | |
| 1607 | def _handle_tool_call_deltas(self, state: _StreamState, deltas: list) -> list[str]: |
| 1608 | events: list[str] = [] |
| 1609 | for tc in deltas: |
| 1610 | tc_index = tc.index |
| 1611 | fn = tc.function |
| 1612 | |
| 1613 | if tc_index not in state.tool_calls_seen: |
| 1614 | # New tool call. Close any open content/tool block first. |
| 1615 | events.extend(self._close_open_blocks(state)) |
| 1616 | state.tool_calls_seen.add(tc_index) |
| 1617 | state.active_tool_index = tc_index |
| 1618 | state.active_block = ContentBlockType.TOOL_CALL |
| 1619 | events.append( |
| 1620 | _emit( |
| 1621 | ToolCallStartEvent( |
| 1622 | index=tc_index, |
| 1623 | delta={ |
| 1624 | "message": { |
| 1625 | "tool_calls": { |
| 1626 | "id": tc.id or "", |
| 1627 | "type": "function", |
| 1628 | "function": { |
| 1629 | "name": (fn.name if fn else "") or "", |
| 1630 | "arguments": (fn.arguments if fn else None) |
| 1631 | or "", |
| 1632 | }, |
| 1633 | } |
| 1634 | } |
| 1635 | }, |
| 1636 | ) |
| 1637 | ) |
| 1638 | ) |
| 1639 | continue |
| 1640 | |
| 1641 | if fn and fn.arguments: |
| 1642 | events.append( |
| 1643 | _emit( |
| 1644 | ToolCallDeltaEvent( |
| 1645 | index=tc_index, |
| 1646 | delta={ |
| 1647 | "message": { |
| 1648 | "tool_calls": { |
| 1649 | "function": { |
| 1650 | "arguments": fn.arguments, |
| 1651 | } |
| 1652 | } |
| 1653 | } |
| 1654 | }, |
| 1655 | ) |
| 1656 | ) |
| 1657 | ) |
| 1658 | return events |
| 1659 | |
| 1660 | def _handle_citation_deltas( |
| 1661 | self, |
| 1662 | state: _StreamState, |
| 1663 | citations: list, |
| 1664 | ) -> list[str]: |
| 1665 | """Emit ``citation-start`` / ``citation-end`` events for a delta. |
| 1666 | |
| 1667 | The reasoning parser populates |
| 1668 | :class:`CohereDeltaMessage.citations` (see |
| 1669 | :mod:`vllm.entrypoints.cohere.cohere_chat_message`) with |
| 1670 | sources whose ``id`` / ``type`` / ``document`` / ``tool_output`` |
| 1671 | are already resolved against the request's document context |
| 1672 | (see :func:`_melody_sources_to_vllm` in the parser). This |
| 1673 | method only coerces to the SDK wire shape, applies the |
| 1674 | ``THINKING_CONTENT`` -> ``PLAN`` rewrite for non-reasoning |
| 1675 | models, and drops citations whose sources didn't resolve -- |
| 1676 | the shared :meth:`_to_wire_citation` handles all three. |
| 1677 | """ |
| 1678 | events: list[str] = [] |
| 1679 | for c in citations: |
| 1680 | citation = self._to_wire_citation(c) |
| 1681 | if citation is None: |
| 1682 | continue |
| 1683 | idx = state.next_citation_index() |
| 1684 | events.append( |
| 1685 | _emit( |
| 1686 | CitationStartEvent( |
| 1687 | index=idx, |
| 1688 | delta={ |
| 1689 | "message": { |
| 1690 | "citations": citation.model_dump(exclude_none=True) |
| 1691 | } |
| 1692 | }, |
| 1693 | ) |
| 1694 | ) |
| 1695 | ) |
| 1696 | events.append(_emit(CitationEndEvent(index=idx))) |
| 1697 | return events |
| 1698 | |
| 1699 | # -- block lifecycle helpers -------------------------------------- |
| 1700 | |
| 1701 | def _close_open_blocks(self, state: _StreamState) -> list[str]: |
| 1702 | """Emit ``content-end`` / ``tool-call-end`` for the currently open |
| 1703 | block (if any) and reset the corresponding ``_StreamState`` slots. |
| 1704 | """ |
| 1705 | events: list[str] = [] |
| 1706 | if state.active_block in (ContentBlockType.TEXT, ContentBlockType.THINKING): |
| 1707 | events.append(_emit(ContentEndEvent(index=state.active_block_index))) |
| 1708 | elif state.active_block == ContentBlockType.TOOL_CALL: |
| 1709 | events.append(_emit(ToolCallEndEvent(index=state.active_tool_index))) |
| 1710 | state.active_block = None |
| 1711 | state.active_block_index = None |
| 1712 | state.active_tool_index = None |
| 1713 | return events |
| 1714 | |
| 1715 | def _build_message_end_event( |
| 1716 | self, |
| 1717 | chunk_id: str, |
| 1718 | finish_reason: str | None, |
| 1719 | stop_reason: int | str | None = None, |
| 1720 | usage_chunk: ChatCompletionStreamResponse | None = None, |
| 1721 | ) -> str: |
| 1722 | delta: dict[str, Any] = { |
| 1723 | "finish_reason": _map_finish_reason(finish_reason, stop_reason), |
| 1724 | } |
| 1725 | if usage_chunk is not None and usage_chunk.usage is not None: |
| 1726 | prompt = usage_chunk.usage.prompt_tokens |
| 1727 | completion = usage_chunk.usage.completion_tokens or 0 |
| 1728 | usage_block: dict[str, Any] = { |
| 1729 | "billed_units": { |
| 1730 | "input_tokens": prompt, |
| 1731 | "output_tokens": completion, |
| 1732 | }, |
| 1733 | "tokens": { |
| 1734 | "input_tokens": prompt, |
| 1735 | "output_tokens": completion, |
| 1736 | }, |
| 1737 | } |
| 1738 | if usage_chunk.usage.prompt_tokens_details is not None: |
| 1739 | cached = usage_chunk.usage.prompt_tokens_details.cached_tokens |
| 1740 | if cached is not None: |
| 1741 | usage_block["cached_tokens"] = cached |
| 1742 | delta["usage"] = usage_block |
| 1743 | return _emit(MessageEndEvent(id=chunk_id, delta=delta)) |
| 1744 | |
| 1745 | # ================================================================== |
| 1746 | # Helpers for the router |
| 1747 | # ================================================================== |
| 1748 | |
| 1749 | @staticmethod |
| 1750 | def create_error_response( |
| 1751 | message: str | Exception, |
| 1752 | err_type: str = "bad_request", |
| 1753 | status_code: HTTPStatus = HTTPStatus.BAD_REQUEST, |
| 1754 | param: str | None = None, |
| 1755 | ) -> ErrorResponse: |
| 1756 | # Override of :meth:`BaseServing.create_error_response` that uses |
| 1757 | # Cohere-flavored defaults (``type="bad_request"``, ``code=400``) |
| 1758 | # so the router can translate the envelope uniformly. ``param`` |
| 1759 | # is accepted for signature parity with the base class but is |
| 1760 | # not surfaced in the Cohere wire format. |
| 1761 | |
| 1762 | del param # unused; kept for signature compatibility |
| 1763 | return ErrorResponse( |
| 1764 | error=ErrorInfo( |
| 1765 | message=str(message), |
| 1766 | type=err_type, |
| 1767 | code=int(status_code), |
| 1768 | ) |
| 1769 | ) |
| 1770 | |
| 1771 | |
| 1772 | # --------------------------------------------------------------------------- |
| 1773 | # Stream state |
| 1774 | # --------------------------------------------------------------------------- |
| 1775 | |
| 1776 | |
| 1777 | class _StreamState: |
| 1778 | """Tracks which Cohere v2 stream block (if any) is currently open.""" |
| 1779 | |
| 1780 | def __init__(self) -> None: |
| 1781 | self.started: bool = False |
| 1782 | self.ended: bool = False |
| 1783 | self.finish_reason: str | None = None |
| 1784 | self.stop_reason: int | str | None = None |
| 1785 | self.last_chunk_id: str = "" |
| 1786 | self.active_block: ContentBlockType | None = None |
| 1787 | self.active_block_index: int | None = None |
| 1788 | self.active_tool_index: int | None = None |
| 1789 | self._next_index: int = 0 |
| 1790 | self._next_citation_index: int = 0 |
| 1791 | self.tool_calls_seen: set[int] = set() |
| 1792 | |
| 1793 | def next_content_index(self) -> int: |
| 1794 | idx = self._next_index |
| 1795 | self._next_index += 1 |
| 1796 | return idx |
| 1797 | |
| 1798 | def next_citation_index(self) -> int: |
| 1799 | idx = self._next_citation_index |
| 1800 | self._next_citation_index += 1 |
| 1801 | return idx |