| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for TokenIDScanner.""" |
| 4 | |
| 5 | from unittest.mock import MagicMock |
| 6 | |
| 7 | import pytest |
| 8 | |
| 9 | from vllm.parser.engine.events import EventType |
| 10 | from vllm.parser.engine.streaming_parser_engine import StreamingParserEngine |
| 11 | from vllm.parser.engine.token_id_scanner import ( |
| 12 | PreLexedTerminal, |
| 13 | TextChunk, |
| 14 | TokenIDScanner, |
| 15 | ) |
| 16 | from vllm.parser.gemma4 import gemma4_config |
| 17 | |
| 18 | CHANNEL_START = "<|channel>" |
| 19 | CHANNEL_END = "<channel|>" |
| 20 | CHANNEL_START_ID = 100 |
| 21 | CHANNEL_END_ID = 101 |
| 22 | REGULAR_TOKEN_ID = 200 |
| 23 | TOOL_START = "<tool_call>" |
| 24 | TOOL_END = "</tool_call>" |
| 25 | TOOL_START_ID = 110 |
| 26 | TOOL_END_ID = 111 |
| 27 | |
| 28 | |
| 29 | @pytest.fixture |
| 30 | def tokenizer(): |
| 31 | tok = MagicMock() |
| 32 | tok.get_vocab.return_value = { |
| 33 | CHANNEL_START: CHANNEL_START_ID, |
| 34 | CHANNEL_END: CHANNEL_END_ID, |
| 35 | } |
| 36 | tok.decode.side_effect = lambda ids: { |
| 37 | CHANNEL_START_ID: CHANNEL_START, |
| 38 | CHANNEL_END_ID: CHANNEL_END, |
| 39 | REGULAR_TOKEN_ID: "regular", |
| 40 | }.get(ids[0], f"<unk:{ids[0]}>") |
| 41 | return tok |
| 42 | |
| 43 | |
| 44 | @pytest.fixture |
| 45 | def scanner(tokenizer): |
| 46 | return TokenIDScanner( |
| 47 | token_id_to_terminal={ |
| 48 | CHANNEL_START_ID: "THINK_START", |
| 49 | CHANNEL_END_ID: "THINK_END", |
| 50 | }, |
| 51 | tokenizer=tokenizer, |
| 52 | ) |
| 53 | |
| 54 | |
| 55 | class TestJoinDecodedTextReturnsStr: |
| 56 | """_join_decoded_text always returns str.""" |
| 57 | |
| 58 | @pytest.fixture |
| 59 | def bare_scanner(self): |
| 60 | return TokenIDScanner({}, tokenizer=None) |
| 61 | |
| 62 | def test_mixed_items(self, bare_scanner): |
| 63 | items = [ |
| 64 | TextChunk("hello "), |
| 65 | PreLexedTerminal("TOOL_START", 42, "<tool_call>"), |
| 66 | TextChunk(" world"), |
| 67 | ] |
| 68 | result = bare_scanner._join_decoded_text(items) |
| 69 | assert isinstance(result, str) |
| 70 | assert result == "hello <tool_call> world" |
| 71 | |
| 72 | def test_empty_list(self, bare_scanner): |
| 73 | result = bare_scanner._join_decoded_text([]) |
| 74 | assert isinstance(result, str) |
| 75 | assert result == "" |
| 76 | |
| 77 | def test_only_text_chunks(self, bare_scanner): |
| 78 | result = bare_scanner._join_decoded_text([TextChunk("abc"), TextChunk("def")]) |
| 79 | assert result == "abcdef" |
| 80 | |
| 81 | |
| 82 | class TestHoldbackTextRecovery: |
| 83 | def test_holdback_text_with_special_token_text_absent(self, scanner): |
| 84 | """Terminal deferred when its text is absent from delta_text.""" |
| 85 | result = scanner.scan( |
| 86 | delta_text="processed is appropriate.", |
| 87 | delta_token_ids=[CHANNEL_END_ID], |
| 88 | ) |
| 89 | |
| 90 | assert len(result) == 0 |
| 91 | |
| 92 | result2 = scanner.scan( |
| 93 | delta_text="<channel|>Understood.", |
| 94 | delta_token_ids=[20, 21], |
| 95 | ) |
| 96 | pre_lexed = [r for r in result2 if isinstance(r, PreLexedTerminal)] |
| 97 | assert len(pre_lexed) == 1 |
| 98 | assert pre_lexed[0].terminal == "THINK_END" |
| 99 | texts = [r.text for r in result2 if isinstance(r, TextChunk)] |
| 100 | combined = "".join(texts) |
| 101 | assert "processed is appropriate." in combined |
| 102 | assert "Understood." in combined |
| 103 | |
| 104 | def test_holdback_text_with_special_token_text_present(self, scanner): |
| 105 | """Hold-back text + special token text both in delta_text.""" |
| 106 | result = scanner.scan( |
| 107 | delta_text="holdback text<channel|>", |
| 108 | delta_token_ids=[CHANNEL_END_ID], |
| 109 | ) |
| 110 | |
| 111 | assert len(result) == 2 |
| 112 | assert isinstance(result[0], TextChunk) |
| 113 | assert result[0].text == "holdback text" |
| 114 | assert isinstance(result[1], PreLexedTerminal) |
| 115 | assert result[1].terminal == "THINK_END" |
| 116 | |
| 117 | def test_no_holdback_text(self, scanner): |
| 118 | """delta_text is exactly the special token text.""" |
| 119 | result = scanner.scan( |
| 120 | delta_text="<channel|>", |
| 121 | delta_token_ids=[CHANNEL_END_ID], |
| 122 | ) |
| 123 | |
| 124 | assert len(result) == 1 |
| 125 | assert isinstance(result[0], PreLexedTerminal) |
| 126 | assert result[0].terminal == "THINK_END" |
| 127 | |
| 128 | def test_empty_delta_text(self, scanner): |
| 129 | """Empty delta_text defers the terminal until text arrives.""" |
| 130 | result = scanner.scan( |
| 131 | delta_text="", |
| 132 | delta_token_ids=[CHANNEL_END_ID], |
| 133 | ) |
| 134 | |
| 135 | assert len(result) == 0 |
| 136 | |
| 137 | flushed = scanner.flush_pending() |
| 138 | assert len(flushed) == 1 |
| 139 | assert isinstance(flushed[0], PreLexedTerminal) |
| 140 | assert flushed[0].terminal == "THINK_END" |
| 141 | |
| 142 | def test_empty_delta_text_drops_individual_decode_text(self, tokenizer): |
| 143 | """Empty delta_text with multiple tokens: all results deferred.""" |
| 144 | tool_start_id = 400 |
| 145 | tok_a = 201 |
| 146 | tok_b = 202 |
| 147 | tokenizer.decode.side_effect = lambda ids: { |
| 148 | tool_start_id: "<|tool_call>", |
| 149 | tok_a: "call:", |
| 150 | tok_b: "get_weather", |
| 151 | }.get(ids[0], "?") |
| 152 | |
| 153 | scanner = TokenIDScanner( |
| 154 | token_id_to_terminal={tool_start_id: "TOOL_START"}, |
| 155 | tokenizer=tokenizer, |
| 156 | ) |
| 157 | |
| 158 | result = scanner.scan( |
| 159 | delta_text="", |
| 160 | delta_token_ids=[tool_start_id, tok_a, tok_b], |
| 161 | ) |
| 162 | |
| 163 | assert len(result) == 0 |
| 164 | |
| 165 | flushed = scanner.flush_pending() |
| 166 | assert len(flushed) == 2 |
| 167 | assert isinstance(flushed[0], PreLexedTerminal) |
| 168 | assert flushed[0].terminal == "TOOL_START" |
| 169 | assert isinstance(flushed[1], TextChunk) |
| 170 | assert flushed[1].text == "" |
| 171 | assert flushed[1].token_count == 2 |
| 172 | |
| 173 | def test_deferred_terminal_preserves_trailing_token_count(self, tokenizer): |
| 174 | tokenizer.decode.side_effect = lambda ids: { |
| 175 | CHANNEL_START_ID: CHANNEL_START, |
| 176 | 201: "alpha", |
| 177 | 202: "beta", |
| 178 | }[ids[0]] |
| 179 | scanner = TokenIDScanner({CHANNEL_START_ID: "THINK_START"}, tokenizer) |
| 180 | |
| 181 | assert scanner.scan("", [CHANNEL_START_ID, 201, 202]) == [] |
| 182 | |
| 183 | result = scanner.scan(f"{CHANNEL_START}alphabeta", []) |
| 184 | assert isinstance(result[0], PreLexedTerminal) |
| 185 | assert isinstance(result[1], TextChunk) |
| 186 | assert result[1].text == "alphabeta" |
| 187 | assert result[1].token_count == 2 |
| 188 | |
| 189 | def test_holdback_before_start_tag(self, scanner): |
| 190 | result = scanner.scan( |
| 191 | delta_text="prefix text<|channel>", |
| 192 | delta_token_ids=[CHANNEL_START_ID], |
| 193 | ) |
| 194 | |
| 195 | assert len(result) == 2 |
| 196 | assert isinstance(result[0], TextChunk) |
| 197 | assert result[0].text == "prefix text" |
| 198 | assert isinstance(result[1], PreLexedTerminal) |
| 199 | assert result[1].terminal == "THINK_START" |
| 200 | |
| 201 | def test_multi_token_batch_special_in_middle(self, scanner, tokenizer): |
| 202 | """Multi-token batch with special token in the middle.""" |
| 203 | tok_a = 201 |
| 204 | tok_b = 202 |
| 205 | tokenizer.decode.side_effect = lambda ids: { |
| 206 | tok_a: "wordA", |
| 207 | tok_b: "wordB", |
| 208 | CHANNEL_END_ID: CHANNEL_END, |
| 209 | }.get(ids[0], "?") |
| 210 | |
| 211 | scanner_multi = TokenIDScanner( |
| 212 | token_id_to_terminal={CHANNEL_END_ID: "THINK_END"}, |
| 213 | tokenizer=tokenizer, |
| 214 | ) |
| 215 | |
| 216 | result = scanner_multi.scan( |
| 217 | delta_text="holdback wordA<channel|> wordB", |
| 218 | delta_token_ids=[tok_a, CHANNEL_END_ID, tok_b], |
| 219 | ) |
| 220 | |
| 221 | texts = [r.text for r in result if isinstance(r, TextChunk)] |
| 222 | terminals = [r.terminal for r in result if isinstance(r, PreLexedTerminal)] |
| 223 | assert "THINK_END" in terminals |
| 224 | assert "holdback wordA" in "".join(texts) |
| 225 | |
| 226 | def test_multi_token_batch_special_token_text_absent(self, scanner, tokenizer): |
| 227 | """Multi-token batch where special token text is absent.""" |
| 228 | tok_a = 201 |
| 229 | tok_b = 202 |
| 230 | tokenizer.decode.side_effect = lambda ids: { |
| 231 | tok_a: "alpha", |
| 232 | tok_b: "beta", |
| 233 | CHANNEL_END_ID: CHANNEL_END, |
| 234 | }.get(ids[0], "?") |
| 235 | |
| 236 | scanner_multi = TokenIDScanner( |
| 237 | token_id_to_terminal={CHANNEL_END_ID: "THINK_END"}, |
| 238 | tokenizer=tokenizer, |
| 239 | ) |
| 240 | |
| 241 | result = scanner_multi.scan( |
| 242 | delta_text="holdback alpha", |
| 243 | delta_token_ids=[tok_a, CHANNEL_END_ID, tok_b], |
| 244 | ) |
| 245 | |
| 246 | assert len(result) == 0 |
| 247 | |
| 248 | result2 = scanner_multi.scan( |
| 249 | delta_text="<channel|> more text", |
| 250 | delta_token_ids=[300], |
| 251 | ) |
| 252 | pre_lexed = [r for r in result2 if isinstance(r, PreLexedTerminal)] |
| 253 | assert len(pre_lexed) == 1 |
| 254 | assert pre_lexed[0].terminal == "THINK_END" |
| 255 | text_chunks = [r for r in result2 if isinstance(r, TextChunk)] |
| 256 | combined = "".join(t.text for t in text_chunks) |
| 257 | assert "holdback alpha" in combined |
| 258 | assert "more text" in combined |
| 259 | |
| 260 | def test_holdback_with_content_after_special_token(self, tokenizer): |
| 261 | """Hold-back + special token + content after in one delta.""" |
| 262 | tok_content = 210 |
| 263 | tokenizer.decode.side_effect = lambda ids: { |
| 264 | CHANNEL_END_ID: CHANNEL_END, |
| 265 | tok_content: "content start", |
| 266 | }.get(ids[0], "?") |
| 267 | |
| 268 | scanner = TokenIDScanner( |
| 269 | token_id_to_terminal={CHANNEL_END_ID: "THINK_END"}, |
| 270 | tokenizer=tokenizer, |
| 271 | ) |
| 272 | |
| 273 | result = scanner.scan( |
| 274 | delta_text="reasoning end.<channel|>content start", |
| 275 | delta_token_ids=[CHANNEL_END_ID, tok_content], |
| 276 | ) |
| 277 | |
| 278 | pre_lexed = [r for r in result if isinstance(r, PreLexedTerminal)] |
| 279 | assert len(pre_lexed) == 1 |
| 280 | assert pre_lexed[0].terminal == "THINK_END" |
| 281 | |
| 282 | text_chunks = [r for r in result if isinstance(r, TextChunk)] |
| 283 | combined = "".join(t.text for t in text_chunks) |
| 284 | assert "reasoning end." in combined |
| 285 | |
| 286 | |
| 287 | class TestEndToEndReasoningHoldback: |
| 288 | """End-to-end engine tests with detokenizer hold-back.""" |
| 289 | |
| 290 | def test_reasoning_content_not_truncated(self): |
| 291 | config = gemma4_config() |
| 292 | tok = MagicMock() |
| 293 | vocab = { |
| 294 | CHANNEL_START: CHANNEL_START_ID, |
| 295 | CHANNEL_END: CHANNEL_END_ID, |
| 296 | } |
| 297 | tok.get_vocab.return_value = vocab |
| 298 | tok.decode.side_effect = lambda ids: { |
| 299 | CHANNEL_START_ID: CHANNEL_START, |
| 300 | CHANNEL_END_ID: CHANNEL_END, |
| 301 | }.get(ids[0], f"tok{ids[0]}") |
| 302 | |
| 303 | engine = StreamingParserEngine(config, tok) |
| 304 | all_events = [] |
| 305 | |
| 306 | all_events.extend(engine.feed(CHANNEL_START, [CHANNEL_START_ID])) |
| 307 | all_events.extend( |
| 308 | engine.feed( |
| 309 | "thought\nThe request was received and ", |
| 310 | [10, 11, 12, 13, 14], |
| 311 | ) |
| 312 | ) |
| 313 | # CHANNEL_END token arrives but its text is held back. |
| 314 | all_events.extend( |
| 315 | engine.feed( |
| 316 | "processed is appropriate.", |
| 317 | [CHANNEL_END_ID], |
| 318 | ) |
| 319 | ) |
| 320 | # Detokenizer flushes the held-back text. |
| 321 | all_events.extend( |
| 322 | engine.feed( |
| 323 | "<channel|>Understood.", |
| 324 | [20, 21], |
| 325 | ) |
| 326 | ) |
| 327 | |
| 328 | all_events.extend(engine.finish()) |
| 329 | |
| 330 | reasoning_text = "".join( |
| 331 | e.value for e in all_events if e.type == EventType.REASONING_CHUNK |
| 332 | ) |
| 333 | content_text = "".join( |
| 334 | e.value for e in all_events if e.type == EventType.TEXT_CHUNK |
| 335 | ) |
| 336 | |
| 337 | assert "processed is appropriate." in reasoning_text |
| 338 | assert "Understood." in content_text |
| 339 | |
| 340 | def test_backtick_content_not_truncated(self): |
| 341 | config = gemma4_config() |
| 342 | tok = MagicMock() |
| 343 | vocab = { |
| 344 | CHANNEL_START: CHANNEL_START_ID, |
| 345 | CHANNEL_END: CHANNEL_END_ID, |
| 346 | } |
| 347 | tok.get_vocab.return_value = vocab |
| 348 | tok.decode.side_effect = lambda ids: { |
| 349 | CHANNEL_START_ID: CHANNEL_START, |
| 350 | CHANNEL_END_ID: CHANNEL_END, |
| 351 | }.get(ids[0], f"tok{ids[0]}") |
| 352 | |
| 353 | engine = StreamingParserEngine(config, tok) |
| 354 | all_events = [] |
| 355 | |
| 356 | all_events.extend(engine.feed(CHANNEL_START, [CHANNEL_START_ID])) |
| 357 | all_events.extend( |
| 358 | engine.feed( |
| 359 | "thought\n1/10 completed. Next: ", |
| 360 | [10, 11, 12, 13], |
| 361 | ) |
| 362 | ) |
| 363 | all_events.extend( |
| 364 | engine.feed( |
| 365 | "`hostname`.\n", |
| 366 | [CHANNEL_END_ID], |
| 367 | ) |
| 368 | ) |
| 369 | all_events.extend( |
| 370 | engine.feed( |
| 371 | "<channel|>tool output", |
| 372 | [20, 21], |
| 373 | ) |
| 374 | ) |
| 375 | |
| 376 | all_events.extend(engine.finish()) |
| 377 | |
| 378 | reasoning_text = "".join( |
| 379 | e.value for e in all_events if e.type == EventType.REASONING_CHUNK |
| 380 | ) |
| 381 | |
| 382 | assert "`hostname`." in reasoning_text |
| 383 | |
| 384 | |
| 385 | _CHANNEL_START_TAG = "<|channel>" |
| 386 | _CHANNEL_END_TAG = "<channel|>" |
| 387 | _TOOL_START_TAG = "<|tool_call>" |
| 388 | _TOOL_END_TAG = "<tool_call|>" |
| 389 | _QUOTE_TAG = '<|"|>' |
| 390 | |
| 391 | _CHANNEL_START_TID = 100 |
| 392 | _CHANNEL_END_TID = 101 |
| 393 | _TOOL_START_TID = 102 |
| 394 | _TOOL_END_TID = 103 |
| 395 | _QUOTE_TID = 104 |
| 396 | _TOK = list(range(200, 215)) |
| 397 | |
| 398 | |
| 399 | def _gemma4_vocab() -> dict[str, int]: |
| 400 | return { |
| 401 | _CHANNEL_START_TAG: _CHANNEL_START_TID, |
| 402 | _CHANNEL_END_TAG: _CHANNEL_END_TID, |
| 403 | _TOOL_START_TAG: _TOOL_START_TID, |
| 404 | _TOOL_END_TAG: _TOOL_END_TID, |
| 405 | _QUOTE_TAG: _QUOTE_TID, |
| 406 | } |
| 407 | |
| 408 | |
| 409 | def _make_gemma4_tokenizer( |
| 410 | extra_decode: dict[int, str] | None = None, |
| 411 | ) -> MagicMock: |
| 412 | special = { |
| 413 | _CHANNEL_START_TID: _CHANNEL_START_TAG, |
| 414 | _CHANNEL_END_TID: _CHANNEL_END_TAG, |
| 415 | _TOOL_START_TID: _TOOL_START_TAG, |
| 416 | _TOOL_END_TID: _TOOL_END_TAG, |
| 417 | _QUOTE_TID: _QUOTE_TAG, |
| 418 | } |
| 419 | decode_map = {**special, **(extra_decode or {})} |
| 420 | |
| 421 | tok = MagicMock() |
| 422 | tok.get_vocab.return_value = _gemma4_vocab() |
| 423 | tok.decode.side_effect = lambda ids: decode_map.get(ids[0], f"tok{ids[0]}") |
| 424 | return tok |
| 425 | |
| 426 | |
| 427 | def _collect_events(engine, deltas): |
| 428 | from vllm.parser.engine.events import SemanticEvent |
| 429 | |
| 430 | all_events: list[SemanticEvent] = [] |
| 431 | for delta_text, delta_token_ids in deltas: |
| 432 | all_events.extend(engine.feed(delta_text, delta_token_ids)) |
| 433 | all_events.extend(engine.finish()) |
| 434 | return all_events |
| 435 | |
| 436 | |
| 437 | def _reasoning_text(events) -> str: |
| 438 | return "".join(e.value for e in events if e.type == EventType.REASONING_CHUNK) |
| 439 | |
| 440 | |
| 441 | def _content_text(events) -> str: |
| 442 | return "".join(e.value for e in events if e.type == EventType.TEXT_CHUNK) |
| 443 | |
| 444 | |
| 445 | def _arg_text(events) -> str: |
| 446 | return "".join(e.value for e in events if e.type == EventType.ARG_VALUE_CHUNK) |
| 447 | |
| 448 | |
| 449 | def _has_event(events, event_type) -> bool: |
| 450 | return any(e.type == event_type for e in events) |
| 451 | |
| 452 | |
| 453 | class TestMultiTokenBoundaryPreservation: |
| 454 | """No text lost at state boundaries with multi-token deltas.""" |
| 455 | |
| 456 | def test_empty_delta_text_at_channel_end_unified(self): |
| 457 | """Empty delta_text when CHANNEL_END arrives; text comes later.""" |
| 458 | tok = _make_gemma4_tokenizer() |
| 459 | engine = StreamingParserEngine(gemma4_config(), tok) |
| 460 | |
| 461 | events = _collect_events( |
| 462 | engine, |
| 463 | [ |
| 464 | ("", [_CHANNEL_START_TID]), |
| 465 | ("<|channel>thought\nSome reasoning.", [_TOK[0], _TOK[1]]), |
| 466 | ("", [_CHANNEL_END_TID]), |
| 467 | ("<channel|>Final answer.", [_TOK[2], _TOK[3]]), |
| 468 | ], |
| 469 | ) |
| 470 | |
| 471 | reasoning = _reasoning_text(events) |
| 472 | content = _content_text(events) |
| 473 | assert "Some reasoning." in reasoning |
| 474 | assert "Final answer." in content |
| 475 | assert _has_event(events, EventType.REASONING_START) |
| 476 | assert _has_event(events, EventType.REASONING_END) |
| 477 | |
| 478 | def test_deferred_channel_end_flushed_at_finish_unified(self): |
| 479 | """Deferred CHANNEL_END flushed at end-of-stream.""" |
| 480 | tok = _make_gemma4_tokenizer() |
| 481 | engine = StreamingParserEngine(gemma4_config(), tok) |
| 482 | |
| 483 | events = _collect_events( |
| 484 | engine, |
| 485 | [ |
| 486 | (_CHANNEL_START_TAG, [_CHANNEL_START_TID]), |
| 487 | ("thought\nReasoning text.", [_TOK[0]]), |
| 488 | (" Final thought.", [_CHANNEL_END_TID]), |
| 489 | ], |
| 490 | ) |
| 491 | |
| 492 | reasoning = _reasoning_text(events) |
| 493 | assert "Reasoning text. Final thought." in reasoning |
| 494 | assert _has_event(events, EventType.REASONING_END) |
| 495 | |
| 496 | def test_reasoning_to_tool_call_handoff_unified(self): |
| 497 | """Full reasoning -> content -> tool call flow.""" |
| 498 | tok = _make_gemma4_tokenizer() |
| 499 | engine = StreamingParserEngine(gemma4_config(), tok) |
| 500 | |
| 501 | events = _collect_events( |
| 502 | engine, |
| 503 | [ |
| 504 | (_CHANNEL_START_TAG, [_CHANNEL_START_TID]), |
| 505 | ("thought\nI need to check the weather.", [_TOK[0], _TOK[1], _TOK[2]]), |
| 506 | (_CHANNEL_END_TAG, [_CHANNEL_END_TID]), |
| 507 | ("Let me call a tool.", [_TOK[3], _TOK[4]]), |
| 508 | (_TOOL_START_TAG, [_TOOL_START_TID]), |
| 509 | ("call:get_weather{city:", [_TOK[5], _TOK[6]]), |
| 510 | ('<|"|>SF<|"|>}', [_QUOTE_TID, _TOK[7], _QUOTE_TID, _TOK[8]]), |
| 511 | (_TOOL_END_TAG, [_TOOL_END_TID]), |
| 512 | ], |
| 513 | ) |
| 514 | |
| 515 | reasoning = _reasoning_text(events) |
| 516 | content = _content_text(events) |
| 517 | |
| 518 | assert "I need to check the weather." in reasoning |
| 519 | assert "Let me call a tool." in content |
| 520 | assert _has_event(events, EventType.REASONING_START) |
| 521 | assert _has_event(events, EventType.REASONING_END) |
| 522 | assert _has_event(events, EventType.TOOL_CALL_START) |
| 523 | assert _has_event(events, EventType.TOOL_CALL_END) |
| 524 | assert "SF" in _arg_text(events) |
| 525 | |
| 526 | def test_multiple_tool_calls_rapid_transitions_unified(self): |
| 527 | """Two back-to-back tool calls with correct tool_index tracking.""" |
| 528 | tok = _make_gemma4_tokenizer() |
| 529 | engine = StreamingParserEngine(gemma4_config(), tok) |
| 530 | |
| 531 | events = _collect_events( |
| 532 | engine, |
| 533 | [ |
| 534 | (_TOOL_START_TAG, [_TOOL_START_TID]), |
| 535 | ("call:get_weather{city:", [_TOK[0], _TOK[1]]), |
| 536 | ('<|"|>NYC<|"|>}', [_QUOTE_TID, _TOK[2], _QUOTE_TID, _TOK[3]]), |
| 537 | (_TOOL_END_TAG, [_TOOL_END_TID]), |
| 538 | (_TOOL_START_TAG, [_TOOL_START_TID]), |
| 539 | ("call:get_time{tz:", [_TOK[4], _TOK[5]]), |
| 540 | ('<|"|>EST<|"|>}', [_QUOTE_TID, _TOK[6], _QUOTE_TID, _TOK[7]]), |
| 541 | (_TOOL_END_TAG, [_TOOL_END_TID]), |
| 542 | ], |
| 543 | ) |
| 544 | |
| 545 | starts = [e for e in events if e.type == EventType.TOOL_CALL_START] |
| 546 | ends = [e for e in events if e.type == EventType.TOOL_CALL_END] |
| 547 | assert len(starts) == 2 |
| 548 | assert len(ends) == 2 |
| 549 | assert starts[0].tool_index == 0 |
| 550 | assert starts[1].tool_index == 1 |
| 551 | |
| 552 | names = "".join(e.value for e in events if e.type == EventType.TOOL_NAME) |
| 553 | assert "get_weather" in names |
| 554 | assert "get_time" in names |
| 555 | |
| 556 | def test_deferred_channel_end_before_tool_call_unified(self): |
| 557 | """Deferred CHANNEL_END followed by a tool call.""" |
| 558 | tok = _make_gemma4_tokenizer() |
| 559 | engine = StreamingParserEngine(gemma4_config(), tok) |
| 560 | |
| 561 | events = _collect_events( |
| 562 | engine, |
| 563 | [ |
| 564 | (_CHANNEL_START_TAG, [_CHANNEL_START_TID]), |
| 565 | ("thought\nNeed to call a tool.", [_TOK[0], _TOK[1]]), |
| 566 | (" Let me proceed.", [_CHANNEL_END_TID]), |
| 567 | (_CHANNEL_END_TAG, [_TOK[2]]), |
| 568 | (_TOOL_START_TAG, [_TOOL_START_TID]), |
| 569 | ("call:get_weather{city:", [_TOK[3], _TOK[4]]), |
| 570 | ('<|"|>Tokyo<|"|>}', [_QUOTE_TID, _TOK[5], _QUOTE_TID, _TOK[6]]), |
| 571 | (_TOOL_END_TAG, [_TOOL_END_TID]), |
| 572 | ], |
| 573 | ) |
| 574 | |
| 575 | reasoning = _reasoning_text(events) |
| 576 | assert "Need to call a tool. Let me proceed." in reasoning |
| 577 | assert _has_event(events, EventType.REASONING_END) |
| 578 | assert _has_event(events, EventType.TOOL_CALL_START) |
| 579 | assert _has_event(events, EventType.TOOL_CALL_END) |
| 580 | assert "Tokyo" in _arg_text(events) |
| 581 | |
| 582 | |
| 583 | class TestStreamInterval10: |
| 584 | """Tests with stream_interval=10 (large multi-token batches).""" |
| 585 | |
| 586 | def test_channel_end_mid_batch_text_present(self): |
| 587 | """<channel|> mid-batch with its text present in delta_text.""" |
| 588 | tok = _make_gemma4_tokenizer({_TOK[i]: f"word{i} " for i in range(15)}) |
| 589 | engine = StreamingParserEngine(gemma4_config(), tok) |
| 590 | |
| 591 | events: list = [] |
| 592 | events.extend( |
| 593 | engine.feed( |
| 594 | "<|channel>thought\nword0 word1 word2 word3 word4 " |
| 595 | "word5 word6 word7 word8 ", |
| 596 | [ |
| 597 | _CHANNEL_START_TID, |
| 598 | _TOK[0], |
| 599 | _TOK[1], |
| 600 | _TOK[2], |
| 601 | _TOK[3], |
| 602 | _TOK[4], |
| 603 | _TOK[5], |
| 604 | _TOK[6], |
| 605 | _TOK[7], |
| 606 | _TOK[8], |
| 607 | ], |
| 608 | ) |
| 609 | ) |
| 610 | |
| 611 | events.extend( |
| 612 | engine.feed( |
| 613 | "word9 word10 word11 <channel|>word12 word13 word14 word0 word1 word2 ", |
| 614 | [ |
| 615 | _TOK[9], |
| 616 | _TOK[10], |
| 617 | _TOK[11], |
| 618 | _CHANNEL_END_TID, |
| 619 | _TOK[12], |
| 620 | _TOK[13], |
| 621 | _TOK[14], |
| 622 | _TOK[0], |
| 623 | _TOK[1], |
| 624 | _TOK[2], |
| 625 | ], |
| 626 | ) |
| 627 | ) |
| 628 | |
| 629 | events.extend(engine.finish()) |
| 630 | |
| 631 | reasoning = _reasoning_text(events) |
| 632 | content = _content_text(events) |
| 633 | |
| 634 | for w in ("word9", "word10", "word11"): |
| 635 | assert w in reasoning, f"{w!r} missing from reasoning" |
| 636 | |
| 637 | for w in ("word12", "word13", "word14"): |
| 638 | assert w in content, f"{w!r} missing from content" |
| 639 | |
| 640 | assert _has_event(events, EventType.REASONING_END) |
| 641 | |
| 642 | def test_channel_end_and_tool_start_same_batch_unified(self): |
| 643 | """Both <channel|> and <|tool_call> in a single batch.""" |
| 644 | tok = _make_gemma4_tokenizer({_TOK[i]: f"w{i} " for i in range(15)}) |
| 645 | engine = StreamingParserEngine(gemma4_config(), tok) |
| 646 | |
| 647 | events: list = [] |
| 648 | |
| 649 | events.extend( |
| 650 | engine.feed( |
| 651 | "<|channel>thought\nw0 w1 w2 w3 w4 w5 w6 w7 w8 ", |
| 652 | [ |
| 653 | _CHANNEL_START_TID, |
| 654 | _TOK[0], |
| 655 | _TOK[1], |
| 656 | _TOK[2], |
| 657 | _TOK[3], |
| 658 | _TOK[4], |
| 659 | _TOK[5], |
| 660 | _TOK[6], |
| 661 | _TOK[7], |
| 662 | _TOK[8], |
| 663 | ], |
| 664 | ) |
| 665 | ) |
| 666 | |
| 667 | events.extend( |
| 668 | engine.feed( |
| 669 | "w9 w10 <channel|>w11 <|tool_call>", |
| 670 | [ |
| 671 | _TOK[9], |
| 672 | _TOK[10], |
| 673 | _CHANNEL_END_TID, |
| 674 | _TOK[11], |
| 675 | _TOOL_START_TID, |
| 676 | _TOK[12], |
| 677 | _TOK[13], |
| 678 | _TOK[14], |
| 679 | _TOK[0], |
| 680 | _TOK[1], |
| 681 | ], |
| 682 | ) |
| 683 | ) |
| 684 | events.extend(engine.finish()) |
| 685 | |
| 686 | reasoning = _reasoning_text(events) |
| 687 | |
| 688 | assert "w9" in reasoning |
| 689 | assert "w10" in reasoning |
| 690 | assert _has_event(events, EventType.REASONING_END) |
| 691 | assert _has_event(events, EventType.TOOL_CALL_START) |
| 692 | |
| 693 | def test_channel_end_mid_batch_text_absent(self): |
| 694 | """<channel|> mid-batch with its text absent from delta_text.""" |
| 695 | tok = _make_gemma4_tokenizer({_TOK[i]: f"word{i} " for i in range(15)}) |
| 696 | engine = StreamingParserEngine(gemma4_config(), tok) |
| 697 | |
| 698 | events: list = [] |
| 699 | events.extend( |
| 700 | engine.feed( |
| 701 | "<|channel>thought\nword0 word1 word2 word3 word4 " |
| 702 | "word5 word6 word7 word8 ", |
| 703 | [ |
| 704 | _CHANNEL_START_TID, |
| 705 | _TOK[0], |
| 706 | _TOK[1], |
| 707 | _TOK[2], |
| 708 | _TOK[3], |
| 709 | _TOK[4], |
| 710 | _TOK[5], |
| 711 | _TOK[6], |
| 712 | _TOK[7], |
| 713 | _TOK[8], |
| 714 | ], |
| 715 | ) |
| 716 | ) |
| 717 | |
| 718 | events.extend( |
| 719 | engine.feed( |
| 720 | "word9 word10 word11 ", |
| 721 | [ |
| 722 | _TOK[9], |
| 723 | _TOK[10], |
| 724 | _TOK[11], |
| 725 | _CHANNEL_END_TID, |
| 726 | _TOK[12], |
| 727 | _TOK[13], |
| 728 | _TOK[14], |
| 729 | _TOK[0], |
| 730 | _TOK[1], |
| 731 | _TOK[2], |
| 732 | ], |
| 733 | ) |
| 734 | ) |
| 735 | |
| 736 | events.extend( |
| 737 | engine.feed( |
| 738 | "<channel|>word12 word13 word14 word0 word1 word2 ", |
| 739 | [_TOK[3], _TOK[4], _TOK[5]], |
| 740 | ) |
| 741 | ) |
| 742 | |
| 743 | events.extend(engine.finish()) |
| 744 | |
| 745 | reasoning = _reasoning_text(events) |
| 746 | content = _content_text(events) |
| 747 | |
| 748 | for w in ("word9", "word10", "word11"): |
| 749 | assert w in reasoning, f"{w!r} missing from reasoning" |
| 750 | |
| 751 | for w in ("word12", "word13", "word14"): |
| 752 | assert w in content, f"{w!r} missing from content" |
| 753 | |
| 754 | assert _has_event(events, EventType.REASONING_END) |
| 755 | |
| 756 | def test_tool_end_mid_batch_text_absent_unified(self): |
| 757 | """<tool_call|> mid-batch with text absent.""" |
| 758 | tok = _make_gemma4_tokenizer({_TOK[i]: f"w{i}" for i in range(15)}) |
| 759 | engine = StreamingParserEngine(gemma4_config(), tok) |
| 760 | |
| 761 | events: list = [] |
| 762 | events.extend( |
| 763 | engine.feed( |
| 764 | _CHANNEL_START_TAG, |
| 765 | [_CHANNEL_START_TID], |
| 766 | ) |
| 767 | ) |
| 768 | events.extend( |
| 769 | engine.feed( |
| 770 | "thought\nNeed a tool.", |
| 771 | [_TOK[0], _TOK[1]], |
| 772 | ) |
| 773 | ) |
| 774 | events.extend( |
| 775 | engine.feed( |
| 776 | _TOOL_START_TAG, |
| 777 | [_TOOL_START_TID], |
| 778 | ) |
| 779 | ) |
| 780 | events.extend( |
| 781 | engine.feed( |
| 782 | "call:get_weather{city:", |
| 783 | [_TOK[2], _TOK[3], _TOK[4]], |
| 784 | ) |
| 785 | ) |
| 786 | |
| 787 | events.extend( |
| 788 | engine.feed( |
| 789 | '<|"|>San Francisco<|"|>}', |
| 790 | [ |
| 791 | _QUOTE_TID, |
| 792 | _TOK[5], |
| 793 | _TOK[6], |
| 794 | _QUOTE_TID, |
| 795 | _TOK[7], |
| 796 | _TOOL_END_TID, |
| 797 | _TOK[8], |
| 798 | _TOK[9], |
| 799 | _TOK[10], |
| 800 | _TOK[11], |
| 801 | ], |
| 802 | ) |
| 803 | ) |
| 804 | |
| 805 | events.extend( |
| 806 | engine.feed( |
| 807 | "<tool_call|>w8w9w10w11w12", |
| 808 | [_TOK[12], _TOK[13]], |
| 809 | ) |
| 810 | ) |
| 811 | |
| 812 | events.extend(engine.finish()) |
| 813 | |
| 814 | assert _has_event(events, EventType.TOOL_CALL_END) |
| 815 | assert "San Francisco" in _arg_text(events) |
| 816 | |
| 817 | def test_large_batch_holdback_spans_two_batches(self): |
| 818 | """Holdback text spanning two batches with <channel|> in the second.""" |
| 819 | tok = _make_gemma4_tokenizer({_TOK[i]: f"w{i} " for i in range(15)}) |
| 820 | engine = StreamingParserEngine(gemma4_config(), tok) |
| 821 | |
| 822 | events: list = [] |
| 823 | |
| 824 | events.extend( |
| 825 | engine.feed( |
| 826 | "<|channel>thought\nThe user asked about machine learning " |
| 827 | "and I need to think about the best approach to", |
| 828 | [ |
| 829 | _CHANNEL_START_TID, |
| 830 | _TOK[0], |
| 831 | _TOK[1], |
| 832 | _TOK[2], |
| 833 | _TOK[3], |
| 834 | _TOK[4], |
| 835 | _TOK[5], |
| 836 | _TOK[6], |
| 837 | _TOK[7], |
| 838 | _TOK[8], |
| 839 | ], |
| 840 | ) |
| 841 | ) |
| 842 | |
| 843 | events.extend( |
| 844 | engine.feed( |
| 845 | " explain this complex topic. Let me organize my thoughts.", |
| 846 | [ |
| 847 | _TOK[9], |
| 848 | _TOK[10], |
| 849 | _TOK[11], |
| 850 | _TOK[12], |
| 851 | _TOK[13], |
| 852 | _TOK[14], |
| 853 | _CHANNEL_END_TID, |
| 854 | _TOK[0], |
| 855 | _TOK[1], |
| 856 | _TOK[2], |
| 857 | ], |
| 858 | ) |
| 859 | ) |
| 860 | |
| 861 | events.extend( |
| 862 | engine.feed( |
| 863 | "<channel|>w0 w1 w2 Here is what I recommend: start with " |
| 864 | "the fundamentals and build up from there.", |
| 865 | [ |
| 866 | _TOK[3], |
| 867 | _TOK[4], |
| 868 | _TOK[5], |
| 869 | _TOK[6], |
| 870 | _TOK[7], |
| 871 | _TOK[8], |
| 872 | _TOK[9], |
| 873 | _TOK[10], |
| 874 | _TOK[11], |
| 875 | _TOK[12], |
| 876 | ], |
| 877 | ) |
| 878 | ) |
| 879 | |
| 880 | events.extend(engine.finish()) |
| 881 | |
| 882 | reasoning = _reasoning_text(events) |
| 883 | content = _content_text(events) |
| 884 | |
| 885 | assert "organize my thoughts." in reasoning |
| 886 | assert "explain" in reasoning |
| 887 | assert "recommend" in content |
| 888 | |
| 889 | assert _has_event(events, EventType.REASONING_START) |
| 890 | assert _has_event(events, EventType.REASONING_END) |
| 891 | |
| 892 | |
| 893 | class TestRebuildFromAnchorsLiteralLookalike: |
| 894 | """Literal token text in prose must not be consumed as an anchor.""" |
| 895 | |
| 896 | @pytest.fixture |
| 897 | def tool_scanner(self): |
| 898 | tok = MagicMock() |
| 899 | tok.get_vocab.return_value = { |
| 900 | TOOL_START: TOOL_START_ID, |
| 901 | TOOL_END: TOOL_END_ID, |
| 902 | } |
| 903 | tok.decode.side_effect = lambda ids: { |
| 904 | TOOL_START_ID: TOOL_START, |
| 905 | TOOL_END_ID: TOOL_END, |
| 906 | }.get(ids[0], f"t{ids[0]}") |
| 907 | return TokenIDScanner( |
| 908 | {TOOL_START_ID: "TOOL_START", TOOL_END_ID: "TOOL_END"}, |
| 909 | tok, |
| 910 | ) |
| 911 | |
| 912 | def test_literal_before_real_anchor(self, tool_scanner): |
| 913 | delta_text = 'Use <tool_call> like this: <tool_call>{"name":"f"}</tool_call>' |
| 914 | delta_token_ids = [1, 2, 3, 4, 5, TOOL_START_ID, 6, 7, TOOL_END_ID] |
| 915 | items = tool_scanner.scan(delta_text, delta_token_ids) |
| 916 | |
| 917 | text_parts = [it.text for it in items if isinstance(it, TextChunk)] |
| 918 | terminals = [it for it in items if isinstance(it, PreLexedTerminal)] |
| 919 | |
| 920 | assert len(terminals) == 2 |
| 921 | assert terminals[0].terminal == "TOOL_START" |
| 922 | assert terminals[1].terminal == "TOOL_END" |
| 923 | |
| 924 | joined_text = "".join(text_parts) |
| 925 | assert "<tool_call>" in joined_text |
| 926 | assert '{"name":"f"}' in joined_text |
| 927 | |
| 928 | def test_multiple_tool_calls_with_literal_between(self, tool_scanner): |
| 929 | delta_text = ( |
| 930 | '<tool_call>{"name":"a"}</tool_call>' |
| 931 | " see <tool_call> syntax " |
| 932 | '<tool_call>{"name":"b"}</tool_call>' |
| 933 | ) |
| 934 | delta_token_ids = [ |
| 935 | TOOL_START_ID, |
| 936 | 1, |
| 937 | TOOL_END_ID, |
| 938 | 2, |
| 939 | 3, |
| 940 | 4, |
| 941 | TOOL_START_ID, |
| 942 | 5, |
| 943 | TOOL_END_ID, |
| 944 | ] |
| 945 | items = tool_scanner.scan(delta_text, delta_token_ids) |
| 946 | |
| 947 | terminals = [it for it in items if isinstance(it, PreLexedTerminal)] |
| 948 | assert len(terminals) == 4 |
| 949 | |
| 950 | text_parts = [it.text for it in items if isinstance(it, TextChunk)] |
| 951 | joined_text = "".join(text_parts) |
| 952 | assert "<tool_call> syntax" in joined_text |
| 953 | |
| 954 | |
| 955 | class TestRebuildFromAnchorsCascadingDeferral: |
| 956 | """Missing middle anchor defers only itself, not subsequent ones.""" |
| 957 | |
| 958 | @pytest.fixture |
| 959 | def bare_scanner(self): |
| 960 | tok = MagicMock() |
| 961 | tok.decode.side_effect = lambda ids: f"t{ids[0]}" |
| 962 | return TokenIDScanner({}, tok) |
| 963 | |
| 964 | def test_middle_anchor_missing_does_not_cascade(self, bare_scanner): |
| 965 | a = PreLexedTerminal("TOOL_START", TOOL_START_ID, TOOL_START) |
| 966 | b = PreLexedTerminal("THINK_END", CHANNEL_END_ID, CHANNEL_END) |
| 967 | c = PreLexedTerminal("TOOL_END", TOOL_END_ID, TOOL_END) |
| 968 | delta_text = f"prefix{TOOL_START}middle{TOOL_END}suffix" |
| 969 | results = [a, b, c] |
| 970 | |
| 971 | rebuilt = bare_scanner._rebuild_from_anchors(delta_text, results) |
| 972 | |
| 973 | terminals = [r for r in rebuilt if isinstance(r, PreLexedTerminal)] |
| 974 | texts = [r for r in rebuilt if isinstance(r, TextChunk)] |
| 975 | joined = "".join(t.text for t in texts) |
| 976 | |
| 977 | assert len(terminals) == 2 |
| 978 | assert terminals[0].terminal == "TOOL_START" |
| 979 | assert terminals[1].terminal == "TOOL_END" |
| 980 | assert "prefix" in joined |
| 981 | assert "middle" in joined |
| 982 | assert "suffix" in joined |
| 983 | assert len(bare_scanner._deferred_terminals) == 1 |
| 984 | assert bare_scanner._deferred_terminals[0].terminal == "THINK_END" |
| 985 | assert bare_scanner._deferred_post_text == "" |
| 986 | |
| 987 | def test_first_anchor_missing_rest_still_emitted(self, bare_scanner): |
| 988 | a = PreLexedTerminal("THINK_END", CHANNEL_END_ID, CHANNEL_END) |
| 989 | b = PreLexedTerminal("TOOL_START", TOOL_START_ID, TOOL_START) |
| 990 | delta_text = f"text{TOOL_START}more" |
| 991 | results = [a, b] |
| 992 | |
| 993 | rebuilt = bare_scanner._rebuild_from_anchors(delta_text, results) |
| 994 | |
| 995 | terminals = [r for r in rebuilt if isinstance(r, PreLexedTerminal)] |
| 996 | assert len(terminals) == 1 |
| 997 | assert terminals[0].terminal == "TOOL_START" |
| 998 | assert len(bare_scanner._deferred_terminals) == 1 |
| 999 | assert bare_scanner._deferred_terminals[0].terminal == "THINK_END" |
| 1000 | |
| 1001 | def test_last_anchor_missing_preceding_still_emitted(self, bare_scanner): |
| 1002 | a = PreLexedTerminal("TOOL_START", TOOL_START_ID, TOOL_START) |
| 1003 | b = PreLexedTerminal("THINK_END", CHANNEL_END_ID, CHANNEL_END) |
| 1004 | delta_text = f"text{TOOL_START}more" |
| 1005 | results = [a, b] |
| 1006 | |
| 1007 | rebuilt = bare_scanner._rebuild_from_anchors(delta_text, results) |
| 1008 | |
| 1009 | terminals = [r for r in rebuilt if isinstance(r, PreLexedTerminal)] |
| 1010 | assert len(terminals) == 1 |
| 1011 | assert terminals[0].terminal == "TOOL_START" |
| 1012 | texts = [r for r in rebuilt if isinstance(r, TextChunk)] |
| 1013 | joined = "".join(t.text for t in texts) |
| 1014 | assert "text" in joined |
| 1015 | assert bare_scanner._deferred_post_text == "more" |
| 1016 | assert len(bare_scanner._deferred_terminals) == 1 |
| 1017 | assert bare_scanner._deferred_terminals[0].terminal == "THINK_END" |