| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Cohere prompt renderer. |
| 4 | |
| 5 | Templates the Cohere Command-family prompt formats (cmd3 / cmd4) using |
| 6 | the ``cohere_melody`` Rust bindings instead of Jinja. Enabled by |
| 7 | passing ``--tokenizer-mode cohere`` on the engine; tokenization itself |
| 8 | still flows through the cached HuggingFace tokenizer. |
| 9 | |
| 10 | The renderer's client surface is ``chat_template_kwargs`` (vLLM's |
| 11 | generic passthrough for template-time inputs). It is fed by two |
| 12 | different code paths, which speak different vocabularies: |
| 13 | |
| 14 | 1. :class:`vllm.entrypoints.cohere.serving.CohereServingChatV2` -- |
| 15 | translates a native `Cohere Chat v2 request |
| 16 | <https://docs.cohere.com/reference/chat>`__ into |
| 17 | ``chat_template_kwargs``. |
| 18 | 2. :class:`OpenAIServingChat` -- forwards whatever the client puts in |
| 19 | ``request.chat_template_kwargs`` directly, without translation. |
| 20 | |
| 21 | Accordingly the renderer accepts keys from both vocabularies. |
| 22 | |
| 23 | **Cohere Chat v2 request fields** (populated automatically by |
| 24 | ``CohereServingChatV2``; a direct ``chat_template_kwargs`` caller may |
| 25 | also pass them): |
| 26 | |
| 27 | * ``documents``: list of documents (v2 ``documents``) |
| 28 | * ``tools``: list of tools (v2 ``tools``) |
| 29 | * ``safety_mode``: v2 ``safety_mode`` (``CONTEXTUAL`` / ``STRICT`` / |
| 30 | ``NONE``) -- forwarded to melody's cmd3 ``safety_mode`` slot; |
| 31 | lowercased. |
| 32 | * ``citation_options``: v2 ``citation_options``. Its ``.mode`` |
| 33 | (``ENABLED`` / ``DISABLED`` / ``FAST`` / ``ACCURATE`` / ``OFF``) is |
| 34 | normalized into melody's cmd3 ``citation_quality`` (``on`` / ``off``) |
| 35 | and cmd4 ``grounding`` (``enabled`` / ``disabled`` / ``unknown``) |
| 36 | slots. cmd4 has no fast/accurate distinction at the prompt-template |
| 37 | layer -- both mean grounding-on. |
| 38 | * ``response_format``: v2 ``response_format`` (``json_object`` / |
| 39 | ``json_schema``) -- mapped onto melody's ``json_mode`` / |
| 40 | ``json_schema`` config slots. |
| 41 | * ``thinking``: v2 ``thinking``; its ``.type`` (``enabled`` / |
| 42 | ``disabled``) becomes melody's ``reasoning_type``. |
| 43 | |
| 44 | **Melody template-config knobs** (accepted for direct-passthrough |
| 45 | callers; not part of the Cohere Chat v2 surface): |
| 46 | |
| 47 | * ``cohere_format``: renderer selector, ``cmd3`` or ``cmd4`` (default |
| 48 | ``cmd4``). |
| 49 | * ``template_id``: pick one of melody's built-in template variants. |
| 50 | * ``available_tools``: melody's own name for tools; takes precedence |
| 51 | over ``tools`` when both are set. |
| 52 | * ``reasoning_type``: direct ``enabled`` / ``disabled`` toggle; |
| 53 | overrides derivation from ``thinking.type``. |
| 54 | * ``dev_instruction``: developer-instruction override. |
| 55 | * ``json_mode`` / ``json_schema``: direct structured-output toggles; |
| 56 | override derivation from ``response_format``. |
| 57 | * cmd3-only: ``citation_quality`` (``on`` / ``off``), ``skip_preamble``. |
| 58 | * cmd4-only: ``grounding`` (``enabled`` / ``disabled`` / ``unknown``; |
| 59 | overrides derivation from ``citation_options.mode``), |
| 60 | ``platform_instruction``. |
| 61 | |
| 62 | **Rejected inputs**: ``template_jinja`` and ``template`` are Cohere's |
| 63 | own inlets for raw Jinja template source. They are valid at Cohere's |
| 64 | API surface but not at vLLM's -- raw template source must flow through |
| 65 | the standard ``chat_template`` request field (guarded by |
| 66 | ``--trust-request-chat-template``), which this renderer forwards to |
| 67 | melody as ``template_jinja``. Setting these keys explicitly in |
| 68 | ``chat_template_kwargs`` raises ``ValueError`` so client |
| 69 | misconfiguration surfaces loudly. |
| 70 | |
| 71 | **Everything else** in ``chat_template_kwargs`` is forwarded verbatim |
| 72 | to melody as ``additional_template_fields`` and becomes accessible as |
| 73 | Jinja variables inside the template. This matches vLLM's documented |
| 74 | contract for ``chat_template_kwargs`` ("kwargs accessible by the |
| 75 | template"), so e.g. ``chat_template_kwargs={"reasoning_effort": |
| 76 | "low"}`` resolves ``{{ reasoning_effort }}`` inside cmd3 / cmd4 |
| 77 | templates. |
| 78 | |
| 79 | Citations produced by Cohere models are surfaced through the |
| 80 | Cohere-scoped ``CohereChatMessage.citations`` / |
| 81 | ``CohereDeltaMessage.citations`` fields (see |
| 82 | :mod:`vllm.entrypoints.cohere.cohere_chat_message`), populated by the |
| 83 | ``cohere2`` reasoning parser. The base OpenAI ``ChatMessage`` / |
| 84 | ``DeltaMessage`` keep their declared schemas unchanged; the response |
| 85 | envelope declares them as ``SerializeAsAny[...]`` so the subclass |
| 86 | fields survive JSON serialization. |
| 87 | """ |
| 88 | |
| 89 | from __future__ import annotations |
| 90 | |
| 91 | import copy |
| 92 | import json |
| 93 | from enum import Enum |
| 94 | from typing import Any |
| 95 | |
| 96 | from vllm.config import VllmConfig |
| 97 | from vllm.entrypoints.chat_utils import ( |
| 98 | ChatCompletionMessageParam, |
| 99 | ConversationMessage, |
| 100 | parse_chat_messages, |
| 101 | parse_chat_messages_async, |
| 102 | ) |
| 103 | from vllm.logger import init_logger |
| 104 | from vllm.tokenizers.hf import HfTokenizer |
| 105 | from vllm.utils.async_utils import make_async |
| 106 | |
| 107 | from .base import BaseRenderer |
| 108 | from .inputs import DictPrompt |
| 109 | from .inputs.preprocess import parse_dec_only_prompt |
| 110 | from .params import ChatParams |
| 111 | |
| 112 | logger = init_logger(__name__) |
| 113 | |
| 114 | |
| 115 | _DEFAULT_FORMAT = "cmd4" |
| 116 | _VALID_FORMATS = ("cmd3", "cmd4") |
| 117 | |
| 118 | |
| 119 | class MelodyContentType(str, Enum): |
| 120 | """Wire-format discriminator for melody content blocks. |
| 121 | |
| 122 | These strings are what the cmd3 / cmd4 Jinja templates check against |
| 123 | in ``message.content[0].type`` (e.g. the ``thinking`` branch in |
| 124 | ``cmd4-v1.jinja``). Keep new values in sync with melody's template |
| 125 | schema. |
| 126 | """ |
| 127 | |
| 128 | TEXT = "text" |
| 129 | THINKING = "thinking" |
| 130 | IMAGE = "image" |
| 131 | DOCUMENT = "document" |
| 132 | |
| 133 | |
| 134 | # Reserved ``chat_template_kwargs`` key that carries |
| 135 | # ``AssistantChatMessageV2.citations`` from a Cohere v2 request through |
| 136 | # to the renderer, keyed by request-message index. Populated by |
| 137 | # ``CohereServingChatV2._apply_cohere_template_kwargs`` and consumed by |
| 138 | # ``_conversation_to_melody_messages``. The value is a |
| 139 | # ``dict[int, list[dict]]`` of melody-shape ``FilterCitation`` dicts. |
| 140 | # The leading underscore signals "internal, not part of the public |
| 141 | # ``chat_template_kwargs`` surface" -- clients should not set this key |
| 142 | # directly. |
| 143 | MESSAGES_CITATIONS_KEY = "_messages_citations" |
| 144 | |
| 145 | |
| 146 | # Reserved ``chat_template_kwargs`` key carrying a request-side |
| 147 | # ``dict[tuple[int, int], CitationSource]`` map from melody's numeric |
| 148 | # ``(tool_call_index, tool_result_index)`` addressing to fully-resolved |
| 149 | # wire-shape sources (type, id, document/tool_output payload). |
| 150 | # Populated by :meth:`CohereServingChatV2._apply_cohere_template_kwargs` |
| 151 | # and consumed by :class:`BaseCohereCommandReasoningParser` at request |
| 152 | # start so the parser can attach real document ids on each melody source |
| 153 | # it emits -- eliminating the numeric-coord round-trip through the |
| 154 | # internal streaming pipeline. Leading underscore signals "internal, |
| 155 | # not part of the public ``chat_template_kwargs`` surface" -- clients |
| 156 | # should not set this key directly. |
| 157 | POSITION_TO_SOURCE_KEY = "_position_to_source" |
| 158 | |
| 159 | |
| 160 | # Reserved ``chat_template_kwargs`` key carrying the original Cohere v2 |
| 161 | # tool-message content, keyed by request-message index. Populated by |
| 162 | # :meth:`CohereServingChatV2._apply_cohere_template_kwargs` and consumed |
| 163 | # by :func:`_conversation_to_melody_messages`, which runs each block |
| 164 | # through :func:`_v2_tool_content_to_melody_block` to produce the |
| 165 | # melody-shape content list the renderer emits. |
| 166 | # |
| 167 | # Motivation: OpenAI's ``ChatCompletionContentPartParam`` union (in |
| 168 | # ``vllm/entrypoints/chat_utils.py``) only allows ``text`` / ``image_url`` |
| 169 | # / ``audio`` / ``video`` / ``file`` content parts on tool messages, so |
| 170 | # ``{"type": "document", "document": {...}}`` blocks -- valid in Cohere |
| 171 | # v2 -- get rejected by vLLM's shared chat validator before they can |
| 172 | # reach the renderer. |
| 173 | # ``CohereServingChatV2._convert_tool_message`` works around that by |
| 174 | # rewriting each content block into a plain ``text`` part (with document |
| 175 | # blocks JSON-serialized inline), but that flattening loses the |
| 176 | # structural distinction the cohere renderer needs to shape tool-result |
| 177 | # payloads for cmd3/cmd4. The serving layer therefore forwards the |
| 178 | # original v2 content blocks verbatim (as plain dicts) under this key, |
| 179 | # and the renderer owns the v2-shape → melody-shape mapping. |
| 180 | # |
| 181 | # Value type: ``dict[int, list[dict[str, Any]]]``. Each list entry is a |
| 182 | # ``model_dump``'d ``TextToolContent`` / ``DocumentToolContent`` block |
| 183 | # (e.g. ``{"type": "text", "text": ...}`` or ``{"type": "document", |
| 184 | # "document": {"id": ..., "data": {...}}}``). Index is the position of |
| 185 | # the tool message in the request's ``messages`` list, which -- since |
| 186 | # ``_convert_messages`` is 1-to-1 -- also matches its position in the |
| 187 | # ``ConversationMessage`` list handed to the renderer. Leading |
| 188 | # underscore signals "internal, not part of the public |
| 189 | # ``chat_template_kwargs`` surface" -- clients should not set this key |
| 190 | # directly. |
| 191 | TOOL_MESSAGE_V2_CONTENT_KEY = "_tool_message_v2_content" |
| 192 | |
| 193 | |
| 194 | # Keys this renderer interprets directly from ``chat_template_kwargs`` and |
| 195 | # maps onto typed melody render-config fields. Everything *not* in this |
| 196 | # set is forwarded verbatim to melody as ``additional_template_fields`` |
| 197 | # (i.e. as Jinja template variables), so callers can write |
| 198 | # ``chat_template_kwargs = {"my_var": "..."}`` and have ``{{ my_var }}`` |
| 199 | # resolve inside the template -- matching vLLM's documented contract |
| 200 | # for ``chat_template_kwargs``. |
| 201 | _RENDERER_CONSUMED_KEYS = frozenset( |
| 202 | { |
| 203 | "cohere_format", |
| 204 | "template_id", |
| 205 | # ``template_jinja`` / ``template`` are Cohere-only inputs for |
| 206 | # raw Jinja source -- valid at Cohere's own API surface but not |
| 207 | # at vLLM's. In vLLM, raw template source must flow through the |
| 208 | # standard ``chat_template`` request field, so a non-``None`` |
| 209 | # value here is rejected in ``_build_render_config``. The keys |
| 210 | # stay in the consumed set so an explicit ``None`` (e.g. from |
| 211 | # ``.model_dump(exclude_none=False)``) is dropped rather than |
| 212 | # surfacing as a stray ``{{ template_jinja }}`` Jinja variable. |
| 213 | "template_jinja", |
| 214 | "template", |
| 215 | "use_jinja", |
| 216 | "documents", |
| 217 | "available_tools", |
| 218 | "tools", |
| 219 | "reasoning_type", |
| 220 | "thinking", |
| 221 | "dev_instruction", |
| 222 | "response_format", |
| 223 | "json_schema", |
| 224 | "json_mode", |
| 225 | "safety_mode", |
| 226 | "citation_quality", |
| 227 | "citation_options", |
| 228 | "skip_preamble", |
| 229 | "grounding", |
| 230 | "platform_instruction", |
| 231 | MESSAGES_CITATIONS_KEY, |
| 232 | POSITION_TO_SOURCE_KEY, |
| 233 | TOOL_MESSAGE_V2_CONTENT_KEY, |
| 234 | } |
| 235 | ) |
| 236 | |
| 237 | |
| 238 | def _try_import_melody(): |
| 239 | try: |
| 240 | import cohere_melody # type: ignore |
| 241 | |
| 242 | return cohere_melody |
| 243 | except ImportError as e: # pragma: no cover - exercised at runtime |
| 244 | raise ImportError( |
| 245 | "The `cohere` tokenizer/renderer mode requires the " |
| 246 | "`cohere_melody` package. Install it via " |
| 247 | "`pip install 'cohere-melody>=0.14.0'` or build from " |
| 248 | "https://github.com/cohere-ai/melody." |
| 249 | ) from e |
| 250 | |
| 251 | |
| 252 | _MELODY_ROLES = frozenset({"system", "user", "chatbot", "tool"}) |
| 253 | |
| 254 | |
| 255 | # Cohere v2 ``citation_options.mode`` -> melody cmd4 ``grounding``. |
| 256 | # v2 surfaces five modes (``ENABLED`` / ``DISABLED`` / ``FAST`` / |
| 257 | # ``ACCURATE`` / ``OFF``) but cmd4's template doesn't differentiate |
| 258 | # fast vs accurate at the prompt layer -- both just turn grounding on, |
| 259 | # and ``ENABLED`` / ``DISABLED`` are direct aliases. Melody itself |
| 260 | # accepts ``unknown`` / ``enabled`` / ``disabled``. |
| 261 | _CMD4_GROUNDING_FROM_MODE = { |
| 262 | "fast": "enabled", |
| 263 | "accurate": "enabled", |
| 264 | "on": "enabled", |
| 265 | "enabled": "enabled", |
| 266 | "off": "disabled", |
| 267 | "disabled": "disabled", |
| 268 | "unknown": "unknown", |
| 269 | } |
| 270 | |
| 271 | |
| 272 | def _normalize_cmd4_grounding(value: Any) -> str: |
| 273 | """Coerce a user-facing grounding/citation mode to melody's vocab. |
| 274 | |
| 275 | Raises ``ValueError`` rather than letting an unrecognized value |
| 276 | slip through to ``render_cmd4`` and surface as a generic |
| 277 | ``Invalid config: grounding`` from melody. |
| 278 | """ |
| 279 | out = _CMD4_GROUNDING_FROM_MODE.get(value.lower()) |
| 280 | if out is None: |
| 281 | raise ValueError( |
| 282 | f"Unrecognized cmd4 grounding value: {value!r}. Expected one " |
| 283 | f"of ENABLED / DISABLED / FAST / ACCURATE / OFF (from " |
| 284 | f"citation_options.mode), or the direct melody values " |
| 285 | f"enabled / disabled / unknown." |
| 286 | ) |
| 287 | return out |
| 288 | |
| 289 | |
| 290 | def _role_to_melody(role: str) -> str: |
| 291 | """Map an OpenAI role to the role string melody expects. |
| 292 | |
| 293 | The cmd3 / cmd4 jinja templates only recognize ``system``, ``user``, |
| 294 | ``assistant``/``chatbot``, and ``tool``: any other role is silently |
| 295 | dropped by the template's role-dispatch chain (no fallback branch), |
| 296 | which produces a malformed prompt without any error. We therefore |
| 297 | refuse unknown roles up front rather than letting them disappear. |
| 298 | |
| 299 | Aliases: |
| 300 | |
| 301 | * ``assistant`` -> ``chatbot`` (Cohere's historical assistant role |
| 302 | name used by the templates). |
| 303 | * ``developer`` -> ``system`` (OpenAI's ``developer`` role is |
| 304 | documented as high-priority instructions, which maps onto the |
| 305 | ``system`` slot in Cohere's prompt format). |
| 306 | """ |
| 307 | role = role.lower() |
| 308 | if role == "assistant": |
| 309 | return "chatbot" |
| 310 | if role == "developer": |
| 311 | return "system" |
| 312 | if role in _MELODY_ROLES: |
| 313 | return role |
| 314 | raise ValueError( |
| 315 | f"Unsupported message role for the cohere renderer: {role!r}. " |
| 316 | "Expected one of: system, developer, user, assistant, chatbot, tool." |
| 317 | ) |
| 318 | |
| 319 | |
| 320 | def _normalize_tool_call(tc: dict[str, Any] | Any) -> dict[str, Any]: |
| 321 | """Normalize an OpenAI tool call dict into melody's tool call shape. |
| 322 | |
| 323 | melody expects ``{id, name, parameters: <json str>}`` whereas OpenAI |
| 324 | delivers ``{id, type, function: {name, arguments: <json|str>}}``. |
| 325 | """ |
| 326 | # Pydantic objects |
| 327 | if hasattr(tc, "model_dump"): |
| 328 | tc = tc.model_dump() |
| 329 | if not isinstance(tc, dict): |
| 330 | raise TypeError(f"Unexpected tool_call value: {tc!r}") |
| 331 | |
| 332 | fn = tc.get("function") or {} |
| 333 | name = fn.get("name") or tc.get("name") or "" |
| 334 | args = fn.get("arguments") |
| 335 | if args is None: |
| 336 | args = tc.get("arguments", {}) |
| 337 | # melody expects a JSON-encoded string |
| 338 | if not isinstance(args, str): |
| 339 | args = json.dumps(args, ensure_ascii=False) |
| 340 | return { |
| 341 | "id": tc.get("id") or "", |
| 342 | "name": name, |
| 343 | "parameters": args, |
| 344 | } |
| 345 | |
| 346 | |
| 347 | def _content_blocks(content: Any) -> list[dict[str, Any]]: |
| 348 | """Convert OpenAI ``ConversationMessage.content`` into melody content blocks. |
| 349 | |
| 350 | The chat_utils ``content_format="openai"`` produces a list of |
| 351 | ``{type: text, text: ...}`` dicts already; we pass those through with |
| 352 | minimal coercion. Plain string content is wrapped in a single text block. |
| 353 | Image / multimodal placeholder dicts (``{"type": "image"}``) are |
| 354 | forwarded as-is so that templates can reference image placeholders that |
| 355 | the upstream tokenizer expands separately. |
| 356 | """ |
| 357 | if content is None: |
| 358 | return [] |
| 359 | if isinstance(content, str): |
| 360 | return [{"type": MelodyContentType.TEXT, "text": content}] |
| 361 | blocks: list[dict[str, Any]] = [] |
| 362 | for part in content: |
| 363 | if isinstance(part, str): |
| 364 | blocks.append({"type": MelodyContentType.TEXT, "text": part}) |
| 365 | continue |
| 366 | if not isinstance(part, dict): |
| 367 | raise TypeError(f"Unexpected content part: {part!r}") |
| 368 | part_type = part.get("type", MelodyContentType.TEXT) |
| 369 | if part_type in ("text", "input_text", "output_text", "refusal"): |
| 370 | blocks.append( |
| 371 | {"type": MelodyContentType.TEXT, "text": part.get("text", "")} |
| 372 | ) |
| 373 | elif part_type == MelodyContentType.THINKING: |
| 374 | blocks.append( |
| 375 | { |
| 376 | "type": MelodyContentType.THINKING, |
| 377 | "thinking": part.get("thinking", ""), |
| 378 | } |
| 379 | ) |
| 380 | elif part_type == MelodyContentType.IMAGE: |
| 381 | blocks.append( |
| 382 | { |
| 383 | "type": MelodyContentType.IMAGE, |
| 384 | "image": { |
| 385 | "template_placeholder": part.get( |
| 386 | "template_placeholder", "<image>" |
| 387 | ), |
| 388 | }, |
| 389 | } |
| 390 | ) |
| 391 | elif part_type == MelodyContentType.DOCUMENT: |
| 392 | doc = part.get("document") |
| 393 | if isinstance(doc, dict): |
| 394 | blocks.append({"type": MelodyContentType.DOCUMENT, "document": doc}) |
| 395 | else: |
| 396 | # Fall back to wrapping arbitrary string as text. |
| 397 | blocks.append({"type": MelodyContentType.TEXT, "text": json.dumps(doc)}) |
| 398 | elif part_type == "tool_reference": |
| 399 | # Tool references are rendered by name; emit as text so the |
| 400 | # renderer downstream is content-format agnostic. |
| 401 | blocks.append( |
| 402 | { |
| 403 | "type": MelodyContentType.TEXT, |
| 404 | "text": part.get("name") or part.get("text", ""), |
| 405 | } |
| 406 | ) |
| 407 | else: |
| 408 | # Unknown block type: render as text fallback. |
| 409 | text = part.get("text") or part.get(part_type) or "" |
| 410 | text_str = text if isinstance(text, str) else json.dumps(text) |
| 411 | blocks.append({"type": MelodyContentType.TEXT, "text": text_str}) |
| 412 | return blocks |
| 413 | |
| 414 | |
| 415 | def _v2_tool_content_to_melody_block(block: dict[str, Any]) -> dict[str, Any] | None: |
| 416 | """Convert one Cohere v2 tool-message content block into a melody |
| 417 | ``document`` content block. |
| 418 | |
| 419 | Cohere's chat v2 admits two shapes for tool-message content parts: |
| 420 | |
| 421 | * ``TextToolContent`` -- ``{"type": "text", "text": "..."}``. |
| 422 | * ``DocumentToolContent`` -- ``{"type": "document", "document": |
| 423 | {"id"?: str, "data": {...}}}``. |
| 424 | |
| 425 | Both are emitted as melody ``document`` blocks, mirroring the cohere |
| 426 | api, which routes every tool content part through the same |
| 427 | "tool output as document" conversion. The rules: |
| 428 | |
| 429 | * Text that parses to a JSON object becomes that object directly |
| 430 | (so structured tool output flows through as first-class fields). |
| 431 | * Text that doesn't parse as an object -- empty, non-JSON, |
| 432 | or JSON scalar/array -- gets wrapped as ``{"content": <text>}`` |
| 433 | so templates have a stable field name to look up. |
| 434 | * ``document`` blocks unwrap ``document.data`` (a flat map) into |
| 435 | the melody payload; ``document.id`` is intentionally dropped |
| 436 | here because outbound citation binding tracks it separately via |
| 437 | ``POSITION_TO_SOURCE_KEY``. |
| 438 | |
| 439 | Only the JSON-object case actually changes what gets rendered. The |
| 440 | cmd3/cmd4 templates already wrap a bare ``text`` content block as |
| 441 | ``{"content": <text>}`` themselves, so routing plain text through |
| 442 | here is a no-op on the prompt -- it is the JSON-object unwrap that |
| 443 | promotes structured tool output from an escaped string blob into |
| 444 | individually addressable (and citable) result fields. |
| 445 | |
| 446 | Note the payload is the *flat* field map: melody's ``Content`` |
| 447 | carries ``document`` as a plain map, so re-nesting it under a |
| 448 | ``data`` key would render a spurious ``{"data": {...}}`` wrapper |
| 449 | into the tool-result block. |
| 450 | |
| 451 | Returns ``None`` for unrecognised shapes; the caller should skip |
| 452 | those blocks (the citation-numbering walk in |
| 453 | ``CohereServingChatV2._walk_citable_positions`` advances its own |
| 454 | cursor over unknown blocks, so slot alignment is preserved). |
| 455 | """ |
| 456 | block_type = block.get("type") |
| 457 | if block_type == "text": |
| 458 | text = block.get("text") or "" |
| 459 | try: |
| 460 | parsed = json.loads(text) |
| 461 | except (json.JSONDecodeError, TypeError): |
| 462 | parsed = None |
| 463 | if isinstance(parsed, dict): |
| 464 | payload: dict[str, Any] = dict(parsed) |
| 465 | else: |
| 466 | payload = {"content": text} |
| 467 | return {"type": MelodyContentType.DOCUMENT, "document": payload} |
| 468 | if block_type == "document": |
| 469 | doc = block.get("document") |
| 470 | if not isinstance(doc, dict): |
| 471 | return None |
| 472 | data = doc.get("data") |
| 473 | payload = dict(data) if isinstance(data, dict) else {} |
| 474 | return {"type": MelodyContentType.DOCUMENT, "document": payload} |
| 475 | return None |
| 476 | |
| 477 | |
| 478 | def _document_to_melody(doc: Any) -> dict[str, Any]: |
| 479 | """Coerce a Cohere v2 document into melody's ``Document`` (dict) shape.""" |
| 480 | match doc: |
| 481 | case str(): |
| 482 | return {"text": doc} |
| 483 | # Cohere v2 wraps documents in {id, data: {...}}; the ``**rest`` |
| 484 | # capture keeps the outer ``id`` (and any other outer keys) so we |
| 485 | # can lift them onto the payload without losing them. |
| 486 | case {"data": dict() as data, **rest}: |
| 487 | payload = dict(data) |
| 488 | if (outer_id := rest.get("id")) is not None and "id" not in payload: |
| 489 | payload["id"] = outer_id |
| 490 | return payload |
| 491 | case dict(): |
| 492 | return dict(doc) |
| 493 | case _: |
| 494 | raise TypeError(f"Unsupported document type: {type(doc).__name__}") |
| 495 | |
| 496 | |
| 497 | def _tool_to_melody(tool: Any) -> dict[str, Any]: |
| 498 | """Coerce a chat completions tool definition into melody's ``Tool`` shape. |
| 499 | |
| 500 | Accepts either the raw OpenAI tool wrapper ``{type:"function", function: |
| 501 | {name, description, parameters}}`` or a flat ``{name, description, |
| 502 | parameters}`` dict (which is what melody itself expects). |
| 503 | """ |
| 504 | if hasattr(tool, "model_dump"): |
| 505 | tool = tool.model_dump() |
| 506 | if not isinstance(tool, dict): |
| 507 | raise TypeError(f"Unsupported tool type: {type(tool).__name__}") |
| 508 | if "function" in tool and isinstance(tool["function"], dict): |
| 509 | fn = tool["function"] |
| 510 | else: |
| 511 | fn = tool |
| 512 | return { |
| 513 | "name": fn.get("name", ""), |
| 514 | "description": fn.get("description", "") or "", |
| 515 | "parameters": fn.get("parameters") or {}, |
| 516 | } |
| 517 | |
| 518 | |
| 519 | def _conversation_to_melody_messages( |
| 520 | conversation: list[ConversationMessage], |
| 521 | messages_citations: dict[int, list[dict[str, Any]]] | None = None, |
| 522 | tool_message_v2_content: dict[int, list[dict[str, Any]]] | None = None, |
| 523 | ) -> list[dict[str, Any]]: |
| 524 | """Convert vLLM ``ConversationMessage``s into melody's message dict shape. |
| 525 | |
| 526 | ``messages_citations`` is an optional index-keyed lookup of |
| 527 | per-message citations from the request, populated by |
| 528 | ``CohereServingChatV2`` under the |
| 529 | :data:`MESSAGES_CITATIONS_KEY` chat_template_kwargs entry. Values |
| 530 | must already be in melody's ``FilterCitation`` dict shape (see |
| 531 | ``CohereServingChatV2._sdk_citation_to_melody``). |
| 532 | |
| 533 | ``tool_message_v2_content`` is an optional index-keyed lookup of |
| 534 | the original Cohere v2 tool-message content blocks (as plain |
| 535 | dicts), populated by ``CohereServingChatV2`` under the |
| 536 | :data:`TOOL_MESSAGE_V2_CONTENT_KEY` entry. When set for a given |
| 537 | index, its blocks -- run through |
| 538 | :func:`_v2_tool_content_to_melody_block` to apply the v2 → melody |
| 539 | mapping -- replace the block list this function would otherwise |
| 540 | derive from ``msg["content"]``. This lets the renderer own the |
| 541 | v2-shape → melody-shape conversion for the structured ``document`` |
| 542 | blocks cohere v2 allows on tool messages (which don't survive |
| 543 | OpenAI's content-parts validator on the way in, so |
| 544 | ``_convert_tool_message`` rewrites them into ``text`` parts). |
| 545 | |
| 546 | Both index lookups key by position in the *input* ``conversation`` |
| 547 | list, which currently maps 1-to-1 with |
| 548 | ``CohereChatV2Request.messages`` for text-only Cohere requests. |
| 549 | This mapping breaks silently if ``parse_chat_messages`` ever |
| 550 | expands a single input message into multiple entries (e.g. |
| 551 | multi-modal splitting), which is not the case today. |
| 552 | """ |
| 553 | out: list[dict[str, Any]] = [] |
| 554 | for i, msg in enumerate(conversation): |
| 555 | role = _role_to_melody(msg.get("role", "user")) |
| 556 | if ( |
| 557 | tool_message_v2_content is not None |
| 558 | and (v2_blocks := tool_message_v2_content.get(i)) is not None |
| 559 | ): |
| 560 | content_blocks = [] |
| 561 | for v2_block in v2_blocks: |
| 562 | mapped = _v2_tool_content_to_melody_block(v2_block) |
| 563 | if mapped is not None: |
| 564 | content_blocks.append(mapped) |
| 565 | else: |
| 566 | content_blocks = _content_blocks(msg.get("content")) |
| 567 | |
| 568 | # Treat reasoning content as a thinking block on assistant turns |
| 569 | # so multi-turn reasoning is preserved in the rendered prompt. |
| 570 | # |
| 571 | # Cohere models output thought as either a ``thinking`` |
| 572 | # block (reasoning models) or a ``tool_plan`` field (older non- |
| 573 | # reasoning Command models), but vLLM's ConversationMessage has a |
| 574 | # single unified ``reasoning`` field that drops that difference. |
| 575 | # The cmd3 / cmd4 jinja templates render ``thinking`` blocks in both |
| 576 | # the case of tool calls and regular thinking blocks, so this is fine |
| 577 | # from the renderer's perspective. |
| 578 | reasoning = msg.get("reasoning") or msg.get("reasoning_content") |
| 579 | if role == "chatbot" and reasoning: |
| 580 | content_blocks.insert( |
| 581 | 0, |
| 582 | {"type": MelodyContentType.THINKING, "thinking": reasoning}, |
| 583 | ) |
| 584 | |
| 585 | tool_calls = [_normalize_tool_call(tc) for tc in (msg.get("tool_calls") or [])] |
| 586 | |
| 587 | out_msg: dict[str, Any] = { |
| 588 | "role": role, |
| 589 | "content": content_blocks, |
| 590 | "tool_calls": tool_calls, |
| 591 | } |
| 592 | tool_call_id = msg.get("tool_call_id") |
| 593 | if tool_call_id: |
| 594 | out_msg["tool_call_id"] = tool_call_id |
| 595 | if messages_citations and (cites := messages_citations.get(i)): |
| 596 | out_msg["citations"] = cites |
| 597 | out.append(out_msg) |
| 598 | return out |
| 599 | |
| 600 | |
| 601 | def _build_render_config( |
| 602 | conversation: list[ConversationMessage], |
| 603 | chat_template_kwargs: dict[str, Any], |
| 604 | chat_template: str | None = None, |
| 605 | ) -> tuple[str, dict[str, Any]]: |
| 606 | """Build the ``render_cmd3`` / ``render_cmd4`` config dict. |
| 607 | |
| 608 | ``chat_template`` is the standard vLLM per-request template override |
| 609 | (from ``ChatCompletionRequest.chat_template`` / the server-side |
| 610 | ``--chat-template`` flag). When non-``None`` it is forwarded to melody |
| 611 | as ``template_jinja`` so a single vLLM-native path drives template |
| 612 | source through the trust guard. |
| 613 | |
| 614 | Returns ``(format, config_dict)`` where ``format`` is either ``"cmd3"`` |
| 615 | or ``"cmd4"``. |
| 616 | """ |
| 617 | fmt = chat_template_kwargs.get("cohere_format", _DEFAULT_FORMAT) |
| 618 | if fmt not in _VALID_FORMATS: |
| 619 | raise ValueError( |
| 620 | f"Invalid cohere_format={fmt!r}; expected one of {_VALID_FORMATS}" |
| 621 | ) |
| 622 | |
| 623 | # Reject Cohere-only inlets for raw Jinja source. These keys are |
| 624 | # valid at Cohere's own API surface but not at vLLM's: callers must |
| 625 | # use the standard vLLM ``chat_template`` request field so the |
| 626 | # ``--trust-request-chat-template`` guard applies uniformly. |
| 627 | for cohere_only_key in ("template_jinja", "template"): |
| 628 | if chat_template_kwargs.get(cohere_only_key) is not None: |
| 629 | raise ValueError( |
| 630 | f"chat_template_kwargs.{cohere_only_key!r} is not accepted " |
| 631 | "in vLLM (it is a Cohere-only field); pass the template " |
| 632 | "body via the standard 'chat_template' request field " |
| 633 | "instead (requires --trust-request-chat-template)." |
| 634 | ) |
| 635 | |
| 636 | config: dict[str, Any] = { |
| 637 | "messages": _conversation_to_melody_messages( |
| 638 | conversation, |
| 639 | chat_template_kwargs.get(MESSAGES_CITATIONS_KEY), |
| 640 | chat_template_kwargs.get(TOOL_MESSAGE_V2_CONTENT_KEY), |
| 641 | ), |
| 642 | } |
| 643 | |
| 644 | # ``template_id`` is a selector for one of melody's built-in template |
| 645 | # variants (not raw source), so it is safe to accept from the client. |
| 646 | if template_id := chat_template_kwargs.get("template_id"): |
| 647 | config["template_id"] = template_id |
| 648 | # Raw template source flows exclusively through the standard vLLM |
| 649 | # ``chat_template`` field so it is subject to the |
| 650 | # ``--trust-request-chat-template`` guard in OnlineRenderer. |
| 651 | if chat_template: |
| 652 | config["template_jinja"] = chat_template |
| 653 | # Only support Jinja with vllm |
| 654 | config["use_jinja"] = True |
| 655 | |
| 656 | # Documents |
| 657 | documents = chat_template_kwargs.get("documents") or [] |
| 658 | if documents: |
| 659 | config["documents"] = [_document_to_melody(d) for d in documents] |
| 660 | |
| 661 | # Tools - prefer explicit ``available_tools``, fall back to OpenAI ``tools`` |
| 662 | tools = ( |
| 663 | chat_template_kwargs.get("available_tools") |
| 664 | or chat_template_kwargs.get("tools") |
| 665 | or [] |
| 666 | ) |
| 667 | if tools: |
| 668 | config["available_tools"] = [_tool_to_melody(t) for t in tools] |
| 669 | |
| 670 | # Reasoning toggle (cmd3 + cmd4) |
| 671 | if (rt := chat_template_kwargs.get("reasoning_type")) is not None: |
| 672 | config["reasoning_type"] = str(rt) |
| 673 | elif "thinking" in chat_template_kwargs: |
| 674 | # Cohere v2 ``thinking: {type: enabled|disabled}`` shorthand |
| 675 | thinking = chat_template_kwargs["thinking"] |
| 676 | t = thinking.get("type") if isinstance(thinking, dict) else thinking |
| 677 | if t in ("enabled", "disabled"): |
| 678 | config["reasoning_type"] = t |
| 679 | |
| 680 | if (di := chat_template_kwargs.get("dev_instruction")) is not None: |
| 681 | config["dev_instruction"] = str(di) |
| 682 | |
| 683 | # JSON / structured outputs |
| 684 | if (rf := chat_template_kwargs.get("response_format")) is not None: |
| 685 | rf = rf.model_dump() if hasattr(rf, "model_dump") else dict(rf) |
| 686 | rf_type = rf.get("type") |
| 687 | if rf_type == "json_object": |
| 688 | config["json_mode"] = True |
| 689 | elif rf_type in ("json_schema", "json"): |
| 690 | schema = rf.get("schema") or rf.get("json_schema") |
| 691 | if isinstance(schema, dict) and "schema" in schema: |
| 692 | schema = schema["schema"] |
| 693 | if schema is not None: |
| 694 | config["json_schema"] = ( |
| 695 | schema if isinstance(schema, str) else json.dumps(schema) |
| 696 | ) |
| 697 | if (js := chat_template_kwargs.get("json_schema")) is not None: |
| 698 | config["json_schema"] = js if isinstance(js, str) else json.dumps(js) |
| 699 | if "json_mode" in chat_template_kwargs: |
| 700 | config["json_mode"] = bool(chat_template_kwargs["json_mode"]) |
| 701 | |
| 702 | if fmt == "cmd3": |
| 703 | if (sm := chat_template_kwargs.get("safety_mode")) is not None: |
| 704 | config["safety_mode"] = str(sm).lower() |
| 705 | # citation_quality: ``on`` / ``off`` |
| 706 | cq = chat_template_kwargs.get("citation_quality") |
| 707 | if cq is None and (co := chat_template_kwargs.get("citation_options")): |
| 708 | mode = co.get("mode") if isinstance(co, dict) else None |
| 709 | if mode is not None: |
| 710 | cq = "on" if str(mode).lower() != "off" else "off" |
| 711 | if cq is not None: |
| 712 | config["citation_quality"] = str(cq).lower() |
| 713 | if "skip_preamble" in chat_template_kwargs: |
| 714 | config["skip_preamble"] = bool(chat_template_kwargs["skip_preamble"]) |
| 715 | else: # cmd4 |
| 716 | # cmd4 uses ``grounding`` rather than safety_mode/citation_quality. |
| 717 | # melody's cmd4 only accepts ``unknown`` / ``enabled`` / ``disabled``, |
| 718 | # so the Cohere v2 ``citation_options.mode`` values |
| 719 | # (``FAST`` / ``ACCURATE`` / ``OFF``) have to be normalized -- a |
| 720 | # raw lowercased passthrough would raise from |
| 721 | # ``render_cmd4`` (cmd4 has no fast/accurate distinction at the |
| 722 | # prompt-template layer; both request grounding-on). |
| 723 | if (gr := chat_template_kwargs.get("grounding")) is not None: |
| 724 | config["grounding"] = _normalize_cmd4_grounding(gr) |
| 725 | elif co := chat_template_kwargs.get("citation_options"): |
| 726 | mode = co.get("mode") if isinstance(co, dict) else None |
| 727 | if mode is not None: |
| 728 | config["grounding"] = _normalize_cmd4_grounding(mode) |
| 729 | if (pi := chat_template_kwargs.get("platform_instruction")) is not None: |
| 730 | config["platform_instruction"] = str(pi) |
| 731 | |
| 732 | # Anything we didn't explicitly interpret above is forwarded to melody |
| 733 | # as a Jinja template variable. This matches vLLM's documented contract |
| 734 | # for ``chat_template_kwargs`` ("kwargs accessible by the template") |
| 735 | # and lets callers write e.g. ``{"reasoning_effort": "low"}`` directly |
| 736 | # without a nested ``additional_template_fields`` wrapper. |
| 737 | extra = { |
| 738 | k: v |
| 739 | for k, v in chat_template_kwargs.items() |
| 740 | if k not in _RENDERER_CONSUMED_KEYS |
| 741 | } |
| 742 | if extra: |
| 743 | config["additional_template_fields"] = extra |
| 744 | |
| 745 | return fmt, config |
| 746 | |
| 747 | |
| 748 | class CohereRenderer(BaseRenderer[HfTokenizer]): |
| 749 | """Renderer that templates Cohere prompts via the melody Rust bindings. |
| 750 | |
| 751 | Tokenization is delegated to the standard HF tokenizer; only the |
| 752 | chat-template step is replaced with ``cohere_melody.render_cmd3`` / |
| 753 | ``render_cmd4``. Enabled via ``--tokenizer-mode cohere``. |
| 754 | """ |
| 755 | |
| 756 | def __init__( |
| 757 | self, |
| 758 | config: VllmConfig, |
| 759 | tokenizer: HfTokenizer | None, |
| 760 | ) -> None: |
| 761 | # Match HfRenderer in not mutating the cached tokenizer instance |
| 762 | tokenizer = copy.copy(tokenizer) |
| 763 | super().__init__(config, tokenizer) |
| 764 | |
| 765 | # Lazy import to keep `cohere_melody` an optional dependency |
| 766 | self._melody = _try_import_melody() |
| 767 | # ``render_cmd3`` / ``render_cmd4`` are pure CPU work; cache the |
| 768 | # thread-pool wrapper once so the async path doesn't allocate a new |
| 769 | # adapter on every request. |
| 770 | self._render_async = make_async(self._render, executor=self._executor) |
| 771 | |
| 772 | def _render(self, fmt: str, config_dict: dict[str, Any]) -> str: |
| 773 | if fmt == "cmd3": |
| 774 | return self._melody.render_cmd3(config_dict) |
| 775 | return self._melody.render_cmd4(config_dict) |
| 776 | |
| 777 | def render_messages( |
| 778 | self, |
| 779 | messages: list[ChatCompletionMessageParam], |
| 780 | params: ChatParams, |
| 781 | ) -> tuple[list[ConversationMessage], DictPrompt]: |
| 782 | conversation, mm_data, mm_uuids = parse_chat_messages( |
| 783 | messages, |
| 784 | self.model_config, |
| 785 | content_format="openai", |
| 786 | media_io_kwargs=params.media_io_kwargs, |
| 787 | mm_processor_kwargs=params.mm_processor_kwargs, |
| 788 | ) |
| 789 | |
| 790 | chat_template_kwargs = dict(params.chat_template_kwargs) |
| 791 | fmt, config_dict = _build_render_config( |
| 792 | conversation, chat_template_kwargs, params.chat_template |
| 793 | ) |
| 794 | prompt_text = self._render(fmt, config_dict) |
| 795 | prompt = parse_dec_only_prompt(prompt_text) |
| 796 | |
| 797 | if mm_data is not None: |
| 798 | prompt["multi_modal_data"] = mm_data |
| 799 | if mm_uuids is not None: |
| 800 | prompt["multi_modal_uuids"] = mm_uuids |
| 801 | |
| 802 | return conversation, prompt |
| 803 | |
| 804 | async def render_messages_async( |
| 805 | self, |
| 806 | messages: list[ChatCompletionMessageParam], |
| 807 | params: ChatParams, |
| 808 | ) -> tuple[list[ConversationMessage], DictPrompt]: |
| 809 | conversation, mm_data, mm_uuids = await parse_chat_messages_async( |
| 810 | messages, |
| 811 | self.model_config, |
| 812 | content_format="openai", |
| 813 | media_io_kwargs=params.media_io_kwargs, |
| 814 | mm_processor_kwargs=params.mm_processor_kwargs, |
| 815 | ) |
| 816 | |
| 817 | chat_template_kwargs = dict(params.chat_template_kwargs) |
| 818 | fmt, config_dict = _build_render_config( |
| 819 | conversation, chat_template_kwargs, params.chat_template |
| 820 | ) |
| 821 | prompt_text = await self._render_async(fmt, config_dict) |
| 822 | prompt = parse_dec_only_prompt(prompt_text) |
| 823 | |
| 824 | if mm_data is not None: |
| 825 | prompt["multi_modal_data"] = mm_data |
| 826 | if mm_uuids is not None: |
| 827 | prompt["multi_modal_uuids"] = mm_uuids |
| 828 | |
| 829 | return conversation, prompt |