# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # Adapted from # https://github.com/vllm-project/vllm/blob/main/vllm/entrypoints/anthropic/serving.py """Cohere Chat v2 API serving handler. Implements ``POST /cohere/v2/chat`` by translating the incoming Cohere v2 request into a standard :class:`ChatCompletionRequest` and delegating to :class:`OpenAIServingChat`. The actual prompt rendering is handled by vLLM's renderer pipeline (``vllm.renderers``): - For Cohere Command-family models, set ``--tokenizer-mode cohere`` and the :class:`vllm.renderers.cohere.CohereRenderer` will template the request via the ``cohere_melody`` library (``render_cmd3`` / ``render_cmd4``) and surface citations through the Cohere-scoped :class:`CohereChatMessage` / :class:`CohereDeltaMessage` subclasses (see :mod:`vllm.entrypoints.cohere.cohere_chat_message`). - For any other model the default Jinja-based :class:`HfRenderer` is used, and the endpoint behaves as a plain v2-shaped wrapper around chat completions. """ from __future__ import annotations import json import logging import time from collections.abc import AsyncGenerator from enum import Enum from http import HTTPStatus from typing import TYPE_CHECKING, Any, NamedTuple, cast from cohere.types import ( AssistantChatMessageV2, AssistantMessageResponse, Citation, Document, SystemChatMessageV2, ToolCallV2, ToolCallV2Function, ToolChatMessageV2, UserChatMessageV2, ) from fastapi import Request from pydantic import BaseModel from vllm.engine.protocol import EngineClient from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption from vllm.entrypoints.cohere.cohere_chat_message import ( CitationSource as InternalCitationSource, ) from vllm.entrypoints.cohere.cohere_chat_message import ( CohereChatMessage, ) from vllm.entrypoints.cohere.protocol import ( CitationEndEvent, CitationStartEvent, CohereChatV2Request, CohereChatV2Response, CohereFinishReason, CohereUsage, CohereUsageBilledUnits, CohereUsageTokens, ContentDeltaEvent, ContentEndEvent, ContentStartEvent, MessageEndEvent, MessageStartEvent, ToolCallDeltaEvent, ToolCallEndEvent, ToolCallStartEvent, ToolPlanDeltaEvent, ) from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionRequest, ChatCompletionResponse, ChatCompletionStreamResponse, ChatCompletionToolsParam, ChatMessage, ) from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat from vllm.entrypoints.openai.engine.protocol import ( ErrorResponse, JsonSchemaResponseFormat, ResponseFormat, StreamOptions, ) from vllm.entrypoints.openai.models.serving import OpenAIServingModels from vllm.entrypoints.serve.exception_handling.utils import sanitize_message from vllm.entrypoints.serve.utils.request_logger import RequestLogger from vllm.parser.abstract_parser import Parser from vllm.renderers.cohere import MESSAGES_CITATIONS_KEY, POSITION_TO_SOURCE_KEY if TYPE_CHECKING: from vllm.renderers.online_renderer import OnlineRenderer logger = logging.getLogger(__name__) # --------------------------------------------------------------------------- # SSE helpers # --------------------------------------------------------------------------- def _sse(data: str) -> str: """Wrap a JSON payload in a Server-Sent Event frame. Cohere's stream uses bare ``data:`` lines (no ``event:`` prefix); each JSON object's ``type`` field carries the event discriminator. """ return f"data: {data}\n\n" # Cohere v2 SSE stream terminator. Declared by the upstream OpenAPI spec # (``x-fern-streaming.terminator: "[DONE]"`` on the ``/v2/chat`` stream # operation) and observed by the cohere-python / Fern-generated clients # (the Python SDK breaks its read loop on ``_sse.data == "[DONE]"``). # Must be emitted after the closing ``message-end`` event. _DONE_FRAME = _sse("[DONE]") def _emit(event: BaseModel) -> str: """Serialize a typed stream event into an SSE frame. The typed event classes in ``vllm.entrypoints.cohere.protocol`` (``MessageStartEvent``, ``ContentStartEvent``, ``CitationStartEvent``, ...) bake the wire-format ``type`` field into the model definition, so a plain ``model_dump_json()`` already carries the discriminator. """ return _sse(event.model_dump_json(exclude_none=True)) class ContentBlockType(str, Enum): """Wire-format / internal discriminator for chat content blocks. ``THINKING`` and ``TEXT`` are the documented Cohere v2 content-block type discriminators on the wire. ``TOOL_CALL`` is reserved for the internal stream state machine (see :class:`_StreamState`) when a tool call is the currently open block; it is never serialized. """ THINKING = "thinking" TEXT = "text" TOOL_CALL = "tool_call" # Mapping of vLLM/OpenAI finish reasons to Cohere's enum. _FINISH_REASON_MAP: dict[str | None, CohereFinishReason] = { "stop": "COMPLETE", "length": "MAX_TOKENS", "tool_calls": "TOOL_CALL", "stop_sequence": "STOP_SEQUENCE", "error": "ERROR", None: "COMPLETE", } def _map_finish_reason( reason: str | None, stop_reason: int | str | None = None ) -> CohereFinishReason: # vLLM reports both EOS and matched stop strings as finish_reason="stop". # The latter is distinguished by the matched string in stop_reason. if reason == "stop" and isinstance(stop_reason, str): return "STOP_SEQUENCE" return _FINISH_REASON_MAP.get(reason, "COMPLETE") class _CitablePosition(NamedTuple): """One entry in melody's citation numbering scheme, resolved to wire shape. Populated by :meth:`CohereServingChatV2._walk_citable_positions` (the single source of truth for melody's numbering rule) and consumed by both the inbound (:meth:`CohereServingChatV2._build_doc_id_to_prompt_position`, which drives citation resolution on the request side) and outbound (:meth:`CohereServingChatV2._build_position_to_source`, whose result is forwarded to the reasoning parser via ``chat_template_kwargs[POSITION_TO_SOURCE_KEY]``) helpers -- so the numbering rule only lives in one place. ``ids`` lists every id under which this position is addressable: top-level documents contribute both their explicit id (if set) and the ``doc_{idx}`` synthetic fallback, in that order so ``setdefault``-style resolution prefers the explicit id; tool-message documents contribute their single ``document.id``. ``source`` is the fully-resolved wire-shape source (``type`` / ``id`` / ``document`` or ``tool_output`` payload) for the outbound direction. """ bucket: int result_idx: int ids: tuple[str, ...] source: InternalCitationSource # --------------------------------------------------------------------------- # Serving class # --------------------------------------------------------------------------- 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 ``text`` 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 ``text`` marker, and silently attributing a citation to the wrong bucket is worse than dropping the ```` 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 ``...`` 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, choice.stop_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": }`` 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, stop_reason=state.stop_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 state.stop_reason = choice.stop_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, stop_reason=state.stop_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, stop_reason: int | str | None = None, usage_chunk: ChatCompletionStreamResponse | None = None, ) -> str: delta: dict[str, Any] = { "finish_reason": _map_finish_reason(finish_reason, stop_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), ) ) # --------------------------------------------------------------------------- # Stream state # --------------------------------------------------------------------------- class _StreamState: """Tracks which Cohere v2 stream block (if any) is currently open.""" def __init__(self) -> None: self.started: bool = False self.ended: bool = False self.finish_reason: str | None = None self.stop_reason: int | str | None = None self.last_chunk_id: str = "" self.active_block: ContentBlockType | None = None self.active_block_index: int | None = None self.active_tool_index: int | None = None self._next_index: int = 0 self._next_citation_index: int = 0 self.tool_calls_seen: set[int] = set() def next_content_index(self) -> int: idx = self._next_index self._next_index += 1 return idx def next_citation_index(self) -> int: idx = self._next_citation_index self._next_citation_index += 1 return idx