From c53f03ead38ce41ec883ef6b92bc16c1852d0c99 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Fri, 25 Sep 2026 15:24:54 +0000 Subject: [PATCH] Prove native conditioning before preserving literal reasoning markers --- docs/features/additional-histories.mdx | 19 + src/art/trajectories/_tokenize.py | 307 ++++++++- .../trajectories/test_literal_thinking_off.py | 631 +++++++++++++++++- .../trajectories/test_prefix_render_cache.py | 34 + 4 files changed, 946 insertions(+), 45 deletions(-) diff --git a/docs/features/additional-histories.mdx b/docs/features/additional-histories.mdx index 24802ace7..6ab1b3c52 100644 --- a/docs/features/additional-histories.mdx +++ b/docs/features/additional-histories.mdx @@ -33,6 +33,25 @@ with `chat_template_kwargs={"preserve_thinking": False}`. Additional histories remain useful for custom or externally managed templates that do not expose a prior-thinking preservation option. +Captured literal text can also contain strings such as `` without being +reasoning. During offline tokenization, ART can protect that text from a +template's inline-reasoning parser when the generation's own recorded request +explicitly disabled thinking and its native output decodes to the complete +content. ART tests a temporary render copy against the template's generation +prefix and empty-message scaffold; it does not change the captured messages or +the inference template. Explicit nonempty `reasoning` / `reasoning_content` +fields retain their meaning, and thinking-on or unknown sources are not +reinterpreted as literal-only generations. + +An accepted render adaptation must preserve the complete native conditioning, +sampled token IDs, logprobs, source ownership and sampled STOP flags of **every** +generation in the history, including later turns. Inconsistent native contexts +raise an error rather than silently training a repaired rendering against a +different sampled context. Missing native evidence or unsupported partial +protocol projections leave the existing rendering path unchanged. This check +does not supply missing logprobs or establish a general training contract for +source-free SFT text. + By splitting each turn into a separate history, you can preserve these tokens for training: ```python diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index bfdebad26..9bdf54bef 100644 --- a/src/art/trajectories/_tokenize.py +++ b/src/art/trajectories/_tokenize.py @@ -113,7 +113,12 @@ def __init__(self, render: _ChatRender) -> None: self.bytes = 0 def for_messages( - self, messages: list[dict[str, Any]], text: str, *, settings: object = None + self, + messages: list[dict[str, Any]], + text: str, + *, + settings: object = None, + full_generation_prompt: bool | None = None, ) -> _ChatRender: try: context = tuple(_render_context_key(message) for message in messages) @@ -129,6 +134,10 @@ def for_messages( if original != current: break common += 1 + if full_generation_prompt is not None and common == len(context) == len( + self.context + ): + self._remember((len(context), full_generation_prompt), text) def render( selected_messages: list[dict[str, Any]], *, add_generation_prompt: bool @@ -148,17 +157,20 @@ def render( value = self.render( selected_messages, add_generation_prompt=add_generation_prompt ) - if len(self.prefixes) < self._MAX_ENTRIES: - prefix = _common_prefix_length(self.text, value) - tail = value[prefix:] - size = 256 + 4 * len(tail) - if self.bytes + size <= self._MAX_BYTES: - self.prefixes[key] = prefix, tail - self.bytes += size + self._remember(key, value) return value return render + def _remember(self, key: tuple[int, bool], value: str) -> None: + if key not in self.prefixes and len(self.prefixes) < self._MAX_ENTRIES: + prefix = _common_prefix_length(self.text, value) + tail = value[prefix:] + size = 256 + 4 * len(tail) + if self.bytes + size <= self._MAX_BYTES: + self.prefixes[key] = prefix, tail + self.bytes += size + class _TokenChatRender(Protocol): def __call__( @@ -4532,32 +4544,50 @@ def _source_covers_complete_sampled_message( ) == normalize_chat_message(projected[0]) +@dataclass +class _NativeRenderSource: + key: _SampledSourceKey + prompt: list[int] + output: list[int] + logprobs: list[float] + stops: int + + def _preserve_literal_thinking_off_content( history: ChatCompletionsHistory, messages: list[dict[str, Any]], - template: object, kwargs: Mapping[str, object], -) -> None: - # This Qwen3.5 template treats any in unstructured content as a - # reasoning separator, even with thinking disabled. Restrict the render-copy - # adaptation to its exact preserved template; other templates may interpret - # an empty reasoning_content field differently. + tokenizer: Tokenizer, + render: _ChatRender, +) -> list[_NativeRenderSource]: + """Prove a literal render view without changing captured message semantics. + + Some templates parse legacy inline reasoning even for a generation recorded + with thinking disabled. An empty structured field is a usable escape only + if it preserves the empty-message scaffold and inserts the *entire* native + content at the generation boundary. Template identity is not this contract. + """ if ( - not isinstance(template, str) - or sha256(template.encode()).hexdigest() - != "098047d425a6673b1fe1a82a197a481616e53a283beaa8cb76cbb74d38ca6644" - or kwargs.get("enable_thinking") is not False + kwargs.get("enable_thinking") is not False or kwargs.get("preserve_thinking") is not True ): - return - for message, source in zip(messages, history.message_sources, strict=True): + return [] + decode = getattr(tokenizer, "decode", None) + if not callable(decode): + return [] + adapted = False + working = list(messages) + native: dict[int, _NativeRenderSource] = {} + for index, (message, source) in enumerate( + zip(messages, history.message_sources, strict=True) + ): if ( source is None or not isinstance(source.exchange, ChatCompletionsExchange) or source.choice_index is None or message.get("role") != "assistant" or not isinstance(content := message.get("content"), str) - or "" not in content + or not content ): continue request_kwargs = source.exchange.request.get("chat_template_kwargs") @@ -4567,6 +4597,11 @@ def _preserve_literal_thinking_off_content( ): continue choice = _chat_choice(source) + if ( + _field(choice, "prompt_token_ids") is None + and _field(source.exchange.response, "prompt_token_ids") is None + ): + continue # Visible-only histories may omit structured reasoning present in the # source response. Preserve both that source and normalized aliases. if any( @@ -4579,9 +4614,135 @@ def _preserve_literal_thinking_off_content( ) ): continue - prompt, output, _ = _chat_choice_tokens(choice, source.exchange.response) - if prompt is not None and output is not None: - message["reasoning_content"] = "" + prefix = working[:index] + try: + generation = render(prefix, add_generation_prompt=True) + completed = render([*prefix, message], add_generation_prompt=False) + except Exception: + continue + if completed.startswith(generation + content): + continue + prompt, output, logprobs = _chat_choice_tokens(choice, source.exchange.response) + if prompt is None or not output: + continue + key = _sampled_source_key(source) + stop_count = _sampled_stop_suffix( + output, + source=source, + source_key=key, + tokenizer=tokenizer, + ) + native[id(source)] = _NativeRenderSource( + key, prompt, output, logprobs, stop_count + ) + body = output[:-stop_count] if stop_count else output + # In particular, do not trim whitespace or strip arbitrary special + # tokens to manufacture agreement with the visible message. + try: + if ( + decode( + body, skip_special_tokens=False, clean_up_tokenization_spaces=False + ) + != content + ): + continue + literal = {**message, "reasoning": "", "reasoning_content": ""} + empty = render( + [*prefix, {**message, "content": ""}], add_generation_prompt=False + ) + literal_empty = render( + [*prefix, {**literal, "content": ""}], add_generation_prompt=False + ) + if empty != literal_empty or not empty.startswith(generation): + continue + literal_completed = render([*prefix, literal], add_generation_prompt=False) + if literal_completed != generation + content + empty[len(generation) :]: + continue + except Exception: + # These are optional capability probes, not the actual render. A + # token-only tokenizer or a template that rejects an empty message + # must retain the original rendering/error path. Control exceptions + # and source-validation errors are deliberately not swallowed. + continue + working[index] = literal + adapted = True + if not adapted: + return [] + # Partial protocol projections or absent native evidence cannot certify a + # changed render view. Decline the adaptation before committing it; leave + # their ordinary generic/SFT behavior to the existing tokenizer. + complete: dict[_SampledSourceKey, _NativeRenderSource] = {} + for source in history.message_sources: + if source is None or not _source_is_sampled(source): + continue + evidence = native.get(id(source)) + if evidence is None: + prompt = _chat_source_prompt_tokens(source) + output, logprobs = _chat_source_full_tokens(source) + if prompt is None or output is None: + return [] + key = _sampled_source_key(source) + evidence = _NativeRenderSource( + key, + prompt, + output, + logprobs, + _sampled_stop_suffix( + output, source=source, source_key=key, tokenizer=tokenizer + ), + ) + complete[evidence.key] = evidence + messages[:] = working + return list(complete.values()) + + +def _require_native_render_conditioning( + sources: list[_NativeRenderSource], + tokenized: TokenizedHistory, + trace: _TraceBuilder, +) -> TokenizedHistory: + """A render adaptation must preserve *every* sampled generation's context. + + Text equality in a probe is not token or conditioning equality. In + particular, changing an earlier message can invalidate a later source that + was generated from its lossy rendering. Generic histories without a render + adaptation retain their existing behavior. + """ + assert trace.trace is not None + observed = trace.trace + expected_keys: list[_SampledSourceKey | None] = [None] * len(tokenized.tokens) + required = ( + TokenFlag.EXACT | TokenFlag.SAMPLED | TokenFlag.ASSISTANT | TokenFlag.OUTPUT + ) + for source in sources: + prompt, output, logprobs = source.prompt, source.output, source.logprobs + start, end = len(prompt), len(prompt) + len(output) + if tokenized.tokens[:start] != prompt or tokenized.tokens[start:end] != output: + raise ValueError( + "Literal render adaptation changes native sampled conditioning" + ) + key, count = source.key, source.stops + if len(logprobs) != len(output): + raise ValueError( + "Literal render adaptation lacks complete sampled logprobs" + ) + for index, logprob in enumerate(logprobs, start): + actual = tokenized.logprobs[index] + flag = tokenized.flags[index] + if ( + expected_keys[index] is not None + or observed.source_keys[index] != key + or flag & required != required + or bool(flag & TokenFlag.STOP) != (index >= end - count) + or not (actual == logprob or math.isnan(actual) and math.isnan(logprob)) + ): + raise ValueError( + "Literal render adaptation changes native sampled ownership, flags, or logprobs" + ) + expected_keys[index] = key + if observed.source_keys != expected_keys: + raise ValueError("Literal render adaptation changes sampled source coverage") + return tokenized def _tokenize_chat_view( @@ -4630,7 +4791,6 @@ def _tokenize_chat_view( **default_chat_template_kwargs_for_template(template), **explicit_kwargs, } - _preserve_literal_thinking_off_content(history, messages, template, kwargs) ends_with_assistant = bool(messages) and messages[-1].get("role") == "assistant" segmented = False @@ -4677,6 +4837,69 @@ def render_text( ) prefix_render_cache = _PrefixChatRenderCache(render_normalized_text) + literal_render = render_text + literal_cache_primed = False + if ( + kwargs.get("enable_thinking") is False + and kwargs.get("preserve_thinking") is True + and any( + isinstance(getattr(source, "exchange", None), ChatCompletionsExchange) + and isinstance(source.exchange.request.get("chat_template_kwargs"), Mapping) + and source.exchange.request["chat_template_kwargs"].get("enable_thinking") + is False + for source in history.message_sources + if source is not None + ) + and cacheable_chat_template( + resolved_tokenizer, template, history.tools, kwargs, messages + ) + ): + try: + # The same admitted cache serves behavioral and boundary probes. + # Changed aliases/normalization invalidate its affected prefixes; + # arbitrary or token-only renderers retain the uncached path. + proof_originals = list(messages) + proof_messages = normalize_tool_call_arguments_for_chat_template( + messages, template + ) + cached_proof = prefix_render_cache.for_messages( + proof_messages, + render_normalized_text( + proof_messages, add_generation_prompt=not ends_with_assistant + ), + full_generation_prompt=not ends_with_assistant, + settings=_render_context_key( + [ + history.tools, + kwargs, + getattr(resolved_tokenizer, "special_tokens_map"), + ] + ), + ) + + literal_cache_primed = True + + def literal_render( + selected_messages: list[dict[str, Any]], *, add_generation_prompt: bool + ) -> str: + if len(selected_messages) <= len(proof_originals) and all( + current is original + for current, original in zip(selected_messages, proof_originals) + ): + return cached_proof( + proof_messages[: len(selected_messages)], + add_generation_prompt=add_generation_prompt, + ) + return render_text( + selected_messages, add_generation_prompt=add_generation_prompt + ) + except Exception: + pass # Optional text proof must not block the raw-render fallback. + literal_sources = _preserve_literal_thinking_off_content( + history, messages, kwargs, resolved_tokenizer, literal_render + ) + if literal_sources and _trace is None: + _trace = _TraceBuilder() def segmented_render( selected_messages: list[dict[str, Any]], *, add_generation_prompt: bool @@ -4690,19 +4913,25 @@ def segmented_render( selected_messages = normalize_tool_call_arguments_for_chat_template( selected_messages, template ) - text = render_normalized_text( + settings = _render_context_key( + [ + history.tools, + kwargs, + getattr(resolved_tokenizer, "special_tokens_map"), + ] + ) + full_render = ( + prefix_render_cache.for_messages( + selected_messages, prefix_render_cache.text, settings=settings + ) + if literal_cache_primed + else render_normalized_text + ) + text = full_render( selected_messages, add_generation_prompt=add_generation_prompt ) span_render = prefix_render_cache.for_messages( - selected_messages, - text, - settings=_render_context_key( - [ - history.tools, - kwargs, - getattr(resolved_tokenizer, "special_tokens_map"), - ] - ), + selected_messages, text, settings=settings ) else: text = render_text( @@ -5590,6 +5819,11 @@ def differing_span(probe: Sequence[int]) -> tuple[int, int] | None: ) ) ): + if literal_sources: + assert _trace is not None + return _require_native_render_conditioning( + literal_sources, exact, _trace + ) return exact sampled_message_count = sum( @@ -6296,6 +6530,9 @@ def differing_span(probe: Sequence[int]) -> tuple[int, int] | None: ) if _trace is not None: _trace.set(tokenized, source_keys, sources) + if literal_sources: + assert _trace is not None + return _require_native_render_conditioning(literal_sources, tokenized, _trace) return tokenized diff --git a/tests/unit/trajectories/test_literal_thinking_off.py b/tests/unit/trajectories/test_literal_thinking_off.py index 16188c733..47cac1338 100644 --- a/tests/unit/trajectories/test_literal_thinking_off.py +++ b/tests/unit/trajectories/test_literal_thinking_off.py @@ -3,11 +3,13 @@ from copy import deepcopy from datetime import UTC, datetime from pathlib import Path +import random import re from typing import Any, cast from jinja2.sandbox import ImmutableSandboxedEnvironment from openai.types.chat import ChatCompletion, ChatCompletionMessageParam +from openai.types.responses import Response import pytest import art.trajectories as tr @@ -64,8 +66,10 @@ def _history( content: str = _LITERAL, reasoning: str | None = None, reasoning_field: str = "reasoning_content", + template: str = _TEMPLATE, ) -> tuple[tr.ChatCompletionsHistory, _TemplateTokenizer]: tokenizer = _TemplateTokenizer() + tokenizer.chat_template = template request_kwargs: dict[str, Any] = {"preserve_thinking": True} if thinking is not None: request_kwargs["enable_thinking"] = thinking @@ -109,7 +113,7 @@ def _history( request=tr.ChatCompletionsRequest( model="public/qwen35", messages=cast(list[ChatCompletionMessageParam], prompt_messages), - chat_template=_TEMPLATE, + chat_template=template, chat_template_kwargs=request_kwargs, ), response=response, @@ -126,7 +130,7 @@ def _history( def _outcome( history: tr.ChatCompletionsHistory, tokenizer: _TemplateTokenizer -) -> object: +) -> tuple[Any, ...]: try: value = history.tokenize(tokenizer=tokenizer) except ValueError as error: @@ -153,7 +157,11 @@ def test_native_thinking_off_retains_literal_content( tokenizer.calls.clear() tokenizer.rendered.clear() tokenized = history.tokenize(tokenizer=tokenizer) - assert content in tokenizer.rendered[0] + # The behavioral proof probes the original view before selecting the + # literal render copy. Its first completed candidate and later real render + # must both retain the content. + assert content in tokenizer.rendered[4] + assert content in tokenizer.rendered[5] sampled = [ i for i, flag in enumerate(tokenized.flags) if flag & tr.TokenFlag.SAMPLED ] @@ -162,8 +170,8 @@ def test_native_thinking_off_retains_literal_content( required = tr.TokenFlag.EXACT | tr.TokenFlag.ASSISTANT | tr.TokenFlag.OUTPUT assert all(tokenized.flags[i] & required == required for i in sampled) assert not any(flag & tr.TokenFlag.STOP for flag in tokenized.flags) - assert tokenizer.calls[0][-1]["reasoning_content"] == "" - assert tokenizer.calls[0][-1]["content"] == content + assert tokenizer.calls[5][-1]["reasoning_content"] == "" + assert tokenizer.calls[5][-1]["content"] == content assert history.model_dump(mode="python") == original @@ -174,7 +182,6 @@ def test_native_thinking_off_retains_literal_content( "source_unknown", "effective_on", "preserve_off", - "other_template", "no_source", "request_source", "no_native_prompt", @@ -201,8 +208,6 @@ def test_unrelated_histories_keep_original_rendering( history.chat_template_kwargs["enable_thinking"] = case == "effective_on" if case == "preserve_off": history.chat_template_kwargs["preserve_thinking"] = False - if case == "other_template": - history.chat_template = _TEMPLATE + "{# different template #}" if case == "no_source": history.message_sources[-1] = None source = history.message_sources[-1] @@ -249,7 +254,7 @@ def test_explicit_empty_reasoning_is_preserved(field: str) -> None: ) == _LITERAL ) - assert tokenizer.calls[0][-1]["reasoning_content"] == "" + assert any(call[-1].get("reasoning_content") == "" for call in tokenizer.calls) assert history.model_dump(mode="python") == original @@ -276,7 +281,13 @@ def test_mixed_history_uses_each_generations_own_request( ) original = history.model_dump(mode="python") _outcome(history, tokenizer) - rendered_messages = tokenizer.calls[0] + rendered_messages = next( + call + for call in tokenizer.calls + if len(call) == 4 + and call[-1].get("reasoning_content") == "" + and call[-1]["content"] == _LITERAL + ) assert "reasoning_content" not in rendered_messages[1] assert rendered_messages[3]["reasoning_content"] == "" assert [message["content"] for message in rendered_messages] == [ @@ -441,3 +452,603 @@ def observe(*args: Any, **kwargs: Any) -> tr.TokenizedHistory | None: for flag in value.flags[start:stop] ) assert history.model_dump(mode="python") == original + + +@pytest.mark.parametrize("alias", ["reasoning_content", "reasoning"]) +@pytest.mark.parametrize( + "tags", [("", ""), ("", "")] +) +def test_literal_contract_is_independent_of_template_hash_and_delimiters( + alias: str, tags: tuple[str, str] +) -> None: + template = ( + (_TEMPLATE + "{# semantically equivalent template revision #}") + .replace("message.reasoning_content", f"message.{alias}") + .replace("", tags[0]) + .replace("", tags[1]) + ) + # Public literal strings include nested, repeated, Unicode and whitespace + # cases. Neither the content nor the delimiter spelling is an allowlist. + rng = random.Random(0) + for middle in [ + "", + "\n", + " π ", + tags[0], + tags[1], + *["".join(rng.choices("ab\n Ω", k=8)) for _ in range(5)], + ]: + content = "head" + tags[1] + middle + tags[1] + "tail" + history, tokenizer = _history(content=content, template=template) + original = history.model_dump(mode="python") + value = history.tokenize(tokenizer=tokenizer) + source = history.message_sources[-1] + assert source is not None and isinstance( + source.exchange, tr.ChatCompletionsExchange + ) + prompt, output, logprobs = _tokenize._chat_choice_tokens( + source.exchange.response.choices[0], source.exchange.response + ) + assert prompt is not None and output is not None + assert value.tokens == prompt + output + assert value.logprobs[len(prompt) :] == logprobs + assert value.flags[len(prompt) :] == [ + tr.TokenFlag.EXACT + | tr.TokenFlag.SAMPLED + | tr.TokenFlag.ASSISTANT + | tr.TokenFlag.OUTPUT + ] * len(output) + assert history.model_dump(mode="python") == original + + +@pytest.mark.parametrize("probe", ["decode_kwargs", "text", "empty", "scaffold"]) +def test_unsupported_literal_proof_keeps_original_path( + probe: str, monkeypatch: pytest.MonkeyPatch +) -> None: + history, tokenizer = _history( + content=_LITERAL if probe != "text" else "ordinary public content" + ) + decode, render = tokenizer.decode, tokenizer.apply_chat_template + branches: list[str] = [] + + def guarded_decode(ids: list[int], **kwargs: Any) -> str: + if probe == "decode_kwargs" and kwargs: + branches.append(probe) + raise TypeError("unsupported decode keyword") + return decode(ids, **kwargs) + + def guarded_render( + messages: list[dict[str, Any]], **kwargs: Any + ) -> str | list[int]: + if probe == "text" and kwargs.get("tokenize") is False: + branches.append(probe) + raise TypeError("token-only template") + if ( + probe == "empty" + and messages + and messages[-1].get("role") == "assistant" + and not messages[-1].get("content") + ): + branches.append(probe) + raise ValueError("empty assistant unsupported") + value = render(messages, **kwargs) + if probe == "scaffold" and any("reasoning_content" in m for m in messages): + branches.append(probe) + return "changed" + value if isinstance(value, str) else [999] + value + return value + + monkeypatch.setattr(tokenizer, "decode", guarded_decode) + monkeypatch.setattr(tokenizer, "apply_chat_template", guarded_render) + candidate = _outcome(history, tokenizer) + assert probe in branches + with monkeypatch.context() as patch: + patch.setattr( + _tokenize, "_preserve_literal_thinking_off_content", lambda *args: False + ) + baseline = _outcome(history, tokenizer) + assert candidate == baseline + if probe in {"decode_kwargs", "text"}: + assert not isinstance(candidate[0], type) + + +@pytest.mark.parametrize("error", [KeyboardInterrupt, SystemExit]) +def test_literal_probe_does_not_swallow_control_exceptions( + error: type[BaseException], monkeypatch: pytest.MonkeyPatch +) -> None: + history, tokenizer = _history() + + def interrupt(*args: Any, **kwargs: Any) -> str: + raise error() + + monkeypatch.setattr(tokenizer, "decode", interrupt) + with pytest.raises(error): + history.tokenize(tokenizer=tokenizer) + + +@pytest.mark.parametrize( + "case", + [ + "missing_output", + "different_output", + "nonempty_reasoning", + "none_reasoning", + "empty_reasoning", + ], +) +def test_native_evidence_and_reasoning_are_independent_authorities( + case: str, monkeypatch: pytest.MonkeyPatch +) -> None: + history, tokenizer = _history( + reasoning="structured" + if case == "nonempty_reasoning" + else "" + if case == "empty_reasoning" + else None + ) + source = history.message_sources[-1] + assert source is not None and isinstance( + source.exchange, tr.ChatCompletionsExchange + ) + choice = source.exchange.response.choices[0] + assert choice.model_extra is not None + if case == "missing_output": + choice.model_extra.pop("token_ids") + choice.logprobs = None + elif case == "different_output": + choice.model_extra["token_ids"] = [12345] + choice.logprobs = None + elif case == "none_reasoning": + assert choice.message.model_extra is not None + choice.message.model_extra["reasoning_content"] = None + original = history.model_dump(mode="python") + value = _outcome(history, tokenizer) + if case in {"none_reasoning", "empty_reasoning"}: + assert not isinstance(value[0], type) + assert any(call[-1].get("reasoning_content") == "" for call in tokenizer.calls) + else: + assert not any( + call[-1].get("reasoning_content") == "" for call in tokenizer.calls + ) + with monkeypatch.context() as patch: + patch.setattr( + _tokenize, "_preserve_literal_thinking_off_content", lambda *args: False + ) + assert _outcome(history, tokenizer) == value + assert history.model_dump(mode="python") == original + + +def _two_turn_literal_history( + *, served_literal: bool +) -> tuple[tr.ChatCompletionsHistory, _TemplateTokenizer]: + first, tokenizer = _history() + first_source = first.message_sources[-1] + assert first_source is not None and isinstance( + first_source.exchange, tr.ChatCompletionsExchange + ) + second, _ = _history(content="second answer") + source = second.message_sources[-1] + assert source is not None and isinstance( + source.exchange, tr.ChatCompletionsExchange + ) + exchange = source.exchange + messages = [ + *[dict(m) for m in deepcopy(first.messages)], + {"role": "user", "content": "follow up"}, + ] + exchange.request["messages"] = cast( + list[ChatCompletionMessageParam], deepcopy(messages) + ) + if served_literal: + messages[1]["reasoning_content"] = "" + prompt = tokenizer.apply_chat_template( + messages, + add_generation_prompt=True, + enable_thinking=False, + preserve_thinking=True, + ) + assert isinstance(prompt, list) + exchange.response.id = "second-source" + assert exchange.response.choices[0].model_extra is not None + exchange.response.choices[0].model_extra["prompt_token_ids"] = prompt + # Explicit History construction also exercises inconsistent native chains + # which automatic history grouping would split into separate histories. + history = tr.ChatCompletionsHistory( + model=first.model, + messages=[ + *first.messages, + {"role": "user", "content": "follow up"}, + second.messages[-1], + ], + message_sources=[*first.message_sources, None, source], + chat_template=first.chat_template, + chat_template_kwargs=first.chat_template_kwargs, + ) + tokenizer.calls.clear() + tokenizer.rendered.clear() + return history, tokenizer + + +@pytest.mark.parametrize( + "case", + [ + "exact", + "lossy_later_prompt", + "missing_later_prompt", + "later_logprobs", + "later_owner", + "later_stop", + ], +) +def test_adaptation_proves_every_turn_not_only_the_repaired_message( + case: str, monkeypatch: pytest.MonkeyPatch +) -> None: + history, tokenizer = _two_turn_literal_history( + served_literal=case != "lossy_later_prompt" + ) + source = history.message_sources[-1] + assert source is not None and isinstance( + source.exchange, tr.ChatCompletionsExchange + ) + if case == "missing_later_prompt": + assert source.exchange.response.choices[0].model_extra is not None + source.exchange.response.choices[0].model_extra.pop("prompt_token_ids") + original = history.model_dump(mode="python") + validate = _tokenize._require_native_render_conditioning + extra = source.exchange.response.choices[0].model_extra + assert extra is not None + last_sampled = len(extra.get("prompt_token_ids", [])) + len(extra["token_ids"]) - 1 + + def corrupt(h: Any, value: Any, trace: Any) -> Any: + if case == "later_logprobs": + value.logprobs[last_sampled] = -0.75 + elif case == "later_owner": + trace.trace.source_keys[last_sampled] = trace.trace.source_keys[0] + elif case == "later_stop": + value.flags[last_sampled] |= tr.TokenFlag.STOP + return validate(h, value, trace) + + monkeypatch.setattr(_tokenize, "_require_native_render_conditioning", corrupt) + if case == "missing_later_prompt": + candidate = _outcome(history, tokenizer) + with monkeypatch.context() as patch: + patch.setattr( + _tokenize, "_preserve_literal_thinking_off_content", lambda *args: [] + ) + assert _outcome(history, tokenizer) == candidate + elif case == "exact": + value = history.tokenize(tokenizer=tokenizer) + assert sum(bool(flag & tr.TokenFlag.SAMPLED) for flag in value.flags) == len( + _LITERAL + ) + len("second answer") + else: + with pytest.raises(ValueError, match="Literal render adaptation"): + history.tokenize(tokenizer=tokenizer) + assert history.model_dump(mode="python") == original + + +def test_behavioral_proof_has_a_fixed_render_budget() -> None: + history, tokenizer = _history() + messages = [dict(message) for message in history.messages] + + def render( + selected_messages: list[dict[str, Any]], *, add_generation_prompt: bool + ) -> str: + value = tokenizer.apply_chat_template( + selected_messages, + tokenize=False, + add_generation_prompt=add_generation_prompt, + enable_thinking=False, + preserve_thinking=True, + ) + assert isinstance(value, str) + return value + + assert _tokenize._preserve_literal_thinking_off_content( + history, + messages, + {"enable_thinking": False, "preserve_thinking": True}, + tokenizer, + render, + ) + assert len(tokenizer.calls) == 5 + assert history.messages[-1].get("reasoning_content") is None + + +@pytest.mark.parametrize("partial", [False, True]) +def test_mixed_responses_source_requires_complete_native_inventory( + partial: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + first, tokenizer = _history() + request_messages = [ + *[dict(m) for m in deepcopy(first.messages)], + {"role": "user", "content": "next"}, + ] + served = deepcopy(request_messages) + served[1]["reasoning_content"] = "" + prompt = tokenizer.apply_chat_template( + served, + add_generation_prompt=True, + enable_thinking=False, + preserve_thinking=True, + ) + assert isinstance(prompt, list) + texts = ["one", "two"] if partial else ["one"] + output = list(map(ord, "".join(texts))) + response = Response.model_validate( + { + "id": "mixed-responses", + "created_at": 0, + "model": "public/qwen35", + "object": "response", + "output": [ + { + "id": f"message-{i}", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [ + { + "type": "output_text", + "text": text, + "annotations": [], + "logprobs": [], + } + ], + } + for i, text in enumerate(texts) + ], + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [], + "token_generations": [ + { + "prompt_token_ids": prompt, + "output_tokens": [ + {"token_id": token, "logprob": -0.25} for token in output + ], + "output_indices": list(range(len(texts))), + } + ], + } + ) + exchange = tr.ResponsesExchange( + request=tr.ResponsesRequest(model="public/qwen35", input="next"), + response=response, + start_time=datetime(2026, 1, 1, tzinfo=UTC), + end_time=datetime(2026, 1, 1, tzinfo=UTC), + ) + projected = ( + tr.Trajectory(exchanges=tr.TrajectoryExchanges(responses=[exchange])) + .responses_history() + .as_chat_completions_history() + ) + history = tr.ChatCompletionsHistory( + model=first.model, + messages=[*first.messages, *projected.messages], + message_sources=[*first.message_sources, *projected.message_sources], + chat_template=first.chat_template, + chat_template_kwargs=first.chat_template_kwargs, + ) + original = history.model_dump(mode="python") + candidate = _outcome(history, tokenizer) + if partial: + with monkeypatch.context() as patch: + patch.setattr( + _tokenize, "_preserve_literal_thinking_off_content", lambda *args: [] + ) + assert _outcome(history, tokenizer) == candidate + else: + assert not isinstance(candidate[0], type) + assert candidate[0][: len(prompt) + len(output)] == prompt + output + assert history.model_dump(mode="python") == original + + +@pytest.mark.parametrize("finish", ["length", "stop"]) +def test_literal_markers_do_not_own_sampled_stop_bits(finish: str) -> None: + history, _ = _history() + tokenizer = _NewlineRunTokenizer() + source = history.message_sources[-1] + assert source is not None and isinstance( + source.exchange, tr.ChatCompletionsExchange + ) + exchange = source.exchange + data = exchange.response.model_dump(mode="python") + choice = data["choices"][0] + prompt = tokenizer.apply_chat_template( + [dict(m) for m in exchange.request["messages"]], + add_generation_prompt=True, + enable_thinking=False, + preserve_thinking=True, + ) + assert isinstance(prompt, list) + output = tokenizer(_LITERAL)["input_ids"] + ( + [tokenizer.eos_token_id] if finish == "stop" else [] + ) + choice.update(prompt_token_ids=prompt, token_ids=output, finish_reason=finish) + choice["logprobs"]["content"] = [ + {"token": f"token_id:{token}", "logprob": -0.5, "bytes": [], "top_logprobs": []} + for token in output + ] + exchange.response = ChatCompletion.model_validate(data) + original = history.model_dump(mode="python") + value = history.tokenize(tokenizer=tokenizer) + assert value.tokens[: len(prompt) + len(output)] == prompt + output + sampled = value.flags[len(prompt) : len(prompt) + len(output)] + assert [bool(flag & tr.TokenFlag.STOP) for flag in sampled] == [False] * ( + len(output) - (finish == "stop") + ) + ([True] if finish == "stop" else []) + assert value.logprobs[len(prompt) : len(prompt) + len(output)] == [-0.5] * len( + output + ) + assert history.model_dump(mode="python") == original + + +@pytest.mark.parametrize("limit", [0, 64, 8 << 20]) +def test_literal_adaptation_reuses_only_unchanged_cache_prefixes( + limit: int, monkeypatch: pytest.MonkeyPatch +) -> None: + history, tokenizer = _two_turn_literal_history(served_literal=True) + expected = _outcome(history, tokenizer) + # This fixture renderer is pure, but not stock HF. Exercise the admitted + # integration explicitly; real HF eligibility is tested separately. + monkeypatch.setattr(_tokenize, "cacheable_chat_template", lambda *args: True) + monkeypatch.setattr(tokenizer, "special_tokens_map", {}, raising=False) + monkeypatch.setattr(_tokenize._PrefixChatRenderCache, "_MAX_BYTES", limit) + original = _tokenize._PrefixChatRenderCache.for_messages + observed = [] + + def check(self: Any, *args: Any, **kwargs: Any) -> Any: + result = original(self, *args, **kwargs) + observed.append(self) + return result + + monkeypatch.setattr(_tokenize._PrefixChatRenderCache, "for_messages", check) + assert _outcome(history, tokenizer) == expected + assert observed and all(cache.bytes <= limit for cache in observed) + + +@pytest.mark.parametrize("reasoning_field", ["reasoning", "reasoning_content"]) +def test_earlier_structured_thinking_and_later_literal_content_remain_distinct( + reasoning_field: str, + monkeypatch: pytest.MonkeyPatch, +) -> None: + first, tokenizer = _history( + thinking=True, + content="reasoned answer", + reasoning="explicit thought", + reasoning_field=reasoning_field, + ) + source = first.message_sources[-1] + assert source is not None and isinstance( + source.exchange, tr.ChatCompletionsExchange + ) + first_exchange = source.exchange + output = list(map(ord, "explicit thought\n\nreasoned answer")) + data = first_exchange.response.model_dump(mode="python") + choice = data["choices"][0] + choice["token_ids"] = output + choice["logprobs"]["content"] = [ + {"token": f"token_id:{t}", "logprob": -0.5, "bytes": [], "top_logprobs": []} + for t in output + ] + first_exchange.response = ChatCompletion.model_validate(data) + second, _ = _history() + source = second.message_sources[-1] + assert source is not None and isinstance( + source.exchange, tr.ChatCompletionsExchange + ) + second_exchange = source.exchange + request = [dict(m) for m in first.messages] + [{"role": "user", "content": "next"}] + second_exchange.request["messages"] = cast( + list[ChatCompletionMessageParam], deepcopy(request) + ) + served = deepcopy(request) + served[1]["reasoning_content"] = "explicit thought" + second_exchange.response.id = "later-literal" + extra = second_exchange.response.choices[0].model_extra + assert extra is not None + extra["prompt_token_ids"] = tokenizer.apply_chat_template( + served, + add_generation_prompt=True, + enable_thinking=False, + preserve_thinking=True, + ) + # Different request thinking modes are normally separate automatic + # histories. Explicit mixed histories are also supported and source checked. + history = tr.ChatCompletionsHistory( + model=first.model, + messages=[ + *first.messages, + {"role": "user", "content": "next"}, + second.messages[-1], + ], + message_sources=[*first.message_sources, None, second.message_sources[-1]], + chat_template=second.chat_template, + chat_template_kwargs=second.chat_template_kwargs, + ) + first_prompt, first_output, _ = _tokenize._chat_choice_tokens( + first_exchange.response.choices[0], first_exchange.response + ) + assert first_prompt is not None and first_output is not None + assert ( + extra["prompt_token_ids"][: len(first_prompt) + len(first_output)] + == first_prompt + first_output + ) + original = history.model_dump(mode="python") + checked = [] + validate = _tokenize._require_native_render_conditioning + + def observe(*args: Any) -> Any: + result = validate(*args) + checked.append(len(args[0])) + return result + + monkeypatch.setattr(_tokenize, "_require_native_render_conditioning", observe) + value = history.tokenize(tokenizer=tokenizer) + assert checked == [2] + for e in (first_exchange, second_exchange): + prompt, native, logprobs = _tokenize._chat_choice_tokens( + e.response.choices[0], e.response + ) + assert prompt is not None and native is not None + assert value.tokens[: len(prompt) + len(native)] == prompt + native + assert value.logprobs[len(prompt) : len(prompt) + len(native)] == logprobs + assert history.model_dump(mode="python") == original + + +@pytest.mark.parametrize("literal", [False, True]) +def test_stock_renderer_shares_proofs_without_stale_alias_prefixes( + literal: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + from test_prefix_render_cache import _tokenizer + + tokenizers = pytest.importorskip("tokenizers") + tokenizer = _tokenizer(_TEMPLATE) + tokenizer.backend_tokenizer.decoder = tokenizers.decoders.Fuse() + history, _ = _two_turn_literal_history(served_literal=True) + if not literal: + # Construct a clean equivalent, rather than changing captured messages. + history, _ = _history(content="ordinary answer") + for source in history.message_sources: + if source is None or source.choice_index is None: + continue + assert isinstance(source.exchange, tr.ChatCompletionsExchange) + choice = source.exchange.response.choices[0] + extra = choice.model_extra + assert extra is not None + for field in ("prompt_token_ids", "token_ids"): + extra[field] = tokenizer( + "".join(map(chr, extra[field])), add_special_tokens=False + )["input_ids"] + assert choice.logprobs is not None and choice.logprobs.content + # Keep actual Pydantic logprob objects, with the new codec's token IDs. + pair = choice.logprobs.content[0] + choice.logprobs.content = [ + pair.model_copy(update={"token": f"token_id:{token}"}) + for token in extra["token_ids"] + ] + kwargs = {"enable_thinking": False, "preserve_thinking": True} + assert _tokenize.cacheable_chat_template( + tokenizer, _TEMPLATE, history.tools, kwargs, list(history.messages) + ) + original = history.model_dump(mode="python") + calls = [] + cache = _tokenize._PrefixChatRenderCache.for_messages + + def observe(self: Any, *args: Any, **kwargs: Any) -> Any: + calls.append(self) + return cache(self, *args, **kwargs) + + with monkeypatch.context() as patch: + patch.setattr(_tokenize._PrefixChatRenderCache, "for_messages", observe) + actual = history.tokenize(tokenizer=tokenizer) + assert len(calls) >= 2 and all(item is calls[0] for item in calls) + with monkeypatch.context() as patch: + patch.setattr(_tokenize, "cacheable_chat_template", lambda *args: False) + uncached = history.tokenize(tokenizer=tokenizer) + assert actual.tokens == uncached.tokens + assert actual.flags == uncached.flags + assert [None if x != x else x for x in actual.logprobs] == [ + None if x != x else x for x in uncached.logprobs + ] + assert history.model_dump(mode="python") == original diff --git a/tests/unit/trajectories/test_prefix_render_cache.py b/tests/unit/trajectories/test_prefix_render_cache.py index 05760040a..ea245409b 100644 --- a/tests/unit/trajectories/test_prefix_render_cache.py +++ b/tests/unit/trajectories/test_prefix_render_cache.py @@ -396,3 +396,37 @@ def counted(*args, **kwargs): assert cached_calls <= calls - turns * (turns - 1) else: assert cached_calls == calls # Plain text already avoids tool probes. + + +@pytest.mark.parametrize("limit", [0, 256, 1 << 20]) +def test_known_full_render_priming_keeps_context_flag_and_limits(monkeypatch, limit): + calls = [] + + def render(selected_messages, *, add_generation_prompt): + calls.append(1) + return json.dumps(selected_messages) + str(add_generation_prompt) + + monkeypatch.setattr(tokenization._PrefixChatRenderCache, "_MAX_BYTES", limit) + messages = [{"role": "assistant", "content": "literal "}] + cache = tokenization._PrefixChatRenderCache(render) + known = render(messages, add_generation_prompt=False) + cached = cache.for_messages(messages, known, full_generation_prompt=False) + before = len(calls) + assert cached(messages, add_generation_prompt=False) == known + assert len(calls) == before + (limit < 256) + # The full completed render does not certify a generation-prompt render. + before = len(calls) + assert cached(messages, add_generation_prompt=True).endswith("True") + assert len(calls) == before + 1 + changed = [{**messages[0], "reasoning_content": ""}] + other = render(changed, add_generation_prompt=False) + probe = cache.for_messages(changed, other, full_generation_prompt=False) + assert probe(changed, add_generation_prompt=False) == other + # Priming a changed alias context must never overwrite the baseline entry. + assert cached(messages, add_generation_prompt=False) == known + assert cache.bytes <= limit + current = cache.for_messages( + changed, other, settings="new", full_generation_prompt=False + ) + assert current(changed, add_generation_prompt=False) == other + assert cache.bytes <= limit