@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 187 lines 5.95 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 """parse_delta-level regression tests for the MuseGlimmer parsers.
5
6 These drive the UNIFIED parser-engine streaming API the serving layer actually
7 uses (``DelegatingParser.parse_delta``), NOT the ``extract_*_streaming`` methods
8 directly. This is the coverage that was missing when a live no-tools streaming
9 request leaked raw harmony framing into ``content``: the phase machine marked
10 ``reasoning_ended=True`` at the prompt boundary (because ``is_reasoning_end`` on
11 a prompt with no ``to=self`` wrongly returned True), skipping the reasoning
12 phase for the whole generation.
13
14 Requires a real MuseGlimmer tokenizer; skipped if the checkpoint is unavailable.
15 """
16
17 import json
18 import os
19
20 import pytest
21
22 CKPT = os.environ.get("MUSE_GLIMMER_CKPT", "")
23
24 pytestmark = pytest.mark.skipif(
25 not os.path.isdir(CKPT), reason=f"MuseGlimmer checkpoint not found at {CKPT}"
26 )
27
28 _FRAMING = [
29 "<|start|>",
30 "<|message|>",
31 "<|eom|>",
32 "<|eot|>",
33 "to=self",
34 "to=user",
35 "<atem:",
36 ]
37
38 PROMPT = "<|start|>user<|message|>hi<|eom|><|start|>assistant"
39
40
41 class _Req:
42 tools = None
43 tool_choice = None
44 include_reasoning = True
45
46
47 @pytest.fixture(scope="module")
48 def tok():
49 from transformers import AutoTokenizer
50
51 return AutoTokenizer.from_pretrained(CKPT, trust_remote_code=True)
52
53
54 @pytest.fixture(scope="module")
55 def parser_cls():
56 from vllm.parser import ParserManager
57
58 return ParserManager.get_parser(
59 tool_parser_name="muse_glimmer",
60 reasoning_parser_name="muse_glimmer",
61 enable_auto_tools=True,
62 )
63
64
65 def _drive(parser_cls, tok, gen_text, req=None):
66 """Feed gen_text token-by-token through parse_delta (as serving does)."""
67 from vllm.parser.abstract_parser import StreamState
68
69 if req is None:
70 req = _Req()
71 parser = parser_cls(tok)
72 parser._stream_state = StreamState()
73 prompt_ids = tok.encode(PROMPT, add_special_tokens=False)
74 gen_ids = tok.encode(gen_text, add_special_tokens=False)
75 reasoning, content, tools = [], [], []
76 for i, tid in enumerate(gen_ids):
77 dm = parser.parse_delta(
78 tok.decode([tid]),
79 [tid],
80 req,
81 prompt_token_ids=prompt_ids if i == 0 else None,
82 finished=(i == len(gen_ids) - 1),
83 )
84 if dm is None:
85 continue
86 if getattr(dm, "reasoning", None):
87 reasoning.append(dm.reasoning)
88 if getattr(dm, "content", None):
89 content.append(dm.content)
90 for tc in getattr(dm, "tool_calls", None) or []:
91 fn = tc.function
92 name = fn.get("name") if isinstance(fn, dict) else getattr(fn, "name", None)
93 args = (
94 fn.get("arguments")
95 if isinstance(fn, dict)
96 else getattr(fn, "arguments", None)
97 )
98 tools.append((tc.index, name, args))
99 return "".join(reasoning), "".join(content), tools
100
101
102 def _assert_no_framing(text):
103 for f in _FRAMING:
104 assert f not in text, f"framing {f!r} leaked: {text!r}"
105
106
107 def test_parse_delta_no_tools_reasoning_then_answer(parser_cls, tok):
108 # The core regression: reasoning phase must stay active for a no-tools turn.
109 reasoning, content, tools = _drive(
110 parser_cls,
111 tok,
112 " to=self<|message|>Let me think step by step about the sum.<|eom|>"
113 "<|start|>assistant to=user<|message|>The answer is 42.<|eot|>",
114 )
115 _assert_no_framing(content)
116 assert reasoning == "Let me think step by step about the sum.", repr(reasoning)
117 assert content == "The answer is 42.", repr(content)
118 assert tools == []
119
120
121 def test_parse_delta_content_only(parser_cls, tok):
122 reasoning, content, tools = _drive(
123 parser_cls,
124 tok,
125 " to=user<|message|>Just a direct answer.<|eot|>",
126 )
127 _assert_no_framing(content)
128 assert content == "Just a direct answer.", repr(content)
129 assert tools == []
130
131
132 def test_parse_delta_tool_call(parser_cls, tok):
133 reasoning, content, tools = _drive(
134 parser_cls,
135 tok,
136 " to=self<|message|>I should read the hostname.<|eom|>"
137 "<|start|>assistant to=read.read<|message|>"
138 '<atem:function_calls>\n<atem:invoke name="read.read">\n'
139 '<atem:parameter name="path">/etc/hostname</atem:parameter>\n'
140 "</atem:invoke>\n</atem:function_calls>",
141 )
142 _assert_no_framing(content)
143 assert reasoning == "I should read the hostname.", repr(reasoning)
144 assert len(tools) == 1, tools
145 idx, name, args = tools[0]
146 assert idx == 0 and name == "read.read"
147 assert json.loads(args) == {"path": "/etc/hostname"}
148
149
150 def test_parse_delta_truncated_cot_no_toolcall(parser_cls, tok):
151 # Contemplated invoke inside an unterminated to=self block: no tool call,
152 # no framing leak into content, partial reasoning recovered.
153 reasoning, content, tools = _drive(
154 parser_cls,
155 tok,
156 " to=self<|message|>Maybe I should call "
157 '<atem:function_calls>\n<atem:invoke name="read.read">\n'
158 '<atem:parameter name="path">/etc/hostname</atem:parameter>\n'
159 "</atem:invoke>\n</atem:function_calls> but wait",
160 )
161 _assert_no_framing(content)
162 assert tools == [], f"contemplated invoke leaked as tool call: {tools}"
163 assert "Maybe I should call" in reasoning, repr(reasoning)
164
165
166 def test_parse_delta_reasoning_suppressed_when_not_requested(parser_cls, tok):
167 class _NoReasonReq:
168 tools = None
169 tool_choice = None
170 include_reasoning = False
171
172 reasoning, content, tools = _drive(
173 parser_cls,
174 tok,
175 " to=self<|message|>secret thoughts<|eom|>"
176 "<|start|>assistant to=user<|message|>Public answer.<|eot|>",
177 req=_NoReasonReq(),
178 )
179 assert reasoning == "", repr(reasoning) # suppressed
180 assert content == "Public answer.", repr(content)
181 _assert_no_framing(content)
182
183
184 if __name__ == "__main__":
185 import sys
186
187 sys.exit(pytest.main([__file__, "-v"]))