vllm.entrypoints.cohere.serving ¶

class CohereServingChatV2(OpenAIServingChat):
    """Handler for the Cohere Chat v2 API (``POST /cohere/v2/chat``).

    The handler is intentionally thin: it converts the v2 request into a
    :class:`ChatCompletionRequest` (preserving Cohere-specific fields such
    as ``documents``, ``safety_mode`` and ``citation_options`` via
    ``chat_template_kwargs``) and delegates to the underlying chat
    completion machinery. All Cohere-specific templating happens in the
    renderer (:class:`vllm.renderers.cohere.CohereRenderer`) when the
    engine is started with ``--tokenizer-mode cohere``.

    Citations flow through ``CohereChatMessage`` / ``CohereDeltaMessage``
    subclasses (see :mod:`vllm.entrypoints.cohere.cohere_chat_message`) so
    the OpenAI-shared :class:`ChatMessage` / :class:`DeltaMessage` keep
    their declared schemas unchanged. Non-streaming responses pick up
    the subclass by overriding :meth:`_create_chat_message`; streaming
    deltas are emitted directly by the Cohere reasoning parser.
    """

    def __init__(
        self,
        engine_client: EngineClient,
        models: OpenAIServingModels,
        response_role: str,
        *,
        online_renderer: OnlineRenderer,
        request_logger: RequestLogger | None,
        chat_template: str | None,
        chat_template_content_format: ChatTemplateContentFormatOption,
        return_tokens_as_token_ids: bool = False,
        reasoning_parser: str = "",
        enable_auto_tools: bool = False,
        tool_parser: str | None = None,
        enable_prompt_tokens_details: bool = False,
        enable_force_include_usage: bool = False,
        default_chat_template_kwargs: dict[str, Any] | None = None,
        is_reasoning_model: bool = True,
    ) -> None:
        super().__init__(
            engine_client=engine_client,
            models=models,
            response_role=response_role,
            online_renderer=online_renderer,
            request_logger=request_logger,
            chat_template=chat_template,
            chat_template_content_format=chat_template_content_format,
            return_tokens_as_token_ids=return_tokens_as_token_ids,
            reasoning_parser=reasoning_parser,
            enable_auto_tools=enable_auto_tools,
            tool_parser=tool_parser,
            enable_prompt_tokens_details=enable_prompt_tokens_details,
            enable_force_include_usage=enable_force_include_usage,
            default_chat_template_kwargs=default_chat_template_kwargs,
        )
        # Controls how the assistant's chain-of-thought is surfaced on
        # turns that also contain tool calls.
        #
        # - ``True`` (default): the model is a reasoning Command-family
        #   model; reasoning is always surfaced as a ``thinking`` content
        #   block (non-streaming) or as ``content-start`` / ``content-
        #   delta`` events for a thinking block (streaming), regardless of
        #   whether tool calls also appear.
        # - ``False``: the model is an older non-reasoning Command model
        #   that uses Cohere's ``tool_plan`` field for its chain-of-
        #   thought before tool calls; reasoning is surfaced as
        #   ``tool_plan`` (non-streaming) or as ``tool-plan-delta`` events
        #   (streaming) on tool-call turns, and the thinking content block
        #   is dropped.
        #
        # TODO: replace this manual flag with automatic detection from the
        # model's capabilities config once that exists.
        self._is_reasoning_model = is_reasoning_model

    # ------------------------------------------------------------------
    # Public entry point
    # ------------------------------------------------------------------

    async def create_chat_v2(
        self,
        request: CohereChatV2Request,
        raw_request: Request | None = None,
    ) -> AsyncGenerator[str, None] | CohereChatV2Response | ErrorResponse:
        """Implements ``POST /cohere/v2/chat``."""
        if logger.isEnabledFor(logging.DEBUG):
            logger.debug(
                "Received Cohere v2 chat request %s", request.model_dump_json()
            )

        chat_req = self._convert_v2_to_chat_completion(request)
        if logger.isEnabledFor(logging.DEBUG):
            logger.debug(
                "Converted Cohere v2 -> ChatCompletion: %s",
                chat_req.model_dump_json(),
            )

        generator = await self.create_chat_completion(chat_req, raw_request)

        match generator:
            case ErrorResponse():
                return generator
            case ChatCompletionResponse():
                return self._chat_completion_to_v2(generator, request)
            case _:
                return self._chat_completion_stream_to_v2(generator, request)

    # ==================================================================
    # Request conversion: Cohere V2 -> ChatCompletionRequest
    # ==================================================================

    @classmethod
    def _convert_v2_to_chat_completion(
        cls, request: CohereChatV2Request
    ) -> ChatCompletionRequest:
        openai_messages: list[dict[str, Any]] = []
        cls._convert_messages(request.messages, openai_messages)

        chat_req = cls._build_base_chat_completion(request, openai_messages)
        cls._apply_streaming_options(chat_req, request)
        cls._apply_response_format(chat_req, request)
        cls._apply_tools(chat_req, request)
        cls._apply_tool_choice(chat_req, request)
        cls._apply_cohere_template_kwargs(chat_req, request)
        return chat_req

    @classmethod
    def _convert_messages(
        cls,
        messages: list,
        openai_messages: list[dict[str, Any]],
    ) -> None:
        for msg in messages:
            match msg:
                case SystemChatMessageV2():
                    openai_messages.append(
                        {
                            "role": "system",
                            "content": cls._coerce_text_content(msg.content),
                        }
                    )
                case UserChatMessageV2():
                    openai_messages.append(cls._convert_user_message(msg))
                case AssistantChatMessageV2():
                    openai_messages.append(cls._convert_assistant_message(msg))
                case ToolChatMessageV2():
                    openai_messages.append(cls._convert_tool_message(msg))
                case _:  # pragma: no cover - guarded by Pydantic discriminator
                    raise ValueError(f"Unsupported Cohere v2 message: {msg!r}")

    @staticmethod
    def _coerce_text_content(content: str | list[Any]) -> str:
        if isinstance(content, str):
            return content
        parts: list[str] = []
        for block in content:
            text = getattr(block, "text", None)
            if text:
                parts.append(text)
        return "".join(parts)

    @classmethod
    def _convert_user_message(cls, msg: UserChatMessageV2) -> dict[str, Any]:
        if isinstance(msg.content, str):
            return {"role": "user", "content": msg.content}

        # Discriminate by ``type`` rather than isinstance so we don't have
        # to import every individual ``*Content`` variant from the SDK -
        # the union (``UserMessageV2Content``) covers both ``TextContent``
        # and ``ImageUrlContent`` and both expose ``type`` as a Literal.
        content_parts: list[dict[str, Any]] = []
        for block in msg.content:
            if block.type == "text":
                content_parts.append({"type": "text", "text": block.text})
            elif block.type == "image_url":
                image_url: dict[str, Any] = {"url": block.image_url.url}
                if getattr(block.image_url, "detail", None) is not None:
                    image_url["detail"] = block.image_url.detail
                content_parts.append(
                    {
                        "type": "image_url",
                        "image_url": image_url,
                    }
                )
        if len(content_parts) == 1 and content_parts[0]["type"] == "text":
            return {"role": "user", "content": content_parts[0]["text"]}
        return {"role": "user", "content": content_parts}

    @classmethod
    def _convert_assistant_message(cls, msg: AssistantChatMessageV2) -> dict[str, Any]:
        out: dict[str, Any] = {"role": "assistant"}

        # Cohere splits reasoning out into ``thinking`` content blocks. We
        # collapse them back into the OpenAI ``reasoning`` field for the
        # downstream chat template, while text blocks become ``content``.
        text_parts: list[str] = []
        thinking_parts: list[str] = []
        if isinstance(msg.content, str):
            text_parts.append(msg.content)
        elif msg.content is not None:
            for block in msg.content:
                if block.type == "text":
                    text_parts.append(block.text)
                elif block.type == "thinking":
                    thinking_parts.append(block.thinking)

        # ``tool_plan`` is Cohere's chain-of-thought emitted alongside tool
        # calls; preserve it as reasoning so templates that expect a
        # planning block still see it.
        if msg.tool_plan:
            thinking_parts.append(msg.tool_plan)

        if text_parts:
            out["content"] = "".join(text_parts)
        if thinking_parts:
            out["reasoning"] = "".join(thinking_parts)

        if msg.tool_calls:
            out["tool_calls"] = [
                {
                    "id": tc.id,
                    "type": "function",
                    "function": {
                        "name": (tc.function.name if tc.function else "") or "",
                        "arguments": (tc.function.arguments if tc.function else None)
                        or "{}",
                    },
                }
                for tc in msg.tool_calls
            ]
        return out

    @classmethod
    def _convert_tool_message(cls, msg: ToolChatMessageV2) -> dict[str, Any]:
        if isinstance(msg.content, str):
            return {
                "role": "tool",
                "tool_call_id": msg.tool_call_id,
                "content": msg.content,
            }

        # When the tool result is text-only, flatten to a string for
        # maximum compatibility with standard chat templates. When it
        # includes documents, preserve them as structured content parts
        # so the cohere renderer can surface them as grounding sources
        # (the :class:`CohereRenderer` understands
        # ``{type: document, document: {...}}`` blocks). Non-cohere
        # renderers may not honor document blocks, but that matches
        # the broader "documents are no-op for OSS models" contract
        # documented on this endpoint.
        #
        # Tool message content uses ``ToolMessageV2Content``, which is a
        # union of ``TextToolContent`` and ``DocumentToolContent`` -
        # distinct from the user-message text/image union. We discriminate
        # on the ``type`` literal so we don't have to import each variant.
        has_documents = any(block.type == "document" for block in msg.content)
        if not has_documents:
            text = "\n".join(
                block.text for block in msg.content if block.type == "text"
            )
            return {
                "role": "tool",
                "tool_call_id": msg.tool_call_id,
                "content": text,
            }

        parts: list[dict[str, Any]] = []
        for block in msg.content:
            if block.type == "text":
                parts.append({"type": "text", "text": block.text})
            elif block.type == "document":
                parts.append(
                    {
                        "type": "document",
                        "document": block.document.model_dump(exclude_none=True),
                    }
                )
        return {
            "role": "tool",
            "tool_call_id": msg.tool_call_id,
            "content": parts,
        }

    @classmethod
    def _build_base_chat_completion(
        cls,
        request: CohereChatV2Request,
        openai_messages: list[dict[str, Any]],
    ) -> ChatCompletionRequest:
        return ChatCompletionRequest(
            model=request.model,
            messages=openai_messages,
            max_tokens=request.max_tokens,
            max_completion_tokens=request.max_tokens,
            stop=request.stop_sequences,
            temperature=request.temperature,
            top_p=request.p,
            top_k=request.k,
            seed=request.seed,
            frequency_penalty=request.frequency_penalty,
            presence_penalty=request.presence_penalty,
            logprobs=request.logprobs,
            priority=request.priority or 0,
            kv_transfer_params=request.kv_transfer_params,
            chat_template_kwargs=request.chat_template_kwargs,
        )

    @classmethod
    def _apply_streaming_options(
        cls,
        chat_req: ChatCompletionRequest,
        request: CohereChatV2Request,
    ) -> None:
        if request.stream:
            chat_req.stream = True
            chat_req.stream_options = StreamOptions.model_validate(
                {"include_usage": True}
            )

    @classmethod
    def _apply_response_format(
        cls,
        chat_req: ChatCompletionRequest,
        request: CohereChatV2Request,
    ) -> None:
        rf = request.response_format
        if rf is None or rf.type == "text":
            return
        chat_req.response_format = ResponseFormat(
            type="json_schema" if rf.json_schema else "json_object",
            json_schema=(
                JsonSchemaResponseFormat(
                    name="cohere_v2_json_schema",
                    json_schema=rf.json_schema,
                )
                if rf.json_schema
                else None
            ),
        )

    @classmethod
    def _apply_tools(
        cls,
        chat_req: ChatCompletionRequest,
        request: CohereChatV2Request,
    ) -> None:
        if not request.tools:
            return
        # Cohere's ``strict_tools`` is the spec-equivalent of OpenAI's
        # per-function ``strict`` flag: when true the API guarantees tool
        # call arguments match the declared JSON schema. We surface it in
        # both places so OpenAI-shaped consumers see it on the function
        # definition and the cohere renderer can read it back from
        # ``chat_template_kwargs`` for cmd3/cmd4 preamble selection.
        strict = bool(request.strict_tools)
        chat_req.tools = [
            ChatCompletionToolsParam.model_validate(
                {
                    "type": "function",
                    "function": {
                        "name": tool.function.name,
                        "description": tool.function.description,
                        "parameters": tool.function.parameters,
                        **({"strict": True} if strict else {}),
                    },
                }
            )
            for tool in request.tools
        ]

    @classmethod
    def _apply_tool_choice(
        cls,
        chat_req: ChatCompletionRequest,
        request: CohereChatV2Request,
    ) -> None:
        # Cohere v2's ``tool_choice`` only admits ``REQUIRED`` / ``NONE``
        # (see :data:`CohereToolChoice`); named-tool selection is not
        # part of the v2 spec. Map both onto OpenAI's equivalents.
        if request.tool_choice == "REQUIRED":
            chat_req.tool_choice = "required"
        elif request.tool_choice == "NONE":
            chat_req.tool_choice = "none"
        elif chat_req.tools:
            # Mirrors Cohere's "free choice" default when tools are present.
            chat_req.tool_choice = "auto"

    @classmethod
    def _apply_cohere_template_kwargs(
        cls,
        chat_req: ChatCompletionRequest,
        request: CohereChatV2Request,
    ) -> None:
        """Forward Cohere-specific request fields into ``chat_template_kwargs``.

        The :class:`vllm.renderers.cohere.CohereRenderer` consumes these
        kwargs to drive ``cohere_melody.render_cmd3`` / ``render_cmd4``.
        Other renderers ignore unknown kwargs, so this is also a no-op for
        non-cohere ``--tokenizer-mode`` settings.
        """
        kwargs = dict(chat_req.chat_template_kwargs or {})

        if request.documents:
            documents: list[dict[str, Any]] = []
            for idx, doc in enumerate(request.documents):
                if isinstance(doc, str):
                    documents.append({"id": f"doc_{idx}", "data": {"text": doc}})
                else:
                    documents.append(
                        {
                            "id": doc.id or f"doc_{idx}",
                            "data": (
                                doc.data
                                if isinstance(doc.data, dict)
                                else {"text": doc.data}
                            ),
                        }
                    )
            kwargs.setdefault("documents", documents)

        if request.safety_mode is not None:
            kwargs.setdefault("safety_mode", str(request.safety_mode).lower())
        if request.citation_options is not None:
            kwargs.setdefault(
                "citation_options",
                request.citation_options.model_dump(exclude_none=True),
            )
        if request.thinking is not None:
            kwargs.setdefault(
                "thinking",
                request.thinking.model_dump(exclude_none=True),
            )

        # ``strict_tools`` is intentionally NOT forwarded here. It's a
        # decoder-guidance flag that ``_apply_tools`` already maps onto
        # per-function OpenAI ``strict: true`` on ``chat_req.tools``;
        # ``cohere_melody.render_cmd3``/``render_cmd4`` have no
        # ``strict_tools`` knob, and the renderer's residual-forward path
        # would otherwise surface it as a Jinja variable ``{{ strict_tools }}``
        # that templates have no defined use for.

        # Citations on assistant messages in the request don't fit
        # OpenAI's ``ChatMessage`` / ``ConversationMessage`` schemas, so
        # we forward them under the reserved ``MESSAGES_CITATIONS_KEY``
        # chat_template_kwargs entry keyed by request-message index.
        # See ``_sdk_citation_to_melody`` for the id-resolution rules.
        citations_by_index = cls._collect_message_citations(request)
        if citations_by_index:
            kwargs.setdefault(MESSAGES_CITATIONS_KEY, citations_by_index)

        # Forward the ``(bucket, result_idx) -> resolved wire source``
        # map to the reasoning parser via chat_template_kwargs. The
        # parser (constructed with the same kwargs) uses it to attach
        # real document ids / types / payloads to every source it
        # emits, so citation deltas already carry SDK-shape sources by
        # the time they reach the streaming loop. See
        # ``_build_position_to_source`` for the numbering rule and
        # ``_melody_sources_to_vllm`` in the reasoning parser for the
        # consumer.
        position_to_source = cls._build_position_to_source(request)
        if position_to_source:
            kwargs.setdefault(POSITION_TO_SOURCE_KEY, position_to_source)

        if kwargs:
            chat_req.chat_template_kwargs = kwargs

    @classmethod
    def _collect_message_citations(
        cls, request: CohereChatV2Request
    ) -> dict[int, list[dict[str, Any]]]:
        """Extract ``AssistantChatMessageV2.citations`` keyed by message index.

        Returns an empty dict when no assistant message in the request
        carries any resolvable citations, so callers can use truthiness
        as the "should I emit the ``MESSAGES_CITATIONS_KEY`` entry?"
        test. Citations whose sources can't be resolved are dropped
        individually; messages that end up with zero surviving
        citations are omitted from the result.
        """
        doc_positions = cls._build_doc_id_to_prompt_position(
            request.messages, request.documents
        )
        out: dict[int, list[dict[str, Any]]] = {}
        for idx, msg in enumerate(request.messages):
            if not isinstance(msg, AssistantChatMessageV2) or not msg.citations:
                continue
            resolved: list[dict[str, Any]] = []
            for c in msg.citations:
                converted = cls._sdk_citation_to_melody(c, doc_positions)
                if converted is not None:
                    resolved.append(converted)
            if resolved:
                out[idx] = resolved
        return out

    @classmethod
    def _walk_citable_positions(
        cls, messages: list, documents: list | None
    ) -> list[_CitablePosition]:
        """Single-pass implementation of melody's citation numbering rule.

        Mirrors ``PromptRenderIds::from_messages`` in
        ``melody/src/templating/util.rs``. Emits one
        :class:`_CitablePosition` per document-bearing slot melody
        assigns while rendering the prompt:

        * **Bucket numbering.** Bucket ``0`` is reserved for the
          top-level ``documents`` array when present, and each unique
          ``tool_call_id`` claims the next integer on first-seen basis
          -- registered from whichever of the assistant's
          ``tool_calls[].id`` or the tool message's ``tool_call_id``
          appears first in the history.
        * **Per-bucket result-index slots.** Consumed by *every*
          content item, not just documents -- a text block occupies a
          slot too, and the source at that slot carries an id-less
          ``tool_output`` payload synthesized via
          :meth:`_text_to_tool_output_payload`, matching the cohere
          api's behavior for text-only tool content. Strings and
          empty contents each consume exactly one slot because the
          cohere renderer wraps them into a single text block before
          handing the message off to melody. Multiple tool messages
          sharing a ``tool_call_id`` accumulate into the same bucket,
          so later docs' slots are offset by the sum of previous
          messages' content lengths.
        * **Doc ids attached to each position.** Top-level docs
          contribute both the explicit ``id`` (if set) and the
          ``doc_{idx}`` fallback used for inbound addressing; tool
          document blocks contribute their single ``document.id``
          when present. The ``ids`` tuple is empty for slots that
          are only reachable positionally (text-only tool content,
          docs without a client-provided id) -- those slots still
          have a fully-populated :class:`InternalCitationSource`,
          just no client-visible id to cite them by.

        Ids alone key the inbound helper
        (:meth:`_build_doc_id_to_prompt_position`, which drives
        request-side citation resolution); the resolved
        :class:`InternalCitationSource` (with ``type`` / ``id`` /
        payload) keys the outbound helper
        (:meth:`_build_position_to_source`, whose result is forwarded
        to the reasoning parser via
        ``chat_template_kwargs[POSITION_TO_SOURCE_KEY]``).

        Cohere's on-wire semantics -- ``ToolSource.id`` and
        ``DocumentSource.id`` both carry a document id, with ``type``
        acting only as a payload-shape hint -- are the invariant this
        walk is mirroring.
        """
        out: list[_CitablePosition] = []
        tool_call_id_to_bucket: dict[str, int] = {}
        # Next unassigned ``tool_result_index`` per bucket. Advances
        # by ``len(content)`` (or 1 for string / missing content) for
        # every tool message, whether or not the content had any doc
        # ids in it -- matching melody's per-slot accounting.
        bucket_result_len: dict[int, int] = {}
        next_bucket = 0

        if documents:
            for idx, doc in enumerate(documents):
                # Every top-level doc is addressable inbound by a
                # synthetic ``doc_{idx}`` fallback (so clients can cite
                # positional docs). On the *outbound* wire we prefer
                # the client's explicit id when present, and omit the
                # ``id`` field entirely when it's not -- matching
                # cohere's ``DocumentSource`` where ``id`` is optional
                # and preserving "no id in, no id out" symmetry.
                fallback_id = f"doc_{idx}"
                if isinstance(doc, str):
                    payload: dict[str, Any] = {"text": doc}
                    explicit_id: str | None = None
                else:
                    payload = (
                        dict(doc.data)
                        if isinstance(doc.data, dict)
                        else {"text": doc.data}
                    )
                    explicit_id = doc.id or None
                    if explicit_id:
                        payload.setdefault("id", explicit_id)
                # Register explicit id first so it wins over the
                # fallback for ``setdefault``-style inbound lookup.
                ids: tuple[str, ...] = (
                    (explicit_id, fallback_id) if explicit_id else (fallback_id,)
                )
                out.append(
                    _CitablePosition(
                        bucket=0,
                        result_idx=idx,
                        ids=ids,
                        source=InternalCitationSource(
                            type="document", id=explicit_id, document=payload
                        ),
                    )
                )
            next_bucket = 1

        def register_bucket(tool_call_id: str) -> int:
            nonlocal next_bucket
            bucket = tool_call_id_to_bucket.get(tool_call_id)
            if bucket is None:
                bucket = next_bucket
                tool_call_id_to_bucket[tool_call_id] = bucket
                next_bucket += 1
            return bucket

        for msg in messages:
            if isinstance(msg, ToolChatMessageV2):
                if not msg.tool_call_id:
                    continue
                bucket = register_bucket(msg.tool_call_id)
                base = bucket_result_len.get(bucket, 0)
                if isinstance(msg.content, list):
                    for offset, block in enumerate(msg.content):
                        info = cls._tool_content_block_to_source(block)
                        if info is None:
                            continue
                        # Only real document ids get registered in the
                        # inbound-facing ``ids`` tuple; text-only tool
                        # content produces a source with ``id=None``
                        # to match the cohere api -- the text is
                        # exposed as ``tool_output.content`` but
                        # there's no id for clients to cite by.
                        pos_ids: tuple[str, ...] = (info.id,) if info.id else ()
                        out.append(
                            _CitablePosition(
                                bucket=bucket,
                                result_idx=base + offset,
                                ids=pos_ids,
                                source=info,
                            )
                        )
                    bucket_result_len[bucket] = base + len(msg.content)
                else:
                    # A string or missing content is wrapped by the
                    # cohere renderer into a single text block before
                    # melody sees it, so it consumes one slot. For
                    # string content, emit a synthetic id-less tool
                    # source carrying the text as its payload so a
                    # model citation against this slot still has
                    # something to attach on the wire (matching the
                    # id-less tool_output behavior for text-only
                    # ``ToolChatMessageV2`` content blocks). Missing
                    # content still advances the slot cursor but
                    # produces no source -- a citation against it
                    # would be dropped by ``_to_wire_citation``.
                    if isinstance(msg.content, str):
                        payload = cls._text_to_tool_output_payload(msg.content)
                        out.append(
                            _CitablePosition(
                                bucket=bucket,
                                result_idx=base,
                                ids=(),
                                source=InternalCitationSource(
                                    type="tool", id=None, tool_output=payload
                                ),
                            )
                        )
                    bucket_result_len[bucket] = base + 1
            elif isinstance(msg, AssistantChatMessageV2) and msg.tool_calls:
                for tc in msg.tool_calls:
                    if tc.id:
                        register_bucket(tc.id)

        return out

    @classmethod
    def _build_doc_id_to_prompt_position(
        cls, messages: list, documents: list | None
    ) -> dict[str, tuple[int, int]]:
        """Map every citable document id to its ``(bucket, result_index)``.

        Inbound helper: converts request-side citation ``source.id``
        refs into melody's numeric ``FilterCitation`` addressing (see
        :meth:`_sdk_citation_to_melody`). First-seen-id wins so an
        explicit id on a top-level document takes precedence over the
        ``doc_{idx}`` fallback that :meth:`_walk_citable_positions`
        also emits for the same slot.

        Thin wrapper over :meth:`_walk_citable_positions`; the
        numbering rule lives there.
        """
        doc_positions: dict[str, tuple[int, int]] = {}
        for pos in cls._walk_citable_positions(messages, documents):
            for pos_id in pos.ids:
                doc_positions.setdefault(pos_id, (pos.bucket, pos.result_idx))
        return doc_positions

    @staticmethod
    def _sdk_citation_to_melody(
        citation: Any,
        doc_positions: dict[str, tuple[int, int]],
    ) -> dict[str, Any] | None:
        """Convert a Cohere SDK ``Citation`` to melody's ``FilterCitation``.

        Melody's ``Source`` uses a two-level numeric address:
        ``tool_call_index`` picks a bucket (``0`` = top-level
        ``documents`` when present, ``1..N`` = tool results in history
        order) and ``tool_result_indices`` picks positions inside that
        bucket. The Cohere ``Citation.sources`` schema is per-document:
        each entry's ``id`` is the id of the specific document that
        grounds the citation (``type`` distinguishes payload shape, not
        id semantics). We resolve each source id against
        ``doc_positions`` and aggregate hits sharing a bucket into one
        melody ``Source`` so the renderer emits compact
        ``<co>text</co: N:[i,j]>`` markers.

        Returns ``None`` when any source id fails to resolve or when
        the input has no sources at all: emitting a citation with an
        empty ``sources=[]`` would render a malformed
        ``<co>text</co: :>`` marker, and silently attributing a
        citation to the wrong bucket is worse than dropping the
        ``<co>`` markup. This matches the cohere api's behavior.

        ``Source.document_ids`` is intentionally NOT populated: it is a
        parser-*output* lookup table (see ``PromptRenderIds`` in
        melody), and melody's deserializer ignores it on input.

        The Cohere ``type`` literal (``TEXT_CONTENT`` / ``THINKING_CONTENT``
        / ``PLAN``) collapses to melody's boolean ``is_thinking``.
        ``PLAN`` is treated as a thinking-style citation because cmd3 /
        cmd4 templates render both inline with the same ``<co>...</co>``
        markup.
        """
        raw_sources = getattr(citation, "sources", None) or []
        if not raw_sources:
            return None

        grouped: dict[int, list[int]] = {}
        for src in raw_sources:
            src_id = getattr(src, "id", None)
            if not src_id:
                return None
            position = doc_positions.get(src_id)
            if position is None:
                return None
            bucket, result_idx = position
            grouped.setdefault(bucket, []).append(result_idx)

        sources = [
            {"tool_call_index": bucket, "tool_result_indices": indices}
            for bucket, indices in grouped.items()
        ]

        cite_type = str(getattr(citation, "type", "") or "").upper()
        is_thinking = cite_type in ("THINKING_CONTENT", "PLAN")
        return {
            "start_index": getattr(citation, "start", None) or 0,
            "end_index": getattr(citation, "end", None) or 0,
            "text": getattr(citation, "text", None) or "",
            "sources": sources,
            "is_thinking": is_thinking,
        }

    # ==================================================================
    # Response conversion: ChatCompletion -> Cohere V2
    # ==================================================================

    def _chat_completion_to_v2(
        self,
        response: ChatCompletionResponse,
        request: CohereChatV2Request,
    ) -> CohereChatV2Response:
        choice = response.choices[0]
        msg = choice.message

        # Build content blocks as dicts; ``AssistantMessageResponse``
        # validates them into the proper discriminated union variants
        # (``Text/ThinkingAssistantMessageResponseContentItem``).
        content_blocks: list[dict[str, Any]] = []
        if msg.reasoning:
            content_blocks.append(
                {"type": ContentBlockType.THINKING, "thinking": msg.reasoning}
            )
        if msg.content:
            content_blocks.append({"type": ContentBlockType.TEXT, "text": msg.content})

        tool_calls: list[ToolCallV2] | None = None
        if msg.tool_calls:
            tool_calls = [
                ToolCallV2(
                    id=tc.id,
                    function=ToolCallV2Function(
                        name=tc.function.name,
                        arguments=tc.function.arguments,
                    ),
                )
                for tc in msg.tool_calls
            ]

        # Cohere's ``tool_plan`` is the planning text emitted before tool
        # calls on older, non-reasoning Command models. For those models
        # we surface ``reasoning`` as ``tool_plan`` and drop the thinking
        # block. Reasoning Command models emit a regular thinking block
        # alongside tool calls, so for them we leave the thinking block
        # in place and never set ``tool_plan``.
        tool_plan: str | None = None
        if not self._is_reasoning_model and tool_calls and msg.reasoning:
            tool_plan = msg.reasoning
            content_blocks = [
                blk
                for blk in content_blocks
                if blk.get("type") != ContentBlockType.THINKING
            ]

        assistant_msg = AssistantMessageResponse(
            content=content_blocks or None,
            tool_calls=tool_calls,
            tool_plan=tool_plan,
            citations=self._extract_citations_if_any(msg),
        )

        usage = self._build_usage(response)

        return CohereChatV2Response(
            id=response.id or f"chat_{int(time.time() * 1000)}",
            finish_reason=_map_finish_reason(choice.finish_reason),
            message=assistant_msg,
            usage=usage,
            kv_transfer_params=response.kv_transfer_params,
        )

    def _create_chat_message(self, *args: Any, **kwargs: Any) -> ChatMessage:
        """Route response construction through the citation-carrying subclass.

        Overrides :meth:`OpenAIServingChat._create_chat_message` so every
        non-streaming construction site in the base full-generator
        produces a :class:`CohereChatMessage`. That lets
        :meth:`_finalize_response_message` below set citations on
        ``message.citations`` directly, without relying on
        ``extra="allow"`` attribute injection.
        """
        return CohereChatMessage(*args, **kwargs)

    def _finalize_response_message(
        self,
        message: ChatMessage,
        *,
        parser: Parser | None,
    ) -> ChatMessage:
        """Copy grounding citations off the reasoning parser onto the message.

        The Cohere reasoning parser
        (:mod:`vllm.reasoning.cohere_command_reasoning_parser`) caches the
        citations produced by its most recent unary ``extract_reasoning``
        call on ``last_unary_citations``. We surface them here on
        :class:`CohereChatMessage` so downstream response conversion can
        pick them up without the base :class:`OpenAIServingChat` having to
        know about citations.
        """
        # ``_create_chat_message`` above guarantees the concrete type at
        # runtime. ``cast`` narrows the declared ``ChatMessage`` for
        # mypy without adding a runtime check we don't need: even in the
        # invariant-violated case (a subclass reset the factory back to
        # plain ``ChatMessage``), ``OpenAIBaseModel``'s ``extra="allow"``
        # config lets the ``.citations`` write survive into
        # ``model_dump`` via the extras bucket, so the wire is correct
        # either way.
        message = cast(CohereChatMessage, message)
        citations = getattr(
            getattr(parser, "reasoning_parser", None),
            "last_unary_citations",
            None,
        )
        if citations:
            message.citations = citations
        return message

    def _extract_citations_if_any(self, msg: Any) -> list[Citation] | None:
        """Coerce ``CohereChatMessage.citations`` into the Cohere v2 wire shape.

        :class:`CohereChatMessage` carries a
        ``citations: list[vllm...Citation] | None`` field populated by
        :meth:`_finalize_response_message` (which reads it off the
        reasoning parser's ``last_unary_citations`` cache). Sources are
        already fully resolved by the parser (see
        :func:`_melody_sources_to_vllm` in
        :mod:`vllm.reasoning.cohere_command_reasoning_parser`) using
        the position map forwarded via ``chat_template_kwargs``. This
        method's remaining job is:

        * Drop citations whose sources all failed to resolve (empty
          ``sources`` list) -- fail-closed policy matching the inbound
          ``_sdk_citation_to_melody``.
        * Rewrite ``THINKING_CONTENT`` -> ``PLAN`` for non-reasoning
          models, since ``_chat_completion_to_v2`` surfaces
          ``msg.reasoning`` as ``tool_plan`` on the same
          ``is_reasoning_model`` flag.
        * Coerce to :class:`cohere.types.Citation` for the wire.
        """
        raw = getattr(msg, "citations", None)
        if not raw:
            return None
        out: list[Citation] = []
        for c in raw:
            resolved = self._to_wire_citation(c)
            if resolved is not None:
                out.append(resolved)
        return out or None

    def _to_wire_citation(self, citation: Any) -> Citation | None:
        """Coerce one parser-produced citation into the SDK wire shape.

        Accepts both :class:`InternalCitation` instances (what the
        unary path receives directly from the parser via
        ``last_unary_citations``) and plain ``dict`` payloads (what
        the streaming path sees, because ``DeltaMessage.citations`` is
        an untyped extras field on the OpenAI wire protocol and dict
        shapes survive ``model_validate_json`` unchanged). Both flows
        end up carrying the same key names, so a small ``_get`` helper
        handles both without a shape-specific code path.
        """

        def _get(obj: Any, key: str, default: Any = None) -> Any:
            if isinstance(obj, dict):
                return obj.get(key, default)
            return getattr(obj, key, default)

        raw_sources = _get(citation, "sources") or []
        wire_sources: list[dict[str, Any]] = []
        for src in raw_sources:
            # Sources arrive pre-resolved from the parser (or as trusted
            # dicts on the streaming path). Match the cohere api's
            # behavior of preserving id-less sources (e.g. text-only
            # tool results): cohere's ``Source`` types treat ``id`` as
            # optional, so we simply drop the field via
            # ``exclude_none=True`` instead of dropping the source.
            if isinstance(src, dict):
                wire = {k: v for k, v in src.items() if v is not None}
                if wire:
                    wire_sources.append(wire)
            elif isinstance(src, InternalCitationSource):
                wire_sources.append(src.model_dump(exclude_none=True))
        if not wire_sources:
            return None

        cite_type = _get(citation, "type")
        if cite_type == "THINKING_CONTENT" and not self._is_reasoning_model:
            # Non-reasoning Command models surface the reasoning block
            # as ``tool_plan`` (see ``_chat_completion_to_v2``), so any
            # ``is_thinking`` citation from melody is actually a PLAN
            # citation on the wire.
            cite_type = "PLAN"

        payload = {
            "start": _get(citation, "start"),
            "end": _get(citation, "end"),
            "text": _get(citation, "text"),
            "sources": wire_sources,
            "type": cite_type,
        }
        try:
            return Citation.model_validate(payload)
        except Exception:  # pragma: no cover - defensive
            logger.warning(
                "Skipping malformed citation payload: %r", payload, exc_info=True
            )
            return None

    @classmethod
    def _build_position_to_source(
        cls, request: CohereChatV2Request
    ) -> dict[tuple[int, int], InternalCitationSource]:
        """Invert the request-side numbering into ``(bucket, idx) -> source``.

        Outbound helper: forwarded to the reasoning parser via
        ``chat_template_kwargs[POSITION_TO_SOURCE_KEY]`` so
        :func:`_melody_sources_to_vllm` can attach real document ids /
        types / payloads to each melody source it emits. Each entry
        carries the fully-resolved wire-shape info: ``type``
        (``"document"`` for the reserved top-level bucket 0, ``"tool"``
        for any tool-call bucket), ``id`` (explicit doc id, the
        ``doc_{idx}`` fallback, or ``None`` for id-less sources), and
        ``document`` / ``tool_output`` payload -- matching the cohere
        api's ``Source`` shape.

        Thin wrapper over :meth:`_walk_citable_positions`; the
        numbering rule lives there.
        """
        return {
            (pos.bucket, pos.result_idx): pos.source
            for pos in cls._walk_citable_positions(request.messages, request.documents)
        }

    @staticmethod
    def _text_to_tool_output_payload(text: str) -> dict[str, Any]:
        """Wrap a bare tool-content text into a ``tool_output`` payload.

        Matches the cohere api's handling of text-only tool content:
        attempt to interpret ``text`` as a JSON object first; fall
        back to ``{"content": <text>}`` so a text-only tool result
        still has a structured payload we can attach to a citation
        source.
        """
        if not text:
            return {}
        try:
            parsed = json.loads(text)
        except (json.JSONDecodeError, TypeError):
            return {"content": text}
        if isinstance(parsed, dict):
            return parsed
        return {"content": text}

    @classmethod
    def _tool_content_block_to_source(
        cls,
        block: Any,
    ) -> InternalCitationSource | None:
        """Extract the wire-shape source for one tool-message content block.

        Handles both cohere ``DocumentToolContent`` and
        ``TextToolContent`` blocks:

        * ``document`` blocks with an explicit id produce a
          ``tool`` source keyed by that id (payload is the document's
          ``data`` dict, or ``{"text": ...}`` for scalar data).
        * ``document`` blocks *without* an id and ``text`` blocks
          produce an id-less ``tool`` source: the cohere api emits
          the payload with an empty id in this case, and cohere's
          SDK treats ``ToolSource.id`` as optional -- so we simply
          omit ``id`` on the wire instead of inventing a synthetic
          one.

        Returns ``None`` only for shapes we don't recognise at all;
        that slot still gets counted (callers advance the bucket
        cursor over the block regardless) so numbering stays aligned
        with melody.
        """
        block_type = getattr(block, "type", None)
        if block_type == "document":
            doc: Document | None = getattr(block, "document", None)
            if doc is None:
                return None
            doc_id = doc.id or None
            data = doc.data
            if isinstance(data, dict):
                payload: dict[str, Any] = dict(data)
            elif data is None:
                payload = {}
            else:
                # Match the top-level ``documents`` handling (see
                # ``_apply_cohere_template_kwargs``): non-dict payloads
                # get wrapped as ``{"text": ...}`` so clients get a
                # renderable body instead of a bare id.
                payload = {"text": data}
            if doc_id:
                payload.setdefault("id", doc_id)
            return InternalCitationSource(
                type="tool",
                id=doc_id,
                tool_output=payload,
            )
        if block_type == "text":
            text = getattr(block, "text", None) or ""
            payload = cls._text_to_tool_output_payload(text)
            return InternalCitationSource(
                type="tool",
                id=None,
                tool_output=payload,
            )
        return None

    @staticmethod
    def _build_usage(response: ChatCompletionResponse) -> CohereUsage | None:
        if response.usage is None:
            return None
        prompt = response.usage.prompt_tokens
        completion = response.usage.completion_tokens or 0
        cached: int | None = None
        if response.usage.prompt_tokens_details is not None:
            cached = response.usage.prompt_tokens_details.cached_tokens
        return CohereUsage(
            billed_units=CohereUsageBilledUnits(
                input_tokens=prompt,
                output_tokens=completion,
            ),
            tokens=CohereUsageTokens(
                input_tokens=prompt,
                output_tokens=completion,
            ),
            cached_tokens=cached,
        )

    # ==================================================================
    # Stream conversion: chat completion stream -> Cohere V2 SSE events
    # ==================================================================

    async def _chat_completion_stream_to_v2(
        self,
        generator: AsyncGenerator[str, None],
        request: CohereChatV2Request,
    ) -> AsyncGenerator[str, None]:
        """Translate an OpenAI-style chat completion SSE stream into Cohere's
        v2 stream-event format.

        Cohere's v2 stream lifecycle is:

            message-start
              [content-start, content-delta..., content-end]*
              [tool-plan-delta]*
              [tool-call-start, tool-call-delta..., tool-call-end]*
              message-end
        """
        state = _StreamState()

        try:
            async for item in generator:
                if not item.startswith("data:"):
                    continue
                data_str = item[len("data:") :].strip().rstrip("\n")
                if not data_str:
                    continue
                if data_str == "[DONE]":
                    # OpenAI's stream terminator. Fall through to the
                    # post-loop cleanup so we always emit ``message-end``
                    # even if the usage-only chunk was skipped.
                    break

                chunk = ChatCompletionStreamResponse.model_validate_json(data_str)
                state.last_chunk_id = chunk.id

                if not state.started:
                    yield _emit(
                        MessageStartEvent(
                            id=chunk.id,
                            delta={"message": {"role": "assistant"}},
                        )
                    )
                    state.started = True

                # The final OpenAI chunk has no choices and only carries usage.
                if not chunk.choices:
                    for ev in self._close_open_blocks(state):
                        yield ev
                    yield self._build_message_end_event(
                        chunk_id=chunk.id,
                        finish_reason=state.finish_reason,
                        usage_chunk=chunk,
                    )
                    state.ended = True
                    continue

                choice = chunk.choices[0]
                if choice.finish_reason is not None:
                    state.finish_reason = choice.finish_reason

                delta = choice.delta

                # Reasoning -> thinking content block
                reasoning = getattr(delta, "reasoning", None) or getattr(
                    delta, "reasoning_content", None
                )
                if reasoning:
                    for ev in self._handle_thinking_delta(state, reasoning):
                        yield ev

                if delta.content:
                    for ev in self._handle_text_delta(state, delta.content):
                        yield ev

                if delta.tool_calls:
                    for ev in self._handle_tool_call_deltas(state, delta.tool_calls):
                        yield ev

                # Citations: a Cohere-specific extension on DeltaMessage that
                # the cohere renderer/parsers may populate.
                delta_citations = getattr(delta, "citations", None)
                if delta_citations:
                    for ev in self._handle_citation_deltas(state, delta_citations):
                        yield ev

        except Exception as exc:
            logger.exception("Error converting chat completion stream to v2")
            if state.started and not state.ended:
                yield _sse(
                    json.dumps(
                        {
                            "type": "message-end",
                            "delta": {
                                "error": sanitize_message(str(exc)),
                                "finish_reason": "ERROR",
                            },
                        }
                    )
                )
                state.ended = True
            yield _DONE_FRAME
            return

        # Normal completion or ``[DONE]``: ensure ``message-end`` is always
        # emitted. Upstream may close the stream without sending the final
        # usage-only chunk (e.g. on shutdown, or when ``[DONE]`` is the only
        # terminator); without this fallback Cohere clients would hang
        # waiting for the closing event.
        if state.started and not state.ended:
            for ev in self._close_open_blocks(state):
                yield ev
            yield self._build_message_end_event(
                chunk_id=state.last_chunk_id,
                finish_reason=state.finish_reason,
                usage_chunk=None,
            )
            state.ended = True

        # Stream terminator. Cohere's v2 SSE protocol ends every stream
        # with ``data: [DONE]\n\n`` after ``message-end``; Fern-generated
        # clients (Go/Java) and cohere-python all key their read loop off
        # this sentinel.
        yield _DONE_FRAME

    # -- per-delta helpers --------------------------------------------

    def _handle_thinking_delta(self, state: _StreamState, delta_text: str) -> list[str]:
        # Non-reasoning Command models: emit ``tool-plan-delta`` events
        # directly instead of opening a thinking content block. The
        # ``tool-plan-delta`` event has no start/end pair around it.
        if not self._is_reasoning_model:
            events: list[str] = list(self._close_open_blocks(state))
            events.append(
                _emit(
                    ToolPlanDeltaEvent(
                        delta={"message": {"tool_plan": delta_text}},
                    )
                )
            )
            return events

        # Reasoning model (default): open / continue a thinking block.
        events = []
        if state.active_block != ContentBlockType.THINKING:
            events.extend(self._close_open_blocks(state))
            idx = state.next_content_index()
            state.active_block = ContentBlockType.THINKING
            state.active_block_index = idx
            events.append(
                _emit(
                    ContentStartEvent(
                        index=idx,
                        delta={
                            "message": {
                                "content": {
                                    "type": ContentBlockType.THINKING,
                                    "thinking": "",
                                }
                            }
                        },
                    )
                )
            )
        events.append(
            _emit(
                ContentDeltaEvent(
                    index=state.active_block_index,
                    delta={"message": {"content": {"thinking": delta_text}}},
                )
            )
        )
        return events

    def _handle_text_delta(self, state: _StreamState, delta_text: str) -> list[str]:
        events: list[str] = []
        if state.active_block != ContentBlockType.TEXT:
            events.extend(self._close_open_blocks(state))
            idx = state.next_content_index()
            state.active_block = ContentBlockType.TEXT
            state.active_block_index = idx
            events.append(
                _emit(
                    ContentStartEvent(
                        index=idx,
                        delta={
                            "message": {
                                "content": {
                                    "type": ContentBlockType.TEXT,
                                    "text": "",
                                }
                            }
                        },
                    )
                )
            )
        events.append(
            _emit(
                ContentDeltaEvent(
                    index=state.active_block_index,
                    delta={"message": {"content": {"text": delta_text}}},
                )
            )
        )
        return events

    def _handle_tool_call_deltas(self, state: _StreamState, deltas: list) -> list[str]:
        events: list[str] = []
        for tc in deltas:
            tc_index = tc.index
            fn = tc.function

            if tc_index not in state.tool_calls_seen:
                # New tool call. Close any open content/tool block first.
                events.extend(self._close_open_blocks(state))
                state.tool_calls_seen.add(tc_index)
                state.active_tool_index = tc_index
                state.active_block = ContentBlockType.TOOL_CALL
                events.append(
                    _emit(
                        ToolCallStartEvent(
                            index=tc_index,
                            delta={
                                "message": {
                                    "tool_calls": {
                                        "id": tc.id or "",
                                        "type": "function",
                                        "function": {
                                            "name": (fn.name if fn else "") or "",
                                            "arguments": (fn.arguments if fn else None)
                                            or "",
                                        },
                                    }
                                }
                            },
                        )
                    )
                )
                continue

            if fn and fn.arguments:
                events.append(
                    _emit(
                        ToolCallDeltaEvent(
                            index=tc_index,
                            delta={
                                "message": {
                                    "tool_calls": {
                                        "function": {
                                            "arguments": fn.arguments,
                                        }
                                    }
                                }
                            },
                        )
                    )
                )
        return events

    def _handle_citation_deltas(
        self,
        state: _StreamState,
        citations: list,
    ) -> list[str]:
        """Emit ``citation-start`` / ``citation-end`` events for a delta.

        The reasoning parser populates
        :class:`CohereDeltaMessage.citations` (see
        :mod:`vllm.entrypoints.cohere.cohere_chat_message`) with
        sources whose ``id`` / ``type`` / ``document`` / ``tool_output``
        are already resolved against the request's document context
        (see :func:`_melody_sources_to_vllm` in the parser). This
        method only coerces to the SDK wire shape, applies the
        ``THINKING_CONTENT`` -> ``PLAN`` rewrite for non-reasoning
        models, and drops citations whose sources didn't resolve --
        the shared :meth:`_to_wire_citation` handles all three.
        """
        events: list[str] = []
        for c in citations:
            citation = self._to_wire_citation(c)
            if citation is None:
                continue
            idx = state.next_citation_index()
            events.append(
                _emit(
                    CitationStartEvent(
                        index=idx,
                        delta={
                            "message": {
                                "citations": citation.model_dump(exclude_none=True)
                            }
                        },
                    )
                )
            )
            events.append(_emit(CitationEndEvent(index=idx)))
        return events

    # -- block lifecycle helpers --------------------------------------

    def _close_open_blocks(self, state: _StreamState) -> list[str]:
        """Emit ``content-end`` / ``tool-call-end`` for the currently open
        block (if any) and reset the corresponding ``_StreamState`` slots.
        """
        events: list[str] = []
        if state.active_block in (ContentBlockType.TEXT, ContentBlockType.THINKING):
            events.append(_emit(ContentEndEvent(index=state.active_block_index)))
        elif state.active_block == ContentBlockType.TOOL_CALL:
            events.append(_emit(ToolCallEndEvent(index=state.active_tool_index)))
        state.active_block = None
        state.active_block_index = None
        state.active_tool_index = None
        return events

    def _build_message_end_event(
        self,
        chunk_id: str,
        finish_reason: str | None,
        usage_chunk: ChatCompletionStreamResponse | None = None,
    ) -> str:
        delta: dict[str, Any] = {
            "finish_reason": _map_finish_reason(finish_reason),
        }
        if usage_chunk is not None and usage_chunk.usage is not None:
            prompt = usage_chunk.usage.prompt_tokens
            completion = usage_chunk.usage.completion_tokens or 0
            usage_block: dict[str, Any] = {
                "billed_units": {
                    "input_tokens": prompt,
                    "output_tokens": completion,
                },
                "tokens": {
                    "input_tokens": prompt,
                    "output_tokens": completion,
                },
            }
            if usage_chunk.usage.prompt_tokens_details is not None:
                cached = usage_chunk.usage.prompt_tokens_details.cached_tokens
                if cached is not None:
                    usage_block["cached_tokens"] = cached
            delta["usage"] = usage_block
        return _emit(MessageEndEvent(id=chunk_id, delta=delta))

    # ==================================================================
    # Helpers for the router
    # ==================================================================

    @staticmethod
    def create_error_response(
        message: str | Exception,
        err_type: str = "bad_request",
        status_code: HTTPStatus = HTTPStatus.BAD_REQUEST,
        param: str | None = None,
    ) -> ErrorResponse:
        # Override of :meth:`BaseServing.create_error_response` that uses
        # Cohere-flavored defaults (``type="bad_request"``, ``code=400``)
        # so the router can translate the envelope uniformly. ``param``
        # is accepted for signature parity with the base class but is
        # not surfaced in the Cohere wire format.
        from vllm.entrypoints.openai.engine.protocol import ErrorInfo

        del param  # unused; kept for signature compatibility
        return ErrorResponse(
            error=ErrorInfo(
                message=str(message),
                type=err_type,
                code=int(status_code),
            )
        )