@tduncan / Nvidia-Optimization-Stack-VLLM / commits / f94dcde5c5

Fix Cohere ChatV2 citation and tool handling issues (#52175)

Signed-off-by: Andrew Berneshawi <andrewberneshawi@cohere.com>

andrewbcohere committed Aug 21, 2026 at 18:23 UTC f94dcde5c53b1930f98776e3793dd460d237d054
7 files changed +1024 -68
tests/entrypoints/cohere/test_serving_conversion.py
+418 -11
@@ -11,6 +11,7 @@ a real engine client, model registry, etc.) — the same pattern used in
11 ``test_serving_streaming.py``.
12 """
13
14 +import json
15 from typing import Any
16
17 import pytest
@@ -34,11 +35,61 @@ from vllm.entrypoints.openai.chat_completion.protocol import (
35 ChatCompletionRequest,
36 ChatCompletionResponse,
37 )
38 +from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat
39 +from vllm.renderers.cohere import (
40 + MESSAGES_CITATIONS_KEY,
41 + POSITION_TO_SOURCE_KEY,
42 + TOOL_MESSAGE_V2_CONTENT_KEY,
43 +)
44 +from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder
45
46 # ----------------------------------------------------------------------
47 # Helpers
48 # ----------------------------------------------------------------------
49
50 +# A full grounded tool round-trip: a tool call, a tool result carrying both
51 +# a document and a bare text block, and an assistant turn citing it. Between
52 +# them these populate every reserved ``_``-prefixed chat_template_kwargs key.
53 +_TOOL_ROUND_TRIP_MESSAGES: list[dict[str, Any]] = [
54 + {"role": "user", "content": "What is the capital?"},
55 + {
56 + "role": "assistant",
57 + "tool_calls": [
58 + {
59 + "id": "call-1",
60 + "type": "function",
61 + "function": {"name": "search", "arguments": '{"q": "capital"}'},
62 + }
63 + ],
64 + },
65 + {
66 + "role": "tool",
67 + "tool_call_id": "call-1",
68 + "content": [
69 + {
70 + "type": "document",
71 + "document": {"id": "doc-9", "data": {"text": "Paris"}},
72 + },
73 + {"type": "text", "text": "Paris is the capital."},
74 + ],
75 + },
76 + {
77 + "role": "assistant",
78 + "content": "Paris.",
79 + "citations": [
80 + {
81 + "start": 0,
82 + "end": 5,
83 + "text": "Paris",
84 + "sources": [
85 + {"type": "tool", "id": "doc-9", "tool_output": {"text": "Paris"}}
86 + ],
87 + }
88 + ],
89 + },
90 + {"role": "user", "content": "Are you sure?"},
91 +]
92 +
93
94 def _make_request(**kwargs) -> CohereChatV2Request:
95 kwargs.setdefault("model", "m")
@@ -447,9 +498,15 @@ class TestConvertToolMessage:
498 "content": "result text",
499 }
500
450 - def test_text_only_list_flattened_to_newline_string(self):
451 - # Text-only tool results are flattened to a single newline-joined
452 - # string for compatibility with standard chat templates.
501 + def test_text_only_list_preserves_per_block_text_parts(self):
502 + # A list of text blocks maps 1-to-1 onto ``text`` content
503 + # parts, mirroring the shape ``_convert_user_message`` emits
504 + # for multi-part user content. vLLM's
505 + # ``_parse_chat_message_content`` will collapse this back to
506 + # ``"\n".join(...)`` for the ``ConversationMessage`` handed to
507 + # renderers (see the tool-role branch in
508 + # ``vllm/entrypoints/chat_utils.py``), so downstream chat
509 + # templates that expect string tool content still work.
510 req = self._request_with_tool_message(
511 [
512 {"type": "text", "text": "line 1"},
@@ -458,11 +515,27 @@ class TestConvertToolMessage:
515 )
516 result = _convert(req)
517 tool_msg = result.messages[-1]
461 - assert tool_msg["content"] == "line 1\nline 2"
518 + assert tool_msg["content"] == [
519 + {"type": "text", "text": "line 1"},
520 + {"type": "text", "text": "line 2"},
521 + ]
522
463 - def test_with_document_preserves_structured_content(self):
464 - # When documents appear in the tool result, we keep the list shape
465 - # so the Cohere renderer can lift them into grounding sources.
523 + def test_document_blocks_serialized_as_text_parts(self):
524 + """Document blocks must not survive on the OpenAI-shaped tool
525 + message.
526 +
527 + vLLM's shared ``_parse_chat_message_content_parts`` runs a
528 + Pydantic ``ValidatorIterator`` over the OpenAI content-part
529 + union, which only knows ``text`` / ``image_url`` / etc. and
530 + rejects ``{"type": "document", ...}`` blocks outright. We
531 + therefore JSON-serialize each document into its own ``text``
532 + content part so the request survives OpenAI validation; the
533 + cohere renderer still resolves outbound citations against
534 + the original cohere v2 messages via ``POSITION_TO_SOURCE_KEY``
535 + and reconstructs structured melody ``document`` blocks from the
536 + original v2 content forwarded via
537 + ``TOOL_MESSAGE_V2_CONTENT_KEY``.
538 + """
539 req = self._request_with_tool_message(
540 [
541 {"type": "text", "text": "see attachment"},
@@ -475,11 +548,75 @@ class TestConvertToolMessage:
548 result = _convert(req)
549 tool_msg = result.messages[-1]
550 assert isinstance(tool_msg["content"], list)
551 + assert len(tool_msg["content"]) == 2
552 + # First block passes through unchanged.
553 assert tool_msg["content"][0] == {"type": "text", "text": "see attachment"}
479 - assert tool_msg["content"][1] == {
480 - "type": "document",
481 - "document": {"data": {"text": "doc text"}, "id": "d1"},
482 - }
554 + # Second block: document JSON-serialized into a ``text`` part.
555 + second = tool_msg["content"][1]
556 + assert second["type"] == "text"
557 + assert json.loads(second["text"]) == {"data": {"text": "doc text"}, "id": "d1"}
558 +
559 + def test_document_only_content_serialized_as_text_part(self):
560 + """Document-only tool results must also produce a text-typed
561 + content part (no ``text`` block above them), so a single-doc
562 + result still passes the OpenAI content-parts validator.
563 + """
564 + req = self._request_with_tool_message(
565 + [
566 + {
567 + "type": "document",
568 + "document": {"data": {"result": "K2 is 8611m."}},
569 + }
570 + ]
571 + )
572 + result = _convert(req)
573 + tool_msg = result.messages[-1]
574 + assert isinstance(tool_msg["content"], list)
575 + (part,) = tool_msg["content"]
576 + assert part["type"] == "text"
577 + assert json.loads(part["text"]) == {"data": {"result": "K2 is 8611m."}}
578 +
579 + def test_converted_tool_message_content_survives_openai_validation(self):
580 + """Regression for the reported crash.
581 +
582 + ``vllm/entrypoints/chat_utils.py:_parse_chat_message_content_parts``
583 + runs a Pydantic ``ValidatorIterator`` over
584 + ``list[ChatCompletionContentPartParam]``, which is a
585 + discriminated union that only accepts ``text`` / ``image_url``
586 + / ``input_audio`` / ``refusal`` / ``file`` types. Previously
587 + the converted tool message carried ``{"type": "document",
588 + "document": {...}}`` blocks, which the validator rejected with:
589 +
590 + 2 validation errors for ValidatorIterator
591 + 0.text Field required
592 + 0.type Input should be 'text', got 'document'
593 +
594 + The fix rewrites every ``DocumentToolContent`` into a
595 + JSON-serialized ``text`` part. This test asserts the invariant
596 + that keeps the validator happy: on the OpenAI-shaped request,
597 + every content entry a tool message carries must be a ``text``
598 + part (or a bare string) -- so no future refactor accidentally
599 + re-introduces a ``document``-typed content part.
600 + """
601 + req = self._request_with_tool_message(
602 + [
603 + {
604 + "type": "document",
605 + "document": {
606 + "data": {"result": "The second tallest mountain is K2."}
607 + },
608 + }
609 + ]
610 + )
611 + result = _convert(req)
612 + tool_msg = result.messages[-1]
613 +
614 + content = tool_msg["content"]
615 + assert isinstance(content, (str, list))
616 + if isinstance(content, list):
617 + for part in content:
618 + assert isinstance(part, dict)
619 + assert part.get("type") == "text"
620
621
622 # ======================================================================
@@ -808,6 +945,276 @@ class TestApplyCohereTemplateKwargs:
945 assert result.chat_template_kwargs is None
946
947
948 +# ======================================================================
949 +# Engine-core hop (_engine_chat_template_kwargs)
950 +# ======================================================================
951 +
952 +
953 +class TestEngineChatTemplateKwargs:
954 + """The internal ``_``-prefixed entries must not reach the engine core.
955 +
956 + ``OpenAIServingChat`` uses one ``chat_template_kwargs`` dict twice:
957 + to build the API-server-side parser, and to populate
958 + ``EngineCoreRequest.reasoning_parser_kwargs``, which crosses ZMQ as
959 + msgpack. This handler's reserved keys are only meaningful to the
960 + former, and ``_position_to_source`` holds ``CitationSource`` models
961 + that msgpack cannot encode at all -- forwarding them fails the
962 + request with a ``TypeError``. The
963 + ``_engine_chat_template_kwargs`` override drops them.
964 + """
965 +
966 + @staticmethod
967 + def _engine_kwargs(request: CohereChatV2Request) -> dict[str, Any]:
968 + full = _convert(request).chat_template_kwargs
969 + assert full, "expected the request to populate chat_template_kwargs"
970 + return _FakeServing()._engine_chat_template_kwargs(full)
971 +
972 + def test_grounded_request_drops_internal_keys_keeps_public_ones(self):
973 + request = _make_request(
974 + documents=[{"id": "d1", "data": {"text": "hi"}}],
975 + citation_options={"mode": "accurate"},
976 + )
977 + full = _convert(request).chat_template_kwargs
978 + # Sanity: the key we expect to be dropped is actually present.
979 + assert POSITION_TO_SOURCE_KEY in full
980 +
981 + engine_kwargs = self._engine_kwargs(request)
982 + assert POSITION_TO_SOURCE_KEY not in engine_kwargs
983 + # Public template inputs the engine-side parser may legitimately
984 + # read must survive.
985 + assert engine_kwargs["citation_options"] == {"mode": "accurate"}
986 + assert engine_kwargs["documents"] == [{"id": "d1", "data": {"text": "hi"}}]
987 +
988 + def test_all_internal_keys_dropped(self):
989 + """Exercises all three reserved keys at once: the tool-result
990 + position map, the inbound assistant citations, and the raw v2 tool
991 + content blocks.
992 + """
993 + request = _make_request(messages=_TOOL_ROUND_TRIP_MESSAGES)
994 + full = _convert(request).chat_template_kwargs
995 + internal = {k for k in full if k.startswith("_")}
996 + # Sanity: this request really does populate all three.
997 + assert internal == {
998 + MESSAGES_CITATIONS_KEY,
999 + POSITION_TO_SOURCE_KEY,
1000 + TOOL_MESSAGE_V2_CONTENT_KEY,
1001 + }
1002 +
1003 + engine_kwargs = self._engine_kwargs(request)
1004 + assert not [k for k in engine_kwargs if k.startswith("_")]
1005 +
1006 + def test_engine_kwargs_are_msgpack_encodable(self):
1007 + """The real failure mode: the full dict raises, the filtered one
1008 + doesn't. Pins that the filter is what makes the engine hop work.
1009 + """
1010 + request = _make_request(messages=_TOOL_ROUND_TRIP_MESSAGES)
1011 + full = _convert(request).chat_template_kwargs
1012 +
1013 + with pytest.raises(TypeError, match="is not serializable"):
1014 + MsgpackEncoder().encode({"chat_template_kwargs": full})
1015 +
1016 + engine_kwargs = self._engine_kwargs(request)
1017 + encoder = MsgpackEncoder()
1018 + decoder = MsgpackDecoder(dict[str, Any])
1019 + round_tripped = decoder.decode(
1020 + encoder.encode({"chat_template_kwargs": engine_kwargs})[0]
1021 + )
1022 + assert round_tripped["chat_template_kwargs"] == engine_kwargs
1023 +
1024 + def test_original_dict_not_mutated(self):
1025 + """The caller still passes the full dict to the in-process parser,
1026 + so the filter must return a copy.
1027 + """
1028 + request = _make_request(documents=[{"id": "d1", "data": {"text": "hi"}}])
1029 + full = _convert(request).chat_template_kwargs
1030 + before = set(full)
1031 +
1032 + self._engine_kwargs(request)
1033 + assert set(full) == before
1034 +
1035 + def test_base_implementation_is_a_passthrough(self):
1036 + """Only the Cohere subclass filters; plain chat completions must
1037 + keep forwarding whatever the request asked for.
1038 + """
1039 + kwargs = {
1040 + "documents": [{"id": "d1"}],
1041 + POSITION_TO_SOURCE_KEY: {(0, 0): object()},
1042 + }
1043 + passed_through = OpenAIServingChat._engine_chat_template_kwargs(
1044 + _FakeServing(), kwargs
1045 + )
1046 + assert passed_through == kwargs
1047 +
1048 +
1049 +# ======================================================================
1050 +# Tool-message content forwarding
1051 +# (chat_template_kwargs["_tool_message_v2_content"])
1052 +# ======================================================================
1053 +
1054 +
1055 +class TestToolMessageV2ContentForwarding:
1056 + """The cohere renderer needs the original v2 tool-message content
1057 + (text / document blocks) to reconstruct melody's ``document``-block
1058 + shape, but the OpenAI content-parts validator on the way in only
1059 + accepts text / image / audio / etc. blocks. So
1060 + :meth:`CohereServingChatV2._convert_tool_message` rewrites each
1061 + content block into a plain ``text`` part on the request that
1062 + survives validation, and
1063 + :meth:`CohereServingChatV2._collect_v2_tool_message_content`
1064 + forwards the *original* v2 content blocks -- as ``model_dump``'d
1065 + dicts -- to the renderer under
1066 + :data:`TOOL_MESSAGE_V2_CONTENT_KEY`. These tests pin the
1067 + forwarding contract (presence, dict-shape, index-alignment); the
1068 + v2 → melody mapping the renderer applies to these blocks lives in
1069 + ``tests/renderers/test_cohere.py``.
1070 + """
1071 +
1072 + def _request_with_tool_message(self, content: Any) -> CohereChatV2Request:
1073 + return _make_request(
1074 + messages=[
1075 + {"role": "user", "content": "hi"},
1076 + {
1077 + "role": "assistant",
1078 + "tool_calls": [
1079 + {
1080 + "id": "c1",
1081 + "type": "function",
1082 + "function": {"name": "f", "arguments": "{}"},
1083 + }
1084 + ],
1085 + },
1086 + {"role": "tool", "tool_call_id": "c1", "content": content},
1087 + ]
1088 + )
1089 +
1090 + def test_absent_when_no_tool_messages(self):
1091 + result = _convert(_make_request())
1092 + # ``chat_template_kwargs`` may be ``None`` here (no cohere-side
1093 + # forwarding happened at all) or a dict without the key; both
1094 + # mean the same thing to the renderer.
1095 + assert not (result.chat_template_kwargs or {}).get(TOOL_MESSAGE_V2_CONTENT_KEY)
1096 +
1097 + def test_absent_for_string_tool_content(self):
1098 + # Bare-string tool content works through the renderer's default
1099 + # ``_content_blocks`` path, so no forwarded blocks are needed.
1100 + req = self._request_with_tool_message("plain text")
1101 + result = _convert(req)
1102 + assert not (result.chat_template_kwargs or {}).get(TOOL_MESSAGE_V2_CONTENT_KEY)
1103 +
1104 + def test_forwarded_blocks_are_plain_v2_shape_dicts(self):
1105 + # The forwarded map must contain the original v2 content
1106 + # blocks verbatim (as plain dicts, not Pydantic instances) so
1107 + # the renderer can drive the v2 → melody mapping and the map
1108 + # stays trivially serializable.
1109 + req = self._request_with_tool_message(
1110 + [
1111 + {"type": "text", "text": "K2 is 8611m."},
1112 + {
1113 + "type": "document",
1114 + "document": {
1115 + "id": "d1",
1116 + "data": {"result": "The second tallest mountain is K2."},
1117 + },
1118 + },
1119 + ]
1120 + )
1121 + blocks_by_index = _forwarded_v2_content(req)
1122 + # Index 2 = the tool message (user=0, assistant=1, tool=2).
1123 + assert list(blocks_by_index.keys()) == [2]
1124 + assert blocks_by_index[2] == [
1125 + {"type": "text", "text": "K2 is 8611m."},
1126 + {
1127 + "type": "document",
1128 + "document": {
1129 + "id": "d1",
1130 + "data": {"result": "The second tallest mountain is K2."},
1131 + },
1132 + },
1133 + ]
1134 + # Guard against a regression where Pydantic instances leak into
1135 + # the forwarded map -- msgpack can't serialize them, and even
1136 + # for keys filtered out of the engine-core payload we want the
1137 + # renderer to see plain data.
1138 + for block in blocks_by_index[2]:
1139 + assert isinstance(block, dict)
1140 +
1141 + def test_forwarded_blocks_preserve_input_order(self):
1142 + # Melody's tool-result rendering keys result slots by their
1143 + # position within a tool message's content list, and the
1144 + # outbound citation source map (built by
1145 + # ``_walk_citable_positions``) uses that same positional index.
1146 + # Reordering would drop citations pointing at slot ``1`` onto
1147 + # the wrong payload.
1148 + req = self._request_with_tool_message(
1149 + [
1150 + {"type": "text", "text": "leading note"},
1151 + {
1152 + "type": "document",
1153 + "document": {"data": {"body": "doc body"}, "id": "d0"},
1154 + },
1155 + {"type": "text", "text": '{"parsed": "yes"}'},
1156 + ]
1157 + )
1158 + blocks = _forwarded_v2_content(req)[2]
1159 + assert [b["type"] for b in blocks] == ["text", "document", "text"]
1160 +
1161 + def test_index_matches_original_request_position(self):
1162 + # Entries are keyed by *request-message* index so the renderer,
1163 + # walking the OpenAI-shaped conversation with the same
1164 + # enumeration, resolves the right slot. If ``_convert_messages``
1165 + # ever splits or merges messages this invariant breaks, so pin
1166 + # the current 1-to-1 behavior with a two-tool-message request.
1167 + req = _make_request(
1168 + messages=[
1169 + {"role": "user", "content": "hi"},
1170 + {
1171 + "role": "assistant",
1172 + "tool_calls": [
1173 + {
1174 + "id": "c1",
1175 + "type": "function",
1176 + "function": {"name": "f", "arguments": "{}"},
1177 + },
1178 + {
1179 + "id": "c2",
1180 + "type": "function",
1181 + "function": {"name": "g", "arguments": "{}"},
1182 + },
1183 + ],
1184 + },
1185 + {
1186 + "role": "tool",
1187 + "tool_call_id": "c1",
1188 + "content": [{"type": "text", "text": "first result"}],
1189 + },
1190 + {
1191 + "role": "tool",
1192 + "tool_call_id": "c2",
1193 + "content": [{"type": "text", "text": "second result"}],
1194 + },
1195 + ]
1196 + )
1197 + blocks_by_index = _forwarded_v2_content(req)
1198 + assert sorted(blocks_by_index.keys()) == [2, 3]
1199 + assert blocks_by_index[2] == [{"type": "text", "text": "first result"}]
1200 + assert blocks_by_index[3] == [{"type": "text", "text": "second result"}]
1201 +
1202 +
1203 +def _forwarded_v2_content(
1204 + request: CohereChatV2Request,
1205 +) -> dict[int, list[dict[str, Any]]]:
1206 + """Extract the forwarded v2 tool-message content from a converted
1207 + request, or fail the test if the key isn't present.
1208 + """
1209 + result = _convert(request)
1210 + kwargs = result.chat_template_kwargs or {}
1211 + assert TOOL_MESSAGE_V2_CONTENT_KEY in kwargs, (
1212 + f"expected {TOOL_MESSAGE_V2_CONTENT_KEY} in chat_template_kwargs, "
1213 + f"got keys {list(kwargs)}"
1214 + )
1215 + return kwargs[TOOL_MESSAGE_V2_CONTENT_KEY]
1216 +
1217 +
1218 # ======================================================================
1219 # Message-citation forwarding (chat_template_kwargs["_messages_citations"])
1220 # ======================================================================
tests/reasoning/test_cohere_command_reasoning_parser.py
+120
@@ -655,3 +655,123 @@ class TestMelodySourceResolution:
655 assert out is not None
656 assert out[0].type == "THINKING_CONTENT"
657 assert out[0].sources[0].id == "d0"
658 +
659 +
660 +class TestParserStreamingEndToEnd:
661 + """End-to-end streaming shape: raw model output text is tokenized
662 + with the mock byte-level tokenizer and fed one token at a time
663 + through :meth:`CohereCommand3ReasoningParser.extract_reasoning_streaming`.
664 + We collect the ordered sequence of ``DeltaMessage``s the parser
665 + hands back -- i.e. what the client would receive on the wire --
666 + and pin the exact shape.
667 +
668 + Combines the three flavors that reach ``delta``: a thinking block
669 + (populates ``reasoning``), a text block (populates ``content``),
670 + and a resolved citation (populates ``citations``) referencing
671 + ``(start, end)`` offsets into the already-streamed content.
672 + """
673 +
674 + _MODEL_OUTPUT = (
675 + "<|START_THINKING|>Let me check.<|END_THINKING|>"
676 + "<|START_RESPONSE|>The capital is <co>Paris</co: 0:[1]>.<|END_RESPONSE|>"
677 + )
678 +
679 + @staticmethod
680 + def _drive_parser(
681 + parser: CohereCommand3ReasoningParser, text: str, tok: MockCohereTokenizer
682 + ) -> list[Any]:
683 + """Feed ``text`` through the parser one token at a time
684 + (byte-level) and return the ordered list of non-``None``
685 + ``DeltaMessage``s.
686 + """
687 + deltas: list[Any] = []
688 + previous_text = ""
689 + previous_token_ids: list[int] = []
690 + for token_str in _token_deltas(tok, text):
691 + current_text = previous_text + token_str
692 + current_token_ids = previous_token_ids + [0]
693 + delta = parser.extract_reasoning_streaming(
694 + previous_text=previous_text,
695 + current_text=current_text,
696 + delta_text=token_str,
697 + previous_token_ids=previous_token_ids,
698 + current_token_ids=current_token_ids,
699 + delta_token_ids=[0],
700 + )
701 + if delta is not None:
702 + deltas.append(delta)
703 + previous_text = current_text
704 + previous_token_ids = current_token_ids
705 + return deltas
706 +
707 + def test_full_wire_stream_shape(self, tokenizer):
708 + from vllm.entrypoints.cohere.cohere_chat_message import CitationSource
709 + from vllm.renderers.cohere import POSITION_TO_SOURCE_KEY
710 +
711 + # Mimic what ``CohereServingChatV2._apply_cohere_template_kwargs``
712 + # installs on a request whose tool_call_index=0 tool result at
713 + # index 1 is the "France" document.
714 + position_to_source = {
715 + (0, 1): CitationSource(
716 + type="document",
717 + id="doc-paris",
718 + document={"id": "doc-paris", "title": "France"},
719 + ),
720 + }
721 +
722 + parser = CohereCommand3ReasoningParser(
723 + tokenizer,
724 + chat_template_kwargs={POSITION_TO_SOURCE_KEY: position_to_source},
725 + )
726 +
727 + deltas = self._drive_parser(parser, self._MODEL_OUTPUT, tokenizer)
728 +
729 + # -- 1. Reasoning and content stream separately and reach
730 + # the wire only as their real payload. Framing tokens
731 + # (``<|START_THINKING|>`` etc.) are consumed by the parser's
732 + # state machine and never produce a delta.
733 + reasoning_stream = [d.reasoning for d in deltas if d.reasoning is not None]
734 + content_stream = [d.content for d in deltas if d.content is not None]
735 + citation_deltas = [d for d in deltas if getattr(d, "citations", None)]
736 +
737 + assert "".join(reasoning_stream) == "Let me check."
738 + assert "".join(content_stream) == "The capital is Paris."
739 +
740 + # -- 2. Exactly one citation reaches the wire, resolved to the
741 + # source from ``position_to_source`` (not the raw
742 + # ``(bucket, idx)`` coordinates).
743 + assert len(citation_deltas) == 1
744 + (cite_delta,) = citation_deltas
745 + assert cite_delta.citations is not None
746 + assert len(cite_delta.citations) == 1
747 + cite = cite_delta.citations[0]
748 + assert cite.text == "Paris"
749 + # Char offsets are into the accumulated ``content``: "The
750 + # capital is " is 15 chars, "Paris" is 5.
751 + assert cite.start == 15
752 + assert cite.end == 20
753 + assert len(cite.sources) == 1
754 + assert cite.sources[0].id == "doc-paris"
755 + assert cite.sources[0].type == "document"
756 +
757 + # -- 3. The citation's ``(start, end)`` offsets point into
758 + # bytes that were already emitted as content-deltas earlier
759 + # in the stream -- i.e. the citation is anchored to real
760 + # already-shipped content.
761 + cite_delta_idx = deltas.index(cite_delta)
762 + earlier_content = "".join(
763 + d.content for d in deltas[:cite_delta_idx] if d.content is not None
764 + )
765 + assert earlier_content[cite.start : cite.end] == "Paris"
766 +
767 + # Everything after the citation delta is the trailing period.
768 + later_content = "".join(
769 + d.content for d in deltas[cite_delta_idx + 1 :] if d.content is not None
770 + )
771 + assert later_content == "."
772 +
773 + # -- 4. Reasoning and content are never mixed on the same
774 + # delta: the parser flips modes on the ``<|END_THINKING|>``
775 + # boundary and never emits a delta with both fields set.
776 + for d in deltas:
777 + assert not (d.reasoning is not None and d.content is not None)
tests/renderers/test_cohere.py
+202
@@ -30,6 +30,7 @@ from vllm.renderers.cohere import (
30 _normalize_tool_call,
31 _role_to_melody,
32 _tool_to_melody,
33 + _v2_tool_content_to_melody_block,
34 )
35 from vllm.tokenizers.hf import HfTokenizer
36
@@ -304,6 +305,108 @@ class TestToolToMelody:
305 _tool_to_melody(42) # type: ignore[arg-type]
306
307
308 +# ======================================================================
309 +# _v2_tool_content_to_melody_block
310 +# ======================================================================
311 +
312 +
313 +class TestV2ToolContentToMelodyBlock:
314 + """Pins the v2 → melody mapping applied to every tool-message
315 + content block. The serving layer forwards the original v2 blocks
316 + verbatim under :data:`TOOL_MESSAGE_V2_CONTENT_KEY`; the renderer
317 + owns this mapping so cmd3 / cmd4 render each result field as its
318 + own entry rather than a single stringified-JSON blob.
319 + """
320 +
321 + def test_document_block_unwraps_data_into_melody_document(self):
322 + # ``document.data`` becomes the melody ``document`` payload
323 + # directly. Keeping the wrapping ``data`` key would produce a
324 + # redundant ``{"data": {...}}`` nesting at render time.
325 + # ``document.id`` is dropped here on purpose: outbound
326 + # citation binding tracks IDs through the separate
327 + # ``POSITION_TO_SOURCE_KEY`` map.
328 + result = _v2_tool_content_to_melody_block(
329 + {
330 + "type": "document",
331 + "document": {
332 + "id": "d1",
333 + "data": {"result": "The second tallest mountain is K2."},
334 + },
335 + }
336 + )
337 + assert result == {
338 + "type": "document",
339 + "document": {"result": "The second tallest mountain is K2."},
340 + }
341 +
342 + def test_document_block_with_missing_data_yields_empty_payload(self):
343 + # A document block whose ``data`` isn't a mapping (e.g.
344 + # missing entirely) still produces a document block so
345 + # per-slot numbering stays aligned with the outbound citation
346 + # walk in ``CohereServingChatV2._walk_citable_positions``.
347 + result = _v2_tool_content_to_melody_block(
348 + {"type": "document", "document": {"id": "d1"}}
349 + )
350 + assert result == {"type": "document", "document": {}}
351 +
352 + def test_json_parseable_text_becomes_document_dict(self):
353 + # A ``text`` block whose payload parses to a JSON object lands
354 + # directly under ``document`` (no ``content`` wrapper) so
355 + # cmd3/cmd4 render each field individually.
356 + result = _v2_tool_content_to_melody_block(
357 + {"type": "text", "text": '{"temperature": "20C"}'}
358 + )
359 + assert result == {"type": "document", "document": {"temperature": "20C"}}
360 +
361 + def test_non_json_text_wrapped_under_content_key(self):
362 + # Plain-text tool output gets wrapped under a ``content`` key
363 + # so cmd3/cmd4 templates have a stable field to look up.
364 + result = _v2_tool_content_to_melody_block(
365 + {"type": "text", "text": "K2 is 8611m."}
366 + )
367 + assert result == {"type": "document", "document": {"content": "K2 is 8611m."}}
368 +
369 + def test_json_text_that_is_not_a_dict_is_treated_as_plain_text(self):
370 + # A ``text`` block whose payload parses as a JSON *array* or
371 + # *scalar* isn't a valid melody document (which must be a
372 + # mapping). Fall back to the ``{"content": <text>}`` wrapper so
373 + # the raw text still reaches the prompt.
374 + result = _v2_tool_content_to_melody_block({"type": "text", "text": "[1, 2, 3]"})
375 + assert result == {"type": "document", "document": {"content": "[1, 2, 3]"}}
376 +
377 + def test_empty_text_still_produces_content_block(self):
378 + # Even an empty text block must produce an entry so melody's
379 + # per-content-block numbering stays aligned with the
380 + # citation-side walk in ``_walk_citable_positions``.
381 + result = _v2_tool_content_to_melody_block({"type": "text", "text": ""})
382 + assert result == {"type": "document", "document": {"content": ""}}
383 +
384 + def test_missing_text_field_treated_as_empty(self):
385 + # ``model_dump(exclude_none=True)`` on a ``TextToolContent``
386 + # with an explicitly empty text still emits ``text: ""``, but
387 + # guard against a future v2 shape that omits it entirely.
388 + result = _v2_tool_content_to_melody_block({"type": "text"})
389 + assert result == {"type": "document", "document": {"content": ""}}
390 +
391 + def test_unknown_block_type_returns_none(self):
392 + # Callers use ``None`` to know they should skip the slot; a
393 + # future v2 addition shouldn't silently render as an empty
394 + # document.
395 + assert _v2_tool_content_to_melody_block({"type": "audio"}) is None
396 + assert _v2_tool_content_to_melody_block({"not": "a block"}) is None
397 +
398 + def test_document_block_with_non_dict_document_returns_none(self):
399 + # Malformed ``document`` field -- e.g. someone forwarding a
400 + # raw string rather than the ``{id, data}`` wrapper -- is
401 + # unusable, so return ``None`` so the caller skips it.
402 + assert (
403 + _v2_tool_content_to_melody_block(
404 + {"type": "document", "document": "not a dict"}
405 + )
406 + is None
407 + )
408 +
409 +
410 # ======================================================================
411 # _conversation_to_melody_messages
412 # ======================================================================
@@ -444,6 +547,71 @@ class TestConversationToMelody:
547 out = _conversation_to_melody_messages(conv, {5: [{"anything": 1}]}) # type: ignore[arg-type]
548 assert "citations" not in out[0]
549
550 + def test_tool_message_v2_content_replaces_default_content(self):
551 + # When ``tool_message_v2_content`` has an entry for a message
552 + # index, the renderer maps those v2 blocks through
553 + # :func:`_v2_tool_content_to_melody_block` and uses the result
554 + # instead of the block list ``_content_blocks`` would derive
555 + # from ``msg["content"]``.
556 + conv = [
557 + {"role": "user", "content": "q"},
558 + {
559 + "role": "tool",
560 + "tool_call_id": "c1",
561 + # ``_convert_tool_message`` on the serving side flattens
562 + # the original ``document`` block to this text part.
563 + # The renderer should ignore this flattened content
564 + # because the forwarded map carries the structural
565 + # truth.
566 + "content": '[{"type":"text","text":"{\\"result\\":\\"K2\\"}"}]',
567 + },
568 + ]
569 + v2_content = {
570 + 1: [
571 + {
572 + "type": "document",
573 + "document": {"id": "d1", "data": {"result": "K2"}},
574 + }
575 + ]
576 + }
577 + out = _conversation_to_melody_messages(
578 + conv, # type: ignore[arg-type]
579 + tool_message_v2_content=v2_content,
580 + )
581 + # User message (no forwarded entry) still goes through the
582 + # default ``_content_blocks`` path.
583 + assert out[0]["content"] == [{"type": "text", "text": "q"}]
584 + # Tool message picks up the mapped melody document block.
585 + assert out[1]["content"] == [{"type": "document", "document": {"result": "K2"}}]
586 + assert out[1]["tool_call_id"] == "c1"
587 +
588 + def test_tool_message_v2_content_skips_unrecognized_blocks(self):
589 + # ``_v2_tool_content_to_melody_block`` returns ``None`` for
590 + # shapes it doesn't recognise; the renderer must skip those
591 + # rather than let a ``None`` reach melody.
592 + conv = [{"role": "tool", "tool_call_id": "c", "content": "ignored"}]
593 + v2_content = {
594 + 0: [
595 + {"type": "text", "text": "keep me"},
596 + {"type": "mystery", "unexpected": True},
597 + ]
598 + }
599 + out = _conversation_to_melody_messages(
600 + conv, # type: ignore[arg-type]
601 + tool_message_v2_content=v2_content,
602 + )
603 + assert out[0]["content"] == [
604 + {"type": "document", "document": {"content": "keep me"}}
605 + ]
606 +
607 + def test_tool_message_v2_content_none_falls_back_to_default(self):
608 + conv = [{"role": "tool", "tool_call_id": "c", "content": "plain text"}]
609 + out = _conversation_to_melody_messages(
610 + conv, # type: ignore[arg-type]
611 + tool_message_v2_content=None,
612 + )
613 + assert out[0]["content"] == [{"type": "text", "text": "plain text"}]
614 +
615
616 # ======================================================================
617 # _build_render_config
@@ -762,6 +930,40 @@ class TestBuildRenderConfig:
930 ) # type: ignore[arg-type]
931 assert "template_jinja" not in cfg
932
933 + def test_tool_message_v2_content_key_picked_up_from_kwargs(self):
934 + # ``_build_render_config`` must pull
935 + # ``TOOL_MESSAGE_V2_CONTENT_KEY`` out of
936 + # ``chat_template_kwargs`` and hand it to
937 + # ``_conversation_to_melody_messages`` so the renderer's
938 + # v2 → melody mapping runs over the original v2 blocks rather
939 + # than the flattened string content on the OpenAI-shape tool
940 + # message.
941 + from vllm.renderers.cohere import TOOL_MESSAGE_V2_CONTENT_KEY
942 +
943 + conv = [
944 + {"role": "user", "content": "q"},
945 + {"role": "tool", "tool_call_id": "c1", "content": '{"r":"k"}'},
946 + ]
947 + kwargs = {
948 + TOOL_MESSAGE_V2_CONTENT_KEY: {
949 + 1: [
950 + {
951 + "type": "document",
952 + "document": {"id": "d1", "data": {"r": "k"}},
953 + }
954 + ]
955 + }
956 + }
957 + _, cfg = _build_render_config(conv, kwargs) # type: ignore[arg-type]
958 + tool_msg = cfg["messages"][1]
959 + assert tool_msg["content"] == [{"type": "document", "document": {"r": "k"}}]
960 + # And the key must be *consumed* -- not surfaced as a stray
961 + # Jinja variable (``_RENDERER_CONSUMED_KEYS`` covers this but
962 + # pin it here so a removal from that set surfaces immediately).
963 + assert TOOL_MESSAGE_V2_CONTENT_KEY not in cfg.get(
964 + "additional_template_fields", {}
965 + )
966 +
967
968 # ======================================================================
969 # End-to-end async rendering (mirrors ``test_mistral.py``)
vllm/entrypoints/cohere/serving.py
+130 -31
@@ -89,7 +89,11 @@ from vllm.entrypoints.openai.models.serving import OpenAIServingModels
89 from vllm.entrypoints.serve.exception_handling.utils import sanitize_message
90 from vllm.entrypoints.serve.utils.request_logger import RequestLogger
91 from vllm.parser.abstract_parser import Parser
92 -from vllm.renderers.cohere import MESSAGES_CITATIONS_KEY, POSITION_TO_SOURCE_KEY
92 +from vllm.renderers.cohere import (
93 + MESSAGES_CITATIONS_KEY,
94 + POSITION_TO_SOURCE_KEY,
95 + TOOL_MESSAGE_V2_CONTENT_KEY,
96 +)
97
98 if TYPE_CHECKING:
99 from vllm.renderers.online_renderer import OnlineRenderer
@@ -97,6 +101,19 @@ if TYPE_CHECKING:
101 logger = logging.getLogger(__name__)
102
103
104 +# Reserved ``chat_template_kwargs`` entries this handler populates for
105 +# the API-server-side parser and renderer only. Stripped before the dict
106 +# is forwarded to the engine core -- see
107 +# :meth:`CohereServingChatV2._engine_chat_template_kwargs`.
108 +_API_SERVER_ONLY_TEMPLATE_KWARGS = frozenset(
109 + {
110 + MESSAGES_CITATIONS_KEY,
111 + POSITION_TO_SOURCE_KEY,
112 + TOOL_MESSAGE_V2_CONTENT_KEY,
113 + }
114 +)
115 +
116 +
117 # ---------------------------------------------------------------------------
118 # SSE helpers
119 # ---------------------------------------------------------------------------
@@ -304,6 +321,28 @@ class CohereServingChatV2(OpenAIServingChat):
321 case _:
322 return self._chat_completion_stream_to_v2(generator, request)
323
324 + def _engine_chat_template_kwargs(
325 + self, chat_template_kwargs: dict[str, Any]
326 + ) -> dict[str, Any]:
327 + """Keep this handler's internal entries out of the engine hop.
328 +
329 + Overrides :meth:`OpenAIServingChat._engine_chat_template_kwargs`.
330 + The reserved keys in :data:`_API_SERVER_ONLY_TEMPLATE_KWARGS`
331 + exist only to hand request-side state to the API-server-side
332 + parser and renderer; the engine-core-side reasoning parser (used
333 + for structured-output reasoning-end gating) never reads them. They
334 + also can't survive the trip: ``POSITION_TO_SOURCE_KEY`` holds
335 + :class:`CitationSource` models, and msgpack -- which encodes
336 + ``reasoning_parser_kwargs`` on the way to the engine core -- has
337 + no native encoding for Pydantic models, so forwarding them fails
338 + the request with a ``TypeError``.
339 + """
340 + return {
341 + k: v
342 + for k, v in chat_template_kwargs.items()
343 + if k not in _API_SERVER_ONLY_TEMPLATE_KWARGS
344 + }
345 +
346 # ==================================================================
347 # Request conversion: Cohere V2 -> ChatCompletionRequest
348 # ==================================================================
@@ -438,46 +477,53 @@ class CohereServingChatV2(OpenAIServingChat):
477 "content": msg.content,
478 }
479
441 - # When the tool result is text-only, flatten to a string for
442 - # maximum compatibility with standard chat templates. When it
443 - # includes documents, preserve them as structured content parts
444 - # so the cohere renderer can surface them as grounding sources
445 - # (the :class:`CohereRenderer` understands
446 - # ``{type: document, document: {...}}`` blocks). Non-cohere
447 - # renderers may not honor document blocks, but that matches
448 - # the broader "documents are no-op for OSS models" contract
449 - # documented on this endpoint.
480 + # Map every content block to a plain ``{"type": "text", "text":
481 + # ...}`` part, mirroring the shape :meth:`_convert_user_message`
482 + # emits for multi-part user content. vLLM's
483 + # ``_parse_chat_message_content`` collapses a text-only tool
484 + # part list back into ``"\n".join(...)`` on the way to the
485 + # renderer, so simple chat templates that expect string tool
486 + # content still work; keeping the list here just preserves the
487 + # per-block structure at the OpenAI-shape request level for
488 + # symmetry with the user-message conversion.
489 #
451 - # Tool message content uses ``ToolMessageV2Content``, which is a
452 - # union of ``TextToolContent`` and ``DocumentToolContent`` -
453 - # distinct from the user-message text/image union. We discriminate
454 - # on the ``type`` literal so we don't have to import each variant.
455 - has_documents = any(block.type == "document" for block in msg.content)
456 - if not has_documents:
457 - text = "\n".join(
458 - block.text for block in msg.content if block.type == "text"
459 - )
460 - return {
461 - "role": "tool",
462 - "tool_call_id": msg.tool_call_id,
463 - "content": text,
464 - }
465 -
466 - parts: list[dict[str, Any]] = []
490 + # Tool message content uses ``ToolMessageV2Content``, a union of
491 + # ``TextToolContent`` and ``DocumentToolContent``. The latter
492 + # produces ``{"type": "document", "document": {...}}`` blocks,
493 + # which vLLM's shared ``_parse_chat_message_content_parts`` in
494 + # ``vllm/entrypoints/chat_utils.py`` rejects because the OpenAI
495 + # ``ChatCompletionContentPartParam`` union only knows text /
496 + # image / audio / etc. Emitting them anyway triggers a Pydantic
497 + # ``ValidatorIterator`` failure at chat-parsing time, so we
498 + # JSON-serialize each document into a text part instead.
499 + #
500 + # The cohere renderer doesn't consume this flattened content:
501 + # :meth:`_collect_v2_tool_message_content` forwards the
502 + # original v2 content blocks under
503 + # :data:`TOOL_MESSAGE_V2_CONTENT_KEY`, and the renderer applies
504 + # :func:`vllm.renderers.cohere._v2_tool_content_to_melody_block`
505 + # to produce the melody-shape ``document`` blocks cmd3/cmd4
506 + # tool-result rendering expects. Outbound citation binding is
507 + # unaffected either way: sources resolve through the
508 + # ``POSITION_TO_SOURCE_KEY`` map (built from the original v2
509 + # messages, not the converted OpenAI messages).
510 + content_parts: list[dict[str, Any]] = []
511 for block in msg.content:
512 if block.type == "text":
469 - parts.append({"type": "text", "text": block.text})
513 + content_parts.append({"type": "text", "text": block.text})
514 elif block.type == "document":
471 - parts.append(
515 + content_parts.append(
516 {
473 - "type": "document",
474 - "document": block.document.model_dump(exclude_none=True),
517 + "type": "text",
518 + "text": json.dumps(
519 + block.document.model_dump(exclude_none=True)
520 + ),
521 }
522 )
523 return {
524 "role": "tool",
525 "tool_call_id": msg.tool_call_id,
480 - "content": parts,
526 + "content": content_parts,
527 }
528
529 @classmethod
@@ -660,6 +706,16 @@ class CohereServingChatV2(OpenAIServingChat):
706 if position_to_source:
707 kwargs.setdefault(POSITION_TO_SOURCE_KEY, position_to_source)
708
709 + # Forward the original v2 tool-message content blocks so the
710 + # renderer can apply its v2 → melody mapping over them (the
711 + # renderer owns the melody-shape rules, this handler only
712 + # ships the raw structural info that OpenAI validation would
713 + # otherwise destroy -- see the flatten comment in
714 + # :meth:`_convert_tool_message`).
715 + v2_tool_content = cls._collect_v2_tool_message_content(request)
716 + if v2_tool_content:
717 + kwargs.setdefault(TOOL_MESSAGE_V2_CONTENT_KEY, v2_tool_content)
718 +
719 if kwargs:
720 chat_req.chat_template_kwargs = kwargs
721
@@ -1262,6 +1318,49 @@ class CohereServingChatV2(OpenAIServingChat):
1318 )
1319 return None
1320
1321 + @classmethod
1322 + def _collect_v2_tool_message_content(
1323 + cls, request: CohereChatV2Request
1324 + ) -> dict[int, list[dict[str, Any]]]:
1325 + """Collect the original v2 tool-message content blocks so the
1326 + renderer can apply its v2 → melody mapping.
1327 +
1328 + Builds the map that
1329 + :meth:`_apply_cohere_template_kwargs` stashes under
1330 + :data:`TOOL_MESSAGE_V2_CONTENT_KEY` for the renderer to
1331 + consume (see the constant definition in
1332 + :mod:`vllm.renderers.cohere` for why the map exists). Walks
1333 + ``request.messages`` -- not the OpenAI-shape messages built by
1334 + :meth:`_convert_messages` -- because that flattening step
1335 + already threw away the structural distinction between text and
1336 + document blocks. Each surviving block is
1337 + :meth:`~pydantic.BaseModel.model_dump`'d to a plain dict so the
1338 + forwarded map carries no Pydantic instances (msgpack- /
1339 + general-serialization-safe on principle even though internal
1340 + underscore-prefixed keys are filtered before crossing the
1341 + engine-core boundary).
1342 +
1343 + Only emits an entry for a tool message whose content is a
1344 + non-empty list of blocks; a bare-string tool message renders
1345 + fine through the default text-block path in
1346 + :func:`vllm.renderers.cohere._content_blocks`, so no forwarded
1347 + blocks are needed there. Returns an empty dict when there is
1348 + nothing to forward, so the caller can use truthiness to decide
1349 + whether to stash the key at all.
1350 + """
1351 + blocks_by_index: dict[int, list[dict[str, Any]]] = {}
1352 + for i, msg in enumerate(request.messages):
1353 + if not isinstance(msg, ToolChatMessageV2):
1354 + continue
1355 + if isinstance(msg.content, str) or not msg.content:
1356 + continue
1357 + v2_blocks: list[dict[str, Any]] = []
1358 + for block in msg.content:
1359 + v2_blocks.append(block.model_dump(exclude_none=True))
1360 + if v2_blocks:
1361 + blocks_by_index[i] = v2_blocks
1362 + return blocks_by_index
1363 +
1364 @staticmethod
1365 def _build_usage(response: ChatCompletionResponse) -> CohereUsage | None:
1366 if response.usage is None:
vllm/entrypoints/openai/chat_completion/serving.py
+21 -1
@@ -196,6 +196,24 @@ class OpenAIServingChat(GenerateBaseServing):
196 .chat_template_kwargs
197 )
198
199 + def _engine_chat_template_kwargs(
200 + self, chat_template_kwargs: dict[str, Any]
201 + ) -> dict[str, Any]:
202 + """Subclass hook to narrow ``chat_template_kwargs`` for the engine.
203 +
204 + The same dict is used twice: to build the API-server-side parser
205 + instance, and to populate
206 + ``EngineCoreRequest.reasoning_parser_kwargs`` for the engine-core
207 + side. The latter crosses ZMQ as msgpack, so a handler that stashes
208 + request-scoped state only the API-server-side parser needs (values
209 + msgpack can't encode, or payloads not worth shipping) can drop
210 + those entries here without affecting the in-process parser.
211 +
212 + Must not mutate the argument -- the caller still needs the full
213 + dict. The default forwards it unchanged.
214 + """
215 + return chat_template_kwargs
216 +
217 async def render_chat_request(
218 self,
219 request: ChatCompletionRequest,
@@ -360,7 +378,9 @@ class OpenAIServingChat(GenerateBaseServing):
378 session_id=session_id,
379 reasoning_ended=reasoning_ended,
380 reasoning_parser_kwargs={
363 - "chat_template_kwargs": chat_template_kwargs,
381 + "chat_template_kwargs": self._engine_chat_template_kwargs(
382 + chat_template_kwargs
383 + ),
384 }
385 if parser is not None and parser.reasoning_parser is not None
386 else None,
vllm/reasoning/cohere_command_reasoning_parser.py
+2 -19
@@ -681,27 +681,12 @@ class BaseCohereCommandReasoningParser(ReasoningParser):
681 return request
682
683
684 -# melody's streaming filter only buffers a partial ``<co: ...>`` citation
685 -# across ``write_decoded`` calls when ``stream_non_grounded_answer`` is
686 -# set: otherwise, the moment an opening ``<co`` is seen without a
687 -# closing ``</co: ...>`` in the same delta, the filter emits the partial
688 -# marker bytes verbatim as plain content. In vLLM's streaming path the
689 -# parser is fed one token (1-4 chars) per call, so an unbuffered filter
690 -# will leak ``<co: 0>``-style markers into ``delta.content`` and never
691 -# emit a ``FilterCitation`` for them. Enabling the flag flips the
692 -# partial-match branch in melody's ``parse_citations`` (see
693 -# ``src/parsing/citations_filter.rs``) to ``return (None, 0)`` -- i.e.
694 -# keep buffering -- which lets a full citation eventually resolve.
695 -# Non-streaming (unary) parsing receives the whole output in one call so
696 -# the flag is a no-op there and we leave ``unary_opts`` alone.
684 class CohereCommand3ReasoningParser(BaseCohereCommandReasoningParser):
685 def __init__(self, tokenizer: TokenizerLike, *args, **kwargs):
686 super().__init__(
687 tokenizer,
688 *args,
702 - streaming_opts=(
703 - PyFilterOptions().cmd3().no_tools().stream_non_grounded_answer()
704 - ),
689 + streaming_opts=PyFilterOptions().cmd3().no_tools(),
690 unary_opts=PyFilterOptions().cmd3().no_tools(),
691 **kwargs,
692 )
@@ -712,9 +697,7 @@ class CohereCommand4ReasoningParser(BaseCohereCommandReasoningParser):
697 super().__init__(
698 tokenizer,
699 *args,
715 - streaming_opts=(
716 - PyFilterOptions().cmd4().no_tools().stream_non_grounded_answer()
717 - ),
700 + streaming_opts=PyFilterOptions().cmd4().no_tools(),
701 unary_opts=PyFilterOptions().cmd4().no_tools(),
702 **kwargs,
703 )
vllm/renderers/cohere.py
+131 -6
@@ -157,6 +157,40 @@ MESSAGES_CITATIONS_KEY = "_messages_citations"
157 POSITION_TO_SOURCE_KEY = "_position_to_source"
158
159
160 +# Reserved ``chat_template_kwargs`` key carrying the original Cohere v2
161 +# tool-message content, keyed by request-message index. Populated by
162 +# :meth:`CohereServingChatV2._apply_cohere_template_kwargs` and consumed
163 +# by :func:`_conversation_to_melody_messages`, which runs each block
164 +# through :func:`_v2_tool_content_to_melody_block` to produce the
165 +# melody-shape content list the renderer emits.
166 +#
167 +# Motivation: OpenAI's ``ChatCompletionContentPartParam`` union (in
168 +# ``vllm/entrypoints/chat_utils.py``) only allows ``text`` / ``image_url``
169 +# / ``audio`` / ``video`` / ``file`` content parts on tool messages, so
170 +# ``{"type": "document", "document": {...}}`` blocks -- valid in Cohere
171 +# v2 -- get rejected by vLLM's shared chat validator before they can
172 +# reach the renderer.
173 +# ``CohereServingChatV2._convert_tool_message`` works around that by
174 +# rewriting each content block into a plain ``text`` part (with document
175 +# blocks JSON-serialized inline), but that flattening loses the
176 +# structural distinction the cohere renderer needs to shape tool-result
177 +# payloads for cmd3/cmd4. The serving layer therefore forwards the
178 +# original v2 content blocks verbatim (as plain dicts) under this key,
179 +# and the renderer owns the v2-shape → melody-shape mapping.
180 +#
181 +# Value type: ``dict[int, list[dict[str, Any]]]``. Each list entry is a
182 +# ``model_dump``'d ``TextToolContent`` / ``DocumentToolContent`` block
183 +# (e.g. ``{"type": "text", "text": ...}`` or ``{"type": "document",
184 +# "document": {"id": ..., "data": {...}}}``). Index is the position of
185 +# the tool message in the request's ``messages`` list, which -- since
186 +# ``_convert_messages`` is 1-to-1 -- also matches its position in the
187 +# ``ConversationMessage`` list handed to the renderer. Leading
188 +# underscore signals "internal, not part of the public
189 +# ``chat_template_kwargs`` surface" -- clients should not set this key
190 +# directly.
191 +TOOL_MESSAGE_V2_CONTENT_KEY = "_tool_message_v2_content"
192 +
193 +
194 # Keys this renderer interprets directly from ``chat_template_kwargs`` and
195 # maps onto typed melody render-config fields. Everything *not* in this
196 # set is forwarded verbatim to melody as ``additional_template_fields``
@@ -196,6 +230,7 @@ _RENDERER_CONSUMED_KEYS = frozenset(
230 "platform_instruction",
231 MESSAGES_CITATIONS_KEY,
232 POSITION_TO_SOURCE_KEY,
233 + TOOL_MESSAGE_V2_CONTENT_KEY,
234 }
235 )
236
@@ -377,6 +412,69 @@ def _content_blocks(content: Any) -> list[dict[str, Any]]:
412 return blocks
413
414
415 +def _v2_tool_content_to_melody_block(block: dict[str, Any]) -> dict[str, Any] | None:
416 + """Convert one Cohere v2 tool-message content block into a melody
417 + ``document`` content block.
418 +
419 + Cohere's chat v2 admits two shapes for tool-message content parts:
420 +
421 + * ``TextToolContent`` -- ``{"type": "text", "text": "..."}``.
422 + * ``DocumentToolContent`` -- ``{"type": "document", "document":
423 + {"id"?: str, "data": {...}}}``.
424 +
425 + Both are emitted as melody ``document`` blocks, mirroring the cohere
426 + api, which routes every tool content part through the same
427 + "tool output as document" conversion. The rules:
428 +
429 + * Text that parses to a JSON object becomes that object directly
430 + (so structured tool output flows through as first-class fields).
431 + * Text that doesn't parse as an object -- empty, non-JSON,
432 + or JSON scalar/array -- gets wrapped as ``{"content": <text>}``
433 + so templates have a stable field name to look up.
434 + * ``document`` blocks unwrap ``document.data`` (a flat map) into
435 + the melody payload; ``document.id`` is intentionally dropped
436 + here because outbound citation binding tracks it separately via
437 + ``POSITION_TO_SOURCE_KEY``.
438 +
439 + Only the JSON-object case actually changes what gets rendered. The
440 + cmd3/cmd4 templates already wrap a bare ``text`` content block as
441 + ``{"content": <text>}`` themselves, so routing plain text through
442 + here is a no-op on the prompt -- it is the JSON-object unwrap that
443 + promotes structured tool output from an escaped string blob into
444 + individually addressable (and citable) result fields.
445 +
446 + Note the payload is the *flat* field map: melody's ``Content``
447 + carries ``document`` as a plain map, so re-nesting it under a
448 + ``data`` key would render a spurious ``{"data": {...}}`` wrapper
449 + into the tool-result block.
450 +
451 + Returns ``None`` for unrecognised shapes; the caller should skip
452 + those blocks (the citation-numbering walk in
453 + ``CohereServingChatV2._walk_citable_positions`` advances its own
454 + cursor over unknown blocks, so slot alignment is preserved).
455 + """
456 + block_type = block.get("type")
457 + if block_type == "text":
458 + text = block.get("text") or ""
459 + try:
460 + parsed = json.loads(text)
461 + except (json.JSONDecodeError, TypeError):
462 + parsed = None
463 + if isinstance(parsed, dict):
464 + payload: dict[str, Any] = dict(parsed)
465 + else:
466 + payload = {"content": text}
467 + return {"type": MelodyContentType.DOCUMENT, "document": payload}
468 + if block_type == "document":
469 + doc = block.get("document")
470 + if not isinstance(doc, dict):
471 + return None
472 + data = doc.get("data")
473 + payload = dict(data) if isinstance(data, dict) else {}
474 + return {"type": MelodyContentType.DOCUMENT, "document": payload}
475 + return None
476 +
477 +
478 def _document_to_melody(doc: Any) -> dict[str, Any]:
479 """Coerce a Cohere v2 document into melody's ``Document`` (dict) shape."""
480 match doc:
@@ -421,6 +519,7 @@ def _tool_to_melody(tool: Any) -> dict[str, Any]:
519 def _conversation_to_melody_messages(
520 conversation: list[ConversationMessage],
521 messages_citations: dict[int, list[dict[str, Any]]] | None = None,
522 + tool_message_v2_content: dict[int, list[dict[str, Any]]] | None = None,
523 ) -> list[dict[str, Any]]:
524 """Convert vLLM ``ConversationMessage``s into melody's message dict shape.
525
@@ -429,17 +528,42 @@ def _conversation_to_melody_messages(
528 ``CohereServingChatV2`` under the
529 :data:`MESSAGES_CITATIONS_KEY` chat_template_kwargs entry. Values
530 must already be in melody's ``FilterCitation`` dict shape (see
432 - ``CohereServingChatV2._sdk_citation_to_melody``). Index is the
433 - position in the *input* ``conversation`` list, which currently maps
434 - 1-to-1 with ``CohereChatV2Request.messages`` for text-only Cohere
435 - requests. This mapping breaks silently if ``parse_chat_messages``
436 - ever expands a single input message into multiple entries (e.g.
531 + ``CohereServingChatV2._sdk_citation_to_melody``).
532 +
533 + ``tool_message_v2_content`` is an optional index-keyed lookup of
534 + the original Cohere v2 tool-message content blocks (as plain
535 + dicts), populated by ``CohereServingChatV2`` under the
536 + :data:`TOOL_MESSAGE_V2_CONTENT_KEY` entry. When set for a given
537 + index, its blocks -- run through
538 + :func:`_v2_tool_content_to_melody_block` to apply the v2 → melody
539 + mapping -- replace the block list this function would otherwise
540 + derive from ``msg["content"]``. This lets the renderer own the
541 + v2-shape → melody-shape conversion for the structured ``document``
542 + blocks cohere v2 allows on tool messages (which don't survive
543 + OpenAI's content-parts validator on the way in, so
544 + ``_convert_tool_message`` rewrites them into ``text`` parts).
545 +
546 + Both index lookups key by position in the *input* ``conversation``
547 + list, which currently maps 1-to-1 with
548 + ``CohereChatV2Request.messages`` for text-only Cohere requests.
549 + This mapping breaks silently if ``parse_chat_messages`` ever
550 + expands a single input message into multiple entries (e.g.
551 multi-modal splitting), which is not the case today.
552 """
553 out: list[dict[str, Any]] = []
554 for i, msg in enumerate(conversation):
555 role = _role_to_melody(msg.get("role", "user"))
442 - content_blocks = _content_blocks(msg.get("content"))
556 + if (
557 + tool_message_v2_content is not None
558 + and (v2_blocks := tool_message_v2_content.get(i)) is not None
559 + ):
560 + content_blocks = []
561 + for v2_block in v2_blocks:
562 + mapped = _v2_tool_content_to_melody_block(v2_block)
563 + if mapped is not None:
564 + content_blocks.append(mapped)
565 + else:
566 + content_blocks = _content_blocks(msg.get("content"))
567
568 # Treat reasoning content as a thinking block on assistant turns
569 # so multi-turn reasoning is preserved in the rendered prompt.
@@ -513,6 +637,7 @@ def _build_render_config(
637 "messages": _conversation_to_melody_messages(
638 conversation,
639 chat_template_kwargs.get(MESSAGES_CITATIONS_KEY),
640 + chat_template_kwargs.get(TOOL_MESSAGE_V2_CONTENT_KEY),
641 ),
642 }
643