main
py 556 lines 18.2 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """PLaMo3 grammar, truncation, and serving-adapter regressions."""
4
5 from __future__ import annotations
6
7 import json
8
9 import pytest
10 import regex as re
11
12 from tests.parser.engine.conftest import make_mock_tokenizer
13 from tests.parser.engine.replay_harness import (
14 DUMMY_TOOLS,
15 _test_request,
16 collect_output,
17 replay_streaming,
18 replay_with_text_holdback,
19 )
20 from vllm.parser.parser_manager import ParserManager
21 from vllm.parser.plamo3 import (
22 BEGIN_THINK,
23 BEGIN_TOOL_ARGUMENTS,
24 BEGIN_TOOL_NAME,
25 BEGIN_TOOL_REQUEST,
26 BEGIN_TOOL_REQUESTS,
27 END_THINK,
28 END_TOOL_ARGUMENTS,
29 END_TOOL_NAME,
30 END_TOOL_REQUEST,
31 END_TOOL_REQUESTS,
32 EOT,
33 PLAMO_MARKER_TOKENS,
34 Plamo3Parser,
35 )
36
37
38 def _tool_header(name="weather"):
39 return (
40 BEGIN_TOOL_REQUEST
41 + BEGIN_TOOL_NAME
42 + name
43 + END_TOOL_NAME
44 + BEGIN_TOOL_ARGUMENTS
45 )
46
47
48 def _tool_call(name, arguments):
49 return _tool_header(name) + arguments + END_TOOL_ARGUMENTS + END_TOOL_REQUEST
50
51
52 @pytest.fixture
53 def mock_tokenizer():
54 special = [*sorted(PLAMO_MARKER_TOKENS), EOT, "<|plamo:bos|>"]
55 atoms = [*special, "<|plamo:constrain|>", "<|plamo:msg|>"]
56 vocab = {token: 100 + i for i, token in enumerate(atoms)}
57 inverse = {i: token for token, i in vocab.items()}
58 pattern = re.compile("|".join(map(re.escape, atoms)) + r"|[\s\S]")
59 tokenizer = make_mock_tokenizer(vocab, special_tokens=special)
60 tokenizer.encode.side_effect = lambda text, **kwargs: [
61 vocab[piece] if piece in vocab else 1000 + ord(piece)
62 for piece in pattern.findall(text)
63 ]
64 tokenizer.decode.side_effect = lambda ids, **kwargs: "".join(
65 inverse[i] if i in inverse else chr(i - 1000) for i in ids
66 )
67 return tokenizer
68
69
70 @pytest.fixture(params=["direct", "registered"])
71 def parser_cls(request):
72 if request.param == "direct":
73 return Plamo3Parser
74 return ParserManager.get_parser(
75 tool_parser_name="plamo3",
76 reasoning_parser_name="plamo3",
77 enable_auto_tools=True,
78 )
79
80
81 def _stream(parser, tokenizer, output, chunk_size=1, finished_on_last=False):
82 tokens = [
83 (i, tokenizer.decode([i]))
84 for i in tokenizer.encode(output, add_special_tokens=False)
85 ]
86 deltas = replay_streaming(
87 parser,
88 tokens,
89 chunk_size=chunk_size,
90 finished_on_last=finished_on_last,
91 tools=DUMMY_TOOLS,
92 )
93 if not finished_on_last:
94 deltas.append(
95 parser.parse_delta("", [], _test_request(DUMMY_TOOLS), finished=True)
96 )
97 return deltas
98
99
100 @pytest.mark.parametrize("thinking", [False, True])
101 @pytest.mark.parametrize("with_markers", [False, True])
102 def test_reasoning_mode(
103 parser_cls, mock_tokenizer, mock_request, thinking, with_markers
104 ):
105 output = (
106 BEGIN_THINK + "Reasoning." + END_THINK + "Answer."
107 if with_markers
108 else "Answer."
109 )
110 if thinking:
111 expected = ("Reasoning.", "Answer.") if with_markers else ("Answer.", "")
112 else:
113 expected = ("", output)
114 kwargs = {"chat_template_kwargs": {"enable_thinking": thinking}}
115 reasoning, content, calls = parser_cls(mock_tokenizer, **kwargs).parse(
116 output, mock_request
117 )
118 assert (reasoning or "", content or "") == expected
119 assert not calls
120 result = collect_output(
121 _stream(parser_cls(mock_tokenizer, **kwargs), mock_tokenizer, output)
122 )
123 assert (result.reasoning, result.content) == expected
124 assert not result.tool_calls
125
126
127 def test_duplicate_think_end_does_not_leak_into_content(
128 parser_cls, mock_tokenizer, mock_request
129 ):
130 output = END_THINK + "A" + END_THINK + "B"
131 reasoning, content, calls = parser_cls(mock_tokenizer).parse(output, mock_request)
132 assert (reasoning, content) == (None, "AB")
133 assert not calls
134
135
136 def test_content_can_start_a_new_reasoning_block(
137 parser_cls, mock_tokenizer, mock_request
138 ):
139 output = END_THINK + "A" + BEGIN_THINK + "R" + END_THINK
140 reasoning, content, calls = parser_cls(mock_tokenizer).parse(output, mock_request)
141 assert (reasoning, content) == ("R", "A")
142 assert not calls
143
144
145 @pytest.mark.parametrize("chunk_size", [1, 7, None])
146 @pytest.mark.parametrize("finished_on_last", [False, True])
147 @pytest.mark.parametrize(
148 "closing",
149 [
150 "",
151 END_TOOL_ARGUMENTS + END_TOOL_REQUEST,
152 END_TOOL_ARGUMENTS + END_TOOL_REQUEST + END_TOOL_REQUESTS + EOT,
153 ],
154 )
155 def test_complete_arguments_survive_missing_closers(
156 parser_cls, mock_tokenizer, mock_request, chunk_size, finished_on_last, closing
157 ):
158 output = (
159 BEGIN_THINK
160 + "Reasoning. "
161 + END_THINK
162 + "Before tools. "
163 + BEGIN_TOOL_REQUESTS
164 + _tool_header()
165 + '{"city":"東京"}'
166 + closing
167 )
168 reasoning, content, calls = parser_cls(mock_tokenizer).parse(
169 output, mock_request, enable_auto_tools=True
170 )
171 assert (reasoning, content) == ("Reasoning. ", "Before tools. ")
172 assert calls and [(c.name, json.loads(c.arguments)) for c in calls] == [
173 ("weather", {"city": "東京"})
174 ]
175 result = collect_output(
176 _stream(
177 parser_cls(mock_tokenizer),
178 mock_tokenizer,
179 output,
180 chunk_size,
181 finished_on_last,
182 )
183 )
184 assert (result.reasoning, result.content) == (reasoning, content)
185 assert result.tool_calls == [{"name": "weather", "arguments": '{"city":"東京"}'}]
186
187
188 def test_tool_name_cutoffs_do_not_leak_markers(
189 parser_cls, mock_tokenizer, mock_request
190 ):
191 header = BEGIN_TOOL_REQUEST + BEGIN_TOOL_NAME + "weather" + END_TOOL_NAME
192 header_ids = mock_tokenizer.encode(header, add_special_tokens=False)
193 name_start = len(
194 mock_tokenizer.encode(
195 BEGIN_TOOL_REQUEST + BEGIN_TOOL_NAME, add_special_tokens=False
196 )
197 )
198 for end in range(len(header_ids)):
199 partial = mock_tokenizer.decode(header_ids[:end])
200 output = END_THINK + "Before tools." + BEGIN_TOOL_REQUESTS + partial
201 reasoning, content, calls = parser_cls(mock_tokenizer).parse(
202 output, mock_request, enable_auto_tools=True
203 )
204 assert (reasoning, content) == (None, "Before tools."), end
205 if end <= name_start:
206 expected_name = ""
207 else:
208 name_text = mock_tokenizer.decode(header_ids[name_start:end])
209 if "<|plamo:" in name_text:
210 name_text = name_text[: name_text.rfind("<|plamo:")]
211 expected_name = name_text[: len("weather")]
212 assert [call.name for call in calls or []] == (
213 [expected_name] if expected_name else []
214 ), end
215 result = collect_output(
216 _stream(parser_cls(mock_tokenizer), mock_tokenizer, output)
217 )
218 assert [call["name"] for call in result.tool_calls] == (
219 [expected_name] if expected_name else []
220 ), end
221 assert all("<" not in call["name"] for call in result.tool_calls), end
222 assert result.content == "Before tools.", end
223
224
225 @pytest.mark.parametrize("chunk_size", [1, 7, None])
226 @pytest.mark.parametrize("finished_on_last", [False, True])
227 def test_all_argument_closer_cutoffs_do_not_leak(
228 parser_cls, mock_tokenizer, mock_request, chunk_size, finished_on_last
229 ):
230 head = END_THINK + BEGIN_TOOL_REQUESTS + _tool_header()
231 closer_ids = mock_tokenizer.encode(END_TOOL_ARGUMENTS, add_special_tokens=False)
232 for end in range(len(closer_ids)):
233 partial = mock_tokenizer.decode(closer_ids[:end])
234 output = head + '{"city":"東京"}' + partial
235 _, _, calls = parser_cls(mock_tokenizer).parse(
236 output, mock_request, enable_auto_tools=True
237 )
238 assert calls and calls[0].arguments == '{"city":"東京"}', end
239 result = collect_output(
240 _stream(
241 parser_cls(mock_tokenizer),
242 mock_tokenizer,
243 output,
244 chunk_size,
245 finished_on_last,
246 )
247 )
248 assert result.tool_calls == [
249 {"name": "weather", "arguments": '{"city":"東京"}'}
250 ], end
251
252
253 @pytest.mark.parametrize(
254 "arguments",
255 [
256 "{}",
257 '{"name":"Alice","arguments":{"x":1},"parameters":true}',
258 '{"text":"literal <|plamo:end_',
259 ],
260 )
261 def test_arguments_are_not_rewritten(
262 parser_cls, mock_tokenizer, mock_request, arguments
263 ):
264 output = END_THINK + BEGIN_TOOL_REQUESTS + _tool_header() + arguments
265 _, _, calls = parser_cls(mock_tokenizer).parse(
266 output, mock_request, enable_auto_tools=True
267 )
268 assert calls and calls[0].arguments == arguments
269 result = collect_output(_stream(parser_cls(mock_tokenizer), mock_tokenizer, output))
270 assert result.tool_calls == [{"name": "weather", "arguments": arguments}]
271
272
273 @pytest.mark.parametrize("truncated_second", [False, True])
274 def test_multiple_calls_keep_indices_ids_and_arguments_separate(
275 parser_cls, mock_tokenizer, mock_request, truncated_second
276 ):
277 second = (
278 BEGIN_TOOL_REQUEST + BEGIN_TOOL_NAME + "clo"
279 if truncated_second
280 else _tool_call("clock", '{"timezone":"Asia/Tokyo"}')
281 )
282 output = (
283 END_THINK
284 + BEGIN_TOOL_REQUESTS
285 + _tool_call("weather", '{"city":"東京"}')
286 + second
287 )
288 expected = [{"name": "weather", "arguments": '{"city":"東京"}'}]
289 if truncated_second:
290 expected.append({"name": "clo", "arguments": "{}"})
291 else:
292 output += END_TOOL_REQUESTS
293 expected.append({"name": "clock", "arguments": '{"timezone":"Asia/Tokyo"}'})
294 stream_expected = expected
295 _, content, calls = parser_cls(mock_tokenizer).parse(
296 output, mock_request, enable_auto_tools=True
297 )
298 assert (
299 calls
300 and [{"name": c.name, "arguments": c.arguments} for c in calls] == expected
301 )
302 deltas = _stream(parser_cls(mock_tokenizer), mock_tokenizer, output)
303 result = collect_output(deltas)
304 assert result.tool_calls == stream_expected
305 assert result.content == (content or "") == ""
306 announced = [
307 call for delta in deltas if delta for call in delta.tool_calls or [] if call.id
308 ]
309 assert [call.index for call in announced] == list(range(len(expected)))
310 assert len({call.id for call in announced}) == len(expected)
311
312
313 def test_arguments_are_available_before_outer_closer(mock_tokenizer, mock_request):
314 parser = Plamo3Parser(mock_tokenizer)
315 output = END_THINK + BEGIN_TOOL_REQUESTS + _tool_call("weather", '{"city":"Tokyo"}')
316 first = parser.parse_delta(output, [], mock_request, finished=False)
317 assert collect_output([first]).tool_calls == [
318 {"name": "weather", "arguments": '{"city":"Tokyo"}'}
319 ]
320 assert (
321 parser.parse_delta(END_TOOL_REQUESTS[:10], [], mock_request, finished=False)
322 is None
323 )
324 assert (
325 parser.parse_delta(END_TOOL_REQUESTS[10:], [], mock_request, finished=True)
326 is None
327 )
328
329
330 def test_missing_tool_request_start_does_not_break_streaming(
331 parser_cls, mock_tokenizer, mock_request
332 ):
333 output = (
334 BEGIN_TOOL_REQUESTS
335 + BEGIN_TOOL_NAME
336 + "weather"
337 + END_TOOL_NAME
338 + BEGIN_TOOL_ARGUMENTS
339 + "{}"
340 + END_TOOL_ARGUMENTS
341 + END_TOOL_REQUESTS
342 )
343 delta = parser_cls(mock_tokenizer).parse_delta(
344 output, [], mock_request, finished=False
345 )
346 assert delta is not None
347 assert collect_output([delta]).tool_calls == [
348 {"name": "weather", "arguments": "{}"}
349 ]
350
351
352 def test_reusing_parser_drops_pending_names(mock_tokenizer, mock_request):
353 parser = Plamo3Parser(mock_tokenizer)
354 parser.parse(
355 END_THINK
356 + BEGIN_TOOL_REQUESTS
357 + BEGIN_TOOL_REQUEST
358 + BEGIN_TOOL_NAME
359 + "partial",
360 mock_request,
361 )
362 output = (
363 END_THINK
364 + BEGIN_TOOL_REQUESTS
365 + _tool_call("weather", "{}")
366 + END_TOOL_REQUESTS
367 )
368 _, _, calls = parser.parse(output, mock_request)
369 assert calls and calls[0].name == "weather"
370
371
372 @pytest.mark.parametrize(
373 "text", ["", "a < b", "literal <", "literal <|plamo:unknown", " text\n"]
374 )
375 def test_plain_content_is_preserved(parser_cls, mock_tokenizer, mock_request, text):
376 output = END_THINK + text + EOT
377 _, content, calls = parser_cls(mock_tokenizer).parse(output, mock_request)
378 assert (content or "") == text
379 assert not calls
380 result = collect_output(_stream(parser_cls(mock_tokenizer), mock_tokenizer, output))
381 assert (result.content, result.tool_calls) == (text, [])
382
383
384 @pytest.mark.parametrize(
385 "text, ended",
386 [
387 ("", False),
388 (BEGIN_THINK + "reasoning", False),
389 (END_THINK + "answer", True),
390 (END_THINK + "answer" + BEGIN_THINK + "new reasoning", False),
391 (END_THINK + "answer" + EOT + "assistant", False),
392 ],
393 )
394 def test_reasoning_end_checks_the_latest_turn(mock_tokenizer, text, ended):
395 parser = Plamo3Parser(mock_tokenizer)
396 ids = mock_tokenizer.encode(text)
397 assert parser.is_reasoning_end(ids) == ended
398 if ended:
399 assert mock_tokenizer.decode(parser.extract_content_ids(ids)) == "answer"
400
401
402 @pytest.mark.parametrize("suffix", ["<|plamo:end_", "<|plamo:end_think"])
403 def test_partial_end_think_is_consistent_between_streaming_and_non_streaming(
404 parser_cls, mock_tokenizer, mock_request, suffix
405 ):
406 output = "visible" + suffix
407 reasoning, content, calls = parser_cls(mock_tokenizer).parse(output, mock_request)
408 assert (reasoning, content) == ("visible", None)
409 assert not calls
410 result = collect_output(_stream(parser_cls(mock_tokenizer), mock_tokenizer, output))
411 assert (result.reasoning, result.content, result.tool_calls) == ("visible", "", [])
412
413
414 @pytest.mark.parametrize(
415 "suffix",
416 [
417 "<|plamo:begin_",
418 "<|plamo:begin_tool_name",
419 "<|plamo:end_tool",
420 ],
421 )
422 def test_unfinished_marker_prefixes_are_stripped(mock_tokenizer, suffix):
423 parser = Plamo3Parser(mock_tokenizer)
424 assert parser._strip_unfinished_marker("visible" + suffix) == "visible"
425
426
427 @pytest.mark.parametrize("suffix", ["<|plamo:", "<|plamo:b"])
428 def test_non_token_prefix_is_not_stripped(mock_tokenizer, suffix):
429 parser = Plamo3Parser(mock_tokenizer)
430 text = "visible" + suffix
431 assert parser._strip_unfinished_marker(text) == text
432
433
434 def test_complete_markers_are_not_stripped(mock_tokenizer):
435 parser = Plamo3Parser(mock_tokenizer)
436 text = "visible" + BEGIN_TOOL_NAME
437 assert parser._strip_unfinished_marker(text) == text
438
439
440 def test_short_prefix_is_not_stripped(mock_tokenizer):
441 parser = Plamo3Parser(mock_tokenizer)
442 text = "visible<"
443 assert parser._strip_unfinished_marker(text) == text
444
445
446 def test_tool_block_can_end_reasoning_implicitly(
447 parser_cls, mock_tokenizer, mock_request
448 ):
449 output = (
450 "reasoning"
451 + BEGIN_TOOL_REQUESTS
452 + _tool_call("weather", "{}")
453 + END_TOOL_REQUESTS
454 )
455 reasoning, content, calls = parser_cls(mock_tokenizer).parse(
456 output, mock_request, enable_auto_tools=True
457 )
458 assert (reasoning, content) == ("reasoning", None)
459 assert calls and calls[0].name == "weather"
460 result = collect_output(_stream(parser_cls(mock_tokenizer), mock_tokenizer, output))
461 assert (result.reasoning, result.content) == ("reasoning", "")
462 assert result.tool_calls == [{"name": "weather", "arguments": "{}"}]
463
464
465 def test_every_cut_point_matches_non_streaming(
466 parser_cls, mock_tokenizer, mock_request
467 ):
468 output = (
469 BEGIN_THINK
470 + "Reasoning. "
471 + END_THINK
472 + "Before tools. "
473 + BEGIN_TOOL_REQUESTS
474 + _tool_call("weather", '{"city":"東京"}')
475 + END_TOOL_REQUESTS
476 + EOT
477 )
478 ids = mock_tokenizer.encode(output, add_special_tokens=False)
479 for k in range(1, len(ids) + 1):
480 prefix = mock_tokenizer.decode(ids[:k])
481 reasoning, content, calls = parser_cls(mock_tokenizer).parse(
482 prefix, mock_request, enable_auto_tools=True
483 )
484 result = collect_output(
485 _stream(parser_cls(mock_tokenizer), mock_tokenizer, prefix)
486 )
487 assert (result.reasoning, result.content) == (
488 (reasoning or ""),
489 (content or ""),
490 ), k
491 expected = [{"name": c.name, "arguments": c.arguments} for c in (calls or [])]
492 assert result.tool_calls == expected, (k, expected, result.tool_calls)
493
494
495 @pytest.mark.parametrize(
496 "choice",
497 [
498 "auto",
499 "required",
500 "none",
501 {"type": "function", "function": {"name": "weather"}},
502 ],
503 )
504 def test_tool_choice_uses_the_plamo_format(parser_cls, mock_tokenizer, choice):
505 request = _test_request(
506 [
507 {
508 "type": "function",
509 "function": {"name": "weather", "parameters": {"type": "object"}},
510 },
511 ]
512 )
513 request = type(request).model_validate(
514 request.model_dump() | {"tool_choice": choice}
515 )
516 output = (
517 END_THINK
518 + "Before tools."
519 + BEGIN_TOOL_REQUESTS
520 + _tool_call("weather", "{}")
521 + END_TOOL_REQUESTS
522 )
523 _, content, calls = parser_cls(mock_tokenizer).parse(
524 output, request, enable_auto_tools=True
525 )
526 assert content == "Before tools."
527 assert [c.name for c in calls or []] == ([] if choice == "none" else ["weather"])
528 parser = parser_cls(mock_tokenizer)
529 delta = parser.parse_delta(output, [], request, finished=True)
530 result = collect_output([delta])
531 assert result.content == "Before tools."
532 assert result.tool_calls == (
533 [] if choice == "none" else [{"name": "weather", "arguments": "{}"}]
534 )
535
536
537 @pytest.mark.parametrize("delay", [1, 3])
538 def test_delayed_text_keeps_split_markers_intact(parser_cls, mock_tokenizer, delay):
539 output = (
540 BEGIN_THINK
541 + "reasoning"
542 + END_THINK
543 + "content"
544 + BEGIN_TOOL_REQUESTS
545 + _tool_call("weather", "{}")
546 + END_TOOL_REQUESTS
547 + EOT
548 )
549 tokens = [(i, mock_tokenizer.decode([i])) for i in mock_tokenizer.encode(output)]
550 result = collect_output(
551 replay_with_text_holdback(
552 parser_cls(mock_tokenizer), tokens, text_delay=delay, tools=DUMMY_TOOLS
553 )
554 )
555 assert (result.reasoning, result.content) == ("reasoning", "content")
556 assert result.tool_calls == [{"name": "weather", "arguments": "{}"}]