From ec74cafd50519394421e6023e3df0b377701b111 Mon Sep 17 00:00:00 2001 From: Andrew Berneshawi Date: Tue, 11 Aug 2026 03:45:57 -0400 Subject: [PATCH 1/6] Add e2e test scripts --- scripts/inspect_cohere_renderer.py | 289 ++++++ scripts/test_cohere_citations_e2e.py | 1303 ++++++++++++++++++++++++++ scripts/test_cohere_output.py | 897 ++++++++++++++++++ scripts/test_cohere_v2_e2e.py | 1187 +++++++++++++++++++++++ 4 files changed, 3676 insertions(+) create mode 100644 scripts/inspect_cohere_renderer.py create mode 100644 scripts/test_cohere_citations_e2e.py create mode 100644 scripts/test_cohere_output.py create mode 100644 scripts/test_cohere_v2_e2e.py diff --git a/scripts/inspect_cohere_renderer.py b/scripts/inspect_cohere_renderer.py new file mode 100644 index 000000000000..a8fae4584f0d --- /dev/null +++ b/scripts/inspect_cohere_renderer.py @@ -0,0 +1,289 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inspect the prompt that the Cohere renderer produces for a chat request. + +This script bypasses the engine entirely: it constructs a `VllmConfig` via +`EngineArgs`, loads the renderer registered for the chosen tokenizer mode, +and runs `render_messages_async` on a sample chat. Use this to iterate on +`vllm/renderers/cohere.py` and the cmd3/cmd4 templates without having to +boot a model. + +Run with: + + .venv/bin/python scripts/inspect_cohere_renderer.py \ + --model hmellor/tiny-random-LlamaForCausalLM \ + --tokenizer-mode cohere + +Pass ``--tokenizer-mode hf`` to compare against the default Jinja-based +renderer for the same model. The model only needs a tokenizer/config on +disk; we never instantiate the engine, so any small public model works. +""" + +from __future__ import annotations + +import argparse +import asyncio +import json +import sys +from typing import Any + +from vllm.engine.arg_utils import EngineArgs +from vllm.entrypoints.cohere.protocol import CohereChatV2Request +from vllm.entrypoints.cohere.serving import CohereServingChatV2 +from vllm.renderers import ChatParams +from vllm.renderers.registry import RENDERER_REGISTRY +from vllm.tokenizers.registry import cached_tokenizer_from_config + +SAMPLE_MESSAGES: list[dict[str, Any]] = [ + {"role": "system", "content": "You answer concisely."}, + {"role": "user", "content": "Who wrote Hamlet, and when?"}, +] + +SAMPLE_DOCUMENTS = [ + { + "id": "doc_0", + "data": { + "title": "Wikipedia: Hamlet", + "text": ( + "Hamlet was written by William Shakespeare some time between" + " 1599 and 1601." + ), + }, + }, + { + "id": "doc_1", + "data": {"title": "Britannica", "text": "Shakespeare lived 1564-1616."}, + }, +] + +SAMPLE_TOOLS = [ + { + "type": "function", + "function": { + "name": "lookup_play_metadata", + "description": "Look up metadata about a Shakespeare play.", + "parameters": { + "type": "object", + "properties": { + "title": {"type": "string", "description": "The play's title"} + }, + "required": ["title"], + }, + }, + } +] + + +def _parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__) + p.add_argument( + "--model", + default="hmellor/tiny-random-LlamaForCausalLM", + help=( + "HF repo id or local path. Any model with a fast tokenizer and a " + "config.json works; we never load weights." + ), + ) + p.add_argument( + "--tokenizer-mode", + default="cohere", + choices=["auto", "hf", "cohere"], + help="Renderer/tokenizer mode to exercise. Defaults to 'cohere'.", + ) + p.add_argument( + "--cohere-format", + default="cmd3", + choices=["cmd3", "cmd4"], + help="When tokenizer-mode=cohere, which template family to render.", + ) + p.add_argument( + "--with-documents", + action="store_true", + help="Include sample documents (grounding) in chat_template_kwargs.", + ) + p.add_argument( + "--with-tools", + action="store_true", + help="Include sample tools in chat_template_kwargs.", + ) + p.add_argument( + "--safety-mode", + default=None, + choices=[None, "contextual", "strict", "none"], + help="cmd3 safety mode forwarded via chat_template_kwargs.", + ) + p.add_argument( + "--reasoning", + default=None, + choices=[None, "enabled", "disabled"], + help="Reasoning toggle forwarded via chat_template_kwargs.", + ) + p.add_argument( + "--show-token-ids", + action="store_true", + help="Print prompt_token_ids alongside the text.", + ) + p.add_argument( + "--trust-remote-code", + action="store_true", + help="Pass through to ModelConfig (needed for some Cohere repos).", + ) + p.add_argument( + "--from-v2-request", + action="store_true", + help=( + "Construct a Cohere v2 request, run the same" + " v2 -> ChatCompletion conversion that POST /cohere/v2/chat does, and" + " render the result through the configured renderer. Demonstrates" + " what a non-cohere (`--tokenizer-mode hf`) model actually sees" + " when called with the v2 input shape." + ), + ) + return p.parse_args() + + +def _build_v2_request(args: argparse.Namespace) -> CohereChatV2Request: + """Build a representative Cohere v2 request mirroring the script's flags.""" + body: dict[str, Any] = { + "model": args.model, + "messages": [ + {"role": "system", "content": "You answer concisely."}, + {"role": "user", "content": "Who wrote Hamlet, and when?"}, + ], + } + if args.with_documents: + body["documents"] = [ + { + "id": "doc_0", + "data": { + "title": "Wikipedia: Hamlet", + "text": ( + "Hamlet was written by William Shakespeare some time" + " between 1599 and 1601." + ), + }, + }, + { + "id": "doc_1", + "data": { + "title": "Britannica", + "text": "Shakespeare lived 1564-1616.", + }, + }, + ] + if args.with_tools: + body["tools"] = [ + { + "type": "function", + "function": { + "name": "lookup_play_metadata", + "description": "Look up metadata about a Shakespeare play.", + "parameters": { + "type": "object", + "properties": { + "title": { + "type": "string", + "description": "The play's title", + } + }, + "required": ["title"], + }, + }, + } + ] + if args.safety_mode: + body["safety_mode"] = args.safety_mode.upper() + if args.reasoning: + body["thinking"] = {"type": args.reasoning} + body["citation_options"] = {"mode": "FAST"} + return CohereChatV2Request.model_validate(body) + + +async def main() -> int: + args = _parse_args() + + # Build a VllmConfig the same way the engine does, but never start the + # engine. ``skip_tokenizer_init=False`` is the default and is what the + # renderer needs. + eng_args = EngineArgs( + model=args.model, + tokenizer_mode=args.tokenizer_mode, + trust_remote_code=args.trust_remote_code, + # Keep this small so HF config validation is fast on tiny models. + max_model_len=4096, + # We don't actually run the engine; just satisfy the validators. + enforce_eager=True, + ) + config = eng_args.create_engine_config() + + tokenizer = cached_tokenizer_from_config(config.model_config) + renderer = RENDERER_REGISTRY.load_renderer( + args.tokenizer_mode if args.tokenizer_mode != "auto" else "hf", + config, + tokenizer, + ) + + chat_template_kwargs: dict[str, Any] = {} + messages: list[dict[str, Any]] = SAMPLE_MESSAGES + + if args.from_v2_request: + v2_req = _build_v2_request(args) + chat_req = CohereServingChatV2._convert_v2_to_chat_completion(v2_req) + messages = [ + m.model_dump(exclude_none=True) if hasattr(m, "model_dump") else m + for m in chat_req.messages + ] + chat_template_kwargs.update(chat_req.chat_template_kwargs or {}) + if args.tokenizer_mode == "cohere": + chat_template_kwargs["cohere_format"] = args.cohere_format + if chat_req.tools: + chat_template_kwargs.setdefault( + "tools", + [t.model_dump(exclude_none=True) for t in chat_req.tools], + ) + else: + if args.tokenizer_mode == "cohere": + chat_template_kwargs["cohere_format"] = args.cohere_format + if args.with_documents: + chat_template_kwargs["documents"] = SAMPLE_DOCUMENTS + if args.with_tools: + chat_template_kwargs["tools"] = SAMPLE_TOOLS + if args.safety_mode is not None: + chat_template_kwargs["safety_mode"] = args.safety_mode + if args.reasoning is not None: + chat_template_kwargs["reasoning_type"] = args.reasoning + + params = ChatParams(chat_template_kwargs=chat_template_kwargs) + + print("=" * 72) + print( + f"model={args.model} tokenizer_mode={args.tokenizer_mode}" + + (f" format={args.cohere_format}" if args.tokenizer_mode == "cohere" else "") + ) + print(f"renderer={type(renderer).__name__}") + if chat_template_kwargs: + print("chat_template_kwargs:") + print(json.dumps(chat_template_kwargs, indent=2, default=str)) + print("=" * 72) + + conversation, prompt = await renderer.render_messages_async(messages, params) + + print("\n--- conversation (post parse_chat_messages) ---") + print(json.dumps(conversation, indent=2, default=str)) + + print("\n--- rendered prompt text ---") + text = prompt.get("prompt") + if text is None and "prompt_token_ids" in prompt: + text = tokenizer.decode(prompt["prompt_token_ids"]) + print(text) + + if args.show_token_ids and "prompt_token_ids" in prompt: + print("\n--- prompt_token_ids ---") + print(prompt["prompt_token_ids"]) + + return 0 + + +if __name__ == "__main__": + sys.exit(asyncio.run(main())) diff --git a/scripts/test_cohere_citations_e2e.py b/scripts/test_cohere_citations_e2e.py new file mode 100644 index 000000000000..03d01680700f --- /dev/null +++ b/scripts/test_cohere_citations_e2e.py @@ -0,0 +1,1303 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""End-to-end smoke test for Cohere Chat v2 documents/citations behavior. + +Drives a live vLLM server with a real Cohere Command-family model, focused +specifically on grounding: multi-shape ``documents``, ``citation_options`` +modes, citing prior tool results, citations coexisting with ``thinking``, +and citations already present in the conversation history. + +This complements ``scripts/test_cohere_v2_e2e.py`` (which covers the +broader v2 surface -- tools, streaming lifecycle, error paths -- with only +a couple of basic grounding checks). This script instead ports one test +case per relevant grounding scenario recorded in blobheart's dory replay +suite (``go/dory/pkg/replay/data_replay/``), so the citation surface has +the same scenario coverage server-side that dory already has. + +Cases NOT ported, and why: + +* ``chat_documents_non_stream`` / ``chat_documents_stream`` / + ``chat_citations_fast`` / ``chat_citations_off`` / ``cmd3_plan_citation`` + -- these exercise dory's v1 ``/v1/chat`` API (``citation_quality`` + string field). vLLM only exposes the v2 ``/cohere/v2/chat`` endpoint; + the same grounding behavior is exercised here via v2's + ``citation_options.mode`` instead (see ``test_citation_options_fast`` + and ``test_citation_options_off``). +* ``debugging_rag`` / ``debugging_rag_stream`` -- exercise dory's + ``enable_debugging`` request field, which has no vLLM equivalent. +* ``return_prompt_rag`` -- exercises dory's ``return_prompt`` field, + which has no vLLM equivalent (there's no "echo the rendered prompt" + knob on ``CohereChatV2Request``). +* ``tool_multihop_citations_off`` / ``tool_calls_result`` / + ``tool_multihop_3`` / ``v2_tool_multihop_3`` -- primarily tool + orchestration tests (multihop tool calling); already covered by the + tool-call tests in ``test_cohere_v2_e2e.py``, so not duplicated here. +* ``cmda_thinking_chat_citations_in_history_from_turns`` / + ``..._turns`` -- the latter is an incomplete/empty recording upstream + in dory (WIP); the former is materially similar to + ``cmda_thinking_chat_citations_in_history`` (ported below as + ``test_citations_in_conversation_history``), just with citations + attached to a ``thinking`` block instead of ``text`` blocks. + +Prerequisites on the GPU host +----------------------------- + +1. Install the optional Cohere SDKs: + + uv pip install cohere cohere-melody + +2. Run this script. If ``--base-url`` isn't already serving a healthy + ``/health`` response, the script launches ``vllm serve`` itself with + the Cohere renderer / tokenizer wired up (``VLLM_ENABLE_COHERE_API=1``, + ``--tokenizer-mode cohere``, ``--enable-auto-tool-choice``, + ``--tool-call-parser cohere2``, ``--reasoning-parser cohere2``), waits + for it to become healthy, then runs the test suite: + + python scripts/test_cohere_citations_e2e.py + + defaults to the smallest available Command A+ checkpoint, + ``CohereLabs/command-a-plus-05-2026-w4a4``. Pass ``--model`` to use a + different one, or point ``--base-url`` at a server you started + yourself (e.g. in another terminal, per the manual invocation below) + and pass ``--no-auto-start-server`` to make that mandatory instead of + just preferred: + + VLLM_ENABLE_COHERE_API=1 vllm serve \\ + --tokenizer-mode cohere \\ + --enable-auto-tool-choice \\ + --tool-call-parser cohere2 \\ + --reasoning-parser cohere2 \\ + --port 8000 + + For non-reasoning Command models (cmd3, older Command R), pass + ``--no-reasoning-model`` to this script; when it owns the server + process, that also adds ``--no-cohere-is-reasoning-model`` to the + ``vllm serve`` invocation. + +Optional flags +-------------- + +* ``--reasoning-model / --no-reasoning-model`` -- whether the server was + (or should be) launched with reasoning enabled. Controls whether the + ``thinking`` cases run, and whether an auto-started server gets + ``--no-cohere-is-reasoning-model``. +* ``--auto-start-server / --no-auto-start-server`` -- whether to launch + ``vllm serve`` ourselves when ``--base-url`` isn't already up + (default: on). +* ``--keep-server / --no-keep-server`` -- leave an auto-started server + running after the tests finish, so re-runs skip the (often lengthy) + model load (default: on -- the PID and log path are printed so you can + stop it manually). +* ``--extra-server-arg`` -- repeatable; extra ``vllm serve`` args (e.g. + ``--extra-server-arg=--tensor-parallel-size=8``) forwarded verbatim + when we own the server process. +* ``--startup-timeout`` -- seconds to wait for an auto-started server to + become healthy (default: 1800; large quantized checkpoints are slow to + load). +* ``--verbose`` -- dump full response bodies / SSE frames. +""" + +from __future__ import annotations + +import argparse +import contextlib +import json +import os +import signal +import subprocess +import sys +import tempfile +import time +from dataclasses import dataclass, field +from typing import Any +from urllib.parse import urlparse + +import httpx + +DONE_LINE = "data: [DONE]" + +# Smallest publicly available Command A+ checkpoint at the time of +# writing; picked as the default so this script is runnable without +# having to hunt down a model id first. +DEFAULT_MODEL = "CohereLabs/command-a-plus-05-2026-w4a4" + +# Reasoning models spend most of their decode budget inside ``thinking`` +# blocks before emitting user-visible text/citations; give grounding +# requests a generous budget so citations actually get a chance to appear. +REASONING_BUDGET = 2048 + +PENGUIN_DOCUMENTS = [ + {"data": {"snippet": "The tallest penguin is the Emperor penguin"}}, + { + "data": { + "snippet": "The latin name for Emperor penguin is Aptenodytes forsteri" + } + }, + {"data": {"snippet": "The smallest penguin is the fairy penguin"}}, + {"data": {"snippet": "The latin name for fairy penguin is Eudyptula minor"}}, +] + + +# ---------------------------------------------------------------------- +# Helpers (mirrors scripts/test_cohere_v2_e2e.py) +# ---------------------------------------------------------------------- + + +def _text_from_content_blocks(content: Any) -> str: + if not isinstance(content, list): + return "" + return "".join( + b.get("text", "") + for b in content + if isinstance(b, dict) and b.get("type") == "text" + ) + + +def _thinking_from_content_blocks(content: Any) -> str: + if not isinstance(content, list): + return "" + return "".join( + b.get("thinking", "") + for b in content + if isinstance(b, dict) and b.get("type") == "thinking" + ) + + +@dataclass +class TestResult: + name: str + passed: bool + detail: str = "" + skipped: bool = False + + +@dataclass +class TestContext: + base_url: str + model: str + is_reasoning_model: bool + verbose: bool + timeout: float + request_id_counter: int = 0 + results: list[TestResult] = field(default_factory=list) + + def next_request_id(self, label: str) -> str: + self.request_id_counter += 1 + return f"e2e-cite-{label}-{self.request_id_counter}-{int(time.time())}" + + def record(self, result: TestResult) -> None: + self.results.append(result) + prefix = "SKIP" if result.skipped else ("PASS" if result.passed else "FAIL") + print(f"[{prefix}] {result.name}") + if result.detail: + for line in result.detail.splitlines(): + print(f" {line}") + + +def _post_json( + ctx: TestContext, + *, + body: dict[str, Any], + request_id: str, + expect_status: int | None = 200, +) -> tuple[int, dict[str, Any] | str]: + url = f"{ctx.base_url.rstrip('/')}/cohere/v2/chat" + headers = {"Content-Type": "application/json", "X-Request-Id": request_id} + with httpx.Client(timeout=ctx.timeout) as client: + resp = client.post(url, json=body, headers=headers) + if ctx.verbose: + print(f" -> POST {url} [status {resp.status_code}]") + print(f" request_id={request_id}") + print(f" body={json.dumps(body)[:800]}") + print(f" resp={resp.text[:2000]}") + parsed: dict[str, Any] | str + try: + parsed = resp.json() + except Exception: + parsed = resp.text + if expect_status is not None and resp.status_code != expect_status: + raise AssertionError( + f"expected status {expect_status}, got {resp.status_code}: {parsed!r}" + ) + return resp.status_code, parsed + + +def _stream_post( + ctx: TestContext, + *, + body: dict[str, Any], + request_id: str, +) -> tuple[list[dict[str, Any]], bool]: + body = {**body, "stream": True} + url = f"{ctx.base_url.rstrip('/')}/cohere/v2/chat" + headers = { + "Content-Type": "application/json", + "Accept": "text/event-stream", + "X-Request-Id": request_id, + } + events: list[dict[str, Any]] = [] + saw_done = False + with ( + httpx.Client(timeout=ctx.timeout) as client, + client.stream("POST", url, json=body, headers=headers) as resp, + ): + if resp.status_code != 200: + resp.read() + raise AssertionError( + f"stream expected 200, got {resp.status_code}: {resp.text[:1000]}" + ) + buffer = "" + for chunk in resp.iter_text(): + if not chunk: + continue + buffer += chunk + while "\n\n" in buffer: + frame, buffer = buffer.split("\n\n", 1) + for line in frame.splitlines(): + line = line.strip() + if not line.startswith("data:"): + continue + payload = line[len("data:") :].strip() + if payload == "[DONE]": + saw_done = True + continue + if not payload: + continue + try: + events.append(json.loads(payload)) + except json.JSONDecodeError as e: + raise AssertionError(f"bad SSE frame: {payload!r}: {e}") from e + if ctx.verbose: + print(f" -> stream {url} events={len(events)} done={saw_done}") + for ev in events: + print(f" {ev.get('type'):<20} {json.dumps(ev)[:200]}") + return events, saw_done + + +def _event_types(events: list[dict[str, Any]]) -> list[str]: + return [ev.get("type", "") for ev in events] + + +def _expect(condition: bool, message: str) -> None: + if not condition: + raise AssertionError(message) + + +def _assert_citation_shape(citations: list[dict[str, Any]]) -> None: + for c in citations: + _expect("start" in c and "end" in c, f"citation missing span: {c}") + sources = c.get("sources") or [] + _expect(sources, f"citation without sources: {c}") + for src in sources: + _expect("type" in src, f"citation source missing type: {src}") + + +# ---------------------------------------------------------------------- +# Optional: manage a vLLM server process ourselves +# ---------------------------------------------------------------------- + + +def _server_is_healthy(base_url: str, timeout: float = 5.0) -> bool: + try: + with httpx.Client(timeout=timeout) as client: + resp = client.get(f"{base_url.rstrip('/')}/health") + return resp.status_code == 200 + except (httpx.HTTPError, OSError): + return False + + +@dataclass +class ManagedServer: + proc: subprocess.Popen + log_path: str + + def stop(self, *, timeout: float = 30.0) -> None: + if self.proc.poll() is not None: + return + pgid: int | None = None + with contextlib.suppress(ProcessLookupError, OSError): + pgid = os.getpgid(self.proc.pid) + target = -pgid if pgid is not None else self.proc.pid + with contextlib.suppress(ProcessLookupError, OSError): + os.kill(target, signal.SIGTERM) + try: + self.proc.wait(timeout=timeout) + except subprocess.TimeoutExpired: + with contextlib.suppress(ProcessLookupError, OSError): + os.kill(target, signal.SIGKILL) + with contextlib.suppress(subprocess.TimeoutExpired): + self.proc.wait(timeout=10) + + +def _build_vllm_serve_command( + *, + model: str, + host: str, + port: int, + is_reasoning_model: bool, + extra_args: list[str], +) -> list[str]: + cmd = [ + "vllm", + "serve", + model, + "--host", + host, + "--port", + str(port), + "--tokenizer-mode", + "cohere", + "--enable-auto-tool-choice", + "--tool-call-parser", + "cohere2", + "--reasoning-parser", + "cohere2", + ] + if not is_reasoning_model: + cmd.append("--no-cohere-is-reasoning-model") + cmd.extend(extra_args) + return cmd + + +def _start_vllm_server( + *, + model: str, + host: str, + port: int, + is_reasoning_model: bool, + extra_args: list[str], +) -> ManagedServer: + cmd = _build_vllm_serve_command( + model=model, + host=host, + port=port, + is_reasoning_model=is_reasoning_model, + extra_args=extra_args, + ) + env = os.environ.copy() + env["VLLM_ENABLE_COHERE_API"] = "1" + + log_fd, log_path = tempfile.mkstemp( + prefix="vllm-cohere-citations-e2e-", suffix=".log" + ) + os.close(log_fd) + log_file = open(log_path, "w") # noqa: SIM115 -- lives as long as the server + + print(f"Server not reachable; starting it ourselves:\n {' '.join(cmd)}") + print(f" logs: {log_path}") + proc = subprocess.Popen( + cmd, + env=env, + stdout=log_file, + stderr=subprocess.STDOUT, + start_new_session=True, + ) + return ManagedServer(proc=proc, log_path=log_path) + + +def _wait_for_server( + *, base_url: str, server: ManagedServer, timeout: float +) -> None: + deadline = time.time() + timeout + while time.time() < deadline: + exit_code = server.proc.poll() + if exit_code is not None: + raise RuntimeError( + f"vLLM server process exited early (code={exit_code}); " + f"see {server.log_path} for details." + ) + if _server_is_healthy(base_url): + print(f"Server is healthy at {base_url}") + return + time.sleep(5.0) + raise TimeoutError( + f"vLLM server did not become healthy within {timeout:.0f}s; " + f"see {server.log_path} for details." + ) + + +def _no_grounding_warning(content: Any) -> str: + text = _text_from_content_blocks(content) or "" + thinking = _thinking_from_content_blocks(content) or "" + return ( + "WARN: assistant did not ground its answer.\n" + f" text head: {text[:200]!r}\n" + f" thinking head: {thinking[:200]!r}\n" + " If the text contains no ' None: + """Basic grounded answer, non-streaming. + + Ports dory's ``v2_chat_documents_non_stream`` replay case: plain + ``{"data": {...}}`` documents (no explicit ``id``), single user turn. + """ + name = "documents: basic grounded answer (non-streaming)" + try: + body = { + "model": ctx.model, + "messages": [ + {"role": "user", "content": "What is the tallest penguin?"}, + ], + "documents": PENGUIN_DOCUMENTS, + "citation_options": {"mode": "ACCURATE"}, + "max_tokens": REASONING_BUDGET, + "temperature": 0.0, + "seed": 0, + } + request_id = ctx.next_request_id("docs-basic") + _, resp = _post_json(ctx, body=body, request_id=request_id) + assert isinstance(resp, dict) + message = resp.get("message") or {} + content = message.get("content") or [] + citations = message.get("citations") or [] + _expect(content, f"empty content: {message}") + if not citations: + ctx.record(TestResult(name, True, detail=_no_grounding_warning(content))) + return + _assert_citation_shape(citations) + doc_ids = { + s.get("id") or s.get("document", {}).get("id") + for c in citations + for s in c.get("sources", []) + } + ctx.record( + TestResult( + name, + True, + detail=( + f"got {len(citations)} citation(s); " + f"doc ids referenced={doc_ids}" + ), + ) + ) + except Exception as e: + ctx.record(TestResult(name, False, detail=str(e))) + + +def test_documents_basic_streaming(ctx: TestContext) -> None: + """Basic grounded answer, streaming. + + Ports dory's ``v2_chat_documents_stream`` replay case. + """ + name = "documents: basic grounded answer (streaming)" + try: + body = { + "model": ctx.model, + "messages": [ + {"role": "user", "content": "What is the tallest penguin?"}, + ], + "documents": PENGUIN_DOCUMENTS, + "citation_options": {"mode": "ACCURATE"}, + "max_tokens": REASONING_BUDGET, + "temperature": 0.0, + "seed": 0, + } + request_id = ctx.next_request_id("docs-basic-stream") + events, saw_done = _stream_post(ctx, body=body, request_id=request_id) + _expect(saw_done, "stream did not terminate with data: [DONE]") + types = _event_types(events) + _expect( + types[0] == "message-start" and types[-1] == "message-end", + f"unexpected envelope: {types[:3]} ... {types[-3:]}", + ) + starts = [i for i, t in enumerate(types) if t == "citation-start"] + ends = [i for i, t in enumerate(types) if t == "citation-end"] + if not starts: + ctx.record( + TestResult( + name, + True, + detail="WARN: model did not emit citation events for this stream.", + ) + ) + return + _expect( + len(starts) == len(ends), + f"unbalanced citation events: starts={len(starts)} ends={len(ends)}", + ) + for s, e in zip(starts, ends): + _expect(s < e, f"citation-start at {s} must precede end at {e}") + detail = f"events={len(events)} citations={len(starts)}" + ctx.record(TestResult(name, True, detail=detail)) + except Exception as e: + ctx.record(TestResult(name, False, detail=str(e))) + + +def test_citation_options_fast(ctx: TestContext) -> None: + """``citation_options.mode=FAST`` still grounds, via a streaming call. + + Ports dory's ``v2_chat_citations_fast`` (and v1 ``chat_citations_fast``, + which is the same scenario against the ``/v1/chat`` API) replay cases. + """ + name = "citation_options: mode=FAST grounds the answer (streaming)" + try: + body = { + "model": ctx.model, + "messages": [ + {"role": "user", "content": "What is the tallest penguin?"}, + ], + "documents": PENGUIN_DOCUMENTS, + "citation_options": {"mode": "FAST"}, + "max_tokens": REASONING_BUDGET, + "temperature": 0.0, + "seed": 0, + } + request_id = ctx.next_request_id("cite-fast") + events, saw_done = _stream_post(ctx, body=body, request_id=request_id) + _expect(saw_done, "stream did not terminate with data: [DONE]") + types = _event_types(events) + starts = [i for i, t in enumerate(types) if t == "citation-start"] + if not starts: + ctx.record( + TestResult( + name, + True, + detail="WARN: mode=FAST did not produce citation events.", + ) + ) + return + ctx.record(TestResult(name, True, detail=f"citations={len(starts)}")) + except Exception as e: + ctx.record(TestResult(name, False, detail=str(e))) + + +def test_citation_options_off(ctx: TestContext) -> None: + """``citation_options.mode=OFF`` must suppress citation generation. + + Ports dory's ``v2_chat_citations_off`` (and v1 ``chat_citations_off``) + replay cases, where the recorded golden response has an *empty* + ``citations`` array/no citation events despite grounding documents + being present. + """ + name = "citation_options: mode=OFF suppresses citations (streaming)" + try: + body = { + "model": ctx.model, + "messages": [ + {"role": "user", "content": "What is the tallest penguin?"}, + ], + "documents": PENGUIN_DOCUMENTS, + "citation_options": {"mode": "OFF"}, + "max_tokens": REASONING_BUDGET, + "temperature": 0.0, + "seed": 0, + } + request_id = ctx.next_request_id("cite-off") + events, saw_done = _stream_post(ctx, body=body, request_id=request_id) + _expect(saw_done, "stream did not terminate with data: [DONE]") + types = _event_types(events) + _expect( + "citation-start" not in types, + f"mode=OFF but citation-start events were emitted: {types}", + ) + text_chunks = [ + (ev.get("delta") or {}).get("message", {}).get("content", {}).get("text") + for ev in events + if ev.get("type") == "content-delta" + ] + text = "".join(t for t in text_chunks if isinstance(t, str)) + ctx.record(TestResult(name, True, detail=f"no citation events; text={text!r}")) + except Exception as e: + ctx.record(TestResult(name, False, detail=str(e))) + + +def test_mixed_document_shapes(ctx: TestContext) -> None: + """Documents can be a ``{id, data}`` object, a bare string, or a JSON + string all in the same request; citations should be able to point at + any of them. + + Ports dory's ``cmd3_documents`` replay case. + """ + name = "documents: mixed shapes (object / plain string / JSON string)" + try: + body = { + "model": ctx.model, + "messages": [ + { + "role": "user", + "content": "What is the tallest mountain on Mars and Venus?", + }, + ], + "documents": [ + { + "id": "mars_1", + "data": { + "mountain": "Olympus Mons", + "location": "Mars", + "height": 21088, + }, + }, + "Skadi Mons is the tallest mountain on Venus", + '{"location": "Earth", "mountain": "Mount Everest", "height": 29029}', + ], + "citation_options": {"mode": "ACCURATE"}, + "max_tokens": REASONING_BUDGET, + "temperature": 0.0, + "seed": 0, + } + request_id = ctx.next_request_id("mixed-docs") + _, resp = _post_json(ctx, body=body, request_id=request_id) + assert isinstance(resp, dict) + message = resp.get("message") or {} + content = message.get("content") or [] + citations = message.get("citations") or [] + _expect(content, f"empty content: {message}") + if not citations: + ctx.record(TestResult(name, True, detail=_no_grounding_warning(content))) + return + _assert_citation_shape(citations) + doc_ids = { + s.get("id") or s.get("document", {}).get("id") + for c in citations + for s in c.get("sources", []) + } + ctx.record( + TestResult( + name, + True, + detail=( + f"got {len(citations)} citation(s) across mixed doc shapes; " + f"doc ids referenced={doc_ids}" + ), + ) + ) + except Exception as e: + ctx.record(TestResult(name, False, detail=str(e))) + + +def test_grounding_against_tool_result_in_history(ctx: TestContext) -> None: + """Ask the model to retroactively cite an answer it already gave, + grounding against a tool result from earlier in the conversation. + + Ports dory's ``cmd3_cite_old_tool_result`` replay case: the history + already contains a completed ``internet_search`` tool call/result and + an *ungrounded* assistant answer; the new user turn asks for the same + answer "with grounding", which should produce citations against the + tool's document-shaped result (and/or the top-level ``documents``). + """ + name = "documents: ground a prior answer against a tool result in history" + try: + body = { + "model": ctx.model, + "documents": ["The tallest mountain is mount everest"], + "messages": [ + { + "role": "user", + "content": "What are the two tallest mountains?", + }, + { + "role": "assistant", + "tool_plan": ( + "I will search for the second tallest mountain." + ), + "tool_calls": [ + { + "id": "internet_search_0123", + "type": "function", + "function": { + "name": "internet_search", + "arguments": json.dumps( + {"query": "second tallest mountain"} + ), + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "internet_search_0123", + "content": [ + { + "type": "document", + "document": { + "data": { + "result": "The second tallest mountain is K2." + } + }, + } + ], + }, + { + "role": "assistant", + "content": "The two tallest mountains are Mount Everest and K2", + }, + { + "role": "user", + "content": "Great. Can you repeat that with grounding?", + }, + ], + "tools": [ + { + "type": "function", + "function": { + "name": "internet_search", + "description": "Searches the internet", + "parameters": { + "type": "object", + "properties": { + "query": { + "type": "string", + "description": ( + "The query to search the internet with" + ), + } + }, + "required": ["query"], + }, + }, + } + ], + "max_tokens": REASONING_BUDGET, + "temperature": 0.0, + "seed": 0, + } + request_id = ctx.next_request_id("cite-old-tool-result") + _, resp = _post_json(ctx, body=body, request_id=request_id) + assert isinstance(resp, dict) + message = resp.get("message") or {} + content = message.get("content") or [] + citations = message.get("citations") or [] + tool_calls = message.get("tool_calls") or [] + # The model might legitimately re-issue the search instead of + # citing the existing result -- treat that as informative, not a + # hard failure, since the wiring under test is "does grounding + # against history work at all", not "does the model always skip + # redundant tool calls". + if tool_calls and not citations: + ctx.record( + TestResult( + name, + True, + detail=( + "WARN: model re-issued a tool call instead of citing " + f"history: {tool_calls}" + ), + ) + ) + return + _expect(content, f"empty content: {message}") + if not citations: + ctx.record(TestResult(name, True, detail=_no_grounding_warning(content))) + return + _assert_citation_shape(citations) + source_types = { + s.get("type") for c in citations for s in c.get("sources", []) + } + ctx.record( + TestResult( + name, + True, + detail=( + f"got {len(citations)} citation(s); source types={source_types} " + "(expect at least one 'tool' or 'document' source)" + ), + ) + ) + except Exception as e: + ctx.record(TestResult(name, False, detail=str(e))) + + +def test_thinking_with_citations_non_stream(ctx: TestContext) -> None: + """Reasoning (``thinking.type=enabled``) and grounding must coexist. + + Ports dory's ``cmda_thinking_with_citations`` replay case. + """ + name = "documents: thinking + citations coexist (non-streaming)" + if not ctx.is_reasoning_model: + ctx.record( + TestResult( + name, True, skipped=True, detail="server is not a reasoning model" + ) + ) + return + try: + body = { + "model": ctx.model, + "messages": [ + {"role": "user", "content": "What is the tallest penguin?"}, + ], + "documents": PENGUIN_DOCUMENTS, + "citation_options": {"mode": "FAST"}, + "thinking": {"type": "enabled"}, + "max_tokens": REASONING_BUDGET, + "temperature": 0.0, + "seed": 0, + } + request_id = ctx.next_request_id("thinking-cite") + _, resp = _post_json(ctx, body=body, request_id=request_id) + assert isinstance(resp, dict) + message = resp.get("message") or {} + content = message.get("content") or [] + citations = message.get("citations") or [] + thinking = _thinking_from_content_blocks(content) + _expect(content, f"empty content: {message}") + if not citations: + ctx.record( + TestResult( + name, + True, + detail=( + f"WARN: no citations; thinking_chars={len(thinking)}\n" + + _no_grounding_warning(content) + ), + ) + ) + return + _assert_citation_shape(citations) + ctx.record( + TestResult( + name, + True, + detail=( + f"got {len(citations)} citation(s); thinking_chars={len(thinking)}" + ), + ) + ) + except Exception as e: + ctx.record(TestResult(name, False, detail=str(e))) + + +def test_thinking_with_citations_streaming(ctx: TestContext) -> None: + """Streaming variant: thinking content-blocks must not disrupt the + citation-start/-end event pairing. + + Ports dory's ``cmda_thinking_stream_with_citations`` replay case. + """ + name = "documents: thinking + citations coexist (streaming)" + if not ctx.is_reasoning_model: + ctx.record( + TestResult( + name, True, skipped=True, detail="server is not a reasoning model" + ) + ) + return + try: + body = { + "model": ctx.model, + "messages": [ + {"role": "user", "content": "What is the tallest penguin?"}, + ], + "documents": PENGUIN_DOCUMENTS, + "citation_options": {"mode": "FAST"}, + "thinking": {"type": "enabled"}, + "max_tokens": REASONING_BUDGET, + "temperature": 0.0, + "seed": 0, + } + request_id = ctx.next_request_id("thinking-cite-stream") + events, saw_done = _stream_post(ctx, body=body, request_id=request_id) + _expect(saw_done, "stream did not terminate with data: [DONE]") + types = _event_types(events) + _expect( + types[0] == "message-start" and types[-1] == "message-end", + f"unexpected envelope: {types[:3]} ... {types[-3:]}", + ) + starts = [i for i, t in enumerate(types) if t == "citation-start"] + ends = [i for i, t in enumerate(types) if t == "citation-end"] + if not starts: + ctx.record( + TestResult( + name, + True, + detail=f"WARN: no citation events; total events={len(events)}", + ) + ) + return + _expect( + len(starts) == len(ends), + f"unbalanced citation events: starts={len(starts)} ends={len(ends)}", + ) + for s, e in zip(starts, ends): + _expect(s < e, f"citation-start at {s} must precede end at {e}") + detail = f"events={len(events)} citations={len(starts)}" + ctx.record(TestResult(name, True, detail=detail)) + except Exception as e: + ctx.record(TestResult(name, False, detail=str(e))) + + +def test_citations_in_conversation_history(ctx: TestContext) -> None: + """A prior assistant turn already carries a populated ``citations`` + array (with both ``document`` and ``tool`` sources, one of them + anchored to a ``PLAN``/thinking content index); the server must accept + that history and continue grounding on the follow-up turn. + + Ports dory's ``cmda_thinking_chat_citations_in_history`` replay case. + """ + name = "documents: history already contains citations; follow-up still grounds" + if not ctx.is_reasoning_model: + ctx.record( + TestResult( + name, True, skipped=True, detail="server is not a reasoning model" + ) + ) + return + try: + body = { + "model": ctx.model, + "documents": [ + { + "data": { + "title": "Weather in Tokyo", + "snippet": "The weather in tokyo is 27 degrees", + } + } + ], + "thinking": {"type": "enabled"}, + "messages": [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": ( + "What is the weather in London, Toronto, and Tokyo?" + ), + } + ], + }, + { + "role": "assistant", + "content": [ + { + "type": "thinking", + "thinking": "I will use the get_weather tool", + } + ], + "tool_calls": [ + { + "id": "tool_call_1", + "type": "function", + "function": { + "name": "get_weather", + "arguments": json.dumps({"location": "Toronto"}), + }, + }, + { + "id": "tool_call_2", + "type": "function", + "function": { + "name": "get_weather", + "arguments": json.dumps({"location": "London"}), + }, + }, + ], + }, + { + "role": "tool", + "tool_call_id": "tool_call_1", + "content": [ + {"type": "text", "text": "it's colder than usual"}, + {"type": "text", "text": "10 degrees"}, + ], + }, + { + "role": "tool", + "tool_call_id": "tool_call_2", + "content": [ + { + "type": "text", + "text": json.dumps({"id": "test_res_id", "degrees": 25}), + } + ], + }, + { + "role": "assistant", + "content": [ + { + "type": "thinking", + "thinking": ( + "The user asked about London, Toronto, and Tokyo. " + "I already have results for London and Toronto from " + "the tool calls, and Tokyo is covered by the " + "provided document." + ), + }, + { + "type": "text", + "text": ( + "The weather in London is 25 degrees. The weather " + "in Toronto is 10 degrees. The weather in Tokyo is " + "27 degrees." + ), + }, + ], + "citations": [ + { + "start": 25, + "end": 36, + "text": "25 degrees.", + "type": "TEXT_CONTENT", + "content_index": 1, + "sources": [ + { + "type": "tool", + "id": "test_res_id", + "tool_output": {"degrees": "25"}, + } + ], + }, + { + "start": 63, + "end": 74, + "text": "10 degrees.", + "type": "TEXT_CONTENT", + "content_index": 1, + "sources": [ + { + "type": "tool", + "id": "tool_call_1:1", + "tool_output": {"content": "10 degrees"}, + } + ], + }, + { + "start": 99, + "end": 110, + "text": "27 degrees.", + "type": "TEXT_CONTENT", + "content_index": 1, + "sources": [ + { + "type": "document", + "id": "doc:0", + "document": { + "id": "doc:0", + "title": "Weather in Tokyo", + "snippet": "The weather in tokyo is 27 degrees", + }, + } + ], + }, + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": ( + "Which city is warmest? Please cite your result as " + "described in the Grounding section" + ), + } + ], + }, + ], + "max_tokens": REASONING_BUDGET, + "temperature": 0.0, + "seed": 0, + } + request_id = ctx.next_request_id("cites-in-history") + _, resp = _post_json(ctx, body=body, request_id=request_id) + assert isinstance(resp, dict) + message = resp.get("message") or {} + content = message.get("content") or [] + citations = message.get("citations") or [] + text = _text_from_content_blocks(content) + _expect(content, f"empty content: {message}") + # "London" (25 degrees) should be identified as the warmest city. + _expect( + "london" in text.lower() or "25" in text, + f"expected the warmest-city answer to reference London/25 degrees: " + f"{text!r}", + ) + if not citations: + ctx.record(TestResult(name, True, detail=_no_grounding_warning(content))) + return + _assert_citation_shape(citations) + ctx.record( + TestResult( + name, + True, + detail=f"answer={text!r} citations={len(citations)}", + ) + ) + except Exception as e: + ctx.record(TestResult(name, False, detail=str(e))) + + +# ---------------------------------------------------------------------- +# Entrypoint +# ---------------------------------------------------------------------- + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--base-url", + default="http://127.0.0.1:8000", + help="vLLM server base URL (default: %(default)s)", + ) + parser.add_argument( + "--model", + default=DEFAULT_MODEL, + help=( + "Model id as registered with the server (matches /v1/models). " + "Also used to launch the server when auto-starting it. " + "(default: %(default)s)" + ), + ) + parser.add_argument( + "--reasoning-model", + dest="is_reasoning_model", + action=argparse.BooleanOptionalAction, + default=True, + help=( + "Whether the server was (or should be) started with " + "--cohere-is-reasoning-model (default true). Controls whether " + "the thinking+citations cases run, and whether an auto-started " + "server gets --no-cohere-is-reasoning-model." + ), + ) + parser.add_argument( + "--timeout", + type=float, + default=120.0, + help="HTTP request timeout, seconds (default: %(default)s)", + ) + parser.add_argument( + "--auto-start-server", + action=argparse.BooleanOptionalAction, + default=True, + help=( + "Launch 'vllm serve --model ' ourselves if --base-url " + "isn't already serving a healthy /health response (default true)." + ), + ) + parser.add_argument( + "--keep-server", + action=argparse.BooleanOptionalAction, + default=True, + help=( + "Leave an auto-started server running after the tests finish, " + "so re-runs skip the model load (default true). Only applies " + "when this script started the server itself." + ), + ) + parser.add_argument( + "--startup-timeout", + type=float, + default=1800.0, + help=( + "Seconds to wait for an auto-started server to become healthy " + "(default: %(default)s). Large quantized checkpoints can take " + "a long time to load." + ), + ) + parser.add_argument( + "--extra-server-arg", + dest="extra_server_args", + action="append", + default=[], + help=( + "Extra 'vllm serve' argument, forwarded verbatim when this " + "script starts the server itself. Repeatable, e.g. " + "--extra-server-arg=--tensor-parallel-size=8" + ), + ) + parser.add_argument( + "--verbose", + "-v", + action="store_true", + help="Print full request/response payloads.", + ) + args = parser.parse_args() + + managed_server: ManagedServer | None = None + if args.auto_start_server and not _server_is_healthy(args.base_url): + parsed = urlparse(args.base_url) + host = parsed.hostname or "127.0.0.1" + port = parsed.port or 8000 + managed_server = _start_vllm_server( + model=args.model, + host=host, + port=port, + is_reasoning_model=args.is_reasoning_model, + extra_args=args.extra_server_args, + ) + try: + _wait_for_server( + base_url=args.base_url, + server=managed_server, + timeout=args.startup_timeout, + ) + except Exception: + managed_server.stop() + raise + elif not _server_is_healthy(args.base_url): + print( + f"WARNING: {args.base_url} does not look healthy and " + f"--no-auto-start-server was passed; tests will likely fail." + ) + + try: + ctx = TestContext( + base_url=args.base_url, + model=args.model, + is_reasoning_model=args.is_reasoning_model, + verbose=args.verbose, + timeout=args.timeout, + ) + + print( + f"Driving Cohere v2 documents/citations tests at " + f"{ctx.base_url}/cohere/v2/chat" + ) + print(f" model={ctx.model}") + print(f" is_reasoning_model={ctx.is_reasoning_model}") + print() + + test_documents_basic_non_stream(ctx) + test_documents_basic_streaming(ctx) + test_citation_options_fast(ctx) + test_citation_options_off(ctx) + test_mixed_document_shapes(ctx) + test_grounding_against_tool_result_in_history(ctx) + test_thinking_with_citations_non_stream(ctx) + test_thinking_with_citations_streaming(ctx) + test_citations_in_conversation_history(ctx) + + print() + print("=" * 72) + passed = sum(1 for r in ctx.results if r.passed and not r.skipped) + failed = sum(1 for r in ctx.results if not r.passed and not r.skipped) + skipped = sum(1 for r in ctx.results if r.skipped) + print(f"Summary: {passed} passed, {failed} failed, {skipped} skipped") + if failed: + print("\nFailures:") + for r in ctx.results: + if not r.passed and not r.skipped: + print(f" - {r.name}") + if r.detail: + for line in r.detail.splitlines(): + print(f" {line}") + return 0 if failed == 0 else 1 + finally: + if managed_server is not None: + if args.keep_server: + print( + f"\nLeaving auto-started server running " + f"(pid={managed_server.proc.pid}, " + f"log={managed_server.log_path}).\n" + f"Stop it with: kill {managed_server.proc.pid}" + ) + else: + print( + f"\nStopping auto-started server " + f"(pid={managed_server.proc.pid})..." + ) + managed_server.stop() + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/test_cohere_output.py b/scripts/test_cohere_output.py new file mode 100644 index 000000000000..a0dd2301bf85 --- /dev/null +++ b/scripts/test_cohere_output.py @@ -0,0 +1,897 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Test the Cohere v2 output conversion paths. + +Drives ``CohereServingChatV2._chat_completion_to_v2`` (non-streaming) and +``CohereServingChatV2._chat_completion_stream_to_v2`` (streaming) with +synthetic upstream chat-completion responses, including reasoning, +tool calls, and citations. No engine boot, no model download. + +Run with: + + .venv/bin/python scripts/test_cohere_output.py + .venv/bin/python scripts/test_cohere_output.py --verbose +""" + +from __future__ import annotations + +import argparse +import asyncio +import json +import sys +from collections.abc import AsyncIterator + +from pydantic import ValidationError + +from vllm.entrypoints.cohere.cohere_chat_message import ( + Citation, + CitationSource, + CohereChatMessage, + CohereDeltaMessage, +) +from vllm.entrypoints.cohere.protocol import CohereChatV2Request +from vllm.entrypoints.cohere.serving import CohereServingChatV2 +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionResponse, + ChatCompletionResponseChoice, + ChatCompletionResponseStreamChoice, + ChatCompletionStreamResponse, + ChatMessage, +) +from vllm.entrypoints.openai.engine.protocol import ( + DeltaFunctionCall, + DeltaToolCall, + PromptTokenUsageInfo, + UsageInfo, +) + + +def _make_handler( + is_reasoning_model: bool = True, +) -> CohereServingChatV2: + """Build a converter without going through ``__init__``. + + The non-streaming and streaming converters only touch a handful of + instance attrs (currently just ``_is_reasoning_model``), which we set + here so we don't need to spin up the full handler. + """ + handler = CohereServingChatV2.__new__(CohereServingChatV2) + handler._is_reasoning_model = is_reasoning_model + return handler + + +def _make_request() -> CohereChatV2Request: + return CohereChatV2Request.model_validate( + { + "model": "test-model", + "messages": [ + {"role": "user", "content": "Who wrote Hamlet, and when?"}, + ], + "documents": [ + { + "id": "doc_0", + "data": { + "title": "Wikipedia: Hamlet", + "text": ( + "Hamlet was written by William Shakespeare around 1600." + ), + }, + }, + ], + "citation_options": {"mode": "ACCURATE"}, + } + ) + + +# ---------------------------------------------------------------------- +# Non-streaming +# ---------------------------------------------------------------------- + + +def test_non_streaming(verbose: bool) -> None: + print("=" * 72) + print("non-streaming: ChatCompletionResponse -> CohereChatV2Response") + print("=" * 72) + + handler = _make_handler() + request = _make_request() + + citation = Citation( + start=22, + end=41, + text="William Shakespeare", + sources=[ + CitationSource( + type="document", + id="doc_0", + document={ + "title": "Wikipedia: Hamlet", + "text": "Hamlet was written by William Shakespeare around 1600.", + }, + ) + ], + content_index=1, + type="TEXT_CONTENT", + ) + + msg = CohereChatMessage( + role="assistant", + content="Hamlet was written by William Shakespeare around 1600.", + reasoning="The user asked about Hamlet's authorship and date.", + citations=[citation], + ) + + response = ChatCompletionResponse( + id="chatcmpl-abc123", + model="test-model", + choices=[ + ChatCompletionResponseChoice(index=0, message=msg, finish_reason="stop") + ], + usage=UsageInfo(prompt_tokens=12, completion_tokens=20, total_tokens=32), + ) + + v2 = handler._chat_completion_to_v2(response, request) + payload = json.loads(v2.model_dump_json(exclude_none=True)) + + if verbose: + print(json.dumps(payload, indent=2)) + print() + + # ------------------------------------------------------------------ + # Assertions + # ------------------------------------------------------------------ + assert payload["id"] == "chatcmpl-abc123", payload["id"] + assert payload["finish_reason"] == "COMPLETE", payload["finish_reason"] + + message = payload["message"] + assert message["role"] == "assistant" + + block_types = [b["type"] for b in message["content"]] + assert block_types == ["thinking", "text"], block_types + assert ( + message["content"][0]["thinking"] + == "The user asked about Hamlet's authorship and date." + ) + assert ( + message["content"][1]["text"] + == "Hamlet was written by William Shakespeare around 1600." + ) + + cits = message["citations"] + assert len(cits) == 1 + c = cits[0] + assert c["text"] == "William Shakespeare" + assert c["start"] == 22 and c["end"] == 41 + assert c["sources"][0]["id"] == "doc_0" + assert c["sources"][0]["type"] == "document" + + usage = payload["usage"] + assert usage["billed_units"]["input_tokens"] == 12 + assert usage["billed_units"]["output_tokens"] == 20 + assert usage["tokens"]["input_tokens"] == 12 + assert usage["tokens"]["output_tokens"] == 20 + + print("OK: thinking + text content blocks") + print("OK: 1 citation (William Shakespeare -> doc_0)") + print("OK: finish_reason mapped 'stop' -> 'COMPLETE'") + print("OK: usage billed_units / tokens populated") + print() + + +def _build_tool_call_response() -> ChatCompletionResponse: + """Build a tool-call response with reasoning attached, for the two + parallel test cases below. + """ + from vllm.entrypoints.openai.engine.protocol import FunctionCall, ToolCall + + msg = ChatMessage( + role="assistant", + content=None, + reasoning="I should look up Hamlet's authorship in the knowledge base.", + tool_calls=[ + ToolCall( + id="call_1", + function=FunctionCall( + name="lookup", + arguments='{"query": "Hamlet authorship"}', + ), + ) + ], + ) + + return ChatCompletionResponse( + id="chatcmpl-tool-1", + model="test-model", + choices=[ + ChatCompletionResponseChoice( + index=0, message=msg, finish_reason="tool_calls" + ) + ], + usage=UsageInfo(prompt_tokens=10, completion_tokens=15, total_tokens=25), + ) + + +def test_non_streaming_with_tool_call_reasoning_default(verbose: bool) -> None: + """Default (reasoning-model) behavior: reasoning is kept as a thinking + content block alongside the tool calls; ``tool_plan`` is never set. + """ + print("=" * 72) + print("non-streaming: tool call (default, reasoning -> thinking block)") + print("=" * 72) + + handler = _make_handler(is_reasoning_model=True) + request = _make_request() + response = _build_tool_call_response() + + v2 = handler._chat_completion_to_v2(response, request) + payload = json.loads(v2.model_dump_json(exclude_none=True)) + + if verbose: + print(json.dumps(payload, indent=2)) + print() + + assert payload["finish_reason"] == "TOOL_CALL" + message = payload["message"] + assert "tool_plan" not in message or message["tool_plan"] is None + assert message["content"] is not None + block_types = [b["type"] for b in message["content"]] + assert block_types == ["thinking"], block_types + assert message["content"][0]["thinking"] == ( + "I should look up Hamlet's authorship in the knowledge base." + ) + assert len(message["tool_calls"]) == 1 + + print("OK: thinking block emitted alongside tool_calls") + print("OK: tool_plan is not set (reasoning model assumption)") + print() + + +def test_non_streaming_with_tool_call_tool_plan_flag(verbose: bool) -> None: + """Flag-enabled (non-reasoning-model) behavior: reasoning is surfaced + as ``tool_plan`` and no thinking content block is emitted. + """ + print("=" * 72) + print("non-streaming: tool call (non-reasoning, reasoning -> tool_plan)") + print("=" * 72) + + handler = _make_handler(is_reasoning_model=False) + request = _make_request() + response = _build_tool_call_response() + + v2 = handler._chat_completion_to_v2(response, request) + payload = json.loads(v2.model_dump_json(exclude_none=True)) + + if verbose: + print(json.dumps(payload, indent=2)) + print() + + assert payload["finish_reason"] == "TOOL_CALL" + message = payload["message"] + assert message.get("content") is None or message["content"] == [] + assert message["tool_plan"] == ( + "I should look up Hamlet's authorship in the knowledge base." + ) + assert len(message["tool_calls"]) == 1 + tc = message["tool_calls"][0] + assert tc["id"] == "call_1" + assert tc["function"]["name"] == "lookup" + assert json.loads(tc["function"]["arguments"]) == {"query": "Hamlet authorship"} + + print("OK: reasoning surfaced as tool_plan, no thinking block emitted") + print("OK: tool_calls preserved with id + function.name + arguments") + print() + + +# ---------------------------------------------------------------------- +# Streaming +# ---------------------------------------------------------------------- + + +def _wrap(chunk: ChatCompletionStreamResponse) -> str: + """Format a chunk like it would appear on the upstream OpenAI SSE stream.""" + return f"data: {chunk.model_dump_json()}\n\n" + + +def _stream_chunk( + delta_kwargs: dict | None = None, + finish_reason: str | None = None, + chunk_id: str = "chatcmpl-stream-1", +) -> ChatCompletionStreamResponse: + return ChatCompletionStreamResponse( + id=chunk_id, + model="test-model", + choices=[ + ChatCompletionResponseStreamChoice( + index=0, + delta=CohereDeltaMessage(**(delta_kwargs or {})), + finish_reason=finish_reason, + ) + ], + ) + + +async def _fake_upstream() -> AsyncIterator[str]: + """Simulate vLLM's upstream chat-completion SSE stream. + + Order chosen to exercise every block-transition path: + + role -> reasoning -> text -> citation -> tool_call -> usage-only -> [DONE] + """ + # Initial role-only delta. + yield _wrap(_stream_chunk({"role": "assistant"})) + + # Reasoning -> emits a 'thinking' content block. + yield _wrap(_stream_chunk({"reasoning": "Thinking about Hamlet..."})) + yield _wrap(_stream_chunk({"reasoning": " It's Shakespeare."})) + + # Visible text -> closes thinking block, opens a text block. + yield _wrap(_stream_chunk({"content": "Hamlet was written by "})) + yield _wrap(_stream_chunk({"content": "William Shakespeare."})) + + # Citation grounding the text we just emitted. + citation = Citation( + start=22, + end=41, + text="William Shakespeare", + sources=[CitationSource(type="document", id="doc_0")], + content_index=1, + type="TEXT_CONTENT", + ) + yield _wrap(_stream_chunk({"citations": [citation]})) + + # Tool call -> closes text block, opens tool_call block. + yield _wrap( + _stream_chunk( + { + "tool_calls": [ + DeltaToolCall( + id="call_0", + type="function", + index=0, + function=DeltaFunctionCall(name="lookup", arguments=""), + ) + ] + } + ) + ) + # Last delta-bearing chunk also carries the OpenAI ``finish_reason``, + # mirroring real upstream behavior. The very last chunk is usage-only. + yield _wrap( + _stream_chunk( + delta_kwargs={ + "tool_calls": [ + DeltaToolCall( + index=0, + function=DeltaFunctionCall(arguments='{"q":"Hamlet"}'), + ) + ] + }, + finish_reason="tool_calls", + ) + ) + + # Final, choices-empty chunk carries usage only. + final_chunk = ChatCompletionStreamResponse( + id="chatcmpl-stream-1", + model="test-model", + choices=[], + usage=UsageInfo(prompt_tokens=10, completion_tokens=15, total_tokens=25), + ) + yield _wrap(final_chunk) + yield "data: [DONE]\n\n" + + +async def test_streaming(verbose: bool) -> None: + print("=" * 72) + print("streaming: default (reasoning -> thinking content blocks)") + print("=" * 72) + + handler = _make_handler(is_reasoning_model=True) + request = _make_request() + + events: list[dict] = [] + frames: list[str] = [] + async for sse in handler._chat_completion_stream_to_v2(_fake_upstream(), request): + assert sse.startswith("data: "), repr(sse) + frames.append(sse) + body = sse[len("data: ") :].strip() + if not body or body == "[DONE]": + continue + events.append(json.loads(body)) + + # Cohere v2 streams must terminate with ``data: [DONE]\n\n``. + assert frames[-1] == "data: [DONE]\n\n", repr(frames[-1]) + + if verbose: + for ev in events: + print(json.dumps(ev, indent=2)) + print() + else: + for ev in events: + t = ev.get("type", "?") + extra = "" + if t in ("content-start", "content-end"): + extra = f" index={ev.get('index')}" + elif t == "content-delta": + content = ev.get("delta", {}).get("message", {}).get("content", {}) + extra = f" index={ev.get('index')} delta={content!r}" + elif t in ( + "tool-call-start", + "tool-call-delta", + "tool-call-end", + "citation-start", + "citation-end", + ): + extra = f" index={ev.get('index')}" + elif t == "message-end": + extra = f" finish_reason={ev.get('delta', {}).get('finish_reason')}" + print(f" {t}{extra}") + + # ------------------------------------------------------------------ + # Assertions + # ------------------------------------------------------------------ + types = [ev["type"] for ev in events] + + # Lifecycle bookends. + assert types[0] == "message-start", types[0] + assert types[-1] == "message-end", types[-1] + + # Each kind of event we care about appears at least once. + for required in ( + "content-start", + "content-delta", + "content-end", + "tool-call-start", + "tool-call-delta", + "tool-call-end", + "citation-start", + "citation-end", + ): + assert required in types, f"missing event type: {required}" + + # Block-ordering invariants: + # thinking opens before any text is opened, text opens before tool_call. + def _content_start_kind(ev: dict) -> str | None: + c = ev.get("delta", {}).get("message", {}).get("content") + return c.get("type") if isinstance(c, dict) else None + + thinking_start = next( + i + for i, e in enumerate(events) + if e["type"] == "content-start" and _content_start_kind(e) == "thinking" + ) + text_start = next( + i + for i, e in enumerate(events) + if e["type"] == "content-start" and _content_start_kind(e) == "text" + ) + tool_start = types.index("tool-call-start") + assert thinking_start < text_start < tool_start, ( + thinking_start, + text_start, + tool_start, + ) + + # The text block has to be closed before the tool block opens. + text_block_index = events[text_start].get("index") + text_end_idx = next( + i + for i, e in enumerate(events) + if e["type"] == "content-end" and e.get("index") == text_block_index + ) + assert text_end_idx < tool_start + + # citation-start / citation-end must come as a pair, in order. + cit_starts = [i for i, t in enumerate(types) if t == "citation-start"] + cit_ends = [i for i, t in enumerate(types) if t == "citation-end"] + assert len(cit_starts) == len(cit_ends) == 1 + assert cit_starts[0] < cit_ends[0] + cit_payload = events[cit_starts[0]]["delta"]["message"]["citations"] + assert cit_payload["text"] == "William Shakespeare" + assert cit_payload["sources"][0]["id"] == "doc_0" + + # tool-call-start carries id + function name; subsequent deltas carry args. + tool_start_ev = events[tool_start] + tc = tool_start_ev["delta"]["message"]["tool_calls"] + assert tc["id"] == "call_0" + assert tc["function"]["name"] == "lookup" + + tool_delta_ev = next(e for e in events if e["type"] == "tool-call-delta") + args = tool_delta_ev["delta"]["message"]["tool_calls"]["function"]["arguments"] + assert args == '{"q":"Hamlet"}' + + # message-end: finish_reason + usage. + end = events[-1] + assert end["delta"]["finish_reason"] == "TOOL_CALL" + usage = end["delta"]["usage"] + assert usage["billed_units"]["input_tokens"] == 10 + assert usage["billed_units"]["output_tokens"] == 15 + assert usage["tokens"]["input_tokens"] == 10 + assert usage["tokens"]["output_tokens"] == 15 + + print() + print("OK: full event lifecycle (message-start ... message-end)") + print("OK: thinking -> text -> tool_call block ordering") + print("OK: citation-start / citation-end pair carries text + source") + print("OK: tool-call-start carries id+name, deltas carry arg fragments") + print("OK: message-end maps finish_reason and usage") + print() + + +async def test_streaming_tool_plan(verbose: bool) -> None: + """With ``is_reasoning_model=False`` (older non-reasoning Command + model), reasoning chunks are emitted as ``tool-plan-delta`` events + instead of opening a thinking content block. + """ + print("=" * 72) + print("streaming: non-reasoning (reasoning -> tool-plan-delta)") + print("=" * 72) + + handler = _make_handler(is_reasoning_model=False) + request = _make_request() + + events: list[dict] = [] + frames: list[str] = [] + async for sse in handler._chat_completion_stream_to_v2(_fake_upstream(), request): + frames.append(sse) + body = sse[len("data: ") :].strip() + if body and body != "[DONE]": + events.append(json.loads(body)) + + assert frames[-1] == "data: [DONE]\n\n", repr(frames[-1]) + + if verbose: + for ev in events: + print(json.dumps(ev, indent=2)) + print() + else: + for ev in events: + print(f" {ev.get('type', '?')}") + + types = [ev["type"] for ev in events] + + # No thinking content block ever opens. + for ev in events: + if ev["type"] == "content-start": + content = ev.get("delta", {}).get("message", {}).get("content") + kind = content.get("type") if isinstance(content, dict) else None + assert kind != "thinking", ( + "thinking content-start should not be emitted when " + "is_reasoning_model=False" + ) + + # ``tool-plan-delta`` events are emitted, with the reasoning text payload. + plan_events = [e for e in events if e["type"] == "tool-plan-delta"] + assert len(plan_events) >= 1 + plan_text = "".join(e["delta"]["message"]["tool_plan"] for e in plan_events) + assert "Thinking about Hamlet" in plan_text + assert "Shakespeare" in plan_text + + # The visible text and tool call still flow correctly. + assert "content-start" in types # for the text block + assert "tool-call-start" in types + assert types[-1] == "message-end" + + print() + print("OK: reasoning chunks emitted as tool-plan-delta events") + print("OK: no thinking content-start event in the stream") + print("OK: text + tool call events flow normally") + print() + + +# ---------------------------------------------------------------------- +# Cached tokens (#11) + fallback message-end (#6) +# ---------------------------------------------------------------------- + + +def test_non_streaming_cached_tokens(verbose: bool) -> None: + """``UsageInfo.prompt_tokens_details.cached_tokens`` must flow through + into ``CohereUsage.cached_tokens`` on the v2 response. + """ + print("=" * 72) + print("non-streaming: cached_tokens plumbed through usage") + print("=" * 72) + + handler = _make_handler() + request = _make_request() + + msg = ChatMessage(role="assistant", content="hi") + response = ChatCompletionResponse( + id="chatcmpl-cache", + model="test-model", + choices=[ + ChatCompletionResponseChoice(index=0, message=msg, finish_reason="stop") + ], + usage=UsageInfo( + prompt_tokens=42, + completion_tokens=8, + total_tokens=50, + prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=37), + ), + ) + + v2 = handler._chat_completion_to_v2(response, request) + payload = json.loads(v2.model_dump_json(exclude_none=True)) + + if verbose: + print(json.dumps(payload, indent=2)) + print() + + usage = payload["usage"] + assert usage["cached_tokens"] == 37, usage + assert usage["billed_units"]["input_tokens"] == 42 + + print("OK: usage.cached_tokens == 37") + print() + + +async def _fake_upstream_with_cached_tokens() -> AsyncIterator[str]: + """Minimal stream that includes cached_tokens in the usage chunk.""" + yield _wrap(_stream_chunk({"role": "assistant"})) + yield _wrap(_stream_chunk({"content": "hi"}, finish_reason="stop")) + yield _wrap( + ChatCompletionStreamResponse( + id="chatcmpl-cache-stream", + model="test-model", + choices=[], + usage=UsageInfo( + prompt_tokens=42, + completion_tokens=8, + total_tokens=50, + prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=37), + ), + ) + ) + yield "data: [DONE]\n\n" + + +async def test_streaming_cached_tokens(verbose: bool) -> None: + print("=" * 72) + print("streaming: cached_tokens propagated onto message-end") + print("=" * 72) + + handler = _make_handler() + request = _make_request() + + events: list[dict] = [] + frames: list[str] = [] + async for sse in handler._chat_completion_stream_to_v2( + _fake_upstream_with_cached_tokens(), request + ): + frames.append(sse) + body = sse[len("data: ") :].strip() + if body and body != "[DONE]": + events.append(json.loads(body)) + + if verbose: + for ev in events: + print(json.dumps(ev, indent=2)) + + end = events[-1] + assert end["type"] == "message-end", end + assert end["delta"]["usage"]["cached_tokens"] == 37, end["delta"]["usage"] + assert frames[-1] == "data: [DONE]\n\n", repr(frames[-1]) + + print("OK: message-end delta.usage.cached_tokens == 37") + print("OK: stream terminates with data: [DONE]") + print() + + +async def _fake_upstream_no_usage_chunk() -> AsyncIterator[str]: + """Upstream that closes with ``[DONE]`` but never sends the usage-only + final chunk (some inference backends do this). + """ + yield _wrap(_stream_chunk({"role": "assistant"})) + yield _wrap(_stream_chunk({"content": "hello"})) + yield _wrap(_stream_chunk({"content": " world"}, finish_reason="stop")) + yield "data: [DONE]\n\n" + + +async def _fake_upstream_no_done() -> AsyncIterator[str]: + """Upstream that just exhausts the iterator without ``[DONE]`` and + without a usage-only chunk (e.g. cancellation, abrupt shutdown). + """ + yield _wrap(_stream_chunk({"role": "assistant"})) + yield _wrap(_stream_chunk({"content": "hello"}, finish_reason="stop")) + + +async def test_streaming_emits_message_end_without_usage_chunk( + verbose: bool, +) -> None: + """Even when upstream skips the usage-only final chunk, the v2 stream + must terminate cleanly with ``message-end``. + """ + print("=" * 72) + print("streaming: fallback message-end (no usage chunk, [DONE] only)") + print("=" * 72) + + handler = _make_handler() + request = _make_request() + + events: list[dict] = [] + frames: list[str] = [] + async for sse in handler._chat_completion_stream_to_v2( + _fake_upstream_no_usage_chunk(), request + ): + frames.append(sse) + body = sse[len("data: ") :].strip() + if body and body != "[DONE]": + events.append(json.loads(body)) + + if verbose: + for ev in events: + print(json.dumps(ev, indent=2)) + + types = [ev["type"] for ev in events] + assert types[0] == "message-start", types[0] + assert types[-1] == "message-end", types[-1] + # Open content block must be closed before message-end. + assert types.count("content-start") == types.count("content-end") + end = events[-1] + assert end["delta"]["finish_reason"] == "COMPLETE", end + # No usage chunk arrived from upstream, so the synthetic message-end + # omits the ``usage`` field. + assert "usage" not in end["delta"], end["delta"] + # Still must terminate with [DONE] even in the fallback path. + assert frames[-1] == "data: [DONE]\n\n", repr(frames[-1]) + + print("OK: message-end emitted despite missing usage chunk") + print("OK: open content blocks closed before message-end") + print("OK: finish_reason mapped from last delta-bearing chunk") + print("OK: stream terminates with data: [DONE]") + print() + + +async def test_streaming_emits_message_end_without_done(verbose: bool) -> None: + """Upstream that exhausts without ``[DONE]`` and without a usage chunk + must still produce a closing ``message-end`` event. + """ + print("=" * 72) + print("streaming: fallback message-end (iterator exhausts, no [DONE])") + print("=" * 72) + + handler = _make_handler() + request = _make_request() + + events: list[dict] = [] + frames: list[str] = [] + async for sse in handler._chat_completion_stream_to_v2( + _fake_upstream_no_done(), request + ): + frames.append(sse) + body = sse[len("data: ") :].strip() + if body and body != "[DONE]": + events.append(json.loads(body)) + + if verbose: + for ev in events: + print(json.dumps(ev, indent=2)) + + types = [ev["type"] for ev in events] + assert types[-1] == "message-end", types + assert events[-1]["delta"]["finish_reason"] == "COMPLETE" + # Even when upstream never sent [DONE], we synthesize it on our side. + assert frames[-1] == "data: [DONE]\n\n", repr(frames[-1]) + + print("OK: message-end emitted on plain iterator exhaustion") + print("OK: stream terminates with data: [DONE]") + print() + + +# ---------------------------------------------------------------------- +# Request validation (#12 id construction, #14 max_tokens=0 acceptance) +# ---------------------------------------------------------------------- + + +def test_request_accepts_max_tokens_zero(verbose: bool) -> None: + """Cohere's API treats ``max_tokens=0`` as valid; the v2 request + validator must accept it (only true negatives are rejected). + """ + print("=" * 72) + print("request: max_tokens=0 accepted, negative rejected") + print("=" * 72) + + req = CohereChatV2Request.model_validate( + { + "model": "test-model", + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 0, + } + ) + assert req.max_tokens == 0, req.max_tokens + + # None still allowed. + req2 = CohereChatV2Request.model_validate( + {"model": "test-model", "messages": [{"role": "user", "content": "hi"}]} + ) + assert req2.max_tokens is None + + # Negative still rejected. + try: + CohereChatV2Request.model_validate( + { + "model": "test-model", + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": -1, + } + ) + except ValidationError as exc: + assert "non-negative" in str(exc), str(exc) + else: + raise AssertionError("max_tokens=-1 should have been rejected") + + print("OK: max_tokens=0 accepted") + print("OK: max_tokens=None accepted") + print("OK: max_tokens=-1 rejected") + print() + + +def test_response_id_synthesized_when_upstream_missing(verbose: bool) -> None: + """When the upstream ``ChatCompletionResponse.id`` is empty, the v2 + converter must synthesize a non-empty id at the call site. + """ + print("=" * 72) + print("response: id synthesized when upstream id is empty") + print("=" * 72) + + handler = _make_handler() + request = _make_request() + + msg = ChatMessage(role="assistant", content="hi") + response = ChatCompletionResponse( + id="", # upstream forgot to set it + model="test-model", + choices=[ + ChatCompletionResponseChoice(index=0, message=msg, finish_reason="stop") + ], + usage=UsageInfo(prompt_tokens=1, completion_tokens=1, total_tokens=2), + ) + + v2 = handler._chat_completion_to_v2(response, request) + assert v2.id, f"expected non-empty id, got {v2.id!r}" + assert v2.id.startswith("chat_"), v2.id + + # And the upstream id is preserved when present. + response.id = "chatcmpl-real-id" + v2b = handler._chat_completion_to_v2(response, request) + assert v2b.id == "chatcmpl-real-id", v2b.id + + print(f"OK: empty upstream id -> synthesized {v2.id!r}") + print("OK: non-empty upstream id preserved") + print() + + +# ---------------------------------------------------------------------- +# Entrypoint +# ---------------------------------------------------------------------- + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--verbose", + "-v", + action="store_true", + help="Dump full JSON for each event", + ) + args = parser.parse_args() + + test_request_accepts_max_tokens_zero(args.verbose) + test_response_id_synthesized_when_upstream_missing(args.verbose) + test_non_streaming(args.verbose) + test_non_streaming_with_tool_call_reasoning_default(args.verbose) + test_non_streaming_with_tool_call_tool_plan_flag(args.verbose) + test_non_streaming_cached_tokens(args.verbose) + asyncio.run(test_streaming(args.verbose)) + asyncio.run(test_streaming_tool_plan(args.verbose)) + asyncio.run(test_streaming_cached_tokens(args.verbose)) + asyncio.run(test_streaming_emits_message_end_without_usage_chunk(args.verbose)) + asyncio.run(test_streaming_emits_message_end_without_done(args.verbose)) + + print("ALL TESTS PASSED") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/test_cohere_v2_e2e.py b/scripts/test_cohere_v2_e2e.py new file mode 100644 index 000000000000..9e3fe83140db --- /dev/null +++ b/scripts/test_cohere_v2_e2e.py @@ -0,0 +1,1187 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""End-to-end smoke test for the Cohere Chat v2 API (``POST /cohere/v2/chat``). + +Drives a live vLLM server with a real Cohere Command-family model and +exercises the v2 wire format end-to-end: non-streaming, streaming, +documents/citations, tools, reasoning, role aliasing, and error paths. + +Prerequisites on the GPU host +----------------------------- + +1. Install the optional Cohere SDKs: + + uv pip install cohere cohere-melody + +2. Start a vLLM server with the Cohere renderer / tokenizer wired up. + ``VLLM_ENABLE_COHERE_API=1`` is required to expose ``/cohere/v2/chat``; + the endpoint is opt-in and stays hidden otherwise: + + VLLM_ENABLE_COHERE_API=1 vllm serve \\ + --tokenizer-mode cohere \\ + --enable-auto-tool-choice \\ + --tool-call-parser cohere2 \\ + --reasoning-parser cohere2 \\ + --port 8000 + + For non-reasoning Command models (cmd3, older Command R), append + ``--no-cohere-is-reasoning-model`` so the renderer surfaces reasoning + as ``tool_plan`` instead of as a ``thinking`` content block. + +3. Run this script: + + python scripts/test_cohere_v2_e2e.py \\ + --base-url http://127.0.0.1:8000 \\ + --model + +Optional flags +-------------- + +* ``--reasoning-model / --no-reasoning-model`` -- whether the server was + launched with reasoning enabled (``--cohere-is-reasoning-model``). + Controls whether reasoning is expected as a ``thinking`` content block + or as a ``tool-plan-delta`` event. +* ``--skip-tools`` / ``--skip-citations`` / ``--skip-streaming`` -- + scope down the test surface when iterating locally. +* ``--verbose`` -- dump full response bodies / SSE frames. +""" + +from __future__ import annotations + +import argparse +import json +import sys +import time +from dataclasses import dataclass, field +from typing import Any + +import httpx + +# ---------------------------------------------------------------------- +# Helpers +# ---------------------------------------------------------------------- + +DONE_LINE = "data: [DONE]" + +# Reasoning Command models spend most of their decode budget inside +# ``thinking`` blocks before emitting any user-visible text. The basic +# correctness probes use a generous budget so the response actually +# reaches the text/tool-call payload we're trying to assert against; +# tighter limits are still applied per-test for the negative paths. +REASONING_BUDGET = 2048 + + +def _text_from_content_blocks(content: Any) -> str: + """Concatenate ``text`` blocks from a message ``content`` payload.""" + if not isinstance(content, list): + return "" + return "".join( + b.get("text", "") + for b in content + if isinstance(b, dict) and b.get("type") == "text" + ) + + +def _thinking_from_content_blocks(content: Any) -> str: + """Concatenate ``thinking`` blocks from a message ``content`` payload.""" + if not isinstance(content, list): + return "" + return "".join( + b.get("thinking", "") + for b in content + if isinstance(b, dict) and b.get("type") == "thinking" + ) + + +def _looks_like_4xx_envelope(parsed: Any) -> bool: + """True for any of the documented Cohere v2 4xx response shapes. + + We accept both: + + * Our native ``CohereError`` (``{"message": ..., "id": ...}``). + * vLLM's generic OpenAI-style envelope used by the request-body + validator at ``vllm.entrypoints.serve.utils.api_utils`` -- + ``{"error": {"message": ..., "type": ..., "code": ...}}``. + * FastAPI's default ``{"detail": [...]}`` for body-shape rejection. + """ + if not isinstance(parsed, dict): + return False + if "message" in parsed and "error" not in parsed: + return True + if isinstance(parsed.get("error"), dict) and "message" in parsed["error"]: + return True + return "detail" in parsed + + +@dataclass +class TestResult: + name: str + passed: bool + detail: str = "" + skipped: bool = False + + +@dataclass +class TestContext: + base_url: str + model: str + is_reasoning_model: bool + verbose: bool + timeout: float + request_id_counter: int = 0 + results: list[TestResult] = field(default_factory=list) + + def next_request_id(self, label: str) -> str: + self.request_id_counter += 1 + return f"e2e-{label}-{self.request_id_counter}-{int(time.time())}" + + def record(self, result: TestResult) -> None: + self.results.append(result) + prefix = "SKIP" if result.skipped else ("PASS" if result.passed else "FAIL") + print(f"[{prefix}] {result.name}") + if result.detail: + for line in result.detail.splitlines(): + print(f" {line}") + + +def _post_json( + ctx: TestContext, + *, + body: dict[str, Any], + request_id: str, + expect_status: int | None = 200, +) -> tuple[int, dict[str, Any] | str]: + """POST a non-streaming JSON request and return (status, parsed_body).""" + url = f"{ctx.base_url.rstrip('/')}/cohere/v2/chat" + headers = {"Content-Type": "application/json", "X-Request-Id": request_id} + with httpx.Client(timeout=ctx.timeout) as client: + resp = client.post(url, json=body, headers=headers) + if ctx.verbose: + print(f" -> POST {url} [status {resp.status_code}]") + print(f" request_id={request_id}") + print(f" body={json.dumps(body)[:500]}") + print(f" resp={resp.text[:1500]}") + parsed: dict[str, Any] | str + try: + parsed = resp.json() + except Exception: + parsed = resp.text + if expect_status is not None and resp.status_code != expect_status: + raise AssertionError( + f"expected status {expect_status}, got {resp.status_code}: {parsed!r}" + ) + return resp.status_code, parsed + + +def _stream_post( + ctx: TestContext, + *, + body: dict[str, Any], + request_id: str, +) -> tuple[list[dict[str, Any]], bool]: + """POST a streaming request and return (events, terminated_with_done).""" + body = {**body, "stream": True} + url = f"{ctx.base_url.rstrip('/')}/cohere/v2/chat" + headers = { + "Content-Type": "application/json", + "Accept": "text/event-stream", + "X-Request-Id": request_id, + } + events: list[dict[str, Any]] = [] + saw_done = False + with ( + httpx.Client(timeout=ctx.timeout) as client, + client.stream("POST", url, json=body, headers=headers) as resp, + ): + if resp.status_code != 200: + resp.read() + raise AssertionError( + f"stream expected 200, got {resp.status_code}: {resp.text[:1000]}" + ) + buffer = "" + for chunk in resp.iter_text(): + if not chunk: + continue + buffer += chunk + while "\n\n" in buffer: + frame, buffer = buffer.split("\n\n", 1) + for line in frame.splitlines(): + line = line.strip() + if not line.startswith("data:"): + continue + payload = line[len("data:") :].strip() + if payload == "[DONE]": + saw_done = True + continue + if not payload: + continue + try: + events.append(json.loads(payload)) + except json.JSONDecodeError as e: + raise AssertionError(f"bad SSE frame: {payload!r}: {e}") from e + if ctx.verbose: + print(f" -> stream {url} events={len(events)} done={saw_done}") + for ev in events: + print(f" {ev.get('type'):<20} {json.dumps(ev)[:200]}") + return events, saw_done + + +def _event_types(events: list[dict[str, Any]]) -> list[str]: + return [ev.get("type", "") for ev in events] + + +def _expect(condition: bool, message: str) -> None: + if not condition: + raise AssertionError(message) + + +# ---------------------------------------------------------------------- +# Test cases +# ---------------------------------------------------------------------- + + +def test_health(ctx: TestContext) -> None: + """Verify the server is up and the model is loaded.""" + name = "health: GET /health and /v1/models" + try: + with httpx.Client(timeout=ctx.timeout) as client: + health = client.get(f"{ctx.base_url.rstrip('/')}/health") + _expect( + health.status_code == 200, + f"/health returned {health.status_code}", + ) + models = client.get(f"{ctx.base_url.rstrip('/')}/v1/models") + _expect( + models.status_code == 200, + f"/v1/models returned {models.status_code}: {models.text[:500]}", + ) + ids = [m["id"] for m in models.json().get("data", [])] + _expect( + any(ctx.model == mid or mid.startswith(ctx.model) for mid in ids), + f"model {ctx.model!r} not found in /v1/models -> {ids}", + ) + ctx.record(TestResult(name, True, detail=f"served models: {ids}")) + except Exception as e: + ctx.record(TestResult(name, False, detail=str(e))) + + +def test_non_streaming_basic(ctx: TestContext) -> None: + """Plain system + user -> assistant turn. + + The wire-format contract we're testing is "the response is a valid + v2 envelope with at least one content block". For reasoning models + we'd ideally see a final ``text`` block, but a ``thinking``-only + response (truncated by ``MAX_TOKENS`` mid-reasoning) still satisfies + the contract -- treat that as a soft pass with a warning. + """ + name = "non-streaming: basic system+user prompt" + try: + body = { + "model": ctx.model, + "messages": [ + {"role": "system", "content": "You are a terse assistant."}, + { + "role": "user", + "content": "Reply with the single word: OK", + }, + ], + "max_tokens": REASONING_BUDGET, + "temperature": 0.0, + } + request_id = ctx.next_request_id("basic") + _, resp = _post_json(ctx, body=body, request_id=request_id) + assert isinstance(resp, dict), resp + _expect("id" in resp and resp["id"], f"missing id: {resp}") + _expect( + resp.get("finish_reason") in {"COMPLETE", "MAX_TOKENS"}, + f"unexpected finish_reason: {resp.get('finish_reason')}", + ) + message = resp.get("message") or {} + _expect( + message.get("role") == "assistant", + f"unexpected role: {message.get('role')}", + ) + content = message.get("content") or [] + _expect(content, f"empty content[]: {message}") + text = _text_from_content_blocks(content).strip() + thinking = _thinking_from_content_blocks(content).strip() + if text: + detail = ( + f"text={text!r} thinking_chars={len(thinking)} " + f"finish={resp.get('finish_reason')} usage={resp.get('usage')}" + ) + else: + detail = ( + f"WARN: only thinking content emitted " + f"({len(thinking)} chars); " + f"finish={resp.get('finish_reason')} usage={resp.get('usage')} " + f"-- bump --timeout or the prompt budget if this recurs" + ) + ctx.record(TestResult(name, True, detail=detail)) + except Exception as e: + ctx.record(TestResult(name, False, detail=str(e))) + + +def test_non_streaming_developer_role(ctx: TestContext) -> None: + """OpenAI ``developer`` role must alias to ``system`` (M5).""" + name = "non-streaming: developer role aliased to system" + try: + body = { + "model": ctx.model, + "messages": [ + { + "role": "developer", + "content": "You only ever respond with the word: HELLO.", + }, + {"role": "user", "content": "Say hi."}, + ], + "max_tokens": REASONING_BUDGET, + "temperature": 0.0, + } + request_id = ctx.next_request_id("dev-role") + _, resp = _post_json(ctx, body=body, request_id=request_id) + assert isinstance(resp, dict) + message = resp.get("message") or {} + content = message.get("content") or [] + _expect( + content, + f"developer role accepted but empty content[]: {message}", + ) + text = _text_from_content_blocks(content).strip() + thinking = _thinking_from_content_blocks(content).strip() + # M5 + the protocol-level normalizer: ``developer`` is rewritten + # to ``system`` before the SDK discriminator runs. The request + # must at minimum *not* 4xx and must produce content (text or + # thinking) -- before the fix the SDK rejected the role + # outright with a 400. + if text: + ctx.record(TestResult(name, True, detail=f"text={text!r}")) + else: + ctx.record( + TestResult( + name, + True, + detail=( + f"WARN: only thinking emitted ({len(thinking)} chars); " + f"role alias accepted (no 4xx)" + ), + ) + ) + except Exception as e: + ctx.record(TestResult(name, False, detail=str(e))) + + +def test_non_streaming_with_documents(ctx: TestContext) -> None: + """Documents in -> citations in the response message.""" + name = "non-streaming: grounded answer with citations" + try: + body = { + "model": ctx.model, + "messages": [ + { + "role": "user", + "content": ( + "Using the provided documents, answer who wrote " + "Hamlet and roughly when. Quote the specific " + "spans from the documents that support each " + "claim and ground them with citation tags." + ), + }, + ], + "documents": [ + { + "id": "doc_shakespeare", + "data": { + "title": "Wikipedia: Hamlet", + "text": ( + "Hamlet is a tragedy written by William " + "Shakespeare around 1600." + ), + }, + }, + { + "id": "doc_irrelevant", + "data": { + "title": "Wikipedia: Compilers", + "text": ( + "A compiler is a translator from one language to another." + ), + }, + }, + ], + "citation_options": {"mode": "ACCURATE"}, + "max_tokens": REASONING_BUDGET, + "temperature": 0.0, + } + request_id = ctx.next_request_id("documents") + _, resp = _post_json(ctx, body=body, request_id=request_id) + assert isinstance(resp, dict) + message = resp.get("message") or {} + content = message.get("content") or [] + citations = message.get("citations") or [] + _expect(content, f"empty content: {message}") + # Citations are best-effort on the model side; flag rather than + # fail if the model declined to ground. + if not citations: + # Surface the raw assistant text so we can tell whether the + # model emitted citation markers at all. `` None: + """Tools in -> a TOOL_CALL finish with structured tool_calls.""" + name = "non-streaming: tool call" + try: + body = { + "model": ctx.model, + "messages": [ + { + "role": "user", + "content": ( + "What's the weather in Tokyo right now? " + "Use the get_weather tool." + ), + }, + ], + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather for a city.", + "parameters": { + "type": "object", + "properties": { + "city": { + "type": "string", + "description": "City name.", + } + }, + "required": ["city"], + }, + }, + } + ], + "max_tokens": REASONING_BUDGET, + "temperature": 0.0, + } + request_id = ctx.next_request_id("tools") + _, resp = _post_json(ctx, body=body, request_id=request_id) + assert isinstance(resp, dict) + finish = resp.get("finish_reason") + message = resp.get("message") or {} + tool_calls = message.get("tool_calls") or [] + # If the model decided not to call the tool, surface a warning + # rather than fail. The wiring is what we're testing. + if not tool_calls: + ctx.record( + TestResult( + name, + True, + detail=( + f"WARN: model returned finish_reason={finish!r} with " + f"no tool_calls; cannot verify tool-call response " + f"shape." + ), + ) + ) + return + _expect( + finish == "TOOL_CALL", + f"expected finish_reason=TOOL_CALL, got {finish!r}", + ) + first = tool_calls[0] + _expect("id" in first, f"tool_call missing id: {first}") + fn = first.get("function") or {} + _expect( + fn.get("name") == "get_weather", + f"unexpected tool name: {fn}", + ) + args = fn.get("arguments") + _expect( + isinstance(args, str) and args.strip(), + f"tool arguments not a JSON string: {fn}", + ) + parsed_args = json.loads(args) + _expect( + isinstance(parsed_args, dict) and "city" in parsed_args, + f"tool arguments missing 'city': {parsed_args}", + ) + thinking_blocks = [ + b + for b in (message.get("content") or []) + if isinstance(b, dict) and b.get("type") == "thinking" + ] + tool_plan = message.get("tool_plan") + if ctx.is_reasoning_model: + detail_note = ( + f"thinking_blocks={len(thinking_blocks)} tool_plan={tool_plan!r}" + ) + else: + detail_note = f"tool_plan={tool_plan!r} (no thinking block expected)" + ctx.record( + TestResult( + name, + True, + detail=( + f"tool_calls=1 name={fn.get('name')} args={parsed_args} " + f"{detail_note}" + ), + ) + ) + except Exception as e: + ctx.record(TestResult(name, False, detail=str(e))) + + +def test_non_streaming_citations_from_tool_result(ctx: TestContext) -> None: + """Completed tool call in history -> citations against the tool result. + + Mirrors the ``test_non_streaming_with_documents`` flow but exercises + the alternate grounding path where the assistant has already issued a + tool call and the user message replays the tool's structured output + via a ``tool``-role message. The server should ground the follow-up + answer against the ``content`` of that tool message. + """ + name = "non-streaming: grounded answer from prior tool result" + try: + tool_call_id = "call_get_weather_tokyo" + body = { + "model": ctx.model, + "messages": [ + { + "role": "user", + "content": ( + "What's the current weather in Tokyo? Use the " + "get_weather tool, then answer the user and " + "ground every factual claim in the tool result " + "with a citation tag." + ), + }, + { + "role": "assistant", + "tool_plan": ( + "I should call get_weather with city=Tokyo and " + "then summarize the result for the user." + ), + "tool_calls": [ + { + "id": tool_call_id, + "type": "function", + "function": { + "name": "get_weather", + "arguments": json.dumps({"city": "Tokyo"}), + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": tool_call_id, + "content": [ + { + "type": "document", + "document": { + "id": "weather_tokyo_now", + "data": { + "city": "Tokyo", + "temperature_c": "22", + "condition": "partly cloudy", + "humidity_pct": "58", + }, + }, + } + ], + }, + ], + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather for a city.", + "parameters": { + "type": "object", + "properties": { + "city": { + "type": "string", + "description": "City name.", + } + }, + "required": ["city"], + }, + }, + } + ], + "citation_options": {"mode": "ACCURATE"}, + "max_tokens": REASONING_BUDGET, + "temperature": 0.0, + } + request_id = ctx.next_request_id("tool-result-citations") + _, resp = _post_json(ctx, body=body, request_id=request_id) + assert isinstance(resp, dict) + message = resp.get("message") or {} + content = message.get("content") or [] + citations = message.get("citations") or [] + _expect(content, f"empty content: {message}") + if not citations: + assistant_text = _text_from_content_blocks(content) or "" + assistant_thinking = _thinking_from_content_blocks(content) or "" + ctx.record( + TestResult( + name, + True, + detail=( + "WARN: assistant did not cite the tool result.\n" + f" text head: {assistant_text[:200]!r}\n" + f" thinking head: {assistant_thinking[:200]!r}\n" + " If the text contains no ' None: + """Streaming: full event lifecycle for a plain text answer. + + Reasoning Command models emit ``content-delta`` events with + ``delta.message.content.thinking`` (not ``.text``) while inside a + thinking block. We accumulate both and require *some* content to + have flowed through; if the budget ran out before the model exited + the thinking block, that's a soft pass. + """ + name = "streaming: basic text response" + try: + body = { + "model": ctx.model, + "messages": [ + {"role": "user", "content": "Count from 1 to 3, comma-separated."}, + ], + "max_tokens": REASONING_BUDGET, + "temperature": 0.0, + } + request_id = ctx.next_request_id("stream-basic") + events, saw_done = _stream_post(ctx, body=body, request_id=request_id) + _expect(saw_done, "stream did not terminate with data: [DONE]") + types = _event_types(events) + _expect( + types[0] == "message-start", + f"expected first event message-start, got {types[:3]}", + ) + _expect( + types[-1] == "message-end", + f"expected last event message-end, got {types[-3:]}", + ) + for ev in events: + _expect( + "type" in ev and ev["type"], + f"event missing discriminator: {ev}", + ) + text_chunks: list[str] = [] + thinking_chunks: list[str] = [] + for ev in events: + if ev.get("type") != "content-delta": + continue + content = (ev.get("delta") or {}).get("message", {}).get("content", {}) + if not isinstance(content, dict): + continue + t = content.get("text") + if isinstance(t, str): + text_chunks.append(t) + th = content.get("thinking") + if isinstance(th, str): + thinking_chunks.append(th) + text = "".join(text_chunks) + thinking = "".join(thinking_chunks) + _expect( + text or thinking, + f"no content accumulated from content-delta events: {types}", + ) + if text: + detail = ( + f"events={len(events)} text={text!r} thinking_chars={len(thinking)}" + ) + else: + detail = ( + f"WARN: only thinking deltas in budget " + f"({len(thinking)} chars); events={len(events)} types={types[:6]}..." + ) + ctx.record(TestResult(name, True, detail=detail)) + except Exception as e: + ctx.record(TestResult(name, False, detail=str(e))) + + +def test_streaming_with_documents(ctx: TestContext) -> None: + """Streaming: documents -> citation-start / citation-end pairs.""" + name = "streaming: grounded answer with citation events" + try: + body = { + "model": ctx.model, + "messages": [ + { + "role": "user", + "content": "Who wrote Hamlet, and when?", + }, + ], + "documents": [ + { + "id": "doc_shakespeare", + "data": { + "title": "Wikipedia: Hamlet", + "text": ( + "Hamlet is a tragedy written by William " + "Shakespeare around 1600." + ), + }, + }, + ], + "citation_options": {"mode": "ACCURATE"}, + "max_tokens": REASONING_BUDGET, + "temperature": 0.0, + } + request_id = ctx.next_request_id("stream-docs") + events, saw_done = _stream_post(ctx, body=body, request_id=request_id) + _expect(saw_done, "stream did not terminate with data: [DONE]") + types = _event_types(events) + _expect( + types[0] == "message-start" and types[-1] == "message-end", + f"unexpected envelope: {types[:3]} ... {types[-3:]}", + ) + starts = [i for i, t in enumerate(types) if t == "citation-start"] + ends = [i for i, t in enumerate(types) if t == "citation-end"] + if not starts: + ctx.record( + TestResult( + name, + True, + detail=( + "WARN: model did not emit citation events; verify " + "grounding is enabled for this model." + ), + ) + ) + return + _expect( + len(starts) == len(ends), + f"unbalanced citation events: starts={len(starts)} ends={len(ends)}", + ) + for s, e in zip(starts, ends): + _expect(s < e, f"citation-start at {s} must precede end at {e}") + ctx.record( + TestResult( + name, + True, + detail=f"events={len(events)} citations={len(starts)}", + ) + ) + except Exception as e: + ctx.record(TestResult(name, False, detail=str(e))) + + +def test_streaming_with_tools(ctx: TestContext) -> None: + """Streaming: tools -> tool-call-start / -delta / -end.""" + name = "streaming: tool call" + try: + body = { + "model": ctx.model, + "messages": [ + { + "role": "user", + "content": "What's the weather in Paris? Use get_weather.", + }, + ], + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather for a city.", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, + } + ], + "max_tokens": REASONING_BUDGET, + "temperature": 0.0, + } + request_id = ctx.next_request_id("stream-tools") + events, saw_done = _stream_post(ctx, body=body, request_id=request_id) + _expect(saw_done, "stream did not terminate with data: [DONE]") + types = _event_types(events) + starts = [i for i, t in enumerate(types) if t == "tool-call-start"] + ends = [i for i, t in enumerate(types) if t == "tool-call-end"] + if not starts: + ctx.record( + TestResult( + name, + True, + detail=( + "WARN: model did not emit tool-call events; cannot " + "verify streaming tool wiring." + ), + ) + ) + return + _expect( + len(starts) == len(ends), + f"unbalanced tool-call events: starts={len(starts)} ends={len(ends)}", + ) + delta_count = sum(1 for t in types if t == "tool-call-delta") + _expect( + delta_count > 0, + f"no tool-call-delta events between start/end: {types}", + ) + # Find the message-end event and check finish_reason maps to TOOL_CALL. + end_events = [ev for ev in events if ev.get("type") == "message-end"] + _expect(end_events, "no message-end event in stream") + end_delta = end_events[-1].get("delta") or {} + finish = end_delta.get("finish_reason") + _expect( + finish == "TOOL_CALL", + f"message-end finish_reason={finish!r}, expected TOOL_CALL", + ) + # Verify the reasoning surface matches the configured flag. + plan_events = [t for t in types if t == "tool-plan-delta"] + content_starts = [ev for ev in events if ev.get("type") == "content-start"] + thinking_blocks = [ + ev + for ev in content_starts + if ( + (ev.get("delta") or {}) + .get("message", {}) + .get("content", {}) + .get("type") + ) + == "thinking" + ] + if ctx.is_reasoning_model: + detail_note = ( + f"thinking_content_starts={len(thinking_blocks)} " + f"tool_plan_deltas={len(plan_events)}" + ) + else: + detail_note = f"tool_plan_deltas={len(plan_events)}" + ctx.record( + TestResult( + name, + True, + detail=( + f"events={len(events)} tool_calls={len(starts)} " + f"deltas={delta_count} finish={finish} {detail_note}" + ), + ) + ) + except Exception as e: + ctx.record(TestResult(name, False, detail=str(e))) + + +def test_error_invalid_request(ctx: TestContext) -> None: + """Pydantic rejection: messages must not be empty.""" + name = "error: empty messages rejected (400/422)" + try: + body = {"model": ctx.model, "messages": []} + request_id = ctx.next_request_id("err-empty") + status, parsed = _post_json( + ctx, + body=body, + request_id=request_id, + expect_status=None, + ) + _expect( + status in (400, 422), + f"expected 400 or 422, got {status}: {str(parsed)[:300]}", + ) + _expect( + _looks_like_4xx_envelope(parsed), + f"unexpected 4xx body shape: {parsed!r}", + ) + ctx.record( + TestResult( + name, + True, + detail=f"status={status} body={str(parsed)[:200]}", + ) + ) + except Exception as e: + ctx.record(TestResult(name, False, detail=str(e))) + + +def test_error_request_id_echoed(ctx: TestContext) -> None: + """Error responses should ideally echo ``X-Request-Id``. + + Three envelope shapes exist in the wild depending on where the + rejection happens: + + * Our :class:`CohereError` (``{"message": ..., "id": ...}``) -- only + this shape carries the request id. + * FastAPI body-validation (``{"detail": [...]}``). + * vLLM's standard error wrapper at + ``vllm.entrypoints.serve.utils.api_utils`` -> + ``{"error": {"message": ..., "type": ..., "code": ...}}``. This + one drops the id today; treat it as a soft pass with a warning so + we can light that gap up later without false-failing here. + """ + name = "error: X-Request-Id echoed in CohereError envelope" + try: + body = { + "model": ctx.model, + "messages": [{"role": "user", "content": "hi"}], + # negative max_tokens trips the Pydantic validator -> 4xx + "max_tokens": -1, + } + request_id = ctx.next_request_id("err-reqid") + status, parsed = _post_json( + ctx, + body=body, + request_id=request_id, + expect_status=None, + ) + _expect( + status >= 400, + f"expected 4xx, got {status}: {parsed!r}", + ) + _expect( + _looks_like_4xx_envelope(parsed), + f"unexpected 4xx body shape: {parsed!r}", + ) + if isinstance(parsed, dict) and parsed.get("id") == request_id: + ctx.record( + TestResult( + name, + True, + detail=f"status={status} CohereError.id matches request", + ) + ) + elif isinstance(parsed, dict) and isinstance(parsed.get("error"), dict): + # vLLM's standard error envelope wraps the message but drops + # the request id today. + ctx.record( + TestResult( + name, + True, + detail=( + f"status={status} WARN: vLLM error envelope dropped " + f"X-Request-Id={request_id!r}; only the 'error' " + f"wrapper is present" + ), + ) + ) + elif isinstance(parsed, dict) and "detail" in parsed: + ctx.record( + TestResult( + name, + True, + detail=( + f"status={status} (FastAPI body validation; id echo " + f"not expected at this layer)" + ), + ) + ) + else: + ctx.record( + TestResult( + name, + False, + detail=( + f"status={status} body={parsed!r} -- expected a " + f"CohereError with id={request_id!r}, a vLLM " + f"'error' wrapper, or a 422 'detail' shape" + ), + ) + ) + except Exception as e: + ctx.record(TestResult(name, False, detail=str(e))) + + +# ---------------------------------------------------------------------- +# Entrypoint +# ---------------------------------------------------------------------- + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--base-url", + default="http://127.0.0.1:8000", + help="vLLM server base URL (default: %(default)s)", + ) + parser.add_argument( + "--model", + required=True, + help="Model id as registered with the server (matches /v1/models).", + ) + parser.add_argument( + "--reasoning-model", + dest="is_reasoning_model", + action=argparse.BooleanOptionalAction, + default=True, + help=( + "Whether the server was started with --cohere-is-reasoning-model " + "(default true). Controls whether reasoning is expected as a " + "thinking content block (reasoning) or as tool-plan-delta events " + "(non-reasoning)." + ), + ) + parser.add_argument( + "--timeout", + type=float, + default=120.0, + help="HTTP request timeout, seconds (default: %(default)s)", + ) + parser.add_argument( + "--skip-streaming", + action="store_true", + help="Skip all streaming tests.", + ) + parser.add_argument( + "--skip-tools", + action="store_true", + help="Skip tool-call tests.", + ) + parser.add_argument( + "--skip-citations", + action="store_true", + help="Skip documents/citations tests.", + ) + parser.add_argument( + "--verbose", + "-v", + action="store_true", + help="Print full request/response payloads.", + ) + args = parser.parse_args() + + ctx = TestContext( + base_url=args.base_url, + model=args.model, + is_reasoning_model=args.is_reasoning_model, + verbose=args.verbose, + timeout=args.timeout, + ) + + print(f"Driving Cohere v2 endpoint at {ctx.base_url}/cohere/v2/chat") + print(f" model={ctx.model}") + print(f" is_reasoning_model={ctx.is_reasoning_model}") + print() + + test_health(ctx) + test_non_streaming_basic(ctx) + test_non_streaming_developer_role(ctx) + if not args.skip_citations: + test_non_streaming_with_documents(ctx) + if not args.skip_tools: + test_non_streaming_with_tools(ctx) + if not args.skip_tools and not args.skip_citations: + test_non_streaming_citations_from_tool_result(ctx) + if not args.skip_streaming: + test_streaming_basic(ctx) + if not args.skip_citations: + test_streaming_with_documents(ctx) + if not args.skip_tools: + test_streaming_with_tools(ctx) + test_error_invalid_request(ctx) + test_error_request_id_echoed(ctx) + + print() + print("=" * 72) + passed = sum(1 for r in ctx.results if r.passed and not r.skipped) + failed = sum(1 for r in ctx.results if not r.passed and not r.skipped) + skipped = sum(1 for r in ctx.results if r.skipped) + print(f"Summary: {passed} passed, {failed} failed, {skipped} skipped") + if failed: + print("\nFailures:") + for r in ctx.results: + if not r.passed and not r.skipped: + print(f" - {r.name}") + if r.detail: + for line in r.detail.splitlines(): + print(f" {line}") + return 0 if failed == 0 else 1 + + +if __name__ == "__main__": + sys.exit(main()) From 073e7ae9ed0ae7343a1e4ec37e69a39af9f1f126 Mon Sep 17 00:00:00 2001 From: Andrew Berneshawi Date: Wed, 12 Aug 2026 17:50:58 -0400 Subject: [PATCH 2/6] move tests to subfolder --- scripts/{ => cohere_e2e}/inspect_cohere_renderer.py | 0 scripts/{ => cohere_e2e}/test_cohere_citations_e2e.py | 0 scripts/{ => cohere_e2e}/test_cohere_output.py | 0 scripts/{ => cohere_e2e}/test_cohere_v2_e2e.py | 0 4 files changed, 0 insertions(+), 0 deletions(-) rename scripts/{ => cohere_e2e}/inspect_cohere_renderer.py (100%) rename scripts/{ => cohere_e2e}/test_cohere_citations_e2e.py (100%) rename scripts/{ => cohere_e2e}/test_cohere_output.py (100%) rename scripts/{ => cohere_e2e}/test_cohere_v2_e2e.py (100%) diff --git a/scripts/inspect_cohere_renderer.py b/scripts/cohere_e2e/inspect_cohere_renderer.py similarity index 100% rename from scripts/inspect_cohere_renderer.py rename to scripts/cohere_e2e/inspect_cohere_renderer.py diff --git a/scripts/test_cohere_citations_e2e.py b/scripts/cohere_e2e/test_cohere_citations_e2e.py similarity index 100% rename from scripts/test_cohere_citations_e2e.py rename to scripts/cohere_e2e/test_cohere_citations_e2e.py diff --git a/scripts/test_cohere_output.py b/scripts/cohere_e2e/test_cohere_output.py similarity index 100% rename from scripts/test_cohere_output.py rename to scripts/cohere_e2e/test_cohere_output.py diff --git a/scripts/test_cohere_v2_e2e.py b/scripts/cohere_e2e/test_cohere_v2_e2e.py similarity index 100% rename from scripts/test_cohere_v2_e2e.py rename to scripts/cohere_e2e/test_cohere_v2_e2e.py From 04a89d37515a386c28e633aa099bdbd5644f0e8e Mon Sep 17 00:00:00 2001 From: Andrew Berneshawi Date: Fri, 21 Aug 2026 04:42:22 -0400 Subject: [PATCH 3/6] refactor out vllm server launching Signed-off-by: Andrew Berneshawi --- .../cohere_e2e/test_cohere_citations_e2e.py | 261 +-------------- scripts/cohere_e2e/vllm_server.py | 299 ++++++++++++++++++ 2 files changed, 308 insertions(+), 252 deletions(-) create mode 100644 scripts/cohere_e2e/vllm_server.py diff --git a/scripts/cohere_e2e/test_cohere_citations_e2e.py b/scripts/cohere_e2e/test_cohere_citations_e2e.py index 03d01680700f..4e1a5a168246 100644 --- a/scripts/cohere_e2e/test_cohere_citations_e2e.py +++ b/scripts/cohere_e2e/test_cohere_citations_e2e.py @@ -101,19 +101,14 @@ from __future__ import annotations import argparse -import contextlib import json -import os -import signal -import subprocess import sys -import tempfile import time from dataclasses import dataclass, field from typing import Any -from urllib.parse import urlparse import httpx +from vllm_server import add_server_args, ensure_server, release_server DONE_LINE = "data: [DONE]" @@ -129,11 +124,7 @@ PENGUIN_DOCUMENTS = [ {"data": {"snippet": "The tallest penguin is the Emperor penguin"}}, - { - "data": { - "snippet": "The latin name for Emperor penguin is Aptenodytes forsteri" - } - }, + {"data": {"snippet": "The latin name for Emperor penguin is Aptenodytes forsteri"}}, {"data": {"snippet": "The smallest penguin is the fairy penguin"}}, {"data": {"snippet": "The latin name for fairy penguin is Eudyptula minor"}}, ] @@ -293,130 +284,6 @@ def _assert_citation_shape(citations: list[dict[str, Any]]) -> None: _expect("type" in src, f"citation source missing type: {src}") -# ---------------------------------------------------------------------- -# Optional: manage a vLLM server process ourselves -# ---------------------------------------------------------------------- - - -def _server_is_healthy(base_url: str, timeout: float = 5.0) -> bool: - try: - with httpx.Client(timeout=timeout) as client: - resp = client.get(f"{base_url.rstrip('/')}/health") - return resp.status_code == 200 - except (httpx.HTTPError, OSError): - return False - - -@dataclass -class ManagedServer: - proc: subprocess.Popen - log_path: str - - def stop(self, *, timeout: float = 30.0) -> None: - if self.proc.poll() is not None: - return - pgid: int | None = None - with contextlib.suppress(ProcessLookupError, OSError): - pgid = os.getpgid(self.proc.pid) - target = -pgid if pgid is not None else self.proc.pid - with contextlib.suppress(ProcessLookupError, OSError): - os.kill(target, signal.SIGTERM) - try: - self.proc.wait(timeout=timeout) - except subprocess.TimeoutExpired: - with contextlib.suppress(ProcessLookupError, OSError): - os.kill(target, signal.SIGKILL) - with contextlib.suppress(subprocess.TimeoutExpired): - self.proc.wait(timeout=10) - - -def _build_vllm_serve_command( - *, - model: str, - host: str, - port: int, - is_reasoning_model: bool, - extra_args: list[str], -) -> list[str]: - cmd = [ - "vllm", - "serve", - model, - "--host", - host, - "--port", - str(port), - "--tokenizer-mode", - "cohere", - "--enable-auto-tool-choice", - "--tool-call-parser", - "cohere2", - "--reasoning-parser", - "cohere2", - ] - if not is_reasoning_model: - cmd.append("--no-cohere-is-reasoning-model") - cmd.extend(extra_args) - return cmd - - -def _start_vllm_server( - *, - model: str, - host: str, - port: int, - is_reasoning_model: bool, - extra_args: list[str], -) -> ManagedServer: - cmd = _build_vllm_serve_command( - model=model, - host=host, - port=port, - is_reasoning_model=is_reasoning_model, - extra_args=extra_args, - ) - env = os.environ.copy() - env["VLLM_ENABLE_COHERE_API"] = "1" - - log_fd, log_path = tempfile.mkstemp( - prefix="vllm-cohere-citations-e2e-", suffix=".log" - ) - os.close(log_fd) - log_file = open(log_path, "w") # noqa: SIM115 -- lives as long as the server - - print(f"Server not reachable; starting it ourselves:\n {' '.join(cmd)}") - print(f" logs: {log_path}") - proc = subprocess.Popen( - cmd, - env=env, - stdout=log_file, - stderr=subprocess.STDOUT, - start_new_session=True, - ) - return ManagedServer(proc=proc, log_path=log_path) - - -def _wait_for_server( - *, base_url: str, server: ManagedServer, timeout: float -) -> None: - deadline = time.time() + timeout - while time.time() < deadline: - exit_code = server.proc.poll() - if exit_code is not None: - raise RuntimeError( - f"vLLM server process exited early (code={exit_code}); " - f"see {server.log_path} for details." - ) - if _server_is_healthy(base_url): - print(f"Server is healthy at {base_url}") - return - time.sleep(5.0) - raise TimeoutError( - f"vLLM server did not become healthy within {timeout:.0f}s; " - f"see {server.log_path} for details." - ) - - def _no_grounding_warning(content: Any) -> str: text = _text_from_content_blocks(content) or "" thinking = _thinking_from_content_blocks(content) or "" @@ -475,8 +342,7 @@ def test_documents_basic_non_stream(ctx: TestContext) -> None: name, True, detail=( - f"got {len(citations)} citation(s); " - f"doc ids referenced={doc_ids}" + f"got {len(citations)} citation(s); doc ids referenced={doc_ids}" ), ) ) @@ -697,9 +563,7 @@ def test_grounding_against_tool_result_in_history(ctx: TestContext) -> None: }, { "role": "assistant", - "tool_plan": ( - "I will search for the second tallest mountain." - ), + "tool_plan": ("I will search for the second tallest mountain."), "tool_calls": [ { "id": "internet_search_0123", @@ -720,9 +584,7 @@ def test_grounding_against_tool_result_in_history(ctx: TestContext) -> None: { "type": "document", "document": { - "data": { - "result": "The second tallest mountain is K2." - } + "data": {"result": "The second tallest mountain is K2."} }, } ], @@ -790,9 +652,7 @@ def test_grounding_against_tool_result_in_history(ctx: TestContext) -> None: ctx.record(TestResult(name, True, detail=_no_grounding_warning(content))) return _assert_citation_shape(citations) - source_types = { - s.get("type") for c in citations for s in c.get("sources", []) - } + source_types = {s.get("type") for c in citations for s in c.get("sources", [])} ctx.record( TestResult( name, @@ -1133,78 +993,13 @@ def test_citations_in_conversation_history(ctx: TestContext) -> None: def main() -> int: parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument( - "--base-url", - default="http://127.0.0.1:8000", - help="vLLM server base URL (default: %(default)s)", - ) - parser.add_argument( - "--model", - default=DEFAULT_MODEL, - help=( - "Model id as registered with the server (matches /v1/models). " - "Also used to launch the server when auto-starting it. " - "(default: %(default)s)" - ), - ) - parser.add_argument( - "--reasoning-model", - dest="is_reasoning_model", - action=argparse.BooleanOptionalAction, - default=True, - help=( - "Whether the server was (or should be) started with " - "--cohere-is-reasoning-model (default true). Controls whether " - "the thinking+citations cases run, and whether an auto-started " - "server gets --no-cohere-is-reasoning-model." - ), - ) + add_server_args(parser, default_model=DEFAULT_MODEL) parser.add_argument( "--timeout", type=float, default=120.0, help="HTTP request timeout, seconds (default: %(default)s)", ) - parser.add_argument( - "--auto-start-server", - action=argparse.BooleanOptionalAction, - default=True, - help=( - "Launch 'vllm serve --model ' ourselves if --base-url " - "isn't already serving a healthy /health response (default true)." - ), - ) - parser.add_argument( - "--keep-server", - action=argparse.BooleanOptionalAction, - default=True, - help=( - "Leave an auto-started server running after the tests finish, " - "so re-runs skip the model load (default true). Only applies " - "when this script started the server itself." - ), - ) - parser.add_argument( - "--startup-timeout", - type=float, - default=1800.0, - help=( - "Seconds to wait for an auto-started server to become healthy " - "(default: %(default)s). Large quantized checkpoints can take " - "a long time to load." - ), - ) - parser.add_argument( - "--extra-server-arg", - dest="extra_server_args", - action="append", - default=[], - help=( - "Extra 'vllm serve' argument, forwarded verbatim when this " - "script starts the server itself. Repeatable, e.g. " - "--extra-server-arg=--tensor-parallel-size=8" - ), - ) parser.add_argument( "--verbose", "-v", @@ -1213,32 +1008,7 @@ def main() -> int: ) args = parser.parse_args() - managed_server: ManagedServer | None = None - if args.auto_start_server and not _server_is_healthy(args.base_url): - parsed = urlparse(args.base_url) - host = parsed.hostname or "127.0.0.1" - port = parsed.port or 8000 - managed_server = _start_vllm_server( - model=args.model, - host=host, - port=port, - is_reasoning_model=args.is_reasoning_model, - extra_args=args.extra_server_args, - ) - try: - _wait_for_server( - base_url=args.base_url, - server=managed_server, - timeout=args.startup_timeout, - ) - except Exception: - managed_server.stop() - raise - elif not _server_is_healthy(args.base_url): - print( - f"WARNING: {args.base_url} does not look healthy and " - f"--no-auto-start-server was passed; tests will likely fail." - ) + managed_server = ensure_server(args, log_prefix="vllm-cohere-citations-e2e-") try: ctx = TestContext( @@ -1283,20 +1053,7 @@ def main() -> int: print(f" {line}") return 0 if failed == 0 else 1 finally: - if managed_server is not None: - if args.keep_server: - print( - f"\nLeaving auto-started server running " - f"(pid={managed_server.proc.pid}, " - f"log={managed_server.log_path}).\n" - f"Stop it with: kill {managed_server.proc.pid}" - ) - else: - print( - f"\nStopping auto-started server " - f"(pid={managed_server.proc.pid})..." - ) - managed_server.stop() + release_server(managed_server, keep=args.keep_server) if __name__ == "__main__": diff --git a/scripts/cohere_e2e/vllm_server.py b/scripts/cohere_e2e/vllm_server.py new file mode 100644 index 000000000000..da111b0be895 --- /dev/null +++ b/scripts/cohere_e2e/vllm_server.py @@ -0,0 +1,299 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Shared harness for the Cohere e2e scripts that need a live vLLM server. + +The scripts in this directory drive a real ``vllm serve`` process over +HTTP. Rather than making the caller start one by hand, they can point at +an existing ``--base-url`` and fall back to launching (and tearing down) +a server themselves. + +Two halves that are meant to be used together: + +* :func:`add_server_args` registers the CLI flags this module consumes. +* :func:`ensure_server` reads the resulting namespace and returns a + :class:`ManagedServer` when it had to start one (``None`` when the + caller-supplied URL was already healthy, so we own nothing). + +A typical ``main()`` looks like:: + + parser = argparse.ArgumentParser(description=__doc__) + add_server_args(parser, default_model=DEFAULT_MODEL) + ... + args = parser.parse_args() + + managed = ensure_server(args, log_prefix="vllm-cohere-foo-e2e-") + try: + ... # run the tests against args.base_url + finally: + release_server(managed, keep=args.keep_server) + +Because these are standalone scripts (no ``__init__.py`` here), import +this as a flat sibling module -- ``from vllm_server import ...`` -- which +resolves via the script directory Python puts on ``sys.path``. +""" + +from __future__ import annotations + +import argparse +import contextlib +import os +import signal +import subprocess +import tempfile +import time +from dataclasses import dataclass +from urllib.parse import urlparse + +import httpx + +DEFAULT_BASE_URL = "http://127.0.0.1:8000" + + +def server_is_healthy(base_url: str, timeout: float = 5.0) -> bool: + try: + with httpx.Client(timeout=timeout) as client: + resp = client.get(f"{base_url.rstrip('/')}/health") + return resp.status_code == 200 + except (httpx.HTTPError, OSError): + return False + + +@dataclass +class ManagedServer: + proc: subprocess.Popen + log_path: str + + def stop(self, *, timeout: float = 30.0) -> None: + if self.proc.poll() is not None: + return + pgid: int | None = None + with contextlib.suppress(ProcessLookupError, OSError): + pgid = os.getpgid(self.proc.pid) + target = -pgid if pgid is not None else self.proc.pid + with contextlib.suppress(ProcessLookupError, OSError): + os.kill(target, signal.SIGTERM) + try: + self.proc.wait(timeout=timeout) + except subprocess.TimeoutExpired: + with contextlib.suppress(ProcessLookupError, OSError): + os.kill(target, signal.SIGKILL) + with contextlib.suppress(subprocess.TimeoutExpired): + self.proc.wait(timeout=10) + + +def build_serve_command( + *, + model: str, + host: str, + port: int, + is_reasoning_model: bool, + extra_args: list[str], +) -> list[str]: + cmd = [ + "vllm", + "serve", + model, + "--host", + host, + "--port", + str(port), + "--tokenizer-mode", + "cohere", + "--enable-auto-tool-choice", + "--tool-call-parser", + "cohere2", + "--reasoning-parser", + "cohere2", + ] + if not is_reasoning_model: + cmd.append("--no-cohere-is-reasoning-model") + cmd.extend(extra_args) + return cmd + + +def start_server( + *, + model: str, + host: str, + port: int, + is_reasoning_model: bool, + extra_args: list[str], + log_prefix: str = "vllm-cohere-e2e-", +) -> ManagedServer: + cmd = build_serve_command( + model=model, + host=host, + port=port, + is_reasoning_model=is_reasoning_model, + extra_args=extra_args, + ) + env = os.environ.copy() + env["VLLM_ENABLE_COHERE_API"] = "1" + + log_fd, log_path = tempfile.mkstemp(prefix=log_prefix, suffix=".log") + os.close(log_fd) + log_file = open(log_path, "w") # noqa: SIM115 -- lives as long as the server + + print(f"Server not reachable; starting it ourselves:\n {' '.join(cmd)}") + print(f" logs: {log_path}") + proc = subprocess.Popen( + cmd, + env=env, + stdout=log_file, + stderr=subprocess.STDOUT, + start_new_session=True, + ) + return ManagedServer(proc=proc, log_path=log_path) + + +def wait_until_healthy(*, base_url: str, server: ManagedServer, timeout: float) -> None: + deadline = time.time() + timeout + while time.time() < deadline: + exit_code = server.proc.poll() + if exit_code is not None: + raise RuntimeError( + f"vLLM server process exited early (code={exit_code}); " + f"see {server.log_path} for details." + ) + if server_is_healthy(base_url): + print(f"Server is healthy at {base_url}") + return + time.sleep(5.0) + raise TimeoutError( + f"vLLM server did not become healthy within {timeout:.0f}s; " + f"see {server.log_path} for details." + ) + + +def add_server_args( + parser: argparse.ArgumentParser, + *, + default_model: str, + default_base_url: str = DEFAULT_BASE_URL, +) -> None: + """Register the flags :func:`ensure_server` reads.""" + parser.add_argument( + "--base-url", + default=default_base_url, + help="vLLM server base URL (default: %(default)s)", + ) + parser.add_argument( + "--model", + default=default_model, + help=( + "Model id as registered with the server (matches /v1/models). " + "Also used to launch the server when auto-starting it. " + "(default: %(default)s)" + ), + ) + parser.add_argument( + "--reasoning-model", + dest="is_reasoning_model", + action=argparse.BooleanOptionalAction, + default=True, + help=( + "Whether the server was (or should be) started with " + "--cohere-is-reasoning-model (default true). Controls whether " + "the thinking cases run, and whether an auto-started server " + "gets --no-cohere-is-reasoning-model." + ), + ) + parser.add_argument( + "--auto-start-server", + action=argparse.BooleanOptionalAction, + default=True, + help=( + "Launch 'vllm serve --model ' ourselves if --base-url " + "isn't already serving a healthy /health response (default true)." + ), + ) + parser.add_argument( + "--keep-server", + action=argparse.BooleanOptionalAction, + default=True, + help=( + "Leave an auto-started server running after the tests finish, " + "so re-runs skip the model load (default true). Only applies " + "when this script started the server itself." + ), + ) + parser.add_argument( + "--startup-timeout", + type=float, + default=1800.0, + help=( + "Seconds to wait for an auto-started server to become healthy " + "(default: %(default)s). Large quantized checkpoints can take " + "a long time to load." + ), + ) + parser.add_argument( + "--extra-server-arg", + dest="extra_server_args", + action="append", + default=[], + help=( + "Extra 'vllm serve' argument, forwarded verbatim when this " + "script starts the server itself. Repeatable, e.g. " + "--extra-server-arg=--tensor-parallel-size=8" + ), + ) + + +def ensure_server( + args: argparse.Namespace, + *, + log_prefix: str = "vllm-cohere-e2e-", +) -> ManagedServer | None: + """Make ``args.base_url`` serviceable, starting a server if needed. + + Returns the :class:`ManagedServer` we own and the caller must release + via :func:`release_server`, or ``None`` when the URL was already + healthy (or auto-start was declined, which only warns -- the tests + then fail on their own and say why). + """ + if not args.auto_start_server: + if not server_is_healthy(args.base_url): + print( + f"WARNING: {args.base_url} does not look healthy and " + f"--no-auto-start-server was passed; tests will likely fail." + ) + return None + + if server_is_healthy(args.base_url): + return None + + parsed = urlparse(args.base_url) + server = start_server( + model=args.model, + host=parsed.hostname or "127.0.0.1", + port=parsed.port or 8000, + is_reasoning_model=args.is_reasoning_model, + extra_args=args.extra_server_args, + log_prefix=log_prefix, + ) + try: + wait_until_healthy( + base_url=args.base_url, + server=server, + timeout=args.startup_timeout, + ) + except Exception: + server.stop() + raise + return server + + +def release_server(server: ManagedServer | None, *, keep: bool) -> None: + """Stop a server we started, or explain how to stop it later.""" + if server is None: + return + if keep: + print( + f"\nLeaving auto-started server running " + f"(pid={server.proc.pid}, log={server.log_path}).\n" + f"Stop it with: kill {server.proc.pid}" + ) + else: + print(f"\nStopping auto-started server (pid={server.proc.pid})...") + server.stop() From 004ce5ff5b861b81275b2ee281f2b345ab1d9d3b Mon Sep 17 00:00:00 2001 From: Andrew Berneshawi Date: Fri, 21 Aug 2026 18:42:09 -0400 Subject: [PATCH 4/6] add render e2e test --- scripts/cohere_e2e/test_cohere_render_e2e.py | 600 +++++++++++++++++++ 1 file changed, 600 insertions(+) create mode 100644 scripts/cohere_e2e/test_cohere_render_e2e.py diff --git a/scripts/cohere_e2e/test_cohere_render_e2e.py b/scripts/cohere_e2e/test_cohere_render_e2e.py new file mode 100644 index 000000000000..b75af6a5a367 --- /dev/null +++ b/scripts/cohere_e2e/test_cohere_render_e2e.py @@ -0,0 +1,600 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""End-to-end checks for ``POST /cohere/v2/chat/render``. + +Drives a live vLLM server with a real Cohere Command-family model. The +render endpoint is the Cohere counterpart to +``POST /v1/chat/completions/render``: it converts the v2 body to a +``ChatCompletionRequest``, runs it through the same renderer pipeline the +chat endpoint uses, and returns the resulting ``GenerateRequest`` +(prompt ``token_ids`` plus ``sampling_params``) *without* generating. + +That "no generation" property is what makes this script cheap: every +case below is a prompt-construction check, so the model never decodes a +token. It still needs a loaded server because tokenization and the +melody templating step both run server-side against the real tokenizer. + +What the cases are actually pinning down: + +* The returned tokens really are a *melody* prompt, not a Jinja one -- + verified by detokenizing them and looking for Cohere turn markers. + Renderer selection is a server-level property (``--tokenizer-mode + cohere`` picks ``CohereRenderer``), so this is the check that the + render route inherits it rather than falling back to ``HfRenderer``. +* Cohere-only request surface (``documents``, ``tools``, ``safety_mode``) + reaches the template. These ride on ``chat_template_kwargs``, which is + the part of the conversion most likely to silently stop being + forwarded. +* v2 sampling fields land in ``sampling_params`` under their OpenAI + names (``p`` -> ``top_p``, ``k`` -> ``top_k``, ``stop_sequences`` -> + ``stop``). +* Byte-for-byte token parity with ``/v1/chat/completions/render`` for a + request that carries no Cohere-specific fields -- both routes share + one ``ServingRender``, so any divergence means the v2 conversion is + perturbing the prompt. + +Usage +----- + +Same server-management flags as the sibling e2e scripts (see +``vllm_server.py``): point ``--base-url`` at a running server, or let +this script start one:: + + python scripts/cohere_e2e/test_cohere_render_e2e.py + + python scripts/cohere_e2e/test_cohere_render_e2e.py \\ + --base-url http://127.0.0.1:8000 --no-auto-start-server + +The server must be started with ``VLLM_ENABLE_COHERE_API=1`` and +``--tokenizer-mode cohere`` for these checks to mean anything; +``vllm_server.py`` does both when it owns the process. +""" + +from __future__ import annotations + +import argparse +import json +import sys +import time +from dataclasses import dataclass, field +from typing import Any + +import httpx +from vllm_server import add_server_args, ensure_server, release_server + +DEFAULT_MODEL = "CohereLabs/command-a-plus-05-2026-w4a4" + +# Turn markers every cmd3/cmd4 melody template emits. A Jinja-rendered +# prompt for a non-Cohere template would not contain these, so they are +# how we tell which renderer actually ran. +START_OF_TURN = "<|START_OF_TURN_TOKEN|>" +USER_TOKEN = "<|USER_TOKEN|>" + +PENGUIN_DOCUMENTS = [ + {"data": {"snippet": "The tallest penguin is the Emperor penguin"}}, + {"data": {"snippet": "The latin name for Emperor penguin is Aptenodytes forsteri"}}, +] + +WEATHER_TOOL = { + "type": "function", + "function": { + "name": "get_weather_for_city", + "description": "Look up the current weather for a city.", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, +} + + +@dataclass +class TestResult: + name: str + passed: bool + detail: str = "" + skipped: bool = False + + +@dataclass +class TestContext: + base_url: str + model: str + is_reasoning_model: bool + verbose: bool + timeout: float + request_id_counter: int = 0 + results: list[TestResult] = field(default_factory=list) + + def next_request_id(self, label: str) -> str: + self.request_id_counter += 1 + return f"e2e-render-{label}-{self.request_id_counter}-{int(time.time())}" + + def record(self, result: TestResult) -> None: + self.results.append(result) + prefix = "SKIP" if result.skipped else ("PASS" if result.passed else "FAIL") + print(f"[{prefix}] {result.name}") + if result.detail: + for line in result.detail.splitlines(): + print(f" {line}") + + +def _expect(condition: bool, message: str) -> None: + if not condition: + raise AssertionError(message) + + +def _looks_like_4xx_envelope(parsed: Any) -> bool: + """True for any of the documented Cohere v2 4xx response shapes. + + Mirrors the helper in ``test_cohere_v2_e2e.py``: the rejection can + come from our ``CohereError``, vLLM's generic OpenAI-style envelope, + or FastAPI's default ``{"detail": [...]}``, depending on how far into + the stack the request got. + """ + if not isinstance(parsed, dict): + return False + if "message" in parsed and "error" not in parsed: + return True + if isinstance(parsed.get("error"), dict) and "message" in parsed["error"]: + return True + return "detail" in parsed + + +def _post( + ctx: TestContext, + *, + path: str, + body: dict[str, Any], + request_id: str, + expect_status: int | None = 200, +) -> tuple[int, dict[str, Any] | str]: + url = f"{ctx.base_url.rstrip('/')}{path}" + headers = {"Content-Type": "application/json", "X-Request-Id": request_id} + with httpx.Client(timeout=ctx.timeout) as client: + resp = client.post(url, json=body, headers=headers) + if ctx.verbose: + print(f" -> POST {url} [status {resp.status_code}]") + print(f" request_id={request_id}") + print(f" body={json.dumps(body)[:800]}") + print(f" resp={resp.text[:2000]}") + parsed: dict[str, Any] | str + try: + parsed = resp.json() + except Exception: + parsed = resp.text + if expect_status is not None and resp.status_code != expect_status: + raise AssertionError( + f"expected status {expect_status}, got {resp.status_code}: {parsed!r}" + ) + return resp.status_code, parsed + + +def _render( + ctx: TestContext, + *, + body: dict[str, Any], + label: str, + expect_status: int | None = 200, +) -> tuple[int, dict[str, Any] | str]: + return _post( + ctx, + path="/cohere/v2/chat/render", + body=body, + request_id=ctx.next_request_id(label), + expect_status=expect_status, + ) + + +def _token_ids(payload: Any) -> list[int]: + _expect(isinstance(payload, dict), f"render response is not an object: {payload!r}") + token_ids = payload.get("token_ids") + _expect( + isinstance(token_ids, list) and bool(token_ids), + f"token_ids missing or empty: {str(payload)[:300]}", + ) + _expect( + all(isinstance(t, int) and t >= 0 for t in token_ids), + "token_ids must be non-negative ints", + ) + return token_ids + + +def _sampling_params(payload: Any) -> dict[str, Any]: + params = payload.get("sampling_params") if isinstance(payload, dict) else None + _expect( + isinstance(params, dict), + f"sampling_params missing or not an object: {str(payload)[:300]}", + ) + return params + + +def _detokenize(ctx: TestContext, token_ids: list[int]) -> str: + """Decode rendered tokens back to the prompt string. + + Uses the server's own ``/detokenize`` so we read the prompt through + the same tokenizer that produced it, rather than loading one here. + """ + _, parsed = _post( + ctx, + path="/detokenize", + body={"model": ctx.model, "tokens": token_ids}, + request_id=ctx.next_request_id("detok"), + ) + _expect(isinstance(parsed, dict), f"detokenize returned non-object: {parsed!r}") + prompt = parsed.get("prompt") + _expect(isinstance(prompt, str), f"detokenize response missing prompt: {parsed!r}") + return prompt + + +def _simple_body(ctx: TestContext, text: str = "What is the tallest penguin?") -> dict: + return { + "model": ctx.model, + "messages": [{"role": "user", "content": text}], + "max_tokens": 32, + } + + +# ---------------------------------------------------------------------- +# Cases +# ---------------------------------------------------------------------- + + +def test_render_basic_shape(ctx: TestContext) -> None: + """A minimal v2 body renders to a GenerateRequest and nothing else.""" + name = "render: returns GenerateRequest shape" + try: + _, parsed = _render(ctx, body=_simple_body(ctx), label="basic") + token_ids = _token_ids(parsed) + assert isinstance(parsed, dict) + + _expect("request_id" in parsed, f"no request_id: {str(parsed)[:200]}") + _sampling_params(parsed) + + # Render must not generate: none of the chat-response keys should + # be here. This is the cheap guard that the route didn't get + # wired to the chat handler by mistake. + for generated_key in ("message", "text", "finish_reason", "citations"): + _expect( + generated_key not in parsed, + f"render response leaked generated field {generated_key!r}", + ) + + ctx.record( + TestResult( + name, + True, + detail=( + f"{len(token_ids)} tokens, request_id={parsed.get('request_id')!r}" + ), + ) + ) + except Exception as e: + ctx.record(TestResult(name, False, detail=str(e))) + + +def test_render_tokens_decode_to_melody_prompt(ctx: TestContext) -> None: + """The rendered tokens are a melody prompt, not a Jinja one.""" + name = "render: tokens decode to a melody prompt" + try: + question = "What is the tallest penguin?" + _, parsed = _render(ctx, body=_simple_body(ctx, question), label="melody") + prompt = _detokenize(ctx, _token_ids(parsed)) + + _expect( + START_OF_TURN in prompt, + f"no {START_OF_TURN} in decoded prompt; renderer may not be " + f"CohereRenderer. prompt={prompt[:400]!r}", + ) + _expect( + USER_TOKEN in prompt, + f"no {USER_TOKEN} in decoded prompt: {prompt[:400]!r}", + ) + _expect( + question in prompt, + f"user message missing from decoded prompt: {prompt[:400]!r}", + ) + + ctx.record(TestResult(name, True, detail=f"prompt starts: {prompt[:160]!r}")) + except Exception as e: + ctx.record(TestResult(name, False, detail=str(e))) + + +def test_render_documents_reach_the_prompt(ctx: TestContext) -> None: + """``documents`` are grounded into the prompt, not dropped.""" + name = "render: documents reach the prompt" + try: + baseline_body = _simple_body(ctx) + _, baseline = _render(ctx, body=baseline_body, label="docs-baseline") + baseline_tokens = _token_ids(baseline) + + grounded_body = {**baseline_body, "documents": PENGUIN_DOCUMENTS} + _, grounded = _render(ctx, body=grounded_body, label="docs") + grounded_tokens = _token_ids(grounded) + + _expect( + len(grounded_tokens) > len(baseline_tokens), + f"documents did not grow the prompt " + f"({len(baseline_tokens)} -> {len(grounded_tokens)} tokens)", + ) + + prompt = _detokenize(ctx, grounded_tokens) + for doc in PENGUIN_DOCUMENTS: + snippet = doc["data"]["snippet"] + _expect( + snippet in prompt, + f"document snippet missing from prompt: {snippet!r}", + ) + + ctx.record( + TestResult( + name, + True, + detail=( + f"{len(baseline_tokens)} -> {len(grounded_tokens)} tokens " + f"with {len(PENGUIN_DOCUMENTS)} documents" + ), + ) + ) + except Exception as e: + ctx.record(TestResult(name, False, detail=str(e))) + + +def test_render_tools_reach_the_prompt(ctx: TestContext) -> None: + """``tools`` are templated into the prompt.""" + name = "render: tools reach the prompt" + try: + body = {**_simple_body(ctx, "What is the weather in Toronto?")} + body["tools"] = [WEATHER_TOOL] + _, parsed = _render(ctx, body=body, label="tools") + prompt = _detokenize(ctx, _token_ids(parsed)) + + tool_name = WEATHER_TOOL["function"]["name"] + _expect( + tool_name in prompt, + f"tool name {tool_name!r} missing from prompt: {prompt[:600]!r}", + ) + + ctx.record(TestResult(name, True, detail=f"found {tool_name!r} in prompt")) + except Exception as e: + ctx.record(TestResult(name, False, detail=str(e))) + + +def test_render_safety_mode_changes_the_prompt(ctx: TestContext) -> None: + """``safety_mode`` flows through chat_template_kwargs into melody. + + Two different modes must produce two different preambles; if the + field stopped being forwarded both would render identically. + """ + name = "render: safety_mode changes the prompt" + try: + base = _simple_body(ctx) + _, strict = _render( + ctx, body={**base, "safety_mode": "STRICT"}, label="safety-strict" + ) + _, contextual = _render( + ctx, body={**base, "safety_mode": "CONTEXTUAL"}, label="safety-ctx" + ) + + strict_tokens = _token_ids(strict) + contextual_tokens = _token_ids(contextual) + _expect( + strict_tokens != contextual_tokens, + "STRICT and CONTEXTUAL safety_mode rendered identical prompts; " + "safety_mode is probably not reaching the template", + ) + + ctx.record( + TestResult( + name, + True, + detail=( + f"STRICT={len(strict_tokens)} tokens, " + f"CONTEXTUAL={len(contextual_tokens)} tokens" + ), + ) + ) + except Exception as e: + ctx.record(TestResult(name, False, detail=str(e))) + + +def test_render_sampling_params_mapped(ctx: TestContext) -> None: + """v2 sampling fields land in sampling_params under OpenAI names.""" + name = "render: v2 sampling params map into sampling_params" + try: + body = { + "model": ctx.model, + "messages": [{"role": "user", "content": "Say hi."}], + "max_tokens": 17, + "temperature": 0.42, + "p": 0.83, + "k": 7, + "stop_sequences": ["STOP_HERE"], + "seed": 1234, + "frequency_penalty": 0.25, + "presence_penalty": 0.5, + } + _, parsed = _render(ctx, body=body, label="sampling") + params = _sampling_params(parsed) + + expected = { + "max_tokens": 17, + "temperature": 0.42, + "top_p": 0.83, + "top_k": 7, + "seed": 1234, + "frequency_penalty": 0.25, + "presence_penalty": 0.5, + } + mismatches = [ + f"{key}: expected {value!r}, got {params.get(key)!r}" + for key, value in expected.items() + if params.get(key) != value + ] + _expect( + not mismatches, "sampling_params mismatch:\n " + "\n ".join(mismatches) + ) + + stop = params.get("stop") + _expect( + stop == ["STOP_HERE"] or stop == "STOP_HERE", + f"stop_sequences did not map to stop: {stop!r}", + ) + + ctx.record( + TestResult(name, True, detail=f"all {len(expected) + 1} fields mapped") + ) + except Exception as e: + ctx.record(TestResult(name, False, detail=str(e))) + + +def test_render_matches_chat_completions_render(ctx: TestContext) -> None: + """A plain request renders identically on both render routes. + + Both go through the same ``ServingRender`` and the same + ``CohereRenderer``, so for a body with no Cohere-specific fields the + token ids should match exactly. A difference means the v2 + conversion is perturbing the prompt. + """ + name = "render: token parity with /v1/chat/completions/render" + try: + question = "What is the tallest penguin?" + _, v2 = _render(ctx, body=_simple_body(ctx, question), label="parity-v2") + v2_tokens = _token_ids(v2) + + _, v1 = _post( + ctx, + path="/v1/chat/completions/render", + body={ + "model": ctx.model, + "messages": [{"role": "user", "content": question}], + "max_tokens": 32, + }, + request_id=ctx.next_request_id("parity-v1"), + ) + v1_tokens = _token_ids(v1) + + if v2_tokens != v1_tokens: + v2_prompt = _detokenize(ctx, v2_tokens) + v1_prompt = _detokenize(ctx, v1_tokens) + raise AssertionError( + f"token ids differ ({len(v2_tokens)} vs {len(v1_tokens)} tokens)\n" + f" v2: {v2_prompt[:300]!r}\n" + f" v1: {v1_prompt[:300]!r}" + ) + + ctx.record( + TestResult(name, True, detail=f"identical {len(v2_tokens)} token prompt") + ) + except Exception as e: + ctx.record(TestResult(name, False, detail=str(e))) + + +def test_render_stream_flag_echoed(ctx: TestContext) -> None: + """``stream`` is carried into the GenerateRequest rather than + turning the render response itself into a stream.""" + name = "render: stream flag carried into GenerateRequest" + try: + _, parsed = _render( + ctx, body={**_simple_body(ctx), "stream": True}, label="stream" + ) + _token_ids(parsed) + assert isinstance(parsed, dict) + _expect( + parsed.get("stream") is True, + f"stream not carried through: {parsed.get('stream')!r}", + ) + ctx.record(TestResult(name, True, detail="stream=True echoed")) + except Exception as e: + ctx.record(TestResult(name, False, detail=str(e))) + + +def test_render_invalid_request_rejected(ctx: TestContext) -> None: + """A malformed v2 body is rejected as 4xx, not a 500.""" + name = "render: empty messages rejected (400/422)" + try: + status, parsed = _render( + ctx, + body={"model": ctx.model, "messages": []}, + label="err-empty", + expect_status=None, + ) + _expect( + status in (400, 422), + f"expected 400 or 422, got {status}: {str(parsed)[:300]}", + ) + _expect( + _looks_like_4xx_envelope(parsed), + f"unexpected 4xx body shape: {parsed!r}", + ) + ctx.record( + TestResult(name, True, detail=f"status={status} body={str(parsed)[:200]}") + ) + except Exception as e: + ctx.record(TestResult(name, False, detail=str(e))) + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + add_server_args(parser, default_model=DEFAULT_MODEL) + parser.add_argument( + "--timeout", + type=float, + default=120.0, + help="HTTP request timeout, seconds (default: %(default)s)", + ) + parser.add_argument( + "--verbose", + "-v", + action="store_true", + help="Print full request/response payloads.", + ) + args = parser.parse_args() + + managed_server = ensure_server(args, log_prefix="vllm-cohere-render-e2e-") + + try: + ctx = TestContext( + base_url=args.base_url, + model=args.model, + is_reasoning_model=args.is_reasoning_model, + verbose=args.verbose, + timeout=args.timeout, + ) + + print(f"Driving Cohere v2 render tests at {ctx.base_url}/cohere/v2/chat/render") + print(f" model={ctx.model}") + print() + + test_render_basic_shape(ctx) + test_render_tokens_decode_to_melody_prompt(ctx) + test_render_documents_reach_the_prompt(ctx) + test_render_tools_reach_the_prompt(ctx) + test_render_safety_mode_changes_the_prompt(ctx) + test_render_sampling_params_mapped(ctx) + test_render_matches_chat_completions_render(ctx) + test_render_stream_flag_echoed(ctx) + test_render_invalid_request_rejected(ctx) + + print() + print("=" * 72) + passed = sum(1 for r in ctx.results if r.passed and not r.skipped) + failed = sum(1 for r in ctx.results if not r.passed and not r.skipped) + skipped = sum(1 for r in ctx.results if r.skipped) + print(f"Summary: {passed} passed, {failed} failed, {skipped} skipped") + if failed: + print("\nFailures:") + for r in ctx.results: + if not r.passed and not r.skipped: + print(f" - {r.name}") + if r.detail: + for line in r.detail.splitlines(): + print(f" {line}") + return 0 if failed == 0 else 1 + finally: + release_server(managed_server, keep=args.keep_server) + + +if __name__ == "__main__": + sys.exit(main()) From 19b3eaf2a30931f284c9b51b038b91ecfc538712 Mon Sep 17 00:00:00 2001 From: Andrew Berneshawi Date: Wed, 26 Aug 2026 09:23:22 -0400 Subject: [PATCH 5/6] Fix bad parser naming --- .../cohere_e2e/test_cohere_citations_e2e.py | 11 +++++-- scripts/cohere_e2e/test_cohere_v2_e2e.py | 7 +++-- scripts/cohere_e2e/vllm_server.py | 29 +++++++++++++++++-- vllm/renderers/cohere.py | 3 +- 4 files changed, 42 insertions(+), 8 deletions(-) diff --git a/scripts/cohere_e2e/test_cohere_citations_e2e.py b/scripts/cohere_e2e/test_cohere_citations_e2e.py index 4e1a5a168246..95022259abcd 100644 --- a/scripts/cohere_e2e/test_cohere_citations_e2e.py +++ b/scripts/cohere_e2e/test_cohere_citations_e2e.py @@ -51,7 +51,8 @@ ``/health`` response, the script launches ``vllm serve`` itself with the Cohere renderer / tokenizer wired up (``VLLM_ENABLE_COHERE_API=1``, ``--tokenizer-mode cohere``, ``--enable-auto-tool-choice``, - ``--tool-call-parser cohere2``, ``--reasoning-parser cohere2``), waits + ``--tool-call-parser cohere_command4``, + ``--reasoning-parser cohere_command4``), waits for it to become healthy, then runs the test suite: python scripts/test_cohere_citations_e2e.py @@ -66,10 +67,14 @@ VLLM_ENABLE_COHERE_API=1 vllm serve \\ --tokenizer-mode cohere \\ --enable-auto-tool-choice \\ - --tool-call-parser cohere2 \\ - --reasoning-parser cohere2 \\ + --tool-call-parser cohere_command4 \\ + --reasoning-parser cohere_command4 \\ --port 8000 + The parser name must match the checkpoint's prompt generation: use + ``cohere_command3`` for cmd3-generation models (and + ``--cohere-parser cohere_command3`` when this script owns the server). + For non-reasoning Command models (cmd3, older Command R), pass ``--no-reasoning-model`` to this script; when it owns the server process, that also adds ``--no-cohere-is-reasoning-model`` to the diff --git a/scripts/cohere_e2e/test_cohere_v2_e2e.py b/scripts/cohere_e2e/test_cohere_v2_e2e.py index 9e3fe83140db..28f193a611ab 100644 --- a/scripts/cohere_e2e/test_cohere_v2_e2e.py +++ b/scripts/cohere_e2e/test_cohere_v2_e2e.py @@ -21,10 +21,13 @@ VLLM_ENABLE_COHERE_API=1 vllm serve \\ --tokenizer-mode cohere \\ --enable-auto-tool-choice \\ - --tool-call-parser cohere2 \\ - --reasoning-parser cohere2 \\ + --tool-call-parser cohere_command4 \\ + --reasoning-parser cohere_command4 \\ --port 8000 + Use ``cohere_command3`` for both parsers when serving a + cmd3-generation checkpoint. + For non-reasoning Command models (cmd3, older Command R), append ``--no-cohere-is-reasoning-model`` so the renderer surfaces reasoning as ``tool_plan`` instead of as a ``thinking`` content block. diff --git a/scripts/cohere_e2e/vllm_server.py b/scripts/cohere_e2e/vllm_server.py index da111b0be895..f5f931ed13ff 100644 --- a/scripts/cohere_e2e/vllm_server.py +++ b/scripts/cohere_e2e/vllm_server.py @@ -48,6 +48,16 @@ DEFAULT_BASE_URL = "http://127.0.0.1:8000" +# Tool-call and reasoning parser name, passed to both ``--tool-call-parser`` +# and ``--reasoning-parser``. Must be a key registered in +# ``vllm.tool_parsers`` / ``vllm.reasoning``: today that means +# ``cohere_command3`` or ``cohere_command4``. cmd4 is the default because +# it matches ``CohereRenderer``'s own default prompt format +# (``_DEFAULT_FORMAT``); serving a cmd3-generation checkpoint needs +# ``--cohere-parser cohere_command3`` so the parsers agree with the +# markup the model was trained to emit. +DEFAULT_COHERE_PARSER = "cohere_command4" + def server_is_healthy(base_url: str, timeout: float = 5.0) -> bool: try: @@ -88,6 +98,7 @@ def build_serve_command( port: int, is_reasoning_model: bool, extra_args: list[str], + cohere_parser: str = DEFAULT_COHERE_PARSER, ) -> list[str]: cmd = [ "vllm", @@ -101,9 +112,9 @@ def build_serve_command( "cohere", "--enable-auto-tool-choice", "--tool-call-parser", - "cohere2", + cohere_parser, "--reasoning-parser", - "cohere2", + cohere_parser, ] if not is_reasoning_model: cmd.append("--no-cohere-is-reasoning-model") @@ -118,6 +129,7 @@ def start_server( port: int, is_reasoning_model: bool, extra_args: list[str], + cohere_parser: str = DEFAULT_COHERE_PARSER, log_prefix: str = "vllm-cohere-e2e-", ) -> ManagedServer: cmd = build_serve_command( @@ -126,6 +138,7 @@ def start_server( port=port, is_reasoning_model=is_reasoning_model, extra_args=extra_args, + cohere_parser=cohere_parser, ) env = os.environ.copy() env["VLLM_ENABLE_COHERE_API"] = "1" @@ -217,6 +230,17 @@ def add_server_args( "when this script started the server itself." ), ) + parser.add_argument( + "--cohere-parser", + default=DEFAULT_COHERE_PARSER, + help=( + "Parser name passed to both --tool-call-parser and " + "--reasoning-parser when this script starts the server. Must " + "be registered in vllm.tool_parsers / vllm.reasoning: " + "cohere_command3 or cohere_command4 " + "(default: %(default)s)." + ), + ) parser.add_argument( "--startup-timeout", type=float, @@ -270,6 +294,7 @@ def ensure_server( port=parsed.port or 8000, is_reasoning_model=args.is_reasoning_model, extra_args=args.extra_server_args, + cohere_parser=args.cohere_parser, log_prefix=log_prefix, ) try: diff --git a/vllm/renderers/cohere.py b/vllm/renderers/cohere.py index dbe5bf0ae5d9..cdbaca613ab1 100644 --- a/vllm/renderers/cohere.py +++ b/vllm/renderers/cohere.py @@ -80,7 +80,8 @@ Cohere-scoped ``CohereChatMessage.citations`` / ``CohereDeltaMessage.citations`` fields (see :mod:`vllm.entrypoints.cohere.cohere_chat_message`), populated by the -``cohere2`` reasoning parser. The base OpenAI ``ChatMessage`` / +``cohere_command3`` / ``cohere_command4`` reasoning parsers. The base +OpenAI ``ChatMessage`` / ``DeltaMessage`` keep their declared schemas unchanged; the response envelope declares them as ``SerializeAsAny[...]`` so the subclass fields survive JSON serialization. From e8845293751e70ef140b1535641b2fd7a2f4c8d8 Mon Sep 17 00:00:00 2001 From: Andrew Berneshawi Date: Wed, 26 Aug 2026 16:10:12 -0400 Subject: [PATCH 6/6] Change citation mode to arg --- scripts/cohere_e2e/inspect_cohere_renderer.py | 23 ++++++++++++++++++- 1 file changed, 22 insertions(+), 1 deletion(-) diff --git a/scripts/cohere_e2e/inspect_cohere_renderer.py b/scripts/cohere_e2e/inspect_cohere_renderer.py index a8fae4584f0d..9b094f3e21e2 100644 --- a/scripts/cohere_e2e/inspect_cohere_renderer.py +++ b/scripts/cohere_e2e/inspect_cohere_renderer.py @@ -119,6 +119,19 @@ def _parse_args() -> argparse.Namespace: choices=[None, "enabled", "disabled"], help="Reasoning toggle forwarded via chat_template_kwargs.", ) + p.add_argument( + "--citation-mode", + default=None, + choices=[None, "enabled", "disabled", "fast", "accurate", "off"], + help=( + "citation_options.mode forwarded via chat_template_kwargs. " + "'enabled'/'disabled' are the current vocabulary; " + "'fast'/'accurate'/'off' are the legacy values, kept here so " + "you can inspect how they normalize (cmd3 -> citation_quality " + "on/off, cmd4 -> grounding enabled/disabled). Omit to leave " + "citation_options unset." + ), + ) p.add_argument( "--show-token-ids", action="store_true", @@ -196,7 +209,8 @@ def _build_v2_request(args: argparse.Namespace) -> CohereChatV2Request: body["safety_mode"] = args.safety_mode.upper() if args.reasoning: body["thinking"] = {"type": args.reasoning} - body["citation_options"] = {"mode": "FAST"} + if args.citation_mode: + body["citation_options"] = {"mode": args.citation_mode.upper()} return CohereChatV2Request.model_validate(body) @@ -253,6 +267,13 @@ async def main() -> int: chat_template_kwargs["safety_mode"] = args.safety_mode if args.reasoning is not None: chat_template_kwargs["reasoning_type"] = args.reasoning + if args.citation_mode is not None: + # The renderer derives cmd3's ``citation_quality`` and cmd4's + # ``grounding`` from this same key, so one flag covers both + # template families. + chat_template_kwargs["citation_options"] = { + "mode": args.citation_mode.upper() + } params = ChatParams(chat_template_kwargs=chat_template_kwargs)