| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Shared streaming simulation helpers for parser engine tests.""" |
| 4 | |
| 5 | from __future__ import annotations |
| 6 | |
| 7 | from typing import Any |
| 8 | |
| 9 | from vllm.entrypoints.generate.base.protocol import DeltaMessage |
| 10 | |
| 11 | |
| 12 | def _build_token_id_map(parser) -> dict[str, int]: |
| 13 | """Map special token text to token IDs from the parser's config.""" |
| 14 | token_id_map: dict[str, int] = {} |
| 15 | cfg = getattr(parser, "parser_engine_config", None) |
| 16 | vocab = getattr(parser, "vocab", None) |
| 17 | if cfg is not None and vocab is not None: |
| 18 | for text in (cfg.token_id_terminals or {}).values(): |
| 19 | tid = vocab.get(text) |
| 20 | if tid is not None: |
| 21 | token_id_map[text] = tid |
| 22 | return token_id_map |
| 23 | |
| 24 | |
| 25 | def simulate_tool_streaming( |
| 26 | parser, |
| 27 | request, |
| 28 | chunks: list[str], |
| 29 | ) -> list[tuple[DeltaMessage | None, str]]: |
| 30 | """Feed text chunks through ``extract_tool_calls_streaming()``.""" |
| 31 | token_id_map = _build_token_id_map(parser) |
| 32 | |
| 33 | results: list[tuple[Any, str]] = [] |
| 34 | previous_text = "" |
| 35 | previous_token_ids: list[int] = [] |
| 36 | |
| 37 | for chunk in chunks: |
| 38 | current_text = previous_text + chunk |
| 39 | |
| 40 | delta_token_ids: list[int] = [ |
| 41 | tid for text, tid in token_id_map.items() if text in chunk |
| 42 | ] |
| 43 | |
| 44 | current_token_ids = previous_token_ids + delta_token_ids |
| 45 | |
| 46 | delta = parser.extract_tool_calls_streaming( |
| 47 | previous_text=previous_text, |
| 48 | current_text=current_text, |
| 49 | delta_text=chunk, |
| 50 | previous_token_ids=tuple(previous_token_ids), |
| 51 | current_token_ids=tuple(current_token_ids), |
| 52 | delta_token_ids=tuple(delta_token_ids), |
| 53 | request=request, |
| 54 | ) |
| 55 | results.append((delta, current_text)) |
| 56 | previous_text = current_text |
| 57 | previous_token_ids = list(current_token_ids) |
| 58 | |
| 59 | return results |
| 60 | |
| 61 | |
| 62 | def collect_tool_arguments( |
| 63 | results: list[tuple[DeltaMessage | None, str]], |
| 64 | ) -> str: |
| 65 | """Concatenate all streamed argument fragments.""" |
| 66 | args_text = "" |
| 67 | for delta, _ in results: |
| 68 | if delta and delta.tool_calls: |
| 69 | for tc in delta.tool_calls: |
| 70 | if tc.function and tc.function.arguments: |
| 71 | args_text += tc.function.arguments |
| 72 | return args_text |
| 73 | |
| 74 | |
| 75 | def collect_content( |
| 76 | results: list[tuple[DeltaMessage | None, str]], |
| 77 | ) -> str: |
| 78 | """Concatenate all streamed content parts.""" |
| 79 | parts: list[str] = [] |
| 80 | for delta, _ in results: |
| 81 | if delta and delta.content: |
| 82 | parts.append(delta.content) |
| 83 | return "".join(parts) |
| 84 | |
| 85 | |
| 86 | def collect_function_name( |
| 87 | results: list[tuple[DeltaMessage | None, str]], |
| 88 | ) -> str | None: |
| 89 | """Return first function name from deltas.""" |
| 90 | for delta, _ in results: |
| 91 | if delta and delta.tool_calls: |
| 92 | for tc in delta.tool_calls: |
| 93 | if tc.function and tc.function.name: |
| 94 | return tc.function.name |
| 95 | return None |
| 96 | |
| 97 | |
| 98 | def simulate_reasoning_streaming( |
| 99 | parser, |
| 100 | chunks: list[str], |
| 101 | delta_token_ids_per_chunk: list[tuple[int, ...]] | None = None, |
| 102 | ) -> tuple[str, str]: |
| 103 | """Feed chunks through ``extract_reasoning_streaming()``. |
| 104 | |
| 105 | Returns ``(reasoning_text, content_text)`` tuple. |
| 106 | """ |
| 107 | token_id_map = ( |
| 108 | _build_token_id_map(parser) if delta_token_ids_per_chunk is None else {} |
| 109 | ) |
| 110 | |
| 111 | reasoning_parts: list[str] = [] |
| 112 | content_parts: list[str] = [] |
| 113 | prev_text = "" |
| 114 | prev_ids: list[int] = [] |
| 115 | for i, chunk in enumerate(chunks): |
| 116 | cur_text = prev_text + chunk |
| 117 | if delta_token_ids_per_chunk is not None: |
| 118 | d_ids = delta_token_ids_per_chunk[i] |
| 119 | else: |
| 120 | d_ids = tuple(tid for text, tid in token_id_map.items() if text in chunk) |
| 121 | cur_ids = prev_ids + list(d_ids) |
| 122 | delta = parser.extract_reasoning_streaming( |
| 123 | previous_text=prev_text, |
| 124 | current_text=cur_text, |
| 125 | delta_text=chunk, |
| 126 | previous_token_ids=tuple(prev_ids), |
| 127 | current_token_ids=tuple(cur_ids), |
| 128 | delta_token_ids=d_ids, |
| 129 | ) |
| 130 | if delta: |
| 131 | if delta.reasoning: |
| 132 | reasoning_parts.append(delta.reasoning) |
| 133 | if delta.content: |
| 134 | content_parts.append(delta.content) |
| 135 | prev_text = cur_text |
| 136 | prev_ids = list(cur_ids) |
| 137 | return "".join(reasoning_parts), "".join(content_parts) |