class TokenIDScanner:
"""Maps special token IDs in the delta to terminals.
Before text-based lexing happens, the scanner checks each token ID
in the delta against a mapping of ``{token_id: terminal_name}``.
Matched tokens are emitted as :class:`PreLexedTerminal` items;
everything else is grouped into :class:`TextChunk` items for the
incremental lexer to process.
When a terminal's text is not yet in ``delta_text`` (held back by
the detokenizer), the terminal is deferred until the text arrives
in a subsequent delta.
"""
def __init__(
self,
token_id_to_terminal: dict[int, str],
tokenizer,
) -> None:
self.token_id_to_terminal = token_id_to_terminal
self.tokenizer = tokenizer
self._token_text_cache: dict[int, str] = {}
self._deferred_terminals: list[PreLexedTerminal] = []
self._deferred_prefix_token_counts: list[int] = []
self._deferred_trailing_token_count = 0
self._deferred_post_text: str = ""
def reset(self) -> None:
"""Clear mutable state for reuse. Preserves the token text cache."""
self._deferred_terminals.clear()
self._deferred_prefix_token_counts.clear()
self._deferred_trailing_token_count = 0
self._deferred_post_text = ""
def _decode_token(self, token_id: int) -> str:
if token_id not in self._token_text_cache:
token_text = self.tokenizer.decode([token_id])
self._token_text_cache[token_id] = (
token_text if isinstance(token_text, str) else ""
)
return self._token_text_cache[token_id]
def _decode_tokens(self, token_ids: Sequence[int]) -> tuple[str, ...]:
if self.tokenizer is None:
return ()
return tuple(self._decode_token(tid) for tid in token_ids)
_EMPTY: tuple[LexerInput, ...] = ()
def scan(
self,
delta_text: str,
delta_token_ids: Sequence[int],
) -> Sequence[LexerInput]:
prefix_items: list[LexerInput] = []
effective_text = delta_text
deferred_trailing_count = 0
if self._deferred_terminals:
prefix_items, effective_text, deferred_trailing_count = (
self._resolve_deferred(delta_text, len(delta_token_ids))
)
if not self.token_id_to_terminal:
if effective_text:
token_texts = self._decode_tokens(delta_token_ids)
prefix_items.append(
TextChunk(
effective_text,
token_texts,
deferred_trailing_count + len(token_texts),
)
)
return prefix_items
has_special = False
token_id_to_terminal = self.token_id_to_terminal
for tid in delta_token_ids:
if tid in token_id_to_terminal:
has_special = True
break
if not has_special:
if effective_text:
if not prefix_items:
token_texts = self._decode_tokens(delta_token_ids)
return [TextChunk(effective_text, token_texts, len(token_texts))]
token_texts = self._decode_tokens(delta_token_ids)
prefix_items.append(
TextChunk(
effective_text,
token_texts,
deferred_trailing_count + len(token_texts),
)
)
return prefix_items or self._EMPTY
decoded_token_texts = [self._decode_token(tid) for tid in delta_token_ids]
results: list[LexerInput] = []
text_accum: list[str] = []
token_text_accum: list[str] = []
for idx, tid in enumerate(delta_token_ids):
terminal = self.token_id_to_terminal.get(tid)
if terminal is not None:
if text_accum:
joined = "".join(text_accum)
if joined:
results.append(
TextChunk(
joined,
tuple(token_text_accum),
len(token_text_accum),
)
)
text_accum.clear()
token_text_accum.clear()
results.append(
PreLexedTerminal(terminal, tid, decoded_token_texts[idx])
)
else:
text_accum.append(decoded_token_texts[idx])
token_text_accum.append(decoded_token_texts[idx])
if text_accum:
joined = "".join(text_accum)
if joined:
results.append(
TextChunk(joined, tuple(token_text_accum), len(token_text_accum))
)
if deferred_trailing_count:
results.insert(0, TextChunk("", token_count=deferred_trailing_count))
if effective_text:
results = self._recover_holdback_text(effective_text, results)
else:
# No detokenizer text to validate against — individually-decoded
# TextChunks are unreliable (context-dependent decoding).
# Defer PreLexedTerminals so the state machine doesn't
# transition before the preceding text has arrived. The
# deferred terminals will be resolved against the actual
# delta_text in a subsequent scan() or flushed by finish().
prefix_token_count = 0
for r in results:
if isinstance(r, TextChunk):
prefix_token_count += r.token_count
if isinstance(r, PreLexedTerminal):
self._deferred_terminals.append(r)
self._deferred_prefix_token_counts.append(prefix_token_count)
prefix_token_count = 0
self._deferred_trailing_token_count += prefix_token_count
results = []
return prefix_items + results
def flush_pending(self) -> list[LexerInput]:
if not self._deferred_terminals and not self._deferred_post_text:
return []
results: list[LexerInput] = []
if self._deferred_post_text:
prefix_count = (
self._deferred_prefix_token_counts[0]
if self._deferred_prefix_token_counts
else 0
)
results.append(
TextChunk(self._deferred_post_text, token_count=prefix_count)
)
self._deferred_post_text = ""
results.extend(self._deferred_terminals)
if self._deferred_trailing_token_count:
results.append(
TextChunk("", token_count=self._deferred_trailing_token_count)
)
self._deferred_terminals.clear()
self._deferred_prefix_token_counts.clear()
self._deferred_trailing_token_count = 0
return results
def _resolve_deferred(
self,
delta_text: str,
current_token_count: int = 0,
) -> tuple[list[LexerInput], str, int]:
"""Resolve deferred terminals against new delta_text.
When a previous ``scan()`` deferred a terminal (its text hadn't
arrived yet), the next delta's text should contain that terminal's
text. Split delta_text at the terminal boundary: text before
belongs to the previous parser state, the terminal triggers the
state transition, and text after belongs to the new state.
Returns ``(prefix_items, remaining_text)`` where prefix_items
are the resolved deferred terminals (with any preceding text)
and remaining_text is the unconsumed portion of delta_text that
should be scanned with the current delta's token IDs.
"""
deferred = self._deferred_terminals
prefix_token_counts = self._deferred_prefix_token_counts
trailing_token_count = self._deferred_trailing_token_count
self._deferred_terminals = []
self._deferred_prefix_token_counts = []
self._deferred_trailing_token_count = 0
results: list[LexerInput] = []
remaining = delta_text
if self._deferred_post_text:
remaining = self._deferred_post_text + remaining
self._deferred_post_text = ""
# Duplicate-text deferred terminals resolve left-to-right via
# find(); correct when each terminal text appears once in sequence.
for idx, terminal in enumerate(deferred):
prefix_token_count = prefix_token_counts[idx]
pos = remaining.find(terminal.text)
if pos > 0:
results.append(
TextChunk(remaining[:pos], token_count=prefix_token_count)
)
results.append(terminal)
remaining = remaining[pos + len(terminal.text) :]
elif pos == 0:
results.append(terminal)
remaining = remaining[len(terminal.text) :]
else:
# Accumulate text until terminal text arrives —
# only the terminal provides a reliable split point.
if remaining:
self._deferred_post_text += remaining
prefix_token_count += current_token_count
remaining = ""
self._deferred_terminals.append(terminal)
self._deferred_prefix_token_counts.append(prefix_token_count)
if self._deferred_terminals:
self._deferred_trailing_token_count = trailing_token_count
trailing_token_count = 0
return results, remaining, trailing_token_count
def _recover_holdback_text(
self,
delta_text: str,
results: list[LexerInput],
) -> list[LexerInput]:
"""Recover detokenizer hold-back text not in delta_token_ids.
The detokenizer may flush previously held-back text in
``delta_text`` that has no corresponding token ID in
``delta_token_ids``. This hold-back text always appears as a
prefix of ``delta_text``.
"""
if not results:
return [TextChunk(delta_text)]
reconstructed = self._join_decoded_text(results)
if not reconstructed:
return [TextChunk(delta_text)] + results
pos = delta_text.find(reconstructed)
if pos > 0:
return [TextChunk(delta_text[:pos])] + results
if pos == 0:
return results
# Fallback: SentencePiece context-dependent decoding mismatch.
# Rebuild from delta_text using PreLexedTerminals as split anchors.
return self._rebuild_from_anchors(delta_text, results)
def _join_decoded_text(self, results: list[LexerInput]) -> str:
"""Join TextChunk and PreLexedTerminal text into one string."""
parts: list[str] = []
for item in results:
if isinstance(item, (TextChunk, PreLexedTerminal)):
parts.append(item.text)
return "".join(parts)
def _rebuild_from_anchors(
self,
delta_text: str,
results: list[LexerInput],
) -> list[LexerInput]:
"""Rebuild results from delta_text using terminals as anchors.
When context-dependent decoding creates a mismatch between
individually-decoded tokens and delta_text, use
PreLexedTerminals as split points and reallocate text from
delta_text. If a terminal's text is not found in delta_text,
it is deferred to the next scan() call.
Anchors are resolved right-to-left with ``rfind`` so that each
anchor binds to the *rightmost* available occurrence of its
text. This prevents earlier literal lookalikes (e.g. a user
mentioning ``<tool_call>`` in prose) from stealing the position
of a real special-token anchor that appears later.
If the same anchor text appears multiple times as real special
tokens (not prose), the rightmost-first binding could misalign.
In practice this doesn't happen: each special token ID maps to
a distinct PreLexedTerminal, and duplicates in prose are resolved
by the token-ID filtering layer above.
"""
anchors = [item for item in results if isinstance(item, PreLexedTerminal)]
if not anchors:
return [TextChunk(delta_text)]
token_groups: list[list[str]] = [[] for _ in range(len(anchors) + 1)]
count_groups = [0] * (len(anchors) + 1)
group_idx = 0
for item in results:
if isinstance(item, PreLexedTerminal):
group_idx += 1
else:
token_groups[group_idx].extend(item.token_texts)
count_groups[group_idx] += item.token_count
# Resolve positions right-to-left: each anchor gets the
# rightmost occurrence that is still before the next anchor.
positions: list[int] = [-1] * len(anchors)
search_end = len(delta_text)
for i in range(len(anchors) - 1, -1, -1):
pos = delta_text.rfind(anchors[i].text, 0, search_end)
if pos >= 0:
positions[i] = pos
search_end = pos
# Build results left-to-right using the resolved positions.
new_results: list[LexerInput] = []
consumed = 0
for i, anchor in enumerate(anchors):
pos = positions[i]
if pos >= consumed:
if pos > consumed:
new_results.append(
TextChunk(
delta_text[consumed:pos],
tuple(token_groups[i]),
count_groups[i],
)
)
new_results.append(anchor)
consumed = pos + len(anchor.text)
else:
has_later_valid = any(p >= 0 for p in positions[i + 1 :])
# DROP anchors (EOS, etc.) may have text that never
# arrives in delta_text (stripped by detokenizer).
# Don't defer remaining content waiting for text
# that will never come.
if (
not has_later_valid
and consumed < len(delta_text)
and anchor.terminal != DROP_TERMINAL
):
self._deferred_post_text += delta_text[consumed:]
consumed = len(delta_text)
self._deferred_terminals.append(anchor)
self._deferred_prefix_token_counts.append(count_groups[i])
if consumed < len(delta_text):
new_results.append(
TextChunk(
delta_text[consumed:],
tuple(token_groups[-1]),
count_groups[-1],
)
)
return new_results