| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Incremental text lexer that converts text chunks into terminal |
| 4 | tokens, with prefix-match buffering for ambiguous boundaries.""" |
| 5 | |
| 6 | from __future__ import annotations |
| 7 | |
| 8 | from collections.abc import Mapping |
| 9 | from dataclasses import dataclass |
| 10 | |
| 11 | import regex as re |
| 12 | |
| 13 | CONTENT_TERMINAL = "__CONTENT__" |
| 14 | |
| 15 | |
| 16 | @dataclass(slots=True) |
| 17 | class TerminalDef: |
| 18 | name: str |
| 19 | pattern: re.Pattern[str] |
| 20 | is_literal: bool = False |
| 21 | literal: str = "" |
| 22 | |
| 23 | |
| 24 | @dataclass(slots=True) |
| 25 | class LexToken: |
| 26 | terminal: str |
| 27 | value: str |
| 28 | token_count: int = 0 |
| 29 | |
| 30 | |
| 31 | class LexerShape: |
| 32 | """Immutable pre-computed data derived from terminal definitions. |
| 33 | |
| 34 | Created once per :class:`ParserEngineConfig` and shared across all |
| 35 | :class:`IncrementalLexer` instances that use the same config. |
| 36 | """ |
| 37 | |
| 38 | __slots__ = ( |
| 39 | "terminals", |
| 40 | "literal_strings", |
| 41 | "max_literal_len", |
| 42 | "literal_first_chars", |
| 43 | "has_only_literals", |
| 44 | "prefix_set", |
| 45 | "literals_by_first", |
| 46 | ) |
| 47 | |
| 48 | def __init__(self, terminals: list[TerminalDef]) -> None: |
| 49 | self.terminals = sorted( |
| 50 | terminals, |
| 51 | key=lambda t: (not t.is_literal, -len(t.pattern.pattern)), |
| 52 | ) |
| 53 | literal_strings: list[tuple[str, str]] = [] |
| 54 | for t in self.terminals: |
| 55 | if t.is_literal: |
| 56 | literal_strings.append((t.literal, t.name)) |
| 57 | |
| 58 | self.literal_strings = literal_strings |
| 59 | max_len = 0 |
| 60 | for lit, _ in literal_strings: |
| 61 | if len(lit) > max_len: |
| 62 | max_len = len(lit) |
| 63 | self.max_literal_len = max_len |
| 64 | self.literal_first_chars = frozenset( |
| 65 | lit[0] for lit, _ in literal_strings if lit |
| 66 | ) |
| 67 | self.has_only_literals = all(t.is_literal for t in terminals) |
| 68 | |
| 69 | prefix_set: set[str] = set() |
| 70 | for lit, _ in literal_strings: |
| 71 | for i in range(1, len(lit)): |
| 72 | prefix_set.add(lit[:i]) |
| 73 | self.prefix_set = frozenset(prefix_set) |
| 74 | |
| 75 | by_first: dict[str, list[tuple[str, str]]] = {} |
| 76 | for lit, name in literal_strings: |
| 77 | if lit: |
| 78 | by_first.setdefault(lit[0], []).append((lit, name)) |
| 79 | self.literals_by_first = by_first |
| 80 | |
| 81 | |
| 82 | class IncrementalLexer: |
| 83 | """Converts streaming text into terminal tokens. |
| 84 | |
| 85 | The key feature is **prefix-match buffering**: when the text in the |
| 86 | buffer could be the start of a multi-character terminal (e.g. |
| 87 | ``"<tool_"`` that could become ``"<tool_call>"``), the lexer holds |
| 88 | the text rather than emitting it. When the next chunk arrives, it |
| 89 | either completes the terminal or flushes the buffered text as |
| 90 | content. |
| 91 | |
| 92 | Terminals are tried in priority order (literals first, then by |
| 93 | descending priority, then by pattern length). |
| 94 | """ |
| 95 | |
| 96 | def __init__( |
| 97 | self, |
| 98 | terminals: list[TerminalDef] | LexerShape, |
| 99 | content_terminal: str = CONTENT_TERMINAL, |
| 100 | ) -> None: |
| 101 | if isinstance(terminals, LexerShape): |
| 102 | shape = terminals |
| 103 | else: |
| 104 | shape = LexerShape(terminals) |
| 105 | self._shape = shape |
| 106 | self.terminals = shape.terminals |
| 107 | self.content_terminal = content_terminal |
| 108 | self.buffer = "" |
| 109 | self._token_counts: list[int] = [] |
| 110 | |
| 111 | self._literal_strings = shape.literal_strings |
| 112 | self._max_literal_len = shape.max_literal_len |
| 113 | self._literal_first_chars = shape.literal_first_chars |
| 114 | self._has_only_literals = shape.has_only_literals |
| 115 | self._prefix_set = shape.prefix_set |
| 116 | self._literals_by_first = shape.literals_by_first |
| 117 | |
| 118 | def reset(self) -> None: |
| 119 | self.buffer = "" |
| 120 | self._token_counts.clear() |
| 121 | |
| 122 | def feed( |
| 123 | self, |
| 124 | text: str, |
| 125 | token_texts: tuple[str, ...] = (), |
| 126 | token_count: int = 0, |
| 127 | ) -> list[LexToken]: |
| 128 | char_token_counts = self._char_token_counts(text, token_texts, token_count) |
| 129 | if not self.buffer and self._has_only_literals and self._literal_first_chars: |
| 130 | for ch in text: |
| 131 | if ch in self._literal_first_chars: |
| 132 | break |
| 133 | else: |
| 134 | return [LexToken(self.content_terminal, text, sum(char_token_counts))] |
| 135 | self.buffer += text |
| 136 | self._token_counts.extend(char_token_counts) |
| 137 | return self._drain() |
| 138 | |
| 139 | def flush(self) -> list[LexToken]: |
| 140 | tokens: list[LexToken] = [] |
| 141 | if self.buffer: |
| 142 | tokens.extend(self._drain(final=True)) |
| 143 | if self.buffer: |
| 144 | tokens.append( |
| 145 | LexToken(self.content_terminal, self.buffer, sum(self._token_counts)) |
| 146 | ) |
| 147 | self.buffer = "" |
| 148 | self._token_counts.clear() |
| 149 | return tokens |
| 150 | |
| 151 | @staticmethod |
| 152 | def _char_token_counts( |
| 153 | text: str, |
| 154 | token_texts: tuple[str, ...], |
| 155 | token_count: int, |
| 156 | ) -> list[int]: |
| 157 | counts = [0] * len(text) |
| 158 | if not text: |
| 159 | return counts |
| 160 | if token_texts: |
| 161 | pos = 0 |
| 162 | assigned = 0 |
| 163 | for token_text in token_texts: |
| 164 | if not token_text: |
| 165 | continue |
| 166 | found = text.find(token_text, pos) |
| 167 | if found < 0: |
| 168 | continue |
| 169 | counts[found] += 1 |
| 170 | assigned += 1 |
| 171 | pos = found + len(token_text) |
| 172 | missing = token_count - assigned |
| 173 | if missing > 0: |
| 174 | counts[0] += missing |
| 175 | elif token_count: |
| 176 | counts[0] = token_count |
| 177 | return counts |
| 178 | |
| 179 | def _pop_token_count(self, length: int) -> int: |
| 180 | token_count = sum(self._token_counts[:length]) |
| 181 | del self._token_counts[:length] |
| 182 | return token_count |
| 183 | |
| 184 | def _drain(self, *, final: bool = False) -> list[LexToken]: |
| 185 | tokens: list[LexToken] = [] |
| 186 | first_chars = self._literal_first_chars |
| 187 | content_terminal = self.content_terminal |
| 188 | has_only_literals = self._has_only_literals |
| 189 | literals_by_first = self._literals_by_first |
| 190 | prefix_set = self._prefix_set |
| 191 | |
| 192 | while self.buffer: |
| 193 | if has_only_literals and first_chars: |
| 194 | has_potential = False |
| 195 | for ch in self.buffer: |
| 196 | if ch in first_chars: |
| 197 | has_potential = True |
| 198 | break |
| 199 | if not has_potential: |
| 200 | tokens.append( |
| 201 | LexToken(content_terminal, self.buffer, sum(self._token_counts)) |
| 202 | ) |
| 203 | self.buffer = "" |
| 204 | self._token_counts.clear() |
| 205 | break |
| 206 | |
| 207 | best_match: tuple[str, str, int] | None = None |
| 208 | |
| 209 | first = self.buffer[0] |
| 210 | for lit, name in literals_by_first.get(first, ()): |
| 211 | if self.buffer.startswith(lit) and ( |
| 212 | best_match is None or len(lit) > best_match[2] |
| 213 | ): |
| 214 | best_match = (name, lit, len(lit)) |
| 215 | |
| 216 | # If the current buffer is both a complete literal and the prefix |
| 217 | # of a longer literal, wait for the next chunk. For example, |
| 218 | # "<invoke name=" should not be emitted before the next chunk |
| 219 | # proves whether this is the quoted form '<invoke name="'. |
| 220 | if self.buffer in prefix_set and not final: |
| 221 | if best_match is not None: |
| 222 | longer_match = False |
| 223 | for lit, _ in literals_by_first.get(first, ()): |
| 224 | if len(lit) > best_match[2] and lit.startswith(self.buffer): |
| 225 | longer_match = True |
| 226 | break |
| 227 | if not longer_match: |
| 228 | tokens.append( |
| 229 | LexToken( |
| 230 | best_match[0], |
| 231 | best_match[1], |
| 232 | self._pop_token_count(best_match[2]), |
| 233 | ) |
| 234 | ) |
| 235 | self.buffer = self.buffer[best_match[2] :] |
| 236 | continue |
| 237 | break |
| 238 | else: |
| 239 | break |
| 240 | |
| 241 | if best_match is not None: |
| 242 | tokens.append( |
| 243 | LexToken( |
| 244 | best_match[0], |
| 245 | best_match[1], |
| 246 | self._pop_token_count(best_match[2]), |
| 247 | ) |
| 248 | ) |
| 249 | self.buffer = self.buffer[best_match[2] :] |
| 250 | else: |
| 251 | content_end = self._find_content_boundary() |
| 252 | if content_end > 0: |
| 253 | tokens.append( |
| 254 | LexToken( |
| 255 | content_terminal, |
| 256 | self.buffer[:content_end], |
| 257 | self._pop_token_count(content_end), |
| 258 | ) |
| 259 | ) |
| 260 | self.buffer = self.buffer[content_end:] |
| 261 | else: |
| 262 | tokens.append( |
| 263 | LexToken( |
| 264 | content_terminal, |
| 265 | self.buffer[0], |
| 266 | self._pop_token_count(1), |
| 267 | ) |
| 268 | ) |
| 269 | self.buffer = self.buffer[1:] |
| 270 | |
| 271 | return tokens |
| 272 | |
| 273 | def _find_content_boundary(self) -> int: |
| 274 | buf = self.buffer |
| 275 | n = len(buf) |
| 276 | first_chars = self._literal_first_chars |
| 277 | for i in range(1, n): |
| 278 | if buf[i] not in first_chars: |
| 279 | continue |
| 280 | remaining = n - i |
| 281 | for lit, _ in self._literal_strings: |
| 282 | check_len = min(remaining, len(lit)) |
| 283 | if buf[i : i + check_len] == lit[:check_len]: |
| 284 | return i |
| 285 | return n |
| 286 | |
| 287 | |
| 288 | def terminals_from_literals( |
| 289 | literals: Mapping[str, str | tuple[str, ...]], |
| 290 | ) -> list[TerminalDef]: |
| 291 | return [ |
| 292 | TerminalDef( |
| 293 | name=name, |
| 294 | pattern=re.compile(re.escape(lit)), |
| 295 | is_literal=True, |
| 296 | literal=lit, |
| 297 | ) |
| 298 | for name, value in literals.items() |
| 299 | for lit in ((value,) if isinstance(value, str) else value) |
| 300 | ] |