| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Citation source resolution and the streamed citation wire shape.""" |
| 4 | |
| 5 | from __future__ import annotations |
| 6 | |
| 7 | from types import SimpleNamespace |
| 8 | from typing import Any |
| 9 | |
| 10 | from vllm.parser.cohere_command import ( |
| 11 | _melody_citations_to_vllm, |
| 12 | _melody_sources_to_vllm, |
| 13 | ) |
| 14 | |
| 15 | from .utils import make_parser, stream_parser |
| 16 | |
| 17 | |
| 18 | class TestMelodySourceResolution: |
| 19 | """Pin the parser-side source resolution (see |
| 20 | ``_melody_sources_to_vllm``). The parser receives melody's numeric |
| 21 | ``(tool_call_index, tool_result_indices)`` addressing and resolves |
| 22 | it against a ``position_to_source`` map handed in by |
| 23 | ``CohereServingChatV2._apply_cohere_template_kwargs`` -- the same |
| 24 | map that would otherwise live in the serving layer's resolver. |
| 25 | """ |
| 26 | |
| 27 | @staticmethod |
| 28 | def _fake_melody_source(bucket: int, indices: list[int]) -> Any: |
| 29 | return SimpleNamespace(tool_call_index=bucket, tool_result_indices=indices) |
| 30 | |
| 31 | def test_multi_index_source_fans_out(self): |
| 32 | from vllm.entrypoints.cohere.cohere_chat_message import CitationSource |
| 33 | |
| 34 | position_map: dict[tuple[int, int], CitationSource] = { |
| 35 | (0, 0): CitationSource(type="document", id="d0", document={"id": "d0"}), |
| 36 | (0, 1): CitationSource(type="document", id="d1", document={"id": "d1"}), |
| 37 | } |
| 38 | raw = [self._fake_melody_source(0, [0, 1])] |
| 39 | out = _melody_sources_to_vllm(raw, position_map) |
| 40 | assert [s.id for s in out] == ["d0", "d1"] |
| 41 | # Verify type / payload were plumbed through, not just the id. |
| 42 | assert out[0].type == "document" |
| 43 | assert out[0].document == {"id": "d0"} |
| 44 | |
| 45 | def test_unresolvable_position_skipped(self): |
| 46 | from vllm.entrypoints.cohere.cohere_chat_message import CitationSource |
| 47 | |
| 48 | position_map: dict[tuple[int, int], CitationSource] = { |
| 49 | (0, 0): CitationSource(type="document", id="d0"), |
| 50 | } |
| 51 | raw = [self._fake_melody_source(9, [0])] |
| 52 | assert _melody_sources_to_vllm(raw, position_map) == [] |
| 53 | |
| 54 | def test_missing_position_map_drops_all_sources(self): |
| 55 | # A parser instance without a position map (parser wired |
| 56 | # outside of ``CohereServingChatV2``) can't attribute anything, |
| 57 | # so every source is dropped. Callers downstream will see the |
| 58 | # citation with empty ``sources`` and drop it entirely. |
| 59 | raw = [self._fake_melody_source(0, [0])] |
| 60 | assert _melody_sources_to_vllm(raw, None) == [] |
| 61 | |
| 62 | def test_citations_pass_through_is_thinking_tag(self): |
| 63 | from vllm.entrypoints.cohere.cohere_chat_message import CitationSource |
| 64 | |
| 65 | position_map: dict[tuple[int, int], CitationSource] = { |
| 66 | (0, 0): CitationSource(type="document", id="d0"), |
| 67 | } |
| 68 | raw = [ |
| 69 | SimpleNamespace( |
| 70 | start_index=0, |
| 71 | end_index=5, |
| 72 | text="hello", |
| 73 | is_thinking=True, |
| 74 | sources=[self._fake_melody_source(0, [0])], |
| 75 | ) |
| 76 | ] |
| 77 | out = _melody_citations_to_vllm(raw, position_map) |
| 78 | assert out is not None |
| 79 | assert out[0].type == "THINKING_CONTENT" |
| 80 | assert out[0].sources[0].id == "d0" |
| 81 | |
| 82 | |
| 83 | class TestParserStreamingEndToEnd: |
| 84 | """End-to-end streaming shape: raw model output text is tokenized |
| 85 | with the mock byte-level tokenizer and fed one token at a time |
| 86 | through :meth:`CohereCommand3ReasoningParser.extract_reasoning_streaming`. |
| 87 | We collect the ordered sequence of ``DeltaMessage``s the parser |
| 88 | hands back -- i.e. what the client would receive on the wire -- |
| 89 | and pin the exact shape. |
| 90 | |
| 91 | Combines the three flavors that reach ``delta``: a thinking block |
| 92 | (populates ``reasoning``), a text block (populates ``content``), |
| 93 | and a resolved citation (populates ``citations``) referencing |
| 94 | ``(start, end)`` offsets into the already-streamed content. |
| 95 | """ |
| 96 | |
| 97 | _MODEL_OUTPUT = ( |
| 98 | "<|START_THINKING|>Let me check.<|END_THINKING|>" |
| 99 | "<|START_RESPONSE|>The capital is <co>Paris</co: 0:[1]>.<|END_RESPONSE|>" |
| 100 | ) |
| 101 | |
| 102 | def test_full_wire_stream_shape(self, tokenizer, request_obj): |
| 103 | from vllm.entrypoints.cohere.cohere_chat_message import CitationSource |
| 104 | from vllm.renderers.cohere import POSITION_TO_SOURCE_KEY |
| 105 | |
| 106 | # Mimic what ``CohereServingChatV2._apply_cohere_template_kwargs`` |
| 107 | # installs on a request whose tool_call_index=0 tool result at |
| 108 | # index 1 is the "France" document. |
| 109 | position_to_source = { |
| 110 | (0, 1): CitationSource( |
| 111 | type="document", |
| 112 | id="doc-paris", |
| 113 | document={"id": "doc-paris", "title": "France"}, |
| 114 | ), |
| 115 | } |
| 116 | |
| 117 | parser = make_parser( |
| 118 | tokenizer, |
| 119 | "cohere_command3", |
| 120 | chat_template_kwargs={POSITION_TO_SOURCE_KEY: position_to_source}, |
| 121 | ) |
| 122 | deltas = stream_parser(parser, request_obj, tokenizer, self._MODEL_OUTPUT) |
| 123 | |
| 124 | # -- 1. Reasoning and content stream separately and reach |
| 125 | # the wire only as their real payload. Framing tokens |
| 126 | # (``<|START_THINKING|>`` etc.) are consumed by the parser's |
| 127 | # state machine and never produce a delta. |
| 128 | reasoning_stream = [d.reasoning for d in deltas if d.reasoning is not None] |
| 129 | content_stream = [d.content for d in deltas if d.content is not None] |
| 130 | citation_deltas = [d for d in deltas if getattr(d, "citations", None)] |
| 131 | |
| 132 | assert "".join(reasoning_stream) == "Let me check." |
| 133 | assert "".join(content_stream) == "The capital is Paris." |
| 134 | |
| 135 | # -- 2. Exactly one citation reaches the wire, resolved to the |
| 136 | # source from ``position_to_source`` (not the raw |
| 137 | # ``(bucket, idx)`` coordinates). |
| 138 | assert len(citation_deltas) == 1 |
| 139 | (cite_delta,) = citation_deltas |
| 140 | assert cite_delta.citations is not None |
| 141 | assert len(cite_delta.citations) == 1 |
| 142 | cite = cite_delta.citations[0] |
| 143 | assert cite.text == "Paris" |
| 144 | # Char offsets are into the accumulated ``content``: "The |
| 145 | # capital is " is 15 chars, "Paris" is 5. |
| 146 | assert cite.start == 15 |
| 147 | assert cite.end == 20 |
| 148 | assert len(cite.sources) == 1 |
| 149 | assert cite.sources[0].id == "doc-paris" |
| 150 | assert cite.sources[0].type == "document" |
| 151 | |
| 152 | # -- 3. The citation's ``(start, end)`` offsets point into |
| 153 | # bytes that were already emitted as content-deltas earlier |
| 154 | # in the stream -- i.e. the citation is anchored to real |
| 155 | # already-shipped content. |
| 156 | cite_delta_idx = deltas.index(cite_delta) |
| 157 | earlier_content = "".join( |
| 158 | d.content for d in deltas[:cite_delta_idx] if d.content is not None |
| 159 | ) |
| 160 | assert earlier_content[cite.start : cite.end] == "Paris" |
| 161 | |
| 162 | # Everything after the citation delta is the trailing period. |
| 163 | later_content = "".join( |
| 164 | d.content for d in deltas[cite_delta_idx + 1 :] if d.content is not None |
| 165 | ) |
| 166 | assert later_content == "." |
| 167 | |
| 168 | # -- 4. Reasoning and content are never mixed on the same |
| 169 | # delta: the parser flips modes on the ``<|END_THINKING|>`` |
| 170 | # boundary and never emits a delta with both fields set. |
| 171 | for d in deltas: |
| 172 | assert not (d.reasoning is not None and d.content is not None) |