| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Data-driven replay harness for parser engine testing. |
| 4 | |
| 5 | Replays token sequences through parsers at different chunk sizes to |
| 6 | verify chunk-size invariance: the same token sequence must produce |
| 7 | identical output regardless of how tokens are batched. |
| 8 | """ |
| 9 | |
| 10 | from __future__ import annotations |
| 11 | |
| 12 | import json |
| 13 | from collections.abc import Sequence |
| 14 | from dataclasses import dataclass, field |
| 15 | |
| 16 | from vllm.entrypoints.generate.base.protocol import DeltaMessage |
| 17 | from vllm.entrypoints.openai.chat_completion.protocol import ( |
| 18 | ChatCompletionRequest, |
| 19 | ) |
| 20 | |
| 21 | |
| 22 | @dataclass |
| 23 | class Sample: |
| 24 | """One test sample loaded from a JSONL file.""" |
| 25 | |
| 26 | id: str |
| 27 | description: str |
| 28 | source: str |
| 29 | vocab: dict[str, int] |
| 30 | tokens: list[tuple[int, str]] |
| 31 | expected_reasoning: str | None |
| 32 | expected_content: str | None |
| 33 | expected_tool_calls: list[dict] | None |
| 34 | tools: list[dict] | None = None |
| 35 | chat_template_kwargs: dict | None = None |
| 36 | prompt_token_ids: list[int] | None = None |
| 37 | content_lstrip: str | None = None |
| 38 | |
| 39 | |
| 40 | @dataclass |
| 41 | class ParseOutput: |
| 42 | """Accumulated parse output from replaying a token stream.""" |
| 43 | |
| 44 | reasoning: str = "" |
| 45 | content: str = "" |
| 46 | tool_calls: list[dict] = field(default_factory=list) |
| 47 | |
| 48 | |
| 49 | class MockTokenizer: |
| 50 | """Lightweight tokenizer mock that avoids unittest.mock overhead. |
| 51 | |
| 52 | Used by ``benchmarks/benchmark_parsers.py`` in tight timing loops, |
| 53 | so hot-path methods (``decode``, ``get_vocab``) must be cheap. |
| 54 | MagicMock's call-recording machinery added ~40% overhead to small- |
| 55 | sample benchmarks, inflating the per-token cost of the parser engine. |
| 56 | """ |
| 57 | |
| 58 | __slots__ = ( |
| 59 | "_vocab", |
| 60 | "_token_ids", |
| 61 | "_token_decode_map", |
| 62 | "_special_ids", |
| 63 | "eos_token_id", |
| 64 | "bos_token_id", |
| 65 | "pad_token_id", |
| 66 | "_all_special_tokens", |
| 67 | ) |
| 68 | |
| 69 | def __init__( |
| 70 | self, |
| 71 | vocab: dict[str, int], |
| 72 | tokens: list[tuple[int, str]], |
| 73 | ) -> None: |
| 74 | self._vocab = vocab |
| 75 | self._token_ids = [tid for tid, _ in tokens] |
| 76 | self._token_decode_map = {tid: text for tid, text in tokens} |
| 77 | self._special_ids = set(vocab.values()) |
| 78 | self._all_special_tokens = list(vocab.keys()) |
| 79 | self.eos_token_id = None |
| 80 | self.bos_token_id = None |
| 81 | self.pad_token_id = None |
| 82 | |
| 83 | def set_vocab(self, vocab: dict[str, int]) -> None: |
| 84 | self._vocab = vocab |
| 85 | |
| 86 | def get_vocab(self) -> dict[str, int]: |
| 87 | return self._vocab |
| 88 | |
| 89 | @property |
| 90 | def all_special_tokens(self) -> list[str]: |
| 91 | return self._all_special_tokens |
| 92 | |
| 93 | @property |
| 94 | def all_special_ids(self) -> list[int]: |
| 95 | return [self._vocab[t] for t in self._all_special_tokens if t in self._vocab] |
| 96 | |
| 97 | def encode(self, text: str, **kwargs) -> list[int]: |
| 98 | return self._token_ids |
| 99 | |
| 100 | def decode(self, ids: list[int], skip_special_tokens: bool = False) -> str: |
| 101 | parts: list[str] = [] |
| 102 | for tid in ids: |
| 103 | if skip_special_tokens and tid in self._special_ids: |
| 104 | continue |
| 105 | text = self._token_decode_map.get(tid, f"?{tid}?") |
| 106 | parts.append(text) |
| 107 | return "".join(parts) |
| 108 | |
| 109 | |
| 110 | CHUNK_SIZES = [1, 2, 3, 5, 11, 23, None] |
| 111 | |
| 112 | |
| 113 | def make_mock_tokenizer(sample: Sample) -> MockTokenizer: |
| 114 | """Build a mock tokenizer from a sample's vocab and token data.""" |
| 115 | return MockTokenizer( |
| 116 | vocab=dict(sample.vocab), |
| 117 | tokens=sample.tokens, |
| 118 | ) |
| 119 | |
| 120 | |
| 121 | def _test_request( |
| 122 | tools: list[dict] | None = None, |
| 123 | ) -> ChatCompletionRequest: |
| 124 | return ChatCompletionRequest( |
| 125 | model="test-model", |
| 126 | messages=[{"role": "user", "content": "test"}], |
| 127 | tools=tools, |
| 128 | ) |
| 129 | |
| 130 | |
| 131 | DUMMY_TOOLS = [ |
| 132 | { |
| 133 | "type": "function", |
| 134 | "function": {"name": "stub", "parameters": {"type": "object"}}, |
| 135 | }, |
| 136 | ] |
| 137 | |
| 138 | |
| 139 | def parse_non_streaming( |
| 140 | parser, |
| 141 | sample: Sample, |
| 142 | request: ChatCompletionRequest, |
| 143 | ) -> ParseOutput: |
| 144 | """Run ``parser.parse()`` and return a :class:`ParseOutput`.""" |
| 145 | full_text = "".join(text for _, text in sample.tokens) |
| 146 | reasoning, content, tool_calls = parser.parse( |
| 147 | full_text, |
| 148 | request, |
| 149 | enable_auto_tools=True, |
| 150 | ) |
| 151 | tc_list: list[dict] = [] |
| 152 | if tool_calls: |
| 153 | for tc in tool_calls: |
| 154 | tc_list.append({"name": tc.name, "arguments": tc.arguments}) |
| 155 | return ParseOutput( |
| 156 | reasoning=reasoning or "", |
| 157 | content=content or "", |
| 158 | tool_calls=tc_list, |
| 159 | ) |
| 160 | |
| 161 | |
| 162 | def replay_streaming( |
| 163 | parser, |
| 164 | tokens: list[tuple[int, str]], |
| 165 | chunk_size: int | None = None, |
| 166 | holdback_chars: int = 0, |
| 167 | finished_on_last: bool = False, |
| 168 | tools: list[dict] | None = None, |
| 169 | prompt_token_ids: list[int] | None = None, |
| 170 | ) -> list[DeltaMessage | None]: |
| 171 | """Feed tokens through ``parser.parse_delta()`` at a given chunk size. |
| 172 | |
| 173 | Args: |
| 174 | parser: A :class:`Parser` instance with ``parse_delta()`` method. |
| 175 | tokens: List of ``(token_id, decoded_text)`` pairs. |
| 176 | chunk_size: Number of tokens per batch. ``None`` means all at once. |
| 177 | holdback_chars: Simulate detokenizer holdback by holding back |
| 178 | this many characters of decoded text between batches. |
| 179 | finished_on_last: When True, pass ``finished=True`` on the last |
| 180 | ``parse_delta()`` call, matching real server behavior. |
| 181 | tools: Optional tool definitions to include on the request, |
| 182 | matching the serving layer where tools set |
| 183 | ``tool_choice`` to ``"auto"``. |
| 184 | |
| 185 | Returns: |
| 186 | List of ``DeltaMessage`` results from each ``parse_delta()`` call. |
| 187 | |
| 188 | """ |
| 189 | if chunk_size is None: |
| 190 | chunk_size = len(tokens) |
| 191 | |
| 192 | results: list[DeltaMessage | None] = [] |
| 193 | all_ids = [tid for tid, _ in tokens] |
| 194 | all_texts = [text for _, text in tokens] |
| 195 | |
| 196 | request = _test_request(tools=tools) |
| 197 | first_prompt_ids = prompt_token_ids if prompt_token_ids is not None else [] |
| 198 | |
| 199 | if holdback_chars <= 0: |
| 200 | chunks = list(range(0, len(tokens), chunk_size)) |
| 201 | for i, start in enumerate(chunks): |
| 202 | batch_end = min(start + chunk_size, len(tokens)) |
| 203 | batch_ids = all_ids[start:batch_end] |
| 204 | delta_text = "".join(all_texts[start:batch_end]) |
| 205 | is_last = i == len(chunks) - 1 |
| 206 | |
| 207 | result = parser.parse_delta( |
| 208 | delta_text, |
| 209 | batch_ids, |
| 210 | request, |
| 211 | prompt_token_ids=first_prompt_ids if start == 0 else None, |
| 212 | finished=finished_on_last and is_last, |
| 213 | ) |
| 214 | results.append(result) |
| 215 | return results |
| 216 | |
| 217 | emitted_up_to = 0 |
| 218 | is_first = True |
| 219 | |
| 220 | for start in range(0, len(tokens), chunk_size): |
| 221 | batch_end = min(start + chunk_size, len(tokens)) |
| 222 | |
| 223 | if batch_end < len(tokens): |
| 224 | held_chars = 0 |
| 225 | safe_end = batch_end |
| 226 | while safe_end > emitted_up_to and held_chars < holdback_chars: |
| 227 | safe_end -= 1 |
| 228 | held_chars += len(all_texts[safe_end]) |
| 229 | else: |
| 230 | safe_end = batch_end |
| 231 | |
| 232 | if safe_end <= emitted_up_to: |
| 233 | continue |
| 234 | |
| 235 | batch_ids = all_ids[emitted_up_to:safe_end] |
| 236 | delta_text = "".join(all_texts[emitted_up_to:safe_end]) |
| 237 | emitted_up_to = safe_end |
| 238 | |
| 239 | is_last_chunk = batch_end >= len(tokens) |
| 240 | result = parser.parse_delta( |
| 241 | delta_text, |
| 242 | batch_ids, |
| 243 | request, |
| 244 | prompt_token_ids=first_prompt_ids if is_first else None, |
| 245 | finished=finished_on_last and is_last_chunk, |
| 246 | ) |
| 247 | results.append(result) |
| 248 | is_first = False |
| 249 | |
| 250 | if emitted_up_to < len(tokens): |
| 251 | batch_ids = all_ids[emitted_up_to:] |
| 252 | delta_text = "".join(all_texts[emitted_up_to:]) |
| 253 | result = parser.parse_delta( |
| 254 | delta_text, |
| 255 | batch_ids, |
| 256 | request, |
| 257 | prompt_token_ids=first_prompt_ids if is_first else None, |
| 258 | finished=finished_on_last, |
| 259 | ) |
| 260 | results.append(result) |
| 261 | |
| 262 | return results |
| 263 | |
| 264 | |
| 265 | def replay_with_text_holdback( |
| 266 | parser, |
| 267 | tokens: list[tuple[int, str]], |
| 268 | text_delay: int = 1, |
| 269 | tools: list[dict] | None = None, |
| 270 | prompt_token_ids: list[int] | None = None, |
| 271 | ) -> list[DeltaMessage | None]: |
| 272 | """Replay token-by-token with text arriving *text_delay* steps late. |
| 273 | |
| 274 | Simulates the production detokenizer holdback where token IDs arrive |
| 275 | immediately but decoded text is delayed. On the last token all |
| 276 | remaining held-back text is flushed, matching real server behavior:: |
| 277 | |
| 278 | step 0: ids=[tok0], text="" (held back) |
| 279 | step 1: ids=[tok1], text=tok0_text (tok0 released) |
| 280 | ... |
| 281 | step N-1: ids=[tokN-1], text=remaining_texts (flush all) |
| 282 | |
| 283 | This exercises the TokenIDScanner deferred-terminal path that |
| 284 | ``replay_streaming`` (which keeps text and IDs aligned) does not. |
| 285 | """ |
| 286 | results: list[DeltaMessage | None] = [] |
| 287 | request = _test_request(tools=tools) |
| 288 | first_prompt_ids = prompt_token_ids if prompt_token_ids is not None else [] |
| 289 | |
| 290 | n = len(tokens) |
| 291 | held_texts: list[str] = [] |
| 292 | |
| 293 | for i in range(n): |
| 294 | token_id = tokens[i][0] |
| 295 | held_texts.append(tokens[i][1]) |
| 296 | |
| 297 | is_last = i == n - 1 |
| 298 | if is_last: |
| 299 | delta_text = "".join(held_texts) |
| 300 | held_texts.clear() |
| 301 | elif len(held_texts) > text_delay: |
| 302 | delta_text = held_texts.pop(0) |
| 303 | else: |
| 304 | delta_text = "" |
| 305 | |
| 306 | result = parser.parse_delta( |
| 307 | delta_text, |
| 308 | [token_id], |
| 309 | request, |
| 310 | prompt_token_ids=first_prompt_ids if i == 0 else None, |
| 311 | finished=is_last, |
| 312 | ) |
| 313 | results.append(result) |
| 314 | |
| 315 | return results |
| 316 | |
| 317 | |
| 318 | def accumulate_deltas( |
| 319 | deltas: Sequence[DeltaMessage | None], |
| 320 | ) -> dict: |
| 321 | reasoning_parts: list[str] = [] |
| 322 | content_parts: list[str] = [] |
| 323 | tool_calls_by_idx: dict[int, dict] = {} |
| 324 | |
| 325 | for delta in deltas: |
| 326 | if delta is None: |
| 327 | continue |
| 328 | if delta.reasoning: |
| 329 | reasoning_parts.append(delta.reasoning) |
| 330 | if delta.content: |
| 331 | content_parts.append(delta.content) |
| 332 | if delta.tool_calls: |
| 333 | for tc in delta.tool_calls: |
| 334 | if tc.function and tc.function.name: |
| 335 | existing = tool_calls_by_idx.get(tc.index) |
| 336 | if existing is None: |
| 337 | tool_calls_by_idx[tc.index] = { |
| 338 | "name": tc.function.name, |
| 339 | "_args_parts": [tc.function.arguments or ""], |
| 340 | } |
| 341 | else: |
| 342 | existing["_args_parts"].append(tc.function.arguments or "") |
| 343 | elif tc.function and tc.function.arguments: |
| 344 | existing = tool_calls_by_idx.get(tc.index) |
| 345 | if existing is not None: |
| 346 | existing["_args_parts"].append(tc.function.arguments) |
| 347 | |
| 348 | return { |
| 349 | "reasoning": "".join(reasoning_parts), |
| 350 | "content": "".join(content_parts), |
| 351 | "tool_calls": [ |
| 352 | {"name": tc["name"], "arguments": "".join(tc["_args_parts"])} |
| 353 | for tc in tool_calls_by_idx.values() |
| 354 | ], |
| 355 | } |
| 356 | |
| 357 | |
| 358 | def collect_output(results: list[DeltaMessage | None]) -> ParseOutput: |
| 359 | """Accumulate ``DeltaMessage`` results into a :class:`ParseOutput`.""" |
| 360 | result = accumulate_deltas(results) |
| 361 | return ParseOutput( |
| 362 | reasoning=result["reasoning"], |
| 363 | content=result["content"], |
| 364 | tool_calls=result["tool_calls"], |
| 365 | ) |
| 366 | |
| 367 | |
| 368 | def assert_parse_output(actual: ParseOutput, sample: Sample) -> None: |
| 369 | """Compare actual parse output against expected values from a sample.""" |
| 370 | if sample.expected_reasoning is not None: |
| 371 | assert actual.reasoning == sample.expected_reasoning, ( |
| 372 | f"Reasoning mismatch:\n" |
| 373 | f" expected: {sample.expected_reasoning!r}\n" |
| 374 | f" actual: {actual.reasoning!r}" |
| 375 | ) |
| 376 | |
| 377 | if sample.expected_content is not None: |
| 378 | assert actual.content == sample.expected_content, ( |
| 379 | f"Content mismatch:\n" |
| 380 | f" expected: {sample.expected_content!r}\n" |
| 381 | f" actual: {actual.content!r}" |
| 382 | ) |
| 383 | if sample.expected_tool_calls is not None: |
| 384 | assert len(actual.tool_calls) == len(sample.expected_tool_calls), ( |
| 385 | f"Tool call count mismatch: " |
| 386 | f"expected {len(sample.expected_tool_calls)}, " |
| 387 | f"got {len(actual.tool_calls)}" |
| 388 | ) |
| 389 | for i, (expected_tc, actual_tc) in enumerate( |
| 390 | zip(sample.expected_tool_calls, actual.tool_calls) |
| 391 | ): |
| 392 | assert actual_tc["name"] == expected_tc["name"], ( |
| 393 | f"Tool call {i} name mismatch: " |
| 394 | f"expected {expected_tc['name']!r}, " |
| 395 | f"got {actual_tc['name']!r}" |
| 396 | ) |
| 397 | if "arguments" in expected_tc: |
| 398 | expected_args = expected_tc["arguments"] |
| 399 | actual_args_str = actual_tc.get("arguments", "{}") |
| 400 | if isinstance(expected_args, dict): |
| 401 | try: |
| 402 | actual_args = json.loads(actual_args_str) |
| 403 | except json.JSONDecodeError as e: |
| 404 | raise AssertionError( |
| 405 | f"Tool call {i} arguments not valid JSON: " |
| 406 | f"{actual_args_str!r}" |
| 407 | ) from e |
| 408 | assert actual_args == expected_args, ( |
| 409 | f"Tool call {i} arguments mismatch:\n" |
| 410 | f" expected: {expected_args}\n" |
| 411 | f" actual: {actual_args}" |
| 412 | ) |
| 413 | |
| 414 | |
| 415 | def assert_no_terminal_leakage( |
| 416 | actual: ParseOutput, |
| 417 | terminals: list[str], |
| 418 | context: str = "", |
| 419 | ) -> None: |
| 420 | """Assert that none of *terminals* appear in reasoning or content.""" |
| 421 | suffix = f" ({context})" if context else "" |
| 422 | for terminal in terminals: |
| 423 | assert terminal not in actual.reasoning, ( |
| 424 | f"{terminal!r} leaked into reasoning{suffix}" |
| 425 | ) |
| 426 | assert terminal not in actual.content, ( |
| 427 | f"{terminal!r} leaked into content{suffix}" |
| 428 | ) |