diff --git a/tests/entrypoints/cohere/__init__.py b/tests/entrypoints/cohere/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/entrypoints/cohere/test_api_router.py b/tests/entrypoints/cohere/test_api_router.py new file mode 100644 index 000000000000..4c88d627c7aa --- /dev/null +++ b/tests/entrypoints/cohere/test_api_router.py @@ -0,0 +1,200 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for ``vllm/entrypoints/cohere/api_router.py``. + +Covers: + +* The optional-import guard: ``attach_router`` is a no-op when the + ``cohere`` SDK isn't installed. +* The router wiring: response shapes (JSON + SSE), error translation, + and the ``cohere_serving_chat_v2 is None`` fallback (501 Not + Implemented). +""" + +import json +import sys +from collections.abc import AsyncGenerator +from http import HTTPStatus + +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from vllm.entrypoints.cohere.api_router import attach_router +from vllm.entrypoints.cohere.protocol import ( + AssistantMessageResponse, + CohereChatV2Response, +) +from vllm.entrypoints.openai.engine.protocol import ErrorInfo, ErrorResponse + +# ---------------------------------------------------------------------- +# Fakes +# ---------------------------------------------------------------------- + + +class _Handler: + """Minimal stand-in for :class:`CohereServingChatV2` used by the + router. Each test sets ``self.result`` to either: + + * a :class:`CohereChatV2Response` (non-streaming JSON path); + * an async generator yielding SSE frames (streaming path); + * an :class:`ErrorResponse` (error envelope path); or + * an exception (router-level 500 path). + """ + + def __init__(self, result): + self.result = result + + async def create_chat_v2(self, request, raw_request): + if isinstance(self.result, Exception): + raise self.result + return self.result + + +def _build_app(handler: _Handler | None) -> FastAPI: + app = FastAPI() + attach_router(app) + app.state.cohere_serving_chat_v2 = handler + return app + + +def _minimal_request_body() -> dict: + return { + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + } + + +# ---------------------------------------------------------------------- +# Optional-import guard +# ---------------------------------------------------------------------- + + +class TestOptionalCohereImport: + def test_attach_router_noop_when_cohere_missing(self, monkeypatch, caplog): + # Simulate ``import cohere`` failing by injecting a sentinel that + # raises ImportError on first attribute access. The simplest + # approach: stash the real module, delete it from sys.modules, + # and shadow it with an entry that raises on next import. + real = sys.modules.pop("cohere", None) + + # Block re-import. + class _RaisingFinder: + def find_spec(self, name, path=None, target=None): + if name == "cohere": + raise ImportError("simulated missing cohere SDK") + return None + + finder = _RaisingFinder() + sys.meta_path.insert(0, finder) + try: + with caplog.at_level("INFO", logger="vllm.entrypoints.cohere.api_router"): + app = FastAPI() + attach_router(app) + # No routes should have been registered. + paths = [r.path for r in app.routes] + assert "/cohere/v2/chat" not in paths + assert any( + "cohere SDK not installed" in rec.message for rec in caplog.records + ) + finally: + sys.meta_path.remove(finder) + if real is not None: + sys.modules["cohere"] = real + + def test_attach_router_registers_route_when_cohere_present(self): + app = _build_app(handler=None) + paths = [getattr(r, "path", None) for r in app.routes] + assert "/cohere/v2/chat" in paths + + +# ---------------------------------------------------------------------- +# Endpoint behavior +# ---------------------------------------------------------------------- + + +class TestEndpoint: + def test_501_when_handler_missing(self): + app = _build_app(handler=None) + with TestClient(app) as client: + r = client.post("/cohere/v2/chat", json=_minimal_request_body()) + assert r.status_code == HTTPStatus.NOT_IMPLEMENTED + body = r.json() + assert "does not support" in body["message"] + assert "id" not in body # excluded by ``exclude_none=True`` + + def test_non_streaming_response_is_json(self): + msg = AssistantMessageResponse(content=[{"type": "text", "text": "hello"}]) + result = CohereChatV2Response(id="r1", finish_reason="COMPLETE", message=msg) + app = _build_app(handler=_Handler(result)) + with TestClient(app) as client: + r = client.post("/cohere/v2/chat", json=_minimal_request_body()) + assert r.status_code == HTTPStatus.OK + assert r.headers["content-type"].startswith("application/json") + body = r.json() + assert body["id"] == "r1" + assert body["finish_reason"] == "COMPLETE" + assert body["message"]["content"][0]["text"] == "hello" + + def test_streaming_response_is_sse(self): + async def _gen() -> AsyncGenerator[str, None]: + yield 'data: {"type":"message-start"}\n\n' + yield "data: [DONE]\n\n" + + app = _build_app(handler=_Handler(_gen())) + with TestClient(app) as client: + r = client.post( + "/cohere/v2/chat", + json={**_minimal_request_body(), "stream": True}, + ) + assert r.status_code == HTTPStatus.OK + assert r.headers["content-type"].startswith("text/event-stream") + body = r.text + assert "message-start" in body + assert body.rstrip().endswith("[DONE]") + + def test_error_response_translated_to_cohere_envelope(self): + err = ErrorResponse( + error=ErrorInfo( + message="bad request", + type="bad_request", + code=400, + ) + ) + app = _build_app(handler=_Handler(err)) + with TestClient(app) as client: + r = client.post("/cohere/v2/chat", json=_minimal_request_body()) + assert r.status_code == HTTPStatus.BAD_REQUEST + body = r.json() + assert body == {"message": "bad request"} + + def test_handler_exception_returns_500_envelope(self): + app = _build_app(handler=_Handler(RuntimeError("kaboom"))) + with TestClient(app) as client: + r = client.post("/cohere/v2/chat", json=_minimal_request_body()) + assert r.status_code == HTTPStatus.INTERNAL_SERVER_ERROR + body = r.json() + assert body == {"message": "kaboom"} + + def test_non_json_content_type_rejected(self): + """The ``validate_json_request`` dependency raises + ``RequestValidationError`` (HTTP 422) for non-JSON content + types, matching the behavior of the other vLLM API routers. + """ + app = _build_app(handler=None) + with TestClient(app) as client: + r = client.post( + "/cohere/v2/chat", + content=json.dumps(_minimal_request_body()), + headers={"content-type": "text/plain"}, + ) + assert r.status_code == HTTPStatus.UNPROCESSABLE_ENTITY + + def test_invalid_body_returns_422(self): + # ``model`` is required; omit it to trip Pydantic validation. + app = _build_app(handler=None) + with TestClient(app) as client: + r = client.post( + "/cohere/v2/chat", + json={"messages": [{"role": "user", "content": "hi"}]}, + ) + assert r.status_code == HTTPStatus.UNPROCESSABLE_ENTITY diff --git a/tests/entrypoints/cohere/test_chat_v2.py b/tests/entrypoints/cohere/test_chat_v2.py new file mode 100644 index 000000000000..505a2f93d836 --- /dev/null +++ b/tests/entrypoints/cohere/test_chat_v2.py @@ -0,0 +1,238 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""End-to-end integration tests for the ``POST /cohere/v2/chat`` endpoint. + +These tests spin up a real ``vllm serve`` process via +:class:`tests.utils.RemoteOpenAIServer` and exercise the Cohere Chat v2 +contract over HTTP. They mirror the pattern used by +:mod:`tests.entrypoints.anthropic.test_messages`. + +Two layers of integration are covered: + +1. Raw HTTP via :mod:`httpx` — always runs, verifies the wire contract. +2. Cohere SDK (``pip install cohere``) — auto-skipped when the optional + dependency isn't installed, verifies SDK-level interop. + +Like ``test_messages.py`` we use a small generic chat model +(``Qwen/Qwen3-0.6B``); the Cohere v2 endpoint is model-agnostic and +just performs the v2 ↔ OpenAI chat-completion translation, so the +choice of model only matters for response shape (not for the +translation logic itself, which is unit-tested elsewhere). +""" + +import json + +import httpx +import pytest +import pytest_asyncio + +from tests.utils import RemoteOpenAIServer + +MODEL_NAME = "Qwen/Qwen3-0.6B" +SERVED_MODEL_NAME = "command-r-plus-08-2024" + + +@pytest.fixture(scope="module") +def server(): + args = [ + "--max-model-len", + "2048", + "--enforce-eager", + "--enable-auto-tool-choice", + "--tool-call-parser", + "hermes", + # Advertise a Cohere model name so Cohere SDK ``model=`` calls + # round-trip without the ``model not found`` check. + "--served-model-name", + SERVED_MODEL_NAME, + # Disable the reasoning-model path so the test doesn't require + # the conversation to surface a thinking block (Qwen3-0.6B + # doesn't emit Cohere-style reasoning tokens out of the box). + "--no-cohere-is-reasoning-model", + ] + + with RemoteOpenAIServer(MODEL_NAME, args) as remote_server: + yield remote_server + + +# ---------------------------------------------------------------------- +# Layer 1: raw HTTP contract (no optional deps) +# ---------------------------------------------------------------------- + + +@pytest_asyncio.fixture +async def httpx_client(server): + async with httpx.AsyncClient( + base_url=server.url_root, timeout=httpx.Timeout(120.0) + ) as client: + yield client + + +@pytest.mark.asyncio +async def test_cohere_v2_chat_non_streaming(httpx_client: httpx.AsyncClient): + resp = await httpx_client.post( + "/cohere/v2/chat", + json={ + "model": SERVED_MODEL_NAME, + "messages": [{"role": "user", "content": "Say hi."}], + "max_tokens": 16, + "stream": False, + }, + ) + assert resp.status_code == 200, resp.text + payload = resp.json() + # The response envelope follows Cohere v2's schema. + assert "message" in payload + assert payload["message"]["role"] == "assistant" + # ``content`` is a list of content blocks; the synthesized + # ``CohereServingChatV2`` should always emit at least one ``text`` + # block (it falls back to an empty block when the model returned + # nothing). + content = payload["message"]["content"] + assert isinstance(content, list) and len(content) >= 1 + assert content[0]["type"] == "text" + # ``usage`` is always populated by the translator. + assert "usage" in payload + assert "billed_units" in payload["usage"] + assert "tokens" in payload["usage"] + # ``finish_reason`` is one of Cohere's enum values. + assert payload["finish_reason"] in { + "COMPLETE", + "MAX_TOKENS", + "STOP_SEQUENCE", + "TOOL_CALL", + "ERROR", + } + + +@pytest.mark.asyncio +async def test_cohere_v2_chat_streaming(httpx_client: httpx.AsyncClient): + """Streaming returns SSE frames in the v2 message-lifecycle shape.""" + events: list[dict] = [] + async with httpx_client.stream( + "POST", + "/cohere/v2/chat", + json={ + "model": SERVED_MODEL_NAME, + "messages": [{"role": "user", "content": "Hi"}], + "max_tokens": 8, + "stream": True, + }, + ) as resp: + assert resp.status_code == 200, await resp.aread() + assert resp.headers["content-type"].startswith("text/event-stream") + async for line in resp.aiter_lines(): + if not line.startswith("data: "): + continue + data = line[len("data: ") :] + if data == "[DONE]": + events.append({"type": "_DONE_"}) + continue + events.append(json.loads(data)) + + types = [ev["type"] for ev in events] + # The lifecycle always starts with message-start and ends with [DONE] + # preceded by message-end. + assert types[0] == "message-start" + assert types[-1] == "_DONE_" + assert types[-2] == "message-end" + # message-start must carry the chunk/message id. + assert events[0].get("id") + + +@pytest.mark.asyncio +async def test_cohere_v2_chat_validation_error_returns_422( + httpx_client: httpx.AsyncClient, +): + # Missing required ``model`` field → FastAPI/Pydantic returns 422. + resp = await httpx_client.post( + "/cohere/v2/chat", + json={"messages": []}, + ) + assert resp.status_code == 422 + + +@pytest.mark.asyncio +async def test_cohere_v2_chat_documents_field_accepted( + httpx_client: httpx.AsyncClient, +): + """The v2 endpoint forwards ``documents`` into chat_template_kwargs. + + We only assert the request is accepted and produces a 200 response — + the renderer-level effect is covered by ``tests/renderers/test_cohere.py``. + """ + resp = await httpx_client.post( + "/cohere/v2/chat", + json={ + "model": SERVED_MODEL_NAME, + "messages": [{"role": "user", "content": "Summarize."}], + "documents": [{"id": "d1", "data": {"title": "T", "snippet": "S"}}], + "max_tokens": 16, + "stream": False, + }, + ) + assert resp.status_code == 200, resp.text + + +# ---------------------------------------------------------------------- +# Layer 2: Cohere SDK round-trip (auto-skipped if SDK absent) +# ---------------------------------------------------------------------- + + +@pytest_asyncio.fixture +async def cohere_async_client(server): + cohere = pytest.importorskip("cohere") + # The vLLM endpoint is mounted at ``/cohere/v2/chat`` while the + # cohere SDK targets ``${base_url}/v2/chat``; point base_url at the + # ``/cohere`` prefix so paths line up. + client = cohere.AsyncClientV2( + api_key="dummy", + base_url=server.url_for("cohere"), + ) + try: + yield client + finally: + # ``AsyncClientV2`` exposes a sync close; if a future version + # adds aclose we still close cleanly. + close = getattr(client, "aclose", None) or getattr(client, "close", None) + if close is not None: + result = close() + if hasattr(result, "__await__"): + await result + + +@pytest.mark.asyncio +async def test_cohere_sdk_non_streaming(cohere_async_client): + resp = await cohere_async_client.chat( + model=SERVED_MODEL_NAME, + messages=[{"role": "user", "content": "Say hi."}], + max_tokens=16, + ) + # SDK parses our JSON into typed objects. + assert resp.message.role == "assistant" + assert resp.message.content is not None + assert len(resp.message.content) >= 1 + assert resp.message.content[0].type == "text" + assert resp.finish_reason in { + "COMPLETE", + "MAX_TOKENS", + "STOP_SEQUENCE", + "TOOL_CALL", + "ERROR", + } + + +@pytest.mark.asyncio +async def test_cohere_sdk_streaming(cohere_async_client): + events: list[str] = [] + stream = cohere_async_client.chat_stream( + model=SERVED_MODEL_NAME, + messages=[{"role": "user", "content": "Hi"}], + max_tokens=8, + ) + async for ev in stream: + events.append(ev.type) + + assert events, "SDK stream yielded no events" + assert events[0] == "message-start" + assert events[-1] == "message-end" diff --git a/tests/entrypoints/cohere/test_protocol.py b/tests/entrypoints/cohere/test_protocol.py new file mode 100644 index 000000000000..c2e56297cd0d --- /dev/null +++ b/tests/entrypoints/cohere/test_protocol.py @@ -0,0 +1,350 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for ``vllm/entrypoints/cohere/protocol.py``. + +The module is mostly a thin wrapper around the official ``cohere`` SDK +types plus a few local additions: + +* :class:`CohereError` envelope. +* :class:`CohereChatV2Request` (model required, ``max_tokens`` non-negative). +* :class:`CohereChatV2Response` plus the usage / logprob helpers. +* The streaming event subclasses that bake a wire-format ``type`` + discriminator into ``model_dump()`` so SSE consumers can demux on it. +""" + +import pytest +from pydantic import ValidationError + +from vllm.entrypoints.cohere.protocol import ( + AssistantChatMessageV2, + AssistantMessageResponse, + CitationEndEvent, + CitationStartEvent, + CohereChatV2Request, + CohereChatV2Response, + CohereError, + CohereLogprobItem, + CohereUsage, + CohereUsageBilledUnits, + CohereUsageTokens, + ContentDeltaEvent, + ContentEndEvent, + ContentStartEvent, + MessageEndEvent, + MessageStartEvent, + SystemChatMessageV2, + ToolCallDeltaEvent, + ToolCallEndEvent, + ToolCallStartEvent, + ToolPlanDeltaEvent, + UserChatMessageV2, +) + +# ====================================================================== +# CohereError +# ====================================================================== + + +class TestCohereError: + def test_message_only(self): + err = CohereError(message="boom") + assert err.message == "boom" + assert err.id is None + + def test_with_id(self): + err = CohereError(message="boom", id="req_123") + assert err.id == "req_123" + + def test_model_dump_excludes_none(self): + err = CohereError(message="boom") + assert err.model_dump(exclude_none=True) == {"message": "boom"} + + +# ====================================================================== +# CohereChatV2Request +# ====================================================================== + + +class TestCohereChatV2Request: + def test_minimal_required_fields(self): + req = CohereChatV2Request( + model="m", messages=[{"role": "user", "content": "hi"}] + ) + assert req.model == "m" + assert req.stream is False + assert req.max_tokens is None + assert req.tools is None + assert req.documents is None + assert req.kv_transfer_params is None + assert req.chat_template_kwargs is None + + def test_empty_model_rejected(self): + with pytest.raises(ValidationError, match="model is required"): + CohereChatV2Request(model="", messages=[{"role": "user", "content": "hi"}]) + + def test_negative_max_tokens_rejected(self): + with pytest.raises(ValidationError, match="non-negative"): + CohereChatV2Request( + model="m", + messages=[{"role": "user", "content": "hi"}], + max_tokens=-1, + ) + + def test_zero_max_tokens_allowed(self): + # Zero is allowed (the docs allow 0 -> return prompt only). + req = CohereChatV2Request( + model="m", + messages=[{"role": "user", "content": "hi"}], + max_tokens=0, + ) + assert req.max_tokens == 0 + + def test_full_schema(self): + """Sanity-check that every field listed in the model accepts a value.""" + req = CohereChatV2Request( + model="m", + stream=True, + messages=[ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": "hello", + "tool_plan": "p", + "tool_calls": [ + { + "id": "c1", + "type": "function", + "function": {"name": "f", "arguments": "{}"}, + } + ], + }, + ], + tools=[ + { + "type": "function", + "function": { + "name": "f", + "description": "d", + "parameters": {}, + }, + } + ], + strict_tools=True, + tool_choice="REQUIRED", + documents=["plain", {"id": "d1", "data": {"text": "t"}}], + citation_options={"mode": "accurate"}, + response_format={"type": "json_object"}, + safety_mode="CONTEXTUAL", + max_tokens=128, + stop_sequences=[""], + temperature=0.5, + seed=42, + frequency_penalty=0.1, + presence_penalty=0.2, + k=50, + p=0.95, + logprobs=True, + thinking={"type": "enabled", "token_budget": 1024}, + priority=1, + kv_transfer_params={"x": 1}, + chat_template_kwargs={"y": 2}, + ) + # Discriminated-union messages resolve to the SDK variants. + assert isinstance(req.messages[0], SystemChatMessageV2) + assert isinstance(req.messages[1], UserChatMessageV2) + assert isinstance(req.messages[2], AssistantChatMessageV2) + assert req.thinking.type == "enabled" + assert req.thinking.token_budget == 1024 + assert req.citation_options.mode == "accurate" + assert req.response_format.type == "json_object" + assert req.safety_mode == "CONTEXTUAL" + assert req.tool_choice == "REQUIRED" + assert req.tools[0].function.name == "f" + # Documents accept str OR Document. + assert isinstance(req.documents[0], str) + assert req.documents[1].id == "d1" + + def test_invalid_tool_choice_rejected(self): + with pytest.raises(ValidationError): + CohereChatV2Request( + model="m", + messages=[{"role": "user", "content": "hi"}], + tool_choice="ANY", # not REQUIRED/NONE + ) + + def test_unknown_field_ignored(self): + # Pydantic by default ignores extra fields. Make sure that contract + # holds so clients sending forward-compatible kwargs don't 422. + req = CohereChatV2Request( + model="m", + messages=[{"role": "user", "content": "hi"}], + some_future_field="x", # type: ignore[arg-type] + ) + assert req.model == "m" + + +# ====================================================================== +# Usage / Logprobs +# ====================================================================== + + +class TestUsage: + def test_billed_units_optional(self): + u = CohereUsageBilledUnits() + assert u.input_tokens is None + assert u.output_tokens is None + + def test_tokens_serialization(self): + u = CohereUsageTokens(input_tokens=10, output_tokens=5) + assert u.model_dump() == {"input_tokens": 10.0, "output_tokens": 5.0} + + def test_usage_envelope(self): + u = CohereUsage( + billed_units=CohereUsageBilledUnits(input_tokens=10, output_tokens=5), + tokens=CohereUsageTokens(input_tokens=10, output_tokens=5), + cached_tokens=3, + ) + assert u.cached_tokens == 3 + assert u.billed_units.input_tokens == 10 + + def test_logprob_item(self): + lp = CohereLogprobItem(text="hi", token_ids=[1, 2], logprobs=[-0.1, -0.2]) + assert lp.token_ids == [1, 2] + assert lp.logprobs == [-0.1, -0.2] + + +# ====================================================================== +# CohereChatV2Response +# ====================================================================== + + +class TestCohereChatV2Response: + def test_minimal(self): + msg = AssistantMessageResponse(content=[{"type": "text", "text": "hi"}]) + resp = CohereChatV2Response( + id="r1", + finish_reason="COMPLETE", + message=msg, + ) + assert resp.id == "r1" + assert resp.usage is None + assert resp.kv_transfer_params is None + + def test_invalid_finish_reason_rejected(self): + msg = AssistantMessageResponse(content=[{"type": "text", "text": "hi"}]) + with pytest.raises(ValidationError): + CohereChatV2Response( + id="r1", + finish_reason="NOT_A_REASON", # type: ignore[arg-type] + message=msg, + ) + + +# ====================================================================== +# Streaming event ``type`` discriminator baked into model_dump() +# ====================================================================== + + +class TestStreamingEventTypeField: + """Each event subclass adds a ``type: Literal[...]`` field with a + default so ``model_dump()`` always emits the wire-format discriminator + (the parent SDK classes don't declare ``type`` as a Pydantic field). + """ + + @pytest.mark.parametrize( + "cls, expected_type, kwargs", + [ + ( + MessageStartEvent, + "message-start", + {"id": "a", "delta": {"message": {"role": "assistant"}}}, + ), + ( + ContentStartEvent, + "content-start", + { + "index": 0, + "delta": {"message": {"content": {"type": "text", "text": ""}}}, + }, + ), + ( + ContentDeltaEvent, + "content-delta", + { + "index": 0, + "delta": {"message": {"content": {"text": "hi"}}}, + }, + ), + (ContentEndEvent, "content-end", {"index": 0}), + ( + ToolPlanDeltaEvent, + "tool-plan-delta", + {"delta": {"message": {"tool_plan": "thinking"}}}, + ), + ( + ToolCallStartEvent, + "tool-call-start", + { + "index": 0, + "delta": { + "message": { + "tool_calls": { + "id": "c1", + "type": "function", + "function": {"name": "f", "arguments": ""}, + } + } + }, + }, + ), + ( + ToolCallDeltaEvent, + "tool-call-delta", + { + "index": 0, + "delta": { + "message": {"tool_calls": {"function": {"arguments": "{}"}}} + }, + }, + ), + (ToolCallEndEvent, "tool-call-end", {"index": 0}), + ( + CitationStartEvent, + "citation-start", + { + "index": 0, + "delta": { + "message": { + "citations": {"start": 0, "end": 5, "text": "hello"} + } + }, + }, + ), + (CitationEndEvent, "citation-end", {"index": 0}), + ( + MessageEndEvent, + "message-end", + {"id": "a", "delta": {"finish_reason": "COMPLETE"}}, + ), + ], + ) + def test_type_field_default(self, cls, expected_type, kwargs): + ev = cls(**kwargs) + # type field is auto-populated from the Literal default. + assert ev.type == expected_type + # The discriminator must be present in the serialized payload so + # clients reading the stream can demux on it. + dumped = ev.model_dump(exclude_none=True) + assert dumped["type"] == expected_type + # Same in JSON form (what ``_emit`` serializes). + assert f'"type":"{expected_type}"' in ev.model_dump_json(exclude_none=True) + + def test_type_field_cannot_be_overridden_to_wrong_value(self): + # Literal types reject any value other than the bake-in default. + with pytest.raises(ValidationError): + MessageStartEvent( + id="a", + delta={"message": {"role": "assistant"}}, + type="other", # type: ignore[arg-type] + ) diff --git a/tests/entrypoints/cohere/test_registry_and_args.py b/tests/entrypoints/cohere/test_registry_and_args.py new file mode 100644 index 000000000000..5e907b0d9340 --- /dev/null +++ b/tests/entrypoints/cohere/test_registry_and_args.py @@ -0,0 +1,91 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Sanity tests for the small registry / CLI-arg / config additions made +by the Cohere v2 chat API: + +* ``vllm/renderers/registry.py``: new ``"cohere"`` renderer entry. +* ``vllm/tokenizers/registry.py``: new ``"cohere"`` tokenizer entry + (aliased to the cached HF tokenizer). +* ``vllm/entrypoints/openai/cli_args.py``: new + ``cohere_is_reasoning_model`` field. +* ``vllm/config/model.py``: ``"cohere"`` added to ``TokenizerMode`` + Literal. +""" + +import dataclasses +import typing + +import pytest + +from vllm.config.model import TokenizerMode +from vllm.entrypoints.openai.cli_args import BaseFrontendArgs, make_arg_parser +from vllm.renderers.registry import RENDERER_REGISTRY +from vllm.tokenizers.registry import TokenizerRegistry +from vllm.utils.argparse_utils import FlexibleArgumentParser + + +class TestRendererRegistry: + def test_cohere_renderer_registered(self): + # The registry resolves to the importable ``CohereRenderer`` class. + cls = RENDERER_REGISTRY.load_renderer_cls("cohere") + assert cls.__name__ == "CohereRenderer" + # Sanity: class lives in the cohere renderer module. + assert cls.__module__ == "vllm.renderers.cohere" + + +class TestTokenizerRegistry: + def test_cohere_aliased_to_cached_hf_tokenizer(self): + # ``cohere`` mode uses the standard HF tokenizer; only the + # renderer stage is replaced. This test guards against accidental + # divergence. + cls = TokenizerRegistry.load_tokenizer_cls("cohere") + assert cls.__name__ == "CachedHfTokenizer" + + +class TestTokenizerModeLiteral: + def test_cohere_is_a_valid_tokenizer_mode(self): + # The Literal must enumerate ``"cohere"`` so engine arg parsing + # accepts ``--tokenizer-mode cohere``. + modes = typing.get_args(TokenizerMode) + assert "cohere" in modes + + +# ---------------------------------------------------------------------- +# ``--cohere-is-reasoning-model`` CLI flag +# ---------------------------------------------------------------------- + + +class TestCohereCliArg: + """Verifies the new ``--cohere-is-reasoning-model`` flag end-to-end + through ``make_arg_parser`` — mirrors the pattern used in + :mod:`tests.entrypoints.openai.test_cli_args`. + """ + + def test_default_value_on_dataclass_is_true(self): + fields = {f.name: f for f in dataclasses.fields(BaseFrontendArgs)} + assert "cohere_is_reasoning_model" in fields + field = fields["cohere_is_reasoning_model"] + assert field.default is True + assert field.type is bool + + @pytest.fixture + def serve_parser(self) -> FlexibleArgumentParser: + parser = FlexibleArgumentParser() + return make_arg_parser(parser) + + def test_default_via_argparse_is_true(self, serve_parser: FlexibleArgumentParser): + # No flag supplied → dataclass default (True) wins. + args = serve_parser.parse_args(["--model", "m"]) + assert args.cohere_is_reasoning_model is True + + def test_explicit_false_via_argparse(self, serve_parser: FlexibleArgumentParser): + # Boolean dataclass fields are wired up as ``--flag value`` / + # ``--no-flag`` pairs by FlexibleArgumentParser. + args = serve_parser.parse_args( + ["--model", "m", "--no-cohere-is-reasoning-model"] + ) + assert args.cohere_is_reasoning_model is False + + def test_explicit_true_via_argparse(self, serve_parser: FlexibleArgumentParser): + args = serve_parser.parse_args(["--model", "m", "--cohere-is-reasoning-model"]) + assert args.cohere_is_reasoning_model is True diff --git a/tests/entrypoints/cohere/test_serving_conversion.py b/tests/entrypoints/cohere/test_serving_conversion.py new file mode 100644 index 000000000000..680ae775dd46 --- /dev/null +++ b/tests/entrypoints/cohere/test_serving_conversion.py @@ -0,0 +1,1036 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for the Cohere v2 -> OpenAI request / response conversion +implemented in ``vllm/entrypoints/cohere/serving.py``. + +These cover the pure-Python classmethods so we don't need an engine. +For the instance methods that read ``self._is_reasoning_model`` we +build a lightweight :class:`_FakeServing` subclass that skips the +heavy ``OpenAIServingChat.__init__`` chain (which would otherwise need +a real engine client, model registry, etc.) — the same pattern used in +``test_serving_streaming.py``. +""" + +from typing import Any + +import pytest + +from vllm.entrypoints.cohere.protocol import ( + CohereChatV2Request, + CohereChatV2Response, +) +from vllm.entrypoints.cohere.serving import ( + _FINISH_REASON_MAP, + CohereServingChatV2, + ContentBlockType, + _map_finish_reason, +) +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionRequest, + ChatCompletionResponse, +) +from vllm.entrypoints.openai.engine.protocol import ( + Citation as VLLMCitation, +) +from vllm.entrypoints.openai.engine.protocol import ( + CitationSource, +) + +# ---------------------------------------------------------------------- +# Helpers +# ---------------------------------------------------------------------- + + +def _make_request(**kwargs) -> CohereChatV2Request: + kwargs.setdefault("model", "m") + kwargs.setdefault("messages", [{"role": "user", "content": "hi"}]) + return CohereChatV2Request(**kwargs) + + +def _convert(request: CohereChatV2Request) -> ChatCompletionRequest: + return CohereServingChatV2._convert_v2_to_chat_completion(request) + + +class _FakeServing(CohereServingChatV2): + """Lightweight stand-in for :class:`CohereServingChatV2` that skips + the heavy ``OpenAIServingChat.__init__`` chain. + + Only ``_is_reasoning_model`` is read by the methods under test + (``_chat_completion_to_v2`` and friends); the rest of the parent + state is dead weight for unit testing. + """ + + def __init__(self, is_reasoning_model: bool = True) -> None: + # Intentionally skipping super().__init__ — see class docstring. + self._is_reasoning_model = is_reasoning_model + + +def _serving(is_reasoning_model: bool = True) -> CohereServingChatV2: + return _FakeServing(is_reasoning_model=is_reasoning_model) + + +def _build_chat_completion_response( + *, + response_id: str = "resp_1", + content: str | None = "hello", + reasoning: str | None = None, + tool_calls: list[dict[str, Any]] | None = None, + finish_reason: str | None = "stop", + citations: list[Any] | None = None, + usage: dict[str, Any] | None = None, + kv_transfer_params: dict[str, Any] | None = None, +) -> ChatCompletionResponse: + message: dict[str, Any] = {"role": "assistant"} + if content is not None: + message["content"] = content + if reasoning is not None: + message["reasoning"] = reasoning + if tool_calls is not None: + message["tool_calls"] = tool_calls + if citations is not None: + message["citations"] = citations + kwargs: dict[str, Any] = dict( + id=response_id, + object="chat.completion", + created=0, + model="m", + choices=[{"index": 0, "message": message, "finish_reason": finish_reason}], + # ``usage`` is a required field on ChatCompletionResponse, but the + # production code defensively handles ``None`` -> no usage block; + # we round-trip that behavior by post-setting the attribute below. + usage={"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}, + ) + if kv_transfer_params is not None: + kwargs["kv_transfer_params"] = kv_transfer_params + resp = ChatCompletionResponse(**kwargs) + if usage is None: + resp.usage = None + else: + # Replace the placeholder with the caller-provided usage. + resp = ChatCompletionResponse.model_validate( + {**resp.model_dump(), "usage": usage} + ) + if kv_transfer_params is not None: + resp.kv_transfer_params = kv_transfer_params + return resp + + +# ====================================================================== +# _map_finish_reason +# ====================================================================== + + +class TestMapFinishReason: + @pytest.mark.parametrize( + "openai, cohere", + [ + ("stop", "COMPLETE"), + ("length", "MAX_TOKENS"), + ("tool_calls", "TOOL_CALL"), + ("stop_sequence", "STOP_SEQUENCE"), + ("error", "ERROR"), + (None, "COMPLETE"), + ], + ) + def test_known_reasons(self, openai, cohere): + assert _map_finish_reason(openai) == cohere + + def test_unknown_reason_defaults_to_complete(self): + assert _map_finish_reason("not_a_real_reason") == "COMPLETE" + + def test_finish_reason_map_is_complete(self): + # Sanity check that the lookup table covers all documented states. + assert set(_FINISH_REASON_MAP) == { + "stop", + "length", + "tool_calls", + "stop_sequence", + "error", + None, + } + + +# ====================================================================== +# _coerce_text_content (system / tool string fallback) +# ====================================================================== + + +class TestCoerceTextContent: + def test_string_passthrough(self): + assert CohereServingChatV2._coerce_text_content("hi") == "hi" + + def test_concatenates_text_blocks(self): + from cohere.types import SystemChatMessageV2 + + sys_msg = SystemChatMessageV2( + content=[ + {"type": "text", "text": "a"}, + {"type": "text", "text": "b"}, + ] + ) + assert CohereServingChatV2._coerce_text_content(sys_msg.content) == "ab" + + +# ====================================================================== +# User message conversion +# ====================================================================== + + +class TestConvertUserMessage: + def test_string_content(self): + req = _make_request(messages=[{"role": "user", "content": "hi"}]) + result = _convert(req) + assert result.messages == [{"role": "user", "content": "hi"}] + + def test_text_only_list_flattened_to_string(self): + # Single-text-block list is flattened back to a string for maximum + # downstream-template compatibility. + req = _make_request( + messages=[ + { + "role": "user", + "content": [{"type": "text", "text": "hi"}], + } + ] + ) + result = _convert(req) + assert result.messages[0] == {"role": "user", "content": "hi"} + + def test_image_url_content_with_detail(self): + req = _make_request( + messages=[ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": { + "url": "data:image/png;base64,xxx", + "detail": "high", + }, + } + ], + } + ] + ) + result = _convert(req) + msg = result.messages[0] + assert msg["role"] == "user" + assert msg["content"] == [ + { + "type": "image_url", + "image_url": { + "url": "data:image/png;base64,xxx", + "detail": "high", + }, + } + ] + + def test_image_url_without_detail_omits_field(self): + req = _make_request( + messages=[ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": {"url": "https://x/i.png"}, + } + ], + } + ] + ) + result = _convert(req) + assert result.messages[0]["content"][0]["image_url"] == { + "url": "https://x/i.png" + } + + def test_text_plus_image_keeps_list(self): + req = _make_request( + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "describe"}, + { + "type": "image_url", + "image_url": {"url": "https://x/i.png"}, + }, + ], + } + ] + ) + result = _convert(req) + content = result.messages[0]["content"] + assert isinstance(content, list) + assert len(content) == 2 + assert content[0] == {"type": "text", "text": "describe"} + assert content[1]["type"] == "image_url" + + +# ====================================================================== +# Assistant message conversion +# ====================================================================== + + +class TestConvertAssistantMessage: + def test_string_content(self): + req = _make_request( + messages=[ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + ] + ) + result = _convert(req) + asst = result.messages[1] + assert asst == {"role": "assistant", "content": "hello"} + + def test_text_and_thinking_blocks(self): + # ``thinking`` blocks collapse back into the ``reasoning`` field on + # the OpenAI message; ``text`` blocks become ``content``. + req = _make_request( + messages=[ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "let me think"}, + {"type": "text", "text": "Hi!"}, + ], + }, + ] + ) + result = _convert(req) + asst = result.messages[1] + assert asst["role"] == "assistant" + assert asst["content"] == "Hi!" + assert asst["reasoning"] == "let me think" + + def test_thinking_only(self): + # Thinking-only assistant messages have no ``content`` set. + req = _make_request( + messages=[ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "ponder"}, + ], + }, + ] + ) + result = _convert(req) + asst = result.messages[1] + assert asst.get("reasoning") == "ponder" + assert "content" not in asst + + def test_multiple_thinking_blocks_concatenated(self): + req = _make_request( + messages=[ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "first."}, + {"type": "thinking", "thinking": "second."}, + {"type": "text", "text": "done."}, + ], + }, + ] + ) + result = _convert(req) + asst = result.messages[1] + assert asst["reasoning"] == "first.second." + assert asst["content"] == "done." + + def test_tool_plan_collapses_into_reasoning(self): + # Cohere's ``tool_plan`` is the older chain-of-thought field; it + # should be appended to ``reasoning`` so the rendered template + # preserves the planning context. + req = _make_request( + messages=[ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": [ + {"type": "text", "text": "I'll call a tool."}, + ], + "tool_plan": "plan: use calculator", + "tool_calls": [ + { + "id": "c1", + "type": "function", + "function": {"name": "calc", "arguments": "{}"}, + } + ], + }, + ] + ) + result = _convert(req) + asst = result.messages[1] + assert asst["content"] == "I'll call a tool." + assert asst["reasoning"] == "plan: use calculator" + assert asst["tool_calls"][0]["function"] == { + "name": "calc", + "arguments": "{}", + } + + def test_tool_calls_with_missing_function_pieces_get_defaults(self): + # The conversion defends against missing function name/arguments + # by emitting empty string / "{}" defaults so downstream + # validation never sees a None. + req = _make_request( + messages=[ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "tool_calls": [ + { + "id": "c1", + "type": "function", + "function": {"name": "", "arguments": ""}, + } + ], + }, + ] + ) + result = _convert(req) + tc = result.messages[1]["tool_calls"][0] + assert tc["id"] == "c1" + assert tc["type"] == "function" + assert tc["function"] == {"name": "", "arguments": "{}"} + + +# ====================================================================== +# Tool message conversion +# ====================================================================== + + +class TestConvertToolMessage: + def _request_with_tool_message(self, content: Any) -> CohereChatV2Request: + return _make_request( + messages=[ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "tool_calls": [ + { + "id": "c1", + "type": "function", + "function": {"name": "f", "arguments": "{}"}, + } + ], + }, + {"role": "tool", "tool_call_id": "c1", "content": content}, + ] + ) + + def test_string_content(self): + req = self._request_with_tool_message("result text") + result = _convert(req) + tool_msg = result.messages[-1] + assert tool_msg == { + "role": "tool", + "tool_call_id": "c1", + "content": "result text", + } + + def test_text_only_list_flattened_to_newline_string(self): + # Text-only tool results are flattened to a single newline-joined + # string for compatibility with vanilla chat templates. + req = self._request_with_tool_message( + [ + {"type": "text", "text": "line 1"}, + {"type": "text", "text": "line 2"}, + ] + ) + result = _convert(req) + tool_msg = result.messages[-1] + assert tool_msg["content"] == "line 1\nline 2" + + def test_with_document_preserves_structured_content(self): + # When documents appear in the tool result, we keep the list shape + # so the Cohere renderer can lift them into grounding sources. + req = self._request_with_tool_message( + [ + {"type": "text", "text": "see attachment"}, + { + "type": "document", + "document": {"data": {"text": "doc text"}, "id": "d1"}, + }, + ] + ) + result = _convert(req) + tool_msg = result.messages[-1] + assert isinstance(tool_msg["content"], list) + assert tool_msg["content"][0] == {"type": "text", "text": "see attachment"} + assert tool_msg["content"][1] == { + "type": "document", + "document": {"data": {"text": "doc text"}, "id": "d1"}, + } + + +# ====================================================================== +# System message +# ====================================================================== + + +class TestSystemMessage: + def test_system_string(self): + req = _make_request( + messages=[ + {"role": "system", "content": "be helpful"}, + {"role": "user", "content": "hi"}, + ] + ) + result = _convert(req) + assert result.messages[0] == { + "role": "system", + "content": "be helpful", + } + + def test_system_text_blocks_concatenated(self): + req = _make_request( + messages=[ + { + "role": "system", + "content": [ + {"type": "text", "text": "part1 "}, + {"type": "text", "text": "part2"}, + ], + }, + {"role": "user", "content": "hi"}, + ] + ) + result = _convert(req) + assert result.messages[0]["content"] == "part1 part2" + + +# ====================================================================== +# Base ChatCompletionRequest field mapping +# ====================================================================== + + +class TestBuildBaseChatCompletion: + def test_sampling_and_limits_mapped(self): + req = _make_request( + max_tokens=128, + stop_sequences=["", "STOP"], + temperature=0.5, + seed=42, + frequency_penalty=0.1, + presence_penalty=0.2, + k=50, + p=0.95, + logprobs=True, + priority=2, + kv_transfer_params={"x": 1}, + chat_template_kwargs={"y": 2}, + ) + result = _convert(req) + assert result.model == "m" + # ``max_tokens`` is deprecated in favor of ``max_completion_tokens`` + # but the serving code intentionally sets both for compatibility. + assert result.max_completion_tokens == 128 + assert result.stop == ["", "STOP"] + assert result.temperature == 0.5 + assert result.seed == 42 + assert result.frequency_penalty == 0.1 + assert result.presence_penalty == 0.2 + assert result.top_k == 50 + assert result.top_p == 0.95 + assert result.logprobs is True + assert result.priority == 2 + assert result.kv_transfer_params == {"x": 1} + # ``chat_template_kwargs`` may be expanded by _apply_cohere_*; the + # base build at least preserves what the caller passed. + assert (result.chat_template_kwargs or {}).get("y") == 2 + + def test_priority_defaults_to_zero(self): + # ChatCompletionRequest.priority defaults to 0; ``None`` Cohere + # priority must be coerced rather than passed through. + req = _make_request() + result = _convert(req) + assert result.priority == 0 + + +# ====================================================================== +# Streaming options +# ====================================================================== + + +class TestStreamingOptions: + def test_no_stream_leaves_defaults(self): + result = _convert(_make_request(stream=False)) + assert not result.stream + assert result.stream_options is None + + def test_stream_enables_usage_options(self): + # The v2 translator forces ``include_usage=True`` so the + # ``message-end`` event can surface ``billed_units`` / ``tokens``; + # ``continuous_usage_stats`` is intentionally left at its + # ``StreamOptions`` default (False) — Cohere v2 only reports + # usage on the terminal event. + result = _convert(_make_request(stream=True)) + assert result.stream is True + assert result.stream_options is not None + assert result.stream_options.include_usage is True + + +# ====================================================================== +# Response format +# ====================================================================== + + +class TestResponseFormat: + def test_text_is_passthrough(self): + result = _convert(_make_request(response_format={"type": "text"})) + assert result.response_format is None + + def test_json_object(self): + result = _convert(_make_request(response_format={"type": "json_object"})) + assert result.response_format is not None + assert result.response_format.type == "json_object" + assert result.response_format.json_schema is None + + def test_json_schema(self): + schema = {"type": "object", "properties": {"a": {"type": "string"}}} + result = _convert( + _make_request( + response_format={"type": "json_object", "json_schema": schema} + ) + ) + assert result.response_format is not None + assert result.response_format.type == "json_schema" + assert result.response_format.json_schema is not None + assert result.response_format.json_schema.name == "cohere_v2_json_schema" + # ``JsonSchemaResponseFormat.json_schema`` has alias=``schema`` on + # the Pydantic field, so we observe the value via the serialized + # payload (which is what downstream consumers actually read). + dumped = result.response_format.json_schema.model_dump(exclude_none=True) + assert dumped["json_schema"] == schema + + +# ====================================================================== +# Tools / tool_choice +# ====================================================================== + + +class TestApplyTools: + def test_no_tools(self): + result = _convert(_make_request()) + assert result.tools is None + + def test_basic_tool(self): + result = _convert( + _make_request( + tools=[ + { + "type": "function", + "function": { + "name": "calc", + "description": "calculator", + "parameters": {"type": "object"}, + }, + } + ] + ) + ) + assert result.tools is not None + assert len(result.tools) == 1 + tool = result.tools[0] + assert tool.type == "function" + assert tool.function.name == "calc" + assert tool.function.description == "calculator" + # ``strict`` is an extra attribute on FunctionDefinition (the + # field is only stamped onto the OpenAI tool when strict_tools is + # set on the request). The default path must not set it. + assert getattr(tool.function, "strict", None) is None + + def test_strict_tools_propagates_to_function(self): + result = _convert( + _make_request( + strict_tools=True, + tools=[ + { + "type": "function", + "function": { + "name": "calc", + "description": "", + "parameters": {}, + }, + } + ], + ) + ) + assert result.tools[0].function.strict is True + + +class TestApplyToolChoice: + def test_required(self): + result = _convert( + _make_request( + tool_choice="REQUIRED", + tools=[ + { + "type": "function", + "function": { + "name": "f", + "description": "", + "parameters": {}, + }, + } + ], + ) + ) + assert result.tool_choice == "required" + + def test_none(self): + result = _convert( + _make_request( + tool_choice="NONE", + tools=[ + { + "type": "function", + "function": { + "name": "f", + "description": "", + "parameters": {}, + }, + } + ], + ) + ) + assert result.tool_choice == "none" + + def test_default_to_auto_when_tools_present(self): + # No explicit ``tool_choice`` + tools present → auto, mirroring + # Cohere's documented "free choice" default. + result = _convert( + _make_request( + tools=[ + { + "type": "function", + "function": { + "name": "f", + "description": "", + "parameters": {}, + }, + } + ] + ) + ) + assert result.tool_choice == "auto" + + def test_no_tools_no_choice_left_unset(self): + # When there are no tools the underlying ChatCompletionRequest + # default applies; we must not stamp ``auto``. + result = _convert(_make_request()) + assert result.tool_choice != "auto" + + +# ====================================================================== +# Cohere-specific template kwargs forwarding +# ====================================================================== + + +class TestApplyCohereTemplateKwargs: + def test_string_documents_wrapped(self): + result = _convert(_make_request(documents=["doc 1", "doc 2"])) + docs = (result.chat_template_kwargs or {}).get("documents") + assert docs == [ + {"id": "doc_0", "data": {"text": "doc 1"}}, + {"id": "doc_1", "data": {"text": "doc 2"}}, + ] + + def test_document_with_explicit_id_preserved(self): + result = _convert( + _make_request( + documents=[ + {"id": "custom", "data": {"text": "t"}}, + {"data": {"text": "t2"}}, # no id -> synthesized + ] + ) + ) + docs = result.chat_template_kwargs["documents"] + assert docs[0] == {"id": "custom", "data": {"text": "t"}} + assert docs[1]["id"] == "doc_1" + + def test_safety_mode_normalized_to_lowercase(self): + result = _convert(_make_request(safety_mode="CONTEXTUAL")) + assert result.chat_template_kwargs["safety_mode"] == "contextual" + + def test_citation_options_forwarded_as_dict(self): + result = _convert(_make_request(citation_options={"mode": "accurate"})) + assert result.chat_template_kwargs["citation_options"] == {"mode": "accurate"} + + def test_thinking_forwarded_as_dict(self): + result = _convert( + _make_request(thinking={"type": "enabled", "token_budget": 16}) + ) + assert result.chat_template_kwargs["thinking"] == { + "type": "enabled", + "token_budget": 16, + } + + def test_strict_tools_forwarded(self): + result = _convert(_make_request(strict_tools=True)) + assert result.chat_template_kwargs["strict_tools"] is True + + def test_existing_chat_template_kwargs_preserved(self): + # User-supplied kwargs should not be clobbered by the v2 fields + # (setdefault semantics). + result = _convert( + _make_request( + chat_template_kwargs={ + "safety_mode": "user-explicit", + "extra": "x", + }, + safety_mode="CONTEXTUAL", + ) + ) + assert result.chat_template_kwargs["safety_mode"] == "user-explicit" + assert result.chat_template_kwargs["extra"] == "x" + + def test_no_template_kwargs_when_no_cohere_fields(self): + # Without any of the Cohere-specific fields and no caller-supplied + # kwargs, we must leave ``chat_template_kwargs`` as None so other + # renderers see a clean request. + result = _convert(_make_request()) + assert result.chat_template_kwargs is None + + +# ====================================================================== +# _chat_completion_to_v2 (non-streaming response builder) +# ====================================================================== + + +class TestChatCompletionToV2: + def test_text_only(self): + serving = _serving(is_reasoning_model=True) + resp = _build_chat_completion_response(content="hello") + v2 = serving._chat_completion_to_v2(resp, _make_request()) + assert isinstance(v2, CohereChatV2Response) + assert v2.id == "resp_1" + assert v2.finish_reason == "COMPLETE" + assert v2.message.role == "assistant" + assert len(v2.message.content) == 1 + assert v2.message.content[0].type == "text" + assert v2.message.content[0].text == "hello" + assert v2.message.tool_calls is None + assert v2.message.tool_plan is None + assert v2.usage is None + + def test_reasoning_model_keeps_thinking_with_tool_calls(self): + # Reasoning Command models: ``thinking`` block stays in + # ``message.content`` and ``tool_plan`` is left unset, even when + # tool calls are present. + serving = _serving(is_reasoning_model=True) + resp = _build_chat_completion_response( + content="resp text", + reasoning="thoughts", + tool_calls=[ + { + "id": "c1", + "type": "function", + "function": {"name": "f", "arguments": "{}"}, + } + ], + finish_reason="tool_calls", + ) + v2 = serving._chat_completion_to_v2(resp, _make_request()) + assert v2.finish_reason == "TOOL_CALL" + assert v2.message.tool_plan is None + types = [c.type for c in v2.message.content] + assert types == ["thinking", "text"] + assert v2.message.content[0].thinking == "thoughts" + assert v2.message.content[1].text == "resp text" + assert v2.message.tool_calls[0].id == "c1" + assert v2.message.tool_calls[0].function.name == "f" + + def test_non_reasoning_model_moves_reasoning_to_tool_plan(self): + # Older Command models surface reasoning as ``tool_plan`` on tool- + # call turns; the thinking block should be dropped from content. + serving = _serving(is_reasoning_model=False) + resp = _build_chat_completion_response( + content=None, + reasoning="plan", + tool_calls=[ + { + "id": "c1", + "type": "function", + "function": {"name": "f", "arguments": "{}"}, + } + ], + finish_reason="tool_calls", + ) + v2 = serving._chat_completion_to_v2(resp, _make_request()) + assert v2.message.tool_plan == "plan" + assert v2.message.content is None + assert v2.message.tool_calls[0].id == "c1" + + def test_non_reasoning_model_keeps_thinking_when_no_tool_calls(self): + # No tool calls => non-reasoning behavior is identical to + # reasoning behavior; the thinking block stays. + serving = _serving(is_reasoning_model=False) + resp = _build_chat_completion_response( + content="answer", reasoning="plan", tool_calls=None + ) + v2 = serving._chat_completion_to_v2(resp, _make_request()) + assert v2.message.tool_plan is None + types = [c.type for c in v2.message.content] + assert types == ["thinking", "text"] + + def test_id_synthesized_when_response_id_missing(self): + serving = _serving() + resp = _build_chat_completion_response(content="hi", response_id="") + v2 = serving._chat_completion_to_v2(resp, _make_request()) + assert v2.id.startswith("chat_") + + def test_kv_transfer_params_propagated(self): + serving = _serving() + resp = _build_chat_completion_response( + content="hi", kv_transfer_params={"k": 1} + ) + v2 = serving._chat_completion_to_v2(resp, _make_request()) + assert v2.kv_transfer_params == {"k": 1} + + +# ====================================================================== +# _build_usage +# ====================================================================== + + +class TestBuildUsage: + def test_none_passthrough(self): + resp = _build_chat_completion_response(content="hi") + # default usage is None + assert CohereServingChatV2._build_usage(resp) is None + + def test_basic_usage(self): + resp = _build_chat_completion_response( + content="hi", + usage={ + "prompt_tokens": 10, + "completion_tokens": 5, + "total_tokens": 15, + }, + ) + usage = CohereServingChatV2._build_usage(resp) + assert usage is not None + assert usage.billed_units.input_tokens == 10 + assert usage.billed_units.output_tokens == 5 + assert usage.tokens.input_tokens == 10 + assert usage.tokens.output_tokens == 5 + assert usage.cached_tokens is None + + def test_completion_tokens_default_to_zero_when_missing(self): + resp = _build_chat_completion_response( + content="hi", + usage={"prompt_tokens": 10, "completion_tokens": 0, "total_tokens": 10}, + ) + usage = CohereServingChatV2._build_usage(resp) + assert usage.billed_units.output_tokens == 0 + + def test_cached_tokens_propagated(self): + resp = _build_chat_completion_response( + content="hi", + usage={ + "prompt_tokens": 10, + "completion_tokens": 5, + "total_tokens": 15, + "prompt_tokens_details": {"cached_tokens": 3}, + }, + ) + usage = CohereServingChatV2._build_usage(resp) + assert usage.cached_tokens == 3 + + +# ====================================================================== +# _extract_citations_if_any +# ====================================================================== + + +class TestExtractCitations: + def test_none_or_empty_returns_none(self): + assert ( + CohereServingChatV2._extract_citations_if_any( + type("M", (), {"citations": None})() + ) + is None + ) + assert ( + CohereServingChatV2._extract_citations_if_any( + type("M", (), {"citations": []})() + ) + is None + ) + # Missing field entirely. + assert ( + CohereServingChatV2._extract_citations_if_any(type("M", (), {})()) is None + ) + + def test_vllm_citation_objects_normalized(self): + msg = type("M", (), {})() + msg.citations = [ + VLLMCitation( + start=0, + end=5, + text="hello", + sources=[CitationSource(type="document", id="d1")], + ) + ] + out = CohereServingChatV2._extract_citations_if_any(msg) + assert out is not None + assert len(out) == 1 + assert out[0].start == 0 + assert out[0].end == 5 + assert out[0].text == "hello" + + def test_dict_citation_payloads_accepted(self): + msg = type("M", (), {})() + msg.citations = [ + { + "start": 0, + "end": 3, + "text": "hi!", + "sources": [{"type": "document", "id": "d1"}], + } + ] + out = CohereServingChatV2._extract_citations_if_any(msg) + assert out is not None + assert out[0].start == 0 + assert out[0].text == "hi!" + + def test_malformed_citation_skipped_not_raised(self): + msg = type("M", (), {})() + msg.citations = [object()] # neither dict nor Pydantic + out = CohereServingChatV2._extract_citations_if_any(msg) + # All citations dropped → None + assert out is None + + +# ====================================================================== +# create_error_response sanity +# ====================================================================== + + +class TestCreateErrorResponse: + def test_envelope_uses_400(self): + serving = _serving() + err = serving.create_error_response("oops") + assert err.error.message == "oops" + assert err.error.code == 400 + assert err.error.type == "bad_request" + + +# ====================================================================== +# ContentBlockType enum +# ====================================================================== + + +class TestContentBlockType: + def test_values(self): + assert ContentBlockType.THINKING == "thinking" + assert ContentBlockType.TEXT == "text" + assert ContentBlockType.TOOL_CALL == "tool_call" diff --git a/tests/entrypoints/cohere/test_serving_streaming.py b/tests/entrypoints/cohere/test_serving_streaming.py new file mode 100644 index 000000000000..baf006cca627 --- /dev/null +++ b/tests/entrypoints/cohere/test_serving_streaming.py @@ -0,0 +1,815 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for the Cohere v2 SSE stream conversion in +``vllm/entrypoints/cohere/serving.py``. + +The stream-translation entry point is +:meth:`CohereServingChatV2._chat_completion_stream_to_v2`, which turns an +async iterable of OpenAI SSE chunks into Cohere's +``message-start → (content|tool-call|citation)* → message-end → [DONE]`` +event stream. + +We test the helpers (``_StreamState``, ``_handle_*_delta``, etc.) in +isolation, plus a handful of end-to-end scenarios that exercise the +state machine. +""" + +import json +from collections.abc import AsyncGenerator +from typing import Any + +import pytest + +from vllm.entrypoints.cohere.protocol import ( + CohereChatV2Request, + MessageStartEvent, +) +from vllm.entrypoints.cohere.serving import ( + _DONE_FRAME, + CohereServingChatV2, + ContentBlockType, + _emit, + _sse, + _StreamState, +) +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionStreamResponse, +) +from vllm.entrypoints.openai.engine.protocol import ( + Citation as VLLMCitation, +) +from vllm.entrypoints.openai.engine.protocol import ( + CitationSource, +) + +# ---------------------------------------------------------------------- +# Helpers +# ---------------------------------------------------------------------- + + +class _FakeServing(CohereServingChatV2): + """Lightweight stand-in for :class:`CohereServingChatV2` that skips + the heavy ``OpenAIServingChat.__init__`` chain (which would need a + real engine client, model registry, render service, etc.). + + Only ``_is_reasoning_model`` is read by the methods under test + (``_chat_completion_stream_to_v2`` and the per-delta handlers); the + rest is dead weight for unit testing. + """ + + def __init__(self, is_reasoning_model: bool = True) -> None: + # Intentionally skipping super().__init__ — see class docstring. + self._is_reasoning_model = is_reasoning_model + + +def _serving(is_reasoning_model: bool = True) -> CohereServingChatV2: + return _FakeServing(is_reasoning_model=is_reasoning_model) + + +def _parse_event(frame: str) -> dict[str, Any]: + """Strip the ``data: ... \\n\\n`` wrapper and parse the JSON payload.""" + assert frame.startswith("data: ") + assert frame.endswith("\n\n") + return json.loads(frame[len("data: ") : -2]) + + +def _make_chunk( + *, + chunk_id: str = "chunk_0", + role: str | None = None, + content: str | None = None, + reasoning: str | None = None, + tool_calls: list[dict[str, Any]] | None = None, + finish_reason: str | None = None, + usage: dict[str, Any] | None = None, + omit_choices: bool = False, + citations: list[Any] | None = None, +) -> str: + """Build the ``data: {...}\\n\\n`` SSE frame the production code + consumes.""" + delta: dict[str, Any] = {} + if role is not None: + delta["role"] = role + if content is not None: + delta["content"] = content + if reasoning is not None: + delta["reasoning"] = reasoning + if tool_calls is not None: + delta["tool_calls"] = tool_calls + if citations is not None: + delta["citations"] = citations + + payload: dict[str, Any] = { + "id": chunk_id, + "object": "chat.completion.chunk", + "created": 0, + "model": "m", + } + if not omit_choices: + payload["choices"] = [ + {"index": 0, "delta": delta, "finish_reason": finish_reason} + ] + else: + payload["choices"] = [] + if usage is not None: + payload["usage"] = usage + # Use ChatCompletionStreamResponse to normalize the payload shape. + chunk = ChatCompletionStreamResponse.model_validate(payload) + return f"data: {chunk.model_dump_json(exclude_none=False)}\n\n" + + +def _make_done() -> str: + return "data: [DONE]\n\n" + + +async def _stream_from(items: list[str]) -> AsyncGenerator[str, None]: + for item in items: + yield item + + +async def _drain(serving: CohereServingChatV2, items: list[str]) -> list[str]: + """Drive ``_chat_completion_stream_to_v2`` over ``items`` and collect + the emitted SSE frames.""" + request = CohereChatV2Request( + model="m", messages=[{"role": "user", "content": "hi"}], stream=True + ) + gen = serving._chat_completion_stream_to_v2(_stream_from(items), request) + return [frame async for frame in gen] + + +# ====================================================================== +# Low-level helpers: _sse, _emit, _DONE_FRAME +# ====================================================================== + + +class TestSSEHelpers: + def test_sse_wraps_payload(self): + assert _sse("hello") == "data: hello\n\n" + + def test_done_frame_constant(self): + # cohere-python and Fern-generated clients key off this exact + # sentinel; keep it byte-for-byte stable. + assert _DONE_FRAME == "data: [DONE]\n\n" + + def test_emit_serializes_event_with_type_discriminator(self): + frame = _emit( + MessageStartEvent(id="abc", delta={"message": {"role": "assistant"}}) + ) + payload = _parse_event(frame) + assert payload["type"] == "message-start" + assert payload["id"] == "abc" + assert payload["delta"] == {"message": {"role": "assistant"}} + + +# ====================================================================== +# _StreamState +# ====================================================================== + + +class TestStreamState: + def test_defaults(self): + st = _StreamState() + assert st.started is False + assert st.ended is False + assert st.finish_reason is None + assert st.last_chunk_id == "" + assert st.active_block is None + assert st.active_block_index is None + assert st.active_tool_index is None + assert st.tool_calls_seen == set() + + def test_content_index_monotonic(self): + st = _StreamState() + assert st.next_content_index() == 0 + assert st.next_content_index() == 1 + assert st.next_content_index() == 2 + + def test_citation_index_separate_from_content_index(self): + st = _StreamState() + st.next_content_index() # 0 + st.next_content_index() # 1 + # Citation indexing is independent of content indexing. + assert st.next_citation_index() == 0 + assert st.next_citation_index() == 1 + + +# ====================================================================== +# _close_open_blocks +# ====================================================================== + + +class TestCloseOpenBlocks: + def test_no_block_open_emits_nothing(self): + serving = _serving() + state = _StreamState() + assert serving._close_open_blocks(state) == [] + + def test_text_block_emits_content_end(self): + serving = _serving() + state = _StreamState() + state.active_block = ContentBlockType.TEXT + state.active_block_index = 1 + out = serving._close_open_blocks(state) + assert len(out) == 1 + payload = _parse_event(out[0]) + assert payload == {"type": "content-end", "index": 1} + assert state.active_block is None + assert state.active_block_index is None + + def test_thinking_block_emits_content_end(self): + serving = _serving() + state = _StreamState() + state.active_block = ContentBlockType.THINKING + state.active_block_index = 3 + out = serving._close_open_blocks(state) + payload = _parse_event(out[0]) + assert payload == {"type": "content-end", "index": 3} + + def test_tool_call_block_emits_tool_call_end(self): + serving = _serving() + state = _StreamState() + state.active_block = ContentBlockType.TOOL_CALL + state.active_tool_index = 7 + out = serving._close_open_blocks(state) + payload = _parse_event(out[0]) + assert payload == {"type": "tool-call-end", "index": 7} + assert state.active_tool_index is None + + +# ====================================================================== +# _handle_text_delta +# ====================================================================== + + +class TestHandleTextDelta: + def test_opens_block_first_time(self): + serving = _serving() + state = _StreamState() + events = serving._handle_text_delta(state, "Hi") + assert len(events) == 2 + start = _parse_event(events[0]) + delta = _parse_event(events[1]) + assert start["type"] == "content-start" + assert start["index"] == 0 + assert start["delta"]["message"]["content"]["type"] == "text" + assert delta["type"] == "content-delta" + assert delta["index"] == 0 + assert delta["delta"]["message"]["content"]["text"] == "Hi" + assert state.active_block == ContentBlockType.TEXT + assert state.active_block_index == 0 + + def test_continues_block_with_just_delta(self): + serving = _serving() + state = _StreamState() + serving._handle_text_delta(state, "Hi") + events = serving._handle_text_delta(state, " there") + # Only a delta event, no new content-start. + assert len(events) == 1 + delta = _parse_event(events[0]) + assert delta["type"] == "content-delta" + assert delta["delta"]["message"]["content"]["text"] == " there" + + def test_switches_from_thinking_block(self): + serving = _serving(is_reasoning_model=True) + state = _StreamState() + # Open a thinking block first, then switch to text. + serving._handle_thinking_delta(state, "ponder") + events = serving._handle_text_delta(state, "answer") + types = [_parse_event(ev)["type"] for ev in events] + assert types == ["content-end", "content-start", "content-delta"] + # The text block gets a new index (1), distinct from thinking's 0. + assert _parse_event(events[1])["index"] == 1 + assert state.active_block == ContentBlockType.TEXT + + +# ====================================================================== +# _handle_thinking_delta +# ====================================================================== + + +class TestHandleThinkingDelta: + def test_reasoning_model_opens_thinking_block(self): + serving = _serving(is_reasoning_model=True) + state = _StreamState() + events = serving._handle_thinking_delta(state, "thought") + assert len(events) == 2 + start = _parse_event(events[0]) + delta = _parse_event(events[1]) + assert start["type"] == "content-start" + assert start["delta"]["message"]["content"]["type"] == "thinking" + assert delta["type"] == "content-delta" + assert delta["delta"]["message"]["content"]["thinking"] == "thought" + assert state.active_block == ContentBlockType.THINKING + + def test_reasoning_model_continues_thinking_block(self): + serving = _serving(is_reasoning_model=True) + state = _StreamState() + serving._handle_thinking_delta(state, "first") + events = serving._handle_thinking_delta(state, " more") + assert len(events) == 1 + assert _parse_event(events[0])["type"] == "content-delta" + + def test_non_reasoning_model_emits_tool_plan_delta(self): + # Older Command models stream reasoning as ``tool_plan`` deltas; + # no content-start/end pair is emitted. + serving = _serving(is_reasoning_model=False) + state = _StreamState() + events = serving._handle_thinking_delta(state, "planning") + assert len(events) == 1 + payload = _parse_event(events[0]) + assert payload["type"] == "tool-plan-delta" + assert payload["delta"]["message"]["tool_plan"] == "planning" + # ``tool_plan`` deltas don't claim an active content block. + assert state.active_block is None + + def test_non_reasoning_model_closes_open_text_block(self): + serving = _serving(is_reasoning_model=False) + state = _StreamState() + # Open a text block first. + serving._handle_text_delta(state, "answer") + events = serving._handle_thinking_delta(state, "rethink") + types = [_parse_event(ev)["type"] for ev in events] + assert types == ["content-end", "tool-plan-delta"] + + +# ====================================================================== +# _handle_tool_call_deltas +# ====================================================================== + + +class TestHandleToolCallDeltas: + def test_new_tool_call_opens_tool_call_start(self): + serving = _serving() + state = _StreamState() + deltas = [ + type( + "Delta", + (), + { + "index": 0, + "id": "c1", + "function": type( + "Fn", (), {"name": "calc", "arguments": '{"x":'} + )(), + }, + )() + ] + events = serving._handle_tool_call_deltas(state, deltas) + assert len(events) == 1 + payload = _parse_event(events[0]) + assert payload["type"] == "tool-call-start" + assert payload["index"] == 0 + tc = payload["delta"]["message"]["tool_calls"] + assert tc["id"] == "c1" + assert tc["function"]["name"] == "calc" + assert tc["function"]["arguments"] == '{"x":' + assert state.active_block == ContentBlockType.TOOL_CALL + assert state.active_tool_index == 0 + assert 0 in state.tool_calls_seen + + def test_subsequent_arguments_emit_delta(self): + serving = _serving() + state = _StreamState() + # First call: start. + first = [ + type( + "Delta", + (), + { + "index": 0, + "id": "c1", + "function": type("Fn", (), {"name": "calc", "arguments": ""})(), + }, + )() + ] + serving._handle_tool_call_deltas(state, first) + # Second call: same index, additional arguments fragment. + more = [ + type( + "Delta", + (), + { + "index": 0, + "id": None, + "function": type("Fn", (), {"name": None, "arguments": "1}"})(), + }, + )() + ] + events = serving._handle_tool_call_deltas(state, more) + assert len(events) == 1 + payload = _parse_event(events[0]) + assert payload["type"] == "tool-call-delta" + assert payload["index"] == 0 + assert ( + payload["delta"]["message"]["tool_calls"]["function"]["arguments"] == "1}" + ) + + def test_new_tool_call_closes_existing_content_block(self): + serving = _serving() + state = _StreamState() + # Open a text block, then start a tool call. + serving._handle_text_delta(state, "I'll call:") + deltas = [ + type( + "Delta", + (), + { + "index": 0, + "id": "c1", + "function": type("Fn", (), {"name": "calc", "arguments": "{}"})(), + }, + )() + ] + events = serving._handle_tool_call_deltas(state, deltas) + types = [_parse_event(ev)["type"] for ev in events] + assert types == ["content-end", "tool-call-start"] + + +# ====================================================================== +# _handle_citation_deltas +# ====================================================================== + + +class TestHandleCitationDeltas: + def test_citation_objects_emit_start_and_end(self): + serving = _serving() + state = _StreamState() + citations = [ + VLLMCitation( + start=0, + end=5, + text="hello", + sources=[CitationSource(type="document", id="d1")], + ) + ] + events = serving._handle_citation_deltas(state, citations) + assert len(events) == 2 + start = _parse_event(events[0]) + end = _parse_event(events[1]) + assert start["type"] == "citation-start" + assert start["index"] == 0 + cit_payload = start["delta"]["message"]["citations"] + assert cit_payload["start"] == 0 + assert cit_payload["end"] == 5 + assert cit_payload["text"] == "hello" + assert end == {"type": "citation-end", "index": 0} + + def test_dict_citations_accepted(self): + serving = _serving() + state = _StreamState() + citations = [{"start": 0, "end": 3, "text": "Hi!", "sources": []}] + events = serving._handle_citation_deltas(state, citations) + assert len(events) == 2 + assert _parse_event(events[0])["type"] == "citation-start" + assert _parse_event(events[1])["type"] == "citation-end" + + def test_malformed_citation_skipped(self): + serving = _serving() + state = _StreamState() + # Neither a dict nor a Pydantic-model-like object. + events = serving._handle_citation_deltas(state, [object()]) + assert events == [] + + +# ====================================================================== +# _build_message_end_event +# ====================================================================== + + +class TestBuildMessageEndEvent: + def test_without_usage(self): + serving = _serving() + frame = serving._build_message_end_event(chunk_id="abc", finish_reason="stop") + payload = _parse_event(frame) + assert payload["type"] == "message-end" + assert payload["id"] == "abc" + assert payload["delta"]["finish_reason"] == "COMPLETE" + assert "usage" not in payload["delta"] + + def test_with_usage(self): + serving = _serving() + chunk = ChatCompletionStreamResponse.model_validate( + { + "id": "abc", + "object": "chat.completion.chunk", + "created": 0, + "model": "m", + "choices": [], + "usage": { + "prompt_tokens": 10, + "completion_tokens": 5, + "total_tokens": 15, + }, + } + ) + frame = serving._build_message_end_event( + chunk_id="abc", finish_reason="length", usage_chunk=chunk + ) + payload = _parse_event(frame) + assert payload["delta"]["finish_reason"] == "MAX_TOKENS" + usage = payload["delta"]["usage"] + assert usage["billed_units"] == {"input_tokens": 10, "output_tokens": 5} + assert usage["tokens"] == {"input_tokens": 10, "output_tokens": 5} + assert "cached_tokens" not in usage + + def test_with_cached_tokens(self): + serving = _serving() + chunk = ChatCompletionStreamResponse.model_validate( + { + "id": "abc", + "object": "chat.completion.chunk", + "created": 0, + "model": "m", + "choices": [], + "usage": { + "prompt_tokens": 10, + "completion_tokens": 5, + "total_tokens": 15, + "prompt_tokens_details": {"cached_tokens": 3}, + }, + } + ) + frame = serving._build_message_end_event( + chunk_id="abc", finish_reason="stop", usage_chunk=chunk + ) + payload = _parse_event(frame) + assert payload["delta"]["usage"]["cached_tokens"] == 3 + + +# ====================================================================== +# End-to-end: _chat_completion_stream_to_v2 +# ====================================================================== + + +class TestChatCompletionStreamToV2: + """End-to-end stream lifecycle tests.""" + + @pytest.mark.asyncio + async def test_text_only_happy_path_emits_full_lifecycle(self): + serving = _serving(is_reasoning_model=True) + items = [ + _make_chunk(role="assistant"), + _make_chunk(content="Hi"), + _make_chunk(content=" there"), + _make_chunk( + omit_choices=True, + usage={ + "prompt_tokens": 4, + "completion_tokens": 2, + "total_tokens": 6, + }, + ), + ] + frames = await _drain(serving, items) + types = [_parse_event(f)["type"] for f in frames[:-1]] + # message-start, content-start, content-delta, content-delta, + # content-end, message-end, then [DONE] as the last frame. + assert types == [ + "message-start", + "content-start", + "content-delta", + "content-delta", + "content-end", + "message-end", + ] + assert frames[-1] == _DONE_FRAME + # message-end should carry usage stats from the trailing chunk. + end_payload = _parse_event(frames[-2]) + assert end_payload["delta"]["finish_reason"] == "COMPLETE" + assert end_payload["delta"]["usage"]["billed_units"]["input_tokens"] == 4 + + @pytest.mark.asyncio + async def test_text_then_tool_call_closes_text_first(self): + serving = _serving(is_reasoning_model=True) + items = [ + _make_chunk(role="assistant"), + _make_chunk(content="planning..."), + _make_chunk( + tool_calls=[ + { + "index": 0, + "id": "c1", + "type": "function", + "function": {"name": "calc", "arguments": "{}"}, + } + ] + ), + _make_chunk(finish_reason="tool_calls"), + _make_chunk( + omit_choices=True, + usage={ + "prompt_tokens": 1, + "completion_tokens": 1, + "total_tokens": 2, + }, + ), + ] + frames = await _drain(serving, items) + types = [_parse_event(f)["type"] for f in frames[:-1]] + # Text block is opened+delta, then closed before the tool-call + # opens, and the final close happens on the usage chunk path. + assert types == [ + "message-start", + "content-start", + "content-delta", + "content-end", + "tool-call-start", + "tool-call-end", + "message-end", + ] + # finish_reason captured from the prior chunk. + end = _parse_event(frames[-2]) + assert end["delta"]["finish_reason"] == "TOOL_CALL" + + @pytest.mark.asyncio + async def test_thinking_then_text_reasoning_model(self): + serving = _serving(is_reasoning_model=True) + items = [ + _make_chunk(role="assistant"), + _make_chunk(reasoning="thinking..."), + _make_chunk(content="answer"), + _make_chunk(finish_reason="stop"), + _make_chunk( + omit_choices=True, + usage={ + "prompt_tokens": 1, + "completion_tokens": 1, + "total_tokens": 2, + }, + ), + ] + frames = await _drain(serving, items) + types = [_parse_event(f)["type"] for f in frames[:-1]] + # Thinking block opens with index 0; text block reopens with index 1. + assert types == [ + "message-start", + "content-start", + "content-delta", + "content-end", + "content-start", + "content-delta", + "content-end", + "message-end", + ] + thinking_start = _parse_event(frames[1]) + assert thinking_start["delta"]["message"]["content"]["type"] == "thinking" + text_start = _parse_event(frames[4]) + assert text_start["delta"]["message"]["content"]["type"] == "text" + + @pytest.mark.asyncio + async def test_reasoning_on_non_reasoning_model_emits_tool_plan_delta(self): + serving = _serving(is_reasoning_model=False) + items = [ + _make_chunk(role="assistant"), + _make_chunk(reasoning="planning"), + _make_chunk( + tool_calls=[ + { + "index": 0, + "id": "c1", + "type": "function", + "function": {"name": "f", "arguments": "{}"}, + } + ] + ), + _make_chunk(finish_reason="tool_calls"), + _make_chunk( + omit_choices=True, + usage={ + "prompt_tokens": 1, + "completion_tokens": 1, + "total_tokens": 2, + }, + ), + ] + frames = await _drain(serving, items) + types = [_parse_event(f)["type"] for f in frames[:-1]] + assert types == [ + "message-start", + "tool-plan-delta", + "tool-call-start", + "tool-call-end", + "message-end", + ] + # No thinking content blocks should be present. + assert "content-start" not in types + assert "content-end" not in types + + @pytest.mark.asyncio + async def test_done_marker_in_middle_closes_open_block(self): + # Some upstreams send [DONE] without a trailing usage-only chunk. + # The translator must still emit message-end before [DONE] so + # Cohere clients don't hang. + serving = _serving(is_reasoning_model=True) + items = [ + _make_chunk(role="assistant"), + _make_chunk(content="Hi"), + _make_done(), + ] + frames = await _drain(serving, items) + types = [_parse_event(f)["type"] for f in frames[:-1]] + assert types == [ + "message-start", + "content-start", + "content-delta", + "content-end", + "message-end", + ] + assert frames[-1] == _DONE_FRAME + + @pytest.mark.asyncio + async def test_skips_empty_and_non_data_lines(self): + serving = _serving() + items = [ + "\n", + "event: ping\n\n", + _make_chunk(role="assistant"), + "data: \n\n", # empty data + _make_chunk(content="Hi"), + _make_chunk( + omit_choices=True, + usage={ + "prompt_tokens": 1, + "completion_tokens": 1, + "total_tokens": 2, + }, + ), + ] + frames = await _drain(serving, items) + assert frames[-1] == _DONE_FRAME + # Should still produce a complete lifecycle. + types = [_parse_event(f)["type"] for f in frames[:-1]] + assert types[0] == "message-start" + assert types[-1] == "message-end" + + @pytest.mark.asyncio + async def test_exception_in_chunk_parsing_emits_error_message_end(self): + # An invalid JSON payload after the first chunk triggers the + # error path: a synthetic message-end with finish_reason=ERROR + # followed by [DONE]. + serving = _serving() + items = [ + _make_chunk(role="assistant"), + "data: {not valid json}\n\n", + ] + frames = await _drain(serving, items) + assert frames[-1] == _DONE_FRAME + # Find the error-shaped message-end. + error_end = _parse_event(frames[-2]) + assert error_end["type"] == "message-end" + assert error_end["delta"]["finish_reason"] == "ERROR" + assert "error" in error_end["delta"] + + @pytest.mark.asyncio + async def test_citations_in_delta_emit_citation_events(self): + serving = _serving() + # Pass citation dicts the way the cohere2 reasoning parser does. + items = [ + _make_chunk(role="assistant"), + _make_chunk(content="hello"), + _make_chunk( + citations=[ + { + "start": 0, + "end": 5, + "text": "hello", + "sources": [{"type": "document", "id": "d1"}], + } + ] + ), + _make_chunk(finish_reason="stop"), + _make_chunk( + omit_choices=True, + usage={ + "prompt_tokens": 1, + "completion_tokens": 1, + "total_tokens": 2, + }, + ), + ] + frames = await _drain(serving, items) + types = [_parse_event(f)["type"] for f in frames[:-1]] + assert "citation-start" in types + assert "citation-end" in types + + @pytest.mark.asyncio + async def test_first_chunk_emits_message_start_with_chunk_id(self): + serving = _serving() + items = [ + _make_chunk(chunk_id="my-id", role="assistant"), + _make_chunk(chunk_id="my-id", content="hi"), + _make_chunk( + chunk_id="my-id", + omit_choices=True, + usage={ + "prompt_tokens": 1, + "completion_tokens": 1, + "total_tokens": 2, + }, + ), + ] + frames = await _drain(serving, items) + ms = _parse_event(frames[0]) + assert ms["type"] == "message-start" + assert ms["id"] == "my-id" + assert ms["delta"]["message"]["role"] == "assistant" diff --git a/tests/renderers/test_cohere.py b/tests/renderers/test_cohere.py new file mode 100644 index 000000000000..4a7a23d2b737 --- /dev/null +++ b/tests/renderers/test_cohere.py @@ -0,0 +1,743 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for ``vllm/renderers/cohere.py``. + +The tests focus on the pure-Python helpers that produce the render-config +dicts passed to ``cohere_melody.render_cmd3`` / ``render_cmd4``. We also +include a class-level instantiation + async-non-blocking test that +mirrors the analogous ``test_mistral.py`` pattern, exercising the +:class:`CohereRenderer` end-to-end with mocked ``model_config`` / +tokenizer / melody bindings. +""" + +import asyncio +import json +import time +from dataclasses import dataclass +from typing import Any +from unittest.mock import Mock + +import pytest + +from vllm.renderers import ChatParams +from vllm.renderers.cohere import ( + CohereRenderer, + MelodyContentType, + _build_render_config, + _content_blocks, + _conversation_to_melody_messages, + _document_to_melody, + _normalize_tool_call, + _role_to_melody, + _tool_to_melody, +) +from vllm.tokenizers.hf import HfTokenizer + +# ====================================================================== +# _role_to_melody +# ====================================================================== + + +class TestRoleToMelody: + def test_assistant_maps_to_chatbot(self): + # melody's templates use the legacy Cohere ``chatbot`` role name. + assert _role_to_melody("assistant") == "chatbot" + + def test_developer_aliases_to_system(self): + # OpenAI's ``developer`` role is documented as high-priority + # instructions; map it onto the ``system`` slot rather than + # letting the templates drop it on the floor. + assert _role_to_melody("developer") == "system" + + @pytest.mark.parametrize("role", ["user", "system", "tool", "chatbot"]) + def test_recognized_roles_passthrough(self, role): + assert _role_to_melody(role) == role + + @pytest.mark.parametrize( + "role,expected", + [ + ("ASSISTANT", "chatbot"), + ("Developer", "system"), + ("User", "user"), + ("SYSTEM", "system"), + ], + ) + def test_role_normalization_is_case_insensitive(self, role, expected): + # cmd3 / cmd4 templates lowercase the role before matching, so + # accept any casing the caller provides. + assert _role_to_melody(role) == expected + + @pytest.mark.parametrize("role", ["function", "moderator", "", "anything"]) + def test_unknown_roles_raise(self, role): + # Silently dropping unknown roles produces malformed prompts + # (the templates' role chain has no else branch). + with pytest.raises(ValueError, match="Unsupported message role"): + _role_to_melody(role) + + def test_non_string_role_rejected(self): + # The function is typed ``role: str`` and the implementation + # relies on Python's attribute lookup (``role.lower()``) to + # reject non-strings — any exception type is acceptable as long + # as we don't silently produce a malformed prompt. + with pytest.raises((AttributeError, TypeError, ValueError)): + _role_to_melody(None) # type: ignore[arg-type] + + +# ====================================================================== +# _normalize_tool_call +# ====================================================================== + + +class TestNormalizeToolCall: + def test_openai_dict_with_dict_arguments_json_encoded(self): + # melody expects ``parameters`` as a JSON-encoded string even when + # OpenAI delivers an already-parsed dict. + out = _normalize_tool_call( + { + "id": "c1", + "type": "function", + "function": {"name": "f", "arguments": {"a": 1}}, + } + ) + assert out == {"id": "c1", "name": "f", "parameters": '{"a": 1}'} + + def test_openai_dict_with_string_arguments_preserved(self): + out = _normalize_tool_call( + { + "id": "c1", + "type": "function", + "function": {"name": "f", "arguments": '{"a":1}'}, + } + ) + assert out["parameters"] == '{"a":1}' + + def test_flat_dict_without_function_wrapper(self): + out = _normalize_tool_call({"id": "c1", "name": "f", "arguments": '{"k": 1}'}) + # Falls back to top-level ``name`` / ``arguments``. + assert out == {"id": "c1", "name": "f", "parameters": '{"k": 1}'} + + def test_missing_id_becomes_empty_string(self): + out = _normalize_tool_call({"function": {"name": "f", "arguments": "{}"}}) + assert out["id"] == "" + + def test_pydantic_model_dump_supported(self): + class _Fake: + def model_dump(self): + return { + "id": "c1", + "type": "function", + "function": {"name": "f", "arguments": "{}"}, + } + + out = _normalize_tool_call(_Fake()) + assert out == {"id": "c1", "name": "f", "parameters": "{}"} + + def test_invalid_type_rejected(self): + with pytest.raises(TypeError, match="Unexpected tool_call value"): + _normalize_tool_call(42) # type: ignore[arg-type] + + +# ====================================================================== +# _content_blocks +# ====================================================================== + + +class TestContentBlocks: + def test_none_returns_empty_list(self): + assert _content_blocks(None) == [] + + def test_string_wrapped_in_text_block(self): + out = _content_blocks("hi") + assert out == [{"type": MelodyContentType.TEXT, "text": "hi"}] + + def test_string_item_in_list_wrapped(self): + out = _content_blocks(["a", "b"]) + assert out == [ + {"type": MelodyContentType.TEXT, "text": "a"}, + {"type": MelodyContentType.TEXT, "text": "b"}, + ] + + @pytest.mark.parametrize( + "part_type", + ["text", "input_text", "output_text", "refusal"], + ) + def test_text_variants_normalized(self, part_type): + out = _content_blocks([{"type": part_type, "text": "hello"}]) + assert out == [{"type": MelodyContentType.TEXT, "text": "hello"}] + + def test_thinking_block(self): + out = _content_blocks([{"type": "thinking", "thinking": "thoughts"}]) + assert out == [{"type": MelodyContentType.THINKING, "thinking": "thoughts"}] + + def test_image_block_with_default_placeholder(self): + out = _content_blocks([{"type": "image"}]) + assert out == [ + { + "type": MelodyContentType.IMAGE, + "image": {"template_placeholder": ""}, + } + ] + + def test_image_block_custom_placeholder(self): + out = _content_blocks([{"type": "image", "template_placeholder": "[[IMG]]"}]) + assert out[0]["image"]["template_placeholder"] == "[[IMG]]" + + def test_document_block_dict_passthrough(self): + out = _content_blocks( + [{"type": "document", "document": {"data": {"text": "doc"}}}] + ) + assert out == [ + { + "type": MelodyContentType.DOCUMENT, + "document": {"data": {"text": "doc"}}, + } + ] + + def test_document_block_with_non_dict_falls_back_to_json_text(self): + out = _content_blocks([{"type": "document", "document": "raw string doc"}]) + assert out[0]["type"] == MelodyContentType.TEXT + # JSON-encoded for safety since melody expects a structured doc. + assert out[0]["text"] == json.dumps("raw string doc") + + def test_tool_reference_emitted_as_text(self): + out = _content_blocks([{"type": "tool_reference", "name": "calc"}]) + assert out == [{"type": MelodyContentType.TEXT, "text": "calc"}] + + def test_unknown_block_type_fallback_to_text(self): + # Unknown block type with a string value is wrapped in a text block. + out = _content_blocks([{"type": "custom", "custom": "value"}]) + assert out == [{"type": MelodyContentType.TEXT, "text": "value"}] + + def test_unknown_block_type_dict_value_json_encoded(self): + out = _content_blocks([{"type": "custom", "custom": {"k": 1}}]) + assert out == [{"type": MelodyContentType.TEXT, "text": json.dumps({"k": 1})}] + + def test_non_string_non_dict_part_rejected(self): + with pytest.raises(TypeError, match="Unexpected content part"): + _content_blocks([42]) # type: ignore[list-item] + + +# ====================================================================== +# _document_to_melody +# ====================================================================== + + +class TestDocumentToMelody: + def test_string_wrapped_in_text_dict(self): + assert _document_to_melody("hello") == {"text": "hello"} + + def test_pure_dict_passthrough(self): + out = _document_to_melody({"text": "x", "id": "d1"}) + assert out == {"text": "x", "id": "d1"} + # Must be a defensive copy (mutating output should not affect input). + out["new_key"] = "value" + + def test_data_wrapper_flattened(self): + # Cohere v2 documents use ``{id, data: {...}}``; melody expects + # the flat shape with ``id`` merged into the payload. + out = _document_to_melody({"id": "d1", "data": {"text": "hello", "title": "t"}}) + assert out == {"id": "d1", "text": "hello", "title": "t"} + + def test_data_wrapper_preserves_inner_id(self): + # If the inner ``data`` already has an ``id``, it wins. + out = _document_to_melody({"id": "outer", "data": {"id": "inner", "text": "x"}}) + assert out["id"] == "inner" + + def test_invalid_type_rejected(self): + with pytest.raises(TypeError, match="Unsupported document type"): + _document_to_melody(42) # type: ignore[arg-type] + + +# ====================================================================== +# _tool_to_melody +# ====================================================================== + + +class TestToolToMelody: + def test_openai_wrapper(self): + out = _tool_to_melody( + { + "type": "function", + "function": { + "name": "calc", + "description": "calculate", + "parameters": {"type": "object"}, + }, + } + ) + assert out == { + "name": "calc", + "description": "calculate", + "parameters": {"type": "object"}, + } + + def test_flat_dict(self): + out = _tool_to_melody({"name": "calc", "description": "d", "parameters": {}}) + assert out["name"] == "calc" + assert out["parameters"] == {} + + def test_pydantic_like_model_dump(self): + class _Fake: + def model_dump(self): + return { + "type": "function", + "function": { + "name": "calc", + "description": "x", + "parameters": {}, + }, + } + + out = _tool_to_melody(_Fake()) + assert out["name"] == "calc" + + def test_missing_description_becomes_empty(self): + out = _tool_to_melody({"name": "calc"}) + assert out["description"] == "" + assert out["parameters"] == {} + + def test_invalid_type_rejected(self): + with pytest.raises(TypeError, match="Unsupported tool type"): + _tool_to_melody(42) # type: ignore[arg-type] + + +# ====================================================================== +# _conversation_to_melody_messages +# ====================================================================== + + +class TestConversationToMelody: + def test_basic_user_assistant_pair(self): + conv = [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + ] + out = _conversation_to_melody_messages(conv) # type: ignore[arg-type] + assert out == [ + { + "role": "user", + "content": [{"type": MelodyContentType.TEXT, "text": "hi"}], + "tool_calls": [], + }, + { + "role": "chatbot", + "content": [{"type": MelodyContentType.TEXT, "text": "hello"}], + "tool_calls": [], + }, + ] + + def test_assistant_reasoning_prepended_as_thinking_block(self): + # ``reasoning`` (or ``reasoning_content``) is prepended as a + # ``thinking`` block on assistant turns, preserving multi-turn + # chain-of-thought across the rendered prompt. + conv = [ + { + "role": "assistant", + "content": "answer", + "reasoning": "thoughts", + } + ] + out = _conversation_to_melody_messages(conv) # type: ignore[arg-type] + assert out[0]["content"] == [ + {"type": MelodyContentType.THINKING, "thinking": "thoughts"}, + {"type": MelodyContentType.TEXT, "text": "answer"}, + ] + + def test_assistant_reasoning_content_alias_accepted(self): + conv = [ + { + "role": "assistant", + "content": "answer", + "reasoning_content": "thoughts", + } + ] + out = _conversation_to_melody_messages(conv) # type: ignore[arg-type] + assert out[0]["content"][0] == { + "type": MelodyContentType.THINKING, + "thinking": "thoughts", + } + + def test_user_reasoning_ignored(self): + # Only assistant turns get reasoning-as-thinking lifting; user + # turns with a ``reasoning`` key (which shouldn't happen in + # practice) must not produce a phantom thinking block. + conv = [ + { + "role": "user", + "content": "hi", + "reasoning": "should be ignored", + } + ] + out = _conversation_to_melody_messages(conv) # type: ignore[arg-type] + assert out[0]["content"] == [{"type": MelodyContentType.TEXT, "text": "hi"}] + + def test_tool_calls_normalized(self): + conv = [ + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "c1", + "type": "function", + "function": {"name": "f", "arguments": '{"a":1}'}, + } + ], + } + ] + out = _conversation_to_melody_messages(conv) # type: ignore[arg-type] + assert out[0]["tool_calls"] == [ + {"id": "c1", "name": "f", "parameters": '{"a":1}'} + ] + + def test_tool_call_id_preserved_on_tool_role(self): + conv = [ + { + "role": "tool", + "content": "result", + "tool_call_id": "c1", + } + ] + out = _conversation_to_melody_messages(conv) # type: ignore[arg-type] + assert out[0]["tool_call_id"] == "c1" + + +# ====================================================================== +# _build_render_config +# ====================================================================== + + +class TestBuildRenderConfig: + def _conv(self): + return [{"role": "user", "content": "hi"}] + + def test_default_format_is_cmd3(self): + fmt, cfg = _build_render_config(self._conv(), {}) # type: ignore[arg-type] + assert fmt == "cmd3" + assert cfg["use_jinja"] is True + assert isinstance(cfg["messages"], list) + # No additional_template_fields when no extra kwargs are set. + assert "additional_template_fields" not in cfg + + def test_explicit_cmd4(self): + fmt, cfg = _build_render_config(self._conv(), {"cohere_format": "cmd4"}) # type: ignore[arg-type] + assert fmt == "cmd4" + + def test_invalid_format_raises(self): + with pytest.raises(ValueError, match="Invalid cohere_format"): + _build_render_config(self._conv(), {"cohere_format": "cmd5"}) # type: ignore[arg-type] + + def test_documents_converted(self): + _, cfg = _build_render_config( + self._conv(), + { + "documents": [ + "doc text", + {"id": "d1", "data": {"text": "wrapped"}}, + ] + }, + ) # type: ignore[arg-type] + assert cfg["documents"] == [ + {"text": "doc text"}, + {"id": "d1", "text": "wrapped"}, + ] + + def test_available_tools_take_precedence_over_tools(self): + _, cfg = _build_render_config( + self._conv(), + { + "tools": [{"type": "function", "function": {"name": "from_tools"}}], + "available_tools": [ + {"type": "function", "function": {"name": "preferred"}} + ], + }, + ) # type: ignore[arg-type] + names = [t["name"] for t in cfg["available_tools"]] + assert names == ["preferred"] + + def test_tools_used_when_no_available_tools(self): + _, cfg = _build_render_config( + self._conv(), + {"tools": [{"type": "function", "function": {"name": "from_tools"}}]}, + ) # type: ignore[arg-type] + assert [t["name"] for t in cfg["available_tools"]] == ["from_tools"] + + @pytest.mark.parametrize("value", ["enabled", "disabled"]) + def test_reasoning_type_direct(self, value): + _, cfg = _build_render_config(self._conv(), {"reasoning_type": value}) # type: ignore[arg-type] + assert cfg["reasoning_type"] == value + + def test_thinking_dict_shorthand_resolves_reasoning_type(self): + _, cfg = _build_render_config(self._conv(), {"thinking": {"type": "enabled"}}) # type: ignore[arg-type] + assert cfg["reasoning_type"] == "enabled" + + def test_thinking_shorthand_ignores_unknown_type(self): + _, cfg = _build_render_config(self._conv(), {"thinking": {"type": "auto"}}) # type: ignore[arg-type] + assert "reasoning_type" not in cfg + + def test_dev_instruction_forwarded(self): + _, cfg = _build_render_config(self._conv(), {"dev_instruction": "be brief"}) # type: ignore[arg-type] + assert cfg["dev_instruction"] == "be brief" + + def test_response_format_json_object_sets_json_mode(self): + _, cfg = _build_render_config( + self._conv(), {"response_format": {"type": "json_object"}} + ) # type: ignore[arg-type] + assert cfg["json_mode"] is True + assert "json_schema" not in cfg + + def test_response_format_json_schema_sets_json_schema(self): + schema = {"type": "object"} + _, cfg = _build_render_config( + self._conv(), + {"response_format": {"type": "json_schema", "schema": schema}}, + ) # type: ignore[arg-type] + # JSON-encoded for melody (string-only schema field). + assert cfg["json_schema"] == json.dumps(schema) + + def test_response_format_nested_json_schema_unwrapped(self): + # When the SDK shape is ``{type: json_schema, schema: {schema: + # {...}}}``, the inner ``schema`` value is used. + inner = {"type": "object"} + _, cfg = _build_render_config( + self._conv(), + { + "response_format": { + "type": "json_schema", + "schema": {"schema": inner}, + } + }, + ) # type: ignore[arg-type] + assert cfg["json_schema"] == json.dumps(inner) + + def test_json_schema_kwarg_direct(self): + # Caller can also pass ``json_schema`` directly, both as dict and + # as a pre-stringified value. + _, cfg = _build_render_config(self._conv(), {"json_schema": {"a": 1}}) # type: ignore[arg-type] + assert cfg["json_schema"] == '{"a": 1}' + _, cfg = _build_render_config( + self._conv(), {"json_schema": "raw-string-schema"} + ) # type: ignore[arg-type] + assert cfg["json_schema"] == "raw-string-schema" + + def test_json_mode_kwarg_overrides(self): + _, cfg = _build_render_config(self._conv(), {"json_mode": True}) # type: ignore[arg-type] + assert cfg["json_mode"] is True + + def test_cmd3_safety_mode_lowercased(self): + _, cfg = _build_render_config(self._conv(), {"safety_mode": "CONTEXTUAL"}) # type: ignore[arg-type] + assert cfg["safety_mode"] == "contextual" + + def test_cmd3_citation_quality_direct(self): + _, cfg = _build_render_config(self._conv(), {"citation_quality": "ACCURATE"}) # type: ignore[arg-type] + assert cfg["citation_quality"] == "accurate" + + def test_cmd3_citation_quality_derived_from_citation_options(self): + # When ``citation_quality`` is unset, ``citation_options.mode`` is + # collapsed to on/off so cmd3's binary toggle has a value. + _, cfg = _build_render_config( + self._conv(), {"citation_options": {"mode": "accurate"}} + ) # type: ignore[arg-type] + assert cfg["citation_quality"] == "on" + + _, cfg = _build_render_config( + self._conv(), {"citation_options": {"mode": "off"}} + ) # type: ignore[arg-type] + assert cfg["citation_quality"] == "off" + + def test_cmd3_skip_preamble_forwarded(self): + _, cfg = _build_render_config(self._conv(), {"skip_preamble": True}) # type: ignore[arg-type] + assert cfg["skip_preamble"] is True + + def test_cmd3_no_grounding_field(self): + # cmd3 should never emit a cmd4-only ``grounding`` field. + _, cfg = _build_render_config( + self._conv(), + {"cohere_format": "cmd3", "grounding": "fast"}, + ) # type: ignore[arg-type] + assert "grounding" not in cfg + + @pytest.mark.parametrize( + "raw,expected", + [ + ("FAST", "enabled"), + ("ACCURATE", "enabled"), + ("OFF", "disabled"), + ("enabled", "enabled"), + ("disabled", "disabled"), + ("unknown", "unknown"), + ], + ) + def test_cmd4_grounding_direct(self, raw, expected): + # melody's cmd4 only accepts ``unknown``/``enabled``/``disabled``, + # so the renderer normalizes any of the v2-facing values into + # that vocab. + _, cfg = _build_render_config( + self._conv(), + {"cohere_format": "cmd4", "grounding": raw}, + ) # type: ignore[arg-type] + assert cfg["grounding"] == expected + + @pytest.mark.parametrize( + "mode,expected", + [ + ("ACCURATE", "enabled"), + ("FAST", "enabled"), + ("OFF", "disabled"), + ], + ) + def test_cmd4_grounding_from_citation_options_mode(self, mode, expected): + _, cfg = _build_render_config( + self._conv(), + { + "cohere_format": "cmd4", + "citation_options": {"mode": mode}, + }, + ) # type: ignore[arg-type] + assert cfg["grounding"] == expected + + def test_cmd4_grounding_rejects_unknown_value(self): + with pytest.raises(ValueError, match="Unrecognized cmd4 grounding"): + _build_render_config( + self._conv(), + {"cohere_format": "cmd4", "grounding": "foobar"}, + ) # type: ignore[arg-type] + + def test_cmd4_platform_instruction(self): + _, cfg = _build_render_config( + self._conv(), + { + "cohere_format": "cmd4", + "platform_instruction": "do this", + }, + ) # type: ignore[arg-type] + assert cfg["platform_instruction"] == "do this" + + def test_cmd4_no_safety_mode_field(self): + # cmd4 should never carry cmd3-only ``safety_mode``/``citation_quality``. + _, cfg = _build_render_config( + self._conv(), + { + "cohere_format": "cmd4", + "safety_mode": "contextual", + "citation_quality": "on", + }, + ) # type: ignore[arg-type] + assert "safety_mode" not in cfg + assert "citation_quality" not in cfg + + def test_extra_kwargs_become_additional_template_fields(self): + # Anything not in the renderer's consumed-keys set is forwarded + # verbatim under ``additional_template_fields`` so jinja templates + # can resolve ``{{ var }}`` directly. + _, cfg = _build_render_config( + self._conv(), + { + "reasoning_effort": "low", + "my_var": "x", + "documents": ["doc"], # consumed, must NOT leak through + }, + ) # type: ignore[arg-type] + extras = cfg["additional_template_fields"] + assert extras == {"reasoning_effort": "low", "my_var": "x"} + # Sanity: the consumed key still produced its dedicated config slot. + assert cfg["documents"] == [{"text": "doc"}] + + def test_template_id_and_template_jinja_passthrough(self): + _, cfg = _build_render_config( + self._conv(), + { + "template_id": "tpl1", + "template_jinja": "raw {{ jinja }}", + }, + ) # type: ignore[arg-type] + assert cfg["template_id"] == "tpl1" + assert cfg["template_jinja"] == "raw {{ jinja }}" + # use_jinja is always True, regardless of caller input. + assert cfg["use_jinja"] is True + + +# ====================================================================== +# End-to-end async rendering (mirrors ``test_mistral.py``) +# ====================================================================== +# +# Verifies that the synchronous melody bindings run on the renderer's +# thread pool so the asyncio event loop stays responsive under +# concurrent load. Mirrors +# ``test_async_mistral_tokenizer_does_not_block_event_loop`` so future +# regressions in either path are caught uniformly. + + +@dataclass +class _MockHFConfig: + model_type: str = "any" + + +@dataclass +class _MockModelConfig: + runner_type = "generate" + model: str = "cohere-test" + tokenizer: str = "cohere-test" + trust_remote_code: bool = False + max_model_len: int = 100 + tokenizer_revision = None + tokenizer_mode = "cohere" + hf_config = _MockHFConfig() + hf_text_config = _MockHFConfig() + encoder_config: dict[str, Any] | None = None + enable_prompt_embeds: bool = True + skip_tokenizer_init: bool = True + is_encoder_decoder: bool = False + is_multimodal_model: bool = False + renderer_num_workers: int = 1 + + +@dataclass +class _MockParallelConfig: + _api_process_rank: int = 0 + + +@dataclass +class _MockVllmConfig: + model_config: _MockModelConfig + parallel_config: _MockParallelConfig + + +@pytest.mark.asyncio +async def test_async_cohere_renderer_does_not_block_event_loop(): + expected_prompt = "MOCK_RENDERED_PROMPT" + + def slow_render(*_a, **_kw): + time.sleep(2) + return expected_prompt + + mock_tokenizer = Mock(spec=HfTokenizer) + renderer = CohereRenderer( + _MockVllmConfig(_MockModelConfig(), _MockParallelConfig()), + tokenizer=mock_tokenizer, + ) + + # Replace the (already-imported) ``cohere_melody`` bindings with a + # blocking mock. ``_render`` reads ``self._melody`` at call time, so + # this works even though ``_render_async`` was bound at __init__. + fake_melody = Mock() + fake_melody.render_cmd3 = slow_render + fake_melody.render_cmd4 = slow_render + renderer._melody = fake_melody + + task = renderer.render_messages_async([], ChatParams()) + + # Ensure the event loop is not blocked while the (blocking) render + # call is in flight on the thread pool. + blocked_count = 0 + for _ in range(20): # ~2 seconds at 0.1s slices + start = time.perf_counter() + await asyncio.sleep(0) + elapsed = time.perf_counter() - start + if elapsed >= 0.5: + blocked_count += 1 + await asyncio.sleep(0.1) + + _, prompt = await task + assert prompt["prompt"] == expected_prompt, "Mocked blocking render was not called" + assert blocked_count == 0, "Event loop blocked during rendering" diff --git a/vllm/config/model.py b/vllm/config/model.py index 245af557df06..ca1572c57c7d 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -77,7 +77,9 @@ RunnerOption = Literal["auto", RunnerType] ConvertType = Literal["none", "embed", "classify"] ConvertOption = Literal["auto", ConvertType] -TokenizerMode = Literal["auto", "hf", "slow", "mistral", "deepseek_v32", "deepseek_v4"] +TokenizerMode = Literal[ + "auto", "hf", "slow", "mistral", "deepseek_v32", "deepseek_v4", "cohere" +] ModelDType = Literal["auto", "half", "float16", "bfloat16", "float", "float32"] LogprobsMode = Literal[ "raw_logits", "raw_logprobs", "processed_logits", "processed_logprobs" @@ -129,6 +131,9 @@ class ModelConfig: - "mistral" will always use the tokenizer from `mistral_common`. - "deepseek_v32" will always use the tokenizer from `deepseek_v32`. - "deepseek_v4" will always use the tokenizer from `deepseek_v4`. + - "cohere" uses the standard HF tokenizer but renders the chat template + via the `cohere_melody` library (cmd3 / cmd4 templates) instead of + Jinja, and surfaces grounded-citation metadata on responses. - Other custom values can be supported via plugins. To swap the Rust BPE backend that powers HF fast tokenizers for the diff --git a/vllm/entrypoints/cohere/__init__.py b/vllm/entrypoints/cohere/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/vllm/entrypoints/cohere/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/entrypoints/cohere/api_router.py b/vllm/entrypoints/cohere/api_router.py new file mode 100644 index 000000000000..f0b0ed586e9e --- /dev/null +++ b/vllm/entrypoints/cohere/api_router.py @@ -0,0 +1,146 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""FastAPI router for the Cohere Chat v2 API (``POST /cohere/v2/chat``). + +The Cohere v2 protocol models are sourced from the official ``cohere`` +Python SDK (``pip install cohere``). To keep that an *optional* +dependency for vLLM, the SDK-dependent imports - and the route handler +itself - are gated on a one-shot probe at module load. If the SDK isn't +installed, :func:`attach_router` becomes a no-op (with an info log) and +vLLM continues to boot normally. + +Note: the handler must live at module scope (not inside +``attach_router``) so that FastAPI's ``typing.get_type_hints`` resolves +the ``CohereChatV2Request`` body annotation against the module's +globals. Defining it locally inside ``attach_router`` would hide the +type from ``get_type_hints``, causing FastAPI to silently degrade the +body parameter into a query parameter and reject every request with +422. +""" + +from fastapi import FastAPI + +from vllm.logger import init_logger + +logger = init_logger(__name__) + + +try: + import cohere # noqa: F401 -- dependency probe +except ImportError: + _SDK_AVAILABLE = False +else: + _SDK_AVAILABLE = True + + +if _SDK_AVAILABLE: + from http import HTTPStatus + + from fastapi import APIRouter, Depends, Request + from fastapi.responses import JSONResponse, StreamingResponse + + from vllm.entrypoints.cohere.protocol import ( + CohereChatV2Request, + CohereChatV2Response, + CohereError, + ) + from vllm.entrypoints.cohere.serving import CohereServingChatV2 + from vllm.entrypoints.openai.engine.protocol import ErrorResponse + from vllm.entrypoints.serve.utils.api_utils import ( + load_aware_call, + validate_json_request, + with_cancellation, + ) + + router = APIRouter() + + def _serving(request: Request) -> CohereServingChatV2 | None: + return getattr(request.app.state, "cohere_serving_chat_v2", None) + + def _request_id(raw_request: Request | None) -> str | None: + """Best-effort lookup of the active request id. + + Prefers the id the underlying chat handler stamped onto + ``raw_request.state.request_metadata`` (if it got that far before + failing), falling back to the ``X-Request-Id`` HTTP header. May + return ``None`` if neither is available, in which case the field + is omitted from the response. + """ + if raw_request is None: + return None + meta = getattr(raw_request.state, "request_metadata", None) + if meta is not None and getattr(meta, "request_id", None): + return meta.request_id + return raw_request.headers.get("X-Request-Id") + + def _error_response( + error: ErrorResponse, + raw_request: Request | None, + *, + fallback_status: int = HTTPStatus.BAD_REQUEST, + ) -> JSONResponse: + """Translate vLLM's internal error envelope into Cohere's shape.""" + info = error.error + status = info.code or fallback_status + return JSONResponse( + status_code=status, + content=CohereError(message=info.message, id=_request_id(raw_request)).model_dump(exclude_none=True), + ) + + @router.post( + "/cohere/v2/chat", + dependencies=[Depends(validate_json_request)], + responses={ + HTTPStatus.OK.value: {"content": {"text/event-stream": {}}}, + HTTPStatus.BAD_REQUEST.value: {"model": CohereError}, + HTTPStatus.NOT_FOUND.value: {"model": CohereError}, + HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": CohereError}, + }, + ) + @with_cancellation + @load_aware_call + async def chat_v2(request: CohereChatV2Request, raw_request: Request): + handler = _serving(raw_request) + if handler is None: + return JSONResponse( + status_code=HTTPStatus.NOT_IMPLEMENTED.value, + content=CohereError( + message="The model does not support the Cohere v2 chat API.", + id=_request_id(raw_request), + ).model_dump(exclude_none=True), + ) + + try: + result = await handler.create_chat_v2(request, raw_request) + except Exception as e: # noqa: BLE001 - report as 500 for parity + logger.exception("Error in /cohere/v2/chat: %s", e) + return JSONResponse( + status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, + content=CohereError( + message=str(e), + id=_request_id(raw_request), + ).model_dump(exclude_none=True), + ) + + if isinstance(result, ErrorResponse): + return _error_response(result, raw_request) + + if isinstance(result, CohereChatV2Response): + return JSONResponse(content=result.model_dump(exclude_none=True)) + + return StreamingResponse(content=result, media_type="text/event-stream") + + +def attach_router(app: FastAPI) -> None: + """Register ``POST /cohere/v2/chat`` on ``app``. + + No-op (with an info log) when the optional ``cohere`` SDK isn't + installed, since the v2 protocol models live there. + """ + if not _SDK_AVAILABLE: + logger.info( + "cohere SDK not installed; /cohere/v2/chat endpoint disabled. " + "Install with `pip install cohere` to enable it." + ) + return + app.include_router(router) diff --git a/vllm/entrypoints/cohere/protocol.py b/vllm/entrypoints/cohere/protocol.py new file mode 100644 index 000000000000..7409dc1709db --- /dev/null +++ b/vllm/entrypoints/cohere/protocol.py @@ -0,0 +1,363 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Cohere Chat v2 API protocol. + +The bulk of the wire types come straight from the official ``cohere`` +Python SDK so we stay in lockstep with the upstream specification and +avoid hand-mirroring the schema. We only own three things locally: + +1. The top-level request body model (the SDK doesn't ship one — its + ``ClientV2.chat`` takes the body as kwargs), with vLLM-specific + extensions (``kv_transfer_params`` / ``chat_template_kwargs``). +2. The non-streaming response envelope (the SDK exposes the message + shape via :class:`AssistantMessageResponse` but no full response + wrapper). +3. The streaming discriminated union (the SDK exports each event type + individually but not as a combined ``Annotated[Union[...], + discriminator]``). + +Importing this module pulls in the ``cohere`` package. The router that +mounts ``POST /cohere/v2/chat`` guards on that import succeeding so vLLM still +boots without the SDK installed. + +See https://docs.cohere.com/reference/chat for the upstream spec. +""" +from __future__ import annotations + +from typing import Any, Literal + +from cohere import types as _sdk +from cohere.types import ( + AssistantChatMessageV2, + AssistantMessageResponse, + ChatMessageV2, + ChatRequestSafetyMode, + Citation, + CitationOptions, + Document, + ResponseFormatV2, + SystemChatMessageV2, + Thinking, + ToolCallV2, + ToolChatMessageV2, + ToolV2, + UserChatMessageV2, +) +from pydantic import BaseModel, Field, field_validator + +# Re-export the SDK wire-format types alongside our local extensions so +# ``vllm.entrypoints.cohere.serving`` and friends can import everything +# they need from this module. +__all__ = [ + "AssistantChatMessageV2", + "AssistantMessageResponse", + "ChatMessageV2", + "Citation", + "CitationEndEvent", + "CitationOptions", + "CitationStartEvent", + "CohereChatV2Request", + "CohereChatV2Response", + "CohereError", + "CohereFinishReason", + "CohereLogprobItem", + "CohereUsage", + "CohereUsageBilledUnits", + "CohereUsageTokens", + "ContentDeltaEvent", + "ContentEndEvent", + "ContentStartEvent", + "Document", + "MessageEndEvent", + "MessageStartEvent", + "ResponseFormatV2", + "SystemChatMessageV2", + "Thinking", + "ToolCallDeltaEvent", + "ToolCallEndEvent", + "ToolCallStartEvent", + "ToolCallV2", + "ToolChatMessageV2", + "ToolPlanDeltaEvent", + "ToolV2", + "UserChatMessageV2", +] + + +# --------------------------------------------------------------------------- +# Errors +# --------------------------------------------------------------------------- + + +class CohereError(BaseModel): + """Top-level error body returned by ``/cohere/v2/chat`` error responses. + + Cohere's documented error schemas are uniform: ``{message, id}``. + """ + + message: str + id: str | None = None + + +# --------------------------------------------------------------------------- +# Tool choice / finish reasons +# --------------------------------------------------------------------------- +# +# These literals aren't first-class enums in the SDK but are documented at +# https://docs.cohere.com/reference/chat. We declare them here so the +# request/response models can validate them. + +CohereToolChoice = Literal["REQUIRED", "NONE"] + +CohereFinishReason = Literal[ + "COMPLETE", + "STOP_SEQUENCE", + "MAX_TOKENS", + "TOOL_CALL", + "ERROR", + "TIMEOUT", +] + + +# --------------------------------------------------------------------------- +# Request +# --------------------------------------------------------------------------- + + +class CohereChatV2Request(BaseModel): + """Cohere Chat v2 request body. + + Mirrors the schema documented at https://docs.cohere.com/reference/chat. + All structured fields delegate to the official SDK types so the body + schema stays in sync with the upstream spec. + """ + + model: str + messages: list[ChatMessageV2] + stream: bool | None = False + + # Tooling + tools: list[ToolV2] | None = None + strict_tools: bool | None = None + tool_choice: CohereToolChoice | None = None + + # Grounding + documents: list[str | Document] | None = None + citation_options: CitationOptions | None = None + + # Output + response_format: ResponseFormatV2 | None = None + safety_mode: ChatRequestSafetyMode | None = None + max_tokens: int | None = None + stop_sequences: list[str] | None = None + + # Sampling + temperature: float | None = None + seed: int | None = None + frequency_penalty: float | None = None + presence_penalty: float | None = None + k: int | None = None + p: float | None = None + logprobs: bool | None = None + + # Reasoning + thinking: Thinking | None = None + + # Scheduling + priority: int | None = None + + # vLLM-specific extensions (not in Cohere spec). These mirror what the + # Anthropic and OpenAI surfaces already expose so V2 callers can reach + # the same engine knobs when needed. + kv_transfer_params: dict[str, Any] | None = Field( + default=None, + description="KVTransfer parameters used for disaggregated serving.", + ) + chat_template_kwargs: dict[str, Any] | None = Field( + default=None, + description=( + "Additional keyword args to pass to the chat template renderer. " + "Will be accessible by the template." + ), + ) + + @field_validator("model") + @classmethod + def _validate_model(cls, v: str) -> str: + if not v: + raise ValueError("model is required") + return v + + @field_validator("max_tokens") + @classmethod + def _validate_max_tokens(cls, v: int | None) -> int | None: + if v is not None and v < 0: + raise ValueError("max_tokens must be non-negative") + return v + + @field_validator("messages", mode="before") + @classmethod + def _normalize_message_roles(cls, v: Any) -> Any: + """Rewrite OpenAI-style ``developer`` roles to ``system``. + + Cohere's v2 ``ChatMessageV2`` discriminated union only admits + the four literal roles ``user`` / ``assistant`` / ``system`` / + ``tool``. OpenAI's ``developer`` role is documented as + high-priority system instructions, so we alias it onto + ``system`` *before* the SDK discriminator runs (otherwise it + rejects the message with a ``literal_error`` against each union + member). Mirrors ``_role_to_melody`` in the renderer so the + same alias is honoured no matter which surface the message + arrives through. + + ``mode="before"`` is required so the rewrite happens before the + ``list[ChatMessageV2]`` coercion runs the SDK's discriminated + union; a default-mode validator would never see ``developer`` + because validation would have already failed. On any structural + surprise (non-iterable input, items without a dict-shaped + ``role`` field, etc.) we hand ``v`` back unchanged and let + Pydantic's normal coercion surface a precise error. + """ + try: + return [ + {**msg, "role": "system"} + if msg.get("role", "").lower() == "developer" + else msg + for msg in v + ] + except (AttributeError, TypeError): + return v + + @field_validator("messages") + @classmethod + def _validate_messages( + cls, v: list[ChatMessageV2] + ) -> list[ChatMessageV2]: + if not v: + raise ValueError("messages must contain at least one message") + return v + + +# --------------------------------------------------------------------------- +# Usage / Logprobs +# --------------------------------------------------------------------------- +# +# The Cohere SDK only exposes a v1 ``ApiMetaBilledUnits``. The v2 usage +# envelope is documented separately in the OpenAPI spec and we declare it +# here. + + +class CohereUsageBilledUnits(BaseModel): + input_tokens: float | None = None + output_tokens: float | None = None + search_units: float | None = None + classifications: float | None = None + + +class CohereUsageTokens(BaseModel): + input_tokens: float | None = None + output_tokens: float | None = None + + +class CohereUsage(BaseModel): + billed_units: CohereUsageBilledUnits | None = None + tokens: CohereUsageTokens | None = None + cached_tokens: float | None = None + + +class CohereLogprobItem(BaseModel): + text: str | None = None + token_ids: list[int] + logprobs: list[float] | None = None + + +# --------------------------------------------------------------------------- +# Non-streaming response +# --------------------------------------------------------------------------- + + +class CohereChatV2Response(BaseModel): + """Cohere Chat v2 non-streaming response body. + + Wraps the SDK :class:`AssistantMessageResponse` (the message shape) in + the documented v2 response envelope (``id``, ``finish_reason``, + ``usage``, ``logprobs``). The single constructor in + :class:`CohereServingChatV2._chat_completion_to_v2` is responsible for + supplying a non-empty ``id`` (falling back to a synthesized one if + the upstream response is missing it) to this model. + """ + + id: str + finish_reason: CohereFinishReason + message: AssistantMessageResponse + usage: CohereUsage | None = None + logprobs: list[CohereLogprobItem] | None = None + + # vLLM-specific extension. + kv_transfer_params: dict[str, Any] | None = Field( + default=None, description="KVTransfer parameters." + ) + + +# --------------------------------------------------------------------------- +# Streaming events +# --------------------------------------------------------------------------- +# +# Cohere V2 streams a sequence of typed JSON events delivered as Server- +# Sent Events; each event's ``type`` field carries the discriminator. The +# SDK exposes a Pydantic model per event but none of them declare ``type`` +# as a field (the SDK relies on its own deserializer for discrimination), +# so a naive ``model_dump_json()`` would silently drop the discriminator +# and break clients that demux on ``type``. +# +# We therefore subclass each SDK event and bake the wire-format ``type`` +# string in as a ``Literal`` field with a default. ``model_dump()`` now +# emits ``type`` for free, and surfaces can simply construct the event +# class and serialize it -- no manual ``type`` parameter needed. +# +# See https://docs.cohere.com/v2/docs/streaming and the OpenAPI +# ``StreamedChatResponseV2`` schema for the wire-format reference. + + +class MessageStartEvent(_sdk.ChatMessageStartEvent): + type: Literal["message-start"] = "message-start" + + +class ContentStartEvent(_sdk.ChatContentStartEvent): + type: Literal["content-start"] = "content-start" + + +class ContentDeltaEvent(_sdk.ChatContentDeltaEvent): + type: Literal["content-delta"] = "content-delta" + + +class ContentEndEvent(_sdk.ChatContentEndEvent): + type: Literal["content-end"] = "content-end" + + +class ToolPlanDeltaEvent(_sdk.ChatToolPlanDeltaEvent): + type: Literal["tool-plan-delta"] = "tool-plan-delta" + + +class ToolCallStartEvent(_sdk.ChatToolCallStartEvent): + type: Literal["tool-call-start"] = "tool-call-start" + + +class ToolCallDeltaEvent(_sdk.ChatToolCallDeltaEvent): + type: Literal["tool-call-delta"] = "tool-call-delta" + + +class ToolCallEndEvent(_sdk.ChatToolCallEndEvent): + type: Literal["tool-call-end"] = "tool-call-end" + + +class CitationStartEvent(_sdk.CitationStartEvent): + type: Literal["citation-start"] = "citation-start" + + +class CitationEndEvent(_sdk.CitationEndEvent): + type: Literal["citation-end"] = "citation-end" + + +class MessageEndEvent(_sdk.ChatMessageEndEvent): + type: Literal["message-end"] = "message-end" diff --git a/vllm/entrypoints/cohere/serving.py b/vllm/entrypoints/cohere/serving.py new file mode 100644 index 000000000000..42df5c699334 --- /dev/null +++ b/vllm/entrypoints/cohere/serving.py @@ -0,0 +1,1095 @@ +# 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 standard ``ChatMessage.citations`` / + ``DeltaMessage.citations`` fields. +- 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 StrEnum +from typing import TYPE_CHECKING, Any + +from cohere.types import ( + AssistantChatMessageV2, + AssistantMessageResponse, + Citation, + 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.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, +) +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.utils.request_logger import RequestLogger + +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(StrEnum): + """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) -> CohereFinishReason: + return _FINISH_REASON_MAP.get(reason, "COMPLETE") + + +# --------------------------------------------------------------------------- +# 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 are read + natively from the resulting :class:`ChatMessage` / :class:`DeltaMessage`. + """ + + 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) + + if isinstance(generator, ErrorResponse): + return generator + if isinstance(generator, ChatCompletionResponse): + return self._chat_completion_to_v2(generator, request) + 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: + if isinstance(msg, SystemChatMessageV2): + openai_messages.append( + { + "role": "system", + "content": cls._coerce_text_content(msg.content), + } + ) + elif isinstance(msg, UserChatMessageV2): + openai_messages.append(cls._convert_user_message(msg)) + elif isinstance(msg, AssistantChatMessageV2): + openai_messages.append(cls._convert_assistant_message(msg)) + elif isinstance(msg, ToolChatMessageV2): + openai_messages.append(cls._convert_tool_message(msg)) + else: # 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 vanilla 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) if request.strict_tools else False + 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: + # TODO need to add support for this + 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. + + if kwargs: + chat_req.chat_template_kwargs = kwargs + + # ================================================================== + # 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, + ) + + @staticmethod + def _extract_citations_if_any(msg: Any) -> list[Citation] | None: + """Coerce ``ChatMessage.citations`` into the Cohere v2 wire shape. + + ``ChatMessage`` natively carries a ``citations: list[Citation] | + None`` field (see :mod:`vllm.entrypoints.openai.engine.protocol`). + Renderers / parsers (the ``cohere_command3`` and + ``cohere_command4`` reasoning parsers) populate it. We map each :class:`vllm...Citation` into the SDK + :class:`cohere.types.Citation` wire model, preserving sources and + span. + """ + raw = getattr(msg, "citations", None) + if not raw: + return None + out: list[Citation] = [] + for c in raw: + try: + if isinstance(c, Citation): + out.append(c) + continue + if hasattr(c, "model_dump"): + payload = c.model_dump(exclude_none=True) + else: + payload = c + if isinstance(payload, dict): + out.append(Citation.model_validate(payload)) + except Exception: # pragma: no cover - defensive + logger.warning("Skipping malformed citation: %r", c, exc_info=True) + return out or 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. + if getattr(delta, "citations", None): + 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": 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 cohere renderer/parsers populate ``DeltaMessage.citations`` with + :class:`vllm.entrypoints.openai.engine.protocol.Citation` instances + once a citation has fully resolved (start/end indices known). We + emit a complete start+end pair for each citation in the delta so + Cohere clients can attach the annotation to the surrounding text. + """ + events: list[str] = [] + for c in citations: + payload = c.model_dump(exclude_none=True) if hasattr(c, "model_dump") else c + if not isinstance(payload, dict): + continue + try: + citation = Citation.model_validate(payload) + except Exception: # pragma: no cover - defensive + logger.warning( + "Skipping malformed streamed citation: %r", + payload, + exc_info=True, + ) + 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 + # ================================================================== + + def create_error_response(self, message: str) -> ErrorResponse: + # Reuse the OpenAI engine's error response so that the router can + # translate it into Cohere's error envelope uniformly. + from vllm.entrypoints.openai.engine.protocol import ErrorInfo + + return ErrorResponse( + error=ErrorInfo( + message=message, + type="bad_request", + code=400, + ) + ) + + +# --------------------------------------------------------------------------- +# 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.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 diff --git a/vllm/entrypoints/generate/api_router.py b/vllm/entrypoints/generate/api_router.py index 38ecdec5ce28..a9c36f66089a 100644 --- a/vllm/entrypoints/generate/api_router.py +++ b/vllm/entrypoints/generate/api_router.py @@ -41,6 +41,12 @@ def register_generate_api_routers(app: FastAPI): register_anthropic_api_router(app) + from vllm.entrypoints.cohere.api_router import ( + attach_router as register_cohere_api_router, + ) + + register_cohere_api_router(app) + from .generative_scoring.api_router import register_generative_scoring_api_router register_generative_scoring_api_router(app) @@ -55,6 +61,16 @@ async def init_generate_state( ): from vllm.entrypoints.anthropic.serving import AnthropicServingMessages from vllm.entrypoints.chat_utils import load_chat_template + + try: + from vllm.entrypoints.cohere.serving import CohereServingChatV2 + except ImportError: + # The Cohere serving handler depends on the optional `cohere` SDK + # for its wire-format protocol models. When it isn't installed, + # `register_cohere_api_router` already skips registering the route, + # so we simply leave `cohere_serving_chat_v2` unset. + CohereServingChatV2 = None # type: ignore[assignment,misc] + from vllm.entrypoints.mcp.tool_server import ( DemoToolServer, MCPToolServer, @@ -87,6 +103,19 @@ async def init_generate_state( tool_server = None resolved_chat_template = load_chat_template(args.chat_template) + # Fold the dedicated ``--cohere-format`` CLI flag into the renderer's + # default chat-template kwargs. The cohere renderer reads + # ``chat_template_kwargs["cohere_format"]`` to pick cmd3 vs cmd4 + # rendering; making this a first-class flag keeps the right format + # discoverable for ``vllm serve --tokenizer-mode cohere`` users + # without forcing them to hand-construct a JSON dict for + # ``--default-chat-template-kwargs``. Per-request overrides still + # take precedence (see ``merge_kwargs`` in + # ``ChatCompletionRequest.build_chat_params``). + default_chat_template_kwargs = dict(args.default_chat_template_kwargs or {}) + if getattr(args, "cohere_format", None): + default_chat_template_kwargs.setdefault("cohere_format", args.cohere_format) + # Render endpoints are always backed by OnlineRenderer so that # /v1/chat/completions/render and /v1/completions/render work on both # generate-mode and render-only servers. Created in init_app_state. @@ -107,7 +136,7 @@ async def init_generate_state( enable_prompt_tokens_details=args.enable_prompt_tokens_details, enable_force_include_usage=args.enable_force_include_usage, enable_log_outputs=args.enable_log_outputs, - default_chat_template_kwargs=args.default_chat_template_kwargs, + default_chat_template_kwargs=default_chat_template_kwargs, ) if "generate" in supported_tasks else None @@ -120,7 +149,7 @@ async def init_generate_state( request_logger=request_logger, chat_template=resolved_chat_template, chat_template_content_format=args.chat_template_content_format, - default_chat_template_kwargs=args.default_chat_template_kwargs, + default_chat_template_kwargs=default_chat_template_kwargs, trust_request_chat_template=args.trust_request_chat_template, return_tokens_as_token_ids=args.return_tokens_as_token_ids, enable_auto_tools=args.enable_auto_tool_choice, @@ -170,11 +199,32 @@ async def init_generate_state( reasoning_parser=args.structured_outputs_config.reasoning_parser, enable_prompt_tokens_details=args.enable_prompt_tokens_details, enable_force_include_usage=args.enable_force_include_usage, - default_chat_template_kwargs=args.default_chat_template_kwargs, + default_chat_template_kwargs=default_chat_template_kwargs, ) if "generate" in supported_tasks else None ) + state.cohere_serving_chat_v2 = ( + CohereServingChatV2( + engine_client, + state.openai_serving_models, + args.response_role, + online_renderer=state.online_renderer, + request_logger=request_logger, + chat_template=resolved_chat_template, + chat_template_content_format=args.chat_template_content_format, + return_tokens_as_token_ids=args.return_tokens_as_token_ids, + enable_auto_tools=args.enable_auto_tool_choice, + tool_parser=args.tool_call_parser, + reasoning_parser=args.structured_outputs_config.reasoning_parser, + enable_prompt_tokens_details=args.enable_prompt_tokens_details, + enable_force_include_usage=args.enable_force_include_usage, + default_chat_template_kwargs=default_chat_template_kwargs, + is_reasoning_model=args.cohere_is_reasoning_model, + ) + if CohereServingChatV2 is not None and "generate" in supported_tasks + else None + ) state.serving_tokens = ( ServingTokens( engine_client, diff --git a/vllm/entrypoints/openai/chat_completion/protocol.py b/vllm/entrypoints/openai/chat_completion/protocol.py index 09ce8bf8dabe..41005065de74 100644 --- a/vllm/entrypoints/openai/chat_completion/protocol.py +++ b/vllm/entrypoints/openai/chat_completion/protocol.py @@ -21,6 +21,7 @@ ) from vllm.entrypoints.openai.engine.protocol import ( AnyResponseFormat, + Citation, DeltaMessage, FunctionCall, FunctionDefinition, @@ -65,6 +66,11 @@ class ChatMessage(OpenAIBaseModel): # vLLM-specific fields that are not in OpenAI spec reasoning: str | None = None + # Citations grounding the message content in source material. Populated + # by parsers/renderers for grounded models (e.g. Cohere Command). Left + # unset (and therefore omitted from JSON) for ungrounded models so + # OpenAI-compatible clients see the standard shape. + citations: list[Citation] | None = None @model_serializer(mode="wrap") def _serialize(self, handler): diff --git a/vllm/entrypoints/openai/chat_completion/serving.py b/vllm/entrypoints/openai/chat_completion/serving.py index 284b2511dbaa..f613f3e1ac28 100644 --- a/vllm/entrypoints/openai/chat_completion/serving.py +++ b/vllm/entrypoints/openai/chat_completion/serving.py @@ -863,10 +863,23 @@ async def chat_completion_full_generator( ) if not request.include_reasoning: reasoning = None + # Reasoning parsers that extract grounding citations + # (Cohere Command family) cache them on the parser + # instance after ``parse``. We surface them here so + # grounded surfaces (e.g. /cohere/v2/chat) can attach + # them to the response message. Non-citation parsers + # leave the attribute absent and ``citations`` stays + # ``None`` so it round-trips out of the OpenAI envelope. + citations = getattr( + getattr(parser, "reasoning_parser", None), + "last_unary_citations", + None, + ) else: reasoning = None content = output.text tool_calls = [] + citations = None auto_tools_called = False is_named_tool_choice = ( @@ -930,6 +943,10 @@ async def chat_completion_full_generator( "completion." ) message = ChatMessage(role=role, reasoning=reasoning, content=content) + + if citations: + message.citations = citations + # In OpenAI's API, when a tool is called, the finish_reason is: # "tool_calls" for "auto" or "required" tool calls, # and "stop" for named tool calls. diff --git a/vllm/entrypoints/openai/cli_args.py b/vllm/entrypoints/openai/cli_args.py index 1533895edcdd..b1274ebfb7ae 100644 --- a/vllm/entrypoints/openai/cli_args.py +++ b/vllm/entrypoints/openai/cli_args.py @@ -149,6 +149,25 @@ class BaseFrontendArgs: """If set to False, output deltas will not be logged. Relevant only if --enable-log-outputs is set. """ + cohere_is_reasoning_model: bool = True + """Cohere ``/cohere/v2/chat`` only. Whether the served model is a + reasoning Command-family model. When True (default), the assistant's + chain-of-thought is surfaced as a ``thinking`` content block (or + ``content-*`` events on the stream). When False, reasoning is + surfaced as Cohere's ``tool_plan`` field (or ``tool-plan-delta`` + events) whenever the model emits tool calls, matching older non- + reasoning Command models. Has no effect on the non-Cohere + endpoints.""" + cohere_format: str = "cmd4" + """Cohere ``--tokenizer-mode cohere`` only. Which Cohere prompt + format to render: ``cmd4`` (current Command A / R+ 2025+ models; + default) or ``cmd3`` (legacy Command R / R+ 2024 models). Selecting + the wrong format silently produces a prompt the model wasn't trained + on, which most commonly manifests as the model emitting text but no + citations / tool calls / thinking blocks. Equivalent to passing + ``--default-chat-template-kwargs '{"cohere_format": "..."}'`` -- any + explicit request-level ``chat_template_kwargs.cohere_format`` still + wins.""" log_error_stack: bool = envs.VLLM_SERVER_DEV_MODE """If set to True, log the stack trace of error responses""" tokens_only: bool = False diff --git a/vllm/entrypoints/openai/engine/protocol.py b/vllm/entrypoints/openai/engine/protocol.py index 084d8d429a67..07a792a13aed 100644 --- a/vllm/entrypoints/openai/engine/protocol.py +++ b/vllm/entrypoints/openai/engine/protocol.py @@ -347,11 +347,58 @@ class ExtractedToolCallInformation(BaseModel): content: str | None = None +class CitationSource(OpenAIBaseModel): + """Source attribution for a :class:`Citation`. + + Mirrors the shape used by Cohere's Chat v2 API. ``type`` is the source + discriminator (``document`` or ``tool``); ``id`` is the citing + document/tool-output identifier; ``document`` and ``tool_output`` carry + the original payload that produced the citation. + """ + + type: Literal["document", "tool"] + id: str | None = None + document: dict[str, Any] | None = None + tool_output: dict[str, Any] | None = None + + +class Citation(OpenAIBaseModel): + """A citation grounding a span of generated text in source material. + + This is a vLLM-specific extension to the OpenAI chat completion shape + so that grounded models (e.g. Cohere Command-family) can surface + citation metadata through the standard chat completion responses. It is + consumed by surfaces that expose citations (such as Cohere's + ``/cohere/v2/chat`` endpoint). + """ + + start: int | None = None + """Start character offset in the surrounding text content.""" + end: int | None = None + """End character offset (exclusive) in the surrounding text content.""" + text: str | None = None + """The cited text snippet.""" + sources: list[CitationSource] = Field(default_factory=list) + """Source documents / tool outputs that ground this citation.""" + content_index: int | None = None + """Index of the content block this citation refers to (when the message + has multiple content blocks).""" + type: Literal["TEXT_CONTENT", "THINKING_CONTENT", "PLAN"] | None = None + """Which kind of content block this citation grounds: the user-visible + text (``TEXT_CONTENT``), a thinking block (``THINKING_CONTENT``), or a + tool-plan block (``PLAN``). ``None`` means unspecified.""" + + class DeltaMessage(OpenAIBaseModel): role: str | None = None content: str | None = None reasoning: str | None = None tool_calls: list[DeltaToolCall] = Field(default_factory=list) + # vLLM-specific extension: citations grounding the streamed content. + # Surfaces (e.g. Cohere Chat v2) consume this when present; OpenAI + # clients ignore it. Default is ``None`` so it is omitted from the wire + # for non-grounded models. + citations: list[Citation] | None = None @model_serializer(mode="wrap") def _serialize(self, handler): diff --git a/vllm/reasoning/cohere_command_reasoning_parser.py b/vllm/reasoning/cohere_command_reasoning_parser.py index 34066ef2d922..2dea67b3c9b2 100644 --- a/vllm/reasoning/cohere_command_reasoning_parser.py +++ b/vllm/reasoning/cohere_command_reasoning_parser.py @@ -25,6 +25,8 @@ ) from vllm.entrypoints.openai.engine.protocol import ( AnyResponseFormat, + Citation, + CitationSource, DeltaFunctionCall, DeltaMessage, DeltaToolCall, @@ -404,6 +406,59 @@ def _schema_dict_from_structured_outputs( ) +def _melody_sources_to_vllm(raw_sources: Any) -> list[CitationSource]: + """Convert melody's ``Source`` objects into :class:`CitationSource`. + + melody's ``Source`` shape is ``{tool_call_index, tool_result_indices, + document_ids}``. ``document_ids`` may not be set; if it is empty and + we have no resolvable identifier then we fall back to a generic + ``tool``-style source carrying the tool-call index for visibility. + """ + out: list[CitationSource] = [] + for s in raw_sources or []: + # TODO Verify the tool vs doc logic + doc_ids: list[str] = list(getattr(s, "document_ids", None) or []) + if doc_ids: + for did in doc_ids: + if did: + out.append(CitationSource(type="document", id=did)) + continue + tool_call_index = getattr(s, "tool_call_index", None) + out.append( + CitationSource( + type="tool", + id=( + str(tool_call_index) + if tool_call_index is not None + else None + ), + ) + ) + return out + + +def _melody_citations_to_vllm(raw_citations: Any) -> list[Citation] | None: + """Convert melody's ``FilterCitation`` objects into :class:`Citation`.""" + if not raw_citations: + return None + out: list[Citation] = [] + for c in raw_citations: + out.append( + Citation( + start=getattr(c, "start_index", None), + end=getattr(c, "end_index", None), + text=getattr(c, "text", None), + sources=_melody_sources_to_vllm(getattr(c, "sources", None)), + type=( + "THINKING_CONTENT" + if getattr(c, "is_thinking", False) + else "TEXT_CONTENT" + ), + ) + ) + return out + + class BaseCohereCommandReasoningParser(ReasoningParser): def __init__( self, @@ -418,6 +473,13 @@ def __init__( self.unary_opts = unary_opts self.melody_unary = PyFilter(unary_opts) self.melody_streaming = PyFilter(streaming_opts) + # Citations extracted by the most recent ``extract_reasoning`` call. + # The non-streaming chat-completion path reads this back from the + # parser instance (which is constructed per-request) and attaches + # the result to ``ChatMessage.citations`` so grounded surfaces + # (e.g. ``/cohere/v2/chat``) can surface them. ``None`` when the + # last parse produced no citations. + self.last_unary_citations: list[Citation] | None = None @property def reasoning_start_str(self) -> str | None: @@ -437,7 +499,13 @@ def extract_reasoning_streaming( delta_token_ids: Sequence[int], ) -> DeltaMessage | None: r = self.melody_streaming.write_decoded(delta_text) - if r.content is None and r.reasoning is None and not r.tool_calls: + citations = _melody_citations_to_vllm(getattr(r, "citations", None)) + if ( + r.content is None + and r.reasoning is None + and not r.tool_calls + and not citations + ): return None msg = DeltaMessage() if r.content is not None: @@ -454,12 +522,22 @@ def extract_reasoning_streaming( ) for tc in r.tool_calls ] + if citations: + msg.citations = citations return msg def extract_reasoning( self, model_output: str, request: ChatCompletionRequest | ResponsesRequest ) -> tuple[str | None, str | None]: result = self.melody_unary.process_full_text(model_output) + # Cache citations so the non-streaming chat-completion path can + # surface them on ``ChatMessage.citations`` (the ``parse`` return + # tuple is locked to ``(reasoning, content, tool_calls)`` across + # all parsers, so we ferry citations via parser-instance state -- + # safe because the parser is constructed per-request). + self.last_unary_citations = _melody_citations_to_vllm( + getattr(result, "citations", None) + ) return result.reasoning, result.content def extract_content_ids(self, input_ids: list[int]) -> list[int]: @@ -533,12 +611,25 @@ def adjust_request( return request +# melody's streaming filter only buffers a partial ```` citation +# across ``write_decoded`` calls when ``stream_non_grounded_answer`` is +# set: otherwise, the moment an opening ```` in the same delta, the filter emits the partial +# marker bytes verbatim as plain content. In vLLM's streaming path the +# parser is fed one token (1-4 chars) per call, so an unbuffered filter +# will leak ````-style markers into ``delta.content`` and never +# emit a ``FilterCitation`` for them. Enabling the flag flips the +# partial-match branch in melody's ``parse_citations`` (see +# ``src/parsing/citations_filter.rs``) to ``return (None, 0)`` -- i.e. +# keep buffering -- which lets a full citation eventually resolve. +# Non-streaming (unary) parsing receives the whole output in one call so +# the flag is a no-op there and we leave ``unary_opts`` alone. class CohereCommand3ReasoningParser(BaseCohereCommandReasoningParser): def __init__(self, tokenizer: TokenizerLike, *args, **kwargs): super().__init__( tokenizer, *args, - streaming_opts=PyFilterOptions().cmd3(), + streaming_opts=PyFilterOptions().cmd3().stream_non_grounded_answer(), unary_opts=PyFilterOptions().cmd3().no_tools(), **kwargs, ) @@ -549,7 +640,7 @@ def __init__(self, tokenizer: TokenizerLike, *args, **kwargs): super().__init__( tokenizer, *args, - streaming_opts=PyFilterOptions().cmd4(), + streaming_opts=PyFilterOptions().cmd4().stream_non_grounded_answer(), unary_opts=PyFilterOptions().cmd4().no_tools(), **kwargs, ) diff --git a/vllm/renderers/cohere.py b/vllm/renderers/cohere.py new file mode 100644 index 000000000000..c6fc5b29608c --- /dev/null +++ b/vllm/renderers/cohere.py @@ -0,0 +1,571 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Cohere prompt renderer. + +Templates the Cohere Command-family prompt formats (cmd3 / cmd4) using the +``cohere_melody`` Rust bindings instead of Jinja. Selecting this renderer is +a matter of setting ``--tokenizer-mode cohere`` on the engine; tokenization +itself still flows through the cached HuggingFace tokenizer. + +This renderer intentionally accepts the same ``chat_template_kwargs`` shape +used by the standard chat completions endpoint, plus a few Cohere-specific +fields that grounding/citation features require: + +* ``documents``: list of document dicts to expose to the model +* ``available_tools``: list of tool dicts (overrides ``tools``) +* ``safety_mode``: cmd3-only safety mode (``contextual`` / ``strict`` / ``none``) +* ``citation_quality``: cmd3 citation toggle (``on`` / ``off``) +* ``citation_options.mode``: cmd4 grounding (``fast`` / ``accurate`` / ``off``) +* ``reasoning_type``: ``enabled`` / ``disabled`` +* ``response_prefix``: optional response prefix +* ``json_schema`` / ``json_mode`` / ``response_format``: structured outputs +* ``template_id`` / ``template`` / ``template_jinja`` / ``use_jinja``: + template overrides +* ``cohere_format``: ``cmd3`` (default) or ``cmd4`` + +Any *other* keys in ``chat_template_kwargs`` are forwarded verbatim to +the melody render config as ``additional_template_fields`` -- i.e. they +become Jinja variables accessible inside the template. This mirrors +vLLM's documented contract for ``chat_template_kwargs`` ("kwargs +accessible by the template"), so e.g. +``chat_template_kwargs={"reasoning_effort": "low"}`` resolves +``{{ reasoning_effort }}`` inside cmd3 / cmd4 templates. + +Citations produced by Cohere models are surfaced through the standard +``ChatMessage.citations`` / ``DeltaMessage.citations`` fields (populated by +the ``cohere2`` reasoning parser). +""" +from __future__ import annotations + +import copy +import json +from enum import StrEnum +from typing import TYPE_CHECKING, Any + +from vllm.config import VllmConfig +from vllm.entrypoints.chat_utils import ( + ChatCompletionMessageParam, + ConversationMessage, + parse_chat_messages, + parse_chat_messages_async, +) +from vllm.logger import init_logger +from vllm.tokenizers.hf import HfTokenizer +from vllm.utils.async_utils import make_async + +from .base import BaseRenderer +from .inputs import DictPrompt +from .inputs.preprocess import parse_dec_only_prompt +from .params import ChatParams + +if TYPE_CHECKING: + pass + +logger = init_logger(__name__) + + +_DEFAULT_FORMAT = "cmd3" +_VALID_FORMATS = ("cmd3", "cmd4") + + +class MelodyContentType(StrEnum): + """Wire-format discriminator for melody content blocks. + + These strings are what the cmd3 / cmd4 Jinja templates check against + in ``message.content[0].type`` (e.g. the ``thinking`` branch in + ``cmd4-v1.jinja``). Keep new values in sync with melody's template + schema. + """ + + TEXT = "text" + THINKING = "thinking" + IMAGE = "image" + DOCUMENT = "document" + +# Keys this renderer interprets directly from ``chat_template_kwargs`` and +# maps onto typed melody render-config fields. Everything *not* in this +# set is forwarded verbatim to melody as ``additional_template_fields`` +# (i.e. as Jinja template variables), so callers can write +# ``chat_template_kwargs = {"my_var": "..."}`` and have ``{{ my_var }}`` +# resolve inside the template -- matching vLLM's documented contract for +# ``chat_template_kwargs`` and avoiding the older nested-namespace form +# (``chat_template_kwargs.additional_template_fields.my_var``). +_RENDERER_CONSUMED_KEYS = frozenset( + { + "cohere_format", + "template_id", + "template_jinja", + "use_jinja", + "documents", + "available_tools", + "tools", + "reasoning_type", + "thinking", + "dev_instruction", + "response_format", + "json_schema", + "json_mode", + "safety_mode", + "citation_quality", + "citation_options", + "skip_preamble", + "grounding", + "platform_instruction", + } +) + + +def _try_import_melody(): + try: + import cohere_melody # type: ignore + + return cohere_melody + except ImportError as e: # pragma: no cover - exercised at runtime + raise ImportError( + "The `cohere` tokenizer/renderer mode requires the " + "`cohere_melody` package. Install it via " + "`pip install cohere-melody` or build from " + "https://github.com/cohere-ai/melody." + ) from e + + +_MELODY_ROLES = frozenset({"system", "user", "chatbot", "tool"}) + + +# Cohere v2 ``citation_options.mode`` -> melody cmd4 ``grounding``. +# v2 surfaces three modes (``FAST`` / ``ACCURATE`` / ``OFF``) but cmd4's +# template doesn't differentiate fast vs accurate at the prompt layer -- +# both just turn grounding on. ``unknown`` / ``enabled`` / ``disabled`` +# are the values melody actually accepts. +_CMD4_GROUNDING_FROM_MODE = { + "fast": "enabled", + "accurate": "enabled", + "on": "enabled", + "enabled": "enabled", + "off": "disabled", + "disabled": "disabled", + "unknown": "unknown", +} + + +def _normalize_cmd4_grounding(value: Any) -> str: + """Coerce a user-facing grounding/citation mode to melody's vocab. + + Raises ``ValueError`` rather than letting an unrecognized value slip + through to ``render_cmd4`` and surface as a generic + ``Invalid config: grounding`` from melody. + """ + out = _CMD4_GROUNDING_FROM_MODE.get(value.lower()) + if out is None: + raise ValueError( + f"Unrecognized cmd4 grounding value: {value!r}. Expected one of " + f"FAST, ACCURATE, OFF (citation_options.mode), or " + f"enabled / disabled / unknown." + ) + return out + + +def _role_to_melody(role: str) -> str: + """Map an OpenAI role to the role string melody expects. + + The cmd3 / cmd4 jinja templates only recognize ``system``, ``user``, + ``assistant``/``chatbot``, and ``tool``: any other role is silently + dropped by the template's role-dispatch chain (no fallback branch), + which produces a malformed prompt without any error. We therefore + refuse unknown roles up front rather than letting them disappear. + + Aliases: + + * ``assistant`` -> ``chatbot`` (Cohere's historical assistant role + name used by the templates). + * ``developer`` -> ``system`` (OpenAI's ``developer`` role is + documented as high-priority instructions, which maps onto the + ``system`` slot in Cohere's prompt format). + """ + role = role.lower() + if role == "assistant": + return "chatbot" + if role == "developer": + return "system" + if role in _MELODY_ROLES: + return role + raise ValueError( + f"Unsupported message role for the cohere renderer: {role!r}. " + "Expected one of: system, developer, user, assistant, chatbot, tool." + ) + + +def _normalize_tool_call(tc: dict[str, Any] | Any) -> dict[str, Any]: + """Normalize an OpenAI tool call dict into melody's tool call shape. + + melody expects ``{id, name, parameters: }`` whereas OpenAI + delivers ``{id, type, function: {name, arguments: }}``. + """ + # Pydantic objects + if hasattr(tc, "model_dump"): + tc = tc.model_dump() + if not isinstance(tc, dict): + raise TypeError(f"Unexpected tool_call value: {tc!r}") + + fn = tc.get("function") or {} + name = fn.get("name") or tc.get("name") or "" + args = fn.get("arguments") + if args is None: + args = tc.get("arguments", {}) + # melody expects a JSON-encoded string + if not isinstance(args, str): + args = json.dumps(args, ensure_ascii=False) + return { + "id": tc.get("id") or "", + "name": name, + "parameters": args, + } + + +def _content_blocks(content: Any) -> list[dict[str, Any]]: + """Convert OpenAI ``ConversationMessage.content`` into melody content blocks. + + The chat_utils ``content_format="openai"`` produces a list of + ``{type: text, text: ...}`` dicts already; we pass those through with + minimal coercion. Plain string content is wrapped in a single text block. + Image / multimodal placeholder dicts (``{"type": "image"}``) are + forwarded as-is so that templates can reference image placeholders that + the upstream tokenizer expands separately. + """ + if content is None: + return [] + if isinstance(content, str): + return [{"type": MelodyContentType.TEXT, "text": content}] + blocks: list[dict[str, Any]] = [] + for part in content: + if isinstance(part, str): + blocks.append({"type": MelodyContentType.TEXT, "text": part}) + continue + if not isinstance(part, dict): + raise TypeError(f"Unexpected content part: {part!r}") + part_type = part.get("type", MelodyContentType.TEXT) + if part_type in ("text", "input_text", "output_text", "refusal"): + blocks.append( + {"type": MelodyContentType.TEXT, "text": part.get("text", "")} + ) + elif part_type == MelodyContentType.THINKING: + blocks.append( + { + "type": MelodyContentType.THINKING, + "thinking": part.get("thinking", ""), + } + ) + elif part_type == MelodyContentType.IMAGE: + blocks.append( + { + "type": MelodyContentType.IMAGE, + "image": { + "template_placeholder": part.get( + "template_placeholder", "" + ), + }, + } + ) + elif part_type == MelodyContentType.DOCUMENT: + doc = part.get("document") + if isinstance(doc, dict): + blocks.append( + {"type": MelodyContentType.DOCUMENT, "document": doc} + ) + else: + # Fall back to wrapping arbitrary string as text. + blocks.append( + {"type": MelodyContentType.TEXT, "text": json.dumps(doc)} + ) + elif part_type == "tool_reference": + # Tool references are rendered by name; emit as text so the + # renderer downstream is content-format agnostic. + blocks.append( + { + "type": MelodyContentType.TEXT, + "text": part.get("name") or part.get("text", ""), + } + ) + else: + # Unknown block type: render as text fallback. + text = part.get("text") or part.get(part_type) or "" + text_str = text if isinstance(text, str) else json.dumps(text) + blocks.append({"type": MelodyContentType.TEXT, "text": text_str}) + return blocks + + +def _document_to_melody(doc: Any) -> dict[str, Any]: + """Coerce a Cohere v2 document into melody's ``Document`` (dict) shape.""" + if isinstance(doc, str): + return {"text": doc} + if isinstance(doc, dict): + # Cohere v2 wraps documents in {id, data: {...}} or pure dicts. + if "data" in doc and isinstance(doc["data"], dict): + payload = dict(doc["data"]) + if "id" in doc and "id" not in payload: + payload["id"] = doc["id"] + return payload + return dict(doc) + raise TypeError(f"Unsupported document type: {type(doc).__name__}") + + +def _tool_to_melody(tool: Any) -> dict[str, Any]: + """Coerce a chat completions tool definition into melody's ``Tool`` shape. + + Accepts either the raw OpenAI tool wrapper ``{type:"function", function: + {name, description, parameters}}`` or a flat ``{name, description, + parameters}`` dict (which is what melody itself expects). + """ + if hasattr(tool, "model_dump"): + tool = tool.model_dump() + if not isinstance(tool, dict): + raise TypeError(f"Unsupported tool type: {type(tool).__name__}") + if "function" in tool and isinstance(tool["function"], dict): + fn = tool["function"] + else: + fn = tool + return { + "name": fn.get("name", ""), + "description": fn.get("description", "") or "", + "parameters": fn.get("parameters") or {}, + } + + +def _conversation_to_melody_messages( + conversation: list[ConversationMessage], +) -> list[dict[str, Any]]: + out: list[dict[str, Any]] = [] + for msg in conversation: + role = _role_to_melody(msg.get("role", "user")) + content_blocks = _content_blocks(msg.get("content")) + + # Treat reasoning content as a thinking block on assistant turns + # so multi-turn reasoning is preserved in the rendered prompt. + # + # Cohere models output thought as either a ``thinking`` + # block (reasoning models) or a ``tool_plan`` field (older non- + # reasoning Command models), but vLLM's ConversationMessage has a + # single unified ``reasoning`` field that drops that difference. + # The cmd3 / cmd4 jinja templates render ``thinking`` blocks in both + # the case of tool calls and regular thinking blocks, so this is fine + # from the renderer's perspective. + reasoning = msg.get("reasoning") or msg.get("reasoning_content") + if role == "chatbot" and reasoning: + content_blocks.insert( + 0, + {"type": MelodyContentType.THINKING, "thinking": reasoning}, + ) + + tool_calls = [ + _normalize_tool_call(tc) for tc in (msg.get("tool_calls") or []) + ] + + out_msg: dict[str, Any] = { + "role": role, + "content": content_blocks, + "tool_calls": tool_calls, + } + tool_call_id = msg.get("tool_call_id") + if tool_call_id: + out_msg["tool_call_id"] = tool_call_id + out.append(out_msg) + return out + + +def _build_render_config( + conversation: list[ConversationMessage], + chat_template_kwargs: dict[str, Any], +) -> tuple[str, dict[str, Any]]: + """Build the ``render_cmd3`` / ``render_cmd4`` config dict. + + Returns ``(format, config_dict)`` where ``format`` is either ``"cmd3"`` + or ``"cmd4"``. + """ + fmt = chat_template_kwargs.get("cohere_format", _DEFAULT_FORMAT) + if fmt not in _VALID_FORMATS: + raise ValueError( + f"Invalid cohere_format={fmt!r}; expected one of {_VALID_FORMATS}" + ) + + config: dict[str, Any] = { + "messages": _conversation_to_melody_messages(conversation), + } + + # Optional template overrides + for k in ("template_id", "template_jinja"): + if v := chat_template_kwargs.get(k): + config[k] = v + # Only support Jinja with vllm + config["use_jinja"] = True + + # Documents + documents = chat_template_kwargs.get("documents") or [] + if documents: + config["documents"] = [_document_to_melody(d) for d in documents] + + # Tools - prefer explicit ``available_tools``, fall back to OpenAI ``tools`` + tools = ( + chat_template_kwargs.get("available_tools") + or chat_template_kwargs.get("tools") + or [] + ) + if tools: + config["available_tools"] = [_tool_to_melody(t) for t in tools] + + # Reasoning toggle (cmd3 + cmd4) + if (rt := chat_template_kwargs.get("reasoning_type")) is not None: + config["reasoning_type"] = str(rt) + elif "thinking" in chat_template_kwargs: + # Cohere v2 ``thinking: {type: enabled|disabled}`` shorthand + thinking = chat_template_kwargs["thinking"] + t = thinking.get("type") if isinstance(thinking, dict) else thinking + if t in ("enabled", "disabled"): + config["reasoning_type"] = t + + if (di := chat_template_kwargs.get("dev_instruction")) is not None: + config["dev_instruction"] = str(di) + + # JSON / structured outputs + if (rf := chat_template_kwargs.get("response_format")) is not None: + rf = rf.model_dump() if hasattr(rf, "model_dump") else dict(rf) + rf_type = rf.get("type") + if rf_type == "json_object": + config["json_mode"] = True + elif rf_type in ("json_schema", "json"): + schema = rf.get("schema") or rf.get("json_schema") + if isinstance(schema, dict) and "schema" in schema: + schema = schema["schema"] + if schema is not None: + config["json_schema"] = ( + schema if isinstance(schema, str) else json.dumps(schema) + ) + if (js := chat_template_kwargs.get("json_schema")) is not None: + config["json_schema"] = js if isinstance(js, str) else json.dumps(js) + if "json_mode" in chat_template_kwargs: + config["json_mode"] = bool(chat_template_kwargs["json_mode"]) + + if fmt == "cmd3": + if (sm := chat_template_kwargs.get("safety_mode")) is not None: + config["safety_mode"] = str(sm).lower() + # citation_quality: ``on`` / ``off`` + cq = chat_template_kwargs.get("citation_quality") + if cq is None and (co := chat_template_kwargs.get("citation_options")): + mode = co.get("mode") if isinstance(co, dict) else None + if mode is not None: + cq = "on" if str(mode).lower() != "off" else "off" + if cq is not None: + config["citation_quality"] = str(cq).lower() + if "skip_preamble" in chat_template_kwargs: + config["skip_preamble"] = bool(chat_template_kwargs["skip_preamble"]) + else: # cmd4 + # cmd4 uses ``grounding`` rather than safety_mode/citation_quality. + # melody's cmd4 only accepts ``unknown`` / ``enabled`` / ``disabled``, + # so the Cohere v2 ``citation_options.mode`` values + # (``FAST`` / ``ACCURATE`` / ``OFF``) have to be normalized -- a + # raw lowercased passthrough would raise from + # ``render_cmd4`` (cmd4 has no fast/accurate distinction at the + # prompt-template layer; both request grounding-on). + if (gr := chat_template_kwargs.get("grounding")) is not None: + config["grounding"] = _normalize_cmd4_grounding(gr) + elif (co := chat_template_kwargs.get("citation_options")): + mode = co.get("mode") if isinstance(co, dict) else None + if mode is not None: + config["grounding"] = _normalize_cmd4_grounding(mode) + if (pi := chat_template_kwargs.get("platform_instruction")) is not None: + config["platform_instruction"] = str(pi) + + # Anything we didn't explicitly interpret above is forwarded to melody + # as a Jinja template variable. This matches vLLM's documented contract + # for ``chat_template_kwargs`` ("kwargs accessible by the template") + # and lets callers write e.g. ``{"reasoning_effort": "low"}`` directly + # without a nested ``additional_template_fields`` wrapper. + extra = { + k: v + for k, v in chat_template_kwargs.items() + if k not in _RENDERER_CONSUMED_KEYS + } + if extra: + config["additional_template_fields"] = extra + + return fmt, config + + +class CohereRenderer(BaseRenderer[HfTokenizer]): + """Renderer that templates Cohere prompts via the melody Rust bindings. + + Tokenization is delegated to the standard HF tokenizer; only the + chat-template step is replaced with ``cohere_melody.render_cmd3`` / + ``render_cmd4``. Picking this renderer is opt-in via + ``--tokenizer-mode cohere``. + """ + + def __init__( + self, + config: VllmConfig, + tokenizer: HfTokenizer | None, + ) -> None: + # Match HfRenderer in not mutating the cached tokenizer instance + tokenizer = copy.copy(tokenizer) + super().__init__(config, tokenizer) + + # Lazy import to keep `cohere_melody` an optional dependency + self._melody = _try_import_melody() + # ``render_cmd3`` / ``render_cmd4`` are pure CPU work; cache the + # thread-pool wrapper once so the async path doesn't allocate a new + # adapter on every request. + self._render_async = make_async(self._render, executor=self._executor) + + def _render(self, fmt: str, config_dict: dict[str, Any]) -> str: + if fmt == "cmd3": + return self._melody.render_cmd3(config_dict) + return self._melody.render_cmd4(config_dict) + + def render_messages( + self, + messages: list[ChatCompletionMessageParam], + params: ChatParams, + ) -> tuple[list[ConversationMessage], DictPrompt]: + conversation, mm_data, mm_uuids = parse_chat_messages( + messages, + self.model_config, + content_format="openai", + media_io_kwargs=params.media_io_kwargs, + mm_processor_kwargs=params.mm_processor_kwargs, + ) + + chat_template_kwargs = dict(params.chat_template_kwargs) + fmt, config_dict = _build_render_config(conversation, chat_template_kwargs) + prompt_text = self._render(fmt, config_dict) + prompt = parse_dec_only_prompt(prompt_text) + + if mm_data is not None: + prompt["multi_modal_data"] = mm_data + if mm_uuids is not None: + prompt["multi_modal_uuids"] = mm_uuids + + return conversation, prompt + + async def render_messages_async( + self, + messages: list[ChatCompletionMessageParam], + params: ChatParams, + ) -> tuple[list[ConversationMessage], DictPrompt]: + conversation, mm_data, mm_uuids = await parse_chat_messages_async( + messages, + self.model_config, + content_format="openai", + media_io_kwargs=params.media_io_kwargs, + mm_processor_kwargs=params.mm_processor_kwargs, + ) + + chat_template_kwargs = dict(params.chat_template_kwargs) + fmt, config_dict = _build_render_config(conversation, chat_template_kwargs) + prompt_text = await self._render_async(fmt, config_dict) + prompt = parse_dec_only_prompt(prompt_text) + + if mm_data is not None: + prompt["multi_modal_data"] = mm_data + if mm_uuids is not None: + prompt["multi_modal_uuids"] = mm_uuids + + return conversation, prompt diff --git a/vllm/renderers/registry.py b/vllm/renderers/registry.py index a6da9ec50178..bbaee67e3aec 100644 --- a/vllm/renderers/registry.py +++ b/vllm/renderers/registry.py @@ -20,6 +20,7 @@ _VLLM_RENDERERS = { + "cohere": ("cohere", "CohereRenderer"), "deepseek_v32": ("deepseek_v32", "DeepseekV32Renderer"), "deepseek_v4": ("deepseek_v4", "DeepseekV4Renderer"), "grok2": ("grok2", "Grok2Renderer"), diff --git a/vllm/tokenizers/registry.py b/vllm/tokenizers/registry.py index d928da3306ef..7eff8073f4ca 100644 --- a/vllm/tokenizers/registry.py +++ b/vllm/tokenizers/registry.py @@ -34,6 +34,9 @@ _MODEL_TYPES_WITH_INCORRECT_TOKENIZER_CLASS: set[str] = {"step3_vl", "step3p7"} _VLLM_TOKENIZERS = { + # ``cohere`` mode uses the standard cached HF tokenizer; only the + # renderer (template stage) is replaced with a melody-based one. + "cohere": ("hf", "CachedHfTokenizer"), "deepseek_v32": ("deepseek_v32", "DeepseekV32Tokenizer"), "deepseek_v4": ("deepseek_v4", "DeepseekV4Tokenizer"), "grok2": ("grok2", "Grok2Tokenizer"),