@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 9c110fa522

[Frontend] Cohere chat v2 api support (#47189)

Signed-off-by: Andrew Berneshawi <andrewberneshawi@cohere.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>

andrewbcohere committed Aug 1, 2026 at 01:53 UTC 9c110fa522f0bdd1ab3df34727b71ce2c9454d7d
31 files changed +8770 -20
requirements/test/cpu.txt
+15 -1
@@ -126,6 +126,8 @@ click==8.4.2
126 # uvicorn
127 cloudpickle==3.1.2
128 # via -r requirements/test/../common.txt
129 +cohere==7.0.8
130 + # via -r requirements/test/cuda.in
131 cohere-melody==0.9.0
132 # via -r requirements/test/cuda.in
133 colorama==0.4.6
@@ -226,6 +228,8 @@ fastar==0.11.0
228 # via
229 # fastapi
230 # fastapi-cloud-cli
231 +fastavro==1.12.2
232 + # via cohere
233 fastparquet==2024.11.0
234 # via genai-perf
235 fastrlock==0.8.2
@@ -326,6 +330,7 @@ httpx==0.27.2
330 # via
331 # -r requirements/test/cuda.in
332 # anthropic
333 + # cohere
334 # fastapi
335 # fastapi-cloud-cli
336 # huggingface-hub
@@ -772,6 +777,7 @@ pydantic==2.12.0
777 # -r requirements/test/../common.txt
778 # albumentations
779 # anthropic
780 + # cohere
781 # compressed-tensors
782 # datamodel-code-generator
783 # fastapi
@@ -789,7 +795,9 @@ pydantic==2.12.0
795 # ray
796 # xgrammar
797 pydantic-core==2.41.1
792 - # via pydantic
798 + # via
799 + # cohere
800 + # pydantic
801 pydantic-extra-types==2.10.5
802 # via
803 # fastapi
@@ -911,6 +919,7 @@ requests==2.32.3
919 # -r requirements/test/../common.txt
920 # azure-core
921 # buildkite-test-collector
922 + # cohere
923 # datasets
924 # docker
925 # evaluate
@@ -1110,6 +1119,7 @@ tokenizers==0.22.2
1119 # via
1120 # -r requirements/test/../common.txt
1121 # -r requirements/test/cuda.in
1122 + # cohere
1123 # transformers
1124 torch==2.13.0+cpu
1125 # via
@@ -1191,6 +1201,8 @@ typer==0.26.8
1201 # fastsafetensors
1202 # perceptron
1203 # transformers
1204 +types-requests==2.33.0.20260712
1205 + # via cohere
1206 typing-extensions==4.15.0
1207 # via
1208 # -r requirements/test/../common.txt
@@ -1204,6 +1216,7 @@ typing-extensions==4.15.0
1216 # azure-identity
1217 # azure-storage-blob
1218 # chz
1219 + # cohere
1220 # fastapi
1221 # grpcio
1222 # huggingface-hub
@@ -1248,6 +1261,7 @@ urllib3==2.2.3
1261 # responses
1262 # sentry-sdk
1263 # tritonclient
1264 + # types-requests
1265 uvicorn==0.35.0
1266 # via
1267 # fastapi
requirements/test/cuda.in
+1
@@ -73,6 +73,7 @@ gpt-oss >= 0.0.7; python_version > '3.11'
73 perceptron # required for isaac test
74 kaldi-native-fbank >= 1.18.7 # required for fireredasr2 test
75 cohere_melody>=0.9.0 # required for cohere command reasoning parser test
76 +cohere>=7.0.0 # required for cohere chat v2 api tests (protocol/serving import cohere.types)
77
78 # Newer versions of datasets require torchcoded, that makes the tests fail in CI because of a missing library.
79 # Older versions are in conflict with teerratorch requirements.
requirements/test/cuda.txt
+15 -1
@@ -131,6 +131,8 @@ click==8.4.2
131 # uvicorn
132 cloudpickle==3.1.2
133 # via -r requirements/test/../common.txt
134 +cohere==7.0.8
135 + # via -r requirements/test/cuda.in
136 cohere-melody==0.9.0
137 # via -r requirements/test/cuda.in
138 colorama==0.4.6
@@ -242,6 +244,8 @@ fastar==0.11.0
244 # via
245 # fastapi
246 # fastapi-cloud-cli
247 +fastavro==1.12.2
248 + # via cohere
249 fastparquet==2024.11.0
250 # via genai-perf
251 fastrlock==0.8.2
@@ -345,6 +349,7 @@ httpx==0.27.2
349 # via
350 # -r requirements/test/cuda.in
351 # anthropic
352 + # cohere
353 # fastapi
354 # fastapi-cloud-cli
355 # huggingface-hub
@@ -860,6 +865,7 @@ pydantic==2.12.0
865 # -r requirements/test/../common.txt
866 # albumentations
867 # anthropic
868 + # cohere
869 # compressed-tensors
870 # datamodel-code-generator
871 # fastapi
@@ -877,7 +883,9 @@ pydantic==2.12.0
883 # ray
884 # xgrammar
885 pydantic-core==2.41.1
880 - # via pydantic
886 + # via
887 + # cohere
888 + # pydantic
889 pydantic-extra-types==2.10.5
890 # via
891 # fastapi
@@ -1002,6 +1010,7 @@ requests==2.32.3
1010 # -r requirements/test/../common.txt
1011 # azure-core
1012 # buildkite-test-collector
1013 + # cohere
1014 # datasets
1015 # docker
1016 # evaluate
@@ -1207,6 +1216,7 @@ tokenizers==0.22.2
1216 # -c requirements/common.txt
1217 # -r requirements/test/../common.txt
1218 # -r requirements/test/cuda.in
1219 + # cohere
1220 # transformers
1221 torch==2.13.0+cu130
1222 # via
@@ -1296,6 +1306,8 @@ typer==0.26.8
1306 # fastsafetensors
1307 # perceptron
1308 # transformers
1309 +types-requests==2.33.0.20260712
1310 + # via cohere
1311 typing-extensions==4.15.0
1312 # via
1313 # -c requirements/common.txt
@@ -1310,6 +1322,7 @@ typing-extensions==4.15.0
1322 # azure-identity
1323 # azure-storage-blob
1324 # chz
1325 + # cohere
1326 # fastapi
1327 # grpcio
1328 # huggingface-hub
@@ -1354,6 +1367,7 @@ urllib3==2.2.3
1367 # responses
1368 # sentry-sdk
1369 # tritonclient
1370 + # types-requests
1371 uvicorn==0.35.0
1372 # via
1373 # fastapi
requirements/test/rocm.in
+1
@@ -71,6 +71,7 @@ gpt-oss>=0.0.7; python_version > '3.11'
71 perceptron # required for isaac test
72 kaldi-native-fbank>=1.18.7 # required for fireredasr2 test
73 cohere_melody>=0.9.0 # required for cohere command reasoning parser test
74 +cohere>=7.0.0 # required for cohere chat v2 api tests (protocol/serving import cohere.types)
75
76 # Newer versions of datasets require torchcoded, that makes the tests fail in CI because of a missing library.
77 # Older versions are in conflict with terratorch requirements.
requirements/test/rocm.txt
+15 -1
@@ -131,6 +131,8 @@ cloudpickle==3.1.2
131 # via
132 # -r requirements/test/../common.txt
133 # tilelang
134 +cohere==7.0.8
135 + # via -r requirements/test/rocm.in
136 cohere-melody==0.9.0
137 # via -r requirements/test/rocm.in
138 colorama==0.4.6
@@ -231,6 +233,8 @@ fastapi-cloud-cli==0.16.1
233 # via fastapi-cli
234 fastar==0.10.0
235 # via fastapi-cloud-cli
236 +fastavro==1.12.2
237 + # via cohere
238 fastparquet==2026.3.0
239 # via genai-perf
240 fastsafetensors==0.3.2
@@ -336,6 +340,7 @@ httpx==0.27.2
340 # via
341 # -r requirements/test/rocm.in
342 # anthropic
343 + # cohere
344 # fastapi
345 # fastapi-cloud-cli
346 # huggingface-hub
@@ -832,6 +837,7 @@ pydantic==2.12.5
837 # -r requirements/test/../common.txt
838 # albumentations
839 # anthropic
840 + # cohere
841 # compressed-tensors
842 # datamodel-code-generator
843 # fastapi
@@ -849,7 +855,9 @@ pydantic==2.12.5
855 # ray
856 # xgrammar
857 pydantic-core==2.41.5
852 - # via pydantic
858 + # via
859 + # cohere
860 + # pydantic
861 pydantic-extra-types==2.11.1
862 # via
863 # fastapi
@@ -978,6 +986,7 @@ requests==2.32.5
986 # -r requirements/test/../common.txt
987 # azure-core
988 # buildkite-test-collector
989 + # cohere
990 # datasets
991 # docker
992 # evaluate
@@ -1193,6 +1202,7 @@ tokenizers==0.22.2
1202 # -c requirements/common.txt
1203 # -r requirements/test/../common.txt
1204 # -r requirements/test/rocm.in
1205 + # cohere
1206 # transformers
1207 torch-c-dlpack-ext==0.1.5
1208 # via tilelang
@@ -1243,6 +1253,8 @@ typer==0.24.1
1253 # fastsafetensors
1254 # perceptron
1255 # transformers
1256 +types-requests==2.33.0.20260712
1257 + # via cohere
1258 typing-extensions==4.15.0
1259 # via
1260 # -c requirements/common.txt
@@ -1257,6 +1269,7 @@ typing-extensions==4.15.0
1269 # azure-identity
1270 # azure-storage-blob
1271 # chz
1272 + # cohere
1273 # fastapi
1274 # grpcio
1275 # huggingface-hub
@@ -1302,6 +1315,7 @@ urllib3==2.6.3
1315 # responses
1316 # sentry-sdk
1317 # tritonclient
1318 + # types-requests
1319 uvicorn==0.42.0
1320 # via
1321 # fastapi
tests/entrypoints/cohere/__init__.py
tests/entrypoints/cohere/test_api_router.py new
+425
@@ -0,0 +1,425 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Unit tests for ``vllm/entrypoints/cohere/api_router.py``.
4 +
5 +Covers:
6 +
7 +* The optional-import guard: ``attach_router`` is a no-op when the
8 + ``cohere`` SDK isn't installed.
9 +* The env-var opt-in gate: ``attach_router`` is a no-op unless
10 + ``VLLM_ENABLE_COHERE_API=1`` is set.
11 +* The router wiring: response shapes (JSON + SSE), error translation,
12 + and the ``cohere_serving_chat_v2 is None`` fallback (501 Not
13 + Implemented).
14 +"""
15 +
16 +import json
17 +from argparse import Namespace
18 +from collections.abc import AsyncGenerator
19 +from http import HTTPStatus
20 +
21 +import pytest
22 +from fastapi import FastAPI, HTTPException
23 +from fastapi.exceptions import RequestValidationError
24 +from fastapi.testclient import TestClient
25 +
26 +from vllm.entrypoints.cohere import api_router as api_router_mod
27 +from vllm.entrypoints.cohere.api_router import attach_router
28 +from vllm.entrypoints.cohere.protocol import (
29 + AssistantMessageResponse,
30 + CohereChatV2Response,
31 +)
32 +from vllm.entrypoints.openai.engine.protocol import ErrorInfo, ErrorResponse
33 +from vllm.entrypoints.serve.utils.server_utils import (
34 + http_exception_handler,
35 + validation_exception_handler,
36 +)
37 +
38 +
39 +@pytest.fixture(autouse=True)
40 +def _enable_cohere_api(monkeypatch):
41 + """Auto-enable the Cohere API gate for every test in this module.
42 +
43 + The endpoint is opt-in in production (``VLLM_ENABLE_COHERE_API=1``);
44 + every test in this file exercises the enabled path *except* the
45 + dedicated gate test in :class:`TestEnvVarGate`, which unsets the
46 + flag inside the test body.
47 + """
48 + monkeypatch.setenv("VLLM_ENABLE_COHERE_API", "1")
49 +
50 +
51 +# ----------------------------------------------------------------------
52 +# Fakes
53 +# ----------------------------------------------------------------------
54 +
55 +
56 +class _Handler:
57 + """Minimal stand-in for :class:`CohereServingChatV2` used by the
58 + router. Each test sets ``self.result`` to either:
59 +
60 + * a :class:`CohereChatV2Response` (non-streaming JSON path);
61 + * an async generator yielding SSE frames (streaming path);
62 + * an :class:`ErrorResponse` (error envelope path); or
63 + * an exception (router-level 500 path).
64 + """
65 +
66 + def __init__(self, result):
67 + self.result = result
68 +
69 + async def create_chat_v2(self, request, raw_request):
70 + if isinstance(self.result, Exception):
71 + raise self.result
72 + return self.result
73 +
74 +
75 +def _build_app(handler: _Handler | None) -> FastAPI:
76 + app = FastAPI()
77 + attach_router(app)
78 + app.state.cohere_serving_chat_v2 = handler
79 + return app
80 +
81 +
82 +def _build_app_with_vllm_handlers(handler: _Handler | None) -> FastAPI:
83 + """Build a FastAPI app that mirrors the real vLLM setup by installing
84 + ``validation_exception_handler`` and ``http_exception_handler``. The
85 + :class:`CohereErrorEnvelopeMiddleware` registered by ``attach_router``
86 + is expected to translate any resulting vLLM ``ErrorResponse`` body
87 + into the ``CohereError`` wire shape.
88 + """
89 + app = FastAPI()
90 + attach_router(app)
91 + app.state.cohere_serving_chat_v2 = handler
92 + # ``validation_exception_handler`` reads ``req.app.state.args``; the
93 + # real cli builds this via argparse.
94 + app.state.args = Namespace(log_error_stack=False)
95 + app.exception_handler(RequestValidationError)(validation_exception_handler)
96 + app.exception_handler(HTTPException)(http_exception_handler)
97 + return app
98 +
99 +
100 +def _minimal_request_body() -> dict:
101 + return {
102 + "model": "m",
103 + "messages": [{"role": "user", "content": "hi"}],
104 + }
105 +
106 +
107 +# ----------------------------------------------------------------------
108 +# Optional-import guard
109 +# ----------------------------------------------------------------------
110 +
111 +
112 +class TestOptionalCohereImport:
113 + """``attach_router`` probes for the SDK once at module load (because
114 + the route handler uses types imported from ``cohere``) and stashes
115 + the result in ``_SDK_AVAILABLE``. Tests simulate the "SDK missing"
116 + state by flipping that flag for the duration of the test.
117 +
118 + ``attach_router`` checks the env-var gate before the SDK probe, so
119 + the SDK-missing branch is only reachable when the operator opts in
120 + via ``VLLM_ENABLE_COHERE_API=1``. The flag-off-and-SDK-missing case
121 + below exists to pin down that ordering — the flag-off short-circuits
122 + """
123 +
124 + def test_flag_off_and_sdk_missing_stays_silent_about_sdk(self, monkeypatch, caplog):
125 + """Flag off doesn't do SDK-missing check.
126 +
127 + When the operator hasn't opted in, ``attach_router`` must not
128 + warn about the ``cohere`` SDK being missing: they never asked
129 + for the endpoint, so surfacing the SDK gap is misleading noise.
130 + Only the flag-off DEBUG message should fire.
131 + """
132 + monkeypatch.delenv("VLLM_ENABLE_COHERE_API", raising=False)
133 + monkeypatch.setattr(api_router_mod, "_SDK_AVAILABLE", False)
134 +
135 + with caplog.at_level("DEBUG", logger="vllm.entrypoints.cohere.api_router"):
136 + app = FastAPI()
137 + attach_router(app)
138 +
139 + paths = [getattr(r, "path", None) for r in app.routes]
140 + assert "/cohere/v2/chat" not in paths
141 + # The flag-off short-circuit ran; the SDK check never did.
142 + assert not any(
143 + "SDK is not installed" in rec.message for rec in caplog.records
144 + ), "SDK-missing log leaked despite the flag being off"
145 +
146 + def test_flag_on_but_sdk_missing_logs_warning(self, monkeypatch, caplog):
147 + """Misconfiguration path: the operator explicitly opted into the
148 + endpoint via ``VLLM_ENABLE_COHERE_API=1`` (already set by the
149 + autouse fixture) but forgot to install ``cohere``.
150 + """
151 + monkeypatch.setattr(api_router_mod, "_SDK_AVAILABLE", False)
152 +
153 + with caplog.at_level("DEBUG", logger="vllm.entrypoints.cohere.api_router"):
154 + app = FastAPI()
155 + attach_router(app)
156 +
157 + paths = [getattr(r, "path", None) for r in app.routes]
158 + assert "/cohere/v2/chat" not in paths
159 + warn_sdk_records = [
160 + rec
161 + for rec in caplog.records
162 + if "VLLM_ENABLE_COHERE_API=1" in rec.message
163 + and "SDK is not installed" in rec.message
164 + ]
165 + assert warn_sdk_records, (
166 + "expected a WARNING that pairs the opt-in flag with the "
167 + "missing SDK so operators notice the misconfiguration"
168 + )
169 + assert all(rec.levelname == "WARNING" for rec in warn_sdk_records)
170 +
171 + def test_attach_router_registers_route_when_cohere_present(self):
172 + app = _build_app(handler=None)
173 + paths = [getattr(r, "path", None) for r in app.routes]
174 + assert "/cohere/v2/chat" in paths
175 +
176 +
177 +# ----------------------------------------------------------------------
178 +# VLLM_ENABLE_COHERE_API gate
179 +# ----------------------------------------------------------------------
180 +
181 +
182 +class TestEnvVarGate:
183 + """The Cohere v2 endpoint is opt-in via ``VLLM_ENABLE_COHERE_API``.
184 +
185 + Even with the SDK installed, :func:`attach_router` must skip route
186 + registration and middleware installation unless the env flag is
187 + set. The autouse fixture on this module enables the flag by
188 + default, so each test here explicitly disables it.
189 + """
190 +
191 + def test_attach_router_noop_when_flag_unset(self, monkeypatch, caplog):
192 + monkeypatch.delenv("VLLM_ENABLE_COHERE_API", raising=False)
193 +
194 + # The flag-off skip logs at DEBUG on purpose: this is the default
195 + # state for every non-Cohere vLLM deployment, so an INFO log on
196 + # every server startup would be pointless noise. The test raises
197 + # caplog's level accordingly.
198 + with caplog.at_level("DEBUG", logger="vllm.entrypoints.cohere.api_router"):
199 + app = FastAPI()
200 + attach_router(app)
201 +
202 + paths = [getattr(r, "path", None) for r in app.routes]
203 + assert "/cohere/v2/chat" not in paths
204 + debug_flag_records = [
205 + rec
206 + for rec in caplog.records
207 + if "VLLM_ENABLE_COHERE_API is not set" in rec.message
208 + ]
209 + assert debug_flag_records, (
210 + "expected a DEBUG message that the cohere flag is off"
211 + )
212 + assert all(rec.levelname == "DEBUG" for rec in debug_flag_records)
213 +
214 + def test_attach_router_noop_when_flag_zero(self, monkeypatch):
215 + monkeypatch.setenv("VLLM_ENABLE_COHERE_API", "0")
216 +
217 + app = FastAPI()
218 + attach_router(app)
219 +
220 + paths = [getattr(r, "path", None) for r in app.routes]
221 + assert "/cohere/v2/chat" not in paths
222 +
223 +
224 +# ----------------------------------------------------------------------
225 +# Endpoint behavior
226 +# ----------------------------------------------------------------------
227 +
228 +
229 +class TestEndpoint:
230 + def test_501_when_handler_missing(self):
231 + app = _build_app(handler=None)
232 + with TestClient(app) as client:
233 + r = client.post("/cohere/v2/chat", json=_minimal_request_body())
234 + assert r.status_code == HTTPStatus.NOT_IMPLEMENTED
235 + body = r.json()
236 + assert "does not support" in body["message"]
237 + assert "id" not in body # excluded by ``exclude_none=True``
238 +
239 + def test_non_streaming_response_is_json(self):
240 + msg = AssistantMessageResponse(content=[{"type": "text", "text": "hello"}])
241 + result = CohereChatV2Response(id="r1", finish_reason="COMPLETE", message=msg)
242 + app = _build_app(handler=_Handler(result))
243 + with TestClient(app) as client:
244 + r = client.post("/cohere/v2/chat", json=_minimal_request_body())
245 + assert r.status_code == HTTPStatus.OK
246 + assert r.headers["content-type"].startswith("application/json")
247 + body = r.json()
248 + assert body["id"] == "r1"
249 + assert body["finish_reason"] == "COMPLETE"
250 + assert body["message"]["content"][0]["text"] == "hello"
251 +
252 + def test_streaming_response_is_sse(self):
253 + async def _gen() -> AsyncGenerator[str, None]:
254 + yield 'data: {"type":"message-start"}\n\n'
255 + yield "data: [DONE]\n\n"
256 +
257 + app = _build_app(handler=_Handler(_gen()))
258 + with TestClient(app) as client:
259 + r = client.post(
260 + "/cohere/v2/chat",
261 + json={**_minimal_request_body(), "stream": True},
262 + )
263 + assert r.status_code == HTTPStatus.OK
264 + assert r.headers["content-type"].startswith("text/event-stream")
265 + body = r.text
266 + assert "message-start" in body
267 + assert body.rstrip().endswith("[DONE]")
268 +
269 + def test_error_response_translated_to_cohere_envelope(self):
270 + err = ErrorResponse(
271 + error=ErrorInfo(
272 + message="bad request",
273 + type="bad_request",
274 + code=400,
275 + )
276 + )
277 + app = _build_app(handler=_Handler(err))
278 + with TestClient(app) as client:
279 + r = client.post("/cohere/v2/chat", json=_minimal_request_body())
280 + assert r.status_code == HTTPStatus.BAD_REQUEST
281 + body = r.json()
282 + assert body == {"message": "bad request"}
283 +
284 + def test_handler_exception_returns_500_envelope(self):
285 + app = _build_app(handler=_Handler(RuntimeError("kaboom")))
286 + with TestClient(app) as client:
287 + r = client.post("/cohere/v2/chat", json=_minimal_request_body())
288 + assert r.status_code == HTTPStatus.INTERNAL_SERVER_ERROR
289 + body = r.json()
290 + assert body == {"message": "kaboom"}
291 +
292 + def test_non_json_content_type_rejected(self):
293 + """The ``validate_json_request`` dependency raises
294 + ``RequestValidationError`` (HTTP 422) for non-JSON content
295 + types, matching the behavior of the other vLLM API routers.
296 + """
297 + app = _build_app(handler=None)
298 + with TestClient(app) as client:
299 + r = client.post(
300 + "/cohere/v2/chat",
301 + content=json.dumps(_minimal_request_body()),
302 + headers={"content-type": "text/plain"},
303 + )
304 + assert r.status_code == HTTPStatus.UNPROCESSABLE_ENTITY
305 +
306 + def test_invalid_body_returns_422(self):
307 + # ``model`` is required; omit it to trip Pydantic validation.
308 + app = _build_app(handler=None)
309 + with TestClient(app) as client:
310 + r = client.post(
311 + "/cohere/v2/chat",
312 + json={"messages": [{"role": "user", "content": "hi"}]},
313 + )
314 + assert r.status_code == HTTPStatus.UNPROCESSABLE_ENTITY
315 +
316 +
317 +# ----------------------------------------------------------------------
318 +# CohereErrorEnvelopeMiddleware
319 +# ----------------------------------------------------------------------
320 +
321 +
322 +class TestCohereErrorEnvelope:
323 + """When the app installs vLLM's global exception handlers, validation
324 + and HTTP errors escape as ``ErrorResponse`` bodies before the route
325 + handler runs. The middleware installed by ``attach_router`` must
326 + normalise those bodies to the ``CohereError`` shape declared on the
327 + endpoint's OpenAPI ``responses`` map so schema-conformance tests
328 + (``test_openai_schema.py``) don't see a mismatch on ``/cohere/*``
329 + responses.
330 + """
331 +
332 + def test_validation_error_body_is_cohere_shaped(self):
333 + # ``model=""`` and ``messages=[]`` trip our custom field
334 + # validators, which raise pydantic ValueErrors and are routed
335 + # through ``validation_exception_handler`` in the real vLLM
336 + # server (producing the ``{"error": {...}}`` shape).
337 + app = _build_app_with_vllm_handlers(handler=None)
338 + with TestClient(app) as client:
339 + r = client.post("/cohere/v2/chat", json={"messages": [], "model": ""})
340 + assert r.status_code == HTTPStatus.BAD_REQUEST
341 + body = r.json()
342 + # ``CohereError`` has ``message`` at the top level, not nested
343 + # under an ``error`` envelope.
344 + assert "error" not in body
345 + assert "message" in body
346 + assert isinstance(body["message"], str) and body["message"]
347 +
348 + def test_http_error_body_is_cohere_shaped(self):
349 + # A raised ``HTTPException`` from anywhere in the request cycle
350 + # is routed through ``http_exception_handler`` (producing the
351 + # ``ErrorResponse`` shape) and must be translated.
352 + app = _build_app_with_vllm_handlers(handler=None)
353 +
354 + @app.get("/cohere/v2/boom")
355 + async def _boom():
356 + raise HTTPException(status_code=418, detail="teapot")
357 +
358 + with TestClient(app) as client:
359 + r = client.get("/cohere/v2/boom")
360 + assert r.status_code == 418
361 + body = r.json()
362 + assert body == {"message": "teapot"}
363 +
364 + def test_non_cohere_path_is_not_translated(self):
365 + app = _build_app_with_vllm_handlers(handler=None)
366 +
367 + @app.get("/v1/other")
368 + async def _other():
369 + raise HTTPException(status_code=400, detail="nope")
370 +
371 + with TestClient(app) as client:
372 + r = client.get("/v1/other")
373 + assert r.status_code == HTTPStatus.BAD_REQUEST
374 + body = r.json()
375 + # Non-cohere paths keep the vLLM ``ErrorResponse`` shape.
376 + assert "error" in body
377 + assert body["error"]["message"] == "nope"
378 +
379 + def test_streaming_response_passes_through(self):
380 + # SSE responses have content-type text/event-stream; the
381 + # middleware must never buffer these (which would break
382 + # streaming) even though they're on ``/cohere/*``.
383 + async def _gen() -> AsyncGenerator[str, None]:
384 + yield 'data: {"type":"message-start"}\n\n'
385 + yield "data: [DONE]\n\n"
386 +
387 + app = _build_app_with_vllm_handlers(handler=_Handler(_gen()))
388 + with TestClient(app) as client:
389 + r = client.post(
390 + "/cohere/v2/chat",
391 + json={**_minimal_request_body(), "stream": True},
392 + )
393 + assert r.status_code == HTTPStatus.OK
394 + assert r.headers["content-type"].startswith("text/event-stream")
395 + assert "message-start" in r.text
396 + assert r.text.rstrip().endswith("[DONE]")
397 +
398 + def test_already_cohere_shaped_body_passes_through(self):
399 + # When the handler returns an ``ErrorResponse`` the route
400 + # itself translates it to ``CohereError``; the middleware sees
401 + # the ``CohereError`` shape and must leave it alone.
402 + err = ErrorResponse(
403 + error=ErrorInfo(message="already cohere", type="Bad Request", code=400)
404 + )
405 + app = _build_app_with_vllm_handlers(handler=_Handler(err))
406 + with TestClient(app) as client:
407 + r = client.post("/cohere/v2/chat", json=_minimal_request_body())
408 + assert r.status_code == HTTPStatus.BAD_REQUEST
409 + body = r.json()
410 + # No ``error`` wrapper: the route already emitted the wire shape.
411 + assert body == {"message": "already cohere"}
412 +
413 + def test_request_id_preserved_in_translated_body(self):
414 + # Client-provided ``X-Request-Id`` should be echoed as
415 + # ``CohereError.id`` so callers can correlate failures.
416 + app = _build_app_with_vllm_handlers(handler=None)
417 + with TestClient(app) as client:
418 + r = client.post(
419 + "/cohere/v2/chat",
420 + json={"messages": [], "model": ""},
421 + headers={"X-Request-Id": "req-abc"},
422 + )
423 + assert r.status_code == HTTPStatus.BAD_REQUEST
424 + body = r.json()
425 + assert body.get("id") == "req-abc"
tests/entrypoints/cohere/test_chat_v2.py new
+259
@@ -0,0 +1,259 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""End-to-end integration tests for the ``POST /cohere/v2/chat`` endpoint.
4 +
5 +These tests spin up a real ``vllm serve`` process via
6 +:class:`tests.utils.RemoteOpenAIServer` and exercise the Cohere Chat v2
7 +contract over HTTP. They mirror the pattern used by
8 +:mod:`tests.entrypoints.anthropic.test_messages`.
9 +
10 +Two layers of integration are covered:
11 +
12 +1. Raw HTTP via :mod:`httpx` — always runs, verifies the wire contract.
13 +2. Cohere SDK (``pip install cohere``) — auto-skipped when the optional
14 + dependency isn't installed, verifies SDK-level interop.
15 +
16 +We use a tiny generic chat model (``HuggingFaceTB/SmolLM2-135M-Instruct``,
17 +≈135M params) so the fixture stays runnable on CPU-only laptops. The
18 +Cohere v2 endpoint is model-agnostic and just performs the v2 ↔ OpenAI
19 +chat-completion translation, so the choice of model only matters for
20 +response *shape* — the translation logic itself is unit-tested
21 +elsewhere.
22 +"""
23 +
24 +import json
25 +
26 +import httpx
27 +import pytest
28 +import pytest_asyncio
29 +
30 +from tests.utils import RemoteOpenAIServer
31 +
32 +# Tiny chat-tuned model so the fixture is cheap to boot on CPU-only
33 +# hosts (Mac, CI runners without GPUs). The Cohere v2 translation is
34 +# completely independent of which underlying chat model is loaded.
35 +MODEL_NAME = "HuggingFaceTB/SmolLM2-135M-Instruct"
36 +SERVED_MODEL_NAME = "command-r-plus-08-2024"
37 +
38 +
39 +@pytest.fixture(scope="module")
40 +def server():
41 + args = [
42 + # 1k tokens is plenty for the prompts these tests send; trimming
43 + # it from the default keeps the KV-cache footprint tiny.
44 + "--max-model-len",
45 + "1024",
46 + "--max-num-seqs",
47 + "4",
48 + "--dtype",
49 + "bfloat16",
50 + "--enforce-eager",
51 + # Advertise a Cohere model name so Cohere SDK ``model=`` calls
52 + # round-trip without the ``model not found`` check.
53 + "--served-model-name",
54 + SERVED_MODEL_NAME,
55 + # Disable the reasoning-model path so the test doesn't require
56 + # the conversation to surface a thinking block (SmolLM2 doesn't
57 + # emit Cohere-style reasoning tokens out of the box).
58 + "--no-cohere-is-reasoning-model",
59 + ]
60 +
61 + # Cap the CPU KV-cache pool the vLLM CPU backend reserves at
62 + # startup. The default (4 GB) trips Mac's RAM check on smaller
63 + # machines; 1 GB is more than enough for max_model_len=1024 and
64 + # max_num_seqs=4.
65 + # ``VLLM_ENABLE_COHERE_API=1`` flips the opt-in gate that
66 + # ``vllm.entrypoints.cohere.api_router.attach_router`` reads to
67 + # decide whether to register ``POST /cohere/v2/chat`` at server
68 + # startup. Without it the fixture would boot a server that returns
69 + # 404 for every request the tests below make.
70 + env_dict = {
71 + "VLLM_CPU_KVCACHE_SPACE": "1",
72 + "VLLM_ENABLE_COHERE_API": "1",
73 + }
74 +
75 + with RemoteOpenAIServer(MODEL_NAME, args, env_dict=env_dict) as remote_server:
76 + yield remote_server
77 +
78 +
79 +# ----------------------------------------------------------------------
80 +# Layer 1: raw HTTP contract (no optional deps)
81 +# ----------------------------------------------------------------------
82 +
83 +
84 +@pytest_asyncio.fixture
85 +async def httpx_client(server):
86 + async with httpx.AsyncClient(
87 + base_url=server.url_root, timeout=httpx.Timeout(120.0)
88 + ) as client:
89 + yield client
90 +
91 +
92 +@pytest.mark.asyncio
93 +async def test_cohere_v2_chat_non_streaming(httpx_client: httpx.AsyncClient):
94 + resp = await httpx_client.post(
95 + "/cohere/v2/chat",
96 + json={
97 + "model": SERVED_MODEL_NAME,
98 + "messages": [{"role": "user", "content": "Say hi."}],
99 + "max_tokens": 16,
100 + "stream": False,
101 + },
102 + )
103 + assert resp.status_code == 200, resp.text
104 + payload = resp.json()
105 + # The response envelope follows Cohere v2's schema.
106 + assert "message" in payload
107 + assert payload["message"]["role"] == "assistant"
108 + # ``content`` is a list of content blocks; the synthesized
109 + # ``CohereServingChatV2`` should always emit at least one ``text``
110 + # block (it falls back to an empty block when the model returned
111 + # nothing).
112 + content = payload["message"]["content"]
113 + assert isinstance(content, list) and len(content) >= 1
114 + assert content[0]["type"] == "text"
115 + # ``usage`` is always populated by the translator.
116 + assert "usage" in payload
117 + assert "billed_units" in payload["usage"]
118 + assert "tokens" in payload["usage"]
119 + # ``finish_reason`` is one of Cohere's enum values.
120 + assert payload["finish_reason"] in {
121 + "COMPLETE",
122 + "MAX_TOKENS",
123 + "STOP_SEQUENCE",
124 + "TOOL_CALL",
125 + "ERROR",
126 + }
127 +
128 +
129 +@pytest.mark.asyncio
130 +async def test_cohere_v2_chat_streaming(httpx_client: httpx.AsyncClient):
131 + """Streaming returns SSE frames in the v2 message-lifecycle shape."""
132 + events: list[dict] = []
133 + async with httpx_client.stream(
134 + "POST",
135 + "/cohere/v2/chat",
136 + json={
137 + "model": SERVED_MODEL_NAME,
138 + "messages": [{"role": "user", "content": "Hi"}],
139 + "max_tokens": 8,
140 + "stream": True,
141 + },
142 + ) as resp:
143 + assert resp.status_code == 200, await resp.aread()
144 + assert resp.headers["content-type"].startswith("text/event-stream")
145 + async for line in resp.aiter_lines():
146 + if not line.startswith("data: "):
147 + continue
148 + data = line[len("data: ") :]
149 + if data == "[DONE]":
150 + events.append({"type": "_DONE_"})
151 + continue
152 + events.append(json.loads(data))
153 +
154 + types = [ev["type"] for ev in events]
155 + # The lifecycle always starts with message-start and ends with [DONE]
156 + # preceded by message-end.
157 + assert types[0] == "message-start"
158 + assert types[-1] == "_DONE_"
159 + assert types[-2] == "message-end"
160 + # message-start must carry the chunk/message id.
161 + assert events[0].get("id")
162 +
163 +
164 +@pytest.mark.asyncio
165 +async def test_cohere_v2_chat_validation_error_returns_400(
166 + httpx_client: httpx.AsyncClient,
167 +):
168 + # Missing required ``model`` field → FastAPI/Pydantic returns 400.
169 + resp = await httpx_client.post(
170 + "/cohere/v2/chat",
171 + json={"messages": []},
172 + )
173 + assert resp.status_code == 400
174 +
175 +
176 +@pytest.mark.asyncio
177 +async def test_cohere_v2_chat_documents_field_accepted(
178 + httpx_client: httpx.AsyncClient,
179 +):
180 + """The v2 endpoint forwards ``documents`` into chat_template_kwargs.
181 +
182 + We only assert the request is accepted and produces a 200 response —
183 + the renderer-level effect is covered by ``tests/renderers/test_cohere.py``.
184 + """
185 + resp = await httpx_client.post(
186 + "/cohere/v2/chat",
187 + json={
188 + "model": SERVED_MODEL_NAME,
189 + "messages": [{"role": "user", "content": "Summarize."}],
190 + "documents": [{"id": "d1", "data": {"title": "T", "snippet": "S"}}],
191 + "max_tokens": 16,
192 + "stream": False,
193 + },
194 + )
195 + assert resp.status_code == 200, resp.text
196 +
197 +
198 +# ----------------------------------------------------------------------
199 +# Layer 2: Cohere SDK round-trip (auto-skipped if SDK absent)
200 +# ----------------------------------------------------------------------
201 +
202 +
203 +@pytest_asyncio.fixture
204 +async def cohere_async_client(server):
205 + cohere = pytest.importorskip("cohere")
206 + # The vLLM endpoint is mounted at ``/cohere/v2/chat`` while the
207 + # cohere SDK targets ``${base_url}/v2/chat``; point base_url at the
208 + # ``/cohere`` prefix so paths line up.
209 + client = cohere.AsyncClientV2(
210 + api_key="dummy",
211 + base_url=server.url_for("cohere"),
212 + )
213 + try:
214 + yield client
215 + finally:
216 + # ``AsyncClientV2`` exposes a sync close; if a future version
217 + # adds aclose we still close cleanly.
218 + close = getattr(client, "aclose", None) or getattr(client, "close", None)
219 + if close is not None:
220 + result = close()
221 + if hasattr(result, "__await__"):
222 + await result
223 +
224 +
225 +@pytest.mark.asyncio
226 +async def test_cohere_sdk_non_streaming(cohere_async_client):
227 + resp = await cohere_async_client.chat(
228 + model=SERVED_MODEL_NAME,
229 + messages=[{"role": "user", "content": "Say hi."}],
230 + max_tokens=16,
231 + )
232 + # SDK parses our JSON into typed objects.
233 + assert resp.message.role == "assistant"
234 + assert resp.message.content is not None
235 + assert len(resp.message.content) >= 1
236 + assert resp.message.content[0].type == "text"
237 + assert resp.finish_reason in {
238 + "COMPLETE",
239 + "MAX_TOKENS",
240 + "STOP_SEQUENCE",
241 + "TOOL_CALL",
242 + "ERROR",
243 + }
244 +
245 +
246 +@pytest.mark.asyncio
247 +async def test_cohere_sdk_streaming(cohere_async_client):
248 + events: list[str] = []
249 + stream = cohere_async_client.chat_stream(
250 + model=SERVED_MODEL_NAME,
251 + messages=[{"role": "user", "content": "Hi"}],
252 + max_tokens=8,
253 + )
254 + async for ev in stream:
255 + events.append(ev.type)
256 +
257 + assert events, "SDK stream yielded no events"
258 + assert events[0] == "message-start"
259 + assert events[-1] == "message-end"
tests/entrypoints/cohere/test_cohere_chat_message.py new
+135
@@ -0,0 +1,135 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Unit tests for :mod:`vllm.entrypoints.cohere.cohere_chat_message`.
4 +
5 +The module hosts the citation-carrying subclasses of the OpenAI chat
6 +completion protocol (``CohereChatMessage`` / ``CohereDeltaMessage``) plus
7 +the shared ``Citation`` / ``CitationSource`` shapes. These tests pin:
8 +
9 +1. Serialization behavior we own (empty-``citations`` cleanup on the
10 + subclass ``_serialize``, ``exclude_none`` shape, round-trip).
11 +2. The constraint that ``type`` discriminators reject unknown values.
12 +"""
13 +
14 +import pytest
15 +from pydantic import ValidationError
16 +
17 +from vllm.entrypoints.cohere.cohere_chat_message import (
18 + Citation,
19 + CitationSource,
20 + CohereChatMessage,
21 + CohereDeltaMessage,
22 +)
23 +
24 +# ======================================================================
25 +# CitationSource
26 +# ======================================================================
27 +
28 +
29 +class TestCitationSource:
30 + def test_invalid_type_rejected(self):
31 + with pytest.raises(ValidationError):
32 + CitationSource(type="other") # type: ignore[arg-type]
33 +
34 + def test_none_fields_excluded_from_dump(self):
35 + s = CitationSource(type="document", id="d1")
36 + assert s.model_dump(exclude_none=True) == {
37 + "type": "document",
38 + "id": "d1",
39 + }
40 +
41 + def test_resolved_source_survives_json_round_trip(self):
42 + # The parser produces fully-resolved sources (see
43 + # ``_melody_sources_to_vllm`` in
44 + # ``vllm/reasoning/cohere_command_reasoning_parser.py``) which
45 + # then flow through the internal streaming pipeline (parser ->
46 + # OpenAI stream chunk -> ``_chat_completion_stream_to_v2``).
47 + # Every field must round-trip through ``model_dump_json`` at
48 + # that layer or the wire event ends up missing pieces.
49 + src = CitationSource(
50 + type="tool",
51 + id="res_a0",
52 + tool_output={"id": "res_a0", "text": "r"},
53 + )
54 + src2 = CitationSource.model_validate_json(src.model_dump_json())
55 + assert src2.type == "tool"
56 + assert src2.id == "res_a0"
57 + assert src2.tool_output == {"id": "res_a0", "text": "r"}
58 +
59 +
60 +# ======================================================================
61 +# Citation
62 +# ======================================================================
63 +
64 +
65 +class TestCitation:
66 + def test_invalid_type_rejected(self):
67 + with pytest.raises(ValidationError):
68 + Citation(type="OTHER") # type: ignore[arg-type]
69 +
70 + def test_dump_excludes_none_fields(self):
71 + c = Citation(start=0, end=5, text="hello")
72 + dumped = c.model_dump(exclude_none=True)
73 + assert dumped == {
74 + "start": 0,
75 + "end": 5,
76 + "text": "hello",
77 + "sources": [],
78 + }
79 +
80 +
81 +# ======================================================================
82 +# CohereDeltaMessage
83 +# ======================================================================
84 +
85 +
86 +class TestCohereDeltaMessage:
87 + def test_citations_omitted_from_dump_when_none(self):
88 + # Non-grounded deltas should never carry ``citations: null`` on
89 + # the wire even when routed through the Cohere subclass.
90 + d = CohereDeltaMessage(content="hello")
91 + dumped = d.model_dump(exclude_none=True)
92 + assert "citations" not in dumped
93 +
94 + def test_citations_dump_round_trip(self):
95 + d = CohereDeltaMessage(citations=[Citation(start=0, end=5, text="hi")])
96 + dumped = d.model_dump(exclude_none=True)
97 + assert dumped["citations"][0]["text"] == "hi"
98 + # Round-trip through the subclass to confirm the field validates.
99 + d2 = CohereDeltaMessage.model_validate(dumped)
100 + assert d2.citations[0].text == "hi"
101 +
102 +
103 +# ======================================================================
104 +# CohereChatMessage
105 +# ======================================================================
106 +
107 +
108 +class TestCohereChatMessage:
109 + def test_citations_omitted_from_dump_when_none(self):
110 + m = CohereChatMessage(role="assistant", content="hi")
111 + dumped = m.model_dump(exclude_none=True)
112 + assert "citations" not in dumped
113 +
114 + def test_citations_dump_round_trip(self):
115 + m = CohereChatMessage(
116 + role="assistant",
117 + content="hello",
118 + citations=[
119 + Citation(
120 + start=0,
121 + end=5,
122 + text="hello",
123 + sources=[CitationSource(type="document", id="d1")],
124 + )
125 + ],
126 + )
127 + dumped = m.model_dump(exclude_none=True)
128 + assert dumped["role"] == "assistant"
129 + assert dumped["content"] == "hello"
130 + assert dumped["citations"][0]["text"] == "hello"
131 + assert dumped["citations"][0]["sources"][0]["id"] == "d1"
132 + # Round-trip through the subclass to confirm the field validates.
133 + m2 = CohereChatMessage.model_validate(dumped)
134 + assert m2.citations[0].text == "hello"
135 + assert m2.citations[0].sources[0].id == "d1"
tests/entrypoints/cohere/test_protocol.py new
+217
@@ -0,0 +1,217 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Unit tests for ``vllm/entrypoints/cohere/protocol.py``.
4 +
5 +The module is mostly a thin wrapper around the official ``cohere`` SDK
6 +types plus a few local additions:
7 +
8 +* :class:`CohereError` envelope.
9 +* :class:`CohereChatV2Request` (model required, ``max_tokens`` non-negative).
10 +* :class:`CohereChatV2Response` plus the usage / logprob helpers.
11 +* The streaming event subclasses that bake a wire-format ``type``
12 + discriminator into ``model_dump()`` so SSE consumers can demux on it.
13 +"""
14 +
15 +import pytest
16 +from pydantic import ValidationError
17 +
18 +from vllm.entrypoints.cohere.protocol import (
19 + AssistantMessageResponse,
20 + CitationEndEvent,
21 + CitationStartEvent,
22 + CohereChatV2Request,
23 + CohereChatV2Response,
24 + CohereError,
25 + CohereUsageTokens,
26 + ContentDeltaEvent,
27 + ContentEndEvent,
28 + ContentStartEvent,
29 + MessageEndEvent,
30 + MessageStartEvent,
31 + ToolCallDeltaEvent,
32 + ToolCallEndEvent,
33 + ToolCallStartEvent,
34 + ToolPlanDeltaEvent,
35 +)
36 +
37 +# ======================================================================
38 +# CohereError
39 +# ======================================================================
40 +
41 +
42 +class TestCohereError:
43 + def test_model_dump_excludes_none(self):
44 + err = CohereError(message="boom")
45 + assert err.model_dump(exclude_none=True) == {"message": "boom"}
46 +
47 +
48 +# ======================================================================
49 +# CohereChatV2Request
50 +# ======================================================================
51 +
52 +
53 +class TestCohereChatV2Request:
54 + def test_empty_model_rejected(self):
55 + with pytest.raises(ValidationError, match="model is required"):
56 + CohereChatV2Request(model="", messages=[{"role": "user", "content": "hi"}])
57 +
58 + def test_negative_max_tokens_rejected(self):
59 + with pytest.raises(ValidationError, match="non-negative"):
60 + CohereChatV2Request(
61 + model="m",
62 + messages=[{"role": "user", "content": "hi"}],
63 + max_tokens=-1,
64 + )
65 +
66 + def test_zero_max_tokens_allowed(self):
67 + # Zero is allowed (the docs allow 0 -> return prompt only).
68 + req = CohereChatV2Request(
69 + model="m",
70 + messages=[{"role": "user", "content": "hi"}],
71 + max_tokens=0,
72 + )
73 + assert req.max_tokens == 0
74 +
75 + def test_invalid_tool_choice_rejected(self):
76 + with pytest.raises(ValidationError):
77 + CohereChatV2Request(
78 + model="m",
79 + messages=[{"role": "user", "content": "hi"}],
80 + tool_choice="ANY", # not REQUIRED/NONE
81 + )
82 +
83 +
84 +# ======================================================================
85 +# Usage / Logprobs
86 +# ======================================================================
87 +
88 +
89 +class TestUsage:
90 + def test_tokens_serialization(self):
91 + u = CohereUsageTokens(input_tokens=10, output_tokens=5)
92 + assert u.model_dump() == {"input_tokens": 10.0, "output_tokens": 5.0}
93 +
94 +
95 +# ======================================================================
96 +# CohereChatV2Response
97 +# ======================================================================
98 +
99 +
100 +class TestCohereChatV2Response:
101 + def test_invalid_finish_reason_rejected(self):
102 + msg = AssistantMessageResponse(content=[{"type": "text", "text": "hi"}])
103 + with pytest.raises(ValidationError):
104 + CohereChatV2Response(
105 + id="r1",
106 + finish_reason="NOT_A_REASON", # type: ignore[arg-type]
107 + message=msg,
108 + )
109 +
110 +
111 +# ======================================================================
112 +# Streaming event ``type`` discriminator baked into model_dump()
113 +# ======================================================================
114 +
115 +
116 +class TestStreamingEventTypeField:
117 + """Each event subclass adds a ``type: Literal[...]`` field with a
118 + default so ``model_dump()`` always emits the wire-format discriminator
119 + (the parent SDK classes don't declare ``type`` as a Pydantic field).
120 + """
121 +
122 + @pytest.mark.parametrize(
123 + "cls, expected_type, kwargs",
124 + [
125 + (
126 + MessageStartEvent,
127 + "message-start",
128 + {"id": "a", "delta": {"message": {"role": "assistant"}}},
129 + ),
130 + (
131 + ContentStartEvent,
132 + "content-start",
133 + {
134 + "index": 0,
135 + "delta": {"message": {"content": {"type": "text", "text": ""}}},
136 + },
137 + ),
138 + (
139 + ContentDeltaEvent,
140 + "content-delta",
141 + {
142 + "index": 0,
143 + "delta": {"message": {"content": {"text": "hi"}}},
144 + },
145 + ),
146 + (ContentEndEvent, "content-end", {"index": 0}),
147 + (
148 + ToolPlanDeltaEvent,
149 + "tool-plan-delta",
150 + {"delta": {"message": {"tool_plan": "thinking"}}},
151 + ),
152 + (
153 + ToolCallStartEvent,
154 + "tool-call-start",
155 + {
156 + "index": 0,
157 + "delta": {
158 + "message": {
159 + "tool_calls": {
160 + "id": "c1",
161 + "type": "function",
162 + "function": {"name": "f", "arguments": ""},
163 + }
164 + }
165 + },
166 + },
167 + ),
168 + (
169 + ToolCallDeltaEvent,
170 + "tool-call-delta",
171 + {
172 + "index": 0,
173 + "delta": {
174 + "message": {"tool_calls": {"function": {"arguments": "{}"}}}
175 + },
176 + },
177 + ),
178 + (ToolCallEndEvent, "tool-call-end", {"index": 0}),
179 + (
180 + CitationStartEvent,
181 + "citation-start",
182 + {
183 + "index": 0,
184 + "delta": {
185 + "message": {
186 + "citations": {"start": 0, "end": 5, "text": "hello"}
187 + }
188 + },
189 + },
190 + ),
191 + (CitationEndEvent, "citation-end", {"index": 0}),
192 + (
193 + MessageEndEvent,
194 + "message-end",
195 + {"id": "a", "delta": {"finish_reason": "COMPLETE"}},
196 + ),
197 + ],
198 + )
199 + def test_type_field_default(self, cls, expected_type, kwargs):
200 + ev = cls(**kwargs)
201 + # type field is auto-populated from the Literal default.
202 + assert ev.type == expected_type
203 + # The discriminator must be present in the serialized payload so
204 + # clients reading the stream can demux on it.
205 + dumped = ev.model_dump(exclude_none=True)
206 + assert dumped["type"] == expected_type
207 + # Same in JSON form (what ``_emit`` serializes).
208 + assert f'"type":"{expected_type}"' in ev.model_dump_json(exclude_none=True)
209 +
210 + def test_type_field_cannot_be_overridden_to_wrong_value(self):
211 + # Literal types reject any value other than the bake-in default.
212 + with pytest.raises(ValidationError):
213 + MessageStartEvent(
214 + id="a",
215 + delta={"message": {"role": "assistant"}},
216 + type="other", # type: ignore[arg-type]
217 + )
tests/entrypoints/cohere/test_registry_and_args.py new
+91
@@ -0,0 +1,91 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Sanity tests for the small registry / CLI-arg / config additions made
4 +by the Cohere v2 chat API:
5 +
6 +* ``vllm/renderers/registry.py``: new ``"cohere"`` renderer entry.
7 +* ``vllm/tokenizers/registry.py``: new ``"cohere"`` tokenizer entry
8 + (aliased to the cached HF tokenizer).
9 +* ``vllm/entrypoints/openai/cli_args.py``: new
10 + ``cohere_is_reasoning_model`` field.
11 +* ``vllm/config/model.py``: ``"cohere"`` added to ``TokenizerMode``
12 + Literal.
13 +"""
14 +
15 +import dataclasses
16 +import typing
17 +
18 +import pytest
19 +
20 +from vllm.config.model import TokenizerMode
21 +from vllm.entrypoints.openai.cli_args import BaseFrontendArgs, make_arg_parser
22 +from vllm.renderers.registry import RENDERER_REGISTRY
23 +from vllm.tokenizers.registry import TokenizerRegistry
24 +from vllm.utils.argparse_utils import FlexibleArgumentParser
25 +
26 +
27 +class TestRendererRegistry:
28 + def test_cohere_renderer_registered(self):
29 + # The registry resolves to the importable ``CohereRenderer`` class.
30 + cls = RENDERER_REGISTRY.load_renderer_cls("cohere")
31 + assert cls.__name__ == "CohereRenderer"
32 + # Sanity: class lives in the cohere renderer module.
33 + assert cls.__module__ == "vllm.renderers.cohere"
34 +
35 +
36 +class TestTokenizerRegistry:
37 + def test_cohere_aliased_to_cached_hf_tokenizer(self):
38 + # ``cohere`` mode uses the standard HF tokenizer; only the
39 + # renderer stage is replaced. This test guards against accidental
40 + # divergence.
41 + cls = TokenizerRegistry.load_tokenizer_cls("cohere")
42 + assert cls.__name__ == "CachedHfTokenizer"
43 +
44 +
45 +class TestTokenizerModeLiteral:
46 + def test_cohere_is_a_valid_tokenizer_mode(self):
47 + # The Literal must enumerate ``"cohere"`` so engine arg parsing
48 + # accepts ``--tokenizer-mode cohere``.
49 + modes = typing.get_args(TokenizerMode)
50 + assert "cohere" in modes
51 +
52 +
53 +# ----------------------------------------------------------------------
54 +# ``--cohere-is-reasoning-model`` CLI flag
55 +# ----------------------------------------------------------------------
56 +
57 +
58 +class TestCohereCliArg:
59 + """Verifies the new ``--cohere-is-reasoning-model`` flag end-to-end
60 + through ``make_arg_parser`` — mirrors the pattern used in
61 + :mod:`tests.entrypoints.openai.test_cli_args`.
62 + """
63 +
64 + def test_default_value_on_dataclass_is_true(self):
65 + fields = {f.name: f for f in dataclasses.fields(BaseFrontendArgs)}
66 + assert "cohere_is_reasoning_model" in fields
67 + field = fields["cohere_is_reasoning_model"]
68 + assert field.default is True
69 + assert field.type is bool
70 +
71 + @pytest.fixture
72 + def serve_parser(self) -> FlexibleArgumentParser:
73 + parser = FlexibleArgumentParser()
74 + return make_arg_parser(parser)
75 +
76 + def test_default_via_argparse_is_true(self, serve_parser: FlexibleArgumentParser):
77 + # No flag supplied → dataclass default (True) wins.
78 + args = serve_parser.parse_args(["--model", "m"])
79 + assert args.cohere_is_reasoning_model is True
80 +
81 + def test_explicit_false_via_argparse(self, serve_parser: FlexibleArgumentParser):
82 + # Boolean dataclass fields are wired up as ``--flag value`` /
83 + # ``--no-flag`` pairs by FlexibleArgumentParser.
84 + args = serve_parser.parse_args(
85 + ["--model", "m", "--no-cohere-is-reasoning-model"]
86 + )
87 + assert args.cohere_is_reasoning_model is False
88 +
89 + def test_explicit_true_via_argparse(self, serve_parser: FlexibleArgumentParser):
90 + args = serve_parser.parse_args(["--model", "m", "--cohere-is-reasoning-model"])
91 + assert args.cohere_is_reasoning_model is True
tests/entrypoints/cohere/test_serving_conversion.py new
+2084
@@ -0,0 +1,2084 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Unit tests for the Cohere v2 -> OpenAI request / response conversion
4 +implemented in ``vllm/entrypoints/cohere/serving.py``.
5 +
6 +These cover the pure-Python classmethods so we don't need an engine.
7 +For the instance methods that read ``self._is_reasoning_model`` we
8 +build a lightweight :class:`_FakeServing` subclass that skips the
9 +heavy ``OpenAIServingChat.__init__`` chain (which would otherwise need
10 +a real engine client, model registry, etc.) — the same pattern used in
11 +``test_serving_streaming.py``.
12 +"""
13 +
14 +from typing import Any
15 +
16 +import pytest
17 +
18 +from vllm.entrypoints.cohere.cohere_chat_message import (
19 + Citation as VLLMCitation,
20 +)
21 +from vllm.entrypoints.cohere.cohere_chat_message import (
22 + CitationSource,
23 +)
24 +from vllm.entrypoints.cohere.protocol import (
25 + CohereChatV2Request,
26 + CohereChatV2Response,
27 +)
28 +from vllm.entrypoints.cohere.serving import (
29 + _FINISH_REASON_MAP,
30 + CohereServingChatV2,
31 + _map_finish_reason,
32 +)
33 +from vllm.entrypoints.openai.chat_completion.protocol import (
34 + ChatCompletionRequest,
35 + ChatCompletionResponse,
36 +)
37 +
38 +# ----------------------------------------------------------------------
39 +# Helpers
40 +# ----------------------------------------------------------------------
41 +
42 +
43 +def _make_request(**kwargs) -> CohereChatV2Request:
44 + kwargs.setdefault("model", "m")
45 + kwargs.setdefault("messages", [{"role": "user", "content": "hi"}])
46 + return CohereChatV2Request(**kwargs)
47 +
48 +
49 +def _convert(request: CohereChatV2Request) -> ChatCompletionRequest:
50 + return CohereServingChatV2._convert_v2_to_chat_completion(request)
51 +
52 +
53 +class _FakeServing(CohereServingChatV2):
54 + """Lightweight stand-in for :class:`CohereServingChatV2` that skips
55 + the heavy ``OpenAIServingChat.__init__`` chain.
56 +
57 + Only ``_is_reasoning_model`` is read by the methods under test
58 + (``_chat_completion_to_v2`` and friends); the rest of the parent
59 + state is dead weight for unit testing.
60 + """
61 +
62 + def __init__(self, is_reasoning_model: bool = True) -> None:
63 + # Intentionally skipping super().__init__ — see class docstring.
64 + self._is_reasoning_model = is_reasoning_model
65 +
66 +
67 +def _serving(is_reasoning_model: bool = True) -> CohereServingChatV2:
68 + return _FakeServing(is_reasoning_model=is_reasoning_model)
69 +
70 +
71 +def _build_chat_completion_response(
72 + *,
73 + response_id: str = "resp_1",
74 + content: str | None = "hello",
75 + reasoning: str | None = None,
76 + tool_calls: list[dict[str, Any]] | None = None,
77 + finish_reason: str | None = "stop",
78 + citations: list[Any] | None = None,
79 + usage: dict[str, Any] | None = None,
80 + kv_transfer_params: dict[str, Any] | None = None,
81 +) -> ChatCompletionResponse:
82 + message: dict[str, Any] = {"role": "assistant"}
83 + if content is not None:
84 + message["content"] = content
85 + if reasoning is not None:
86 + message["reasoning"] = reasoning
87 + if tool_calls is not None:
88 + message["tool_calls"] = tool_calls
89 + if citations is not None:
90 + message["citations"] = citations
91 + kwargs: dict[str, Any] = dict(
92 + id=response_id,
93 + object="chat.completion",
94 + created=0,
95 + model="m",
96 + choices=[{"index": 0, "message": message, "finish_reason": finish_reason}],
97 + # ``usage`` is a required field on ChatCompletionResponse, but the
98 + # production code defensively handles ``None`` -> no usage block;
99 + # we round-trip that behavior by post-setting the attribute below.
100 + usage={"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0},
101 + )
102 + if kv_transfer_params is not None:
103 + kwargs["kv_transfer_params"] = kv_transfer_params
104 + resp = ChatCompletionResponse(**kwargs)
105 + if usage is None:
106 + resp.usage = None
107 + else:
108 + # Replace the placeholder with the caller-provided usage.
109 + resp = ChatCompletionResponse.model_validate(
110 + {**resp.model_dump(), "usage": usage}
111 + )
112 + if kv_transfer_params is not None:
113 + resp.kv_transfer_params = kv_transfer_params
114 + return resp
115 +
116 +
117 +# ======================================================================
118 +# _map_finish_reason
119 +# ======================================================================
120 +
121 +
122 +class TestMapFinishReason:
123 + @pytest.mark.parametrize(
124 + "openai, cohere",
125 + [
126 + ("stop", "COMPLETE"),
127 + ("length", "MAX_TOKENS"),
128 + ("tool_calls", "TOOL_CALL"),
129 + ("stop_sequence", "STOP_SEQUENCE"),
130 + ("error", "ERROR"),
131 + (None, "COMPLETE"),
132 + ],
133 + )
134 + def test_known_reasons(self, openai, cohere):
135 + assert _map_finish_reason(openai) == cohere
136 +
137 + def test_unknown_reason_defaults_to_complete(self):
138 + assert _map_finish_reason("not_a_real_reason") == "COMPLETE"
139 +
140 + def test_finish_reason_map_is_complete(self):
141 + # Sanity check that the lookup table covers all documented states.
142 + assert set(_FINISH_REASON_MAP) == {
143 + "stop",
144 + "length",
145 + "tool_calls",
146 + "stop_sequence",
147 + "error",
148 + None,
149 + }
150 +
151 +
152 +# ======================================================================
153 +# _coerce_text_content (system / tool string fallback)
154 +# ======================================================================
155 +
156 +
157 +class TestCoerceTextContent:
158 + def test_string_passthrough(self):
159 + assert CohereServingChatV2._coerce_text_content("hi") == "hi"
160 +
161 + def test_concatenates_text_blocks(self):
162 + from cohere.types import SystemChatMessageV2
163 +
164 + sys_msg = SystemChatMessageV2(
165 + content=[
166 + {"type": "text", "text": "a"},
167 + {"type": "text", "text": "b"},
168 + ]
169 + )
170 + assert CohereServingChatV2._coerce_text_content(sys_msg.content) == "ab"
171 +
172 +
173 +# ======================================================================
174 +# User message conversion
175 +# ======================================================================
176 +
177 +
178 +class TestConvertUserMessage:
179 + def test_string_content(self):
180 + req = _make_request(messages=[{"role": "user", "content": "hi"}])
181 + result = _convert(req)
182 + assert result.messages == [{"role": "user", "content": "hi"}]
183 +
184 + def test_text_only_list_flattened_to_string(self):
185 + # Single-text-block list is flattened back to a string for maximum
186 + # downstream-template compatibility.
187 + req = _make_request(
188 + messages=[
189 + {
190 + "role": "user",
191 + "content": [{"type": "text", "text": "hi"}],
192 + }
193 + ]
194 + )
195 + result = _convert(req)
196 + assert result.messages[0] == {"role": "user", "content": "hi"}
197 +
198 + def test_image_url_content_with_detail(self):
199 + req = _make_request(
200 + messages=[
201 + {
202 + "role": "user",
203 + "content": [
204 + {
205 + "type": "image_url",
206 + "image_url": {
207 + "url": "data:image/png;base64,xxx",
208 + "detail": "high",
209 + },
210 + }
211 + ],
212 + }
213 + ]
214 + )
215 + result = _convert(req)
216 + msg = result.messages[0]
217 + assert msg["role"] == "user"
218 + assert msg["content"] == [
219 + {
220 + "type": "image_url",
221 + "image_url": {
222 + "url": "data:image/png;base64,xxx",
223 + "detail": "high",
224 + },
225 + }
226 + ]
227 +
228 + def test_image_url_without_detail_omits_field(self):
229 + req = _make_request(
230 + messages=[
231 + {
232 + "role": "user",
233 + "content": [
234 + {
235 + "type": "image_url",
236 + "image_url": {"url": "https://x/i.png"},
237 + }
238 + ],
239 + }
240 + ]
241 + )
242 + result = _convert(req)
243 + assert result.messages[0]["content"][0]["image_url"] == {
244 + "url": "https://x/i.png"
245 + }
246 +
247 + def test_text_plus_image_keeps_list(self):
248 + req = _make_request(
249 + messages=[
250 + {
251 + "role": "user",
252 + "content": [
253 + {"type": "text", "text": "describe"},
254 + {
255 + "type": "image_url",
256 + "image_url": {"url": "https://x/i.png"},
257 + },
258 + ],
259 + }
260 + ]
261 + )
262 + result = _convert(req)
263 + content = result.messages[0]["content"]
264 + assert isinstance(content, list)
265 + assert len(content) == 2
266 + assert content[0] == {"type": "text", "text": "describe"}
267 + assert content[1]["type"] == "image_url"
268 +
269 +
270 +# ======================================================================
271 +# Assistant message conversion
272 +# ======================================================================
273 +
274 +
275 +class TestConvertAssistantMessage:
276 + def test_string_content(self):
277 + req = _make_request(
278 + messages=[
279 + {"role": "user", "content": "hi"},
280 + {"role": "assistant", "content": "hello"},
281 + ]
282 + )
283 + result = _convert(req)
284 + asst = result.messages[1]
285 + assert asst == {"role": "assistant", "content": "hello"}
286 +
287 + def test_text_and_thinking_blocks(self):
288 + # ``thinking`` blocks collapse back into the ``reasoning`` field on
289 + # the OpenAI message; ``text`` blocks become ``content``.
290 + req = _make_request(
291 + messages=[
292 + {"role": "user", "content": "hi"},
293 + {
294 + "role": "assistant",
295 + "content": [
296 + {"type": "thinking", "thinking": "let me think"},
297 + {"type": "text", "text": "Hi!"},
298 + ],
299 + },
300 + ]
301 + )
302 + result = _convert(req)
303 + asst = result.messages[1]
304 + assert asst["role"] == "assistant"
305 + assert asst["content"] == "Hi!"
306 + assert asst["reasoning"] == "let me think"
307 +
308 + def test_thinking_only(self):
309 + # Thinking-only assistant messages have no ``content`` set.
310 + req = _make_request(
311 + messages=[
312 + {"role": "user", "content": "hi"},
313 + {
314 + "role": "assistant",
315 + "content": [
316 + {"type": "thinking", "thinking": "ponder"},
317 + ],
318 + },
319 + ]
320 + )
321 + result = _convert(req)
322 + asst = result.messages[1]
323 + assert asst.get("reasoning") == "ponder"
324 + assert "content" not in asst
325 +
326 + def test_multiple_thinking_blocks_concatenated(self):
327 + req = _make_request(
328 + messages=[
329 + {"role": "user", "content": "hi"},
330 + {
331 + "role": "assistant",
332 + "content": [
333 + {"type": "thinking", "thinking": "first."},
334 + {"type": "thinking", "thinking": "second."},
335 + {"type": "text", "text": "done."},
336 + ],
337 + },
338 + ]
339 + )
340 + result = _convert(req)
341 + asst = result.messages[1]
342 + assert asst["reasoning"] == "first.second."
343 + assert asst["content"] == "done."
344 +
345 + def test_tool_plan_collapses_into_reasoning(self):
346 + # Cohere's ``tool_plan`` is the older chain-of-thought field; it
347 + # should be appended to ``reasoning`` so the rendered template
348 + # preserves the planning context.
349 + req = _make_request(
350 + messages=[
351 + {"role": "user", "content": "hi"},
352 + {
353 + "role": "assistant",
354 + "content": [
355 + {"type": "text", "text": "I'll call a tool."},
356 + ],
357 + "tool_plan": "plan: use calculator",
358 + "tool_calls": [
359 + {
360 + "id": "c1",
361 + "type": "function",
362 + "function": {"name": "calc", "arguments": "{}"},
363 + }
364 + ],
365 + },
366 + ]
367 + )
368 + result = _convert(req)
369 + asst = result.messages[1]
370 + assert asst["content"] == "I'll call a tool."
371 + assert asst["reasoning"] == "plan: use calculator"
372 + assert asst["tool_calls"][0]["function"] == {
373 + "name": "calc",
374 + "arguments": "{}",
375 + }
376 +
377 + def test_tool_calls_with_missing_function_pieces_get_defaults(self):
378 + # The conversion defends against missing function name/arguments
379 + # by emitting empty string / "{}" defaults so downstream
380 + # validation never sees a None.
381 + req = _make_request(
382 + messages=[
383 + {"role": "user", "content": "hi"},
384 + {
385 + "role": "assistant",
386 + "tool_calls": [
387 + {
388 + "id": "c1",
389 + "type": "function",
390 + "function": {"name": "", "arguments": ""},
391 + }
392 + ],
393 + },
394 + ]
395 + )
396 + result = _convert(req)
397 + tc = result.messages[1]["tool_calls"][0]
398 + assert tc["id"] == "c1"
399 + assert tc["type"] == "function"
400 + assert tc["function"] == {"name": "", "arguments": "{}"}
401 +
402 +
403 +# ======================================================================
404 +# Tool message conversion
405 +# ======================================================================
406 +
407 +
408 +class TestConvertToolMessage:
409 + def _request_with_tool_message(self, content: Any) -> CohereChatV2Request:
410 + return _make_request(
411 + messages=[
412 + {"role": "user", "content": "hi"},
413 + {
414 + "role": "assistant",
415 + "tool_calls": [
416 + {
417 + "id": "c1",
418 + "type": "function",
419 + "function": {"name": "f", "arguments": "{}"},
420 + }
421 + ],
422 + },
423 + {"role": "tool", "tool_call_id": "c1", "content": content},
424 + ]
425 + )
426 +
427 + def test_string_content(self):
428 + req = self._request_with_tool_message("result text")
429 + result = _convert(req)
430 + tool_msg = result.messages[-1]
431 + assert tool_msg == {
432 + "role": "tool",
433 + "tool_call_id": "c1",
434 + "content": "result text",
435 + }
436 +
437 + def test_text_only_list_flattened_to_newline_string(self):
438 + # Text-only tool results are flattened to a single newline-joined
439 + # string for compatibility with standard chat templates.
440 + req = self._request_with_tool_message(
441 + [
442 + {"type": "text", "text": "line 1"},
443 + {"type": "text", "text": "line 2"},
444 + ]
445 + )
446 + result = _convert(req)
447 + tool_msg = result.messages[-1]
448 + assert tool_msg["content"] == "line 1\nline 2"
449 +
450 + def test_with_document_preserves_structured_content(self):
451 + # When documents appear in the tool result, we keep the list shape
452 + # so the Cohere renderer can lift them into grounding sources.
453 + req = self._request_with_tool_message(
454 + [
455 + {"type": "text", "text": "see attachment"},
456 + {
457 + "type": "document",
458 + "document": {"data": {"text": "doc text"}, "id": "d1"},
459 + },
460 + ]
461 + )
462 + result = _convert(req)
463 + tool_msg = result.messages[-1]
464 + assert isinstance(tool_msg["content"], list)
465 + assert tool_msg["content"][0] == {"type": "text", "text": "see attachment"}
466 + assert tool_msg["content"][1] == {
467 + "type": "document",
468 + "document": {"data": {"text": "doc text"}, "id": "d1"},
469 + }
470 +
471 +
472 +# ======================================================================
473 +# System message
474 +# ======================================================================
475 +
476 +
477 +class TestSystemMessage:
478 + def test_system_string(self):
479 + req = _make_request(
480 + messages=[
481 + {"role": "system", "content": "be helpful"},
482 + {"role": "user", "content": "hi"},
483 + ]
484 + )
485 + result = _convert(req)
486 + assert result.messages[0] == {
487 + "role": "system",
488 + "content": "be helpful",
489 + }
490 +
491 + def test_system_text_blocks_concatenated(self):
492 + req = _make_request(
493 + messages=[
494 + {
495 + "role": "system",
496 + "content": [
497 + {"type": "text", "text": "part1 "},
498 + {"type": "text", "text": "part2"},
499 + ],
500 + },
501 + {"role": "user", "content": "hi"},
502 + ]
503 + )
504 + result = _convert(req)
505 + assert result.messages[0]["content"] == "part1 part2"
506 +
507 +
508 +# ======================================================================
509 +# Base ChatCompletionRequest field mapping
510 +# ======================================================================
511 +
512 +
513 +class TestBuildBaseChatCompletion:
514 + def test_sampling_and_limits_mapped(self):
515 + req = _make_request(
516 + max_tokens=128,
517 + stop_sequences=["</s>", "STOP"],
518 + temperature=0.5,
519 + seed=42,
520 + frequency_penalty=0.1,
521 + presence_penalty=0.2,
522 + k=50,
523 + p=0.95,
524 + logprobs=True,
525 + priority=2,
526 + kv_transfer_params={"x": 1},
527 + chat_template_kwargs={"y": 2},
528 + )
529 + result = _convert(req)
530 + assert result.model == "m"
531 + # ``max_tokens`` is deprecated in favor of ``max_completion_tokens``
532 + # but the serving code intentionally sets both for compatibility.
533 + assert result.max_completion_tokens == 128
534 + assert result.stop == ["</s>", "STOP"]
535 + assert result.temperature == 0.5
536 + assert result.seed == 42
537 + assert result.frequency_penalty == 0.1
538 + assert result.presence_penalty == 0.2
539 + assert result.top_k == 50
540 + assert result.top_p == 0.95
541 + assert result.logprobs is True
542 + assert result.priority == 2
543 + assert result.kv_transfer_params == {"x": 1}
544 + # ``chat_template_kwargs`` may be expanded by _apply_cohere_*; the
545 + # base build at least preserves what the caller passed.
546 + assert (result.chat_template_kwargs or {}).get("y") == 2
547 +
548 + def test_priority_defaults_to_zero(self):
549 + # ChatCompletionRequest.priority defaults to 0; ``None`` Cohere
550 + # priority must be coerced rather than passed through.
551 + req = _make_request()
552 + result = _convert(req)
553 + assert result.priority == 0
554 +
555 +
556 +# ======================================================================
557 +# Streaming options
558 +# ======================================================================
559 +
560 +
561 +class TestStreamingOptions:
562 + def test_no_stream_leaves_defaults(self):
563 + result = _convert(_make_request(stream=False))
564 + assert not result.stream
565 + assert result.stream_options is None
566 +
567 + def test_stream_enables_usage_options(self):
568 + # The v2 translator forces ``include_usage=True`` so the
569 + # ``message-end`` event can surface ``billed_units`` / ``tokens``;
570 + # ``continuous_usage_stats`` is intentionally left at its
571 + # ``StreamOptions`` default (False) — Cohere v2 only reports
572 + # usage on the terminal event.
573 + result = _convert(_make_request(stream=True))
574 + assert result.stream is True
575 + assert result.stream_options is not None
576 + assert result.stream_options.include_usage is True
577 +
578 +
579 +# ======================================================================
580 +# Response format
581 +# ======================================================================
582 +
583 +
584 +class TestResponseFormat:
585 + def test_text_is_passthrough(self):
586 + result = _convert(_make_request(response_format={"type": "text"}))
587 + assert result.response_format is None
588 +
589 + def test_json_object(self):
590 + result = _convert(_make_request(response_format={"type": "json_object"}))
591 + assert result.response_format is not None
592 + assert result.response_format.type == "json_object"
593 + assert result.response_format.json_schema is None
594 +
595 + def test_json_schema(self):
596 + schema = {"type": "object", "properties": {"a": {"type": "string"}}}
597 + result = _convert(
598 + _make_request(
599 + response_format={"type": "json_object", "json_schema": schema}
600 + )
601 + )
602 + assert result.response_format is not None
603 + assert result.response_format.type == "json_schema"
604 + assert result.response_format.json_schema is not None
605 + assert result.response_format.json_schema.name == "cohere_v2_json_schema"
606 + # ``JsonSchemaResponseFormat.json_schema`` has alias=``schema`` on
607 + # the Pydantic field, so we observe the value via the serialized
608 + # payload (which is what downstream consumers actually read).
609 + dumped = result.response_format.json_schema.model_dump(exclude_none=True)
610 + assert dumped["json_schema"] == schema
611 +
612 +
613 +# ======================================================================
614 +# Tools / tool_choice
615 +# ======================================================================
616 +
617 +
618 +class TestApplyTools:
619 + def test_no_tools(self):
620 + result = _convert(_make_request())
621 + assert result.tools is None
622 +
623 + def test_basic_tool(self):
624 + result = _convert(
625 + _make_request(
626 + tools=[
627 + {
628 + "type": "function",
629 + "function": {
630 + "name": "calc",
631 + "description": "calculator",
632 + "parameters": {"type": "object"},
633 + },
634 + }
635 + ]
636 + )
637 + )
638 + assert result.tools is not None
639 + assert len(result.tools) == 1
640 + tool = result.tools[0]
641 + assert tool.type == "function"
642 + assert tool.function.name == "calc"
643 + assert tool.function.description == "calculator"
644 + # ``strict`` is an extra attribute on FunctionDefinition (the
645 + # field is only stamped onto the OpenAI tool when strict_tools is
646 + # set on the request). The default path must not set it.
647 + assert getattr(tool.function, "strict", None) is None
648 +
649 + def test_strict_tools_propagates_to_function(self):
650 + result = _convert(
651 + _make_request(
652 + strict_tools=True,
653 + tools=[
654 + {
655 + "type": "function",
656 + "function": {
657 + "name": "calc",
658 + "description": "",
659 + "parameters": {},
660 + },
661 + }
662 + ],
663 + )
664 + )
665 + assert result.tools[0].function.strict is True
666 +
667 +
668 +class TestApplyToolChoice:
669 + def test_required(self):
670 + result = _convert(
671 + _make_request(
672 + tool_choice="REQUIRED",
673 + tools=[
674 + {
675 + "type": "function",
676 + "function": {
677 + "name": "f",
678 + "description": "",
679 + "parameters": {},
680 + },
681 + }
682 + ],
683 + )
684 + )
685 + assert result.tool_choice == "required"
686 +
687 + def test_none(self):
688 + result = _convert(
689 + _make_request(
690 + tool_choice="NONE",
691 + tools=[
692 + {
693 + "type": "function",
694 + "function": {
695 + "name": "f",
696 + "description": "",
697 + "parameters": {},
698 + },
699 + }
700 + ],
701 + )
702 + )
703 + assert result.tool_choice == "none"
704 +
705 + def test_default_to_auto_when_tools_present(self):
706 + # No explicit ``tool_choice`` + tools present → auto, mirroring
707 + # Cohere's documented "free choice" default.
708 + result = _convert(
709 + _make_request(
710 + tools=[
711 + {
712 + "type": "function",
713 + "function": {
714 + "name": "f",
715 + "description": "",
716 + "parameters": {},
717 + },
718 + }
719 + ]
720 + )
721 + )
722 + assert result.tool_choice == "auto"
723 +
724 + def test_no_tools_no_choice_left_unset(self):
725 + # When there are no tools the underlying ChatCompletionRequest
726 + # default applies; we must not stamp ``auto``.
727 + result = _convert(_make_request())
728 + assert result.tool_choice != "auto"
729 +
730 +
731 +# ======================================================================
732 +# Cohere-specific template kwargs forwarding
733 +# ======================================================================
734 +
735 +
736 +class TestApplyCohereTemplateKwargs:
737 + def test_string_documents_wrapped(self):
738 + result = _convert(_make_request(documents=["doc 1", "doc 2"]))
739 + docs = (result.chat_template_kwargs or {}).get("documents")
740 + assert docs == [
741 + {"id": "doc_0", "data": {"text": "doc 1"}},
742 + {"id": "doc_1", "data": {"text": "doc 2"}},
743 + ]
744 +
745 + def test_document_with_explicit_id_preserved(self):
746 + result = _convert(
747 + _make_request(
748 + documents=[
749 + {"id": "custom", "data": {"text": "t"}},
750 + {"data": {"text": "t2"}}, # no id -> synthesized
751 + ]
752 + )
753 + )
754 + docs = result.chat_template_kwargs["documents"]
755 + assert docs[0] == {"id": "custom", "data": {"text": "t"}}
756 + assert docs[1]["id"] == "doc_1"
757 +
758 + def test_safety_mode_normalized_to_lowercase(self):
759 + result = _convert(_make_request(safety_mode="CONTEXTUAL"))
760 + assert result.chat_template_kwargs["safety_mode"] == "contextual"
761 +
762 + def test_citation_options_forwarded_as_dict(self):
763 + result = _convert(_make_request(citation_options={"mode": "accurate"}))
764 + assert result.chat_template_kwargs["citation_options"] == {"mode": "accurate"}
765 +
766 + def test_thinking_forwarded_as_dict(self):
767 + result = _convert(
768 + _make_request(thinking={"type": "enabled", "token_budget": 16})
769 + )
770 + assert result.chat_template_kwargs["thinking"] == {
771 + "type": "enabled",
772 + "token_budget": 16,
773 + }
774 +
775 + def test_existing_chat_template_kwargs_preserved(self):
776 + # User-supplied kwargs should not be clobbered by the v2 fields
777 + # (setdefault semantics).
778 + result = _convert(
779 + _make_request(
780 + chat_template_kwargs={
781 + "safety_mode": "user-explicit",
782 + "extra": "x",
783 + },
784 + safety_mode="CONTEXTUAL",
785 + )
786 + )
787 + assert result.chat_template_kwargs["safety_mode"] == "user-explicit"
788 + assert result.chat_template_kwargs["extra"] == "x"
789 +
790 + def test_no_template_kwargs_when_no_cohere_fields(self):
791 + # Without any of the Cohere-specific fields and no caller-supplied
792 + # kwargs, we must leave ``chat_template_kwargs`` as None so other
793 + # renderers see a clean request.
794 + result = _convert(_make_request())
795 + assert result.chat_template_kwargs is None
796 +
797 +
798 +# ======================================================================
799 +# Message-citation forwarding (chat_template_kwargs["_messages_citations"])
800 +# ======================================================================
801 +
802 +
803 +class TestMessageCitations:
804 + """Covers the ``_messages_citations`` chat_template_kwargs entry
805 + that carries ``AssistantChatMessageV2.citations`` from the request
806 + through to the Cohere renderer. Response-shape ``Citation`` ->
807 + melody ``FilterCitation`` conversion is best-effort: some
808 + melody-side fields have no equivalent on the Cohere wire and get
809 + defaulted.
810 + """
811 +
812 + def test_absent_when_no_assistant_citations(self):
813 + result = _convert(
814 + _make_request(
815 + messages=[
816 + {"role": "user", "content": "hi"},
817 + {"role": "assistant", "content": "hello"},
818 + ]
819 + )
820 + )
821 + assert result.chat_template_kwargs is None or (
822 + "_messages_citations" not in (result.chat_template_kwargs or {})
823 + )
824 +
825 + def test_document_citation_forwarded_by_index(self):
826 + # Assistant at request-index 1 cites the *second* top-level
827 + # document (``doc_x``, position 1). Melody addresses documents
828 + # as ``tool_call_index=0`` (the reserved documents bucket) with
829 + # ``tool_result_indices`` selecting positions inside the
830 + # top-level ``documents`` list.
831 + result = _convert(
832 + _make_request(
833 + documents=[
834 + {"id": "doc_a", "data": {"text": "irrelevant"}},
835 + {"id": "doc_x", "data": {"text": "cited doc"}},
836 + ],
837 + messages=[
838 + {"role": "user", "content": "q"},
839 + {
840 + "role": "assistant",
841 + "content": "a",
842 + "citations": [
843 + {
844 + "start": 0,
845 + "end": 1,
846 + "text": "a",
847 + "sources": [{"type": "document", "id": "doc_x"}],
848 + "type": "TEXT_CONTENT",
849 + }
850 + ],
851 + },
852 + ],
853 + )
854 + )
855 + citations_by_index = result.chat_template_kwargs["_messages_citations"]
856 + assert list(citations_by_index.keys()) == [1]
857 + cite = citations_by_index[1][0]
858 + assert cite["start_index"] == 0
859 + assert cite["end_index"] == 1
860 + assert cite["text"] == "a"
861 + assert cite["is_thinking"] is False
862 + # ``document_ids`` is a parser-*output* field in melody (see
863 + # ``PromptRenderIds`` in src/templating/util.rs); on the input
864 + # side melody expects the id resolved into
865 + # ``tool_result_indices``. Setting a ``document_ids`` list here
866 + # would have no effect at all -- the renderer would still
867 + # render ``</co: 0:[]>`` (no anchor).
868 + assert cite["sources"] == [{"tool_call_index": 0, "tool_result_indices": [1]}]
869 +
870 + def test_document_citation_with_unresolvable_id_drops_citation(self):
871 + # A citation source pointing at an id that doesn't appear in
872 + # the request's ``documents`` array leaves the citation with
873 + # nothing to anchor to. Emitting it with ``sources=[]`` would
874 + # render a malformed ``<co>text</co: :>`` marker, and silently
875 + # misattributing to ``documents[0]`` would be worse -- so we
876 + # drop the whole citation, matching the cohere api's behavior.
877 + # Since this message's only citation was dropped, the message
878 + # index shouldn't appear in the forwarded map at all.
879 + result = _convert(
880 + _make_request(
881 + documents=[{"id": "doc_present", "data": {"text": "x"}}],
882 + messages=[
883 + {"role": "user", "content": "q"},
884 + {
885 + "role": "assistant",
886 + "content": "a",
887 + "citations": [
888 + {
889 + "start": 0,
890 + "end": 1,
891 + "text": "a",
892 + "sources": [{"type": "document", "id": "doc_ghost"}],
893 + "type": "TEXT_CONTENT",
894 + }
895 + ],
896 + },
897 + ],
898 + )
899 + )
900 + forwarded = (result.chat_template_kwargs or {}).get("_messages_citations")
901 + assert not forwarded
902 +
903 + def test_document_citation_resolves_auto_assigned_fallback_id(self):
904 + # Requests may include documents without an explicit ``id``;
905 + # ``_apply_cohere_template_kwargs`` synthesizes ``doc_{idx}``
906 + # ids for them. A citation that targets that fallback id must
907 + # still resolve to the right position.
908 + result = _convert(
909 + _make_request(
910 + documents=[
911 + {"data": {"text": "unnamed 0"}},
912 + {"data": {"text": "unnamed 1"}},
913 + ],
914 + messages=[
915 + {"role": "user", "content": "q"},
916 + {
917 + "role": "assistant",
918 + "content": "a",
919 + "citations": [
920 + {
921 + "start": 0,
922 + "end": 1,
923 + "text": "a",
924 + "sources": [{"type": "document", "id": "doc_1"}],
925 + "type": "TEXT_CONTENT",
926 + }
927 + ],
928 + },
929 + ],
930 + )
931 + )
932 + (cite,) = result.chat_template_kwargs["_messages_citations"][1]
933 + assert cite["sources"] == [{"tool_call_index": 0, "tool_result_indices": [1]}]
934 +
935 + def test_tool_citation_resolves_to_bucket_and_result_index(self):
936 + # ``ToolSource.id`` on the wire is the id of a specific document
937 + # inside a tool result (NOT the tool_call_id -- ``type`` is a
938 + # payload-shape hint, not an id-space discriminator). Here the
939 + # citation points at the second doc in the second tool call's
940 + # results, which melody addresses as bucket ``1`` (first tool
941 + # call, no top-level documents present so it takes bucket 0
942 + # ... wait: buckets are per unique tool_call_id, so with no
943 + # top-level docs, ``call_a`` = 0 and ``call_b`` = 1) and
944 + # ``tool_result_indices=[1]`` (2nd doc inside that bucket).
945 + result = _convert(
946 + _make_request(
947 + messages=[
948 + {"role": "user", "content": "q"},
949 + {
950 + "role": "assistant",
951 + "tool_plan": "plan",
952 + "tool_calls": [
953 + {
954 + "id": "call_a",
955 + "type": "function",
956 + "function": {"name": "f", "arguments": "{}"},
957 + }
958 + ],
959 + },
960 + {
961 + "role": "tool",
962 + "tool_call_id": "call_a",
963 + "content": [
964 + {
965 + "type": "document",
966 + "document": {"id": "res_a0", "data": {"text": "r_a0"}},
967 + },
968 + ],
969 + },
970 + {
971 + "role": "assistant",
972 + "tool_plan": "plan2",
973 + "tool_calls": [
974 + {
975 + "id": "call_b",
976 + "type": "function",
977 + "function": {"name": "g", "arguments": "{}"},
978 + }
979 + ],
980 + },
981 + {
982 + "role": "tool",
983 + "tool_call_id": "call_b",
984 + "content": [
985 + {
986 + "type": "document",
987 + "document": {"id": "res_b0", "data": {"text": "r_b0"}},
988 + },
989 + {
990 + "type": "document",
991 + "document": {"id": "res_b1", "data": {"text": "r_b1"}},
992 + },
993 + ],
994 + },
995 + {
996 + "role": "assistant",
997 + "content": "final",
998 + "citations": [
999 + {
1000 + "start": 0,
1001 + "end": 5,
1002 + "text": "final",
1003 + "sources": [{"type": "tool", "id": "res_b1"}],
1004 + "type": "TEXT_CONTENT",
1005 + }
1006 + ],
1007 + },
1008 + ]
1009 + )
1010 + )
1011 + (cite,) = result.chat_template_kwargs["_messages_citations"][5]
1012 + assert cite["sources"] == [
1013 + {"tool_call_index": 1, "tool_result_indices": [1]},
1014 + ]
1015 +
1016 + def test_tool_citation_shifts_bucket_when_documents_present(self):
1017 + # When the request has a top-level ``documents`` array, melody
1018 + # reserves ``tool_call_index=0`` for it (see
1019 + # ``PromptRenderIds::new`` in src/templating/util.rs), so the
1020 + # first tool-call bucket becomes 1. A regression that forgets
1021 + # this shift would silently attribute tool citations to the
1022 + # documents bucket instead.
1023 + result = _convert(
1024 + _make_request(
1025 + documents=[{"id": "d0", "data": {"text": "x"}}],
1026 + messages=[
1027 + {"role": "user", "content": "q"},
1028 + {
1029 + "role": "assistant",
1030 + "tool_plan": "plan",
1031 + "tool_calls": [
1032 + {
1033 + "id": "call_a",
1034 + "type": "function",
1035 + "function": {"name": "f", "arguments": "{}"},
1036 + }
1037 + ],
1038 + },
1039 + {
1040 + "role": "tool",
1041 + "tool_call_id": "call_a",
1042 + "content": [
1043 + {
1044 + "type": "document",
1045 + "document": {"id": "res_a0", "data": {"text": "r"}},
1046 + },
1047 + ],
1048 + },
1049 + {
1050 + "role": "assistant",
1051 + "content": "a",
1052 + "citations": [
1053 + {
1054 + "start": 0,
1055 + "end": 1,
1056 + "text": "a",
1057 + "sources": [{"type": "tool", "id": "res_a0"}],
1058 + "type": "TEXT_CONTENT",
1059 + }
1060 + ],
1061 + },
1062 + ],
1063 + )
1064 + )
1065 + (cite,) = result.chat_template_kwargs["_messages_citations"][3]
1066 + assert cite["sources"] == [
1067 + {"tool_call_index": 1, "tool_result_indices": [0]},
1068 + ]
1069 +
1070 + def test_multiple_sources_in_same_bucket_aggregate(self):
1071 + # Two citation sources both point at docs inside the same tool
1072 + # result. Melody's ``<co>text</co: N:[i,j]>`` marker packs them
1073 + # into a single ``Source`` with a list of ``tool_result_indices``
1074 + # rather than two separate ``Source`` entries; the converter
1075 + # groups per-bucket to match that -- matching the cohere api's
1076 + # inbound citation shape.
1077 + result = _convert(
1078 + _make_request(
1079 + messages=[
1080 + {"role": "user", "content": "q"},
1081 + {
1082 + "role": "assistant",
1083 + "tool_plan": "plan",
1084 + "tool_calls": [
1085 + {
1086 + "id": "call_a",
1087 + "type": "function",
1088 + "function": {"name": "f", "arguments": "{}"},
1089 + }
1090 + ],
1091 + },
1092 + {
1093 + "role": "tool",
1094 + "tool_call_id": "call_a",
1095 + "content": [
1096 + {
1097 + "type": "document",
1098 + "document": {"id": "res0", "data": {"text": "0"}},
1099 + },
1100 + {
1101 + "type": "document",
1102 + "document": {"id": "res1", "data": {"text": "1"}},
1103 + },
1104 + ],
1105 + },
1106 + {
1107 + "role": "assistant",
1108 + "content": "ans",
1109 + "citations": [
1110 + {
1111 + "start": 0,
1112 + "end": 3,
1113 + "text": "ans",
1114 + "sources": [
1115 + {"type": "tool", "id": "res0"},
1116 + {"type": "tool", "id": "res1"},
1117 + ],
1118 + "type": "TEXT_CONTENT",
1119 + }
1120 + ],
1121 + },
1122 + ]
1123 + )
1124 + )
1125 + (cite,) = result.chat_template_kwargs["_messages_citations"][3]
1126 + assert cite["sources"] == [
1127 + {"tool_call_index": 0, "tool_result_indices": [0, 1]},
1128 + ]
1129 +
1130 + def test_tool_citation_with_unresolvable_id_drops_citation(self):
1131 + # Same policy as ``test_document_citation_with_unresolvable_id
1132 + # _drops_citation``: any unresolvable source id drops the whole
1133 + # citation rather than emitting a malformed
1134 + # ``<co>text</co: :>`` marker.
1135 + result = _convert(
1136 + _make_request(
1137 + messages=[
1138 + {"role": "user", "content": "q"},
1139 + {
1140 + "role": "assistant",
1141 + "content": "a",
1142 + "citations": [
1143 + {
1144 + "start": 0,
1145 + "end": 1,
1146 + "text": "a",
1147 + "sources": [{"type": "tool", "id": "ghost_doc"}],
1148 + "type": "TEXT_CONTENT",
1149 + }
1150 + ],
1151 + },
1152 + ]
1153 + )
1154 + )
1155 + forwarded = (result.chat_template_kwargs or {}).get("_messages_citations")
1156 + assert not forwarded
1157 +
1158 + def test_partial_unresolvable_source_drops_whole_citation(self):
1159 + # If any source in a citation fails to resolve, drop the entire
1160 + # citation -- mixing resolved and unresolved sources produces a
1161 + # partially-correct ``<co>`` marker, which is worse than none.
1162 + # This matches the cohere api's behavior.
1163 + result = _convert(
1164 + _make_request(
1165 + documents=[{"id": "doc_real", "data": {"text": "x"}}],
1166 + messages=[
1167 + {"role": "user", "content": "q"},
1168 + {
1169 + "role": "assistant",
1170 + "content": "a",
1171 + "citations": [
1172 + {
1173 + "start": 0,
1174 + "end": 1,
1175 + "text": "a",
1176 + "sources": [
1177 + {"type": "document", "id": "doc_real"},
1178 + {"type": "document", "id": "doc_ghost"},
1179 + ],
1180 + "type": "TEXT_CONTENT",
1181 + }
1182 + ],
1183 + },
1184 + ],
1185 + )
1186 + )
1187 + forwarded = (result.chat_template_kwargs or {}).get("_messages_citations")
1188 + assert not forwarded
1189 +
1190 + def test_thinking_content_type_maps_to_is_thinking(self):
1191 + result = _convert(
1192 + _make_request(
1193 + documents=[{"id": "d", "data": {"text": "x"}}],
1194 + messages=[
1195 + {"role": "user", "content": "q"},
1196 + {
1197 + "role": "assistant",
1198 + "content": [
1199 + {"type": "thinking", "thinking": "chain of thought"},
1200 + {"type": "text", "text": "answer"},
1201 + ],
1202 + "citations": [
1203 + {
1204 + "start": 0,
1205 + "end": 5,
1206 + "text": "chain",
1207 + "sources": [{"type": "document", "id": "d"}],
1208 + "type": "THINKING_CONTENT",
1209 + }
1210 + ],
1211 + },
1212 + ],
1213 + )
1214 + )
1215 + (cite,) = result.chat_template_kwargs["_messages_citations"][1]
1216 + assert cite["is_thinking"] is True
1217 + # And the source resolved (sanity: is_thinking is meaningful
1218 + # only if the citation actually renders).
1219 + assert cite["sources"] == [{"tool_call_index": 0, "tool_result_indices": [0]}]
1220 +
1221 + def test_plan_type_maps_to_is_thinking(self):
1222 + # PLAN citations sit on tool-plan-style non-reasoning assistant
1223 + # turns; melody treats them like THINKING for template purposes.
1224 + result = _convert(
1225 + _make_request(
1226 + documents=[{"id": "d", "data": {"text": "x"}}],
1227 + messages=[
1228 + {"role": "user", "content": "q"},
1229 + {
1230 + "role": "assistant",
1231 + "tool_plan": "plan",
1232 + "citations": [
1233 + {
1234 + "start": 0,
1235 + "end": 4,
1236 + "text": "plan",
1237 + "sources": [{"type": "document", "id": "d"}],
1238 + "type": "PLAN",
1239 + }
1240 + ],
1241 + },
1242 + ],
1243 + )
1244 + )
1245 + (cite,) = result.chat_template_kwargs["_messages_citations"][1]
1246 + assert cite["is_thinking"] is True
1247 +
1248 + def test_tool_content_string_registers_bucket_but_has_no_doc_ids(self):
1249 + # A string tool result has no citable substructure. The bucket
1250 + # still gets registered (so a later tool call's bucket doesn't
1251 + # slide into its slot), but any citation targeting the string
1252 + # message fails to resolve and the citation is dropped.
1253 + result = _convert(
1254 + _make_request(
1255 + messages=[
1256 + {"role": "user", "content": "q"},
1257 + {
1258 + "role": "assistant",
1259 + "tool_plan": "plan",
1260 + "tool_calls": [
1261 + {
1262 + "id": "call_a",
1263 + "type": "function",
1264 + "function": {"name": "f", "arguments": "{}"},
1265 + }
1266 + ],
1267 + },
1268 + {"role": "tool", "tool_call_id": "call_a", "content": "plain text"},
1269 + {
1270 + "role": "assistant",
1271 + "content": "a",
1272 + "citations": [
1273 + {
1274 + "start": 0,
1275 + "end": 1,
1276 + "text": "a",
1277 + "sources": [{"type": "tool", "id": "call_a"}],
1278 + "type": "TEXT_CONTENT",
1279 + }
1280 + ],
1281 + },
1282 + ]
1283 + )
1284 + )
1285 + forwarded = (result.chat_template_kwargs or {}).get("_messages_citations")
1286 + assert not forwarded
1287 +
1288 + def test_tool_content_all_text_blocks_registers_bucket_but_no_docs(self):
1289 + # Same as the string case, but with a structured content list
1290 + # of pure text blocks. No doc ids are exposed, so any citation
1291 + # targeting the message drops.
1292 + result = _convert(
1293 + _make_request(
1294 + messages=[
1295 + {"role": "user", "content": "q"},
1296 + {
1297 + "role": "assistant",
1298 + "tool_plan": "plan",
1299 + "tool_calls": [
1300 + {
1301 + "id": "call_a",
1302 + "type": "function",
1303 + "function": {"name": "f", "arguments": "{}"},
1304 + }
1305 + ],
1306 + },
1307 + {
1308 + "role": "tool",
1309 + "tool_call_id": "call_a",
1310 + "content": [
1311 + {"type": "text", "text": "part 1"},
1312 + {"type": "text", "text": "part 2"},
1313 + ],
1314 + },
1315 + {
1316 + "role": "assistant",
1317 + "content": "a",
1318 + "citations": [
1319 + {
1320 + "start": 0,
1321 + "end": 1,
1322 + "text": "a",
1323 + "sources": [{"type": "tool", "id": "part 1"}],
1324 + "type": "TEXT_CONTENT",
1325 + }
1326 + ],
1327 + },
1328 + ]
1329 + )
1330 + )
1331 + forwarded = (result.chat_template_kwargs or {}).get("_messages_citations")
1332 + assert not forwarded
1333 +
1334 + def test_tool_content_mixed_text_and_docs_uses_content_array_position(self):
1335 + # Melody advances ``tool_result_index`` for *every* content
1336 + # item (text items push an empty id into the row; see
1337 + # ``push_tool_message_contents`` in melody/src/templating/
1338 + # util.rs). Docs at positions 1 and 3 in a 4-item mixed
1339 + # content must resolve to result_indices 1 and 3, not the
1340 + # "position among documents only" numbering (which would give
1341 + # 0 and 1).
1342 + result = _convert(
1343 + _make_request(
1344 + messages=[
1345 + {"role": "user", "content": "q"},
1346 + {
1347 + "role": "assistant",
1348 + "tool_plan": "plan",
1349 + "tool_calls": [
1350 + {
1351 + "id": "call_a",
1352 + "type": "function",
1353 + "function": {"name": "f", "arguments": "{}"},
1354 + }
1355 + ],
1356 + },
1357 + {
1358 + "role": "tool",
1359 + "tool_call_id": "call_a",
1360 + "content": [
1361 + {"type": "text", "text": "before"},
1362 + {
1363 + "type": "document",
1364 + "document": {"id": "doc1", "data": {"text": "d1"}},
1365 + },
1366 + {"type": "text", "text": "between"},
1367 + {
1368 + "type": "document",
1369 + "document": {"id": "doc2", "data": {"text": "d2"}},
1370 + },
1371 + ],
1372 + },
1373 + {
1374 + "role": "assistant",
1375 + "content": "ans",
1376 + "citations": [
1377 + {
1378 + "start": 0,
1379 + "end": 3,
1380 + "text": "ans",
1381 + "sources": [
1382 + {"type": "tool", "id": "doc1"},
1383 + {"type": "tool", "id": "doc2"},
1384 + ],
1385 + "type": "TEXT_CONTENT",
1386 + }
1387 + ],
1388 + },
1389 + ]
1390 + )
1391 + )
1392 + (cite,) = result.chat_template_kwargs["_messages_citations"][3]
1393 + assert cite["sources"] == [
1394 + {"tool_call_index": 0, "tool_result_indices": [1, 3]},
1395 + ]
1396 +
1397 + def test_same_tool_call_id_across_messages_offsets_result_indices(self):
1398 + # If a client splits one tool call's outputs across multiple
1399 + # tool messages with the same ``tool_call_id``, melody
1400 + # accumulates them into the same bucket (each message's
1401 + # content array is *appended* to the bucket's slot list). The
1402 + # second message's docs must therefore be offset by the first
1403 + # message's content length. Dory's melody adapter relies on
1404 + # exactly this behavior when it flattens each tool-result
1405 + # document into its own melody ``Message``.
1406 + result = _convert(
1407 + _make_request(
1408 + messages=[
1409 + {"role": "user", "content": "q"},
1410 + {
1411 + "role": "assistant",
1412 + "tool_plan": "plan",
1413 + "tool_calls": [
1414 + {
1415 + "id": "call_a",
1416 + "type": "function",
1417 + "function": {"name": "f", "arguments": "{}"},
1418 + }
1419 + ],
1420 + },
1421 + {
1422 + "role": "tool",
1423 + "tool_call_id": "call_a",
1424 + "content": [
1425 + {
1426 + "type": "document",
1427 + "document": {"id": "d0", "data": {"text": "0"}},
1428 + },
1429 + {
1430 + "type": "document",
1431 + "document": {"id": "d1", "data": {"text": "1"}},
1432 + },
1433 + ],
1434 + },
1435 + {
1436 + "role": "tool",
1437 + "tool_call_id": "call_a",
1438 + "content": [
1439 + {
1440 + "type": "document",
1441 + "document": {"id": "d2", "data": {"text": "2"}},
1442 + },
1443 + ],
1444 + },
1445 + {
1446 + "role": "assistant",
1447 + "content": "ans",
1448 + "citations": [
1449 + {
1450 + "start": 0,
1451 + "end": 3,
1452 + "text": "ans",
1453 + "sources": [{"type": "tool", "id": "d2"}],
1454 + "type": "TEXT_CONTENT",
1455 + }
1456 + ],
1457 + },
1458 + ]
1459 + )
1460 + )
1461 + (cite,) = result.chat_template_kwargs["_messages_citations"][4]
1462 + # ``d2`` is at position 0 within its own message's content,
1463 + # but the previous tool message contributed 2 slots, so its
1464 + # result index in bucket 0 is 2.
1465 + assert cite["sources"] == [
1466 + {"tool_call_index": 0, "tool_result_indices": [2]},
1467 + ]
1468 +
1469 +
1470 +# ======================================================================
1471 +# _chat_completion_to_v2 (non-streaming response builder)
1472 +# ======================================================================
1473 +
1474 +
1475 +class TestChatCompletionToV2:
1476 + def test_text_only(self):
1477 + serving = _serving(is_reasoning_model=True)
1478 + resp = _build_chat_completion_response(content="hello")
1479 + v2 = serving._chat_completion_to_v2(resp, _make_request())
1480 + assert isinstance(v2, CohereChatV2Response)
1481 + assert v2.id == "resp_1"
1482 + assert v2.finish_reason == "COMPLETE"
1483 + assert v2.message.role == "assistant"
1484 + assert len(v2.message.content) == 1
1485 + assert v2.message.content[0].type == "text"
1486 + assert v2.message.content[0].text == "hello"
1487 + assert v2.message.tool_calls is None
1488 + assert v2.message.tool_plan is None
1489 + assert v2.usage is None
1490 +
1491 + def test_reasoning_model_keeps_thinking_with_tool_calls(self):
1492 + # Reasoning Command models: ``thinking`` block stays in
1493 + # ``message.content`` and ``tool_plan`` is left unset, even when
1494 + # tool calls are present.
1495 + serving = _serving(is_reasoning_model=True)
1496 + resp = _build_chat_completion_response(
1497 + content="resp text",
1498 + reasoning="thoughts",
1499 + tool_calls=[
1500 + {
1501 + "id": "c1",
1502 + "type": "function",
1503 + "function": {"name": "f", "arguments": "{}"},
1504 + }
1505 + ],
1506 + finish_reason="tool_calls",
1507 + )
1508 + v2 = serving._chat_completion_to_v2(resp, _make_request())
1509 + assert v2.finish_reason == "TOOL_CALL"
1510 + assert v2.message.tool_plan is None
1511 + types = [c.type for c in v2.message.content]
1512 + assert types == ["thinking", "text"]
1513 + assert v2.message.content[0].thinking == "thoughts"
1514 + assert v2.message.content[1].text == "resp text"
1515 + assert v2.message.tool_calls[0].id == "c1"
1516 + assert v2.message.tool_calls[0].function.name == "f"
1517 +
1518 + def test_non_reasoning_model_moves_reasoning_to_tool_plan(self):
1519 + # Older Command models surface reasoning as ``tool_plan`` on tool-
1520 + # call turns; the thinking block should be dropped from content.
1521 + serving = _serving(is_reasoning_model=False)
1522 + resp = _build_chat_completion_response(
1523 + content=None,
1524 + reasoning="plan",
1525 + tool_calls=[
1526 + {
1527 + "id": "c1",
1528 + "type": "function",
1529 + "function": {"name": "f", "arguments": "{}"},
1530 + }
1531 + ],
1532 + finish_reason="tool_calls",
1533 + )
1534 + v2 = serving._chat_completion_to_v2(resp, _make_request())
1535 + assert v2.message.tool_plan == "plan"
1536 + assert v2.message.content is None
1537 + assert v2.message.tool_calls[0].id == "c1"
1538 +
1539 + def test_non_reasoning_model_keeps_thinking_when_no_tool_calls(self):
1540 + # No tool calls => non-reasoning behavior is identical to
1541 + # reasoning behavior; the thinking block stays.
1542 + serving = _serving(is_reasoning_model=False)
1543 + resp = _build_chat_completion_response(
1544 + content="answer", reasoning="plan", tool_calls=None
1545 + )
1546 + v2 = serving._chat_completion_to_v2(resp, _make_request())
1547 + assert v2.message.tool_plan is None
1548 + types = [c.type for c in v2.message.content]
1549 + assert types == ["thinking", "text"]
1550 +
1551 + def test_id_synthesized_when_response_id_missing(self):
1552 + serving = _serving()
1553 + resp = _build_chat_completion_response(content="hi", response_id="")
1554 + v2 = serving._chat_completion_to_v2(resp, _make_request())
1555 + assert v2.id.startswith("chat_")
1556 +
1557 + def test_kv_transfer_params_propagated(self):
1558 + serving = _serving()
1559 + resp = _build_chat_completion_response(
1560 + content="hi", kv_transfer_params={"k": 1}
1561 + )
1562 + v2 = serving._chat_completion_to_v2(resp, _make_request())
1563 + assert v2.kv_transfer_params == {"k": 1}
1564 +
1565 +
1566 +# ======================================================================
1567 +# _build_usage
1568 +# ======================================================================
1569 +
1570 +
1571 +class TestBuildUsage:
1572 + def test_none_passthrough(self):
1573 + resp = _build_chat_completion_response(content="hi")
1574 + # default usage is None
1575 + assert CohereServingChatV2._build_usage(resp) is None
1576 +
1577 + def test_basic_usage(self):
1578 + resp = _build_chat_completion_response(
1579 + content="hi",
1580 + usage={
1581 + "prompt_tokens": 10,
1582 + "completion_tokens": 5,
1583 + "total_tokens": 15,
1584 + },
1585 + )
1586 + usage = CohereServingChatV2._build_usage(resp)
1587 + assert usage is not None
1588 + assert usage.billed_units.input_tokens == 10
1589 + assert usage.billed_units.output_tokens == 5
1590 + assert usage.tokens.input_tokens == 10
1591 + assert usage.tokens.output_tokens == 5
1592 + assert usage.cached_tokens is None
1593 +
1594 + def test_completion_tokens_default_to_zero_when_missing(self):
1595 + resp = _build_chat_completion_response(
1596 + content="hi",
1597 + usage={"prompt_tokens": 10, "completion_tokens": 0, "total_tokens": 10},
1598 + )
1599 + usage = CohereServingChatV2._build_usage(resp)
1600 + assert usage.billed_units.output_tokens == 0
1601 +
1602 + def test_cached_tokens_propagated(self):
1603 + resp = _build_chat_completion_response(
1604 + content="hi",
1605 + usage={
1606 + "prompt_tokens": 10,
1607 + "completion_tokens": 5,
1608 + "total_tokens": 15,
1609 + "prompt_tokens_details": {"cached_tokens": 3},
1610 + },
1611 + )
1612 + usage = CohereServingChatV2._build_usage(resp)
1613 + assert usage.cached_tokens == 3
1614 +
1615 +
1616 +# ======================================================================
1617 +# _extract_citations_if_any
1618 +# ======================================================================
1619 +
1620 +
1621 +class TestExtractCitations:
1622 + """Test the parser-output -> wire-shape citation coercion.
1623 +
1624 + Sources are already fully resolved by the reasoning parser (see
1625 + ``_melody_sources_to_vllm`` in
1626 + ``vllm/reasoning/cohere_command_reasoning_parser.py``, which
1627 + consumes the position map forwarded via ``chat_template_kwargs``).
1628 + The serving layer's remaining responsibilities on the outbound
1629 + path are: coerce to :class:`cohere.types.Citation`, apply the
1630 + ``THINKING_CONTENT`` -> ``PLAN`` rewrite for non-reasoning models,
1631 + and drop citations whose ``sources`` list is entirely empty
1632 + (unattributable anchors). Id-less sources with a payload are
1633 + preserved to match the cohere api's optional-``id`` source shape.
1634 + """
1635 +
1636 + @staticmethod
1637 + def _make_citation(**overrides: Any) -> VLLMCitation:
1638 + base = dict(
1639 + start=0,
1640 + end=11,
1641 + text="Shakespeare",
1642 + sources=[
1643 + CitationSource(
1644 + type="document",
1645 + id="doc_hamlet",
1646 + document={"id": "doc_hamlet", "text": "Hamlet."},
1647 + ),
1648 + ],
1649 + type="TEXT_CONTENT",
1650 + )
1651 + base.update(overrides)
1652 + return VLLMCitation(**base)
1653 +
1654 + def test_none_or_empty_returns_none(self):
1655 + serving = _serving()
1656 + assert (
1657 + serving._extract_citations_if_any(type("M", (), {"citations": None})())
1658 + is None
1659 + )
1660 + assert (
1661 + serving._extract_citations_if_any(type("M", (), {"citations": []})())
1662 + is None
1663 + )
1664 + assert serving._extract_citations_if_any(type("M", (), {})()) is None
1665 +
1666 + def test_resolved_document_source_survives_to_wire(self):
1667 + # A parser-resolved document source flows through unchanged
1668 + # (aside from removing ``None`` fields).
1669 + serving = _serving()
1670 + msg = type("M", (), {})()
1671 + msg.citations = [self._make_citation()]
1672 + out = serving._extract_citations_if_any(msg)
1673 + assert out is not None and len(out) == 1
1674 + wire = out[0].model_dump(exclude_none=True)
1675 + assert wire["type"] == "TEXT_CONTENT"
1676 + assert wire["sources"] == [
1677 + {
1678 + "type": "document",
1679 + "id": "doc_hamlet",
1680 + "document": {"id": "doc_hamlet", "text": "Hamlet."},
1681 + }
1682 + ]
1683 +
1684 + def test_resolved_tool_source_survives_to_wire(self):
1685 + # ``type="tool"`` sources carry ``tool_output``; validated by
1686 + # the SDK's discriminated union.
1687 + serving = _serving()
1688 + msg = type("M", (), {})()
1689 + msg.citations = [
1690 + self._make_citation(
1691 + sources=[
1692 + CitationSource(
1693 + type="tool",
1694 + id="res_a0",
1695 + tool_output={"id": "res_a0", "text": "r"},
1696 + ),
1697 + ],
1698 + )
1699 + ]
1700 + out = serving._extract_citations_if_any(msg)
1701 + assert out is not None
1702 + wire = out[0].model_dump(exclude_none=True)
1703 + assert wire["sources"] == [
1704 + {
1705 + "type": "tool",
1706 + "id": "res_a0",
1707 + "tool_output": {"id": "res_a0", "text": "r"},
1708 + }
1709 + ]
1710 +
1711 + def test_source_without_id_preserved_with_payload(self):
1712 + # Matches the cohere api: text-only tool results (and top-level
1713 + # docs without a client-provided id) surface on the wire
1714 + # without an ``id`` field, but the payload still rides through
1715 + # so the client can see what was cited.
1716 + serving = _serving()
1717 + msg = type("M", (), {})()
1718 + msg.citations = [
1719 + self._make_citation(
1720 + sources=[
1721 + CitationSource(type="tool", tool_output={"content": "hi"}),
1722 + ]
1723 + )
1724 + ]
1725 + out = serving._extract_citations_if_any(msg)
1726 + assert out is not None
1727 + wire = out[0].model_dump(exclude_none=True)
1728 + assert wire["sources"] == [{"type": "tool", "tool_output": {"content": "hi"}}]
1729 + assert "id" not in wire["sources"][0]
1730 +
1731 + def test_citation_with_no_surviving_sources_dropped(self):
1732 + # If every source in a citation is empty (e.g. parser produced
1733 + # a stray citation with no resolvable positions), the whole
1734 + # citation is dropped so we don't emit unattributed anchors.
1735 + serving = _serving()
1736 + msg = type("M", (), {})()
1737 + msg.citations = [self._make_citation(sources=[])]
1738 + assert serving._extract_citations_if_any(msg) is None
1739 +
1740 + def test_multiple_sources_survive(self):
1741 + serving = _serving()
1742 + msg = type("M", (), {})()
1743 + msg.citations = [
1744 + self._make_citation(
1745 + sources=[
1746 + CitationSource(type="document", id="d0", document={"id": "d0"}),
1747 + CitationSource(type="document", id="d1", document={"id": "d1"}),
1748 + ],
1749 + )
1750 + ]
1751 + out = serving._extract_citations_if_any(msg)
1752 + assert out is not None
1753 + wire = out[0].model_dump(exclude_none=True)
1754 + assert [s["id"] for s in wire["sources"]] == ["d0", "d1"]
1755 +
1756 + def test_plan_type_on_non_reasoning_model(self):
1757 + # ``is_thinking=True`` melody citations arrive tagged as
1758 + # ``THINKING_CONTENT``; on non-reasoning models the reasoning
1759 + # block is surfaced as ``tool_plan`` (see
1760 + # ``_chat_completion_to_v2``), so citations on it are
1761 + # ``PLAN`` on the wire.
1762 + serving = _serving(is_reasoning_model=False)
1763 + msg = type("M", (), {})()
1764 + msg.citations = [self._make_citation(type="THINKING_CONTENT")]
1765 + out = serving._extract_citations_if_any(msg)
1766 + assert out is not None
1767 + assert out[0].type == "PLAN"
1768 +
1769 + def test_thinking_type_preserved_on_reasoning_model(self):
1770 + serving = _serving(is_reasoning_model=True)
1771 + msg = type("M", (), {})()
1772 + msg.citations = [self._make_citation(type="THINKING_CONTENT")]
1773 + out = serving._extract_citations_if_any(msg)
1774 + assert out is not None
1775 + assert out[0].type == "THINKING_CONTENT"
1776 +
1777 + def test_dict_shape_from_streaming_pipeline(self):
1778 + # The streaming path receives citations as dicts (because
1779 + # ``DeltaMessage.citations`` is an untyped extras field on the
1780 + # OpenAI wire protocol). ``_to_wire_citation`` must accept
1781 + # both dict and object shapes -- unary uses objects.
1782 + serving = _serving()
1783 + msg = type("M", (), {})()
1784 + msg.citations = [
1785 + {
1786 + "start": 0,
1787 + "end": 3,
1788 + "text": "abc",
1789 + "sources": [
1790 + {"type": "document", "id": "d0", "document": {"id": "d0"}},
1791 + ],
1792 + "type": "TEXT_CONTENT",
1793 + }
1794 + ]
1795 + out = serving._extract_citations_if_any(msg)
1796 + assert out is not None
1797 + assert out[0].sources[0].id == "d0"
1798 +
1799 +
1800 +# ======================================================================
1801 +# _build_position_to_source
1802 +# ======================================================================
1803 +
1804 +
1805 +class TestBuildPositionToSource:
1806 + """Pin the outbound numbering invariant.
1807 +
1808 + ``_build_position_to_source`` inverts the same numbering rule that
1809 + ``_build_doc_id_to_prompt_position`` follows on the inbound path
1810 + (melody's ``PromptRenderIds::from_messages`` -- see
1811 + ``melody/src/templating/util.rs``). These tests exist to catch
1812 + numbering drift between the two helpers, which would cause
1813 + outbound citations to attribute to the wrong document.
1814 + """
1815 +
1816 + def test_top_level_documents_populate_bucket_zero(self):
1817 + request = _make_request(
1818 + documents=[
1819 + {"id": "d0", "data": {"text": "zero"}},
1820 + {"id": "d1", "data": {"text": "one"}},
1821 + ],
1822 + )
1823 + result = CohereServingChatV2._build_position_to_source(request)
1824 + assert set(result.keys()) == {(0, 0), (0, 1)}
1825 + assert result[(0, 0)].type == "document"
1826 + assert result[(0, 0)].id == "d0"
1827 + assert result[(0, 0)].document == {"id": "d0", "text": "zero"}
1828 + assert result[(0, 1)].id == "d1"
1829 +
1830 + def test_tool_result_documents_start_at_bucket_one_when_docs_present(self):
1831 + request = _make_request(
1832 + documents=[{"id": "d0", "data": {"text": "top"}}],
1833 + messages=[
1834 + {"role": "user", "content": "q"},
1835 + {
1836 + "role": "assistant",
1837 + "tool_plan": "plan",
1838 + "tool_calls": [
1839 + {
1840 + "id": "call_a",
1841 + "type": "function",
1842 + "function": {"name": "f", "arguments": "{}"},
1843 + }
1844 + ],
1845 + },
1846 + {
1847 + "role": "tool",
1848 + "tool_call_id": "call_a",
1849 + "content": [
1850 + {
1851 + "type": "document",
1852 + "document": {"id": "res_a0", "data": {"text": "r"}},
1853 + }
1854 + ],
1855 + },
1856 + ],
1857 + )
1858 + result = CohereServingChatV2._build_position_to_source(request)
1859 + assert set(result.keys()) == {(0, 0), (1, 0)}
1860 + assert result[(1, 0)].type == "tool"
1861 + assert result[(1, 0)].id == "res_a0"
1862 + assert result[(1, 0)].tool_output == {"id": "res_a0", "text": "r"}
1863 +
1864 + def test_tool_call_id_registered_from_assistant_message(self):
1865 + # Bucket assignment must match the cohere api's rule: the
1866 + # first-seen ``tool_call_id`` (whether on the assistant
1867 + # tool_calls or on the tool message) claims the next integer.
1868 + request = _make_request(
1869 + messages=[
1870 + {"role": "user", "content": "q"},
1871 + {
1872 + "role": "assistant",
1873 + "tool_plan": "plan",
1874 + "tool_calls": [
1875 + {
1876 + "id": "first_call",
1877 + "type": "function",
1878 + "function": {"name": "f", "arguments": "{}"},
1879 + },
1880 + {
1881 + "id": "second_call",
1882 + "type": "function",
1883 + "function": {"name": "g", "arguments": "{}"},
1884 + },
1885 + ],
1886 + },
1887 + {
1888 + "role": "tool",
1889 + "tool_call_id": "second_call",
1890 + "content": [
1891 + {
1892 + "type": "document",
1893 + "document": {"id": "res_second", "data": {"text": "s"}},
1894 + }
1895 + ],
1896 + },
1897 + {
1898 + "role": "tool",
1899 + "tool_call_id": "first_call",
1900 + "content": [
1901 + {
1902 + "type": "document",
1903 + "document": {"id": "res_first", "data": {"text": "f"}},
1904 + }
1905 + ],
1906 + },
1907 + ],
1908 + )
1909 + result = CohereServingChatV2._build_position_to_source(request)
1910 + # ``first_call`` was seen first -> bucket 0; ``second_call``
1911 + # -> bucket 1. The subsequent tool messages fill their
1912 + # respective buckets, not the message-order buckets.
1913 + assert result[(0, 0)].id == "res_first"
1914 + assert result[(1, 0)].id == "res_second"
1915 +
1916 + def test_tool_content_text_and_document_blocks_both_populate_map(self):
1917 + # Mixed text + document content: text blocks now surface in
1918 + # the map as id-less ``tool`` sources (matching the cohere
1919 + # api's handling of text-only tool content) so a model
1920 + # citation pointing at the text slot still has a payload to
1921 + # attach.
1922 + request = _make_request(
1923 + messages=[
1924 + {"role": "user", "content": "q"},
1925 + {
1926 + "role": "assistant",
1927 + "tool_plan": "plan",
1928 + "tool_calls": [
1929 + {
1930 + "id": "call_a",
1931 + "type": "function",
1932 + "function": {"name": "f", "arguments": "{}"},
1933 + }
1934 + ],
1935 + },
1936 + {
1937 + "role": "tool",
1938 + "tool_call_id": "call_a",
1939 + "content": [
1940 + {"type": "text", "text": "prelude"},
1941 + {
1942 + "type": "document",
1943 + "document": {"id": "res_x", "data": {"text": "x"}},
1944 + },
1945 + ],
1946 + },
1947 + ],
1948 + )
1949 + result = CohereServingChatV2._build_position_to_source(request)
1950 + assert set(result.keys()) == {(0, 0), (0, 1)}
1951 + assert result[(0, 0)].type == "tool"
1952 + assert result[(0, 0)].id is None
1953 + assert result[(0, 0)].tool_output == {"content": "prelude"}
1954 + assert result[(0, 1)].id == "res_x"
1955 +
1956 + def test_tool_content_json_text_parsed_as_tool_output(self):
1957 + # The cohere api tries JSON first for text tool content; parity
1958 + # here means a JSON-object text block is spread into
1959 + # ``tool_output`` directly instead of nested under ``content``.
1960 + request = _make_request(
1961 + messages=[
1962 + {"role": "user", "content": "q"},
1963 + {
1964 + "role": "assistant",
1965 + "tool_plan": "plan",
1966 + "tool_calls": [
1967 + {
1968 + "id": "call_a",
1969 + "type": "function",
1970 + "function": {"name": "f", "arguments": "{}"},
1971 + }
1972 + ],
1973 + },
1974 + {
1975 + "role": "tool",
1976 + "tool_call_id": "call_a",
1977 + "content": [{"type": "text", "text": '{"title": "hi"}'}],
1978 + },
1979 + ],
1980 + )
1981 + result = CohereServingChatV2._build_position_to_source(request)
1982 + assert result[(0, 0)].tool_output == {"title": "hi"}
1983 +
1984 + def test_top_level_document_without_id_omits_id_but_keeps_payload(self):
1985 + # Client didn't provide a doc id: match cohere's optional-id
1986 + # source shape by leaving ``id=None`` on the wire, but keep
1987 + # the payload so the citation is still meaningful.
1988 + request = _make_request(
1989 + documents=[{"data": {"text": "anon"}}],
1990 + )
1991 + result = CohereServingChatV2._build_position_to_source(request)
1992 + assert result[(0, 0)].type == "document"
1993 + assert result[(0, 0)].id is None
1994 + assert result[(0, 0)].document == {"text": "anon"}
1995 +
1996 + def test_top_level_string_document_has_no_wire_id(self):
1997 + request = _make_request(documents=["raw text"])
1998 + result = CohereServingChatV2._build_position_to_source(request)
1999 + assert result[(0, 0)].id is None
2000 + assert result[(0, 0)].document == {"text": "raw text"}
2001 +
2002 + def test_tool_content_string_synthesizes_idless_source(self):
2003 + # ``ToolChatMessageV2.content`` may be a bare string; the
2004 + # cohere renderer wraps it as a single text block, so we
2005 + # synthesize a matching id-less source for the slot.
2006 + request = _make_request(
2007 + messages=[
2008 + {"role": "user", "content": "q"},
2009 + {
2010 + "role": "assistant",
2011 + "tool_plan": "plan",
2012 + "tool_calls": [
2013 + {
2014 + "id": "call_a",
2015 + "type": "function",
2016 + "function": {"name": "f", "arguments": "{}"},
2017 + }
2018 + ],
2019 + },
2020 + {"role": "tool", "tool_call_id": "call_a", "content": "just text"},
2021 + ],
2022 + )
2023 + result = CohereServingChatV2._build_position_to_source(request)
2024 + assert result[(0, 0)].type == "tool"
2025 + assert result[(0, 0)].id is None
2026 + assert result[(0, 0)].tool_output == {"content": "just text"}
2027 +
2028 + def test_multiple_tool_messages_same_call_id_offset_slots(self):
2029 + # Tool results streamed across multiple messages accumulate
2030 + # into the same bucket, so later docs are offset by the sum
2031 + # of previous messages' content lengths.
2032 + request = _make_request(
2033 + messages=[
2034 + {"role": "user", "content": "q"},
2035 + {
2036 + "role": "assistant",
2037 + "tool_plan": "plan",
2038 + "tool_calls": [
2039 + {
2040 + "id": "call_a",
2041 + "type": "function",
2042 + "function": {"name": "f", "arguments": "{}"},
2043 + }
2044 + ],
2045 + },
2046 + {
2047 + "role": "tool",
2048 + "tool_call_id": "call_a",
2049 + "content": [
2050 + {
2051 + "type": "document",
2052 + "document": {"id": "res_first", "data": {"text": "1"}},
2053 + },
2054 + ],
2055 + },
2056 + {
2057 + "role": "tool",
2058 + "tool_call_id": "call_a",
2059 + "content": [
2060 + {
2061 + "type": "document",
2062 + "document": {"id": "res_second", "data": {"text": "2"}},
2063 + },
2064 + ],
2065 + },
2066 + ],
2067 + )
2068 + result = CohereServingChatV2._build_position_to_source(request)
2069 + assert result[(0, 0)].id == "res_first"
2070 + assert result[(0, 1)].id == "res_second"
2071 +
2072 +
2073 +# ======================================================================
2074 +# create_error_response
2075 +# ======================================================================
2076 +
2077 +
2078 +class TestCreateErrorResponse:
2079 + def test_envelope_uses_400(self):
2080 + serving = _serving()
2081 + err = serving.create_error_response("oops")
2082 + assert err.error.message == "oops"
2083 + assert err.error.code == 400
2084 + assert err.error.type == "bad_request"
tests/entrypoints/cohere/test_serving_streaming.py new
+944
@@ -0,0 +1,944 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Unit tests for the Cohere v2 SSE stream conversion in
4 +``vllm/entrypoints/cohere/serving.py``.
5 +
6 +The stream-translation entry point is
7 +:meth:`CohereServingChatV2._chat_completion_stream_to_v2`, which turns an
8 +async iterable of OpenAI SSE chunks into Cohere's
9 +``message-start → (content|tool-call|citation)* → message-end → [DONE]``
10 +event stream.
11 +
12 +We test the helpers (``_StreamState``, ``_handle_*_delta``, etc.) in
13 +isolation, plus a handful of end-to-end scenarios that exercise the
14 +state machine.
15 +"""
16 +
17 +import json
18 +from collections.abc import AsyncGenerator
19 +from typing import Any
20 +
21 +import pytest
22 +
23 +from vllm.entrypoints.cohere.cohere_chat_message import (
24 + Citation as VLLMCitation,
25 +)
26 +from vllm.entrypoints.cohere.cohere_chat_message import (
27 + CitationSource,
28 +)
29 +from vllm.entrypoints.cohere.protocol import (
30 + CohereChatV2Request,
31 + MessageStartEvent,
32 +)
33 +from vllm.entrypoints.cohere.serving import (
34 + _DONE_FRAME,
35 + CohereServingChatV2,
36 + ContentBlockType,
37 + _emit,
38 + _sse,
39 + _StreamState,
40 +)
41 +from vllm.entrypoints.openai.chat_completion.protocol import (
42 + ChatCompletionStreamResponse,
43 +)
44 +
45 +# ----------------------------------------------------------------------
46 +# Helpers
47 +# ----------------------------------------------------------------------
48 +
49 +
50 +class _FakeServing(CohereServingChatV2):
51 + """Lightweight stand-in for :class:`CohereServingChatV2` that skips
52 + the heavy ``OpenAIServingChat.__init__`` chain (which would need a
53 + real engine client, model registry, render service, etc.).
54 +
55 + Only ``_is_reasoning_model`` is read by the methods under test
56 + (``_chat_completion_stream_to_v2`` and the per-delta handlers); the
57 + rest is dead weight for unit testing.
58 + """
59 +
60 + def __init__(self, is_reasoning_model: bool = True) -> None:
61 + # Intentionally skipping super().__init__ — see class docstring.
62 + self._is_reasoning_model = is_reasoning_model
63 +
64 +
65 +def _serving(is_reasoning_model: bool = True) -> CohereServingChatV2:
66 + return _FakeServing(is_reasoning_model=is_reasoning_model)
67 +
68 +
69 +def _parse_event(frame: str) -> dict[str, Any]:
70 + """Strip the ``data: ... \\n\\n`` wrapper and parse the JSON payload."""
71 + assert frame.startswith("data: ")
72 + assert frame.endswith("\n\n")
73 + return json.loads(frame[len("data: ") : -2])
74 +
75 +
76 +def _make_chunk(
77 + *,
78 + chunk_id: str = "chunk_0",
79 + role: str | None = None,
80 + content: str | None = None,
81 + reasoning: str | None = None,
82 + tool_calls: list[dict[str, Any]] | None = None,
83 + finish_reason: str | None = None,
84 + usage: dict[str, Any] | None = None,
85 + omit_choices: bool = False,
86 + citations: list[Any] | None = None,
87 +) -> str:
88 + """Build the ``data: {...}\\n\\n`` SSE frame the production code
89 + consumes."""
90 + delta: dict[str, Any] = {}
91 + if role is not None:
92 + delta["role"] = role
93 + if content is not None:
94 + delta["content"] = content
95 + if reasoning is not None:
96 + delta["reasoning"] = reasoning
97 + if tool_calls is not None:
98 + delta["tool_calls"] = tool_calls
99 + if citations is not None:
100 + delta["citations"] = citations
101 +
102 + payload: dict[str, Any] = {
103 + "id": chunk_id,
104 + "object": "chat.completion.chunk",
105 + "created": 0,
106 + "model": "m",
107 + }
108 + if not omit_choices:
109 + payload["choices"] = [
110 + {"index": 0, "delta": delta, "finish_reason": finish_reason}
111 + ]
112 + else:
113 + payload["choices"] = []
114 + if usage is not None:
115 + payload["usage"] = usage
116 + # Use ChatCompletionStreamResponse to normalize the payload shape.
117 + chunk = ChatCompletionStreamResponse.model_validate(payload)
118 + return f"data: {chunk.model_dump_json(exclude_none=False)}\n\n"
119 +
120 +
121 +def _make_done() -> str:
122 + return "data: [DONE]\n\n"
123 +
124 +
125 +async def _stream_from(items: list[str]) -> AsyncGenerator[str, None]:
126 + for item in items:
127 + yield item
128 +
129 +
130 +async def _drain(serving: CohereServingChatV2, items: list[str]) -> list[str]:
131 + """Drive ``_chat_completion_stream_to_v2`` over ``items`` and collect
132 + the emitted SSE frames."""
133 + request = CohereChatV2Request(
134 + model="m", messages=[{"role": "user", "content": "hi"}], stream=True
135 + )
136 + gen = serving._chat_completion_stream_to_v2(_stream_from(items), request)
137 + return [frame async for frame in gen]
138 +
139 +
140 +# ======================================================================
141 +# Low-level helpers: _sse, _emit, _DONE_FRAME
142 +# ======================================================================
143 +
144 +
145 +class TestSSEHelpers:
146 + def test_sse_wraps_payload(self):
147 + assert _sse("hello") == "data: hello\n\n"
148 +
149 + def test_done_frame_constant(self):
150 + # cohere-python and Fern-generated clients key off this exact
151 + # sentinel; keep it byte-for-byte stable.
152 + assert _DONE_FRAME == "data: [DONE]\n\n"
153 +
154 + def test_emit_serializes_event_with_type_discriminator(self):
155 + frame = _emit(
156 + MessageStartEvent(id="abc", delta={"message": {"role": "assistant"}})
157 + )
158 + payload = _parse_event(frame)
159 + assert payload["type"] == "message-start"
160 + assert payload["id"] == "abc"
161 + assert payload["delta"] == {"message": {"role": "assistant"}}
162 +
163 +
164 +# ======================================================================
165 +# _StreamState
166 +# ======================================================================
167 +
168 +
169 +class TestStreamState:
170 + def test_content_index_monotonic(self):
171 + st = _StreamState()
172 + assert st.next_content_index() == 0
173 + assert st.next_content_index() == 1
174 + assert st.next_content_index() == 2
175 +
176 + def test_citation_index_separate_from_content_index(self):
177 + st = _StreamState()
178 + st.next_content_index() # 0
179 + st.next_content_index() # 1
180 + # Citation indexing is independent of content indexing.
181 + assert st.next_citation_index() == 0
182 + assert st.next_citation_index() == 1
183 +
184 +
185 +# ======================================================================
186 +# _close_open_blocks
187 +# ======================================================================
188 +
189 +
190 +class TestCloseOpenBlocks:
191 + def test_no_block_open_emits_nothing(self):
192 + serving = _serving()
193 + state = _StreamState()
194 + assert serving._close_open_blocks(state) == []
195 +
196 + def test_text_block_emits_content_end(self):
197 + serving = _serving()
198 + state = _StreamState()
199 + state.active_block = ContentBlockType.TEXT
200 + state.active_block_index = 1
201 + out = serving._close_open_blocks(state)
202 + assert len(out) == 1
203 + payload = _parse_event(out[0])
204 + assert payload == {"type": "content-end", "index": 1}
205 + assert state.active_block is None
206 + assert state.active_block_index is None
207 +
208 + def test_thinking_block_emits_content_end(self):
209 + serving = _serving()
210 + state = _StreamState()
211 + state.active_block = ContentBlockType.THINKING
212 + state.active_block_index = 3
213 + out = serving._close_open_blocks(state)
214 + payload = _parse_event(out[0])
215 + assert payload == {"type": "content-end", "index": 3}
216 +
217 + def test_tool_call_block_emits_tool_call_end(self):
218 + serving = _serving()
219 + state = _StreamState()
220 + state.active_block = ContentBlockType.TOOL_CALL
221 + state.active_tool_index = 7
222 + out = serving._close_open_blocks(state)
223 + payload = _parse_event(out[0])
224 + assert payload == {"type": "tool-call-end", "index": 7}
225 + assert state.active_tool_index is None
226 +
227 +
228 +# ======================================================================
229 +# _handle_text_delta
230 +# ======================================================================
231 +
232 +
233 +class TestHandleTextDelta:
234 + def test_opens_block_first_time(self):
235 + serving = _serving()
236 + state = _StreamState()
237 + events = serving._handle_text_delta(state, "Hi")
238 + assert len(events) == 2
239 + start = _parse_event(events[0])
240 + delta = _parse_event(events[1])
241 + assert start["type"] == "content-start"
242 + assert start["index"] == 0
243 + assert start["delta"]["message"]["content"]["type"] == "text"
244 + assert delta["type"] == "content-delta"
245 + assert delta["index"] == 0
246 + assert delta["delta"]["message"]["content"]["text"] == "Hi"
247 + assert state.active_block == ContentBlockType.TEXT
248 + assert state.active_block_index == 0
249 +
250 + def test_continues_block_with_just_delta(self):
251 + serving = _serving()
252 + state = _StreamState()
253 + serving._handle_text_delta(state, "Hi")
254 + events = serving._handle_text_delta(state, " there")
255 + # Only a delta event, no new content-start.
256 + assert len(events) == 1
257 + delta = _parse_event(events[0])
258 + assert delta["type"] == "content-delta"
259 + assert delta["delta"]["message"]["content"]["text"] == " there"
260 +
261 + def test_switches_from_thinking_block(self):
262 + serving = _serving(is_reasoning_model=True)
263 + state = _StreamState()
264 + # Open a thinking block first, then switch to text.
265 + serving._handle_thinking_delta(state, "ponder")
266 + events = serving._handle_text_delta(state, "answer")
267 + types = [_parse_event(ev)["type"] for ev in events]
268 + assert types == ["content-end", "content-start", "content-delta"]
269 + # The text block gets a new index (1), distinct from thinking's 0.
270 + assert _parse_event(events[1])["index"] == 1
271 + assert state.active_block == ContentBlockType.TEXT
272 +
273 +
274 +# ======================================================================
275 +# _handle_thinking_delta
276 +# ======================================================================
277 +
278 +
279 +class TestHandleThinkingDelta:
280 + def test_reasoning_model_opens_thinking_block(self):
281 + serving = _serving(is_reasoning_model=True)
282 + state = _StreamState()
283 + events = serving._handle_thinking_delta(state, "thought")
284 + assert len(events) == 2
285 + start = _parse_event(events[0])
286 + delta = _parse_event(events[1])
287 + assert start["type"] == "content-start"
288 + assert start["delta"]["message"]["content"]["type"] == "thinking"
289 + assert delta["type"] == "content-delta"
290 + assert delta["delta"]["message"]["content"]["thinking"] == "thought"
291 + assert state.active_block == ContentBlockType.THINKING
292 +
293 + def test_reasoning_model_continues_thinking_block(self):
294 + serving = _serving(is_reasoning_model=True)
295 + state = _StreamState()
296 + serving._handle_thinking_delta(state, "first")
297 + events = serving._handle_thinking_delta(state, " more")
298 + assert len(events) == 1
299 + assert _parse_event(events[0])["type"] == "content-delta"
300 +
301 + def test_non_reasoning_model_emits_tool_plan_delta(self):
302 + # Older Command models stream reasoning as ``tool_plan`` deltas;
303 + # no content-start/end pair is emitted.
304 + serving = _serving(is_reasoning_model=False)
305 + state = _StreamState()
306 + events = serving._handle_thinking_delta(state, "planning")
307 + assert len(events) == 1
308 + payload = _parse_event(events[0])
309 + assert payload["type"] == "tool-plan-delta"
310 + assert payload["delta"]["message"]["tool_plan"] == "planning"
311 + # ``tool_plan`` deltas don't claim an active content block.
312 + assert state.active_block is None
313 +
314 + def test_non_reasoning_model_closes_open_text_block(self):
315 + serving = _serving(is_reasoning_model=False)
316 + state = _StreamState()
317 + # Open a text block first.
318 + serving._handle_text_delta(state, "answer")
319 + events = serving._handle_thinking_delta(state, "rethink")
320 + types = [_parse_event(ev)["type"] for ev in events]
321 + assert types == ["content-end", "tool-plan-delta"]
322 +
323 +
324 +# ======================================================================
325 +# _handle_tool_call_deltas
326 +# ======================================================================
327 +
328 +
329 +class TestHandleToolCallDeltas:
330 + def test_new_tool_call_opens_tool_call_start(self):
331 + serving = _serving()
332 + state = _StreamState()
333 + deltas = [
334 + type(
335 + "Delta",
336 + (),
337 + {
338 + "index": 0,
339 + "id": "c1",
340 + "function": type(
341 + "Fn", (), {"name": "calc", "arguments": '{"x":'}
342 + )(),
343 + },
344 + )()
345 + ]
346 + events = serving._handle_tool_call_deltas(state, deltas)
347 + assert len(events) == 1
348 + payload = _parse_event(events[0])
349 + assert payload["type"] == "tool-call-start"
350 + assert payload["index"] == 0
351 + tc = payload["delta"]["message"]["tool_calls"]
352 + assert tc["id"] == "c1"
353 + assert tc["function"]["name"] == "calc"
354 + assert tc["function"]["arguments"] == '{"x":'
355 + assert state.active_block == ContentBlockType.TOOL_CALL
356 + assert state.active_tool_index == 0
357 + assert 0 in state.tool_calls_seen
358 +
359 + def test_subsequent_arguments_emit_delta(self):
360 + serving = _serving()
361 + state = _StreamState()
362 + # First call: start.
363 + first = [
364 + type(
365 + "Delta",
366 + (),
367 + {
368 + "index": 0,
369 + "id": "c1",
370 + "function": type("Fn", (), {"name": "calc", "arguments": ""})(),
371 + },
372 + )()
373 + ]
374 + serving._handle_tool_call_deltas(state, first)
375 + # Second call: same index, additional arguments fragment.
376 + more = [
377 + type(
378 + "Delta",
379 + (),
380 + {
381 + "index": 0,
382 + "id": None,
383 + "function": type("Fn", (), {"name": None, "arguments": "1}"})(),
384 + },
385 + )()
386 + ]
387 + events = serving._handle_tool_call_deltas(state, more)
388 + assert len(events) == 1
389 + payload = _parse_event(events[0])
390 + assert payload["type"] == "tool-call-delta"
391 + assert payload["index"] == 0
392 + assert (
393 + payload["delta"]["message"]["tool_calls"]["function"]["arguments"] == "1}"
394 + )
395 +
396 + def test_new_tool_call_closes_existing_content_block(self):
397 + serving = _serving()
398 + state = _StreamState()
399 + # Open a text block, then start a tool call.
400 + serving._handle_text_delta(state, "I'll call:")
401 + deltas = [
402 + type(
403 + "Delta",
404 + (),
405 + {
406 + "index": 0,
407 + "id": "c1",
408 + "function": type("Fn", (), {"name": "calc", "arguments": "{}"})(),
409 + },
410 + )()
411 + ]
412 + events = serving._handle_tool_call_deltas(state, deltas)
413 + types = [_parse_event(ev)["type"] for ev in events]
414 + assert types == ["content-end", "tool-call-start"]
415 +
416 +
417 +# ======================================================================
418 +# _handle_citation_deltas
419 +# ======================================================================
420 +
421 +
422 +class TestHandleCitationDeltas:
423 + """Test the per-delta emitter. Sources are pre-resolved by the
424 + reasoning parser via the position map forwarded through
425 + ``chat_template_kwargs`` (see :func:`_melody_sources_to_vllm` in
426 + ``vllm/reasoning/cohere_command_reasoning_parser.py``), so the
427 + handler only has to coerce to the SDK wire shape, apply the
428 + ``THINKING_CONTENT`` -> ``PLAN`` rewrite, and drop citations with
429 + no attributable sources.
430 + """
431 +
432 + @staticmethod
433 + def _resolved_source(**overrides: Any) -> CitationSource:
434 + base = dict(
435 + type="document",
436 + id="doc_hamlet",
437 + document={"id": "doc_hamlet", "text": "Hamlet by Shakespeare."},
438 + )
439 + base.update(overrides)
440 + return CitationSource(**base)
441 +
442 + def test_resolved_citation_emits_start_and_end_with_real_id(self):
443 + serving = _serving()
444 + state = _StreamState()
445 + citations = [
446 + VLLMCitation(
447 + start=0,
448 + end=11,
449 + text="Shakespeare",
450 + sources=[self._resolved_source()],
451 + type="TEXT_CONTENT",
452 + )
453 + ]
454 + events = serving._handle_citation_deltas(state, citations)
455 + assert len(events) == 2
456 + start = _parse_event(events[0])
457 + end = _parse_event(events[1])
458 + assert start["type"] == "citation-start"
459 + assert start["index"] == 0
460 + cit_payload = start["delta"]["message"]["citations"]
461 + assert cit_payload["start"] == 0
462 + assert cit_payload["end"] == 11
463 + assert cit_payload["text"] == "Shakespeare"
464 + assert cit_payload["sources"] == [
465 + {
466 + "type": "document",
467 + "id": "doc_hamlet",
468 + "document": {"id": "doc_hamlet", "text": "Hamlet by Shakespeare."},
469 + }
470 + ]
471 + assert end == {"type": "citation-end", "index": 0}
472 +
473 + def test_multiple_sources_survive_to_wire(self):
474 + serving = _serving()
475 + state = _StreamState()
476 + citations = [
477 + VLLMCitation(
478 + start=0,
479 + end=5,
480 + text="works",
481 + sources=[
482 + self._resolved_source(id="d0", document={"id": "d0"}),
483 + self._resolved_source(id="d1", document={"id": "d1"}),
484 + ],
485 + type="TEXT_CONTENT",
486 + )
487 + ]
488 + events = serving._handle_citation_deltas(state, citations)
489 + cit_payload = _parse_event(events[0])["delta"]["message"]["citations"]
490 + ids = [s["id"] for s in cit_payload["sources"]]
491 + assert ids == ["d0", "d1"]
492 +
493 + def test_source_without_id_streamed_with_payload(self):
494 + # Match the cohere api: id-less tool sources (from text-only
495 + # tool results) still ride through with their ``tool_output``
496 + # payload; only the ``id`` field is omitted.
497 + serving = _serving()
498 + state = _StreamState()
499 + citations = [
500 + VLLMCitation(
501 + start=0,
502 + end=5,
503 + text="hello",
504 + sources=[CitationSource(type="tool", tool_output={"content": "hi"})],
505 + )
506 + ]
507 + events = serving._handle_citation_deltas(state, citations)
508 + assert len(events) >= 1
509 + cit_payload = _parse_event(events[0])["delta"]["message"]["citations"]
510 + assert cit_payload["sources"] == [
511 + {"type": "tool", "tool_output": {"content": "hi"}}
512 + ]
513 + assert "id" not in cit_payload["sources"][0]
514 +
515 + def test_citation_with_no_sources_streamed_dropped(self):
516 + # Citations whose source list is entirely empty are still
517 + # dropped so we don't emit an unattributed anchor.
518 + serving = _serving()
519 + state = _StreamState()
520 + citations = [
521 + VLLMCitation(
522 + start=0,
523 + end=5,
524 + text="hello",
525 + sources=[],
526 + )
527 + ]
528 + events = serving._handle_citation_deltas(state, citations)
529 + assert events == []
530 +
531 + def test_plan_rewrite_on_non_reasoning_model(self):
532 + serving = _serving(is_reasoning_model=False)
533 + state = _StreamState()
534 + citations = [
535 + VLLMCitation(
536 + start=0,
537 + end=5,
538 + text="hello",
539 + sources=[self._resolved_source()],
540 + type="THINKING_CONTENT",
541 + )
542 + ]
543 + events = serving._handle_citation_deltas(state, citations)
544 + cit_payload = _parse_event(events[0])["delta"]["message"]["citations"]
545 + assert cit_payload["type"] == "PLAN"
546 +
547 +
548 +# ======================================================================
549 +# _build_message_end_event
550 +# ======================================================================
551 +
552 +
553 +class TestBuildMessageEndEvent:
554 + def test_without_usage(self):
555 + serving = _serving()
556 + frame = serving._build_message_end_event(chunk_id="abc", finish_reason="stop")
557 + payload = _parse_event(frame)
558 + assert payload["type"] == "message-end"
559 + assert payload["id"] == "abc"
560 + assert payload["delta"]["finish_reason"] == "COMPLETE"
561 + assert "usage" not in payload["delta"]
562 +
563 + def test_with_usage(self):
564 + serving = _serving()
565 + chunk = ChatCompletionStreamResponse.model_validate(
566 + {
567 + "id": "abc",
568 + "object": "chat.completion.chunk",
569 + "created": 0,
570 + "model": "m",
571 + "choices": [],
572 + "usage": {
573 + "prompt_tokens": 10,
574 + "completion_tokens": 5,
575 + "total_tokens": 15,
576 + },
577 + }
578 + )
579 + frame = serving._build_message_end_event(
580 + chunk_id="abc", finish_reason="length", usage_chunk=chunk
581 + )
582 + payload = _parse_event(frame)
583 + assert payload["delta"]["finish_reason"] == "MAX_TOKENS"
584 + usage = payload["delta"]["usage"]
585 + assert usage["billed_units"] == {"input_tokens": 10, "output_tokens": 5}
586 + assert usage["tokens"] == {"input_tokens": 10, "output_tokens": 5}
587 + assert "cached_tokens" not in usage
588 +
589 + def test_with_cached_tokens(self):
590 + serving = _serving()
591 + chunk = ChatCompletionStreamResponse.model_validate(
592 + {
593 + "id": "abc",
594 + "object": "chat.completion.chunk",
595 + "created": 0,
596 + "model": "m",
597 + "choices": [],
598 + "usage": {
599 + "prompt_tokens": 10,
600 + "completion_tokens": 5,
601 + "total_tokens": 15,
602 + "prompt_tokens_details": {"cached_tokens": 3},
603 + },
604 + }
605 + )
606 + frame = serving._build_message_end_event(
607 + chunk_id="abc", finish_reason="stop", usage_chunk=chunk
608 + )
609 + payload = _parse_event(frame)
610 + assert payload["delta"]["usage"]["cached_tokens"] == 3
611 +
612 +
613 +# ======================================================================
614 +# End-to-end: _chat_completion_stream_to_v2
615 +# ======================================================================
616 +
617 +
618 +class TestChatCompletionStreamToV2:
619 + """End-to-end stream lifecycle tests."""
620 +
621 + @pytest.mark.asyncio
622 + async def test_text_only_happy_path_emits_full_lifecycle(self):
623 + serving = _serving(is_reasoning_model=True)
624 + items = [
625 + _make_chunk(role="assistant"),
626 + _make_chunk(content="Hi"),
627 + _make_chunk(content=" there"),
628 + _make_chunk(
629 + omit_choices=True,
630 + usage={
631 + "prompt_tokens": 4,
632 + "completion_tokens": 2,
633 + "total_tokens": 6,
634 + },
635 + ),
636 + ]
637 + frames = await _drain(serving, items)
638 + types = [_parse_event(f)["type"] for f in frames[:-1]]
639 + # message-start, content-start, content-delta, content-delta,
640 + # content-end, message-end, then [DONE] as the last frame.
641 + assert types == [
642 + "message-start",
643 + "content-start",
644 + "content-delta",
645 + "content-delta",
646 + "content-end",
647 + "message-end",
648 + ]
649 + assert frames[-1] == _DONE_FRAME
650 + # message-end should carry usage stats from the trailing chunk.
651 + end_payload = _parse_event(frames[-2])
652 + assert end_payload["delta"]["finish_reason"] == "COMPLETE"
653 + assert end_payload["delta"]["usage"]["billed_units"]["input_tokens"] == 4
654 +
655 + @pytest.mark.asyncio
656 + async def test_text_then_tool_call_closes_text_first(self):
657 + serving = _serving(is_reasoning_model=True)
658 + items = [
659 + _make_chunk(role="assistant"),
660 + _make_chunk(content="planning..."),
661 + _make_chunk(
662 + tool_calls=[
663 + {
664 + "index": 0,
665 + "id": "c1",
666 + "type": "function",
667 + "function": {"name": "calc", "arguments": "{}"},
668 + }
669 + ]
670 + ),
671 + _make_chunk(finish_reason="tool_calls"),
672 + _make_chunk(
673 + omit_choices=True,
674 + usage={
675 + "prompt_tokens": 1,
676 + "completion_tokens": 1,
677 + "total_tokens": 2,
678 + },
679 + ),
680 + ]
681 + frames = await _drain(serving, items)
682 + types = [_parse_event(f)["type"] for f in frames[:-1]]
683 + # Text block is opened+delta, then closed before the tool-call
684 + # opens, and the final close happens on the usage chunk path.
685 + assert types == [
686 + "message-start",
687 + "content-start",
688 + "content-delta",
689 + "content-end",
690 + "tool-call-start",
691 + "tool-call-end",
692 + "message-end",
693 + ]
694 + # finish_reason captured from the prior chunk.
695 + end = _parse_event(frames[-2])
696 + assert end["delta"]["finish_reason"] == "TOOL_CALL"
697 +
698 + @pytest.mark.asyncio
699 + async def test_thinking_then_text_reasoning_model(self):
700 + serving = _serving(is_reasoning_model=True)
701 + items = [
702 + _make_chunk(role="assistant"),
703 + _make_chunk(reasoning="thinking..."),
704 + _make_chunk(content="answer"),
705 + _make_chunk(finish_reason="stop"),
706 + _make_chunk(
707 + omit_choices=True,
708 + usage={
709 + "prompt_tokens": 1,
710 + "completion_tokens": 1,
711 + "total_tokens": 2,
712 + },
713 + ),
714 + ]
715 + frames = await _drain(serving, items)
716 + types = [_parse_event(f)["type"] for f in frames[:-1]]
717 + # Thinking block opens with index 0; text block reopens with index 1.
718 + assert types == [
719 + "message-start",
720 + "content-start",
721 + "content-delta",
722 + "content-end",
723 + "content-start",
724 + "content-delta",
725 + "content-end",
726 + "message-end",
727 + ]
728 + thinking_start = _parse_event(frames[1])
729 + assert thinking_start["delta"]["message"]["content"]["type"] == "thinking"
730 + text_start = _parse_event(frames[4])
731 + assert text_start["delta"]["message"]["content"]["type"] == "text"
732 +
733 + @pytest.mark.asyncio
734 + async def test_reasoning_on_non_reasoning_model_emits_tool_plan_delta(self):
735 + serving = _serving(is_reasoning_model=False)
736 + items = [
737 + _make_chunk(role="assistant"),
738 + _make_chunk(reasoning="planning"),
739 + _make_chunk(
740 + tool_calls=[
741 + {
742 + "index": 0,
743 + "id": "c1",
744 + "type": "function",
745 + "function": {"name": "f", "arguments": "{}"},
746 + }
747 + ]
748 + ),
749 + _make_chunk(finish_reason="tool_calls"),
750 + _make_chunk(
751 + omit_choices=True,
752 + usage={
753 + "prompt_tokens": 1,
754 + "completion_tokens": 1,
755 + "total_tokens": 2,
756 + },
757 + ),
758 + ]
759 + frames = await _drain(serving, items)
760 + types = [_parse_event(f)["type"] for f in frames[:-1]]
761 + assert types == [
762 + "message-start",
763 + "tool-plan-delta",
764 + "tool-call-start",
765 + "tool-call-end",
766 + "message-end",
767 + ]
768 + # No thinking content blocks should be present.
769 + assert "content-start" not in types
770 + assert "content-end" not in types
771 +
772 + @pytest.mark.asyncio
773 + async def test_done_marker_in_middle_closes_open_block(self):
774 + # Some upstreams send [DONE] without a trailing usage-only chunk.
775 + # The translator must still emit message-end before [DONE] so
776 + # Cohere clients don't hang.
777 + serving = _serving(is_reasoning_model=True)
778 + items = [
779 + _make_chunk(role="assistant"),
780 + _make_chunk(content="Hi"),
781 + _make_done(),
782 + ]
783 + frames = await _drain(serving, items)
784 + types = [_parse_event(f)["type"] for f in frames[:-1]]
785 + assert types == [
786 + "message-start",
787 + "content-start",
788 + "content-delta",
789 + "content-end",
790 + "message-end",
791 + ]
792 + assert frames[-1] == _DONE_FRAME
793 +
794 + @pytest.mark.asyncio
795 + async def test_skips_empty_and_non_data_lines(self):
796 + serving = _serving()
797 + items = [
798 + "\n",
799 + "event: ping\n\n",
800 + _make_chunk(role="assistant"),
801 + "data: \n\n", # empty data
802 + _make_chunk(content="Hi"),
803 + _make_chunk(
804 + omit_choices=True,
805 + usage={
806 + "prompt_tokens": 1,
807 + "completion_tokens": 1,
808 + "total_tokens": 2,
809 + },
810 + ),
811 + ]
812 + frames = await _drain(serving, items)
813 + assert frames[-1] == _DONE_FRAME
814 + # Should still produce a complete lifecycle.
815 + types = [_parse_event(f)["type"] for f in frames[:-1]]
816 + assert types[0] == "message-start"
817 + assert types[-1] == "message-end"
818 +
819 + @pytest.mark.asyncio
820 + async def test_exception_in_chunk_parsing_emits_error_message_end(self):
821 + # An invalid JSON payload after the first chunk triggers the
822 + # error path: a synthetic message-end with finish_reason=ERROR
823 + # followed by [DONE].
824 + serving = _serving()
825 + items = [
826 + _make_chunk(role="assistant"),
827 + "data: {not valid json}\n\n",
828 + ]
829 + frames = await _drain(serving, items)
830 + assert frames[-1] == _DONE_FRAME
831 + # Find the error-shaped message-end.
832 + error_end = _parse_event(frames[-2])
833 + assert error_end["type"] == "message-end"
834 + assert error_end["delta"]["finish_reason"] == "ERROR"
835 + assert "error" in error_end["delta"]
836 +
837 + @pytest.mark.asyncio
838 + async def test_error_message_end_is_sanitized(self):
839 + # The client-visible ``error`` field on the terminal message-end
840 + # must be routed through ``sanitize_message`` so tracebacks,
841 + # host filesystem paths, and Python object memory addresses
842 + # can't leak from a mid-stream exception into the SSE payload.
843 + serving = _serving()
844 +
845 + async def _raises_mid_stream() -> AsyncGenerator[str, None]:
846 + yield _make_chunk(role="assistant")
847 + raise RuntimeError(
848 + "boom <Foo at 0x7fabc0deface> reading /Users/dev/vllm/x.py"
849 + )
850 +
851 + request = CohereChatV2Request(
852 + model="m",
853 + messages=[{"role": "user", "content": "hi"}],
854 + stream=True,
855 + )
856 + frames = [
857 + f
858 + async for f in serving._chat_completion_stream_to_v2(
859 + _raises_mid_stream(), request
860 + )
861 + ]
862 + assert frames[-1] == _DONE_FRAME
863 + error_end = _parse_event(frames[-2])
864 + assert error_end["delta"]["finish_reason"] == "ERROR"
865 + err_msg = error_end["delta"]["error"]
866 + # Memory address stripped and path replaced with ``<path>``.
867 + assert "0x7fabc0deface" not in err_msg
868 + assert "/Users/dev/vllm/x.py" not in err_msg
869 + assert "<path>" in err_msg
870 + # The human-facing prefix should still survive sanitization.
871 + assert "boom" in err_msg
872 +
873 + @pytest.mark.asyncio
874 + async def test_citations_in_delta_emit_citation_events(self):
875 + # End-to-end streaming citation flow. The parser hands us
876 + # already-resolved sources (see ``_melody_sources_to_vllm``
877 + # in ``vllm/reasoning/cohere_command_reasoning_parser.py``);
878 + # the streaming loop just has to coerce and emit.
879 + serving = _serving()
880 + items = [
881 + _make_chunk(role="assistant"),
882 + _make_chunk(content="hello"),
883 + _make_chunk(
884 + citations=[
885 + {
886 + "start": 0,
887 + "end": 5,
888 + "text": "hello",
889 + "sources": [
890 + {
891 + "type": "document",
892 + "id": "d1",
893 + "document": {"id": "d1", "text": "cited"},
894 + }
895 + ],
896 + "type": "TEXT_CONTENT",
897 + }
898 + ]
899 + ),
900 + _make_chunk(finish_reason="stop"),
901 + _make_chunk(
902 + omit_choices=True,
903 + usage={
904 + "prompt_tokens": 1,
905 + "completion_tokens": 1,
906 + "total_tokens": 2,
907 + },
908 + ),
909 + ]
910 + frames = await _drain(serving, items)
911 + parsed = [_parse_event(f) for f in frames[:-1]]
912 + types = [p["type"] for p in parsed]
913 + assert "citation-start" in types
914 + assert "citation-end" in types
915 + cit_start = next(p for p in parsed if p["type"] == "citation-start")
916 + assert cit_start["delta"]["message"]["citations"]["sources"] == [
917 + {
918 + "type": "document",
919 + "id": "d1",
920 + "document": {"id": "d1", "text": "cited"},
921 + }
922 + ]
923 +
924 + @pytest.mark.asyncio
925 + async def test_first_chunk_emits_message_start_with_chunk_id(self):
926 + serving = _serving()
927 + items = [
928 + _make_chunk(chunk_id="my-id", role="assistant"),
929 + _make_chunk(chunk_id="my-id", content="hi"),
930 + _make_chunk(
931 + chunk_id="my-id",
932 + omit_choices=True,
933 + usage={
934 + "prompt_tokens": 1,
935 + "completion_tokens": 1,
936 + "total_tokens": 2,
937 + },
938 + ),
939 + ]
940 + frames = await _drain(serving, items)
941 + ms = _parse_event(frames[0])
942 + assert ms["type"] == "message-start"
943 + assert ms["id"] == "my-id"
944 + assert ms["delta"]["message"]["role"] == "assistant"
tests/reasoning/test_cohere_command_reasoning_parser.py
+67
@@ -25,6 +25,8 @@ from vllm.reasoning.cohere_command_reasoning_parser import (
25 CohereCommand3ReasoningParser,
26 CohereCommand4ReasoningParser,
27 _has_effective_tools,
28 + _melody_citations_to_vllm,
29 + _melody_sources_to_vllm,
30 _response_format_type,
31 _schema_dict_from_structured_outputs,
32 convert_schema_to_structural_tags,
@@ -623,3 +625,68 @@ class TestAdjustRequestTools:
625 types = _content_types(o.structured_outputs.structural_tag)
626 assert "grammar" in types
627 assert "json_schema" in types
628 +
629 +
630 +class TestMelodySourceResolution:
631 + """Pin the parser-side source resolution (see
632 + ``_melody_sources_to_vllm``). The parser receives melody's numeric
633 + ``(tool_call_index, tool_result_indices)`` addressing and resolves
634 + it against a ``position_to_source`` map handed in by
635 + ``CohereServingChatV2._apply_cohere_template_kwargs`` -- the same
636 + map that would otherwise live in the serving layer's resolver.
637 + """
638 +
639 + @staticmethod
640 + def _fake_melody_source(bucket: int, indices: list[int]) -> Any:
641 + return SimpleNamespace(tool_call_index=bucket, tool_result_indices=indices)
642 +
643 + def test_multi_index_source_fans_out(self):
644 + from vllm.entrypoints.cohere.cohere_chat_message import CitationSource
645 +
646 + position_map: dict[tuple[int, int], CitationSource] = {
647 + (0, 0): CitationSource(type="document", id="d0", document={"id": "d0"}),
648 + (0, 1): CitationSource(type="document", id="d1", document={"id": "d1"}),
649 + }
650 + raw = [self._fake_melody_source(0, [0, 1])]
651 + out = _melody_sources_to_vllm(raw, position_map)
652 + assert [s.id for s in out] == ["d0", "d1"]
653 + # Verify type / payload were plumbed through, not just the id.
654 + assert out[0].type == "document"
655 + assert out[0].document == {"id": "d0"}
656 +
657 + def test_unresolvable_position_skipped(self):
658 + from vllm.entrypoints.cohere.cohere_chat_message import CitationSource
659 +
660 + position_map: dict[tuple[int, int], CitationSource] = {
661 + (0, 0): CitationSource(type="document", id="d0"),
662 + }
663 + raw = [self._fake_melody_source(9, [0])]
664 + assert _melody_sources_to_vllm(raw, position_map) == []
665 +
666 + def test_missing_position_map_drops_all_sources(self):
667 + # A parser instance without a position map (parser wired
668 + # outside of ``CohereServingChatV2``) can't attribute anything,
669 + # so every source is dropped. Callers downstream will see the
670 + # citation with empty ``sources`` and drop it entirely.
671 + raw = [self._fake_melody_source(0, [0])]
672 + assert _melody_sources_to_vllm(raw, None) == []
673 +
674 + def test_citations_pass_through_is_thinking_tag(self):
675 + from vllm.entrypoints.cohere.cohere_chat_message import CitationSource
676 +
677 + position_map: dict[tuple[int, int], CitationSource] = {
678 + (0, 0): CitationSource(type="document", id="d0"),
679 + }
680 + raw = [
681 + SimpleNamespace(
682 + start_index=0,
683 + end_index=5,
684 + text="hello",
685 + is_thinking=True,
686 + sources=[self._fake_melody_source(0, [0])],
687 + )
688 + ]
689 + out = _melody_citations_to_vllm(raw, position_map)
690 + assert out is not None
691 + assert out[0].type == "THINKING_CONTENT"
692 + assert out[0].sources[0].id == "d0"
tests/renderers/test_cohere.py new
+1049
@@ -0,0 +1,1049 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Unit tests for ``vllm/renderers/cohere.py``.
4 +
5 +The tests focus on the pure-Python helpers that produce the render-config
6 +dicts passed to ``cohere_melody.render_cmd3`` / ``render_cmd4``. We also
7 +include a class-level instantiation + async-non-blocking test that
8 +mirrors the analogous ``test_mistral.py`` pattern, exercising the
9 +:class:`CohereRenderer` end-to-end with mocked ``model_config`` /
10 +tokenizer / melody bindings.
11 +"""
12 +
13 +import asyncio
14 +import json
15 +import time
16 +from dataclasses import dataclass
17 +from typing import Any
18 +from unittest.mock import Mock
19 +
20 +import pytest
21 +
22 +from vllm.renderers import ChatParams
23 +from vllm.renderers.cohere import (
24 + CohereRenderer,
25 + MelodyContentType,
26 + _build_render_config,
27 + _content_blocks,
28 + _conversation_to_melody_messages,
29 + _document_to_melody,
30 + _normalize_tool_call,
31 + _role_to_melody,
32 + _tool_to_melody,
33 +)
34 +from vllm.tokenizers.hf import HfTokenizer
35 +
36 +# ======================================================================
37 +# _role_to_melody
38 +# ======================================================================
39 +
40 +
41 +class TestRoleToMelody:
42 + def test_assistant_maps_to_chatbot(self):
43 + # melody's templates use the legacy Cohere ``chatbot`` role name.
44 + assert _role_to_melody("assistant") == "chatbot"
45 +
46 + def test_developer_aliases_to_system(self):
47 + # OpenAI's ``developer`` role is documented as high-priority
48 + # instructions; map it onto the ``system`` slot rather than
49 + # letting the templates drop it on the floor.
50 + assert _role_to_melody("developer") == "system"
51 +
52 + @pytest.mark.parametrize("role", ["user", "system", "tool", "chatbot"])
53 + def test_recognized_roles_passthrough(self, role):
54 + assert _role_to_melody(role) == role
55 +
56 + @pytest.mark.parametrize(
57 + "role,expected",
58 + [
59 + ("ASSISTANT", "chatbot"),
60 + ("Developer", "system"),
61 + ("User", "user"),
62 + ("SYSTEM", "system"),
63 + ],
64 + )
65 + def test_role_normalization_is_case_insensitive(self, role, expected):
66 + # cmd3 / cmd4 templates lowercase the role before matching, so
67 + # accept any casing the caller provides.
68 + assert _role_to_melody(role) == expected
69 +
70 + @pytest.mark.parametrize("role", ["function", "moderator", "", "anything"])
71 + def test_unknown_roles_raise(self, role):
72 + # Silently dropping unknown roles produces malformed prompts
73 + # (the templates' role chain has no else branch).
74 + with pytest.raises(ValueError, match="Unsupported message role"):
75 + _role_to_melody(role)
76 +
77 + def test_non_string_role_rejected(self):
78 + # The function is typed ``role: str`` and the implementation
79 + # relies on Python's attribute lookup (``role.lower()``) to
80 + # reject non-strings — any exception type is acceptable as long
81 + # as we don't silently produce a malformed prompt.
82 + with pytest.raises((AttributeError, TypeError, ValueError)):
83 + _role_to_melody(None) # type: ignore[arg-type]
84 +
85 +
86 +# ======================================================================
87 +# _normalize_tool_call
88 +# ======================================================================
89 +
90 +
91 +class TestNormalizeToolCall:
92 + def test_openai_dict_with_dict_arguments_json_encoded(self):
93 + # melody expects ``parameters`` as a JSON-encoded string even when
94 + # OpenAI delivers an already-parsed dict.
95 + out = _normalize_tool_call(
96 + {
97 + "id": "c1",
98 + "type": "function",
99 + "function": {"name": "f", "arguments": {"a": 1}},
100 + }
101 + )
102 + assert out == {"id": "c1", "name": "f", "parameters": '{"a": 1}'}
103 +
104 + def test_openai_dict_with_string_arguments_preserved(self):
105 + out = _normalize_tool_call(
106 + {
107 + "id": "c1",
108 + "type": "function",
109 + "function": {"name": "f", "arguments": '{"a":1}'},
110 + }
111 + )
112 + assert out["parameters"] == '{"a":1}'
113 +
114 + def test_flat_dict_without_function_wrapper(self):
115 + out = _normalize_tool_call({"id": "c1", "name": "f", "arguments": '{"k": 1}'})
116 + # Falls back to top-level ``name`` / ``arguments``.
117 + assert out == {"id": "c1", "name": "f", "parameters": '{"k": 1}'}
118 +
119 + def test_missing_id_becomes_empty_string(self):
120 + out = _normalize_tool_call({"function": {"name": "f", "arguments": "{}"}})
121 + assert out["id"] == ""
122 +
123 + def test_pydantic_model_dump_supported(self):
124 + class _Fake:
125 + def model_dump(self):
126 + return {
127 + "id": "c1",
128 + "type": "function",
129 + "function": {"name": "f", "arguments": "{}"},
130 + }
131 +
132 + out = _normalize_tool_call(_Fake())
133 + assert out == {"id": "c1", "name": "f", "parameters": "{}"}
134 +
135 + def test_invalid_type_rejected(self):
136 + with pytest.raises(TypeError, match="Unexpected tool_call value"):
137 + _normalize_tool_call(42) # type: ignore[arg-type]
138 +
139 +
140 +# ======================================================================
141 +# _content_blocks
142 +# ======================================================================
143 +
144 +
145 +class TestContentBlocks:
146 + def test_none_returns_empty_list(self):
147 + assert _content_blocks(None) == []
148 +
149 + def test_string_wrapped_in_text_block(self):
150 + out = _content_blocks("hi")
151 + assert out == [{"type": MelodyContentType.TEXT, "text": "hi"}]
152 +
153 + def test_string_item_in_list_wrapped(self):
154 + out = _content_blocks(["a", "b"])
155 + assert out == [
156 + {"type": MelodyContentType.TEXT, "text": "a"},
157 + {"type": MelodyContentType.TEXT, "text": "b"},
158 + ]
159 +
160 + @pytest.mark.parametrize(
161 + "part_type",
162 + ["text", "input_text", "output_text", "refusal"],
163 + )
164 + def test_text_variants_normalized(self, part_type):
165 + out = _content_blocks([{"type": part_type, "text": "hello"}])
166 + assert out == [{"type": MelodyContentType.TEXT, "text": "hello"}]
167 +
168 + def test_thinking_block(self):
169 + out = _content_blocks([{"type": "thinking", "thinking": "thoughts"}])
170 + assert out == [{"type": MelodyContentType.THINKING, "thinking": "thoughts"}]
171 +
172 + def test_image_block_with_default_placeholder(self):
173 + out = _content_blocks([{"type": "image"}])
174 + assert out == [
175 + {
176 + "type": MelodyContentType.IMAGE,
177 + "image": {"template_placeholder": "<image>"},
178 + }
179 + ]
180 +
181 + def test_image_block_custom_placeholder(self):
182 + out = _content_blocks([{"type": "image", "template_placeholder": "[[IMG]]"}])
183 + assert out[0]["image"]["template_placeholder"] == "[[IMG]]"
184 +
185 + def test_document_block_dict_passthrough(self):
186 + out = _content_blocks(
187 + [{"type": "document", "document": {"data": {"text": "doc"}}}]
188 + )
189 + assert out == [
190 + {
191 + "type": MelodyContentType.DOCUMENT,
192 + "document": {"data": {"text": "doc"}},
193 + }
194 + ]
195 +
196 + def test_document_block_with_non_dict_falls_back_to_json_text(self):
197 + out = _content_blocks([{"type": "document", "document": "raw string doc"}])
198 + assert out[0]["type"] == MelodyContentType.TEXT
199 + # JSON-encoded for safety since melody expects a structured doc.
200 + assert out[0]["text"] == json.dumps("raw string doc")
201 +
202 + def test_tool_reference_emitted_as_text(self):
203 + out = _content_blocks([{"type": "tool_reference", "name": "calc"}])
204 + assert out == [{"type": MelodyContentType.TEXT, "text": "calc"}]
205 +
206 + def test_unknown_block_type_fallback_to_text(self):
207 + # Unknown block type with a string value is wrapped in a text block.
208 + out = _content_blocks([{"type": "custom", "custom": "value"}])
209 + assert out == [{"type": MelodyContentType.TEXT, "text": "value"}]
210 +
211 + def test_unknown_block_type_dict_value_json_encoded(self):
212 + out = _content_blocks([{"type": "custom", "custom": {"k": 1}}])
213 + assert out == [{"type": MelodyContentType.TEXT, "text": json.dumps({"k": 1})}]
214 +
215 + def test_non_string_non_dict_part_rejected(self):
216 + with pytest.raises(TypeError, match="Unexpected content part"):
217 + _content_blocks([42]) # type: ignore[list-item]
218 +
219 +
220 +# ======================================================================
221 +# _document_to_melody
222 +# ======================================================================
223 +
224 +
225 +class TestDocumentToMelody:
226 + def test_string_wrapped_in_text_dict(self):
227 + assert _document_to_melody("hello") == {"text": "hello"}
228 +
229 + def test_pure_dict_passthrough(self):
230 + inp = {"text": "x", "id": "d1"}
231 + out = _document_to_melody(inp)
232 + assert out == {"text": "x", "id": "d1"}
233 + # Must be a defensive copy so caller-side mutations of the
234 + # returned dict don't leak back into the input.
235 + out["new_key"] = "value"
236 + assert "new_key" not in inp
237 +
238 + def test_data_wrapper_flattened(self):
239 + # Cohere v2 documents use ``{id, data: {...}}``; melody expects
240 + # the flat shape with ``id`` merged into the payload.
241 + out = _document_to_melody({"id": "d1", "data": {"text": "hello", "title": "t"}})
242 + assert out == {"id": "d1", "text": "hello", "title": "t"}
243 +
244 + def test_data_wrapper_preserves_inner_id(self):
245 + # If the inner ``data`` already has an ``id``, it wins.
246 + out = _document_to_melody({"id": "outer", "data": {"id": "inner", "text": "x"}})
247 + assert out["id"] == "inner"
248 +
249 + def test_invalid_type_rejected(self):
250 + with pytest.raises(TypeError, match="Unsupported document type"):
251 + _document_to_melody(42) # type: ignore[arg-type]
252 +
253 +
254 +# ======================================================================
255 +# _tool_to_melody
256 +# ======================================================================
257 +
258 +
259 +class TestToolToMelody:
260 + def test_openai_wrapper(self):
261 + out = _tool_to_melody(
262 + {
263 + "type": "function",
264 + "function": {
265 + "name": "calc",
266 + "description": "calculate",
267 + "parameters": {"type": "object"},
268 + },
269 + }
270 + )
271 + assert out == {
272 + "name": "calc",
273 + "description": "calculate",
274 + "parameters": {"type": "object"},
275 + }
276 +
277 + def test_flat_dict(self):
278 + out = _tool_to_melody({"name": "calc", "description": "d", "parameters": {}})
279 + assert out["name"] == "calc"
280 + assert out["parameters"] == {}
281 +
282 + def test_pydantic_like_model_dump(self):
283 + class _Fake:
284 + def model_dump(self):
285 + return {
286 + "type": "function",
287 + "function": {
288 + "name": "calc",
289 + "description": "x",
290 + "parameters": {},
291 + },
292 + }
293 +
294 + out = _tool_to_melody(_Fake())
295 + assert out["name"] == "calc"
296 +
297 + def test_missing_description_becomes_empty(self):
298 + out = _tool_to_melody({"name": "calc"})
299 + assert out["description"] == ""
300 + assert out["parameters"] == {}
301 +
302 + def test_invalid_type_rejected(self):
303 + with pytest.raises(TypeError, match="Unsupported tool type"):
304 + _tool_to_melody(42) # type: ignore[arg-type]
305 +
306 +
307 +# ======================================================================
308 +# _conversation_to_melody_messages
309 +# ======================================================================
310 +
311 +
312 +class TestConversationToMelody:
313 + def test_basic_user_assistant_pair(self):
314 + conv = [
315 + {"role": "user", "content": "hi"},
316 + {"role": "assistant", "content": "hello"},
317 + ]
318 + out = _conversation_to_melody_messages(conv) # type: ignore[arg-type]
319 + assert out == [
320 + {
321 + "role": "user",
322 + "content": [{"type": MelodyContentType.TEXT, "text": "hi"}],
323 + "tool_calls": [],
324 + },
325 + {
326 + "role": "chatbot",
327 + "content": [{"type": MelodyContentType.TEXT, "text": "hello"}],
328 + "tool_calls": [],
329 + },
330 + ]
331 +
332 + def test_assistant_reasoning_prepended_as_thinking_block(self):
333 + # ``reasoning`` (or ``reasoning_content``) is prepended as a
334 + # ``thinking`` block on assistant turns, preserving multi-turn
335 + # chain-of-thought across the rendered prompt.
336 + conv = [
337 + {
338 + "role": "assistant",
339 + "content": "answer",
340 + "reasoning": "thoughts",
341 + }
342 + ]
343 + out = _conversation_to_melody_messages(conv) # type: ignore[arg-type]
344 + assert out[0]["content"] == [
345 + {"type": MelodyContentType.THINKING, "thinking": "thoughts"},
346 + {"type": MelodyContentType.TEXT, "text": "answer"},
347 + ]
348 +
349 + def test_assistant_reasoning_content_alias_accepted(self):
350 + conv = [
351 + {
352 + "role": "assistant",
353 + "content": "answer",
354 + "reasoning_content": "thoughts",
355 + }
356 + ]
357 + out = _conversation_to_melody_messages(conv) # type: ignore[arg-type]
358 + assert out[0]["content"][0] == {
359 + "type": MelodyContentType.THINKING,
360 + "thinking": "thoughts",
361 + }
362 +
363 + def test_user_reasoning_ignored(self):
364 + # Only assistant turns get reasoning-as-thinking lifting; user
365 + # turns with a ``reasoning`` key (which shouldn't happen in
366 + # practice) must not produce a phantom thinking block.
367 + conv = [
368 + {
369 + "role": "user",
370 + "content": "hi",
371 + "reasoning": "should be ignored",
372 + }
373 + ]
374 + out = _conversation_to_melody_messages(conv) # type: ignore[arg-type]
375 + assert out[0]["content"] == [{"type": MelodyContentType.TEXT, "text": "hi"}]
376 +
377 + def test_tool_calls_normalized(self):
378 + conv = [
379 + {
380 + "role": "assistant",
381 + "content": "",
382 + "tool_calls": [
383 + {
384 + "id": "c1",
385 + "type": "function",
386 + "function": {"name": "f", "arguments": '{"a":1}'},
387 + }
388 + ],
389 + }
390 + ]
391 + out = _conversation_to_melody_messages(conv) # type: ignore[arg-type]
392 + assert out[0]["tool_calls"] == [
393 + {"id": "c1", "name": "f", "parameters": '{"a":1}'}
394 + ]
395 +
396 + def test_tool_call_id_preserved_on_tool_role(self):
397 + conv = [
398 + {
399 + "role": "tool",
400 + "content": "result",
401 + "tool_call_id": "c1",
402 + }
403 + ]
404 + out = _conversation_to_melody_messages(conv) # type: ignore[arg-type]
405 + assert out[0]["tool_call_id"] == "c1"
406 +
407 + def test_messages_citations_attached_by_index(self):
408 + # ``messages_citations`` is a dict keyed by message index. Only
409 + # the message at the matching index should receive the
410 + # citations; other messages must be unaffected.
411 + conv = [
412 + {"role": "user", "content": "q"},
413 + {"role": "assistant", "content": "a"},
414 + ]
415 + citations = {
416 + 1: [
417 + {
418 + "start_index": 0,
419 + "end_index": 1,
420 + "text": "a",
421 + "sources": [
422 + {
423 + "tool_call_index": 0,
424 + "tool_result_indices": [0],
425 + }
426 + ],
427 + "is_thinking": False,
428 + }
429 + ]
430 + }
431 + out = _conversation_to_melody_messages(conv, citations) # type: ignore[arg-type]
432 + assert "citations" not in out[0]
433 + assert out[1]["citations"] == citations[1]
434 +
435 + def test_messages_citations_none_is_a_no_op(self):
436 + conv = [{"role": "assistant", "content": "a"}]
437 + out = _conversation_to_melody_messages(conv, None) # type: ignore[arg-type]
438 + assert "citations" not in out[0]
439 +
440 + def test_messages_citations_missing_index_is_a_no_op(self):
441 + # A ``messages_citations`` dict whose key doesn't hit any
442 + # message must not attach anything (and must not raise).
443 + conv = [{"role": "assistant", "content": "a"}]
444 + out = _conversation_to_melody_messages(conv, {5: [{"anything": 1}]}) # type: ignore[arg-type]
445 + assert "citations" not in out[0]
446 +
447 +
448 +# ======================================================================
449 +# _build_render_config
450 +# ======================================================================
451 +
452 +
453 +class TestBuildRenderConfig:
454 + def _conv(self):
455 + return [{"role": "user", "content": "hi"}]
456 +
457 + def test_default_format_is_cmd4(self):
458 + # Bare kwargs -> cmd4 (the current Command A+ prompt format).
459 + # Mirrors ``_DEFAULT_FORMAT`` in ``vllm/renderers/cohere.py`` and
460 + # the ``--cohere-format`` CLI default.
461 + fmt, cfg = _build_render_config(self._conv(), {}) # type: ignore[arg-type]
462 + assert fmt == "cmd4"
463 + assert cfg["use_jinja"] is True
464 + assert isinstance(cfg["messages"], list)
465 + # No additional_template_fields when no extra kwargs are set.
466 + assert "additional_template_fields" not in cfg
467 +
468 + def test_explicit_cmd3(self):
469 + fmt, cfg = _build_render_config(self._conv(), {"cohere_format": "cmd3"}) # type: ignore[arg-type]
470 + assert fmt == "cmd3"
471 +
472 + def test_invalid_format_raises(self):
473 + with pytest.raises(ValueError, match="Invalid cohere_format"):
474 + _build_render_config(self._conv(), {"cohere_format": "cmd5"}) # type: ignore[arg-type]
475 +
476 + def test_documents_converted(self):
477 + _, cfg = _build_render_config(
478 + self._conv(),
479 + {
480 + "documents": [
481 + "doc text",
482 + {"id": "d1", "data": {"text": "wrapped"}},
483 + ]
484 + },
485 + ) # type: ignore[arg-type]
486 + assert cfg["documents"] == [
487 + {"text": "doc text"},
488 + {"id": "d1", "text": "wrapped"},
489 + ]
490 +
491 + def test_available_tools_take_precedence_over_tools(self):
492 + _, cfg = _build_render_config(
493 + self._conv(),
494 + {
495 + "tools": [{"type": "function", "function": {"name": "from_tools"}}],
496 + "available_tools": [
497 + {"type": "function", "function": {"name": "preferred"}}
498 + ],
499 + },
500 + ) # type: ignore[arg-type]
501 + names = [t["name"] for t in cfg["available_tools"]]
502 + assert names == ["preferred"]
503 +
504 + def test_tools_used_when_no_available_tools(self):
505 + _, cfg = _build_render_config(
506 + self._conv(),
507 + {"tools": [{"type": "function", "function": {"name": "from_tools"}}]},
508 + ) # type: ignore[arg-type]
509 + assert [t["name"] for t in cfg["available_tools"]] == ["from_tools"]
510 +
511 + @pytest.mark.parametrize("value", ["enabled", "disabled"])
512 + def test_reasoning_type_direct(self, value):
513 + _, cfg = _build_render_config(self._conv(), {"reasoning_type": value}) # type: ignore[arg-type]
514 + assert cfg["reasoning_type"] == value
515 +
516 + def test_thinking_dict_shorthand_resolves_reasoning_type(self):
517 + _, cfg = _build_render_config(self._conv(), {"thinking": {"type": "enabled"}}) # type: ignore[arg-type]
518 + assert cfg["reasoning_type"] == "enabled"
519 +
520 + def test_thinking_shorthand_ignores_unknown_type(self):
521 + _, cfg = _build_render_config(self._conv(), {"thinking": {"type": "auto"}}) # type: ignore[arg-type]
522 + assert "reasoning_type" not in cfg
523 +
524 + def test_dev_instruction_forwarded(self):
525 + _, cfg = _build_render_config(self._conv(), {"dev_instruction": "be brief"}) # type: ignore[arg-type]
526 + assert cfg["dev_instruction"] == "be brief"
527 +
528 + def test_response_format_json_object_sets_json_mode(self):
529 + _, cfg = _build_render_config(
530 + self._conv(), {"response_format": {"type": "json_object"}}
531 + ) # type: ignore[arg-type]
532 + assert cfg["json_mode"] is True
533 + assert "json_schema" not in cfg
534 +
535 + def test_response_format_json_schema_sets_json_schema(self):
536 + schema = {"type": "object"}
537 + _, cfg = _build_render_config(
538 + self._conv(),
539 + {"response_format": {"type": "json_schema", "schema": schema}},
540 + ) # type: ignore[arg-type]
541 + # JSON-encoded for melody (string-only schema field).
542 + assert cfg["json_schema"] == json.dumps(schema)
543 +
544 + def test_response_format_nested_json_schema_unwrapped(self):
545 + # When the SDK shape is ``{type: json_schema, schema: {schema:
546 + # {...}}}``, the inner ``schema`` value is used.
547 + inner = {"type": "object"}
548 + _, cfg = _build_render_config(
549 + self._conv(),
550 + {
551 + "response_format": {
552 + "type": "json_schema",
553 + "schema": {"schema": inner},
554 + }
555 + },
556 + ) # type: ignore[arg-type]
557 + assert cfg["json_schema"] == json.dumps(inner)
558 +
559 + def test_json_schema_kwarg_direct(self):
560 + # Caller can also pass ``json_schema`` directly, both as dict and
561 + # as a pre-stringified value.
562 + _, cfg = _build_render_config(self._conv(), {"json_schema": {"a": 1}}) # type: ignore[arg-type]
563 + assert cfg["json_schema"] == '{"a": 1}'
564 + _, cfg = _build_render_config(
565 + self._conv(), {"json_schema": "raw-string-schema"}
566 + ) # type: ignore[arg-type]
567 + assert cfg["json_schema"] == "raw-string-schema"
568 +
569 + def test_json_mode_kwarg_overrides(self):
570 + _, cfg = _build_render_config(self._conv(), {"json_mode": True}) # type: ignore[arg-type]
571 + assert cfg["json_mode"] is True
572 +
573 + def test_cmd3_safety_mode_lowercased(self):
574 + _, cfg = _build_render_config(
575 + self._conv(),
576 + {"cohere_format": "cmd3", "safety_mode": "CONTEXTUAL"},
577 + ) # type: ignore[arg-type]
578 + assert cfg["safety_mode"] == "contextual"
579 +
580 + def test_cmd3_citation_quality_direct(self):
581 + _, cfg = _build_render_config(
582 + self._conv(),
583 + {"cohere_format": "cmd3", "citation_quality": "ACCURATE"},
584 + ) # type: ignore[arg-type]
585 + assert cfg["citation_quality"] == "accurate"
586 +
587 + def test_cmd3_citation_quality_derived_from_citation_options(self):
588 + # When ``citation_quality`` is unset, ``citation_options.mode`` is
589 + # collapsed to on/off so cmd3's binary toggle has a value.
590 + _, cfg = _build_render_config(
591 + self._conv(),
592 + {"cohere_format": "cmd3", "citation_options": {"mode": "accurate"}},
593 + ) # type: ignore[arg-type]
594 + assert cfg["citation_quality"] == "on"
595 +
596 + _, cfg = _build_render_config(
597 + self._conv(),
598 + {"cohere_format": "cmd3", "citation_options": {"mode": "off"}},
599 + ) # type: ignore[arg-type]
600 + assert cfg["citation_quality"] == "off"
601 +
602 + def test_cmd3_skip_preamble_forwarded(self):
603 + _, cfg = _build_render_config(
604 + self._conv(),
605 + {"cohere_format": "cmd3", "skip_preamble": True},
606 + ) # type: ignore[arg-type]
607 + assert cfg["skip_preamble"] is True
608 +
609 + def test_cmd3_no_grounding_field(self):
610 + # cmd3 should never emit a cmd4-only ``grounding`` field.
611 + _, cfg = _build_render_config(
612 + self._conv(),
613 + {"cohere_format": "cmd3", "grounding": "fast"},
614 + ) # type: ignore[arg-type]
615 + assert "grounding" not in cfg
616 +
617 + @pytest.mark.parametrize(
618 + "raw,expected",
619 + [
620 + ("FAST", "enabled"),
621 + ("ACCURATE", "enabled"),
622 + ("OFF", "disabled"),
623 + ("enabled", "enabled"),
624 + ("disabled", "disabled"),
625 + ("unknown", "unknown"),
626 + ],
627 + )
628 + def test_cmd4_grounding_direct(self, raw, expected):
629 + # melody's cmd4 only accepts ``unknown``/``enabled``/``disabled``,
630 + # so the renderer normalizes any of the v2-facing values into
631 + # that vocab.
632 + _, cfg = _build_render_config(
633 + self._conv(),
634 + {"cohere_format": "cmd4", "grounding": raw},
635 + ) # type: ignore[arg-type]
636 + assert cfg["grounding"] == expected
637 +
638 + @pytest.mark.parametrize(
639 + "mode,expected",
640 + [
641 + ("ACCURATE", "enabled"),
642 + ("FAST", "enabled"),
643 + ("OFF", "disabled"),
644 + ],
645 + )
646 + def test_cmd4_grounding_from_citation_options_mode(self, mode, expected):
647 + _, cfg = _build_render_config(
648 + self._conv(),
649 + {
650 + "cohere_format": "cmd4",
651 + "citation_options": {"mode": mode},
652 + },
653 + ) # type: ignore[arg-type]
654 + assert cfg["grounding"] == expected
655 +
656 + def test_cmd4_grounding_rejects_unknown_value(self):
657 + with pytest.raises(ValueError, match="Unrecognized cmd4 grounding"):
658 + _build_render_config(
659 + self._conv(),
660 + {"cohere_format": "cmd4", "grounding": "foobar"},
661 + ) # type: ignore[arg-type]
662 +
663 + def test_cmd4_platform_instruction(self):
664 + _, cfg = _build_render_config(
665 + self._conv(),
666 + {
667 + "cohere_format": "cmd4",
668 + "platform_instruction": "do this",
669 + },
670 + ) # type: ignore[arg-type]
671 + assert cfg["platform_instruction"] == "do this"
672 +
673 + def test_cmd4_no_safety_mode_field(self):
674 + # cmd4 should never carry cmd3-only ``safety_mode``/``citation_quality``.
675 + _, cfg = _build_render_config(
676 + self._conv(),
677 + {
678 + "cohere_format": "cmd4",
679 + "safety_mode": "contextual",
680 + "citation_quality": "on",
681 + },
682 + ) # type: ignore[arg-type]
683 + assert "safety_mode" not in cfg
684 + assert "citation_quality" not in cfg
685 +
686 + def test_extra_kwargs_become_additional_template_fields(self):
687 + # Anything not in the renderer's consumed-keys set is forwarded
688 + # verbatim under ``additional_template_fields`` so jinja templates
689 + # can resolve ``{{ var }}`` directly.
690 + _, cfg = _build_render_config(
691 + self._conv(),
692 + {
693 + "reasoning_effort": "low",
694 + "my_var": "x",
695 + "documents": ["doc"], # consumed, must NOT leak through
696 + },
697 + ) # type: ignore[arg-type]
698 + extras = cfg["additional_template_fields"]
699 + assert extras == {"reasoning_effort": "low", "my_var": "x"}
700 + # Sanity: the consumed key still produced its dedicated config slot.
701 + assert cfg["documents"] == [{"text": "doc"}]
702 +
703 + def test_template_id_passthrough(self):
704 + # ``template_id`` is a safe selector for one of melody's built-in
705 + # template variants (not raw source) and is still accepted via
706 + # ``chat_template_kwargs``.
707 + _, cfg = _build_render_config(
708 + self._conv(),
709 + {"template_id": "tpl1"},
710 + ) # type: ignore[arg-type]
711 + assert cfg["template_id"] == "tpl1"
712 + # use_jinja is always True, regardless of caller input.
713 + assert cfg["use_jinja"] is True
714 +
715 + @pytest.mark.parametrize("cohere_only_key", ["template_jinja", "template"])
716 + def test_cohere_only_template_kwargs_are_rejected(self, cohere_only_key):
717 + # ``chat_template_kwargs.template_jinja`` / ``.template`` are
718 + # accepted at Cohere's own API surface but not at vLLM's -- raw
719 + # template source in vLLM must flow through the standard
720 + # ``chat_template`` request field so the
721 + # ``--trust-request-chat-template`` guard applies uniformly.
722 + # Silently dropping these keys would hide client misconfiguration,
723 + # so we reject them loudly instead.
724 + with pytest.raises(
725 + ValueError,
726 + match=f"chat_template_kwargs.{cohere_only_key!r}",
727 + ):
728 + _build_render_config(
729 + self._conv(),
730 + {cohere_only_key: "raw {{ jinja }}"},
731 + ) # type: ignore[arg-type]
732 +
733 + @pytest.mark.parametrize("cohere_only_key", ["template_jinja", "template"])
734 + def test_cohere_only_template_kwargs_none_is_tolerated(self, cohere_only_key):
735 + # An explicit ``None`` (e.g. from ``.model_dump(exclude_none=False)``
736 + # on an optional field) is treated as absent -- we neither raise
737 + # nor let it fall through as a Jinja variable.
738 + _, cfg = _build_render_config(
739 + self._conv(),
740 + {cohere_only_key: None},
741 + ) # type: ignore[arg-type]
742 + assert cohere_only_key not in cfg
743 + assert "additional_template_fields" not in cfg
744 +
745 + def test_chat_template_arg_populates_template_jinja(self):
746 + # The standard vLLM ``chat_template`` request field is the sole
747 + # supported channel for raw template source and is forwarded to
748 + # melody under its ``template_jinja`` config key.
749 + _, cfg = _build_render_config(
750 + self._conv(),
751 + {},
752 + "raw {{ jinja }}",
753 + ) # type: ignore[arg-type]
754 + assert cfg["template_jinja"] == "raw {{ jinja }}"
755 + assert cfg["use_jinja"] is True
756 +
757 + def test_chat_template_arg_none_leaves_template_jinja_unset(self):
758 + _, cfg = _build_render_config(
759 + self._conv(),
760 + {},
761 + None,
762 + ) # type: ignore[arg-type]
763 + assert "template_jinja" not in cfg
764 +
765 +
766 +# ======================================================================
767 +# End-to-end async rendering (mirrors ``test_mistral.py``)
768 +# ======================================================================
769 +#
770 +# Verifies that the synchronous melody bindings run on the renderer's
771 +# thread pool so the asyncio event loop stays responsive under
772 +# concurrent load. Mirrors
773 +# ``test_async_mistral_tokenizer_does_not_block_event_loop`` so future
774 +# regressions in either path are caught uniformly.
775 +
776 +
777 +@dataclass
778 +class _MockHFConfig:
779 + model_type: str = "any"
780 +
781 +
782 +@dataclass
783 +class _MockModelConfig:
784 + runner_type = "generate"
785 + model: str = "cohere-test"
786 + tokenizer: str = "cohere-test"
787 + trust_remote_code: bool = False
788 + max_model_len: int = 100
789 + tokenizer_revision = None
790 + tokenizer_mode = "cohere"
791 + hf_config = _MockHFConfig()
792 + hf_text_config = _MockHFConfig()
793 + encoder_config: dict[str, Any] | None = None
794 + enable_prompt_embeds: bool = True
795 + skip_tokenizer_init: bool = True
796 + is_encoder_decoder: bool = False
797 + is_multimodal_model: bool = False
798 + renderer_num_workers: int = 1
799 +
800 +
801 +@dataclass
802 +class _MockParallelConfig:
803 + _api_process_rank: int = 0
804 +
805 +
806 +@dataclass
807 +class _MockVllmConfig:
808 + model_config: _MockModelConfig
809 + parallel_config: _MockParallelConfig
810 +
811 +
812 +@pytest.mark.asyncio
813 +async def test_async_cohere_renderer_does_not_block_event_loop():
814 + expected_prompt = "MOCK_RENDERED_PROMPT"
815 +
816 + def slow_render(*_a, **_kw):
817 + time.sleep(2)
818 + return expected_prompt
819 +
820 + mock_tokenizer = Mock(spec=HfTokenizer)
821 + renderer = CohereRenderer(
822 + _MockVllmConfig(_MockModelConfig(), _MockParallelConfig()),
823 + tokenizer=mock_tokenizer,
824 + )
825 +
826 + # Replace the (already-imported) ``cohere_melody`` bindings with a
827 + # blocking mock. ``_render`` reads ``self._melody`` at call time, so
828 + # this works even though ``_render_async`` was bound at __init__.
829 + fake_melody = Mock()
830 + fake_melody.render_cmd3 = slow_render
831 + fake_melody.render_cmd4 = slow_render
832 + renderer._melody = fake_melody
833 +
834 + task = renderer.render_messages_async([], ChatParams())
835 +
836 + # Ensure the event loop is not blocked while the (blocking) render
837 + # call is in flight on the thread pool.
838 + blocked_count = 0
839 + for _ in range(20): # ~2 seconds at 0.1s slices
840 + start = time.perf_counter()
841 + await asyncio.sleep(0)
842 + elapsed = time.perf_counter() - start
843 + if elapsed >= 0.5:
844 + blocked_count += 1
845 + await asyncio.sleep(0.1)
846 +
847 + _, prompt = await task
848 + assert prompt["prompt"] == expected_prompt, "Mocked blocking render was not called"
849 + assert blocked_count == 0, "Event loop blocked during rendering"
850 +
851 +
852 +# ======================================================================
853 +# End-to-end: request-level citations -> rendered prompt markup
854 +# ======================================================================
855 +#
856 +# Ties the whole pipeline together: a Cohere v2 request carrying an
857 +# assistant message with citations must produce a rendered prompt that
858 +# contains melody's inline ``<co>...</co: <id>>`` markup around the
859 +# cited span. Preserving this invariant is the whole point of the
860 +# ``_messages_citations`` chat_template_kwargs entry.
861 +
862 +
863 +class TestRequestCitationsReachRenderedPrompt:
864 + """End-to-end verification that assistant-message citations on the
865 + request survive the OpenAI-shape round-trip and land in the melody-
866 + rendered prompt as inline ``<co>...</co>`` markers.
867 +
868 + The chain covered here:
869 +
870 + CohereChatV2Request
871 + -> CohereServingChatV2._convert_v2_to_chat_completion
872 + -> ChatCompletionRequest.chat_template_kwargs["_messages_citations"]
873 + -> _build_render_config (reads the entry)
874 + -> cohere_melody.render_cmd4 (renders <co>...</co>)
875 +
876 + A regression that drops the citations or mangles the melody
877 + ``FilterCitation`` payload would remove the citation markers from
878 + the output, so a text-in / text-out assertion is enough.
879 + """
880 +
881 + @staticmethod
882 + def _melody():
883 + # Import locally so missing optional deps skip this test class
884 + # rather than failing collection. Two guards:
885 + # * ``cohere_melody`` -- the Rust binding this test drives.
886 + # * ``cohere`` -- required transitively by
887 + # ``vllm.entrypoints.cohere.{protocol,serving}`` (both
888 + # unconditionally ``from cohere.types import ...`` at module
889 + # scope) which every test in this class imports locally.
890 + pytest.importorskip("cohere")
891 + pytest.importorskip("cohere_melody")
892 + import cohere_melody
893 +
894 + return cohere_melody
895 +
896 + @staticmethod
897 + def _openai_msgs_to_conversation(
898 + openai_messages: list[dict[str, Any]],
899 + ) -> list[dict[str, Any]]:
900 + """Cheap stand-in for ``parse_chat_messages`` on text-only inputs.
901 +
902 + ``_convert_v2_to_chat_completion`` emits OpenAI-shape assistant
903 + dicts with ``content`` as a plain string. The real chat-utils
904 + pipeline calls ``parse_chat_messages`` (which requires a full
905 + model config / mm tracker), but for text-only content it's
906 + essentially a passthrough -- we normalize ``content`` to the
907 + list-of-parts shape ``_conversation_to_melody_messages`` expects
908 + and preserve every other key the renderer reads.
909 + """
910 + conv: list[dict[str, Any]] = []
911 + for m in openai_messages:
912 + entry: dict[str, Any] = dict(m)
913 + content = entry.get("content")
914 + if isinstance(content, str):
915 + entry["content"] = [{"type": "text", "text": content}]
916 + elif content is None:
917 + entry["content"] = []
918 + conv.append(entry)
919 + return conv
920 +
921 + def test_document_citation_survives_to_prompt(self):
922 + # Prevent this test from being collected/run when the melody
923 + # extension isn't importable.
924 + melody = self._melody()
925 +
926 + from vllm.entrypoints.cohere.protocol import CohereChatV2Request
927 + from vllm.entrypoints.cohere.serving import CohereServingChatV2
928 + from vllm.renderers.cohere import _build_render_config
929 +
930 + request = CohereChatV2Request(
931 + model="m",
932 + messages=[
933 + {"role": "user", "content": "Who wrote Hamlet?"},
934 + {
935 + "role": "assistant",
936 + "content": "Shakespeare wrote it around 1600.",
937 + "citations": [
938 + {
939 + "start": 0,
940 + "end": 11, # "Shakespeare"
941 + "text": "Shakespeare",
942 + "sources": [{"type": "document", "id": "doc_shakespeare"}],
943 + "type": "TEXT_CONTENT",
944 + }
945 + ],
946 + },
947 + {"role": "user", "content": "and what year exactly?"},
948 + ],
949 + documents=[
950 + {
951 + "id": "doc_shakespeare",
952 + "data": {"text": "Hamlet was written by Shakespeare c. 1600."},
953 + }
954 + ],
955 + )
956 +
957 + # Step 1: v2 -> ChatCompletionRequest. Citations must land in
958 + # ``chat_template_kwargs["_messages_citations"]``.
959 + chat_req = CohereServingChatV2._convert_v2_to_chat_completion(request)
960 + assert chat_req.chat_template_kwargs is not None
961 + assert "_messages_citations" in chat_req.chat_template_kwargs
962 +
963 + # Step 2: build the melody render config. The renderer helper
964 + # folds the per-message citations onto the melody message
965 + # dicts.
966 + conversation = self._openai_msgs_to_conversation(chat_req.messages)
967 + fmt, config = _build_render_config(conversation, chat_req.chat_template_kwargs)
968 +
969 + assistant_msg = config["messages"][1]
970 + assert assistant_msg["role"] == "chatbot"
971 + assert "citations" in assistant_msg, (
972 + "citations were not attached to the melody assistant message dict"
973 + )
974 +
975 + # Step 3: hand the config to melody and check the rendered
976 + # prompt actually contains inline citation markup around the
977 + # cited span.
978 + #
979 + # The exact id is deterministic for this input. Melody builds
980 + # ``</co: <tool_call_index>:[<tool_result_indices>]>`` where
981 + # ``tool_call_index=0`` is the reserved bucket for the top-level
982 + # ``documents`` array and ``tool_result_indices`` are positions
983 + # inside it (see ``PromptRenderIds`` in melody/src/templating/
984 + # util.rs). ``doc_shakespeare`` sits at position 0 in the
985 + # request's ``documents`` list, so we expect ``0:[0]``. Two
986 + # historical regressions this pins:
987 + # * ``0:[]`` -- the source id was never resolved to an index
988 + # (documents didn't flow through) and melody had nothing to
989 + # anchor the marker on.
990 + # * ``1:[0]`` -- the citation was routed through the wrong
991 + # tool-call bucket while documents were present, so it
992 + # pointed at the wrong prompt slot.
993 + #
994 + # Note the same rendered prompt also contains an example
995 + # ``<co>span</co: 0:[1,2],1:[0]>`` marker baked into melody's
996 + # system-prompt boilerplate (placeholder text ``"span"``); the
997 + # substring below is specific enough to only match the marker
998 + # around the cited text.
999 + if fmt == "cmd4":
1000 + rendered = melody.render_cmd4(config)
1001 + else:
1002 + rendered = melody.render_cmd3(config)
1003 +
1004 + assert "<co>Shakespeare</co: 0:[0]>" in rendered, (
1005 + f"expected inline citation markup around the cited span; "
1006 + f"tail of rendered prompt: {rendered[-400:]!r}"
1007 + )
1008 +
1009 + # And the cited document's text itself must be in the prompt --
1010 + # otherwise the model would have no way to satisfy the citation.
1011 + assert "Hamlet" in rendered
1012 +
1013 + def test_no_markup_when_no_citations(self):
1014 + # Control: the same request shape without any citations must
1015 + # NOT contain ``<co>`` anywhere in the rendered prompt. Guards
1016 + # against a false-positive where melody injects citation
1017 + # markers regardless of what we passed in.
1018 + melody = self._melody()
1019 +
1020 + from vllm.entrypoints.cohere.protocol import CohereChatV2Request
1021 + from vllm.entrypoints.cohere.serving import CohereServingChatV2
1022 + from vllm.renderers.cohere import _build_render_config
1023 +
1024 + request = CohereChatV2Request(
1025 + model="m",
1026 + messages=[
1027 + {"role": "user", "content": "Who wrote Hamlet?"},
1028 + {
1029 + "role": "assistant",
1030 + "content": "Shakespeare wrote it around 1600.",
1031 + },
1032 + ],
1033 + )
1034 +
1035 + chat_req = CohereServingChatV2._convert_v2_to_chat_completion(request)
1036 + assert (chat_req.chat_template_kwargs or {}).get("_messages_citations") is None
1037 +
1038 + conversation = self._openai_msgs_to_conversation(chat_req.messages)
1039 + fmt, config = _build_render_config(
1040 + conversation, chat_req.chat_template_kwargs or {}
1041 + )
1042 +
1043 + if fmt == "cmd4":
1044 + rendered = melody.render_cmd4(config)
1045 + else:
1046 + rendered = melody.render_cmd3(config)
1047 +
1048 + assert "<co>" not in rendered
1049 + assert "</co:" not in rendered
vllm/config/model.py
+4
@@ -93,6 +93,7 @@ TokenizerMode = Literal[
93 "deepseek_v4",
94 "inkling",
95 "kimi_k3",
96 + "cohere",
97 ]
98 ModelDType = Literal["auto", "half", "float16", "bfloat16", "float", "float32"]
99 LogprobsMode = Literal[
@@ -151,6 +152,9 @@ class ModelConfig:
152 - "deepseek_v4" will always use the tokenizer from `deepseek_v4`.
153 - "kimi_k3" will always use the "hf" tokenizer but render chat prompts
154 with Kimi K3's Python XTML encoding instead of a Jinja template.
155 + - "cohere" uses the standard HF tokenizer but renders the chat template
156 + via the `cohere_melody` library (cmd3 / cmd4 templates) instead of
157 + Jinja, and surfaces grounded-citation metadata on responses.
158 - Other custom values can be supported via plugins.
159
160 To swap the Rust BPE backend that powers HF fast tokenizers for the
vllm/entrypoints/cohere/__init__.py new
+2
@@ -0,0 +1,2 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm/entrypoints/cohere/api_router.py new
+244
@@ -0,0 +1,244 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""FastAPI router for the Cohere Chat v2 API (``POST /cohere/v2/chat``).
4 +
5 +The Cohere v2 protocol models are sourced from the official ``cohere``
6 +Python SDK (``pip install cohere``). To keep that an *optional*
7 +dependency for vLLM, the SDK-dependent imports - and the route handler
8 +itself - are gated on a one-shot probe at module load. If the SDK isn't
9 +installed, :func:`attach_router` becomes a no-op (with an info log) and
10 +vLLM continues to boot normally.
11 +
12 +Even when the SDK is installed, :func:`attach_router` also requires
13 +``VLLM_ENABLE_COHERE_API=1`` in the environment before it will expose
14 +the route. This keeps non-Cohere deployments that pull in the SDK for
15 +unrelated reasons (e.g. test dependencies) from accidentally exposing
16 +the api.
17 +
18 +Note: the handler must live at module scope (not inside
19 +``attach_router``) so that FastAPI's ``typing.get_type_hints`` resolves
20 +the ``CohereChatV2Request`` body annotation against the module's
21 +globals. Defining it locally inside ``attach_router`` would hide the
22 +type from ``get_type_hints``, causing FastAPI to silently degrade the
23 +body parameter into a query parameter and reject every request with
24 +422.
25 +"""
26 +
27 +import json
28 +from http import HTTPStatus
29 +
30 +from fastapi import APIRouter, Depends, FastAPI, Request
31 +from fastapi.responses import JSONResponse, Response, StreamingResponse
32 +from starlette.middleware.base import BaseHTTPMiddleware
33 +
34 +import vllm.envs as envs
35 +from vllm.entrypoints.openai.engine.protocol import ErrorResponse
36 +from vllm.entrypoints.serve.utils.api_utils import (
37 + load_aware_call,
38 + sanitize_message,
39 + validate_json_request,
40 + with_cancellation,
41 +)
42 +from vllm.logger import init_logger
43 +
44 +_COHERE_PATH_PREFIX = "/cohere/"
45 +
46 +logger = init_logger(__name__)
47 +
48 +
49 +try:
50 + import cohere # noqa: F401 -- dependency probe
51 +except ImportError:
52 + _SDK_AVAILABLE = False
53 +else:
54 + _SDK_AVAILABLE = True
55 +
56 +
57 +if _SDK_AVAILABLE:
58 + from vllm.entrypoints.cohere.protocol import (
59 + CohereChatV2Request,
60 + CohereChatV2Response,
61 + CohereError,
62 + )
63 + from vllm.entrypoints.cohere.serving import CohereServingChatV2
64 +
65 + router = APIRouter()
66 +
67 + def _serving(request: Request) -> CohereServingChatV2 | None:
68 + return getattr(request.app.state, "cohere_serving_chat_v2", None)
69 +
70 + def _request_id(raw_request: Request | None) -> str | None:
71 + """Best-effort lookup of the active request id.
72 +
73 + Prefers the id the underlying chat handler stamped onto
74 + ``raw_request.state.request_metadata`` (if it got that far before
75 + failing), falling back to the ``X-Request-Id`` HTTP header. May
76 + return ``None`` if neither is available, in which case the field
77 + is omitted from the response.
78 + """
79 + if raw_request is None:
80 + return None
81 + meta = getattr(raw_request.state, "request_metadata", None)
82 + if meta is not None and getattr(meta, "request_id", None):
83 + return meta.request_id
84 + return raw_request.headers.get("X-Request-Id")
85 +
86 + def _error_response(
87 + error: ErrorResponse,
88 + raw_request: Request | None,
89 + *,
90 + fallback_status: int = HTTPStatus.BAD_REQUEST,
91 + ) -> JSONResponse:
92 + """Translate vLLM's internal error envelope into Cohere's shape."""
93 + info = error.error
94 + status = info.code or fallback_status
95 + return JSONResponse(
96 + status_code=status,
97 + content=CohereError(
98 + message=sanitize_message(info.message),
99 + id=_request_id(raw_request),
100 + ).model_dump(exclude_none=True),
101 + )
102 +
103 + @router.post(
104 + "/cohere/v2/chat",
105 + dependencies=[Depends(validate_json_request)],
106 + responses={
107 + HTTPStatus.OK.value: {"content": {"text/event-stream": {}}},
108 + HTTPStatus.BAD_REQUEST.value: {"model": CohereError},
109 + HTTPStatus.NOT_FOUND.value: {"model": CohereError},
110 + HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": CohereError},
111 + },
112 + )
113 + @with_cancellation
114 + @load_aware_call
115 + async def chat_v2(request: CohereChatV2Request, raw_request: Request):
116 + handler = _serving(raw_request)
117 + if handler is None:
118 + return JSONResponse(
119 + status_code=HTTPStatus.NOT_IMPLEMENTED.value,
120 + content=CohereError(
121 + message="The model does not support the Cohere v2 chat API.",
122 + id=_request_id(raw_request),
123 + ).model_dump(exclude_none=True),
124 + )
125 +
126 + try:
127 + result = await handler.create_chat_v2(request, raw_request)
128 + except Exception as e: # noqa: BLE001 - report as 500 for parity
129 + logger.exception("Error in /cohere/v2/chat: %s", e)
130 + return JSONResponse(
131 + status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value,
132 + content=CohereError(
133 + message=sanitize_message(str(e)),
134 + id=_request_id(raw_request),
135 + ).model_dump(exclude_none=True),
136 + )
137 +
138 + match result:
139 + case ErrorResponse():
140 + return _error_response(result, raw_request)
141 + case CohereChatV2Response():
142 + return JSONResponse(content=result.model_dump(exclude_none=True))
143 + case _:
144 + return StreamingResponse(content=result, media_type="text/event-stream")
145 +
146 + class CohereErrorEnvelopeMiddleware(BaseHTTPMiddleware):
147 + """Rewrite vLLM error bodies into the Cohere ``{message, id}`` shape.
148 +
149 + The endpoint handler above already returns :class:`CohereError` for
150 + errors it owns, but globally-registered exception handlers (e.g.
151 + :func:`validation_exception_handler` for pydantic body errors,
152 + :func:`http_exception_handler`, engine error handlers) fire *before*
153 + the handler runs and produce vLLM's internal
154 + ``ErrorResponse`` shape (``{"error": {"message": ...}}``). That
155 + shape doesn't match the ``CohereError`` schema advertised on the
156 + route's OpenAPI ``responses``, so clients (and schema-conformance
157 + tests like ``test_openai_schema.py``) would see a mismatch on
158 + those paths. This middleware normalises any error body on
159 + ``/cohere/*`` responses to :class:`CohereError`.
160 + """
161 +
162 + async def dispatch(self, request: Request, call_next):
163 + response = await call_next(request)
164 + if not request.url.path.startswith(_COHERE_PATH_PREFIX):
165 + return response
166 + if response.status_code < 400:
167 + return response
168 + content_type = response.headers.get("content-type", "")
169 + if not content_type.startswith("application/json"):
170 + return response
171 +
172 + body = b"".join([chunk async for chunk in response.body_iterator])
173 + translated = _translate_vllm_error_body(body, request)
174 + if translated is not None:
175 + return translated
176 + passthrough_headers = {
177 + k: v
178 + for k, v in response.headers.items()
179 + if k.lower() != "content-length"
180 + }
181 + return Response(
182 + content=body,
183 + status_code=response.status_code,
184 + headers=passthrough_headers,
185 + media_type=content_type,
186 + )
187 +
188 + def _translate_vllm_error_body(raw: bytes, request: Request) -> JSONResponse | None:
189 + """Translate a vLLM ``ErrorResponse`` body to a ``CohereError`` body.
190 +
191 + Returns ``None`` if ``raw`` does not match the vLLM error envelope
192 + (which signals the middleware to pass the body through unchanged).
193 + """
194 + try:
195 + data = json.loads(raw)
196 + except (json.JSONDecodeError, TypeError, ValueError):
197 + return None
198 + if not (
199 + isinstance(data, dict)
200 + and isinstance(data.get("error"), dict)
201 + and "message" in data["error"]
202 + ):
203 + return None
204 + try:
205 + err = ErrorResponse.model_validate(data)
206 + except Exception: # noqa: BLE001 - malformed envelope; pass through
207 + return None
208 + return _error_response(err, request)
209 +
210 +
211 +def attach_router(app: FastAPI) -> None:
212 + """Register ``POST /cohere/v2/chat`` on ``app``.
213 +
214 + No-op when either:
215 +
216 + * the ``VLLM_ENABLE_COHERE_API`` env var isn't set to ``1``. The
217 + Cohere v2 endpoint is opt-in because it carries Cohere-specific
218 + request/response semantics (grounding citations, tool_plan,
219 + PLAN/THINKING_CONTENT blocks) that are only meaningful when
220 + serving a Cohere Command-family model.
221 + * the optional ``cohere`` SDK isn't installed (the v2 protocol
222 + models live there)
223 +
224 + The two skip paths log at different levels: an operator who set
225 + ``VLLM_ENABLE_COHERE_API=1`` but forgot to install ``cohere`` sees
226 + a WARNING (they explicitly asked for the endpoint and it's silently
227 + absent), whereas the default-off skip logs at debug.
228 + """
229 + enabled = envs.VLLM_ENABLE_COHERE_API
230 + if not enabled:
231 + logger.debug(
232 + "VLLM_ENABLE_COHERE_API is not set; /cohere/v2/chat endpoint "
233 + "disabled. Set VLLM_ENABLE_COHERE_API=1 to enable it."
234 + )
235 + return
236 + if not _SDK_AVAILABLE:
237 + logger.warning(
238 + "VLLM_ENABLE_COHERE_API=1 but the `cohere` SDK is not "
239 + "installed; /cohere/v2/chat will not be exposed. Install "
240 + "with `pip install cohere` to enable the endpoint."
241 + )
242 + return
243 + app.include_router(router)
244 + app.add_middleware(CohereErrorEnvelopeMiddleware)
vllm/entrypoints/cohere/cohere_chat_message.py new
+147
@@ -0,0 +1,147 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""vLLM chat-protocol extensions for grounded (citation-carrying) models.
4 +
5 +This module keeps the OpenAI chat completion protocol classes
6 +(:class:`ChatMessage` / :class:`DeltaMessage`) free of Cohere-specific fields.
7 +It provides two things:
8 +
9 +* :class:`Citation` / :class:`CitationSource`: the vLLM-internal citation
10 + representation produced by the Cohere reasoning parser
11 + (:mod:`vllm.reasoning.cohere_command_reasoning_parser`) and consumed by
12 + :mod:`vllm.entrypoints.cohere.serving`. This is *not* the on-wire Cohere
13 + SDK shape -- that conversion happens in the serving layer.
14 +* :class:`CohereChatMessage` / :class:`CohereDeltaMessage`: :class:`ChatMessage`
15 + / :class:`DeltaMessage` subclasses that add a ``citations`` field. They
16 + are only instantiated by the citation-aware serving handler
17 + (:class:`vllm.entrypoints.cohere.serving.CohereServingChatV2`) and by the
18 + Cohere reasoning parser's streaming path. Non-Cohere code paths continue
19 + to construct plain :class:`ChatMessage` / :class:`DeltaMessage`.
20 +
21 +The module intentionally does *not* import the ``cohere`` Python SDK so the
22 +reasoning parser (which is Cohere-model-specific but ships without the SDK
23 +dependency) can import from here freely.
24 +"""
25 +
26 +from __future__ import annotations
27 +
28 +from typing import Any, Literal
29 +
30 +from pydantic import Field, model_serializer
31 +
32 +from vllm.entrypoints.openai.chat_completion.protocol import ChatMessage
33 +from vllm.entrypoints.openai.engine.protocol import (
34 + DeltaMessage,
35 + OpenAIBaseModel,
36 +)
37 +
38 +
39 +class CitationSource(OpenAIBaseModel):
40 + """Source attribution for a :class:`Citation`.
41 +
42 + Mirrors the shape used by Cohere's Chat v2 API. ``type`` is the source
43 + discriminator (``document`` or ``tool``); ``id`` is the citing
44 + document/tool-output identifier; ``document`` and ``tool_output`` carry
45 + the original payload that produced the citation.
46 +
47 + Sources are fully resolved by the reasoning parser at emit time
48 + (see :func:`_melody_sources_to_vllm` in
49 + :mod:`vllm.reasoning.cohere_command_reasoning_parser`) using the
50 + ``POSITION_TO_SOURCE_KEY`` map forwarded through
51 + ``chat_template_kwargs`` by
52 + :meth:`vllm.entrypoints.cohere.serving.CohereServingChatV2._apply_cohere_template_kwargs`.
53 + That means every instance that reaches the serving / wire layer
54 + already has ``type`` / ``id`` populated; the serving layer only has
55 + to coerce to :class:`cohere.types.Citation` and rewrite
56 + ``THINKING_CONTENT`` -> ``PLAN`` for non-reasoning models.
57 + """
58 +
59 + type: Literal["document", "tool"] | None = None
60 + id: str | None = None
61 + document: dict[str, Any] | None = None
62 + tool_output: dict[str, Any] | None = None
63 +
64 +
65 +class Citation(OpenAIBaseModel):
66 + """A citation grounding a span of generated text in source material.
67 +
68 + vLLM-internal representation used by the Cohere reasoning parser and
69 + the Cohere v2 serving layer. This is not the on-wire Cohere SDK shape:
70 + conversion to the SDK's ``cohere.types.Citation`` happens in
71 + :mod:`vllm.entrypoints.cohere.serving`.
72 + """
73 +
74 + start: int | None = None
75 + """Start character offset in the surrounding text content."""
76 + end: int | None = None
77 + """End character offset (exclusive) in the surrounding text content."""
78 + text: str | None = None
79 + """The cited text snippet."""
80 + sources: list[CitationSource] = Field(default_factory=list)
81 + """Source documents / tool outputs that ground this citation."""
82 + content_index: int | None = None
83 + """Index of the content block this citation refers to (when the message
84 + has multiple content blocks)."""
85 + type: Literal["TEXT_CONTENT", "THINKING_CONTENT", "PLAN"] | None = None
86 + """Which kind of content block this citation grounds: the user-visible
87 + text (``TEXT_CONTENT``), a thinking block (``THINKING_CONTENT``), or a
88 + tool-plan block (``PLAN``). ``None`` means unspecified."""
89 +
90 +
91 +class CohereChatMessage(ChatMessage):
92 + """:class:`ChatMessage` extension carrying grounding citations.
93 +
94 + Only instantiated by :class:`CohereServingChatV2` (and any other
95 + :class:`OpenAIServingChat` subclass that opts in by overriding
96 + ``_create_chat_message``). Regular OpenAI-compatible handlers keep
97 + emitting plain :class:`ChatMessage` so their response schema is
98 + unchanged.
99 +
100 + The response envelope declares ``message: SerializeAsAny[ChatMessage]``,
101 + so pydantic serializes this subclass with its own schema (including
102 + ``citations``) when it flows through
103 + :class:`ChatCompletionResponseChoice`.
104 + """
105 +
106 + citations: list[Citation] | None = None
107 +
108 + @model_serializer(mode="wrap")
109 + def _serialize(self, handler):
110 + # ``mode="wrap"`` fully overrides (rather than chains) the
111 + # parent's ``@model_serializer``, so we explicitly delegate via
112 + # ``super()._serialize(handler)`` to preserve the parent's
113 + # cleanup (e.g. stripping empty ``tool_calls``). Then we drop an
114 + # unset ``citations`` field so the wire matches the OpenAI-style
115 + # contract that optional vLLM extensions are omitted rather than
116 + # serialized as ``null``.
117 + data = super()._serialize(handler)
118 + if not data.get("citations"):
119 + data.pop("citations", None)
120 + return data
121 +
122 +
123 +class CohereDeltaMessage(DeltaMessage):
124 + """:class:`DeltaMessage` extension carrying grounding citations for streaming.
125 +
126 + Emitted by
127 + :class:`vllm.reasoning.cohere_command_reasoning_parser.BaseCohereCommandReasoningParser`
128 + on delta events whose payload includes citations. Non-Cohere parsers
129 + return plain :class:`DeltaMessage`, so their streamed shape is
130 + unchanged.
131 +
132 + The response envelope declares
133 + ``delta: SerializeAsAny[DeltaMessage]``, so pydantic serializes this
134 + subclass with its own schema when it flows through
135 + :class:`ChatCompletionResponseStreamChoice`.
136 + """
137 +
138 + citations: list[Citation] | None = None
139 +
140 + @model_serializer(mode="wrap")
141 + def _serialize(self, handler):
142 + # See ``CohereChatMessage._serialize`` for why we delegate via
143 + # ``super()`` instead of calling ``handler(self)`` directly.
144 + data = super()._serialize(handler)
145 + if not data.get("citations"):
146 + data.pop("citations", None)
147 + return data
vllm/entrypoints/cohere/protocol.py new
+362
@@ -0,0 +1,362 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Cohere Chat v2 API protocol.
4 +
5 +The bulk of the wire types come straight from the official ``cohere``
6 +Python SDK so we stay in sync with the upstream specification and
7 +avoid re-declaring the schema. We only own three things locally:
8 +
9 +1. The top-level request body model (the SDK doesn't ship one — its
10 + ``ClientV2.chat`` takes the body as kwargs), with vLLM-specific
11 + extensions (``kv_transfer_params`` / ``chat_template_kwargs``).
12 +2. The non-streaming response envelope (the SDK exposes the message
13 + shape via :class:`AssistantMessageResponse` but no full response
14 + wrapper).
15 +3. The streaming discriminated union (the SDK exports each event type
16 + individually but not as a combined ``Annotated[Union[...],
17 + discriminator]``).
18 +
19 +Importing this module pulls in the ``cohere`` package. The router that
20 +mounts ``POST /cohere/v2/chat`` guards on that import succeeding so vLLM still
21 +boots without the SDK installed.
22 +
23 +See https://docs.cohere.com/reference/chat for the upstream spec.
24 +"""
25 +
26 +from __future__ import annotations
27 +
28 +from typing import Any, Literal
29 +
30 +from cohere import types as _sdk
31 +from cohere.types import (
32 + AssistantChatMessageV2,
33 + AssistantMessageResponse,
34 + ChatMessageV2,
35 + ChatRequestSafetyMode,
36 + Citation,
37 + CitationOptions,
38 + Document,
39 + ResponseFormatV2,
40 + SystemChatMessageV2,
41 + Thinking,
42 + ToolCallV2,
43 + ToolChatMessageV2,
44 + ToolV2,
45 + UserChatMessageV2,
46 +)
47 +from pydantic import BaseModel, Field, field_validator
48 +
49 +# Re-export the SDK wire-format types alongside our local extensions so
50 +# ``vllm.entrypoints.cohere.serving`` and friends can import everything
51 +# they need from this module.
52 +__all__ = [
53 + "AssistantChatMessageV2",
54 + "AssistantMessageResponse",
55 + "ChatMessageV2",
56 + "Citation",
57 + "CitationEndEvent",
58 + "CitationOptions",
59 + "CitationStartEvent",
60 + "CohereChatV2Request",
61 + "CohereChatV2Response",
62 + "CohereError",
63 + "CohereFinishReason",
64 + "CohereLogprobItem",
65 + "CohereUsage",
66 + "CohereUsageBilledUnits",
67 + "CohereUsageTokens",
68 + "ContentDeltaEvent",
69 + "ContentEndEvent",
70 + "ContentStartEvent",
71 + "Document",
72 + "MessageEndEvent",
73 + "MessageStartEvent",
74 + "ResponseFormatV2",
75 + "SystemChatMessageV2",
76 + "Thinking",
77 + "ToolCallDeltaEvent",
78 + "ToolCallEndEvent",
79 + "ToolCallStartEvent",
80 + "ToolCallV2",
81 + "ToolChatMessageV2",
82 + "ToolPlanDeltaEvent",
83 + "ToolV2",
84 + "UserChatMessageV2",
85 +]
86 +
87 +
88 +# ---------------------------------------------------------------------------
89 +# Errors
90 +# ---------------------------------------------------------------------------
91 +
92 +
93 +class CohereError(BaseModel):
94 + """Top-level error body returned by ``/cohere/v2/chat`` error responses.
95 +
96 + Cohere's documented error schemas are uniform: ``{message, id}``.
97 + """
98 +
99 + message: str
100 + id: str | None = None
101 +
102 +
103 +# ---------------------------------------------------------------------------
104 +# Tool choice / finish reasons
105 +# ---------------------------------------------------------------------------
106 +#
107 +# These literals aren't first-class enums in the SDK but are documented at
108 +# https://docs.cohere.com/reference/chat. We declare them here so the
109 +# request/response models can validate them.
110 +
111 +CohereToolChoice = Literal["REQUIRED", "NONE"]
112 +
113 +CohereFinishReason = Literal[
114 + "COMPLETE",
115 + "STOP_SEQUENCE",
116 + "MAX_TOKENS",
117 + "TOOL_CALL",
118 + "ERROR",
119 + "TIMEOUT",
120 +]
121 +
122 +
123 +# ---------------------------------------------------------------------------
124 +# Request
125 +# ---------------------------------------------------------------------------
126 +
127 +
128 +class CohereChatV2Request(BaseModel):
129 + """Cohere Chat v2 request body.
130 +
131 + Mirrors the schema documented at https://docs.cohere.com/reference/chat.
132 + All structured fields delegate to the official SDK types so the body
133 + schema stays in sync with the upstream spec.
134 + """
135 +
136 + model: str
137 + messages: list[ChatMessageV2]
138 + stream: bool | None = False
139 +
140 + # Tooling
141 + tools: list[ToolV2] | None = None
142 + strict_tools: bool | None = None
143 + tool_choice: CohereToolChoice | None = None
144 +
145 + # Grounding
146 + documents: list[str | Document] | None = None
147 + citation_options: CitationOptions | None = None
148 +
149 + # Output
150 + response_format: ResponseFormatV2 | None = None
151 + safety_mode: ChatRequestSafetyMode | None = None
152 + max_tokens: int | None = None
153 + stop_sequences: list[str] | None = None
154 +
155 + # Sampling
156 + temperature: float | None = None
157 + seed: int | None = None
158 + frequency_penalty: float | None = None
159 + presence_penalty: float | None = None
160 + k: int | None = None
161 + p: float | None = None
162 + logprobs: bool | None = None
163 +
164 + # Reasoning
165 + thinking: Thinking | None = None
166 +
167 + # Scheduling
168 + priority: int | None = None
169 +
170 + # vLLM-specific extensions (not in Cohere spec). These mirror what the
171 + # Anthropic and OpenAI surfaces already expose so V2 callers can reach
172 + # the same engine knobs when needed.
173 + kv_transfer_params: dict[str, Any] | None = Field(
174 + default=None,
175 + description="KVTransfer parameters used for disaggregated serving.",
176 + )
177 + chat_template_kwargs: dict[str, Any] | None = Field(
178 + default=None,
179 + description=(
180 + "Additional keyword args to pass to the chat template renderer. "
181 + "Will be accessible by the template."
182 + ),
183 + )
184 +
185 + @field_validator("model")
186 + @classmethod
187 + def _validate_model(cls, v: str) -> str:
188 + if not v:
189 + raise ValueError("model is required")
190 + return v
191 +
192 + @field_validator("max_tokens")
193 + @classmethod
194 + def _validate_max_tokens(cls, v: int | None) -> int | None:
195 + if v is not None and v < 0:
196 + raise ValueError("max_tokens must be non-negative")
197 + return v
198 +
199 + @field_validator("messages", mode="before")
200 + @classmethod
201 + def _normalize_message_roles(cls, v: Any) -> Any:
202 + """Rewrite OpenAI-style ``developer`` roles to ``system``.
203 +
204 + Cohere's v2 ``ChatMessageV2`` discriminated union only admits
205 + the four literal roles ``user`` / ``assistant`` / ``system`` /
206 + ``tool``. OpenAI's ``developer`` role is documented as
207 + high-priority system instructions, so we alias it onto
208 + ``system`` *before* the SDK discriminator runs (otherwise it
209 + rejects the message with a ``literal_error`` against each union
210 + member). Mirrors ``_role_to_melody`` in the renderer so the
211 + same alias is honoured no matter which surface the message
212 + arrives through.
213 +
214 + ``mode="before"`` is required so the rewrite happens before the
215 + ``list[ChatMessageV2]`` coercion runs the SDK's discriminated
216 + union; a default-mode validator would never see ``developer``
217 + because validation would have already failed. On any
218 + structural malformation (non-iterable input, items without a
219 + dict-shaped ``role`` field, etc.) we hand ``v`` back unchanged
220 + and let Pydantic's normal coercion surface a precise error.
221 + """
222 + try:
223 + return [
224 + {**msg, "role": "system"}
225 + if msg.get("role", "").lower() == "developer"
226 + else msg
227 + for msg in v
228 + ]
229 + except (AttributeError, TypeError):
230 + return v
231 +
232 + @field_validator("messages")
233 + @classmethod
234 + def _validate_messages(cls, v: list[ChatMessageV2]) -> list[ChatMessageV2]:
235 + if not v:
236 + raise ValueError("messages must contain at least one message")
237 + return v
238 +
239 +
240 +# ---------------------------------------------------------------------------
241 +# Usage / Logprobs
242 +# ---------------------------------------------------------------------------
243 +#
244 +# The Cohere SDK only exposes a v1 ``ApiMetaBilledUnits``. The v2 usage
245 +# envelope is documented separately in the OpenAPI spec and we declare it
246 +# here.
247 +
248 +
249 +class CohereUsageBilledUnits(BaseModel):
250 + input_tokens: float | None = None
251 + output_tokens: float | None = None
252 + search_units: float | None = None
253 + classifications: float | None = None
254 +
255 +
256 +class CohereUsageTokens(BaseModel):
257 + input_tokens: float | None = None
258 + output_tokens: float | None = None
259 +
260 +
261 +class CohereUsage(BaseModel):
262 + billed_units: CohereUsageBilledUnits | None = None
263 + tokens: CohereUsageTokens | None = None
264 + cached_tokens: float | None = None
265 +
266 +
267 +class CohereLogprobItem(BaseModel):
268 + text: str | None = None
269 + token_ids: list[int]
270 + logprobs: list[float] | None = None
271 +
272 +
273 +# ---------------------------------------------------------------------------
274 +# Non-streaming response
275 +# ---------------------------------------------------------------------------
276 +
277 +
278 +class CohereChatV2Response(BaseModel):
279 + """Cohere Chat v2 non-streaming response body.
280 +
281 + Wraps the SDK :class:`AssistantMessageResponse` (the message shape) in
282 + the documented v2 response envelope (``id``, ``finish_reason``,
283 + ``usage``, ``logprobs``). The single constructor in
284 + :class:`CohereServingChatV2._chat_completion_to_v2` is responsible for
285 + supplying a non-empty ``id`` (falling back to a synthesized one if
286 + the upstream response is missing it) to this model.
287 + """
288 +
289 + id: str
290 + finish_reason: CohereFinishReason
291 + message: AssistantMessageResponse
292 + usage: CohereUsage | None = None
293 + logprobs: list[CohereLogprobItem] | None = None
294 +
295 + # vLLM-specific extension.
296 + kv_transfer_params: dict[str, Any] | None = Field(
297 + default=None, description="KVTransfer parameters."
298 + )
299 +
300 +
301 +# ---------------------------------------------------------------------------
302 +# Streaming events
303 +# ---------------------------------------------------------------------------
304 +#
305 +# Cohere V2 streams a sequence of typed JSON events delivered as Server-
306 +# Sent Events; each event's ``type`` field carries the discriminator. The
307 +# SDK exposes a Pydantic model per event but none of them declare ``type``
308 +# as a field (the SDK relies on its own deserializer for discrimination),
309 +# so a naive ``model_dump_json()`` would silently drop the discriminator
310 +# and break clients that demux on ``type``.
311 +#
312 +# We therefore subclass each SDK event and bake the wire-format ``type``
313 +# string in as a ``Literal`` field with a default. ``model_dump()`` now
314 +# emits ``type`` for free, and surfaces can simply construct the event
315 +# class and serialize it -- no manual ``type`` parameter needed.
316 +#
317 +# See https://docs.cohere.com/v2/docs/streaming and the OpenAPI
318 +# ``StreamedChatResponseV2`` schema for the wire-format reference.
319 +
320 +
321 +class MessageStartEvent(_sdk.ChatMessageStartEvent):
322 + type: Literal["message-start"] = "message-start"
323 +
324 +
325 +class ContentStartEvent(_sdk.ChatContentStartEvent):
326 + type: Literal["content-start"] = "content-start"
327 +
328 +
329 +class ContentDeltaEvent(_sdk.ChatContentDeltaEvent):
330 + type: Literal["content-delta"] = "content-delta"
331 +
332 +
333 +class ContentEndEvent(_sdk.ChatContentEndEvent):
334 + type: Literal["content-end"] = "content-end"
335 +
336 +
337 +class ToolPlanDeltaEvent(_sdk.ChatToolPlanDeltaEvent):
338 + type: Literal["tool-plan-delta"] = "tool-plan-delta"
339 +
340 +
341 +class ToolCallStartEvent(_sdk.ChatToolCallStartEvent):
342 + type: Literal["tool-call-start"] = "tool-call-start"
343 +
344 +
345 +class ToolCallDeltaEvent(_sdk.ChatToolCallDeltaEvent):
346 + type: Literal["tool-call-delta"] = "tool-call-delta"
347 +
348 +
349 +class ToolCallEndEvent(_sdk.ChatToolCallEndEvent):
350 + type: Literal["tool-call-end"] = "tool-call-end"
351 +
352 +
353 +class CitationStartEvent(_sdk.CitationStartEvent):
354 + type: Literal["citation-start"] = "citation-start"
355 +
356 +
357 +class CitationEndEvent(_sdk.CitationEndEvent):
358 + type: Literal["citation-end"] = "citation-end"
359 +
360 +
361 +class MessageEndEvent(_sdk.ChatMessageEndEvent):
362 + type: Literal["message-end"] = "message-end"
vllm/entrypoints/cohere/serving.py new
+1678
@@ -0,0 +1,1678 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +# Adapted from
4 +# https://github.com/vllm-project/vllm/blob/main/vllm/entrypoints/anthropic/serving.py
5 +"""Cohere Chat v2 API serving handler.
6 +
7 +Implements ``POST /cohere/v2/chat`` by translating the incoming Cohere v2 request
8 +into a standard :class:`ChatCompletionRequest` and delegating to
9 +:class:`OpenAIServingChat`. The actual prompt rendering is handled by
10 +vLLM's renderer pipeline (``vllm.renderers``):
11 +
12 +- For Cohere Command-family models, set ``--tokenizer-mode cohere`` and the
13 + :class:`vllm.renderers.cohere.CohereRenderer` will template the request
14 + via the ``cohere_melody`` library (``render_cmd3`` / ``render_cmd4``)
15 + and surface citations through the Cohere-scoped
16 + :class:`CohereChatMessage` / :class:`CohereDeltaMessage` subclasses
17 + (see :mod:`vllm.entrypoints.cohere.cohere_chat_message`).
18 +- For any other model the default Jinja-based :class:`HfRenderer` is used,
19 + and the endpoint behaves as a plain v2-shaped wrapper around chat
20 + completions.
21 +"""
22 +
23 +from __future__ import annotations
24 +
25 +import json
26 +import logging
27 +import time
28 +from collections.abc import AsyncGenerator
29 +from enum import Enum
30 +from http import HTTPStatus
31 +from typing import TYPE_CHECKING, Any, NamedTuple, cast
32 +
33 +from cohere.types import (
34 + AssistantChatMessageV2,
35 + AssistantMessageResponse,
36 + Citation,
37 + Document,
38 + SystemChatMessageV2,
39 + ToolCallV2,
40 + ToolCallV2Function,
41 + ToolChatMessageV2,
42 + UserChatMessageV2,
43 +)
44 +from fastapi import Request
45 +from pydantic import BaseModel
46 +
47 +from vllm.engine.protocol import EngineClient
48 +from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption
49 +from vllm.entrypoints.cohere.cohere_chat_message import (
50 + CitationSource as InternalCitationSource,
51 +)
52 +from vllm.entrypoints.cohere.cohere_chat_message import (
53 + CohereChatMessage,
54 +)
55 +from vllm.entrypoints.cohere.protocol import (
56 + CitationEndEvent,
57 + CitationStartEvent,
58 + CohereChatV2Request,
59 + CohereChatV2Response,
60 + CohereFinishReason,
61 + CohereUsage,
62 + CohereUsageBilledUnits,
63 + CohereUsageTokens,
64 + ContentDeltaEvent,
65 + ContentEndEvent,
66 + ContentStartEvent,
67 + MessageEndEvent,
68 + MessageStartEvent,
69 + ToolCallDeltaEvent,
70 + ToolCallEndEvent,
71 + ToolCallStartEvent,
72 + ToolPlanDeltaEvent,
73 +)
74 +from vllm.entrypoints.openai.chat_completion.protocol import (
75 + ChatCompletionRequest,
76 + ChatCompletionResponse,
77 + ChatCompletionStreamResponse,
78 + ChatCompletionToolsParam,
79 + ChatMessage,
80 +)
81 +from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat
82 +from vllm.entrypoints.openai.engine.protocol import (
83 + ErrorResponse,
84 + JsonSchemaResponseFormat,
85 + ResponseFormat,
86 + StreamOptions,
87 +)
88 +from vllm.entrypoints.openai.models.serving import OpenAIServingModels
89 +from vllm.entrypoints.serve.utils.api_utils import sanitize_message
90 +from vllm.entrypoints.serve.utils.request_logger import RequestLogger
91 +from vllm.parser.abstract_parser import Parser
92 +from vllm.renderers.cohere import MESSAGES_CITATIONS_KEY, POSITION_TO_SOURCE_KEY
93 +
94 +if TYPE_CHECKING:
95 + from vllm.renderers.online_renderer import OnlineRenderer
96 +
97 +logger = logging.getLogger(__name__)
98 +
99 +
100 +# ---------------------------------------------------------------------------
101 +# SSE helpers
102 +# ---------------------------------------------------------------------------
103 +
104 +
105 +def _sse(data: str) -> str:
106 + """Wrap a JSON payload in a Server-Sent Event frame.
107 +
108 + Cohere's stream uses bare ``data:`` lines (no ``event:`` prefix); each
109 + JSON object's ``type`` field carries the event discriminator.
110 + """
111 + return f"data: {data}\n\n"
112 +
113 +
114 +# Cohere v2 SSE stream terminator. Declared by the upstream OpenAPI spec
115 +# (``x-fern-streaming.terminator: "[DONE]"`` on the ``/v2/chat`` stream
116 +# operation) and observed by the cohere-python / Fern-generated clients
117 +# (the Python SDK breaks its read loop on ``_sse.data == "[DONE]"``).
118 +# Must be emitted after the closing ``message-end`` event.
119 +_DONE_FRAME = _sse("[DONE]")
120 +
121 +
122 +def _emit(event: BaseModel) -> str:
123 + """Serialize a typed stream event into an SSE frame.
124 +
125 + The typed event classes in ``vllm.entrypoints.cohere.protocol``
126 + (``MessageStartEvent``, ``ContentStartEvent``, ``CitationStartEvent``,
127 + ...) bake the wire-format ``type`` field into the model definition,
128 + so a plain ``model_dump_json()`` already carries the discriminator.
129 + """
130 + return _sse(event.model_dump_json(exclude_none=True))
131 +
132 +
133 +class ContentBlockType(str, Enum):
134 + """Wire-format / internal discriminator for chat content blocks.
135 +
136 + ``THINKING`` and ``TEXT`` are the documented Cohere v2 content-block
137 + type discriminators on the wire. ``TOOL_CALL`` is reserved for the
138 + internal stream state machine (see :class:`_StreamState`) when a
139 + tool call is the currently open block; it is never serialized.
140 + """
141 +
142 + THINKING = "thinking"
143 + TEXT = "text"
144 + TOOL_CALL = "tool_call"
145 +
146 +
147 +# Mapping of vLLM/OpenAI finish reasons to Cohere's enum.
148 +_FINISH_REASON_MAP: dict[str | None, CohereFinishReason] = {
149 + "stop": "COMPLETE",
150 + "length": "MAX_TOKENS",
151 + "tool_calls": "TOOL_CALL",
152 + "stop_sequence": "STOP_SEQUENCE",
153 + "error": "ERROR",
154 + None: "COMPLETE",
155 +}
156 +
157 +
158 +def _map_finish_reason(reason: str | None) -> CohereFinishReason:
159 + return _FINISH_REASON_MAP.get(reason, "COMPLETE")
160 +
161 +
162 +class _CitablePosition(NamedTuple):
163 + """One entry in melody's citation numbering scheme, resolved to wire shape.
164 +
165 + Populated by :meth:`CohereServingChatV2._walk_citable_positions`
166 + (the single source of truth for melody's numbering rule) and
167 + consumed by both the inbound
168 + (:meth:`CohereServingChatV2._build_doc_id_to_prompt_position`,
169 + which drives citation resolution on the request side) and outbound
170 + (:meth:`CohereServingChatV2._build_position_to_source`, whose
171 + result is forwarded to the reasoning parser via
172 + ``chat_template_kwargs[POSITION_TO_SOURCE_KEY]``) helpers -- so
173 + the numbering rule only lives in one place.
174 +
175 + ``ids`` lists every id under which this position is addressable:
176 + top-level documents contribute both their explicit id (if set) and
177 + the ``doc_{idx}`` synthetic fallback, in that order so
178 + ``setdefault``-style resolution prefers the explicit id;
179 + tool-message documents contribute their single ``document.id``.
180 + ``source`` is the fully-resolved wire-shape source (``type`` /
181 + ``id`` / ``document`` or ``tool_output`` payload) for the
182 + outbound direction.
183 + """
184 +
185 + bucket: int
186 + result_idx: int
187 + ids: tuple[str, ...]
188 + source: InternalCitationSource
189 +
190 +
191 +# ---------------------------------------------------------------------------
192 +# Serving class
193 +# ---------------------------------------------------------------------------
194 +
195 +
196 +class CohereServingChatV2(OpenAIServingChat):
197 + """Handler for the Cohere Chat v2 API (``POST /cohere/v2/chat``).
198 +
199 + The handler is intentionally thin: it converts the v2 request into a
200 + :class:`ChatCompletionRequest` (preserving Cohere-specific fields such
201 + as ``documents``, ``safety_mode`` and ``citation_options`` via
202 + ``chat_template_kwargs``) and delegates to the underlying chat
203 + completion machinery. All Cohere-specific templating happens in the
204 + renderer (:class:`vllm.renderers.cohere.CohereRenderer`) when the
205 + engine is started with ``--tokenizer-mode cohere``.
206 +
207 + Citations flow through ``CohereChatMessage`` / ``CohereDeltaMessage``
208 + subclasses (see :mod:`vllm.entrypoints.cohere.cohere_chat_message`) so
209 + the OpenAI-shared :class:`ChatMessage` / :class:`DeltaMessage` keep
210 + their declared schemas unchanged. Non-streaming responses pick up
211 + the subclass by overriding :meth:`_create_chat_message`; streaming
212 + deltas are emitted directly by the Cohere reasoning parser.
213 + """
214 +
215 + def __init__(
216 + self,
217 + engine_client: EngineClient,
218 + models: OpenAIServingModels,
219 + response_role: str,
220 + *,
221 + online_renderer: OnlineRenderer,
222 + request_logger: RequestLogger | None,
223 + chat_template: str | None,
224 + chat_template_content_format: ChatTemplateContentFormatOption,
225 + return_tokens_as_token_ids: bool = False,
226 + reasoning_parser: str = "",
227 + enable_auto_tools: bool = False,
228 + tool_parser: str | None = None,
229 + enable_prompt_tokens_details: bool = False,
230 + enable_force_include_usage: bool = False,
231 + default_chat_template_kwargs: dict[str, Any] | None = None,
232 + is_reasoning_model: bool = True,
233 + ) -> None:
234 + super().__init__(
235 + engine_client=engine_client,
236 + models=models,
237 + response_role=response_role,
238 + online_renderer=online_renderer,
239 + request_logger=request_logger,
240 + chat_template=chat_template,
241 + chat_template_content_format=chat_template_content_format,
242 + return_tokens_as_token_ids=return_tokens_as_token_ids,
243 + reasoning_parser=reasoning_parser,
244 + enable_auto_tools=enable_auto_tools,
245 + tool_parser=tool_parser,
246 + enable_prompt_tokens_details=enable_prompt_tokens_details,
247 + enable_force_include_usage=enable_force_include_usage,
248 + default_chat_template_kwargs=default_chat_template_kwargs,
249 + )
250 + # Controls how the assistant's chain-of-thought is surfaced on
251 + # turns that also contain tool calls.
252 + #
253 + # - ``True`` (default): the model is a reasoning Command-family
254 + # model; reasoning is always surfaced as a ``thinking`` content
255 + # block (non-streaming) or as ``content-start`` / ``content-
256 + # delta`` events for a thinking block (streaming), regardless of
257 + # whether tool calls also appear.
258 + # - ``False``: the model is an older non-reasoning Command model
259 + # that uses Cohere's ``tool_plan`` field for its chain-of-
260 + # thought before tool calls; reasoning is surfaced as
261 + # ``tool_plan`` (non-streaming) or as ``tool-plan-delta`` events
262 + # (streaming) on tool-call turns, and the thinking content block
263 + # is dropped.
264 + #
265 + # TODO: replace this manual flag with automatic detection from the
266 + # model's capabilities config once that exists.
267 + self._is_reasoning_model = is_reasoning_model
268 +
269 + # ------------------------------------------------------------------
270 + # Public entry point
271 + # ------------------------------------------------------------------
272 +
273 + async def create_chat_v2(
274 + self,
275 + request: CohereChatV2Request,
276 + raw_request: Request | None = None,
277 + ) -> AsyncGenerator[str, None] | CohereChatV2Response | ErrorResponse:
278 + """Implements ``POST /cohere/v2/chat``."""
279 + if logger.isEnabledFor(logging.DEBUG):
280 + logger.debug(
281 + "Received Cohere v2 chat request %s", request.model_dump_json()
282 + )
283 +
284 + chat_req = self._convert_v2_to_chat_completion(request)
285 + if logger.isEnabledFor(logging.DEBUG):
286 + logger.debug(
287 + "Converted Cohere v2 -> ChatCompletion: %s",
288 + chat_req.model_dump_json(),
289 + )
290 +
291 + generator = await self.create_chat_completion(chat_req, raw_request)
292 +
293 + match generator:
294 + case ErrorResponse():
295 + return generator
296 + case ChatCompletionResponse():
297 + return self._chat_completion_to_v2(generator, request)
298 + case _:
299 + return self._chat_completion_stream_to_v2(generator, request)
300 +
301 + # ==================================================================
302 + # Request conversion: Cohere V2 -> ChatCompletionRequest
303 + # ==================================================================
304 +
305 + @classmethod
306 + def _convert_v2_to_chat_completion(
307 + cls, request: CohereChatV2Request
308 + ) -> ChatCompletionRequest:
309 + openai_messages: list[dict[str, Any]] = []
310 + cls._convert_messages(request.messages, openai_messages)
311 +
312 + chat_req = cls._build_base_chat_completion(request, openai_messages)
313 + cls._apply_streaming_options(chat_req, request)
314 + cls._apply_response_format(chat_req, request)
315 + cls._apply_tools(chat_req, request)
316 + cls._apply_tool_choice(chat_req, request)
317 + cls._apply_cohere_template_kwargs(chat_req, request)
318 + return chat_req
319 +
320 + @classmethod
321 + def _convert_messages(
322 + cls,
323 + messages: list,
324 + openai_messages: list[dict[str, Any]],
325 + ) -> None:
326 + for msg in messages:
327 + match msg:
328 + case SystemChatMessageV2():
329 + openai_messages.append(
330 + {
331 + "role": "system",
332 + "content": cls._coerce_text_content(msg.content),
333 + }
334 + )
335 + case UserChatMessageV2():
336 + openai_messages.append(cls._convert_user_message(msg))
337 + case AssistantChatMessageV2():
338 + openai_messages.append(cls._convert_assistant_message(msg))
339 + case ToolChatMessageV2():
340 + openai_messages.append(cls._convert_tool_message(msg))
341 + case _: # pragma: no cover - guarded by Pydantic discriminator
342 + raise ValueError(f"Unsupported Cohere v2 message: {msg!r}")
343 +
344 + @staticmethod
345 + def _coerce_text_content(content: str | list[Any]) -> str:
346 + if isinstance(content, str):
347 + return content
348 + parts: list[str] = []
349 + for block in content:
350 + text = getattr(block, "text", None)
351 + if text:
352 + parts.append(text)
353 + return "".join(parts)
354 +
355 + @classmethod
356 + def _convert_user_message(cls, msg: UserChatMessageV2) -> dict[str, Any]:
357 + if isinstance(msg.content, str):
358 + return {"role": "user", "content": msg.content}
359 +
360 + # Discriminate by ``type`` rather than isinstance so we don't have
361 + # to import every individual ``*Content`` variant from the SDK -
362 + # the union (``UserMessageV2Content``) covers both ``TextContent``
363 + # and ``ImageUrlContent`` and both expose ``type`` as a Literal.
364 + content_parts: list[dict[str, Any]] = []
365 + for block in msg.content:
366 + if block.type == "text":
367 + content_parts.append({"type": "text", "text": block.text})
368 + elif block.type == "image_url":
369 + image_url: dict[str, Any] = {"url": block.image_url.url}
370 + if getattr(block.image_url, "detail", None) is not None:
371 + image_url["detail"] = block.image_url.detail
372 + content_parts.append(
373 + {
374 + "type": "image_url",
375 + "image_url": image_url,
376 + }
377 + )
378 + if len(content_parts) == 1 and content_parts[0]["type"] == "text":
379 + return {"role": "user", "content": content_parts[0]["text"]}
380 + return {"role": "user", "content": content_parts}
381 +
382 + @classmethod
383 + def _convert_assistant_message(cls, msg: AssistantChatMessageV2) -> dict[str, Any]:
384 + out: dict[str, Any] = {"role": "assistant"}
385 +
386 + # Cohere splits reasoning out into ``thinking`` content blocks. We
387 + # collapse them back into the OpenAI ``reasoning`` field for the
388 + # downstream chat template, while text blocks become ``content``.
389 + text_parts: list[str] = []
390 + thinking_parts: list[str] = []
391 + if isinstance(msg.content, str):
392 + text_parts.append(msg.content)
393 + elif msg.content is not None:
394 + for block in msg.content:
395 + if block.type == "text":
396 + text_parts.append(block.text)
397 + elif block.type == "thinking":
398 + thinking_parts.append(block.thinking)
399 +
400 + # ``tool_plan`` is Cohere's chain-of-thought emitted alongside tool
401 + # calls; preserve it as reasoning so templates that expect a
402 + # planning block still see it.
403 + if msg.tool_plan:
404 + thinking_parts.append(msg.tool_plan)
405 +
406 + if text_parts:
407 + out["content"] = "".join(text_parts)
408 + if thinking_parts:
409 + out["reasoning"] = "".join(thinking_parts)
410 +
411 + if msg.tool_calls:
412 + out["tool_calls"] = [
413 + {
414 + "id": tc.id,
415 + "type": "function",
416 + "function": {
417 + "name": (tc.function.name if tc.function else "") or "",
418 + "arguments": (tc.function.arguments if tc.function else None)
419 + or "{}",
420 + },
421 + }
422 + for tc in msg.tool_calls
423 + ]
424 + return out
425 +
426 + @classmethod
427 + def _convert_tool_message(cls, msg: ToolChatMessageV2) -> dict[str, Any]:
428 + if isinstance(msg.content, str):
429 + return {
430 + "role": "tool",
431 + "tool_call_id": msg.tool_call_id,
432 + "content": msg.content,
433 + }
434 +
435 + # When the tool result is text-only, flatten to a string for
436 + # maximum compatibility with standard chat templates. When it
437 + # includes documents, preserve them as structured content parts
438 + # so the cohere renderer can surface them as grounding sources
439 + # (the :class:`CohereRenderer` understands
440 + # ``{type: document, document: {...}}`` blocks). Non-cohere
441 + # renderers may not honor document blocks, but that matches
442 + # the broader "documents are no-op for OSS models" contract
443 + # documented on this endpoint.
444 + #
445 + # Tool message content uses ``ToolMessageV2Content``, which is a
446 + # union of ``TextToolContent`` and ``DocumentToolContent`` -
447 + # distinct from the user-message text/image union. We discriminate
448 + # on the ``type`` literal so we don't have to import each variant.
449 + has_documents = any(block.type == "document" for block in msg.content)
450 + if not has_documents:
451 + text = "\n".join(
452 + block.text for block in msg.content if block.type == "text"
453 + )
454 + return {
455 + "role": "tool",
456 + "tool_call_id": msg.tool_call_id,
457 + "content": text,
458 + }
459 +
460 + parts: list[dict[str, Any]] = []
461 + for block in msg.content:
462 + if block.type == "text":
463 + parts.append({"type": "text", "text": block.text})
464 + elif block.type == "document":
465 + parts.append(
466 + {
467 + "type": "document",
468 + "document": block.document.model_dump(exclude_none=True),
469 + }
470 + )
471 + return {
472 + "role": "tool",
473 + "tool_call_id": msg.tool_call_id,
474 + "content": parts,
475 + }
476 +
477 + @classmethod
478 + def _build_base_chat_completion(
479 + cls,
480 + request: CohereChatV2Request,
481 + openai_messages: list[dict[str, Any]],
482 + ) -> ChatCompletionRequest:
483 + return ChatCompletionRequest(
484 + model=request.model,
485 + messages=openai_messages,
486 + max_tokens=request.max_tokens,
487 + max_completion_tokens=request.max_tokens,
488 + stop=request.stop_sequences,
489 + temperature=request.temperature,
490 + top_p=request.p,
491 + top_k=request.k,
492 + seed=request.seed,
493 + frequency_penalty=request.frequency_penalty,
494 + presence_penalty=request.presence_penalty,
495 + logprobs=request.logprobs,
496 + priority=request.priority or 0,
497 + kv_transfer_params=request.kv_transfer_params,
498 + chat_template_kwargs=request.chat_template_kwargs,
499 + )
500 +
501 + @classmethod
502 + def _apply_streaming_options(
503 + cls,
504 + chat_req: ChatCompletionRequest,
505 + request: CohereChatV2Request,
506 + ) -> None:
507 + if request.stream:
508 + chat_req.stream = True
509 + chat_req.stream_options = StreamOptions.model_validate(
510 + {"include_usage": True}
511 + )
512 +
513 + @classmethod
514 + def _apply_response_format(
515 + cls,
516 + chat_req: ChatCompletionRequest,
517 + request: CohereChatV2Request,
518 + ) -> None:
519 + rf = request.response_format
520 + if rf is None or rf.type == "text":
521 + return
522 + chat_req.response_format = ResponseFormat(
523 + type="json_schema" if rf.json_schema else "json_object",
524 + json_schema=(
525 + JsonSchemaResponseFormat(
526 + name="cohere_v2_json_schema",
527 + json_schema=rf.json_schema,
528 + )
529 + if rf.json_schema
530 + else None
531 + ),
532 + )
533 +
534 + @classmethod
535 + def _apply_tools(
536 + cls,
537 + chat_req: ChatCompletionRequest,
538 + request: CohereChatV2Request,
539 + ) -> None:
540 + if not request.tools:
541 + return
542 + # Cohere's ``strict_tools`` is the spec-equivalent of OpenAI's
543 + # per-function ``strict`` flag: when true the API guarantees tool
544 + # call arguments match the declared JSON schema. We surface it in
545 + # both places so OpenAI-shaped consumers see it on the function
546 + # definition and the cohere renderer can read it back from
547 + # ``chat_template_kwargs`` for cmd3/cmd4 preamble selection.
548 + strict = bool(request.strict_tools)
549 + chat_req.tools = [
550 + ChatCompletionToolsParam.model_validate(
551 + {
552 + "type": "function",
553 + "function": {
554 + "name": tool.function.name,
555 + "description": tool.function.description,
556 + "parameters": tool.function.parameters,
557 + **({"strict": True} if strict else {}),
558 + },
559 + }
560 + )
561 + for tool in request.tools
562 + ]
563 +
564 + @classmethod
565 + def _apply_tool_choice(
566 + cls,
567 + chat_req: ChatCompletionRequest,
568 + request: CohereChatV2Request,
569 + ) -> None:
570 + # Cohere v2's ``tool_choice`` only admits ``REQUIRED`` / ``NONE``
571 + # (see :data:`CohereToolChoice`); named-tool selection is not
572 + # part of the v2 spec. Map both onto OpenAI's equivalents.
573 + if request.tool_choice == "REQUIRED":
574 + chat_req.tool_choice = "required"
575 + elif request.tool_choice == "NONE":
576 + chat_req.tool_choice = "none"
577 + elif chat_req.tools:
578 + # Mirrors Cohere's "free choice" default when tools are present.
579 + chat_req.tool_choice = "auto"
580 +
581 + @classmethod
582 + def _apply_cohere_template_kwargs(
583 + cls,
584 + chat_req: ChatCompletionRequest,
585 + request: CohereChatV2Request,
586 + ) -> None:
587 + """Forward Cohere-specific request fields into ``chat_template_kwargs``.
588 +
589 + The :class:`vllm.renderers.cohere.CohereRenderer` consumes these
590 + kwargs to drive ``cohere_melody.render_cmd3`` / ``render_cmd4``.
591 + Other renderers ignore unknown kwargs, so this is also a no-op for
592 + non-cohere ``--tokenizer-mode`` settings.
593 + """
594 + kwargs = dict(chat_req.chat_template_kwargs or {})
595 +
596 + if request.documents:
597 + documents: list[dict[str, Any]] = []
598 + for idx, doc in enumerate(request.documents):
599 + if isinstance(doc, str):
600 + documents.append({"id": f"doc_{idx}", "data": {"text": doc}})
601 + else:
602 + documents.append(
603 + {
604 + "id": doc.id or f"doc_{idx}",
605 + "data": (
606 + doc.data
607 + if isinstance(doc.data, dict)
608 + else {"text": doc.data}
609 + ),
610 + }
611 + )
612 + kwargs.setdefault("documents", documents)
613 +
614 + if request.safety_mode is not None:
615 + kwargs.setdefault("safety_mode", str(request.safety_mode).lower())
616 + if request.citation_options is not None:
617 + kwargs.setdefault(
618 + "citation_options",
619 + request.citation_options.model_dump(exclude_none=True),
620 + )
621 + if request.thinking is not None:
622 + kwargs.setdefault(
623 + "thinking",
624 + request.thinking.model_dump(exclude_none=True),
625 + )
626 +
627 + # ``strict_tools`` is intentionally NOT forwarded here. It's a
628 + # decoder-guidance flag that ``_apply_tools`` already maps onto
629 + # per-function OpenAI ``strict: true`` on ``chat_req.tools``;
630 + # ``cohere_melody.render_cmd3``/``render_cmd4`` have no
631 + # ``strict_tools`` knob, and the renderer's residual-forward path
632 + # would otherwise surface it as a Jinja variable ``{{ strict_tools }}``
633 + # that templates have no defined use for.
634 +
635 + # Citations on assistant messages in the request don't fit
636 + # OpenAI's ``ChatMessage`` / ``ConversationMessage`` schemas, so
637 + # we forward them under the reserved ``MESSAGES_CITATIONS_KEY``
638 + # chat_template_kwargs entry keyed by request-message index.
639 + # See ``_sdk_citation_to_melody`` for the id-resolution rules.
640 + citations_by_index = cls._collect_message_citations(request)
641 + if citations_by_index:
642 + kwargs.setdefault(MESSAGES_CITATIONS_KEY, citations_by_index)
643 +
644 + # Forward the ``(bucket, result_idx) -> resolved wire source``
645 + # map to the reasoning parser via chat_template_kwargs. The
646 + # parser (constructed with the same kwargs) uses it to attach
647 + # real document ids / types / payloads to every source it
648 + # emits, so citation deltas already carry SDK-shape sources by
649 + # the time they reach the streaming loop. See
650 + # ``_build_position_to_source`` for the numbering rule and
651 + # ``_melody_sources_to_vllm`` in the reasoning parser for the
652 + # consumer.
653 + position_to_source = cls._build_position_to_source(request)
654 + if position_to_source:
655 + kwargs.setdefault(POSITION_TO_SOURCE_KEY, position_to_source)
656 +
657 + if kwargs:
658 + chat_req.chat_template_kwargs = kwargs
659 +
660 + @classmethod
661 + def _collect_message_citations(
662 + cls, request: CohereChatV2Request
663 + ) -> dict[int, list[dict[str, Any]]]:
664 + """Extract ``AssistantChatMessageV2.citations`` keyed by message index.
665 +
666 + Returns an empty dict when no assistant message in the request
667 + carries any resolvable citations, so callers can use truthiness
668 + as the "should I emit the ``MESSAGES_CITATIONS_KEY`` entry?"
669 + test. Citations whose sources can't be resolved are dropped
670 + individually; messages that end up with zero surviving
671 + citations are omitted from the result.
672 + """
673 + doc_positions = cls._build_doc_id_to_prompt_position(
674 + request.messages, request.documents
675 + )
676 + out: dict[int, list[dict[str, Any]]] = {}
677 + for idx, msg in enumerate(request.messages):
678 + if not isinstance(msg, AssistantChatMessageV2) or not msg.citations:
679 + continue
680 + resolved: list[dict[str, Any]] = []
681 + for c in msg.citations:
682 + converted = cls._sdk_citation_to_melody(c, doc_positions)
683 + if converted is not None:
684 + resolved.append(converted)
685 + if resolved:
686 + out[idx] = resolved
687 + return out
688 +
689 + @classmethod
690 + def _walk_citable_positions(
691 + cls, messages: list, documents: list | None
692 + ) -> list[_CitablePosition]:
693 + """Single-pass implementation of melody's citation numbering rule.
694 +
695 + Mirrors ``PromptRenderIds::from_messages`` in
696 + ``melody/src/templating/util.rs``. Emits one
697 + :class:`_CitablePosition` per document-bearing slot melody
698 + assigns while rendering the prompt:
699 +
700 + * **Bucket numbering.** Bucket ``0`` is reserved for the
701 + top-level ``documents`` array when present, and each unique
702 + ``tool_call_id`` claims the next integer on first-seen basis
703 + -- registered from whichever of the assistant's
704 + ``tool_calls[].id`` or the tool message's ``tool_call_id``
705 + appears first in the history.
706 + * **Per-bucket result-index slots.** Consumed by *every*
707 + content item, not just documents -- a text block occupies a
708 + slot too, and the source at that slot carries an id-less
709 + ``tool_output`` payload synthesized via
710 + :meth:`_text_to_tool_output_payload`, matching the cohere
711 + api's behavior for text-only tool content. Strings and
712 + empty contents each consume exactly one slot because the
713 + cohere renderer wraps them into a single text block before
714 + handing the message off to melody. Multiple tool messages
715 + sharing a ``tool_call_id`` accumulate into the same bucket,
716 + so later docs' slots are offset by the sum of previous
717 + messages' content lengths.
718 + * **Doc ids attached to each position.** Top-level docs
719 + contribute both the explicit ``id`` (if set) and the
720 + ``doc_{idx}`` fallback used for inbound addressing; tool
721 + document blocks contribute their single ``document.id``
722 + when present. The ``ids`` tuple is empty for slots that
723 + are only reachable positionally (text-only tool content,
724 + docs without a client-provided id) -- those slots still
725 + have a fully-populated :class:`InternalCitationSource`,
726 + just no client-visible id to cite them by.
727 +
728 + Ids alone key the inbound helper
729 + (:meth:`_build_doc_id_to_prompt_position`, which drives
730 + request-side citation resolution); the resolved
731 + :class:`InternalCitationSource` (with ``type`` / ``id`` /
732 + payload) keys the outbound helper
733 + (:meth:`_build_position_to_source`, whose result is forwarded
734 + to the reasoning parser via
735 + ``chat_template_kwargs[POSITION_TO_SOURCE_KEY]``).
736 +
737 + Cohere's on-wire semantics -- ``ToolSource.id`` and
738 + ``DocumentSource.id`` both carry a document id, with ``type``
739 + acting only as a payload-shape hint -- are the invariant this
740 + walk is mirroring.
741 + """
742 + out: list[_CitablePosition] = []
743 + tool_call_id_to_bucket: dict[str, int] = {}
744 + # Next unassigned ``tool_result_index`` per bucket. Advances
745 + # by ``len(content)`` (or 1 for string / missing content) for
746 + # every tool message, whether or not the content had any doc
747 + # ids in it -- matching melody's per-slot accounting.
748 + bucket_result_len: dict[int, int] = {}
749 + next_bucket = 0
750 +
751 + if documents:
752 + for idx, doc in enumerate(documents):
753 + # Every top-level doc is addressable inbound by a
754 + # synthetic ``doc_{idx}`` fallback (so clients can cite
755 + # positional docs). On the *outbound* wire we prefer
756 + # the client's explicit id when present, and omit the
757 + # ``id`` field entirely when it's not -- matching
758 + # cohere's ``DocumentSource`` where ``id`` is optional
759 + # and preserving "no id in, no id out" symmetry.
760 + fallback_id = f"doc_{idx}"
761 + if isinstance(doc, str):
762 + payload: dict[str, Any] = {"text": doc}
763 + explicit_id: str | None = None
764 + else:
765 + payload = (
766 + dict(doc.data)
767 + if isinstance(doc.data, dict)
768 + else {"text": doc.data}
769 + )
770 + explicit_id = doc.id or None
771 + if explicit_id:
772 + payload.setdefault("id", explicit_id)
773 + # Register explicit id first so it wins over the
774 + # fallback for ``setdefault``-style inbound lookup.
775 + ids: tuple[str, ...] = (
776 + (explicit_id, fallback_id) if explicit_id else (fallback_id,)
777 + )
778 + out.append(
779 + _CitablePosition(
780 + bucket=0,
781 + result_idx=idx,
782 + ids=ids,
783 + source=InternalCitationSource(
784 + type="document", id=explicit_id, document=payload
785 + ),
786 + )
787 + )
788 + next_bucket = 1
789 +
790 + def register_bucket(tool_call_id: str) -> int:
791 + nonlocal next_bucket
792 + bucket = tool_call_id_to_bucket.get(tool_call_id)
793 + if bucket is None:
794 + bucket = next_bucket
795 + tool_call_id_to_bucket[tool_call_id] = bucket
796 + next_bucket += 1
797 + return bucket
798 +
799 + for msg in messages:
800 + if isinstance(msg, ToolChatMessageV2):
801 + if not msg.tool_call_id:
802 + continue
803 + bucket = register_bucket(msg.tool_call_id)
804 + base = bucket_result_len.get(bucket, 0)
805 + if isinstance(msg.content, list):
806 + for offset, block in enumerate(msg.content):
807 + info = cls._tool_content_block_to_source(block)
808 + if info is None:
809 + continue
810 + # Only real document ids get registered in the
811 + # inbound-facing ``ids`` tuple; text-only tool
812 + # content produces a source with ``id=None``
813 + # to match the cohere api -- the text is
814 + # exposed as ``tool_output.content`` but
815 + # there's no id for clients to cite by.
816 + pos_ids: tuple[str, ...] = (info.id,) if info.id else ()
817 + out.append(
818 + _CitablePosition(
819 + bucket=bucket,
820 + result_idx=base + offset,
821 + ids=pos_ids,
822 + source=info,
823 + )
824 + )
825 + bucket_result_len[bucket] = base + len(msg.content)
826 + else:
827 + # A string or missing content is wrapped by the
828 + # cohere renderer into a single text block before
829 + # melody sees it, so it consumes one slot. For
830 + # string content, emit a synthetic id-less tool
831 + # source carrying the text as its payload so a
832 + # model citation against this slot still has
833 + # something to attach on the wire (matching the
834 + # id-less tool_output behavior for text-only
835 + # ``ToolChatMessageV2`` content blocks). Missing
836 + # content still advances the slot cursor but
837 + # produces no source -- a citation against it
838 + # would be dropped by ``_to_wire_citation``.
839 + if isinstance(msg.content, str):
840 + payload = cls._text_to_tool_output_payload(msg.content)
841 + out.append(
842 + _CitablePosition(
843 + bucket=bucket,
844 + result_idx=base,
845 + ids=(),
846 + source=InternalCitationSource(
847 + type="tool", id=None, tool_output=payload
848 + ),
849 + )
850 + )
851 + bucket_result_len[bucket] = base + 1
852 + elif isinstance(msg, AssistantChatMessageV2) and msg.tool_calls:
853 + for tc in msg.tool_calls:
854 + if tc.id:
855 + register_bucket(tc.id)
856 +
857 + return out
858 +
859 + @classmethod
860 + def _build_doc_id_to_prompt_position(
861 + cls, messages: list, documents: list | None
862 + ) -> dict[str, tuple[int, int]]:
863 + """Map every citable document id to its ``(bucket, result_index)``.
864 +
865 + Inbound helper: converts request-side citation ``source.id``
866 + refs into melody's numeric ``FilterCitation`` addressing (see
867 + :meth:`_sdk_citation_to_melody`). First-seen-id wins so an
868 + explicit id on a top-level document takes precedence over the
869 + ``doc_{idx}`` fallback that :meth:`_walk_citable_positions`
870 + also emits for the same slot.
871 +
872 + Thin wrapper over :meth:`_walk_citable_positions`; the
873 + numbering rule lives there.
874 + """
875 + doc_positions: dict[str, tuple[int, int]] = {}
876 + for pos in cls._walk_citable_positions(messages, documents):
877 + for pos_id in pos.ids:
878 + doc_positions.setdefault(pos_id, (pos.bucket, pos.result_idx))
879 + return doc_positions
880 +
881 + @staticmethod
882 + def _sdk_citation_to_melody(
883 + citation: Any,
884 + doc_positions: dict[str, tuple[int, int]],
885 + ) -> dict[str, Any] | None:
886 + """Convert a Cohere SDK ``Citation`` to melody's ``FilterCitation``.
887 +
888 + Melody's ``Source`` uses a two-level numeric address:
889 + ``tool_call_index`` picks a bucket (``0`` = top-level
890 + ``documents`` when present, ``1..N`` = tool results in history
891 + order) and ``tool_result_indices`` picks positions inside that
892 + bucket. The Cohere ``Citation.sources`` schema is per-document:
893 + each entry's ``id`` is the id of the specific document that
894 + grounds the citation (``type`` distinguishes payload shape, not
895 + id semantics). We resolve each source id against
896 + ``doc_positions`` and aggregate hits sharing a bucket into one
897 + melody ``Source`` so the renderer emits compact
898 + ``<co>text</co: N:[i,j]>`` markers.
899 +
900 + Returns ``None`` when any source id fails to resolve or when
901 + the input has no sources at all: emitting a citation with an
902 + empty ``sources=[]`` would render a malformed
903 + ``<co>text</co: :>`` marker, and silently attributing a
904 + citation to the wrong bucket is worse than dropping the
905 + ``<co>`` markup. This matches the cohere api's behavior.
906 +
907 + ``Source.document_ids`` is intentionally NOT populated: it is a
908 + parser-*output* lookup table (see ``PromptRenderIds`` in
909 + melody), and melody's deserializer ignores it on input.
910 +
911 + The Cohere ``type`` literal (``TEXT_CONTENT`` / ``THINKING_CONTENT``
912 + / ``PLAN``) collapses to melody's boolean ``is_thinking``.
913 + ``PLAN`` is treated as a thinking-style citation because cmd3 /
914 + cmd4 templates render both inline with the same ``<co>...</co>``
915 + markup.
916 + """
917 + raw_sources = getattr(citation, "sources", None) or []
918 + if not raw_sources:
919 + return None
920 +
921 + grouped: dict[int, list[int]] = {}
922 + for src in raw_sources:
923 + src_id = getattr(src, "id", None)
924 + if not src_id:
925 + return None
926 + position = doc_positions.get(src_id)
927 + if position is None:
928 + return None
929 + bucket, result_idx = position
930 + grouped.setdefault(bucket, []).append(result_idx)
931 +
932 + sources = [
933 + {"tool_call_index": bucket, "tool_result_indices": indices}
934 + for bucket, indices in grouped.items()
935 + ]
936 +
937 + cite_type = str(getattr(citation, "type", "") or "").upper()
938 + is_thinking = cite_type in ("THINKING_CONTENT", "PLAN")
939 + return {
940 + "start_index": getattr(citation, "start", None) or 0,
941 + "end_index": getattr(citation, "end", None) or 0,
942 + "text": getattr(citation, "text", None) or "",
943 + "sources": sources,
944 + "is_thinking": is_thinking,
945 + }
946 +
947 + # ==================================================================
948 + # Response conversion: ChatCompletion -> Cohere V2
949 + # ==================================================================
950 +
951 + def _chat_completion_to_v2(
952 + self,
953 + response: ChatCompletionResponse,
954 + request: CohereChatV2Request,
955 + ) -> CohereChatV2Response:
956 + choice = response.choices[0]
957 + msg = choice.message
958 +
959 + # Build content blocks as dicts; ``AssistantMessageResponse``
960 + # validates them into the proper discriminated union variants
961 + # (``Text/ThinkingAssistantMessageResponseContentItem``).
962 + content_blocks: list[dict[str, Any]] = []
963 + if msg.reasoning:
964 + content_blocks.append(
965 + {"type": ContentBlockType.THINKING, "thinking": msg.reasoning}
966 + )
967 + if msg.content:
968 + content_blocks.append({"type": ContentBlockType.TEXT, "text": msg.content})
969 +
970 + tool_calls: list[ToolCallV2] | None = None
971 + if msg.tool_calls:
972 + tool_calls = [
973 + ToolCallV2(
974 + id=tc.id,
975 + function=ToolCallV2Function(
976 + name=tc.function.name,
977 + arguments=tc.function.arguments,
978 + ),
979 + )
980 + for tc in msg.tool_calls
981 + ]
982 +
983 + # Cohere's ``tool_plan`` is the planning text emitted before tool
984 + # calls on older, non-reasoning Command models. For those models
985 + # we surface ``reasoning`` as ``tool_plan`` and drop the thinking
986 + # block. Reasoning Command models emit a regular thinking block
987 + # alongside tool calls, so for them we leave the thinking block
988 + # in place and never set ``tool_plan``.
989 + tool_plan: str | None = None
990 + if not self._is_reasoning_model and tool_calls and msg.reasoning:
991 + tool_plan = msg.reasoning
992 + content_blocks = [
993 + blk
994 + for blk in content_blocks
995 + if blk.get("type") != ContentBlockType.THINKING
996 + ]
997 +
998 + assistant_msg = AssistantMessageResponse(
999 + content=content_blocks or None,
1000 + tool_calls=tool_calls,
1001 + tool_plan=tool_plan,
1002 + citations=self._extract_citations_if_any(msg),
1003 + )
1004 +
1005 + usage = self._build_usage(response)
1006 +
1007 + return CohereChatV2Response(
1008 + id=response.id or f"chat_{int(time.time() * 1000)}",
1009 + finish_reason=_map_finish_reason(choice.finish_reason),
1010 + message=assistant_msg,
1011 + usage=usage,
1012 + kv_transfer_params=response.kv_transfer_params,
1013 + )
1014 +
1015 + def _create_chat_message(self, *args: Any, **kwargs: Any) -> ChatMessage:
1016 + """Route response construction through the citation-carrying subclass.
1017 +
1018 + Overrides :meth:`OpenAIServingChat._create_chat_message` so every
1019 + non-streaming construction site in the base full-generator
1020 + produces a :class:`CohereChatMessage`. That lets
1021 + :meth:`_finalize_response_message` below set citations on
1022 + ``message.citations`` directly, without relying on
1023 + ``extra="allow"`` attribute injection.
1024 + """
1025 + return CohereChatMessage(*args, **kwargs)
1026 +
1027 + def _finalize_response_message(
1028 + self,
1029 + message: ChatMessage,
1030 + *,
1031 + parser: Parser | None,
1032 + ) -> ChatMessage:
1033 + """Copy grounding citations off the reasoning parser onto the message.
1034 +
1035 + The Cohere reasoning parser
1036 + (:mod:`vllm.reasoning.cohere_command_reasoning_parser`) caches the
1037 + citations produced by its most recent unary ``extract_reasoning``
1038 + call on ``last_unary_citations``. We surface them here on
1039 + :class:`CohereChatMessage` so downstream response conversion can
1040 + pick them up without the base :class:`OpenAIServingChat` having to
1041 + know about citations.
1042 + """
1043 + # ``_create_chat_message`` above guarantees the concrete type at
1044 + # runtime. ``cast`` narrows the declared ``ChatMessage`` for
1045 + # mypy without adding a runtime check we don't need: even in the
1046 + # invariant-violated case (a subclass reset the factory back to
1047 + # plain ``ChatMessage``), ``OpenAIBaseModel``'s ``extra="allow"``
1048 + # config lets the ``.citations`` write survive into
1049 + # ``model_dump`` via the extras bucket, so the wire is correct
1050 + # either way.
1051 + message = cast(CohereChatMessage, message)
1052 + citations = getattr(
1053 + getattr(parser, "reasoning_parser", None),
1054 + "last_unary_citations",
1055 + None,
1056 + )
1057 + if citations:
1058 + message.citations = citations
1059 + return message
1060 +
1061 + def _extract_citations_if_any(self, msg: Any) -> list[Citation] | None:
1062 + """Coerce ``CohereChatMessage.citations`` into the Cohere v2 wire shape.
1063 +
1064 + :class:`CohereChatMessage` carries a
1065 + ``citations: list[vllm...Citation] | None`` field populated by
1066 + :meth:`_finalize_response_message` (which reads it off the
1067 + reasoning parser's ``last_unary_citations`` cache). Sources are
1068 + already fully resolved by the parser (see
1069 + :func:`_melody_sources_to_vllm` in
1070 + :mod:`vllm.reasoning.cohere_command_reasoning_parser`) using
1071 + the position map forwarded via ``chat_template_kwargs``. This
1072 + method's remaining job is:
1073 +
1074 + * Drop citations whose sources all failed to resolve (empty
1075 + ``sources`` list) -- fail-closed policy matching the inbound
1076 + ``_sdk_citation_to_melody``.
1077 + * Rewrite ``THINKING_CONTENT`` -> ``PLAN`` for non-reasoning
1078 + models, since ``_chat_completion_to_v2`` surfaces
1079 + ``msg.reasoning`` as ``tool_plan`` on the same
1080 + ``is_reasoning_model`` flag.
1081 + * Coerce to :class:`cohere.types.Citation` for the wire.
1082 + """
1083 + raw = getattr(msg, "citations", None)
1084 + if not raw:
1085 + return None
1086 + out: list[Citation] = []
1087 + for c in raw:
1088 + resolved = self._to_wire_citation(c)
1089 + if resolved is not None:
1090 + out.append(resolved)
1091 + return out or None
1092 +
1093 + def _to_wire_citation(self, citation: Any) -> Citation | None:
1094 + """Coerce one parser-produced citation into the SDK wire shape.
1095 +
1096 + Accepts both :class:`InternalCitation` instances (what the
1097 + unary path receives directly from the parser via
1098 + ``last_unary_citations``) and plain ``dict`` payloads (what
1099 + the streaming path sees, because ``DeltaMessage.citations`` is
1100 + an untyped extras field on the OpenAI wire protocol and dict
1101 + shapes survive ``model_validate_json`` unchanged). Both flows
1102 + end up carrying the same key names, so a small ``_get`` helper
1103 + handles both without a shape-specific code path.
1104 + """
1105 +
1106 + def _get(obj: Any, key: str, default: Any = None) -> Any:
1107 + if isinstance(obj, dict):
1108 + return obj.get(key, default)
1109 + return getattr(obj, key, default)
1110 +
1111 + raw_sources = _get(citation, "sources") or []
1112 + wire_sources: list[dict[str, Any]] = []
1113 + for src in raw_sources:
1114 + # Sources arrive pre-resolved from the parser (or as trusted
1115 + # dicts on the streaming path). Match the cohere api's
1116 + # behavior of preserving id-less sources (e.g. text-only
1117 + # tool results): cohere's ``Source`` types treat ``id`` as
1118 + # optional, so we simply drop the field via
1119 + # ``exclude_none=True`` instead of dropping the source.
1120 + if isinstance(src, dict):
1121 + wire = {k: v for k, v in src.items() if v is not None}
1122 + if wire:
1123 + wire_sources.append(wire)
1124 + elif isinstance(src, InternalCitationSource):
1125 + wire_sources.append(src.model_dump(exclude_none=True))
1126 + if not wire_sources:
1127 + return None
1128 +
1129 + cite_type = _get(citation, "type")
1130 + if cite_type == "THINKING_CONTENT" and not self._is_reasoning_model:
1131 + # Non-reasoning Command models surface the reasoning block
1132 + # as ``tool_plan`` (see ``_chat_completion_to_v2``), so any
1133 + # ``is_thinking`` citation from melody is actually a PLAN
1134 + # citation on the wire.
1135 + cite_type = "PLAN"
1136 +
1137 + payload = {
1138 + "start": _get(citation, "start"),
1139 + "end": _get(citation, "end"),
1140 + "text": _get(citation, "text"),
1141 + "sources": wire_sources,
1142 + "type": cite_type,
1143 + }
1144 + try:
1145 + return Citation.model_validate(payload)
1146 + except Exception: # pragma: no cover - defensive
1147 + logger.warning(
1148 + "Skipping malformed citation payload: %r", payload, exc_info=True
1149 + )
1150 + return None
1151 +
1152 + @classmethod
1153 + def _build_position_to_source(
1154 + cls, request: CohereChatV2Request
1155 + ) -> dict[tuple[int, int], InternalCitationSource]:
1156 + """Invert the request-side numbering into ``(bucket, idx) -> source``.
1157 +
1158 + Outbound helper: forwarded to the reasoning parser via
1159 + ``chat_template_kwargs[POSITION_TO_SOURCE_KEY]`` so
1160 + :func:`_melody_sources_to_vllm` can attach real document ids /
1161 + types / payloads to each melody source it emits. Each entry
1162 + carries the fully-resolved wire-shape info: ``type``
1163 + (``"document"`` for the reserved top-level bucket 0, ``"tool"``
1164 + for any tool-call bucket), ``id`` (explicit doc id, the
1165 + ``doc_{idx}`` fallback, or ``None`` for id-less sources), and
1166 + ``document`` / ``tool_output`` payload -- matching the cohere
1167 + api's ``Source`` shape.
1168 +
1169 + Thin wrapper over :meth:`_walk_citable_positions`; the
1170 + numbering rule lives there.
1171 + """
1172 + return {
1173 + (pos.bucket, pos.result_idx): pos.source
1174 + for pos in cls._walk_citable_positions(request.messages, request.documents)
1175 + }
1176 +
1177 + @staticmethod
1178 + def _text_to_tool_output_payload(text: str) -> dict[str, Any]:
1179 + """Wrap a bare tool-content text into a ``tool_output`` payload.
1180 +
1181 + Matches the cohere api's handling of text-only tool content:
1182 + attempt to interpret ``text`` as a JSON object first; fall
1183 + back to ``{"content": <text>}`` so a text-only tool result
1184 + still has a structured payload we can attach to a citation
1185 + source.
1186 + """
1187 + if not text:
1188 + return {}
1189 + try:
1190 + parsed = json.loads(text)
1191 + except (json.JSONDecodeError, TypeError):
1192 + return {"content": text}
1193 + if isinstance(parsed, dict):
1194 + return parsed
1195 + return {"content": text}
1196 +
1197 + @classmethod
1198 + def _tool_content_block_to_source(
1199 + cls,
1200 + block: Any,
1201 + ) -> InternalCitationSource | None:
1202 + """Extract the wire-shape source for one tool-message content block.
1203 +
1204 + Handles both cohere ``DocumentToolContent`` and
1205 + ``TextToolContent`` blocks:
1206 +
1207 + * ``document`` blocks with an explicit id produce a
1208 + ``tool`` source keyed by that id (payload is the document's
1209 + ``data`` dict, or ``{"text": ...}`` for scalar data).
1210 + * ``document`` blocks *without* an id and ``text`` blocks
1211 + produce an id-less ``tool`` source: the cohere api emits
1212 + the payload with an empty id in this case, and cohere's
1213 + SDK treats ``ToolSource.id`` as optional -- so we simply
1214 + omit ``id`` on the wire instead of inventing a synthetic
1215 + one.
1216 +
1217 + Returns ``None`` only for shapes we don't recognise at all;
1218 + that slot still gets counted (callers advance the bucket
1219 + cursor over the block regardless) so numbering stays aligned
1220 + with melody.
1221 + """
1222 + block_type = getattr(block, "type", None)
1223 + if block_type == "document":
1224 + doc: Document | None = getattr(block, "document", None)
1225 + if doc is None:
1226 + return None
1227 + doc_id = doc.id or None
1228 + data = doc.data
1229 + if isinstance(data, dict):
1230 + payload: dict[str, Any] = dict(data)
1231 + elif data is None:
1232 + payload = {}
1233 + else:
1234 + # Match the top-level ``documents`` handling (see
1235 + # ``_apply_cohere_template_kwargs``): non-dict payloads
1236 + # get wrapped as ``{"text": ...}`` so clients get a
1237 + # renderable body instead of a bare id.
1238 + payload = {"text": data}
1239 + if doc_id:
1240 + payload.setdefault("id", doc_id)
1241 + return InternalCitationSource(
1242 + type="tool",
1243 + id=doc_id,
1244 + tool_output=payload,
1245 + )
1246 + if block_type == "text":
1247 + text = getattr(block, "text", None) or ""
1248 + payload = cls._text_to_tool_output_payload(text)
1249 + return InternalCitationSource(
1250 + type="tool",
1251 + id=None,
1252 + tool_output=payload,
1253 + )
1254 + return None
1255 +
1256 + @staticmethod
1257 + def _build_usage(response: ChatCompletionResponse) -> CohereUsage | None:
1258 + if response.usage is None:
1259 + return None
1260 + prompt = response.usage.prompt_tokens
1261 + completion = response.usage.completion_tokens or 0
1262 + cached: int | None = None
1263 + if response.usage.prompt_tokens_details is not None:
1264 + cached = response.usage.prompt_tokens_details.cached_tokens
1265 + return CohereUsage(
1266 + billed_units=CohereUsageBilledUnits(
1267 + input_tokens=prompt,
1268 + output_tokens=completion,
1269 + ),
1270 + tokens=CohereUsageTokens(
1271 + input_tokens=prompt,
1272 + output_tokens=completion,
1273 + ),
1274 + cached_tokens=cached,
1275 + )
1276 +
1277 + # ==================================================================
1278 + # Stream conversion: chat completion stream -> Cohere V2 SSE events
1279 + # ==================================================================
1280 +
1281 + async def _chat_completion_stream_to_v2(
1282 + self,
1283 + generator: AsyncGenerator[str, None],
1284 + request: CohereChatV2Request,
1285 + ) -> AsyncGenerator[str, None]:
1286 + """Translate an OpenAI-style chat completion SSE stream into Cohere's
1287 + v2 stream-event format.
1288 +
1289 + Cohere's v2 stream lifecycle is:
1290 +
1291 + message-start
1292 + [content-start, content-delta..., content-end]*
1293 + [tool-plan-delta]*
1294 + [tool-call-start, tool-call-delta..., tool-call-end]*
1295 + message-end
1296 + """
1297 + state = _StreamState()
1298 +
1299 + try:
1300 + async for item in generator:
1301 + if not item.startswith("data:"):
1302 + continue
1303 + data_str = item[len("data:") :].strip().rstrip("\n")
1304 + if not data_str:
1305 + continue
1306 + if data_str == "[DONE]":
1307 + # OpenAI's stream terminator. Fall through to the
1308 + # post-loop cleanup so we always emit ``message-end``
1309 + # even if the usage-only chunk was skipped.
1310 + break
1311 +
1312 + chunk = ChatCompletionStreamResponse.model_validate_json(data_str)
1313 + state.last_chunk_id = chunk.id
1314 +
1315 + if not state.started:
1316 + yield _emit(
1317 + MessageStartEvent(
1318 + id=chunk.id,
1319 + delta={"message": {"role": "assistant"}},
1320 + )
1321 + )
1322 + state.started = True
1323 +
1324 + # The final OpenAI chunk has no choices and only carries usage.
1325 + if not chunk.choices:
1326 + for ev in self._close_open_blocks(state):
1327 + yield ev
1328 + yield self._build_message_end_event(
1329 + chunk_id=chunk.id,
1330 + finish_reason=state.finish_reason,
1331 + usage_chunk=chunk,
1332 + )
1333 + state.ended = True
1334 + continue
1335 +
1336 + choice = chunk.choices[0]
1337 + if choice.finish_reason is not None:
1338 + state.finish_reason = choice.finish_reason
1339 +
1340 + delta = choice.delta
1341 +
1342 + # Reasoning -> thinking content block
1343 + reasoning = getattr(delta, "reasoning", None) or getattr(
1344 + delta, "reasoning_content", None
1345 + )
1346 + if reasoning:
1347 + for ev in self._handle_thinking_delta(state, reasoning):
1348 + yield ev
1349 +
1350 + if delta.content:
1351 + for ev in self._handle_text_delta(state, delta.content):
1352 + yield ev
1353 +
1354 + if delta.tool_calls:
1355 + for ev in self._handle_tool_call_deltas(state, delta.tool_calls):
1356 + yield ev
1357 +
1358 + # Citations: a Cohere-specific extension on DeltaMessage that
1359 + # the cohere renderer/parsers may populate.
1360 + delta_citations = getattr(delta, "citations", None)
1361 + if delta_citations:
1362 + for ev in self._handle_citation_deltas(state, delta_citations):
1363 + yield ev
1364 +
1365 + except Exception as exc:
1366 + logger.exception("Error converting chat completion stream to v2")
1367 + if state.started and not state.ended:
1368 + yield _sse(
1369 + json.dumps(
1370 + {
1371 + "type": "message-end",
1372 + "delta": {
1373 + "error": sanitize_message(str(exc)),
1374 + "finish_reason": "ERROR",
1375 + },
1376 + }
1377 + )
1378 + )
1379 + state.ended = True
1380 + yield _DONE_FRAME
1381 + return
1382 +
1383 + # Normal completion or ``[DONE]``: ensure ``message-end`` is always
1384 + # emitted. Upstream may close the stream without sending the final
1385 + # usage-only chunk (e.g. on shutdown, or when ``[DONE]`` is the only
1386 + # terminator); without this fallback Cohere clients would hang
1387 + # waiting for the closing event.
1388 + if state.started and not state.ended:
1389 + for ev in self._close_open_blocks(state):
1390 + yield ev
1391 + yield self._build_message_end_event(
1392 + chunk_id=state.last_chunk_id,
1393 + finish_reason=state.finish_reason,
1394 + usage_chunk=None,
1395 + )
1396 + state.ended = True
1397 +
1398 + # Stream terminator. Cohere's v2 SSE protocol ends every stream
1399 + # with ``data: [DONE]\n\n`` after ``message-end``; Fern-generated
1400 + # clients (Go/Java) and cohere-python all key their read loop off
1401 + # this sentinel.
1402 + yield _DONE_FRAME
1403 +
1404 + # -- per-delta helpers --------------------------------------------
1405 +
1406 + def _handle_thinking_delta(self, state: _StreamState, delta_text: str) -> list[str]:
1407 + # Non-reasoning Command models: emit ``tool-plan-delta`` events
1408 + # directly instead of opening a thinking content block. The
1409 + # ``tool-plan-delta`` event has no start/end pair around it.
1410 + if not self._is_reasoning_model:
1411 + events: list[str] = list(self._close_open_blocks(state))
1412 + events.append(
1413 + _emit(
1414 + ToolPlanDeltaEvent(
1415 + delta={"message": {"tool_plan": delta_text}},
1416 + )
1417 + )
1418 + )
1419 + return events
1420 +
1421 + # Reasoning model (default): open / continue a thinking block.
1422 + events = []
1423 + if state.active_block != ContentBlockType.THINKING:
1424 + events.extend(self._close_open_blocks(state))
1425 + idx = state.next_content_index()
1426 + state.active_block = ContentBlockType.THINKING
1427 + state.active_block_index = idx
1428 + events.append(
1429 + _emit(
1430 + ContentStartEvent(
1431 + index=idx,
1432 + delta={
1433 + "message": {
1434 + "content": {
1435 + "type": ContentBlockType.THINKING,
1436 + "thinking": "",
1437 + }
1438 + }
1439 + },
1440 + )
1441 + )
1442 + )
1443 + events.append(
1444 + _emit(
1445 + ContentDeltaEvent(
1446 + index=state.active_block_index,
1447 + delta={"message": {"content": {"thinking": delta_text}}},
1448 + )
1449 + )
1450 + )
1451 + return events
1452 +
1453 + def _handle_text_delta(self, state: _StreamState, delta_text: str) -> list[str]:
1454 + events: list[str] = []
1455 + if state.active_block != ContentBlockType.TEXT:
1456 + events.extend(self._close_open_blocks(state))
1457 + idx = state.next_content_index()
1458 + state.active_block = ContentBlockType.TEXT
1459 + state.active_block_index = idx
1460 + events.append(
1461 + _emit(
1462 + ContentStartEvent(
1463 + index=idx,
1464 + delta={
1465 + "message": {
1466 + "content": {
1467 + "type": ContentBlockType.TEXT,
1468 + "text": "",
1469 + }
1470 + }
1471 + },
1472 + )
1473 + )
1474 + )
1475 + events.append(
1476 + _emit(
1477 + ContentDeltaEvent(
1478 + index=state.active_block_index,
1479 + delta={"message": {"content": {"text": delta_text}}},
1480 + )
1481 + )
1482 + )
1483 + return events
1484 +
1485 + def _handle_tool_call_deltas(self, state: _StreamState, deltas: list) -> list[str]:
1486 + events: list[str] = []
1487 + for tc in deltas:
1488 + tc_index = tc.index
1489 + fn = tc.function
1490 +
1491 + if tc_index not in state.tool_calls_seen:
1492 + # New tool call. Close any open content/tool block first.
1493 + events.extend(self._close_open_blocks(state))
1494 + state.tool_calls_seen.add(tc_index)
1495 + state.active_tool_index = tc_index
1496 + state.active_block = ContentBlockType.TOOL_CALL
1497 + events.append(
1498 + _emit(
1499 + ToolCallStartEvent(
1500 + index=tc_index,
1501 + delta={
1502 + "message": {
1503 + "tool_calls": {
1504 + "id": tc.id or "",
1505 + "type": "function",
1506 + "function": {
1507 + "name": (fn.name if fn else "") or "",
1508 + "arguments": (fn.arguments if fn else None)
1509 + or "",
1510 + },
1511 + }
1512 + }
1513 + },
1514 + )
1515 + )
1516 + )
1517 + continue
1518 +
1519 + if fn and fn.arguments:
1520 + events.append(
1521 + _emit(
1522 + ToolCallDeltaEvent(
1523 + index=tc_index,
1524 + delta={
1525 + "message": {
1526 + "tool_calls": {
1527 + "function": {
1528 + "arguments": fn.arguments,
1529 + }
1530 + }
1531 + }
1532 + },
1533 + )
1534 + )
1535 + )
1536 + return events
1537 +
1538 + def _handle_citation_deltas(
1539 + self,
1540 + state: _StreamState,
1541 + citations: list,
1542 + ) -> list[str]:
1543 + """Emit ``citation-start`` / ``citation-end`` events for a delta.
1544 +
1545 + The reasoning parser populates
1546 + :class:`CohereDeltaMessage.citations` (see
1547 + :mod:`vllm.entrypoints.cohere.cohere_chat_message`) with
1548 + sources whose ``id`` / ``type`` / ``document`` / ``tool_output``
1549 + are already resolved against the request's document context
1550 + (see :func:`_melody_sources_to_vllm` in the parser). This
1551 + method only coerces to the SDK wire shape, applies the
1552 + ``THINKING_CONTENT`` -> ``PLAN`` rewrite for non-reasoning
1553 + models, and drops citations whose sources didn't resolve --
1554 + the shared :meth:`_to_wire_citation` handles all three.
1555 + """
1556 + events: list[str] = []
1557 + for c in citations:
1558 + citation = self._to_wire_citation(c)
1559 + if citation is None:
1560 + continue
1561 + idx = state.next_citation_index()
1562 + events.append(
1563 + _emit(
1564 + CitationStartEvent(
1565 + index=idx,
1566 + delta={
1567 + "message": {
1568 + "citations": citation.model_dump(exclude_none=True)
1569 + }
1570 + },
1571 + )
1572 + )
1573 + )
1574 + events.append(_emit(CitationEndEvent(index=idx)))
1575 + return events
1576 +
1577 + # -- block lifecycle helpers --------------------------------------
1578 +
1579 + def _close_open_blocks(self, state: _StreamState) -> list[str]:
1580 + """Emit ``content-end`` / ``tool-call-end`` for the currently open
1581 + block (if any) and reset the corresponding ``_StreamState`` slots.
1582 + """
1583 + events: list[str] = []
1584 + if state.active_block in (ContentBlockType.TEXT, ContentBlockType.THINKING):
1585 + events.append(_emit(ContentEndEvent(index=state.active_block_index)))
1586 + elif state.active_block == ContentBlockType.TOOL_CALL:
1587 + events.append(_emit(ToolCallEndEvent(index=state.active_tool_index)))
1588 + state.active_block = None
1589 + state.active_block_index = None
1590 + state.active_tool_index = None
1591 + return events
1592 +
1593 + def _build_message_end_event(
1594 + self,
1595 + chunk_id: str,
1596 + finish_reason: str | None,
1597 + usage_chunk: ChatCompletionStreamResponse | None = None,
1598 + ) -> str:
1599 + delta: dict[str, Any] = {
1600 + "finish_reason": _map_finish_reason(finish_reason),
1601 + }
1602 + if usage_chunk is not None and usage_chunk.usage is not None:
1603 + prompt = usage_chunk.usage.prompt_tokens
1604 + completion = usage_chunk.usage.completion_tokens or 0
1605 + usage_block: dict[str, Any] = {
1606 + "billed_units": {
1607 + "input_tokens": prompt,
1608 + "output_tokens": completion,
1609 + },
1610 + "tokens": {
1611 + "input_tokens": prompt,
1612 + "output_tokens": completion,
1613 + },
1614 + }
1615 + if usage_chunk.usage.prompt_tokens_details is not None:
1616 + cached = usage_chunk.usage.prompt_tokens_details.cached_tokens
1617 + if cached is not None:
1618 + usage_block["cached_tokens"] = cached
1619 + delta["usage"] = usage_block
1620 + return _emit(MessageEndEvent(id=chunk_id, delta=delta))
1621 +
1622 + # ==================================================================
1623 + # Helpers for the router
1624 + # ==================================================================
1625 +
1626 + @staticmethod
1627 + def create_error_response(
1628 + message: str | Exception,
1629 + err_type: str = "bad_request",
1630 + status_code: HTTPStatus = HTTPStatus.BAD_REQUEST,
1631 + param: str | None = None,
1632 + ) -> ErrorResponse:
1633 + # Override of :meth:`BaseServing.create_error_response` that uses
1634 + # Cohere-flavored defaults (``type="bad_request"``, ``code=400``)
1635 + # so the router can translate the envelope uniformly. ``param``
1636 + # is accepted for signature parity with the base class but is
1637 + # not surfaced in the Cohere wire format.
1638 + from vllm.entrypoints.openai.engine.protocol import ErrorInfo
1639 +
1640 + del param # unused; kept for signature compatibility
1641 + return ErrorResponse(
1642 + error=ErrorInfo(
1643 + message=str(message),
1644 + type=err_type,
1645 + code=int(status_code),
1646 + )
1647 + )
1648 +
1649 +
1650 +# ---------------------------------------------------------------------------
1651 +# Stream state
1652 +# ---------------------------------------------------------------------------
1653 +
1654 +
1655 +class _StreamState:
1656 + """Tracks which Cohere v2 stream block (if any) is currently open."""
1657 +
1658 + def __init__(self) -> None:
1659 + self.started: bool = False
1660 + self.ended: bool = False
1661 + self.finish_reason: str | None = None
1662 + self.last_chunk_id: str = ""
1663 + self.active_block: ContentBlockType | None = None
1664 + self.active_block_index: int | None = None
1665 + self.active_tool_index: int | None = None
1666 + self._next_index: int = 0
1667 + self._next_citation_index: int = 0
1668 + self.tool_calls_seen: set[int] = set()
1669 +
1670 + def next_content_index(self) -> int:
1671 + idx = self._next_index
1672 + self._next_index += 1
1673 + return idx
1674 +
1675 + def next_citation_index(self) -> int:
1676 + idx = self._next_citation_index
1677 + self._next_citation_index += 1
1678 + return idx
vllm/entrypoints/generate/api_router.py
+62 -3
@@ -4,6 +4,8 @@ from typing import TYPE_CHECKING
4
5 from fastapi import FastAPI
6
7 +import vllm.envs as envs
8 +
9 if TYPE_CHECKING:
10 from argparse import Namespace
11
@@ -41,6 +43,12 @@ def register_generate_api_routers(app: FastAPI):
43
44 register_anthropic_api_router(app)
45
46 + from vllm.entrypoints.cohere.api_router import (
47 + attach_router as register_cohere_api_router,
48 + )
49 +
50 + register_cohere_api_router(app)
51 +
52 from .generative_scoring.api_router import register_generative_scoring_api_router
53
54 register_generative_scoring_api_router(app)
@@ -55,6 +63,23 @@ async def init_generate_state(
63 ):
64 from vllm.entrypoints.anthropic.serving import AnthropicServingMessages
65 from vllm.entrypoints.chat_utils import load_chat_template
66 +
67 + # The Cohere serving handler depends on the optional `cohere` SDK for
68 + # its wire-format protocol models, and is additionally gated on the
69 + # `VLLM_ENABLE_COHERE_API` env flag (see
70 + # `vllm.entrypoints.cohere.api_router.attach_router`). Skip the import
71 + # entirely when the endpoint isn't going to be exposed, both because
72 + # the SDK may not be installed and because the serving object holds
73 + # nontrivial state (chat handler, warmup) that would otherwise be
74 + # unused.
75 + if envs.VLLM_ENABLE_COHERE_API:
76 + try:
77 + from vllm.entrypoints.cohere.serving import CohereServingChatV2
78 + except ImportError:
79 + CohereServingChatV2 = None # type: ignore[assignment,misc]
80 + else:
81 + CohereServingChatV2 = None # type: ignore[assignment,misc]
82 +
83 from vllm.entrypoints.mcp.tool_server import (
84 DemoToolServer,
85 MCPToolServer,
@@ -86,6 +111,19 @@ async def init_generate_state(
111 tool_server = None
112 resolved_chat_template = load_chat_template(args.chat_template)
113
114 + # Fold the dedicated ``--cohere-format`` CLI flag into the renderer's
115 + # default chat-template kwargs. The cohere renderer reads
116 + # ``chat_template_kwargs["cohere_format"]`` to pick cmd3 vs cmd4
117 + # rendering; making this a first-class flag keeps the right format
118 + # discoverable for ``vllm serve --tokenizer-mode cohere`` users
119 + # without forcing them to hand-construct a JSON dict for
120 + # ``--default-chat-template-kwargs``. Per-request overrides still
121 + # take precedence (see ``merge_kwargs`` in
122 + # ``ChatCompletionRequest.build_chat_params``).
123 + default_chat_template_kwargs = dict(args.default_chat_template_kwargs or {})
124 + if getattr(args, "cohere_format", None):
125 + default_chat_template_kwargs.setdefault("cohere_format", args.cohere_format)
126 +
127 # Render endpoints are always backed by OnlineRenderer so that
128 # /v1/chat/completions/render and /v1/completions/render work on both
129 # generate-mode and render-only servers. Created in init_app_state.
@@ -106,7 +144,7 @@ async def init_generate_state(
144 enable_prompt_tokens_details=args.enable_prompt_tokens_details,
145 enable_force_include_usage=args.enable_force_include_usage,
146 enable_log_outputs=args.enable_log_outputs,
109 - default_chat_template_kwargs=args.default_chat_template_kwargs,
147 + default_chat_template_kwargs=default_chat_template_kwargs,
148 )
149 if "generate" in supported_tasks
150 else None
@@ -119,7 +157,7 @@ async def init_generate_state(
157 request_logger=request_logger,
158 chat_template=resolved_chat_template,
159 chat_template_content_format=args.chat_template_content_format,
122 - default_chat_template_kwargs=args.default_chat_template_kwargs,
160 + default_chat_template_kwargs=default_chat_template_kwargs,
161 trust_request_chat_template=args.trust_request_chat_template,
162 return_tokens_as_token_ids=args.return_tokens_as_token_ids,
163 enable_auto_tools=args.enable_auto_tool_choice,
@@ -169,11 +207,32 @@ async def init_generate_state(
207 reasoning_parser=args.structured_outputs_config.reasoning_parser,
208 enable_prompt_tokens_details=args.enable_prompt_tokens_details,
209 enable_force_include_usage=args.enable_force_include_usage,
172 - default_chat_template_kwargs=args.default_chat_template_kwargs,
210 + default_chat_template_kwargs=default_chat_template_kwargs,
211 )
212 if "generate" in supported_tasks
213 else None
214 )
215 + state.cohere_serving_chat_v2 = (
216 + CohereServingChatV2(
217 + engine_client,
218 + state.openai_serving_models,
219 + args.response_role,
220 + online_renderer=state.online_renderer,
221 + request_logger=request_logger,
222 + chat_template=resolved_chat_template,
223 + chat_template_content_format=args.chat_template_content_format,
224 + return_tokens_as_token_ids=args.return_tokens_as_token_ids,
225 + enable_auto_tools=args.enable_auto_tool_choice,
226 + tool_parser=args.tool_call_parser,
227 + reasoning_parser=args.structured_outputs_config.reasoning_parser,
228 + enable_prompt_tokens_details=args.enable_prompt_tokens_details,
229 + enable_force_include_usage=args.enable_force_include_usage,
230 + default_chat_template_kwargs=default_chat_template_kwargs,
231 + is_reasoning_model=args.cohere_is_reasoning_model,
232 + )
233 + if CohereServingChatV2 is not None and "generate" in supported_tasks
234 + else None
235 + )
236
237 from .generative_scoring.serving import ServingGenerativeScoring
238
vllm/entrypoints/openai/chat_completion/protocol.py
+18 -3
@@ -10,7 +10,13 @@ from openai.types.chat.chat_completion_audio import (
10 ChatCompletionAudio as OpenAIChatCompletionAudio,
11 )
12 from openai.types.chat.chat_completion_message import Annotation as OpenAIAnnotation
13 -from pydantic import Field, PrivateAttr, model_serializer, model_validator
13 +from pydantic import (
14 + Field,
15 + PrivateAttr,
16 + SerializeAsAny,
17 + model_serializer,
18 + model_validator,
19 +)
20
21 from vllm.config import ModelConfig
22 from vllm.entrypoints.chat_utils import (
@@ -91,7 +97,12 @@ class ChatCompletionLogProbs(OpenAIBaseModel):
97
98 class ChatCompletionResponseChoice(OpenAIBaseModel):
99 index: int
94 - message: ChatMessage
100 + # ``SerializeAsAny`` lets pydantic honor subclasses of ``ChatMessage``
101 + # (e.g. ``vllm.entrypoints.cohere.cohere_chat_message.CohereChatMessage``)
102 + # so that added fields like ``citations`` survive JSON serialization
103 + # instead of being stripped down to the base schema. Plain
104 + # ``ChatMessage`` instances serialize identically to before.
105 + message: SerializeAsAny[ChatMessage]
106 logprobs: ChatCompletionLogProbs | None = None
107 # per OpenAI spec this is the default
108 finish_reason: str | None = "stop"
@@ -139,7 +150,11 @@ class ChatCompletionResponse(OpenAIBaseModel):
150
151 class ChatCompletionResponseStreamChoice(OpenAIBaseModel):
152 index: int
142 - delta: DeltaMessage
153 + # ``SerializeAsAny`` lets pydantic honor subclasses of ``DeltaMessage``
154 + # (e.g. ``vllm.entrypoints.cohere.cohere_chat_message.CohereDeltaMessage``)
155 + # so streaming ``citations`` survive JSON serialization. Plain
156 + # ``DeltaMessage`` instances serialize identically to before.
157 + delta: SerializeAsAny[DeltaMessage]
158 logprobs: ChatCompletionLogProbs | None = None
159 finish_reason: str | None = None
160 stop_reason: int | str | None = None
vllm/entrypoints/openai/chat_completion/serving.py
+52 -6
@@ -391,6 +391,34 @@ class OpenAIServingChat(GenerateBaseServing):
391 return self.response_role
392 return request.messages[-1]["role"]
393
394 + def _create_chat_message(self, *args: Any, **kwargs: Any) -> ChatMessage:
395 + """Construct the response :class:`ChatMessage` for the non-streaming path.
396 +
397 + The full-generator calls this at every construction site so
398 + subclasses can swap in a specialized :class:`ChatMessage`
399 + subclass (e.g. :class:`CohereServingChatV2` returning
400 + :class:`CohereChatMessage`) without duplicating the branchy
401 + tool-choice / auto-tools logic that decides which fields are
402 + populated. The default returns a plain :class:`ChatMessage`.
403 + """
404 + return ChatMessage(*args, **kwargs)
405 +
406 + def _finalize_response_message(
407 + self,
408 + message: ChatMessage,
409 + *,
410 + parser: Parser | None,
411 + ) -> ChatMessage:
412 + """Subclass hook to enrich a fully-constructed :class:`ChatMessage`.
413 +
414 + Default is a no-op. Subclasses that need to surface parser-side
415 + extras (e.g. :class:`CohereServingChatV2` reading grounding
416 + citations off the reasoning parser and populating
417 + :class:`CohereChatMessage.citations`) override this to inspect
418 + ``parser`` and mutate/replace ``message``.
419 + """
420 + return message
421 +
422 async def chat_completion_stream_generator(
423 self,
424 request: ChatCompletionRequest,
@@ -894,13 +922,19 @@ class OpenAIServingChat(GenerateBaseServing):
922 )
923 is_required_tool_choice = request.tool_choice == "required"
924
925 + # All six construction sites route through ``self._create_chat_message``
926 + # so subclasses can swap in a specialized :class:`ChatMessage`
927 + # (e.g. the Cohere v2 handler's ``CohereChatMessage``) without
928 + # having to duplicate this branch logic.
929 if (not self.enable_auto_tools or not tool_parser_cls) and (
930 not is_named_tool_choice and not is_required_tool_choice
931 ):
900 - message = ChatMessage(role=role, reasoning=reasoning, content=content)
932 + message = self._create_chat_message(
933 + role=role, reasoning=reasoning, content=content
934 + )
935
936 elif is_named_tool_choice or is_required_tool_choice:
903 - message = ChatMessage(
937 + message = self._create_chat_message(
938 role=role,
939 reasoning=reasoning,
940 content=content or "",
@@ -913,7 +947,9 @@ class OpenAIServingChat(GenerateBaseServing):
947 # if the request doesn't use tool choice
948 # OR specifies to not use a tool
949 elif not request.tool_choice or request.tool_choice == "none":
916 - message = ChatMessage(role=role, reasoning=reasoning, content=content)
950 + message = self._create_chat_message(
951 + role=role, reasoning=reasoning, content=content
952 + )
953
954 # handle when there are tools and tool choice is auto
955 elif (
@@ -924,7 +960,7 @@ class OpenAIServingChat(GenerateBaseServing):
960 ):
961 auto_tools_called = tool_calls is not None and len(tool_calls) > 0
962 if tool_calls:
927 - message = ChatMessage(
963 + message = self._create_chat_message(
964 role=role,
965 reasoning=reasoning,
966 content=content,
@@ -935,7 +971,7 @@ class OpenAIServingChat(GenerateBaseServing):
971 )
972
973 else:
938 - message = ChatMessage(
974 + message = self._create_chat_message(
975 role=role,
976 reasoning=reasoning,
977 content=content,
@@ -948,7 +984,17 @@ class OpenAIServingChat(GenerateBaseServing):
984 " if tools should be extracted. Returning a standard chat "
985 "completion."
986 )
951 - message = ChatMessage(role=role, reasoning=reasoning, content=content)
987 + message = self._create_chat_message(
988 + role=role, reasoning=reasoning, content=content
989 + )
990 +
991 + # Subclass hook: enrich the constructed message with any
992 + # parser-side extras that don't fit through the plain
993 + # ``(reasoning, content, tool_calls)`` tuple. Base is a no-op;
994 + # citation-aware handlers use this to surface grounding
995 + # metadata cached on the reasoning parser.
996 + message = self._finalize_response_message(message, parser=parser)
997 +
998 # In OpenAI's API, when a tool is called, the finish_reason is:
999 # "tool_calls" for "auto" or "required" tool calls,
1000 # and "stop" for named tool calls.
vllm/entrypoints/openai/cli_args.py
+19
@@ -148,6 +148,25 @@ class BaseFrontendArgs:
148 """If set to False, output deltas will not be logged. Relevant only if
149 --enable-log-outputs is set.
150 """
151 + cohere_is_reasoning_model: bool = True
152 + """Cohere ``/cohere/v2/chat`` only. Whether the served model is a
153 + reasoning Command-family model. When True (default), the assistant's
154 + chain-of-thought is surfaced as a ``thinking`` content block (or
155 + ``content-*`` events on the stream). When False, reasoning is
156 + surfaced as Cohere's ``tool_plan`` field (or ``tool-plan-delta``
157 + events) whenever the model emits tool calls, matching older non-
158 + reasoning Command models. Has no effect on the non-Cohere
159 + endpoints."""
160 + cohere_format: str = "cmd4"
161 + """Cohere ``--tokenizer-mode cohere`` only. Which Cohere prompt
162 + format to render: ``cmd4`` (current Command A+ models; default) or
163 + ``cmd3`` (earlier Cmd-A and Cmd-A reasoning models). Selecting the
164 + wrong format silently produces a prompt the model wasn't trained
165 + on, which most commonly manifests as the model emitting text but no
166 + citations / tool calls / thinking blocks. Equivalent to passing
167 + ``--default-chat-template-kwargs '{"cohere_format": "..."}'`` -- any
168 + explicit request-level ``chat_template_kwargs.cohere_format`` takes
169 + priority."""
170 log_error_stack: bool = envs.VLLM_SERVER_DEV_MODE
171 """If set to True, log the stack trace of error responses"""
172 tokens_only: bool = False
vllm/entrypoints/serve/utils/server_utils.py
+1 -1
@@ -39,7 +39,7 @@ from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
39 logger = init_logger("vllm.entrypoints.openai.server_utils")
40
41
42 -GUARDED_PREFIX = ("/v1", "/v2", "/inference")
42 +GUARDED_PREFIX = ("/v1", "/v2", "/inference", "/cohere")
43
44
45 class AuthenticationMiddleware:
vllm/envs.py
+6
@@ -240,6 +240,7 @@ if TYPE_CHECKING:
240 VLLM_LOOPBACK_IP: str = ""
241 VLLM_ALLOW_CHUNKED_LOCAL_ATTN_WITH_HYBRID_KV_CACHE: bool = True
242 VLLM_ENABLE_RESPONSES_API_STORE: bool = False
243 + VLLM_ENABLE_COHERE_API: bool = False
244 VLLM_HAS_FLASHINFER_CUBIN: bool = False
245 VLLM_ROCM_FP8_MFMA_PAGE_ATTN: bool = False
246 VLLM_ALLREDUCE_USE_SYMM_MEM: bool = True
@@ -1758,6 +1759,11 @@ environment_variables: dict[str, Callable[[], Any]] = {
1759 "VLLM_ENABLE_RESPONSES_API_STORE": lambda: bool(
1760 int(os.getenv("VLLM_ENABLE_RESPONSES_API_STORE", "0"))
1761 ),
1762 + # If set to 1, expose the Cohere Chat v2 API at ``POST /cohere/v2/chat``.
1763 + # Default off
1764 + "VLLM_ENABLE_COHERE_API": lambda: bool(
1765 + int(os.getenv("VLLM_ENABLE_COHERE_API", "0"))
1766 + ),
1767 # If set, use the fp8 mfma in rocm paged attention.
1768 "VLLM_ROCM_FP8_MFMA_PAGE_ATTN": lambda: bool(
1769 int(os.getenv("VLLM_ROCM_FP8_MFMA_PAGE_ATTN", "0"))
vllm/reasoning/cohere_command_reasoning_parser.py
+149 -4
@@ -20,6 +20,11 @@ except ImportError as e:
20 ) from e
21
22
23 +from vllm.entrypoints.cohere.cohere_chat_message import (
24 + Citation,
25 + CitationSource,
26 + CohereDeltaMessage,
27 +)
28 from vllm.entrypoints.openai.chat_completion.protocol import (
29 ChatCompletionRequest,
30 )
@@ -404,6 +409,83 @@ def _schema_dict_from_structured_outputs(
409 )
410
411
412 +def _melody_sources_to_vllm(
413 + raw_sources: Any,
414 + position_to_source: Mapping[tuple[int, int], CitationSource] | None,
415 +) -> list[CitationSource]:
416 + """Convert melody's ``Source`` objects into resolved :class:`CitationSource`.
417 +
418 + Melody's ``Source`` shape is
419 + ``{tool_call_index, tool_result_indices, document_ids}``. Without
420 + :meth:`PyFilterOptions.with_message_history` the ``document_ids``
421 + list is empty, so we resolve the numeric address ourselves using
422 + ``position_to_source``, which is
423 + :meth:`vllm.entrypoints.cohere.serving.CohereServingChatV2._build_position_to_source`
424 + applied to the inbound request and forwarded through
425 + ``chat_template_kwargs`` at parser-construction time. Each melody
426 + source with ``tool_result_indices=[i, j, ...]`` fans out to one
427 + :class:`CitationSource` per resolved position, populating the
428 + wire-shape ``type`` / ``id`` / ``document`` / ``tool_output``
429 + fields; unresolved positions are silently skipped (the calling
430 + citation is dropped at wire-coercion time if every source in it
431 + ends up unresolved -- see ``_to_wire_citation`` in
432 + :mod:`vllm.entrypoints.cohere.serving`).
433 +
434 + When ``position_to_source`` is ``None`` (parser wired outside of
435 + :class:`CohereServingChatV2`, e.g. a plain OpenAI chat completion
436 + request that happens to hit a cohere reasoning parser) sources
437 + can't be resolved and get dropped; citations then never reach any
438 + citation-aware handler downstream.
439 + """
440 + if not position_to_source:
441 + return []
442 + out: list[CitationSource] = []
443 + for s in raw_sources or []:
444 + bucket = getattr(s, "tool_call_index", None)
445 + indices = list(getattr(s, "tool_result_indices", None) or [])
446 + if bucket is None or not indices:
447 + continue
448 + for idx in indices:
449 + info = position_to_source.get((bucket, idx))
450 + if info is not None:
451 + out.append(info)
452 + return out
453 +
454 +
455 +def _melody_citations_to_vllm(
456 + raw_citations: Any,
457 + position_to_source: Mapping[tuple[int, int], CitationSource] | None,
458 +) -> list[Citation] | None:
459 + """Convert melody's ``FilterCitation`` objects into :class:`Citation`.
460 +
461 + Resolves each source in-place via ``position_to_source`` (see
462 + :func:`_melody_sources_to_vllm`). Citations whose sources all fail
463 + to resolve are still emitted here with ``sources=[]``; the serving
464 + layer drops them at wire-coercion time so the fail-closed policy
465 + lives in exactly one place.
466 + """
467 + if not raw_citations:
468 + return None
469 + out: list[Citation] = []
470 + for c in raw_citations:
471 + out.append(
472 + Citation(
473 + start=getattr(c, "start_index", None),
474 + end=getattr(c, "end_index", None),
475 + text=getattr(c, "text", None),
476 + sources=_melody_sources_to_vllm(
477 + getattr(c, "sources", None), position_to_source
478 + ),
479 + type=(
480 + "THINKING_CONTENT"
481 + if getattr(c, "is_thinking", False)
482 + else "TEXT_CONTENT"
483 + ),
484 + )
485 + )
486 + return out
487 +
488 +
489 class BaseCohereCommandReasoningParser(ReasoningParser):
490 def __init__(
491 self,
@@ -420,6 +502,32 @@ class BaseCohereCommandReasoningParser(ReasoningParser):
502 self.unary_opts = unary_opts
503 self.melody_unary = PyFilter(unary_opts)
504 self.melody_streaming = PyFilter(streaming_opts)
505 + # Citations extracted by the most recent ``extract_reasoning`` call.
506 + # Citation-aware handlers (e.g. :class:`CohereServingChatV2` via its
507 + # ``_finalize_response_message`` hook) read this back from the
508 + # parser instance (which is constructed per-request) and attach
509 + # the result to :class:`CohereChatMessage.citations`. ``None`` when
510 + # the last parse produced no citations.
511 + self.last_unary_citations: list[Citation] | None = None
512 + # Request-scoped ``(tool_call_index, tool_result_idx) -> resolved
513 + # CitationSource`` map, forwarded from
514 + # :meth:`CohereServingChatV2._apply_cohere_template_kwargs` via
515 + # ``chat_template_kwargs``. See :func:`_melody_sources_to_vllm`
516 + # for how it's consumed. Absent for parser instances that
517 + # weren't created by ``CohereServingChatV2`` (e.g. an OpenAI
518 + # chat completion request routed through a cohere reasoning
519 + # parser); citations from those requests are dropped since we
520 + # have no way to attribute their sources.
521 + ctk = kwargs.get("chat_template_kwargs") or {}
522 + # Lazy import: ``renderers.cohere`` transitively imports
523 + # ``cohere_melody`` and other heavy deps; hoisting the constant
524 + # to module scope would tie parser importability to that graph.
525 + from vllm.renderers.cohere import POSITION_TO_SOURCE_KEY
526 +
527 + raw_map = ctk.get(POSITION_TO_SOURCE_KEY)
528 + self._position_to_source: Mapping[tuple[int, int], CitationSource] | None = (
529 + raw_map if isinstance(raw_map, Mapping) else None
530 + )
531
532 @property
533 def reasoning_start_str(self) -> str | None:
@@ -439,9 +547,22 @@ class BaseCohereCommandReasoningParser(ReasoningParser):
547 delta_token_ids: Sequence[int],
548 ) -> DeltaMessage | None:
549 r = self.melody_streaming.write_decoded(delta_text)
442 - if r.content is None and r.reasoning is None and not r.tool_calls:
550 + citations = _melody_citations_to_vllm(
551 + getattr(r, "citations", None), self._position_to_source
552 + )
553 + if (
554 + r.content is None
555 + and r.reasoning is None
556 + and not r.tool_calls
557 + and not citations
558 + ):
559 return None
444 - msg = DeltaMessage()
560 + # Always emit CohereDeltaMessage so citations reach the wire via
561 + # ``SerializeAsAny[DeltaMessage]`` on the streaming choice envelope.
562 + # When ``citations`` is unset, ``CohereDeltaMessage._serialize``
563 + # drops the field so the emitted shape matches plain
564 + # :class:`DeltaMessage`
565 + msg = CohereDeltaMessage()
566 if r.content is not None:
567 msg.content = r.content
568 if r.reasoning is not None:
@@ -456,12 +577,23 @@ class BaseCohereCommandReasoningParser(ReasoningParser):
577 )
578 for tc in r.tool_calls
579 ]
580 + if citations:
581 + msg.citations = citations
582 return msg
583
584 def extract_reasoning(
585 self, model_output: str, request: ChatCompletionRequest | ResponsesRequest
586 ) -> tuple[str | None, str | None]:
587 result = self.melody_unary.process_full_text(model_output)
588 + # Cache citations so citation-aware handlers can surface them
589 + # on :class:`CohereChatMessage.citations`. The base
590 + # :meth:`ReasoningParser.extract_reasoning` contract only
591 + # returns ``(reasoning, content)``, so citations are passed
592 + # back via parser-instance state -- safe because the parser is
593 + # constructed per-request.
594 + self.last_unary_citations = _melody_citations_to_vllm(
595 + getattr(result, "citations", None), self._position_to_source
596 + )
597 return result.reasoning, result.content
598
599 def extract_content_ids(self, input_ids: list[int]) -> list[int]:
@@ -549,12 +681,25 @@ class BaseCohereCommandReasoningParser(ReasoningParser):
681 return request
682
683
684 +# melody's streaming filter only buffers a partial ``<co: ...>`` citation
685 +# across ``write_decoded`` calls when ``stream_non_grounded_answer`` is
686 +# set: otherwise, the moment an opening ``<co`` is seen without a
687 +# closing ``</co: ...>`` in the same delta, the filter emits the partial
688 +# marker bytes verbatim as plain content. In vLLM's streaming path the
689 +# parser is fed one token (1-4 chars) per call, so an unbuffered filter
690 +# will leak ``<co: 0>``-style markers into ``delta.content`` and never
691 +# emit a ``FilterCitation`` for them. Enabling the flag flips the
692 +# partial-match branch in melody's ``parse_citations`` (see
693 +# ``src/parsing/citations_filter.rs``) to ``return (None, 0)`` -- i.e.
694 +# keep buffering -- which lets a full citation eventually resolve.
695 +# Non-streaming (unary) parsing receives the whole output in one call so
696 +# the flag is a no-op there and we leave ``unary_opts`` alone.
697 class CohereCommand3ReasoningParser(BaseCohereCommandReasoningParser):
698 def __init__(self, tokenizer: TokenizerLike, *args, **kwargs):
699 super().__init__(
700 tokenizer,
701 *args,
557 - streaming_opts=PyFilterOptions().cmd3(),
702 + streaming_opts=PyFilterOptions().cmd3().stream_non_grounded_answer(),
703 unary_opts=PyFilterOptions().cmd3().no_tools(),
704 **kwargs,
705 )
@@ -565,7 +710,7 @@ class CohereCommand4ReasoningParser(BaseCohereCommandReasoningParser):
710 super().__init__(
711 tokenizer,
712 *args,
568 - streaming_opts=PyFilterOptions().cmd4(),
713 + streaming_opts=PyFilterOptions().cmd4().stream_non_grounded_answer(),
714 unary_opts=PyFilterOptions().cmd4().no_tools(),
715 **kwargs,
716 )
vllm/renderers/cohere.py new
+704
@@ -0,0 +1,704 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Cohere prompt renderer.
4 +
5 +Templates the Cohere Command-family prompt formats (cmd3 / cmd4) using
6 +the ``cohere_melody`` Rust bindings instead of Jinja. Enabled by
7 +passing ``--tokenizer-mode cohere`` on the engine; tokenization itself
8 +still flows through the cached HuggingFace tokenizer.
9 +
10 +The renderer's client surface is ``chat_template_kwargs`` (vLLM's
11 +generic passthrough for template-time inputs). It is fed by two
12 +different code paths, which speak different vocabularies:
13 +
14 +1. :class:`vllm.entrypoints.cohere.serving.CohereServingChatV2` --
15 + translates a native `Cohere Chat v2 request
16 + <https://docs.cohere.com/reference/chat>`__ into
17 + ``chat_template_kwargs``.
18 +2. :class:`OpenAIServingChat` -- forwards whatever the client puts in
19 + ``request.chat_template_kwargs`` directly, without translation.
20 +
21 +Accordingly the renderer accepts keys from both vocabularies.
22 +
23 +**Cohere Chat v2 request fields** (populated automatically by
24 +``CohereServingChatV2``; a direct ``chat_template_kwargs`` caller may
25 +also pass them):
26 +
27 +* ``documents``: list of documents (v2 ``documents``)
28 +* ``tools``: list of tools (v2 ``tools``)
29 +* ``safety_mode``: v2 ``safety_mode`` (``CONTEXTUAL`` / ``STRICT`` /
30 + ``NONE``) -- forwarded to melody's cmd3 ``safety_mode`` slot;
31 + lowercased.
32 +* ``citation_options``: v2 ``citation_options``. Its ``.mode``
33 + (``ENABLED`` / ``DISABLED`` / ``FAST`` / ``ACCURATE`` / ``OFF``) is
34 + normalized into melody's cmd3 ``citation_quality`` (``on`` / ``off``)
35 + and cmd4 ``grounding`` (``enabled`` / ``disabled`` / ``unknown``)
36 + slots. cmd4 has no fast/accurate distinction at the prompt-template
37 + layer -- both mean grounding-on.
38 +* ``response_format``: v2 ``response_format`` (``json_object`` /
39 + ``json_schema``) -- mapped onto melody's ``json_mode`` /
40 + ``json_schema`` config slots.
41 +* ``thinking``: v2 ``thinking``; its ``.type`` (``enabled`` /
42 + ``disabled``) becomes melody's ``reasoning_type``.
43 +
44 +**Melody template-config knobs** (accepted for direct-passthrough
45 +callers; not part of the Cohere Chat v2 surface):
46 +
47 +* ``cohere_format``: renderer selector, ``cmd3`` or ``cmd4`` (default
48 + ``cmd4``).
49 +* ``template_id``: pick one of melody's built-in template variants.
50 +* ``available_tools``: melody's own name for tools; takes precedence
51 + over ``tools`` when both are set.
52 +* ``reasoning_type``: direct ``enabled`` / ``disabled`` toggle;
53 + overrides derivation from ``thinking.type``.
54 +* ``dev_instruction``: developer-instruction override.
55 +* ``json_mode`` / ``json_schema``: direct structured-output toggles;
56 + override derivation from ``response_format``.
57 +* cmd3-only: ``citation_quality`` (``on`` / ``off``), ``skip_preamble``.
58 +* cmd4-only: ``grounding`` (``enabled`` / ``disabled`` / ``unknown``;
59 + overrides derivation from ``citation_options.mode``),
60 + ``platform_instruction``.
61 +
62 +**Rejected inputs**: ``template_jinja`` and ``template`` are Cohere's
63 +own inlets for raw Jinja template source. They are valid at Cohere's
64 +API surface but not at vLLM's -- raw template source must flow through
65 +the standard ``chat_template`` request field (guarded by
66 +``--trust-request-chat-template``), which this renderer forwards to
67 +melody as ``template_jinja``. Setting these keys explicitly in
68 +``chat_template_kwargs`` raises ``ValueError`` so client
69 +misconfiguration surfaces loudly.
70 +
71 +**Everything else** in ``chat_template_kwargs`` is forwarded verbatim
72 +to melody as ``additional_template_fields`` and becomes accessible as
73 +Jinja variables inside the template. This matches vLLM's documented
74 +contract for ``chat_template_kwargs`` ("kwargs accessible by the
75 +template"), so e.g. ``chat_template_kwargs={"reasoning_effort":
76 +"low"}`` resolves ``{{ reasoning_effort }}`` inside cmd3 / cmd4
77 +templates.
78 +
79 +Citations produced by Cohere models are surfaced through the
80 +Cohere-scoped ``CohereChatMessage.citations`` /
81 +``CohereDeltaMessage.citations`` fields (see
82 +:mod:`vllm.entrypoints.cohere.cohere_chat_message`), populated by the
83 +``cohere2`` reasoning parser. The base OpenAI ``ChatMessage`` /
84 +``DeltaMessage`` keep their declared schemas unchanged; the response
85 +envelope declares them as ``SerializeAsAny[...]`` so the subclass
86 +fields survive JSON serialization.
87 +"""
88 +
89 +from __future__ import annotations
90 +
91 +import copy
92 +import json
93 +from enum import Enum
94 +from typing import Any
95 +
96 +from vllm.config import VllmConfig
97 +from vllm.entrypoints.chat_utils import (
98 + ChatCompletionMessageParam,
99 + ConversationMessage,
100 + parse_chat_messages,
101 + parse_chat_messages_async,
102 +)
103 +from vllm.logger import init_logger
104 +from vllm.tokenizers.hf import HfTokenizer
105 +from vllm.utils.async_utils import make_async
106 +
107 +from .base import BaseRenderer
108 +from .inputs import DictPrompt
109 +from .inputs.preprocess import parse_dec_only_prompt
110 +from .params import ChatParams
111 +
112 +logger = init_logger(__name__)
113 +
114 +
115 +_DEFAULT_FORMAT = "cmd4"
116 +_VALID_FORMATS = ("cmd3", "cmd4")
117 +
118 +
119 +class MelodyContentType(str, Enum):
120 + """Wire-format discriminator for melody content blocks.
121 +
122 + These strings are what the cmd3 / cmd4 Jinja templates check against
123 + in ``message.content[0].type`` (e.g. the ``thinking`` branch in
124 + ``cmd4-v1.jinja``). Keep new values in sync with melody's template
125 + schema.
126 + """
127 +
128 + TEXT = "text"
129 + THINKING = "thinking"
130 + IMAGE = "image"
131 + DOCUMENT = "document"
132 +
133 +
134 +# Reserved ``chat_template_kwargs`` key that carries
135 +# ``AssistantChatMessageV2.citations`` from a Cohere v2 request through
136 +# to the renderer, keyed by request-message index. Populated by
137 +# ``CohereServingChatV2._apply_cohere_template_kwargs`` and consumed by
138 +# ``_conversation_to_melody_messages``. The value is a
139 +# ``dict[int, list[dict]]`` of melody-shape ``FilterCitation`` dicts.
140 +# The leading underscore signals "internal, not part of the public
141 +# ``chat_template_kwargs`` surface" -- clients should not set this key
142 +# directly.
143 +MESSAGES_CITATIONS_KEY = "_messages_citations"
144 +
145 +
146 +# Reserved ``chat_template_kwargs`` key carrying a request-side
147 +# ``dict[tuple[int, int], CitationSource]`` map from melody's numeric
148 +# ``(tool_call_index, tool_result_index)`` addressing to fully-resolved
149 +# wire-shape sources (type, id, document/tool_output payload).
150 +# Populated by :meth:`CohereServingChatV2._apply_cohere_template_kwargs`
151 +# and consumed by :class:`BaseCohereCommandReasoningParser` at request
152 +# start so the parser can attach real document ids on each melody source
153 +# it emits -- eliminating the numeric-coord round-trip through the
154 +# internal streaming pipeline. Leading underscore signals "internal,
155 +# not part of the public ``chat_template_kwargs`` surface" -- clients
156 +# should not set this key directly.
157 +POSITION_TO_SOURCE_KEY = "_position_to_source"
158 +
159 +
160 +# Keys this renderer interprets directly from ``chat_template_kwargs`` and
161 +# maps onto typed melody render-config fields. Everything *not* in this
162 +# set is forwarded verbatim to melody as ``additional_template_fields``
163 +# (i.e. as Jinja template variables), so callers can write
164 +# ``chat_template_kwargs = {"my_var": "..."}`` and have ``{{ my_var }}``
165 +# resolve inside the template -- matching vLLM's documented contract
166 +# for ``chat_template_kwargs``.
167 +_RENDERER_CONSUMED_KEYS = frozenset(
168 + {
169 + "cohere_format",
170 + "template_id",
171 + # ``template_jinja`` / ``template`` are Cohere-only inputs for
172 + # raw Jinja source -- valid at Cohere's own API surface but not
173 + # at vLLM's. In vLLM, raw template source must flow through the
174 + # standard ``chat_template`` request field, so a non-``None``
175 + # value here is rejected in ``_build_render_config``. The keys
176 + # stay in the consumed set so an explicit ``None`` (e.g. from
177 + # ``.model_dump(exclude_none=False)``) is dropped rather than
178 + # surfacing as a stray ``{{ template_jinja }}`` Jinja variable.
179 + "template_jinja",
180 + "template",
181 + "use_jinja",
182 + "documents",
183 + "available_tools",
184 + "tools",
185 + "reasoning_type",
186 + "thinking",
187 + "dev_instruction",
188 + "response_format",
189 + "json_schema",
190 + "json_mode",
191 + "safety_mode",
192 + "citation_quality",
193 + "citation_options",
194 + "skip_preamble",
195 + "grounding",
196 + "platform_instruction",
197 + MESSAGES_CITATIONS_KEY,
198 + POSITION_TO_SOURCE_KEY,
199 + }
200 +)
201 +
202 +
203 +def _try_import_melody():
204 + try:
205 + import cohere_melody # type: ignore
206 +
207 + return cohere_melody
208 + except ImportError as e: # pragma: no cover - exercised at runtime
209 + raise ImportError(
210 + "The `cohere` tokenizer/renderer mode requires the "
211 + "`cohere_melody` package. Install it via "
212 + "`pip install cohere-melody` or build from "
213 + "https://github.com/cohere-ai/melody."
214 + ) from e
215 +
216 +
217 +_MELODY_ROLES = frozenset({"system", "user", "chatbot", "tool"})
218 +
219 +
220 +# Cohere v2 ``citation_options.mode`` -> melody cmd4 ``grounding``.
221 +# v2 surfaces five modes (``ENABLED`` / ``DISABLED`` / ``FAST`` /
222 +# ``ACCURATE`` / ``OFF``) but cmd4's template doesn't differentiate
223 +# fast vs accurate at the prompt layer -- both just turn grounding on,
224 +# and ``ENABLED`` / ``DISABLED`` are direct aliases. Melody itself
225 +# accepts ``unknown`` / ``enabled`` / ``disabled``.
226 +_CMD4_GROUNDING_FROM_MODE = {
227 + "fast": "enabled",
228 + "accurate": "enabled",
229 + "on": "enabled",
230 + "enabled": "enabled",
231 + "off": "disabled",
232 + "disabled": "disabled",
233 + "unknown": "unknown",
234 +}
235 +
236 +
237 +def _normalize_cmd4_grounding(value: Any) -> str:
238 + """Coerce a user-facing grounding/citation mode to melody's vocab.
239 +
240 + Raises ``ValueError`` rather than letting an unrecognized value
241 + slip through to ``render_cmd4`` and surface as a generic
242 + ``Invalid config: grounding`` from melody.
243 + """
244 + out = _CMD4_GROUNDING_FROM_MODE.get(value.lower())
245 + if out is None:
246 + raise ValueError(
247 + f"Unrecognized cmd4 grounding value: {value!r}. Expected one "
248 + f"of ENABLED / DISABLED / FAST / ACCURATE / OFF (from "
249 + f"citation_options.mode), or the direct melody values "
250 + f"enabled / disabled / unknown."
251 + )
252 + return out
253 +
254 +
255 +def _role_to_melody(role: str) -> str:
256 + """Map an OpenAI role to the role string melody expects.
257 +
258 + The cmd3 / cmd4 jinja templates only recognize ``system``, ``user``,
259 + ``assistant``/``chatbot``, and ``tool``: any other role is silently
260 + dropped by the template's role-dispatch chain (no fallback branch),
261 + which produces a malformed prompt without any error. We therefore
262 + refuse unknown roles up front rather than letting them disappear.
263 +
264 + Aliases:
265 +
266 + * ``assistant`` -> ``chatbot`` (Cohere's historical assistant role
267 + name used by the templates).
268 + * ``developer`` -> ``system`` (OpenAI's ``developer`` role is
269 + documented as high-priority instructions, which maps onto the
270 + ``system`` slot in Cohere's prompt format).
271 + """
272 + role = role.lower()
273 + if role == "assistant":
274 + return "chatbot"
275 + if role == "developer":
276 + return "system"
277 + if role in _MELODY_ROLES:
278 + return role
279 + raise ValueError(
280 + f"Unsupported message role for the cohere renderer: {role!r}. "
281 + "Expected one of: system, developer, user, assistant, chatbot, tool."
282 + )
283 +
284 +
285 +def _normalize_tool_call(tc: dict[str, Any] | Any) -> dict[str, Any]:
286 + """Normalize an OpenAI tool call dict into melody's tool call shape.
287 +
288 + melody expects ``{id, name, parameters: <json str>}`` whereas OpenAI
289 + delivers ``{id, type, function: {name, arguments: <json|str>}}``.
290 + """
291 + # Pydantic objects
292 + if hasattr(tc, "model_dump"):
293 + tc = tc.model_dump()
294 + if not isinstance(tc, dict):
295 + raise TypeError(f"Unexpected tool_call value: {tc!r}")
296 +
297 + fn = tc.get("function") or {}
298 + name = fn.get("name") or tc.get("name") or ""
299 + args = fn.get("arguments")
300 + if args is None:
301 + args = tc.get("arguments", {})
302 + # melody expects a JSON-encoded string
303 + if not isinstance(args, str):
304 + args = json.dumps(args, ensure_ascii=False)
305 + return {
306 + "id": tc.get("id") or "",
307 + "name": name,
308 + "parameters": args,
309 + }
310 +
311 +
312 +def _content_blocks(content: Any) -> list[dict[str, Any]]:
313 + """Convert OpenAI ``ConversationMessage.content`` into melody content blocks.
314 +
315 + The chat_utils ``content_format="openai"`` produces a list of
316 + ``{type: text, text: ...}`` dicts already; we pass those through with
317 + minimal coercion. Plain string content is wrapped in a single text block.
318 + Image / multimodal placeholder dicts (``{"type": "image"}``) are
319 + forwarded as-is so that templates can reference image placeholders that
320 + the upstream tokenizer expands separately.
321 + """
322 + if content is None:
323 + return []
324 + if isinstance(content, str):
325 + return [{"type": MelodyContentType.TEXT, "text": content}]
326 + blocks: list[dict[str, Any]] = []
327 + for part in content:
328 + if isinstance(part, str):
329 + blocks.append({"type": MelodyContentType.TEXT, "text": part})
330 + continue
331 + if not isinstance(part, dict):
332 + raise TypeError(f"Unexpected content part: {part!r}")
333 + part_type = part.get("type", MelodyContentType.TEXT)
334 + if part_type in ("text", "input_text", "output_text", "refusal"):
335 + blocks.append(
336 + {"type": MelodyContentType.TEXT, "text": part.get("text", "")}
337 + )
338 + elif part_type == MelodyContentType.THINKING:
339 + blocks.append(
340 + {
341 + "type": MelodyContentType.THINKING,
342 + "thinking": part.get("thinking", ""),
343 + }
344 + )
345 + elif part_type == MelodyContentType.IMAGE:
346 + blocks.append(
347 + {
348 + "type": MelodyContentType.IMAGE,
349 + "image": {
350 + "template_placeholder": part.get(
351 + "template_placeholder", "<image>"
352 + ),
353 + },
354 + }
355 + )
356 + elif part_type == MelodyContentType.DOCUMENT:
357 + doc = part.get("document")
358 + if isinstance(doc, dict):
359 + blocks.append({"type": MelodyContentType.DOCUMENT, "document": doc})
360 + else:
361 + # Fall back to wrapping arbitrary string as text.
362 + blocks.append({"type": MelodyContentType.TEXT, "text": json.dumps(doc)})
363 + elif part_type == "tool_reference":
364 + # Tool references are rendered by name; emit as text so the
365 + # renderer downstream is content-format agnostic.
366 + blocks.append(
367 + {
368 + "type": MelodyContentType.TEXT,
369 + "text": part.get("name") or part.get("text", ""),
370 + }
371 + )
372 + else:
373 + # Unknown block type: render as text fallback.
374 + text = part.get("text") or part.get(part_type) or ""
375 + text_str = text if isinstance(text, str) else json.dumps(text)
376 + blocks.append({"type": MelodyContentType.TEXT, "text": text_str})
377 + return blocks
378 +
379 +
380 +def _document_to_melody(doc: Any) -> dict[str, Any]:
381 + """Coerce a Cohere v2 document into melody's ``Document`` (dict) shape."""
382 + match doc:
383 + case str():
384 + return {"text": doc}
385 + # Cohere v2 wraps documents in {id, data: {...}}; the ``**rest``
386 + # capture keeps the outer ``id`` (and any other outer keys) so we
387 + # can lift them onto the payload without losing them.
388 + case {"data": dict() as data, **rest}:
389 + payload = dict(data)
390 + if (outer_id := rest.get("id")) is not None and "id" not in payload:
391 + payload["id"] = outer_id
392 + return payload
393 + case dict():
394 + return dict(doc)
395 + case _:
396 + raise TypeError(f"Unsupported document type: {type(doc).__name__}")
397 +
398 +
399 +def _tool_to_melody(tool: Any) -> dict[str, Any]:
400 + """Coerce a chat completions tool definition into melody's ``Tool`` shape.
401 +
402 + Accepts either the raw OpenAI tool wrapper ``{type:"function", function:
403 + {name, description, parameters}}`` or a flat ``{name, description,
404 + parameters}`` dict (which is what melody itself expects).
405 + """
406 + if hasattr(tool, "model_dump"):
407 + tool = tool.model_dump()
408 + if not isinstance(tool, dict):
409 + raise TypeError(f"Unsupported tool type: {type(tool).__name__}")
410 + if "function" in tool and isinstance(tool["function"], dict):
411 + fn = tool["function"]
412 + else:
413 + fn = tool
414 + return {
415 + "name": fn.get("name", ""),
416 + "description": fn.get("description", "") or "",
417 + "parameters": fn.get("parameters") or {},
418 + }
419 +
420 +
421 +def _conversation_to_melody_messages(
422 + conversation: list[ConversationMessage],
423 + messages_citations: dict[int, list[dict[str, Any]]] | None = None,
424 +) -> list[dict[str, Any]]:
425 + """Convert vLLM ``ConversationMessage``s into melody's message dict shape.
426 +
427 + ``messages_citations`` is an optional index-keyed lookup of
428 + per-message citations from the request, populated by
429 + ``CohereServingChatV2`` under the
430 + :data:`MESSAGES_CITATIONS_KEY` chat_template_kwargs entry. Values
431 + must already be in melody's ``FilterCitation`` dict shape (see
432 + ``CohereServingChatV2._sdk_citation_to_melody``). Index is the
433 + position in the *input* ``conversation`` list, which currently maps
434 + 1-to-1 with ``CohereChatV2Request.messages`` for text-only Cohere
435 + requests. This mapping breaks silently if ``parse_chat_messages``
436 + ever expands a single input message into multiple entries (e.g.
437 + multi-modal splitting), which is not the case today.
438 + """
439 + out: list[dict[str, Any]] = []
440 + for i, msg in enumerate(conversation):
441 + role = _role_to_melody(msg.get("role", "user"))
442 + content_blocks = _content_blocks(msg.get("content"))
443 +
444 + # Treat reasoning content as a thinking block on assistant turns
445 + # so multi-turn reasoning is preserved in the rendered prompt.
446 + #
447 + # Cohere models output thought as either a ``thinking``
448 + # block (reasoning models) or a ``tool_plan`` field (older non-
449 + # reasoning Command models), but vLLM's ConversationMessage has a
450 + # single unified ``reasoning`` field that drops that difference.
451 + # The cmd3 / cmd4 jinja templates render ``thinking`` blocks in both
452 + # the case of tool calls and regular thinking blocks, so this is fine
453 + # from the renderer's perspective.
454 + reasoning = msg.get("reasoning") or msg.get("reasoning_content")
455 + if role == "chatbot" and reasoning:
456 + content_blocks.insert(
457 + 0,
458 + {"type": MelodyContentType.THINKING, "thinking": reasoning},
459 + )
460 +
461 + tool_calls = [_normalize_tool_call(tc) for tc in (msg.get("tool_calls") or [])]
462 +
463 + out_msg: dict[str, Any] = {
464 + "role": role,
465 + "content": content_blocks,
466 + "tool_calls": tool_calls,
467 + }
468 + tool_call_id = msg.get("tool_call_id")
469 + if tool_call_id:
470 + out_msg["tool_call_id"] = tool_call_id
471 + if messages_citations and (cites := messages_citations.get(i)):
472 + out_msg["citations"] = cites
473 + out.append(out_msg)
474 + return out
475 +
476 +
477 +def _build_render_config(
478 + conversation: list[ConversationMessage],
479 + chat_template_kwargs: dict[str, Any],
480 + chat_template: str | None = None,
481 +) -> tuple[str, dict[str, Any]]:
482 + """Build the ``render_cmd3`` / ``render_cmd4`` config dict.
483 +
484 + ``chat_template`` is the standard vLLM per-request template override
485 + (from ``ChatCompletionRequest.chat_template`` / the server-side
486 + ``--chat-template`` flag). When non-``None`` it is forwarded to melody
487 + as ``template_jinja`` so a single vLLM-native path drives template
488 + source through the trust guard.
489 +
490 + Returns ``(format, config_dict)`` where ``format`` is either ``"cmd3"``
491 + or ``"cmd4"``.
492 + """
493 + fmt = chat_template_kwargs.get("cohere_format", _DEFAULT_FORMAT)
494 + if fmt not in _VALID_FORMATS:
495 + raise ValueError(
496 + f"Invalid cohere_format={fmt!r}; expected one of {_VALID_FORMATS}"
497 + )
498 +
499 + # Reject Cohere-only inlets for raw Jinja source. These keys are
500 + # valid at Cohere's own API surface but not at vLLM's: callers must
501 + # use the standard vLLM ``chat_template`` request field so the
502 + # ``--trust-request-chat-template`` guard applies uniformly.
503 + for cohere_only_key in ("template_jinja", "template"):
504 + if chat_template_kwargs.get(cohere_only_key) is not None:
505 + raise ValueError(
506 + f"chat_template_kwargs.{cohere_only_key!r} is not accepted "
507 + "in vLLM (it is a Cohere-only field); pass the template "
508 + "body via the standard 'chat_template' request field "
509 + "instead (requires --trust-request-chat-template)."
510 + )
511 +
512 + config: dict[str, Any] = {
513 + "messages": _conversation_to_melody_messages(
514 + conversation,
515 + chat_template_kwargs.get(MESSAGES_CITATIONS_KEY),
516 + ),
517 + }
518 +
519 + # ``template_id`` is a selector for one of melody's built-in template
520 + # variants (not raw source), so it is safe to accept from the client.
521 + if template_id := chat_template_kwargs.get("template_id"):
522 + config["template_id"] = template_id
523 + # Raw template source flows exclusively through the standard vLLM
524 + # ``chat_template`` field so it is subject to the
525 + # ``--trust-request-chat-template`` guard in OnlineRenderer.
526 + if chat_template:
527 + config["template_jinja"] = chat_template
528 + # Only support Jinja with vllm
529 + config["use_jinja"] = True
530 +
531 + # Documents
532 + documents = chat_template_kwargs.get("documents") or []
533 + if documents:
534 + config["documents"] = [_document_to_melody(d) for d in documents]
535 +
536 + # Tools - prefer explicit ``available_tools``, fall back to OpenAI ``tools``
537 + tools = (
538 + chat_template_kwargs.get("available_tools")
539 + or chat_template_kwargs.get("tools")
540 + or []
541 + )
542 + if tools:
543 + config["available_tools"] = [_tool_to_melody(t) for t in tools]
544 +
545 + # Reasoning toggle (cmd3 + cmd4)
546 + if (rt := chat_template_kwargs.get("reasoning_type")) is not None:
547 + config["reasoning_type"] = str(rt)
548 + elif "thinking" in chat_template_kwargs:
549 + # Cohere v2 ``thinking: {type: enabled|disabled}`` shorthand
550 + thinking = chat_template_kwargs["thinking"]
551 + t = thinking.get("type") if isinstance(thinking, dict) else thinking
552 + if t in ("enabled", "disabled"):
553 + config["reasoning_type"] = t
554 +
555 + if (di := chat_template_kwargs.get("dev_instruction")) is not None:
556 + config["dev_instruction"] = str(di)
557 +
558 + # JSON / structured outputs
559 + if (rf := chat_template_kwargs.get("response_format")) is not None:
560 + rf = rf.model_dump() if hasattr(rf, "model_dump") else dict(rf)
561 + rf_type = rf.get("type")
562 + if rf_type == "json_object":
563 + config["json_mode"] = True
564 + elif rf_type in ("json_schema", "json"):
565 + schema = rf.get("schema") or rf.get("json_schema")
566 + if isinstance(schema, dict) and "schema" in schema:
567 + schema = schema["schema"]
568 + if schema is not None:
569 + config["json_schema"] = (
570 + schema if isinstance(schema, str) else json.dumps(schema)
571 + )
572 + if (js := chat_template_kwargs.get("json_schema")) is not None:
573 + config["json_schema"] = js if isinstance(js, str) else json.dumps(js)
574 + if "json_mode" in chat_template_kwargs:
575 + config["json_mode"] = bool(chat_template_kwargs["json_mode"])
576 +
577 + if fmt == "cmd3":
578 + if (sm := chat_template_kwargs.get("safety_mode")) is not None:
579 + config["safety_mode"] = str(sm).lower()
580 + # citation_quality: ``on`` / ``off``
581 + cq = chat_template_kwargs.get("citation_quality")
582 + if cq is None and (co := chat_template_kwargs.get("citation_options")):
583 + mode = co.get("mode") if isinstance(co, dict) else None
584 + if mode is not None:
585 + cq = "on" if str(mode).lower() != "off" else "off"
586 + if cq is not None:
587 + config["citation_quality"] = str(cq).lower()
588 + if "skip_preamble" in chat_template_kwargs:
589 + config["skip_preamble"] = bool(chat_template_kwargs["skip_preamble"])
590 + else: # cmd4
591 + # cmd4 uses ``grounding`` rather than safety_mode/citation_quality.
592 + # melody's cmd4 only accepts ``unknown`` / ``enabled`` / ``disabled``,
593 + # so the Cohere v2 ``citation_options.mode`` values
594 + # (``FAST`` / ``ACCURATE`` / ``OFF``) have to be normalized -- a
595 + # raw lowercased passthrough would raise from
596 + # ``render_cmd4`` (cmd4 has no fast/accurate distinction at the
597 + # prompt-template layer; both request grounding-on).
598 + if (gr := chat_template_kwargs.get("grounding")) is not None:
599 + config["grounding"] = _normalize_cmd4_grounding(gr)
600 + elif co := chat_template_kwargs.get("citation_options"):
601 + mode = co.get("mode") if isinstance(co, dict) else None
602 + if mode is not None:
603 + config["grounding"] = _normalize_cmd4_grounding(mode)
604 + if (pi := chat_template_kwargs.get("platform_instruction")) is not None:
605 + config["platform_instruction"] = str(pi)
606 +
607 + # Anything we didn't explicitly interpret above is forwarded to melody
608 + # as a Jinja template variable. This matches vLLM's documented contract
609 + # for ``chat_template_kwargs`` ("kwargs accessible by the template")
610 + # and lets callers write e.g. ``{"reasoning_effort": "low"}`` directly
611 + # without a nested ``additional_template_fields`` wrapper.
612 + extra = {
613 + k: v
614 + for k, v in chat_template_kwargs.items()
615 + if k not in _RENDERER_CONSUMED_KEYS
616 + }
617 + if extra:
618 + config["additional_template_fields"] = extra
619 +
620 + return fmt, config
621 +
622 +
623 +class CohereRenderer(BaseRenderer[HfTokenizer]):
624 + """Renderer that templates Cohere prompts via the melody Rust bindings.
625 +
626 + Tokenization is delegated to the standard HF tokenizer; only the
627 + chat-template step is replaced with ``cohere_melody.render_cmd3`` /
628 + ``render_cmd4``. Enabled via ``--tokenizer-mode cohere``.
629 + """
630 +
631 + def __init__(
632 + self,
633 + config: VllmConfig,
634 + tokenizer: HfTokenizer | None,
635 + ) -> None:
636 + # Match HfRenderer in not mutating the cached tokenizer instance
637 + tokenizer = copy.copy(tokenizer)
638 + super().__init__(config, tokenizer)
639 +
640 + # Lazy import to keep `cohere_melody` an optional dependency
641 + self._melody = _try_import_melody()
642 + # ``render_cmd3`` / ``render_cmd4`` are pure CPU work; cache the
643 + # thread-pool wrapper once so the async path doesn't allocate a new
644 + # adapter on every request.
645 + self._render_async = make_async(self._render, executor=self._executor)
646 +
647 + def _render(self, fmt: str, config_dict: dict[str, Any]) -> str:
648 + if fmt == "cmd3":
649 + return self._melody.render_cmd3(config_dict)
650 + return self._melody.render_cmd4(config_dict)
651 +
652 + def render_messages(
653 + self,
654 + messages: list[ChatCompletionMessageParam],
655 + params: ChatParams,
656 + ) -> tuple[list[ConversationMessage], DictPrompt]:
657 + conversation, mm_data, mm_uuids = parse_chat_messages(
658 + messages,
659 + self.model_config,
660 + content_format="openai",
661 + media_io_kwargs=params.media_io_kwargs,
662 + mm_processor_kwargs=params.mm_processor_kwargs,
663 + )
664 +
665 + chat_template_kwargs = dict(params.chat_template_kwargs)
666 + fmt, config_dict = _build_render_config(
667 + conversation, chat_template_kwargs, params.chat_template
668 + )
669 + prompt_text = self._render(fmt, config_dict)
670 + prompt = parse_dec_only_prompt(prompt_text)
671 +
672 + if mm_data is not None:
673 + prompt["multi_modal_data"] = mm_data
674 + if mm_uuids is not None:
675 + prompt["multi_modal_uuids"] = mm_uuids
676 +
677 + return conversation, prompt
678 +
679 + async def render_messages_async(
680 + self,
681 + messages: list[ChatCompletionMessageParam],
682 + params: ChatParams,
683 + ) -> tuple[list[ConversationMessage], DictPrompt]:
684 + conversation, mm_data, mm_uuids = await parse_chat_messages_async(
685 + messages,
686 + self.model_config,
687 + content_format="openai",
688 + media_io_kwargs=params.media_io_kwargs,
689 + mm_processor_kwargs=params.mm_processor_kwargs,
690 + )
691 +
692 + chat_template_kwargs = dict(params.chat_template_kwargs)
693 + fmt, config_dict = _build_render_config(
694 + conversation, chat_template_kwargs, params.chat_template
695 + )
696 + prompt_text = await self._render_async(fmt, config_dict)
697 + prompt = parse_dec_only_prompt(prompt_text)
698 +
699 + if mm_data is not None:
700 + prompt["multi_modal_data"] = mm_data
701 + if mm_uuids is not None:
702 + prompt["multi_modal_uuids"] = mm_uuids
703 +
704 + return conversation, prompt
vllm/renderers/registry.py
+1
@@ -20,6 +20,7 @@ logger = init_logger(__name__)
20
21
22 _VLLM_RENDERERS = {
23 + "cohere": ("cohere", "CohereRenderer"),
24 "deepseek_v32": ("deepseek_v32", "DeepseekV32Renderer"),
25 "deepseek_v4": ("deepseek_v4", "DeepseekV4Renderer"),
26 "hf": ("hf", "HfRenderer"),
vllm/tokenizers/registry.py
+3
@@ -40,6 +40,9 @@ _MODEL_TYPES_WITH_INCORRECT_TOKENIZER_CLASS: set[str] = {
40 }
41
42 _VLLM_TOKENIZERS = {
43 + # ``cohere`` mode uses the standard cached HF tokenizer; only the
44 + # renderer (template stage) is replaced with a melody-based one.
45 + "cohere": ("hf", "CachedHfTokenizer"),
46 "deepseek_v32": ("deepseek_v32", "DeepseekV32Tokenizer"),
47 "deepseek_v4": ("deepseek_v4", "DeepseekV4Tokenizer"),
48 "hf": ("hf", "CachedHfTokenizer"),