From 11f31d4beb269025783c6a914b350f39f0844dfa Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Thu, 23 Jul 2026 15:05:49 +0000 Subject: [PATCH 01/58] Add protocol-native trajectory histories --- src/art/__init__.py | 34 +- src/art/langgraph/llm_wrapper.py | 4 +- src/art/preprocessing/tokenize.py | 160 +- src/art/tinker/server.py | 4 +- src/art/tinker_native/data.py | 23 +- src/art/trajectories/__init__.py | 512 +++++-- src/art/trajectories/_history.py | 1123 +++++++++++--- src/art/trajectories/_protocols.py | 55 +- src/art/trajectories/_tokenize.py | 1353 ++++++++++++++++- src/art/utils/trajectory_logging.py | 8 +- src/art/utils/trajectory_migration.py | 8 +- .../model_support/chat_template_rollout.py | 6 +- .../chat_template_conformance_cases.py | 8 +- tests/unit/test_tinker_native_exchanges.py | 14 +- tests/unit/test_trajectory_parquet.py | 8 +- tests/unit/trajectories/test_capture.py | 13 +- tests/unit/trajectories/test_history.py | 337 +++- tests/unit/trajectories/test_tokenize.py | 607 +++++++- 18 files changed, 3598 insertions(+), 679 deletions(-) diff --git a/src/art/__init__.py b/src/art/__init__.py index f317da1e8..1f8cf15f2 100644 --- a/src/art/__init__.py +++ b/src/art/__init__.py @@ -79,17 +79,28 @@ from .serverless import ServerlessBackend from .trajectories import ( AnthropicMessagesHistory, + AnthropicMessageSource, ChatCompletionsExchange, ChatCompletionsHistory, + ChatCompletionsMessageSource, CompletionsExchange, - CompletionsHistory, + CompletionsSource, + CompletionsStringHistory, + CompletionsStringSourceSpan, + CompletionsTokenHistory, + CompletionsTokenSourceSpan, History, + LegacyHistory, MessagesExchange, ResponsesExchange, ResponsesHistory, + ResponsesItemSource, TokenFlag, + TokenizedHistory, + TokenizedMultiHistoryTrajectory, TokenizedTrajectory, TokenizedTrajectoryGroup, + Tokenizer, Trajectory, TrajectoryExchanges, TrajectoryGroup, @@ -98,10 +109,6 @@ capture_auto_trajectory, # ty: ignore[deprecated] current_trajectory, current_trajectory_group, - tokenize_trajectories, - tokenize_trajectory, - tokenize_trajectory_group, - tokenize_trajectory_groups, trajectory, trajectory_group, ) @@ -158,24 +165,31 @@ "TrajectoryExchanges", "TrajectoryGroup", "History", + "LegacyHistory", "TrajectoryHistory", "ChatCompletionsExchange", "ChatCompletionsHistory", + "ChatCompletionsMessageSource", "CompletionsExchange", - "CompletionsHistory", + "CompletionsSource", + "CompletionsTokenHistory", + "CompletionsStringHistory", + "CompletionsTokenSourceSpan", + "CompletionsStringSourceSpan", "ResponsesExchange", "ResponsesHistory", + "ResponsesItemSource", "MessagesExchange", "AnthropicMessagesHistory", + "AnthropicMessageSource", "TokenFlag", + "Tokenizer", + "TokenizedHistory", + "TokenizedMultiHistoryTrajectory", "TokenizedTrajectory", "TokenizedTrajectoryGroup", "trajectory", "trajectory_group", - "tokenize_trajectory", - "tokenize_trajectories", - "tokenize_trajectory_group", - "tokenize_trajectory_groups", "capture_yielded_trajectory", "yield_trajectory", ] diff --git a/src/art/langgraph/llm_wrapper.py b/src/art/langgraph/llm_wrapper.py index 86e17dc2c..676eaf1f6 100644 --- a/src/art/langgraph/llm_wrapper.py +++ b/src/art/langgraph/llm_wrapper.py @@ -14,7 +14,7 @@ from langchain_core.utils.function_calling import convert_to_openai_tool from langchain_openai import ChatOpenAI -from art.trajectories import History, Trajectory +from art.trajectories import LegacyHistory, Trajectory from .logging import FileLogger from .message_utils import convert_langgraph_messages @@ -89,7 +89,7 @@ def create_messages_from_logs(logger: FileLogger, trajectory: Trajectory): trajectory.tools = tools[idx] else: trajectory.additional_histories.append( - History(messages_and_choices=converted, tools=tools[idx]) + LegacyHistory(messages_and_choices=converted, tools=tools[idx]) ) except Exception: pass diff --git a/src/art/preprocessing/tokenize.py b/src/art/preprocessing/tokenize.py index b466af12a..03208901a 100644 --- a/src/art/preprocessing/tokenize.py +++ b/src/art/preprocessing/tokenize.py @@ -19,7 +19,8 @@ from ..trajectories import ( ChatCompletionsExchange, - History, + ChatCompletionsHistory, + LegacyHistory, TokenFlag, Trajectory, TrajectoryGroup, @@ -448,7 +449,7 @@ def _tokenized_result_from_vllm_choices( def assemble_vllm_training_sequences( *, tokenizer: PreTrainedTokenizerBase, - histories: list[History], + histories: list[LegacyHistory], advantage: float, allow_training_without_logprobs: bool, trajectory: Trajectory, @@ -567,94 +568,97 @@ def tokenize_trajectory_groups( if advantage == 0 and drop_zero_advantage_trajectories: continue if trajectory.exchanges: - from ..trajectories._tokenize import ( - _as_tokenizer, - _exchange_list, - _exchange_tokens, - tokenize_one, - ) + from ..trajectories._tokenize import _as_tokenizer - exchange_result = tokenize_one( - trajectory, - tokenizer.name_or_path, - model=None, + exchange_results = trajectory.tokenize( + multi_history=True, + base_model=tokenizer.name_or_path, + tokenizer=_as_tokenizer(tokenizer), chat_template=None, chat_template_kwargs=chat_template_kwargs, - tokenizer_instance=_as_tokenizer(tokenizer), ) - sampled = [ - bool(flag & TokenFlag.SAMPLED) for flag in exchange_result.flags - ] - sampled_spans = _flag_spans(exchange_result.flags, TokenFlag.SAMPLED) - choice_spans = sampled_spans - if not allow_training_without_logprobs and any( - trainable and math.isnan(logprob) - for trainable, logprob in zip( - sampled, - exchange_result.logprobs, - strict=True, - ) + histories = trajectory.histories() + trajectory_results = [] + for exchange_result, history in zip( + exchange_results.histories, histories, strict=True ): - raise RuntimeError( - "Exchange trajectory is missing logprobs for trainable tokens" - ) - exchanges = _exchange_list(trajectory, None) - chat_choices = [ - exchange.response.choices[0] - for exchange in exchanges - if isinstance(exchange, ChatCompletionsExchange) - ] - if len(chat_choices) == len(exchanges): - exact_spans: list[tuple[int, int]] = [] - for exchange in exchanges: - prompt, completion, _ = _exchange_tokens(exchange) - if prompt is None or completion is None: - exact_spans = [] - break - exact_spans.append((len(prompt), len(prompt) + len(completion))) - choice_spans = exact_spans or choice_spans - moe_routes, moe_stats = align_choice_routes_to_tokenized_result( - token_ids=exchange_result.token_ids, - choices=chat_choices, - choice_offsets=[start for start, _ in choice_spans], - choice_token_lengths=[ - end - start for start, end in choice_spans - ], - ) - else: - if any( - choice_moe_routing_metadata(choice) is not None - for choice in chat_choices + sampled = [ + bool(flag & TokenFlag.SAMPLED) for flag in exchange_result.flags + ] + choice_spans = _flag_spans(exchange_result.flags, TokenFlag.SAMPLED) + if not allow_training_without_logprobs and any( + trainable and math.isnan(logprob) + for trainable, logprob in zip( + sampled, + exchange_result.logprobs, + strict=True, + ) ): raise RuntimeError( - "MoE routing replay requires an all-Chat-Completions " - "exchange trajectory" + "Exchange trajectory is missing logprobs for trainable tokens" + ) + chat_choices = [] + if isinstance(history, ChatCompletionsHistory): + seen_choices: set[tuple[int, int]] = set() + for source in history.message_sources: + if source is None or source.choice_index is None: + continue + key = (id(source.exchange), source.choice_index) + if key in seen_choices: + continue + seen_choices.add(key) + if not isinstance(source.exchange, ChatCompletionsExchange): + continue + chat_choices.append( + next( + choice + for choice in source.exchange.response.choices + if choice.index == source.choice_index + ) + ) + if len(chat_choices) == len(choice_spans): + moe_routes, moe_stats = align_choice_routes_to_tokenized_result( + token_ids=exchange_result.token_ids, + choices=chat_choices, + choice_offsets=[start for start, _ in choice_spans], + choice_token_lengths=[ + end - start for start, end in choice_spans + ], + ) + else: + if any( + choice_moe_routing_metadata(choice) is not None + for choice in chat_choices + ): + raise RuntimeError( + "MoE routing replay requires complete Chat Completions " + "history sources" + ) + moe_routes = None + moe_stats = MoeRoutingAlignmentStats() + trajectory_results.append( + TokenizedResult( + advantage=advantage, + chat="", + token_ids=exchange_result.token_ids, + input_pos=list(range(len(exchange_result.token_ids))), + assistant_mask=[int(value) for value in sampled], + logprobs=exchange_result.logprobs, + pixel_values=None, + image_grid_thw=None, + trajectory=trajectory, + choice_offsets=[start for start, _ in choice_spans], + extra_logprobs={}, + moe_routed_experts=moe_routes, + moe_routing_alignment_stats=moe_stats, + _tokenizer=tokenizer, ) - moe_routes = None - moe_stats = MoeRoutingAlignmentStats() - trajectory_results = [ - TokenizedResult( - advantage=advantage, - chat="", - token_ids=exchange_result.token_ids, - input_pos=list(range(len(exchange_result.token_ids))), - assistant_mask=[int(value) for value in sampled], - logprobs=exchange_result.logprobs, - pixel_values=None, - image_grid_thw=None, - trajectory=trajectory, - choice_offsets=[start for start, _ in choice_spans], - extra_logprobs={}, - moe_routed_experts=moe_routes, - moe_routing_alignment_stats=moe_stats, - _tokenizer=tokenizer, ) - ] else: trajectory_results = assemble_vllm_training_sequences( tokenizer=tokenizer, histories=[ - History( + LegacyHistory( messages_and_choices=trajectory.messages_and_choices, tools=trajectory.tools, ), @@ -707,7 +711,7 @@ def tokenize_trajectory_groups( def tokenize_trajectory( tokenizer: "PreTrainedTokenizerBase", image_processor: BaseImageProcessor | None, - history: History, + history: LegacyHistory, advantage: float, allow_training_without_logprobs: bool, trajectory: Trajectory, diff --git a/src/art/tinker/server.py b/src/art/tinker/server.py index 7abe85865..1767fa868 100644 --- a/src/art/tinker/server.py +++ b/src/art/tinker/server.py @@ -572,12 +572,12 @@ async def messages_and_choices_prompt_tokens_and_choice_offsets( tools: Tools | None, ) -> tuple[list[int], list[int]] | None: from art.preprocessing.tokenize import tokenize_trajectory - from art.trajectories import History, Trajectory + from art.trajectories import LegacyHistory, Trajectory result = tokenize_trajectory( tokenizer=self._get_renderer(base_model).tokenizer, image_processor=None, - history=History( + history=LegacyHistory( messages_and_choices=messages_and_choices, tools=tools, ), diff --git a/src/art/tinker_native/data.py b/src/art/tinker_native/data.py index 4fd14e8ef..9c1428e9b 100644 --- a/src/art/tinker_native/data.py +++ b/src/art/tinker_native/data.py @@ -10,13 +10,12 @@ import torch from ..trajectories import ( - History, + LegacyHistory, TokenFlag, - TokenizedTrajectory, + TokenizedHistory, Trajectory, TrajectoryGroup, get_messages, - tokenize_trajectory, ) from ..types import MessagesAndChoices @@ -163,11 +162,13 @@ def trajectory_groups_to_datums( continue for trajectory, advantage in zip(group.trajectories, advantages): if trajectory.exchanges: - datum = _tokenized_trajectory_to_datum( - tokenize_trajectory(trajectory, base_model=base_model), advantage + tokenized = trajectory.tokenize( + multi_history=True, base_model=base_model ) - if datum is not None: - datums.append(datum) + for history in tokenized.histories: + datum = _tokenized_trajectory_to_datum(history, advantage) + if datum is not None: + datums.append(datum) continue for history in iter_trajectory_histories(trajectory): datum = history_to_datum(history, advantage, renderer, tokenizer) @@ -177,8 +178,8 @@ def trajectory_groups_to_datums( return datums -def iter_trajectory_histories(trajectory: Trajectory) -> Iterable[History]: - yield History( +def iter_trajectory_histories(trajectory: Trajectory) -> Iterable[LegacyHistory]: + yield LegacyHistory( messages_and_choices=trajectory.messages_and_choices, tools=trajectory.tools, ) @@ -222,7 +223,7 @@ def extract_logprobs_from_choice( def history_to_datum( - history: History, + history: LegacyHistory, advantage: float, renderer: Any, tokenizer: Any, @@ -283,7 +284,7 @@ def build_datum( def _tokenized_trajectory_to_datum( - tokenized: TokenizedTrajectory, advantage: float + tokenized: TokenizedHistory, advantage: float ) -> tinker.Datum | None: sampled = [bool(flag & TokenFlag.SAMPLED) for flag in tokenized.flags] if not (len(tokenized.token_ids) == len(tokenized.logprobs) == len(sampled)): diff --git a/src/art/trajectories/__init__.py b/src/art/trajectories/__init__.py index fa0c9a87a..87067648d 100644 --- a/src/art/trajectories/__init__.py +++ b/src/art/trajectories/__init__.py @@ -9,11 +9,21 @@ Mapping, ) from contextlib import asynccontextmanager +from dataclasses import dataclass from datetime import datetime from enum import IntFlag import time from types import TracebackType -from typing import Annotated, Any, Literal, TypeAlias, overload +from typing import ( + Annotated, + Any, + Generic, + Literal, + Protocol, + TypeAlias, + TypeVar, + overload, +) from anthropic.types import ( Message as AnthropicMessage, @@ -43,6 +53,9 @@ from openai.types.responses import ( ToolParam as ResponsesToolParam, ) +from openai.types.responses.response_create_params import ( + Conversation as ResponsesConversation, +) import pydantic from typing_extensions import TypedDict, deprecated @@ -57,6 +70,23 @@ MetadataValue = Any +class Tokenizer(Protocol): + """Minimal tokenizer surface used by trajectory tokenization.""" + + def __call__(self, text: str, *, add_special_tokens: bool = False) -> object: ... + + def apply_chat_template( + self, + messages: list[dict[str, Any]], + *, + tools: object, + tokenize: bool, + add_generation_prompt: bool, + chat_template: str | None = None, + **kwargs: object, + ) -> object: ... + + class TokenFlag(IntFlag): """Independent facts about a token; members may be combined.""" @@ -96,6 +126,7 @@ class CompletionsRequest(TypedDict, total=False, extra_items=Any): echo: bool stop: str | list[str] seed: int + suffix: str class ResponsesRequest(TypedDict, total=False, extra_items=Any): @@ -105,6 +136,7 @@ class ResponsesRequest(TypedDict, total=False, extra_items=Any): input: str | ResponseInputParam instructions: str previous_response_id: str + conversation: ResponsesConversation stream: bool tools: list[ResponsesToolParam] max_output_tokens: int @@ -216,7 +248,7 @@ class PydanticException(pydantic.BaseModel): traceback: str -class History(pydantic.BaseModel): +class LegacyHistory(pydantic.BaseModel): messages_and_choices: MessagesAndChoices tools: Tools | None = None @@ -227,18 +259,112 @@ def serialize_messages_and_choices(self, value: MessagesAndChoices) -> list[Any] def messages(self) -> Messages: return get_messages(self.messages_and_choices) - def as_chat_completions_history(self) -> ChatCompletionsHistory: + def as_chat_completions_history( + self, *, model: str | None = None + ) -> ChatCompletionsHistory: from ._history import legacy_as_chat_completions_history - return legacy_as_chat_completions_history(self) + return legacy_as_chat_completions_history(self, model=model) + + def tokenize( + self, + *, + model: str, + base_model: str | None = None, + tokenizer: Tokenizer | None = None, + chat_template: str | None = None, + chat_template_kwargs: Mapping[str, object] | None = None, + ) -> TokenizedHistory: + from ._tokenize import tokenize_history + + return tokenize_history( + self, + model=model, + base_model=base_model, + tokenizer=tokenizer, + chat_template=chat_template, + chat_template_kwargs=chat_template_kwargs, + ) -class ChatCompletionsHistory(pydantic.BaseModel): +class History: + """Mutable, protocol-native view of one tokenizable sequence.""" + model: str | None - messages: Annotated[pydantic.SerializeAsAny[Messages], pydantic.SkipValidation] - tools: Annotated[pydantic.SerializeAsAny[Tools | None], pydantic.SkipValidation] = ( - None - ) + + def as_chat_completions_history(self) -> ChatCompletionsHistory: + raise ValueError( + f"{type(self).__name__} cannot be represented as Chat Completions" + ) + + def tokenize( + self, + *, + base_model: str | None = None, + tokenizer: Tokenizer | None = None, + chat_template: str | None = None, + chat_template_kwargs: Mapping[str, object] | None = None, + ) -> TokenizedHistory: + from ._tokenize import tokenize_history + + return tokenize_history( + self, + model=self.model, + base_model=base_model, + tokenizer=tokenizer, + chat_template=chat_template, + chat_template_kwargs=chat_template_kwargs, + ) + + +@dataclass(frozen=True, slots=True) +class ChatCompletionsMessageSource: + exchange: ChatCompletionsExchange | MessagesExchange | ResponsesExchange + request_index: int | None = None + choice_index: int | None = None + output_index: int | None = None + + +@dataclass(frozen=True, slots=True) +class AnthropicMessageSource: + exchange: MessagesExchange + request_index: int | None = None + + +@dataclass(frozen=True, slots=True) +class ResponsesItemSource: + exchange: ResponsesExchange + request_index: int | None = None + output_index: int | None = None + + +@dataclass(frozen=True, slots=True) +class CompletionsSource: + exchange: CompletionsExchange + prompt_index: int + choice_index: int | None = None + + +@dataclass(frozen=True, slots=True) +class CompletionsTokenSourceSpan: + start: int + end: int + source: CompletionsSource | None + + +@dataclass(frozen=True, slots=True) +class CompletionsStringSourceSpan: + start: int + end: int + source: CompletionsSource | None + + +@dataclass +class ChatCompletionsHistory(History): + model: str | None + messages: Messages + message_sources: list[ChatCompletionsMessageSource | None] + tools: Tools | None = None chat_template: str | None = None chat_template_kwargs: dict[str, Any] | None = None @@ -246,23 +372,14 @@ def as_chat_completions_history(self) -> ChatCompletionsHistory: return self -class AnthropicMessagesHistory(pydantic.BaseModel): +@dataclass +class AnthropicMessagesHistory(History): model: str - system: Annotated[ - pydantic.SerializeAsAny[str | list[AnthropicTextBlockParam] | None], - pydantic.SkipValidation, - ] = None - messages: Annotated[ - pydantic.SerializeAsAny[list[AnthropicMessageParam]], pydantic.SkipValidation - ] - tools: Annotated[ - pydantic.SerializeAsAny[list[AnthropicToolParam] | None], - pydantic.SkipValidation, - ] = None - thinking: Annotated[ - pydantic.SerializeAsAny[AnthropicThinkingConfigParam | None], - pydantic.SkipValidation, - ] = None + messages: list[AnthropicMessageParam] + message_sources: list[AnthropicMessageSource | None] + system: str | list[AnthropicTextBlockParam] | None = None + system_source: MessagesExchange | None = None + tools: list[AnthropicToolParam] | None = None chat_template: str | None = None chat_template_kwargs: dict[str, Any] | None = None @@ -272,16 +389,16 @@ def as_chat_completions_history(self) -> ChatCompletionsHistory: return anthropic_as_chat_completions_history(self) -class ResponsesHistory(pydantic.BaseModel): +@dataclass +class ResponsesHistory(History): model: str - input: Annotated[ - pydantic.SerializeAsAny[ResponseInputParam], pydantic.SkipValidation - ] + input: ResponseInputParam + input_sources: list[ResponsesItemSource | None] instructions: str | None = None - tools: Annotated[ - pydantic.SerializeAsAny[list[ResponsesToolParam] | None], - pydantic.SkipValidation, - ] = None + instructions_source: ResponsesExchange | None = None + tools: list[ResponsesToolParam] | None = None + conversation: ResponsesConversation | None = None + previous_response_id: str | None = None chat_template: str | None = None chat_template_kwargs: dict[str, Any] | None = None @@ -291,9 +408,22 @@ def as_chat_completions_history(self) -> ChatCompletionsHistory: return responses_as_chat_completions_history(self) -class CompletionsHistory(pydantic.BaseModel): +@dataclass +class CompletionsTokenHistory(History): model: str - token_ids: list[int] + prompt: list[int] + prompt_sources: list[CompletionsTokenSourceSpan] + sampled_spans: list[tuple[int, int]] + + def as_chat_completions_history(self) -> ChatCompletionsHistory: + raise ValueError("Raw Completions history has no chat-message structure") + + +@dataclass +class CompletionsStringHistory(History): + model: str + prompt: str + prompt_sources: list[CompletionsStringSourceSpan] sampled_spans: list[tuple[int, int]] def as_chat_completions_history(self) -> ChatCompletionsHistory: @@ -301,11 +431,12 @@ def as_chat_completions_history(self) -> ChatCompletionsHistory: TrajectoryHistory: TypeAlias = ( - History + LegacyHistory | ChatCompletionsHistory | AnthropicMessagesHistory | ResponsesHistory - | CompletionsHistory + | CompletionsTokenHistory + | CompletionsStringHistory ) @@ -316,7 +447,7 @@ class Trajectory(_CompactModel): exclude_if=lambda value: not value, ) tools: Tools | None = None - additional_histories: list[History] = pydantic.Field( + additional_histories: list[LegacyHistory] = pydantic.Field( default_factory=list, ) reward: float = 0.0 @@ -385,6 +516,13 @@ def chat_completions_history( return chat_completions_history(self, model=model) + def chat_completions_histories( + self, *, model: str | None = None + ) -> list[ChatCompletionsHistory]: + from ._history import chat_completions_histories + + return chat_completions_histories(self, model=model) + def anthropic_messages_history( self, *, model: str | None = None ) -> AnthropicMessagesHistory: @@ -392,21 +530,109 @@ def anthropic_messages_history( return anthropic_messages_history(self, model=model) + def anthropic_messages_histories( + self, *, model: str | None = None + ) -> list[AnthropicMessagesHistory]: + from ._history import anthropic_messages_histories + + return anthropic_messages_histories(self, model=model) + def responses_history(self, *, model: str | None = None) -> ResponsesHistory: from ._history import responses_history return responses_history(self, model=model) - def completions_history(self, *, model: str | None = None) -> CompletionsHistory: - from ._history import completions_history + def responses_histories( + self, *, model: str | None = None + ) -> list[ResponsesHistory]: + from ._history import responses_histories + + return responses_histories(self, model=model) + + def completions_token_history( + self, *, model: str | None = None + ) -> CompletionsTokenHistory: + from ._history import completions_token_history + + return completions_token_history(self, model=model) + + def completions_token_histories( + self, *, model: str | None = None + ) -> list[CompletionsTokenHistory]: + from ._history import completions_token_histories + + return completions_token_histories(self, model=model) + + def completions_string_history( + self, *, model: str | None = None + ) -> CompletionsStringHistory: + from ._history import completions_string_history + + return completions_string_history(self, model=model) - return completions_history(self, model=model) + def completions_string_histories( + self, *, model: str | None = None + ) -> list[CompletionsStringHistory]: + from ._history import completions_string_histories + + return completions_string_histories(self, model=model) def history(self, *, model: str | None = None) -> TrajectoryHistory: from ._history import trajectory_history return trajectory_history(self, model=model) + def histories(self, *, model: str | None = None) -> list[TrajectoryHistory]: + from ._history import trajectory_histories + + return trajectory_histories(self, model=model) + + @overload + def tokenize( + self, + *, + multi_history: Literal[False] = False, + model: str | None = None, + base_model: str | None = None, + tokenizer: Tokenizer | None = None, + chat_template: str | None = None, + chat_template_kwargs: Mapping[str, object] | None = None, + ) -> TokenizedTrajectory: ... + + @overload + def tokenize( + self, + *, + multi_history: Literal[True], + model: str | None = None, + base_model: str | None = None, + tokenizer: Tokenizer | None = None, + chat_template: str | None = None, + chat_template_kwargs: Mapping[str, object] | None = None, + ) -> TokenizedMultiHistoryTrajectory: ... + + def tokenize( + self, + *, + multi_history: bool = False, + model: str | None = None, + base_model: str | None = None, + tokenizer: Tokenizer | None = None, + chat_template: str | None = None, + chat_template_kwargs: Mapping[str, object] | None = None, + ) -> TokenizedTrajectory | TokenizedMultiHistoryTrajectory: + from ._tokenize import tokenize_trajectory + + return tokenize_trajectory( + self, + multi_history=multi_history, + model=model, + base_model=base_model, + tokenizer=tokenizer, + chat_template=chat_template, + chat_template_kwargs=chat_template_kwargs, + ) + def messages(self) -> Messages: from ._history import trajectory_messages @@ -527,47 +753,121 @@ def __iter__(self) -> Iterator[Trajectory]: # ty: ignore[invalid-method-overrid def __len__(self) -> int: return len(self.trajectories) + @overload + def tokenize( + self, + *, + multi_history: Literal[False] = False, + model: str | None = None, + base_model: str | None = None, + tokenizer: Tokenizer | None = None, + chat_template: str | None = None, + chat_template_kwargs: Mapping[str, object] | None = None, + ) -> TokenizedTrajectoryGroup[TokenizedTrajectory]: ... -class TokenizedTrajectory(pydantic.BaseModel): + @overload + def tokenize( + self, + *, + multi_history: Literal[True], + model: str | None = None, + base_model: str | None = None, + tokenizer: Tokenizer | None = None, + chat_template: str | None = None, + chat_template_kwargs: Mapping[str, object] | None = None, + ) -> TokenizedTrajectoryGroup[TokenizedMultiHistoryTrajectory]: ... + + def tokenize( + self, + *, + multi_history: bool = False, + model: str | None = None, + base_model: str | None = None, + tokenizer: Tokenizer | None = None, + chat_template: str | None = None, + chat_template_kwargs: Mapping[str, object] | None = None, + ) -> ( + TokenizedTrajectoryGroup[TokenizedTrajectory] + | TokenizedTrajectoryGroup[TokenizedMultiHistoryTrajectory] + ): + from ._tokenize import tokenize_group + + return tokenize_group( + self, + multi_history=multi_history, + model=model, + base_model=base_model, + tokenizer=tokenizer, + chat_template=chat_template, + chat_template_kwargs=chat_template_kwargs, + ) + + +class TokenizedHistory(pydantic.BaseModel): + model: str token_ids: list[int] logprobs: list[float] flags: list[TokenFlag] - underlying: Trajectory + @pydantic.model_validator(mode="after") + def validate_tokenwise_lengths(self) -> TokenizedHistory: + if not (len(self.token_ids) == len(self.logprobs) == len(self.flags)): + raise ValueError("Tokenized history fields differ in length") + return self -class TokenizedTrajectoryGroup(pydantic.BaseModel): - trajectories: list[TokenizedTrajectory] - underlying: TrajectoryGroup + +class TokenizedTrajectory(TokenizedHistory): + reward: float + metrics: dict[str, float | int | bool] + metadata: dict[str, MetadataValue] + + +class TokenizedMultiHistoryTrajectory(pydantic.BaseModel): + histories: list[TokenizedHistory] + reward: float + metrics: dict[str, float | int | bool] + metadata: dict[str, MetadataValue] + + +TokenizedTrajectoryT = TypeVar( + "TokenizedTrajectoryT", TokenizedTrajectory, TokenizedMultiHistoryTrajectory +) + + +class TokenizedTrajectoryGroup(pydantic.BaseModel, Generic[TokenizedTrajectoryT]): + trajectories: list[TokenizedTrajectoryT] + metrics: dict[str, float | int | bool] + metadata: dict[str, MetadataValue] @overload -def current_trajectory(*, required: Literal[True]) -> Trajectory: ... +def current_trajectory(*, require: Literal[True]) -> Trajectory: ... @overload -def current_trajectory(*, required: Literal[False] = False) -> Trajectory | None: ... +def current_trajectory(*, require: Literal[False] = False) -> Trajectory | None: ... -def current_trajectory(*, required: bool = False) -> Trajectory | None: +def current_trajectory(*, require: bool = False) -> Trajectory | None: from ._scope import get_current_trajectory - return get_current_trajectory(required=required) + return get_current_trajectory(required=require) @overload -def current_trajectory_group(*, required: Literal[True]) -> TrajectoryGroup: ... +def current_trajectory_group(*, require: Literal[True]) -> TrajectoryGroup: ... @overload def current_trajectory_group( - *, required: Literal[False] = False + *, require: Literal[False] = False ) -> TrajectoryGroup | None: ... -def current_trajectory_group(*, required: bool = False) -> TrajectoryGroup | None: +def current_trajectory_group(*, require: bool = False) -> TrajectoryGroup | None: from ._scope import get_current_trajectory_group - return get_current_trajectory_group(required=required) + return get_current_trajectory_group(required=require) async def trajectory(coroutine: Coroutine[Any, Any, object]) -> Trajectory: @@ -589,98 +889,19 @@ async def trajectory_group( ) -def tokenize_trajectory( - trajectory: Trajectory, - *, - base_model: str | None = None, - model: str | None = None, - chat_template: str | None = None, - chat_template_kwargs: Mapping[str, object] | None = None, -) -> TokenizedTrajectory: - from ._tokenize import tokenize_one - - return tokenize_one( - trajectory, - base_model, - model=model, - chat_template=chat_template, - chat_template_kwargs=chat_template_kwargs, - ) - - -def tokenize_trajectories( - trajectories: Iterable[Trajectory], - *, - base_model: str | None = None, - model: str | None = None, - chat_template: str | None = None, - chat_template_kwargs: Mapping[str, object] | None = None, -) -> list[TokenizedTrajectory]: - return [ - tokenize_trajectory( - item, - base_model=base_model, - model=model, - chat_template=chat_template, - chat_template_kwargs=chat_template_kwargs, - ) - for item in trajectories - ] - - -def tokenize_trajectory_group( - group: TrajectoryGroup, - *, - base_model: str | None = None, - model: str | None = None, - chat_template: str | None = None, - chat_template_kwargs: Mapping[str, object] | None = None, -) -> TokenizedTrajectoryGroup: - return TokenizedTrajectoryGroup( - trajectories=tokenize_trajectories( - group, - base_model=base_model, - model=model, - chat_template=chat_template, - chat_template_kwargs=chat_template_kwargs, - ), - underlying=group, - ) - - -def tokenize_trajectory_groups( - groups: Iterable[TrajectoryGroup], - *, - base_model: str | None = None, - model: str | None = None, - chat_template: str | None = None, - chat_template_kwargs: Mapping[str, object] | None = None, -) -> list[TokenizedTrajectoryGroup]: - return [ - tokenize_trajectory_group( - group, - base_model=base_model, - model=model, - chat_template=chat_template, - chat_template_kwargs=chat_template_kwargs, - ) - for group in groups - ] - - @overload @deprecated("Use current_trajectory() instead.") -def auto_trajectory(*, required: Literal[True]) -> Trajectory: ... +def auto_trajectory(*, require: Literal[True]) -> Trajectory: ... @overload @deprecated("Use current_trajectory() instead.") -def auto_trajectory(*, required: Literal[False] = False) -> Trajectory | None: ... +def auto_trajectory(*, require: Literal[False] = False) -> Trajectory | None: ... @deprecated("Use current_trajectory() instead.") -def auto_trajectory(*, required: bool = False) -> Trajectory | None: - return current_trajectory(required=required) +def auto_trajectory(*, require: bool = False) -> Trajectory | None: + return current_trajectory(require=require) @deprecated("Use trajectory() instead.") @@ -708,25 +929,32 @@ def get_messages(messages_and_choices: MessagesAndChoices) -> Messages: "TrajectoryExchanges", "PydanticException", "History", + "LegacyHistory", + "ChatCompletionsMessageSource", + "AnthropicMessageSource", + "ResponsesItemSource", + "CompletionsSource", + "CompletionsTokenSourceSpan", + "CompletionsStringSourceSpan", "ChatCompletionsHistory", "AnthropicMessagesHistory", "ResponsesHistory", - "CompletionsHistory", + "CompletionsTokenHistory", + "CompletionsStringHistory", "TrajectoryHistory", "Trajectory", "TrajectoryGroup", "TokenizedTrajectory", + "TokenizedHistory", + "TokenizedMultiHistoryTrajectory", "TokenizedTrajectoryGroup", + "Tokenizer", "TokenFlag", "MetadataValue", "current_trajectory", "current_trajectory_group", "trajectory", "trajectory_group", - "tokenize_trajectory", - "tokenize_trajectories", - "tokenize_trajectory_group", - "tokenize_trajectory_groups", "auto_trajectory", "capture_auto_trajectory", "get_messages", diff --git a/src/art/trajectories/_history.py b/src/art/trajectories/_history.py index 48300348b..112d555df 100644 --- a/src/art/trajectories/_history.py +++ b/src/art/trajectories/_history.py @@ -1,23 +1,48 @@ from __future__ import annotations -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import Callable, Iterable, Mapping, Sequence import copy +from dataclasses import dataclass from datetime import datetime -from typing import Protocol, TypeVar +import json +from typing import Generic, Protocol, TypeVar, cast +from anthropic.types import ( + MessageParam as AnthropicMessageParam, +) +from anthropic.types import ( + TextBlockParam as AnthropicTextBlockParam, +) +from anthropic.types import ( + ToolUnionParam as AnthropicToolParam, +) +from openai.types import CompletionChoice +from openai.types.responses import ResponseInputParam +from openai.types.responses import ToolParam as ResponsesToolParam +from openai.types.responses.response_create_params import ( + Conversation as ResponsesConversation, +) +from openai.types.responses.response_input_param import ResponseInputItemParam import pydantic from ..types import Message, Messages, Tools from . import ( AnthropicMessagesHistory, + AnthropicMessageSource, ChatCompletionsExchange, ChatCompletionsHistory, + ChatCompletionsMessageSource, CompletionsExchange, - CompletionsHistory, - History, + CompletionsSource, + CompletionsStringHistory, + CompletionsStringSourceSpan, + CompletionsTokenHistory, + CompletionsTokenSourceSpan, + LegacyHistory, MessagesExchange, ResponsesExchange, ResponsesHistory, + ResponsesItemSource, Trajectory, TrajectoryHistory, ) @@ -33,71 +58,195 @@ def model(self) -> str | None: ... _ExchangeT = TypeVar("_ExchangeT", bound=_ModelledExchange) _ItemT = TypeVar("_ItemT") +_SourceT = TypeVar("_SourceT") +_ContextT = TypeVar("_ContextT") _MESSAGES = pydantic.TypeAdapter(Messages) _MESSAGE = pydantic.TypeAdapter(Message) _TOOLS = pydantic.TypeAdapter(Tools | None) +_CHAT_KWARGS = pydantic.TypeAdapter(dict[str, object] | None) +_ANTHROPIC_SYSTEM = pydantic.TypeAdapter(str | list[AnthropicTextBlockParam] | None) +_ANTHROPIC_TOOLS = pydantic.TypeAdapter(list[AnthropicToolParam] | None) +_RESPONSE_TOOLS = pydantic.TypeAdapter(list[ResponsesToolParam] | None) +_RESPONSE_CONVERSATION = pydantic.TypeAdapter(ResponsesConversation | None) + + +@dataclass(frozen=True) +class _ChatContext: + tools: Tools | None + template: str | None + kwargs: dict[str, object] | None + + +@dataclass(frozen=True) +class _AnthropicContext: + system: str | list[AnthropicTextBlockParam] | None + tools: list[AnthropicToolParam] | None + template: str | None + kwargs: dict[str, object] | None + + +@dataclass(frozen=True) +class _ResponsesContext: + instructions: str | None + tools: list[ResponsesToolParam] | None + conversation: ResponsesConversation | None + previous_response_id: str | None + template: str | None + kwargs: dict[str, object] | None + +@dataclass +class _Branch(Generic[_ItemT, _SourceT, _ContextT]): + items: list[_ItemT] + sources: list[_SourceT | None] + context: _ContextT + order: tuple[int, ...] + first_time: datetime -def _select( + +def _selected_models( exchanges: Sequence[_ExchangeT], model: str | None, protocol: str -) -> list[_ExchangeT]: +) -> list[tuple[str, list[_ExchangeT]]]: selected = [ exchange for exchange in exchanges if model is None or exchange.model == model ] if not selected: suffix = f" for model {model!r}" if model is not None else "" raise ValueError(f"Trajectory contains no {protocol} exchanges{suffix}") - models = {exchange.model for exchange in selected} - if None in models: + if any(exchange.model is None for exchange in selected): raise ValueError(f"Every {protocol} exchange must identify its model") - if len(models) != 1: + grouped: dict[str, list[_ExchangeT]] = {} + for exchange in sorted(selected, key=lambda item: (item.start_time, item.end_time)): + if exchange.model is None: + raise AssertionError("model identity was checked above") + grouped.setdefault(exchange.model, []).append(exchange) + return list(grouped.items()) + + +def _one(history: Sequence[_ItemT], protocol: str) -> _ItemT: + if len(history) != 1: raise ValueError( - f"{protocol} history requires exactly one model; pass model= to select one" + f"{protocol} requires exactly one history; found {len(history)}" ) - return sorted( - selected, key=lambda exchange: (exchange.start_time, exchange.end_time) - ) + return history[0] -def _require_constant(values: Iterable[object], field: str) -> None: - iterator = iter(values) - first = next(iterator) - if any(value != first for value in iterator): - raise ValueError(f"Exchanges with different {field} form different histories") +def _is_prefix(prefix: Sequence[object], value: Sequence[object]) -> bool: + return len(prefix) <= len(value) and list(value[: len(prefix)]) == list(prefix) -def _require_context( - requests: Sequence[Mapping[str, object]], fields: Sequence[str] +def _extend_branches( + branches: list[_Branch[_ItemT, _SourceT, _ContextT]], + *, + prompt: Sequence[_ItemT], + prompt_sources: Sequence[_SourceT | None], + outputs: Sequence[tuple[int, Sequence[_ItemT], Sequence[_SourceT | None]]], + context: _ContextT, + sequence: int, + start_time: datetime, + continuation: Callable[[_Branch[_ItemT, _SourceT, _ContextT]], bool] | None = None, ) -> None: - for field in fields: - _require_constant((request.get(field) for request in requests), field) + if len(prompt) != len(prompt_sources): + raise AssertionError("prompt sources must parallel prompt items") + candidates = [ + (index, branch) + for index, branch in enumerate(branches) + if branch.context == context + and _is_prefix(branch.items, prompt) + and (continuation is None or continuation(branch)) + ] + parent: _Branch[_ItemT, _SourceT, _ContextT] | None = None + remove_parent = False + if candidates: + index, parent = min( + candidates, + key=lambda item: (-len(item[1].items), item[1].order), + ) + remove_parent = True + branches.pop(index) + sources = [ + *parent.sources, + *prompt_sources[len(parent.items) :], + ] + else: + regenerations = [ + branch + for branch in branches + if branch.context == context and _is_prefix(prompt, branch.items) + ] + if regenerations: + parent = min(regenerations, key=lambda branch: branch.order) + sources = copy.copy(parent.sources[: len(prompt)]) + else: + sources = copy.copy(prompt_sources) + base_order = parent.order if parent is not None else (sequence,) + first_time = parent.first_time if parent is not None else start_time + created = [ + _Branch( + items=[*copy.deepcopy(prompt), *copy.deepcopy(output)], + sources=[*sources, *output_sources], + context=copy.deepcopy(context), + order=(*base_order, choice_index), + first_time=first_time, + ) + for choice_index, output, output_sources in outputs + ] + if not created and remove_parent and parent is not None: + branches.append(parent) + branches.extend(created) -def _extend( - history: list[_ItemT] | None, - prompt: Sequence[_ItemT], - completion: Sequence[_ItemT], - protocol: str, -) -> list[_ItemT]: - if history is None: - history = copy.deepcopy(list(prompt)) - elif len(prompt) < len(history) or list(prompt[: len(history)]) != history: - raise ValueError(f"{protocol} exchanges do not form one append-only history") - else: - history.extend(copy.deepcopy(list(prompt[len(history) :]))) - history.extend(copy.deepcopy(list(completion))) - return history +def _contains_tokens(tokens: Sequence[int], sampled: Sequence[int]) -> bool: + return any( + list(tokens[start : start + len(sampled)]) == list(sampled) + for start in range(len(tokens) - len(sampled) + 1) + ) -def _only_choice(exchange: ChatCompletionsExchange | CompletionsExchange) -> None: - if len(exchange.response.choices) != 1: - raise ValueError("Multiple response choices form multiple histories") +def _chat_retains_sampled_reasoning( + branch: _Branch[Message, ChatCompletionsMessageSource, _ChatContext], + exchange: ChatCompletionsExchange, + prompt_length: int, +) -> bool: + """Split histories when a template drops a prior sampled reasoning span.""" -def _model(exchange: _ModelledExchange) -> str: - if exchange.model is None: - raise AssertionError("_select returned an exchange without a model") - return exchange.model + if not exchange.response.choices: + return True + from ._tokenize import _chat_choice_tokens + + response_data = exchange.response.model_dump(mode="python") + prompt_ids, _, _ = _chat_choice_tokens(exchange.response.choices[0], response_data) + if prompt_ids is None: + return True + for message, source in zip( + branch.items[:prompt_length], branch.sources[:prompt_length], strict=True + ): + reasoning = message.get("reasoning") or message.get("reasoning_content") + if ( + not reasoning + or source is None + or source.choice_index is None + or not isinstance(source.exchange, ChatCompletionsExchange) + ): + continue + source_response = source.exchange.response + choice = next( + ( + item + for item in source_response.choices + if item.index == source.choice_index + ), + None, + ) + if choice is None: + continue + _, sampled_ids, _ = _chat_choice_tokens( + choice, source_response.model_dump(mode="python") + ) + if sampled_ids and not _contains_tokens(prompt_ids, sampled_ids): + return False + return True def _require_unmixed(trajectory: Trajectory) -> None: @@ -111,201 +260,643 @@ def _require_unmixed(trajectory: Trajectory) -> None: ) -def legacy_as_chat_completions_history(history: History) -> ChatCompletionsHistory: +def legacy_as_chat_completions_history( + history: LegacyHistory, *, model: str | None +) -> ChatCompletionsHistory: + messages = history.messages() return ChatCompletionsHistory( - model=None, - messages=history.messages(), + model=model, + messages=messages, + message_sources=[None] * len(messages), tools=copy.deepcopy(history.tools), ) -def chat_completions_history( +def chat_completions_histories( trajectory: Trajectory, *, model: str | None -) -> ChatCompletionsHistory: +) -> list[ChatCompletionsHistory]: _require_unmixed(trajectory) if not trajectory.exchanges: - if trajectory.additional_histories: - raise ValueError("Trajectory contains multiple legacy histories") - return ChatCompletionsHistory( - model=model, - messages=History( - messages_and_choices=trajectory.messages_and_choices, - tools=trajectory.tools, - ).messages(), - tools=copy.deepcopy(trajectory.tools), - ) - exchanges = _select( + return [ + history.as_chat_completions_history(model=model) + for history in [ + LegacyHistory( + messages_and_choices=trajectory.messages_and_choices, + tools=trajectory.tools, + ), + *trajectory.additional_histories, + ] + ] + + histories: list[ChatCompletionsHistory] = [] + for selected_model, exchanges in _selected_models( trajectory.exchanges.chat_completions, model, "Chat Completions" - ) - _require_context( - [exchange.request for exchange in exchanges], - ("tools", "chat_template", "chat_template_kwargs", "cache_salt"), - ) - messages: Messages | None = None - for exchange in exchanges: - _only_choice(exchange) - prompt = _MESSAGES.validate_python(exchange.request.get("messages", [])) - response = _MESSAGE.validate_python( - exchange.response.choices[0].message.model_dump( - mode="python", exclude_none=True + ): + branches: list[ + _Branch[Message, ChatCompletionsMessageSource, _ChatContext] + ] = [] + for sequence, exchange in enumerate(exchanges): + prompt = _MESSAGES.validate_python(exchange.request.get("messages", [])) + prompt_sources = [ + ChatCompletionsMessageSource(exchange=exchange, request_index=index) + for index in range(len(prompt)) + ] + outputs: list[ + tuple[ + int, + list[Message], + list[ChatCompletionsMessageSource | None], + ] + ] = [] + for choice in sorted( + exchange.response.choices, key=lambda item: item.index + ): + response = _MESSAGE.validate_python( + choice.message.model_dump(mode="python", exclude_none=True) + ) + outputs.append( + ( + choice.index, + [response], + [ + ChatCompletionsMessageSource( + exchange=exchange, choice_index=choice.index + ) + ], + ) + ) + _extend_branches( + branches, + prompt=prompt, + prompt_sources=prompt_sources, + outputs=outputs, + context=_ChatContext( + tools=_TOOLS.validate_python(exchange.request.get("tools")), + template=exchange.request.get("chat_template"), + kwargs=_CHAT_KWARGS.validate_python( + exchange.request.get("chat_template_kwargs") + ), + ), + sequence=sequence, + start_time=exchange.start_time, + continuation=lambda branch: _chat_retains_sampled_reasoning( + branch, exchange, len(prompt) + ), ) + for branch in sorted(branches, key=lambda item: (item.first_time, item.order)): + histories.append( + ChatCompletionsHistory( + model=selected_model, + messages=copy.deepcopy(branch.items), + message_sources=copy.copy(branch.sources), + tools=copy.deepcopy(branch.context.tools), + chat_template=branch.context.template, + chat_template_kwargs=copy.deepcopy(branch.context.kwargs), + ) + ) + return histories + + +def chat_completions_history( + trajectory: Trajectory, *, model: str | None +) -> ChatCompletionsHistory: + histories = chat_completions_histories(trajectory, model=model) + if model is None and len({history.model for history in histories}) > 1: + raise ValueError( + "Chat Completions history requires exactly one model; pass model= to select one" ) - messages = _extend(messages, prompt, [response], "Chat Completions") - first = exchanges[0].request - return ChatCompletionsHistory( - model=_model(exchanges[0]), - messages=messages or [], - tools=copy.deepcopy(first.get("tools")), - chat_template=first.get("chat_template"), - chat_template_kwargs=copy.deepcopy(first.get("chat_template_kwargs")), - ) + return _one(histories, "Chat Completions history") + + +def anthropic_messages_histories( + trajectory: Trajectory, *, model: str | None +) -> list[AnthropicMessagesHistory]: + _require_unmixed(trajectory) + histories: list[AnthropicMessagesHistory] = [] + for selected_model, exchanges in _selected_models( + trajectory.exchanges.messages, model, "Anthropic Messages" + ): + introduced: dict[tuple[object, ...], AnthropicMessageSource] = {} + branches: list[ + _Branch[AnthropicMessageParam, AnthropicMessageSource, _AnthropicContext] + ] = [] + for sequence, exchange in enumerate(exchanges): + prompt = [ + copy.deepcopy(message) + for message in exchange.request.get("messages", []) + ] + prompt_sources = [ + introduced.get( + _anthropic_message_key(message), + AnthropicMessageSource(exchange=exchange, request_index=index), + ) + for index, message in enumerate(prompt) + ] + response = cast( + AnthropicMessageParam, + { + "role": "assistant", + "content": [ + block.model_dump(mode="json", exclude_none=True) + for block in exchange.response.content + ], + }, + ) + response_source = AnthropicMessageSource(exchange=exchange) + introduced.setdefault(_anthropic_message_key(response), response_source) + introduced.setdefault( + _anthropic_message_key(response, visible_only=True), + response_source, + ) + _extend_branches( + branches, + prompt=prompt, + prompt_sources=prompt_sources, + outputs=[ + ( + 0, + [response], + [response_source], + ) + ], + context=_AnthropicContext( + system=_ANTHROPIC_SYSTEM.validate_python( + exchange.request.get("system") + ), + tools=_ANTHROPIC_TOOLS.validate_python( + exchange.request.get("tools") + ), + template=exchange.request.get("chat_template"), + kwargs=_CHAT_KWARGS.validate_python( + exchange.request.get("chat_template_kwargs") + ), + ), + sequence=sequence, + start_time=exchange.start_time, + ) + for branch in branches: + first_source = next( + (source for source in branch.sources if source is not None), None + ) + histories.append( + AnthropicMessagesHistory( + model=selected_model, + messages=copy.deepcopy(branch.items), + message_sources=copy.copy(branch.sources), + system=copy.deepcopy(branch.context.system), + system_source=( + first_source.exchange if first_source is not None else None + ), + tools=copy.deepcopy(branch.context.tools), + chat_template=branch.context.template, + chat_template_kwargs=copy.deepcopy(branch.context.kwargs), + ) + ) + return histories + + +def _anthropic_message_key( + message: AnthropicMessageParam, *, visible_only: bool = False +) -> tuple[object, ...]: + content = message.get("content") + if isinstance(content, str): + blocks: tuple[object, ...] = (("text", content),) + else: + normalized: list[object] = [] + for block in content: + if isinstance(block, pydantic.BaseModel): + data = block.model_dump(mode="json", exclude_none=True) + elif isinstance(block, Mapping): + data = dict(block) + else: + data = {"type": type(block).__name__, "value": str(block)} + kind = data.get("type") + if visible_only and kind in {"thinking", "redacted_thinking"}: + continue + if kind == "text": + normalized.append(("text", data.get("text"))) + else: + normalized.append(json.dumps(data, sort_keys=True, default=str)) + blocks = tuple(normalized) + return (message.get("role"), *blocks) def anthropic_messages_history( trajectory: Trajectory, *, model: str | None ) -> AnthropicMessagesHistory: - _require_unmixed(trajectory) - exchanges = _select(trajectory.exchanges.messages, model, "Anthropic Messages") - _require_context( - [exchange.request for exchange in exchanges], - ( - "system", - "tools", - "thinking", - "chat_template", - "chat_template_kwargs", - "cache_salt", - ), - ) - messages: list[object] | None = None - for exchange in exchanges: - prompt = copy.deepcopy(exchange.request.get("messages", [])) - response = { - "role": "assistant", - "content": [ - block.model_dump(mode="python", exclude_none=True) - for block in exchange.response.content - ], - } - messages = _extend(messages, prompt, [response], "Anthropic Messages") - first = exchanges[0].request - return AnthropicMessagesHistory( - model=_model(exchanges[0]), - system=copy.deepcopy(first.get("system")), - messages=messages or [], - tools=copy.deepcopy(first.get("tools")), - thinking=copy.deepcopy(first.get("thinking")), - chat_template=first.get("chat_template"), - chat_template_kwargs=copy.deepcopy(first.get("chat_template_kwargs")), - ) + histories = anthropic_messages_histories(trajectory, model=model) + if model is None and len({history.model for history in histories}) > 1: + raise ValueError( + "Anthropic Messages history requires exactly one model; pass model= to select one" + ) + return _one(histories, "Anthropic Messages history") -def _responses_input(value: object) -> list[object]: +def _responses_input(value: object) -> ResponseInputParam: if isinstance(value, str): - return [{"role": "user", "content": value}] + return [cast(ResponseInputItemParam, {"role": "user", "content": value})] if isinstance(value, list): - return [copy.deepcopy(item) for item in value] + return [_copy_response_item(item) for item in value] if value is None: return [] raise ValueError("Responses input must be text or a list of input items") -def responses_history(trajectory: Trajectory, *, model: str | None) -> ResponsesHistory: +def _copy_response_item(value: object) -> ResponseInputItemParam: + if not isinstance(value, dict): + raise ValueError("Responses input items must be JSON objects") + # OpenAI models this as a large TypedDict union. Capture already validated + # the wire shape; this boundary makes the detached protocol-native copy. + return cast(ResponseInputItemParam, copy.deepcopy(value)) + + +def responses_histories( + trajectory: Trajectory, *, model: str | None +) -> list[ResponsesHistory]: _require_unmixed(trajectory) - exchanges = _select(trajectory.exchanges.responses, model, "Responses") - _require_context( - [exchange.request for exchange in exchanges], - ( - "instructions", - "tools", - "chat_template", - "chat_template_kwargs", - "cache_salt", - ), - ) - items: list[object] | None = None - previous_response_id: str | None = None - for exchange in exchanges: - request = exchange.request - prompt = _responses_input(request.get("input")) - previous = request.get("previous_response_id") - if previous is not None: - if previous != previous_response_id or items is None: - raise ValueError( - "Responses exchange refers to a response outside this history" + histories: list[ResponsesHistory] = [] + for selected_model, exchanges in _selected_models( + trajectory.exchanges.responses, model, "Responses" + ): + branches: list[ + _Branch[ResponseInputItemParam, ResponsesItemSource, _ResponsesContext] + ] = [] + responses: dict[ + str, tuple[ResponseInputParam, list[ResponsesItemSource | None]] + ] = {} + for sequence, exchange in enumerate(exchanges): + request = exchange.request + prompt = _responses_input(request.get("input")) + prompt_sources: list[ResponsesItemSource | None] = [ + ResponsesItemSource(exchange=exchange, request_index=index) + for index in range(len(prompt)) + ] + previous = request.get("previous_response_id") + if isinstance(previous, str) and previous in responses: + prior_items, prior_sources = responses[previous] + prompt = [*copy.deepcopy(prior_items), *prompt] + prompt_sources = [*copy.copy(prior_sources), *prompt_sources] + output = [ + cast( + ResponseInputItemParam, + item.model_dump(mode="json", exclude_none=True), ) - prompt = [*items, *prompt] - output = [ - item.model_dump(mode="python", exclude_none=True) - for item in exchange.response.output - ] - items = _extend(items, prompt, output, "Responses") - previous_response_id = exchange.response.id - first = exchanges[0].request - return ResponsesHistory( - model=_model(exchanges[0]), - input=items or [], - instructions=first.get("instructions"), - tools=copy.deepcopy(first.get("tools")), - chat_template=first.get("chat_template"), - chat_template_kwargs=copy.deepcopy(first.get("chat_template_kwargs")), - ) + for item in exchange.response.output + ] + output_sources: list[ResponsesItemSource | None] = [ + ResponsesItemSource(exchange=exchange, output_index=index) + for index in range(len(output)) + ] + instructions = request.get("instructions") + if instructions is not None and not isinstance(instructions, str): + raise ValueError("Responses instructions must be text") + external_previous = ( + previous + if isinstance(previous, str) and previous not in responses + else None + ) + _extend_branches( + branches, + prompt=prompt, + prompt_sources=prompt_sources, + outputs=[(0, output, output_sources)], + context=_ResponsesContext( + instructions=instructions, + tools=_RESPONSE_TOOLS.validate_python(request.get("tools")), + conversation=_RESPONSE_CONVERSATION.validate_python( + request.get("conversation") + ), + previous_response_id=external_previous, + template=request.get("chat_template"), + kwargs=_CHAT_KWARGS.validate_python( + request.get("chat_template_kwargs") + ), + ), + sequence=sequence, + start_time=exchange.start_time, + ) + responses[exchange.response.id] = ( + [*copy.deepcopy(prompt), *copy.deepcopy(output)], + [*copy.copy(prompt_sources), *copy.copy(output_sources)], + ) + for branch in branches: + first_source = next( + (source for source in branch.sources if source is not None), None + ) + histories.append( + ResponsesHistory( + model=selected_model, + input=copy.deepcopy(branch.items), + input_sources=copy.copy(branch.sources), + instructions=branch.context.instructions, + instructions_source=( + first_source.exchange if first_source is not None else None + ), + tools=copy.deepcopy(branch.context.tools), + conversation=copy.deepcopy(branch.context.conversation), + previous_response_id=branch.context.previous_response_id, + chat_template=branch.context.template, + chat_template_kwargs=copy.deepcopy(branch.context.kwargs), + ) + ) + return histories -def completions_history( - trajectory: Trajectory, *, model: str | None -) -> CompletionsHistory: +def responses_history(trajectory: Trajectory, *, model: str | None) -> ResponsesHistory: + histories = responses_histories(trajectory, model=model) + if model is None and len({history.model for history in histories}) > 1: + raise ValueError( + "Responses history requires exactly one model; pass model= to select one" + ) + return _one(histories, "Responses history") + + +def _completion_exact_tokens( + exchange: CompletionsExchange, choice_index: int +) -> tuple[list[int] | None, list[int] | None]: from ._tokenize import _completion_tokens, _exact_token_ids - _require_unmixed(trajectory) - exchanges = _select(trajectory.exchanges.completions, model, "Completions") - _require_context([exchange.request for exchange in exchanges], ("cache_salt",)) - token_ids: list[int] = [] - sampled_spans: list[tuple[int, int]] = [] - for index, exchange in enumerate(exchanges): - _only_choice(exchange) - if exchange.request.get("echo") is True: - raise ValueError("Completions history does not support echo=True") - prompt, completion, _ = _completion_tokens(exchange.response) - request_prompt = exchange.request.get("prompt") - if prompt is None and isinstance(request_prompt, list): + choice = next( + choice for choice in exchange.response.choices if choice.index == choice_index + ) + prompt, completion, _ = _completion_tokens( + exchange.response.model_copy(update={"choices": [choice]}), + echo=exchange.request.get("echo") is True, + ) + request_prompt = exchange.request.get("prompt") + if prompt is None and isinstance(request_prompt, list): + try: prompt = _exact_token_ids( request_prompt, field="Completions request prompt" ) - if prompt is None or completion is None: - raise ValueError( - "Completions history requires exact prompt and output token IDs" + except ValueError: + pass + return prompt, completion + + +def completions_token_histories( + trajectory: Trajectory, *, model: str | None +) -> list[CompletionsTokenHistory]: + _require_unmixed(trajectory) + histories: list[CompletionsTokenHistory] = [] + for selected_model, exchanges in _selected_models( + trajectory.exchanges.completions, model, "Completions" + ): + branches: list[_Branch[int, CompletionsSource, tuple[()]]] = [] + for sequence, exchange in enumerate(exchanges): + if exchange.request.get("suffix") is not None: + raise ValueError("Completions suffix is not supported") + for prompt_index, request_prompt in enumerate( + _completion_prompts(exchange.request.get("prompt")) + ): + choices = _completion_choices_for_prompt(exchange, prompt_index) + for choice in choices: + prompt, completion = _completion_exact_tokens( + exchange, choice.index + ) + if prompt is None: + if not isinstance(request_prompt, list): + continue + prompt = copy.deepcopy(request_prompt) + if completion is None: + continue + prompt_source = CompletionsSource( + exchange=exchange, prompt_index=prompt_index + ) + output_source = CompletionsSource( + exchange=exchange, + prompt_index=prompt_index, + choice_index=choice.index, + ) + _extend_branches( + branches, + prompt=prompt, + prompt_sources=[prompt_source] * len(prompt), + outputs=[ + ( + choice.index, + completion, + [output_source] * len(completion), + ) + ], + context=(), + sequence=sequence, + start_time=exchange.start_time, + ) + for branch in branches: + histories.append( + CompletionsTokenHistory( + model=selected_model, + prompt=copy.deepcopy(branch.items), + prompt_sources=_token_source_spans(branch.sources), + sampled_spans=_sampled_spans(branch.sources), + ) ) - if index == 0: - token_ids.extend(prompt) - elif len(prompt) < len(token_ids) or prompt[: len(token_ids)] != token_ids: - raise ValueError( - "Completions exchanges do not form one append-only token history" + if not histories: + raise ValueError("Completions token history requires exact token IDs") + return histories + + +def completions_token_history( + trajectory: Trajectory, *, model: str | None +) -> CompletionsTokenHistory: + return _one( + completions_token_histories(trajectory, model=model), + "Completions token history", + ) + + +def _completion_prompts(value: object) -> list[str | list[int]]: + if isinstance(value, str): + return [value] + if isinstance(value, list): + if all(isinstance(item, int) and not isinstance(item, bool) for item in value): + return [_checked_token_ids(value)] + prompts: list[str | list[int]] = [] + for item in value: + if isinstance(item, str): + prompts.append(item) + elif isinstance(item, list) and all( + isinstance(token, int) and not isinstance(token, bool) for token in item + ): + prompts.append(_checked_token_ids(item)) + else: + raise ValueError("Invalid batched Completions prompt") + return prompts + raise ValueError("Completions prompt must be text or token IDs") + + +def _checked_token_ids(value: Sequence[object]) -> list[int]: + result: list[int] = [] + for token in value: + if not isinstance(token, int) or isinstance(token, bool): + raise ValueError("Completions token prompts must contain integers") + result.append(token) + return result + + +def _completion_choices_for_prompt( + exchange: CompletionsExchange, prompt_index: int +) -> list[CompletionChoice]: + count = len(_completion_prompts(exchange.request.get("prompt"))) + choices = sorted(exchange.response.choices, key=lambda item: item.index) + if count == 1: + return choices + if len(choices) % count: + raise ValueError("Cannot associate Completions choices with batched prompts") + per_prompt = len(choices) // count + start = prompt_index * per_prompt + expected = list(range(start, start + per_prompt)) + selected = choices[start : start + per_prompt] + if [choice.index for choice in selected] != expected: + raise ValueError("Ambiguous Completions choice-to-prompt association") + return selected + + +def completions_string_histories( + trajectory: Trajectory, *, model: str | None +) -> list[CompletionsStringHistory]: + _require_unmixed(trajectory) + histories: list[CompletionsStringHistory] = [] + for selected_model, exchanges in _selected_models( + trajectory.exchanges.completions, model, "Completions" + ): + branches: list[_Branch[str, CompletionsSource, tuple[()]]] = [] + for sequence, exchange in enumerate(exchanges): + if exchange.request.get("suffix") is not None: + raise ValueError("Completions suffix is not supported") + for prompt_index, prompt in enumerate( + _completion_prompts(exchange.request.get("prompt")) + ): + if not isinstance(prompt, str): + continue + for choice in _completion_choices_for_prompt(exchange, prompt_index): + text = choice.text + if exchange.request.get("echo") is True: + if not text.startswith(prompt): + raise ValueError( + "Cannot locate echoed Completions prompt boundary" + ) + text = text[len(prompt) :] + prompt_source = CompletionsSource( + exchange=exchange, prompt_index=prompt_index + ) + output_source = CompletionsSource( + exchange=exchange, + prompt_index=prompt_index, + choice_index=choice.index, + ) + _extend_branches( + branches, + prompt=list(prompt), + prompt_sources=[prompt_source] * len(prompt), + outputs=[ + ( + choice.index, + list(text), + [output_source] * len(text), + ) + ], + context=(), + sequence=sequence, + start_time=exchange.start_time, + ) + histories.extend( + CompletionsStringHistory( + model=selected_model, + prompt="".join(branch.items), + prompt_sources=_string_source_spans(branch.sources), + sampled_spans=_sampled_spans(branch.sources), ) - else: - token_ids.extend(prompt[len(token_ids) :]) - start = len(token_ids) - token_ids.extend(completion) - sampled_spans.append((start, len(token_ids))) - return CompletionsHistory( - model=_model(exchanges[0]), - token_ids=token_ids, - sampled_spans=sampled_spans, + for branch in branches + ) + if not histories: + raise ValueError("Completions string history requires text prompts") + return histories + + +def completions_string_history( + trajectory: Trajectory, *, model: str | None +) -> CompletionsStringHistory: + return _one( + completions_string_histories(trajectory, model=model), + "Completions string history", ) +def _source_runs( + sources: Sequence[CompletionsSource | None], +) -> Iterable[tuple[int, int, CompletionsSource | None]]: + if not sources: + return + start = 0 + source = sources[0] + for index, item in enumerate(sources[1:], 1): + if item != source: + yield start, index, source + start, source = index, item + yield start, len(sources), source + + +def _token_source_spans( + sources: Sequence[CompletionsSource | None], +) -> list[CompletionsTokenSourceSpan]: + return [ + CompletionsTokenSourceSpan(start=start, end=end, source=source) + for start, end, source in _source_runs(sources) + ] + + +def _string_source_spans( + sources: Sequence[CompletionsSource | None], +) -> list[CompletionsStringSourceSpan]: + return [ + CompletionsStringSourceSpan(start=start, end=end, source=source) + for start, end, source in _source_runs(sources) + ] + + +def _sampled_spans( + sources: Sequence[CompletionsSource | None], +) -> list[tuple[int, int]]: + return [ + (start, end) + for start, end, source in _source_runs(sources) + if source is not None and source.choice_index is not None + ] + + def anthropic_as_chat_completions_history( history: AnthropicMessagesHistory, ) -> ChatCompletionsHistory: from ._tokenize import _anthropic_messages, _openai_tools - messages = _anthropic_messages( - {"system": history.system, "messages": history.messages} - ) + messages: list[dict[str, object]] = [] + sources: list[ChatCompletionsMessageSource | None] = [] + if history.system: + messages.extend(_anthropic_messages({"system": history.system, "messages": []})) + sources.append( + ChatCompletionsMessageSource(exchange=history.system_source) + if history.system_source is not None + else None + ) + for message, source in zip(history.messages, history.message_sources, strict=True): + converted_messages = _anthropic_messages({"messages": [message]}) + messages.extend(converted_messages) + converted = ( + ChatCompletionsMessageSource( + exchange=source.exchange, + request_index=source.request_index, + ) + if source is not None + else None + ) + sources.extend([converted] * len(converted_messages)) tools = _TOOLS.validate_python(_openai_tools(history.tools, dialect="messages")) return ChatCompletionsHistory( model=history.model, messages=_MESSAGES.validate_python(messages), + message_sources=sources, tools=tools, chat_template=history.chat_template, chat_template_kwargs=copy.deepcopy(history.chat_template_kwargs), @@ -315,65 +906,157 @@ def anthropic_as_chat_completions_history( def responses_as_chat_completions_history( history: ResponsesHistory, ) -> ChatCompletionsHistory: + if history.previous_response_id is not None or history.conversation is not None: + raise ValueError( + "Opaque Responses context cannot be represented as Chat Completions" + ) from ._tokenize import _openai_tools, _responses_messages messages = _responses_messages( {"instructions": history.instructions, "input": history.input} ) + sources: list[ChatCompletionsMessageSource | None] = [] + if history.instructions: + sources.append( + ChatCompletionsMessageSource(exchange=history.instructions_source) + if history.instructions_source is not None + else None + ) + + def converted( + source: ResponsesItemSource | None, + ) -> ChatCompletionsMessageSource | None: + return ( + ChatCompletionsMessageSource( + exchange=source.exchange, + request_index=source.request_index, + output_index=source.output_index, + ) + if source is not None + else None + ) + + pending_reasoning: list[ResponsesItemSource | None] = [] + tool_message_open = False + for item, source in zip(history.input, history.input_sources, strict=True): + kind = item.get("type") + if kind == "reasoning": + tool_message_open = False + pending_reasoning.append(source) + continue + if kind == "function_call": + if not tool_message_open: + first = next( + (item for item in pending_reasoning if item is not None), source + ) + sources.append(converted(first)) + pending_reasoning.clear() + tool_message_open = True + continue + tool_message_open = False + if pending_reasoning: + first = next( + (item for item in pending_reasoning if item is not None), source + ) + sources.append(converted(first)) + if not (kind in {None, "message"} and item.get("role") == "assistant"): + sources.append(converted(source)) + pending_reasoning.clear() + else: + sources.append(converted(source)) + if pending_reasoning: + sources.append( + converted( + next((item for item in pending_reasoning if item is not None), None) + ) + ) + if len(sources) != len(messages): + raise AssertionError("Responses conversion sources must parallel messages") tools = _TOOLS.validate_python(_openai_tools(history.tools, dialect="responses")) return ChatCompletionsHistory( model=history.model, messages=_MESSAGES.validate_python(messages), + message_sources=sources, tools=tools, chat_template=history.chat_template, chat_template_kwargs=copy.deepcopy(history.chat_template_kwargs), ) -def trajectory_history( +def trajectory_histories( trajectory: Trajectory, *, model: str | None -) -> TrajectoryHistory: +) -> list[TrajectoryHistory]: _require_unmixed(trajectory) if not trajectory.exchanges: - if trajectory.additional_histories: - raise ValueError("Trajectory contains multiple legacy histories") - if model is not None: - raise ValueError("Legacy trajectory histories do not identify a model") - return History( - messages_and_choices=copy.deepcopy(trajectory.messages_and_choices), - tools=copy.deepcopy(trajectory.tools), - ) + return [ + LegacyHistory( + messages_and_choices=copy.deepcopy(trajectory.messages_and_choices), + tools=copy.deepcopy(trajectory.tools), + ), + *copy.deepcopy(trajectory.additional_histories), + ] - candidates = [ - name - for name, exchanges in ( - ("chat_completions", trajectory.exchanges.chat_completions), - ("completions", trajectory.exchanges.completions), - ("responses", trajectory.exchanges.responses), - ("messages", trajectory.exchanges.messages), - ) - if any(model is None or exchange.model == model for exchange in exchanges) - ] + candidates: list[TrajectoryHistory] = [] + if any( + model is None or exchange.model == model + for exchange in trajectory.exchanges.chat_completions + ): + candidates.extend(chat_completions_histories(trajectory, model=model)) + if any( + model is None or exchange.model == model + for exchange in trajectory.exchanges.completions + ): + try: + candidates.extend(completions_token_histories(trajectory, model=model)) + except ValueError as error: + if "requires exact token IDs" not in str(error): + raise + candidates.extend(completions_string_histories(trajectory, model=model)) + if any( + model is None or exchange.model == model + for exchange in trajectory.exchanges.responses + ): + candidates.extend(responses_histories(trajectory, model=model)) + if any( + model is None or exchange.model == model + for exchange in trajectory.exchanges.messages + ): + candidates.extend(anthropic_messages_histories(trajectory, model=model)) if not candidates: suffix = f" for model {model!r}" if model is not None else "" raise ValueError(f"Trajectory contains no exchanges{suffix}") - if len(candidates) != 1: + protocols = {type(history) for history in candidates} + if len(protocols) != 1: raise ValueError( "Trajectory resolves to multiple protocol histories; use a protocol-specific method" ) - protocol = candidates[0] - if protocol == "chat_completions": - return chat_completions_history(trajectory, model=model) - if protocol == "completions": - return completions_history(trajectory, model=model) - if protocol == "responses": - return responses_history(trajectory, model=model) - return anthropic_messages_history(trajectory, model=model) + return candidates + + +def trajectory_history( + trajectory: Trajectory, *, model: str | None +) -> TrajectoryHistory: + histories = trajectory_histories(trajectory, model=model) + if ( + model is None + and len( + { + history.model + for history in histories + if not isinstance(history, LegacyHistory) + } + ) + > 1 + ): + raise ValueError( + "Trajectory history requires exactly one model; pass model= to select one" + ) + return _one(histories, "Trajectory history") def trajectory_messages(trajectory: Trajectory) -> Messages: if not trajectory.exchanges: - return History( + return LegacyHistory( messages_and_choices=trajectory.messages_and_choices, tools=trajectory.tools, ).messages() diff --git a/src/art/trajectories/_protocols.py b/src/art/trajectories/_protocols.py index db89be737..704df4bbb 100644 --- a/src/art/trajectories/_protocols.py +++ b/src/art/trajectories/_protocols.py @@ -221,6 +221,9 @@ def _messages_response(body: bytes, *, stream: bool) -> Message: complete = False token_ids: list[int] = [] logprobs: list[float] = [] + prompt_token_ids: list[int] = [] + block_token_ids: dict[int, list[int]] = {} + block_logprobs: dict[int, list[float]] = {} for event_name, payload in _sse_events(body): if not isinstance(payload, dict): continue @@ -232,7 +235,21 @@ def _messages_response(body: bytes, *, stream: bool) -> Message: current_snapshot=snapshot, output_format=NOT_GIVEN, ) - if event.type == "message_delta": + if event.type == "message_start": + message = payload.get("message") + values = ( + message.get("prompt_token_ids") if isinstance(message, dict) else None + ) + if isinstance(values, list) and all( + isinstance(value, int) for value in values + ): + prompt_token_ids = values + elif event.type == "message_delta": + values = payload.get("prompt_token_ids") + if isinstance(values, list) and all( + isinstance(value, int) for value in values + ): + prompt_token_ids = values event_token_ids = payload.get("token_ids") event_logprobs = payload.get("logprobs") if isinstance(event_token_ids, list) and all( @@ -243,14 +260,50 @@ def _messages_response(body: bytes, *, stream: bool) -> Message: isinstance(value, (int, float)) for value in event_logprobs ): logprobs = [float(value) for value in event_logprobs] + elif event.type in {"content_block_start", "content_block_delta"}: + index = payload.get("index") + event_token_ids = payload.get("token_ids") + event_logprobs = payload.get("logprobs") + if ( + isinstance(index, int) + and isinstance(event_token_ids, list) + and all(isinstance(value, int) for value in event_token_ids) + ): + block_token_ids.setdefault(index, []).extend(event_token_ids) + if ( + isinstance(index, int) + and isinstance(event_logprobs, list) + and all(isinstance(value, (int, float)) for value in event_logprobs) + ): + block_logprobs.setdefault(index, []).extend( + float(value) for value in event_logprobs + ) complete = complete or event.type == "message_stop" if snapshot is None or not complete: raise ValueError("Incomplete Messages stream") data = snapshot.model_dump(mode="python") + content = data.get("content") + if isinstance(content, list): + rebuilt_content: list[object] = [] + for index, raw_block in enumerate(content): + if not isinstance(raw_block, dict): + rebuilt_content.append(raw_block) + continue + block: dict[str, Any] = { + str(key): value for key, value in raw_block.items() + } + if values := block_token_ids.get(index): + block["token_ids"] = values + if values := block_logprobs.get(index): + block["logprobs"] = values + rebuilt_content.append(block) + data["content"] = rebuilt_content if token_ids: data["token_ids"] = token_ids if logprobs: data["logprobs"] = logprobs + if prompt_token_ids: + data["prompt_token_ids"] = prompt_token_ids return Message.model_validate(data) diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index a02199872..0d2f0584b 100644 --- a/src/art/trajectories/_tokenize.py +++ b/src/art/trajectories/_tokenize.py @@ -1,12 +1,14 @@ from __future__ import annotations -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from dataclasses import dataclass +from functools import lru_cache import math import re -from typing import Any, Protocol, cast +from typing import TYPE_CHECKING, Any, cast +import warnings -from anthropic.types import Message +from anthropic.types import Message, MessageParam, TextBlock from openai.types import Completion from openai.types.chat import ChatCompletion from openai.types.chat.chat_completion import Choice @@ -14,17 +16,35 @@ from pydantic import BaseModel from . import ( + AnthropicMessagesHistory, ChatCompletionsExchange, + ChatCompletionsHistory, CompletionsExchange, + CompletionsSource, + CompletionsStringHistory, + CompletionsTokenHistory, + History, + LegacyHistory, MessagesExchange, ResponsesExchange, + ResponsesHistory, TokenFlag, + TokenizedHistory, + TokenizedMultiHistoryTrajectory, TokenizedTrajectory, + TokenizedTrajectoryGroup, + Tokenizer, Trajectory, + TrajectoryGroup, + TrajectoryHistory, ) from ._protocols import Exchange +if TYPE_CHECKING: + from transformers import PreTrainedTokenizerBase + _TOKEN_ID = re.compile(r"token_id:(\d+)$") +_WARNED_PREFIX_RETOKENIZATION = False @dataclass @@ -35,27 +55,12 @@ class _TokenizerConfig: chat_template_kwargs: Mapping[str, object] | None = None -class _Tokenizer(Protocol): - def __call__(self, text: str, *, add_special_tokens: bool = False) -> object: ... - - def apply_chat_template( - self, - messages: list[dict[str, Any]], - *, - tools: object, - tokenize: bool, - add_generation_prompt: bool, - chat_template: str | None = None, - **kwargs: object, - ) -> object: ... - - -def _as_tokenizer(tokenizer: object) -> _Tokenizer: +def _as_tokenizer(tokenizer: object) -> Tokenizer: # Transformers' annotation permits only string-valued message dictionaries, # although its runtime API supports the structured content ART must tokenize. # Exact-token paths may only need decode(); fallback paths exercise these # capabilities directly and report the missing method at that point. - return cast(_Tokenizer, tokenizer) + return cast(Tokenizer, tokenizer) def _string_dict(value: object) -> dict[str, Any] | None: @@ -186,6 +191,8 @@ def _chat_choice_tokens( if token_ids is not None and pair_ids and token_ids != pair_ids: raise ValueError("Response token IDs disagree with choice logprobs") selected = token_ids if token_ids is not None else pair_ids or None + if selected is not None and values and len(logprobs) != len(selected): + raise ValueError("Chat Completions token IDs and logprobs differ in length") return ( prompt_ids, selected, @@ -201,9 +208,11 @@ def _chat_tokens( return _chat_choice_tokens(response.choices[0], _dump(response)) -def _completion_tokens( +def _completion_evidence( response: Completion, -) -> tuple[list[int] | None, list[int] | None, list[float]]: + *, + echo: bool = False, +) -> tuple[list[int] | None, list[int] | None, list[float], list[float]]: if len(response.choices) != 1: raise ValueError("Trajectory tokenization requires exactly one response choice") choice = response.choices[0] @@ -239,17 +248,50 @@ def _completion_tokens( for value in logprobs.get("token_logprobs") or [] ] if token_ids is not None and pair_ids and token_ids != pair_ids: - raise ValueError("Response token IDs disagree with completion logprobs") + if not ( + echo and prompt_ids is not None and pair_ids == [*prompt_ids, *token_ids] + ): + raise ValueError("Response token IDs disagree with completion logprobs") selected = token_ids if token_ids is not None else pair_ids or None - if selected is not None and len(pair_logprobs) != len(selected): - pair_logprobs = [math.nan] * len(selected) - return prompt_ids, selected, pair_logprobs + prompt_logprobs: list[float] = [] + completion_logprobs = pair_logprobs + if echo and prompt_ids is not None and selected is not None: + if selected[: len(prompt_ids)] == prompt_ids: + if pair_logprobs and len(pair_logprobs) != len(selected): + raise ValueError("Completions token IDs and logprobs differ in length") + prompt_logprobs = pair_logprobs[: len(prompt_ids)] + selected = selected[len(prompt_ids) :] + completion_logprobs = pair_logprobs[len(prompt_ids) :] + elif len(pair_logprobs) == len(prompt_ids) + len(selected): + prompt_logprobs = pair_logprobs[: len(prompt_ids)] + completion_logprobs = pair_logprobs[len(prompt_ids) :] + elif pair_logprobs and len(pair_logprobs) != len(selected): + raise ValueError("Completions token IDs and logprobs differ in length") + elif selected is not None and pair_logprobs and len(pair_logprobs) != len(selected): + raise ValueError("Completions token IDs and logprobs differ in length") + if selected is not None and not completion_logprobs: + completion_logprobs = [math.nan] * len(selected) + return prompt_ids, selected, prompt_logprobs, completion_logprobs + + +def _completion_tokens( + response: Completion, + *, + echo: bool = False, +) -> tuple[list[int] | None, list[int] | None, list[float]]: + prompt, completion, _, completion_logprobs = _completion_evidence( + response, echo=echo + ) + return prompt, completion, completion_logprobs def _responses_tokens( response: Response, -) -> tuple[None, list[int] | None, list[float]]: +) -> tuple[list[int] | None, list[int] | None, list[float]]: data = _dump(response) + prompt_ids = _exact_token_ids( + data.get("prompt_token_ids"), field="Responses prompt_token_ids" + ) if "raw_output_tokens" in data: token_ids, logprobs = _pairs( data["raw_output_tokens"], @@ -257,7 +299,7 @@ def _responses_tokens( field="Responses raw_output_tokens", ) if token_ids or not data.get("output"): - return None, token_ids, logprobs + return prompt_ids, token_ids, logprobs token_ids: list[int] = [] logprobs: list[float] = [] saw_rendered_output = False @@ -282,14 +324,17 @@ def _responses_tokens( token_ids.extend(pair_ids) logprobs.extend(pair_logprobs) if saw_rendered_output and complete: - return None, token_ids, logprobs - return None, None, [] + return prompt_ids, token_ids, logprobs + return prompt_ids, None, [] def _messages_tokens( response: Message, -) -> tuple[None, list[int] | None, list[float]]: +) -> tuple[list[int] | None, list[int] | None, list[float]]: data = _dump(response) + prompt_ids = _exact_token_ids( + data.get("prompt_token_ids"), field="Messages prompt_token_ids" + ) token_ids = _exact_token_ids(data.get("token_ids"), field="Messages token_ids") if token_ids == [] and data.get("content"): token_ids = None @@ -297,9 +342,11 @@ def _messages_tokens( float(value) if isinstance(value, (int, float)) else math.nan for value in data.get("logprobs") or [] ] - if token_ids is None or len(logprobs) != len(token_ids): + if token_ids is not None and logprobs and len(logprobs) != len(token_ids): + raise ValueError("Messages token IDs and logprobs differ in length") + if token_ids is None or not logprobs: logprobs = [math.nan] * len(token_ids or []) - return None, token_ids, logprobs + return prompt_ids, token_ids, logprobs def _exchange_list(trajectory: Trajectory, model: str | None) -> list[Exchange]: @@ -325,22 +372,59 @@ def _exchange_list(trajectory: Trajectory, model: str | None) -> list[Exchange]: ) -def _artifact_config(model: str) -> _TokenizerConfig: +def _artifact_name(model: str) -> str: + return model.removeprefix("wandb-artifact:///") + + +def _artifact_identity(model: str) -> str: + """Return an alias/version-independent checkpoint identity for base-model cache.""" + + path = _artifact_name(model) + name = path.rsplit("/", 1)[-1] + return path[: -len(name)] + name.split(":", 1)[0] + + +_ARTIFACT_BASE_MODELS: dict[str, str] = {} + + +def _artifact_base_model(identity: str) -> str: + if cached := _ARTIFACT_BASE_MODELS.get(identity): + return cached from wandb.apis.public import Api - artifact_path = model.removeprefix("wandb-artifact:///") + artifact_path = identity if ":" not in artifact_path.rsplit("/", 1)[-1]: artifact_path = f"{artifact_path}:latest" artifact = Api().artifact(artifact_path) metadata = artifact.metadata base_model = metadata.get("base_model") or metadata.get("wandb.base_model") if not isinstance(base_model, str): - raise ValueError(f"Checkpoint {model!r} does not identify its base model") + raise ValueError(f"Checkpoint {identity!r} does not identify its base model") + if len(_ARTIFACT_BASE_MODELS) >= 1024: + _ARTIFACT_BASE_MODELS.pop(next(iter(_ARTIFACT_BASE_MODELS))) + _ARTIFACT_BASE_MODELS[identity] = base_model + return base_model + + +def _artifact_config(model: str) -> _TokenizerConfig: + from wandb.apis.public import Api + + artifact_path = _artifact_name(model) + if ":" not in artifact_path.rsplit("/", 1)[-1]: + artifact_path = f"{artifact_path}:latest" + artifact = Api().artifact(artifact_path) + metadata = artifact.metadata + base_model = metadata.get("base_model") or metadata.get("wandb.base_model") + identity = _artifact_identity(model) + if isinstance(base_model, str): + if len(_ARTIFACT_BASE_MODELS) >= 1024 and identity not in _ARTIFACT_BASE_MODELS: + _ARTIFACT_BASE_MODELS.pop(next(iter(_ARTIFACT_BASE_MODELS))) + _ARTIFACT_BASE_MODELS[identity] = base_model renderer = metadata.get("renderer") renderer = renderer if isinstance(renderer, dict) else {} kwargs = renderer.get("chat_template_kwargs") return _TokenizerConfig( - base_model=base_model, + base_model=_artifact_base_model(identity), revision=( renderer.get("tokenizer_revision") if isinstance(renderer.get("tokenizer_revision"), str) @@ -368,7 +452,8 @@ def _tokenizer_config(model: str, base_model: str | None) -> _TokenizerConfig: return _TokenizerConfig(model) -def _load_tokenizer(config: _TokenizerConfig) -> _Tokenizer: +@lru_cache(maxsize=8) +def _cached_tokenizer(base_model: str, revision: str | None) -> Tokenizer: try: from transformers import AutoTokenizer except ImportError as exc: @@ -376,18 +461,25 @@ def _load_tokenizer(config: _TokenizerConfig) -> _Tokenizer: "Tokenizer fallback requires ART's backend or tinker dependencies" ) from exc try: - return _as_tokenizer( - AutoTokenizer.from_pretrained( - config.base_model, - revision=config.revision, - ) + tokenizer = AutoTokenizer.from_pretrained( + base_model, + revision=revision, ) + if base_model.startswith("deepseek-ai/DeepSeek-V4-"): + from ..megatron.dsv4.tokenizer import get_dsv4_tokenizer + + tokenizer = get_dsv4_tokenizer(cast("PreTrainedTokenizerBase", tokenizer)) + return _as_tokenizer(tokenizer) except Exception as exc: raise ValueError( - f"Could not load tokenizer for {config.base_model!r}; pass base_model explicitly" + f"Could not load tokenizer for {base_model!r}; pass base_model explicitly" ) from exc +def _load_tokenizer(config: _TokenizerConfig) -> Tokenizer: + return _cached_tokenizer(config.base_model, config.revision) + + def _ids(value: object) -> list[int]: if (input_ids := getattr(value, "input_ids", None)) is not None: value = input_ids @@ -745,7 +837,7 @@ def _response_message( def _template_ids( - tokenizer: _Tokenizer, + tokenizer: Tokenizer, exchange: Exchange, *, completed: bool, @@ -803,7 +895,9 @@ def _exchange_tokens( if isinstance(exchange, ChatCompletionsExchange): return _chat_tokens(exchange.response) if isinstance(exchange, CompletionsExchange): - return _completion_tokens(exchange.response) + return _completion_tokens( + exchange.response, echo=exchange.request.get("echo") is True + ) if isinstance(exchange, ResponsesExchange): return _responses_tokens(exchange.response) if isinstance(exchange, MessagesExchange): @@ -849,8 +943,81 @@ def _visible_logprobs(exchange: Exchange) -> list[tuple[str, float]]: return values +def _sampled_text(exchange: Exchange) -> str | None: + visible = "".join(text for text, _ in _visible_logprobs(exchange)) + if visible: + return visible + if isinstance(exchange, ChatCompletionsExchange): + content = exchange.response.choices[0].message.content + return content if isinstance(content, str) else None + if isinstance(exchange, CompletionsExchange): + return exchange.response.choices[0].text + if isinstance(exchange, MessagesExchange): + parts = [ + block.text + for block in exchange.response.content + if isinstance(block, TextBlock) + ] + return "".join(parts) if parts else None + if isinstance(exchange, ResponsesExchange): + parts: list[str] = [] + for output in _dump(exchange.response).get("output") or []: + data = _dump(output) + if data.get("type") == "message": + parts.append(_responses_output_text(data.get("content"))) + return "".join(parts) if parts else None + return None + + +def _replace_unique( + values: list[int], old: list[int], new: list[int] +) -> list[int] | None: + if not old: + return None + starts = [ + index + for index in range(len(values) - len(old) + 1) + if values[index : index + len(old)] == old + ] + if len(starts) != 1: + return None + start = starts[0] + return [*values[:start], *new, *values[start + len(old) :]] + + +def _preserve_sampled_prefix( + prompt: list[int], + canonical_prefix: list[int], + sampled_outputs: list[tuple[str, list[int]]], + tokenizer: Tokenizer, +) -> list[int] | None: + repaired = prompt + for text, exact_ids in sampled_outputs: + rendered_ids = _ids(tokenizer(text, add_special_tokens=False)) + if rendered_ids == exact_ids: + continue + replaced = _replace_unique(repaired, rendered_ids, exact_ids) + if replaced is None: + return None + repaired = replaced + return repaired if repaired[: len(canonical_prefix)] == canonical_prefix else None + + +def _warn_prefix_retokenization() -> None: + global _WARNED_PREFIX_RETOKENIZATION + if _WARNED_PREFIX_RETOKENIZATION: + return + _WARNED_PREFIX_RETOKENIZATION = True + warnings.warn( + "Inference prompt token IDs retokenized an earlier sampled response; ART " + "preserved the original sampled token IDs and logprobs. Prefer a service " + "with prefix token-ID preservation, such as Caladan.", + stacklevel=3, + ) + + def _align_visible_logprobs( - tokenizer: _Tokenizer | None, completion: list[int], exchange: Exchange + tokenizer: Tokenizer | None, completion: list[int], exchange: Exchange ) -> list[float] | None: values = _visible_logprobs(exchange) if not values or tokenizer is None: @@ -895,18 +1062,14 @@ def _align_visible_logprobs( def _legacy_tokenize( - trajectory: Trajectory, - base_model: str | None, + history: LegacyHistory, *, - chat_template: str | None, - chat_template_kwargs: Mapping[str, object] | None, -) -> TokenizedTrajectory: - if trajectory.additional_histories: - raise ValueError("Tokenization requires one history") + model: str, +) -> TokenizedHistory: token_ids: list[int] = [] logprobs: list[float] = [] flags: list[TokenFlag] = [] - for item in trajectory.messages_and_choices: + for item in history.messages_and_choices: if not isinstance(item, Choice): continue prompt, completion, completion_logprobs = _chat_choice_tokens(item, {}) @@ -932,23 +1095,23 @@ def _legacy_tokenize( flags.extend([TokenFlag.EXACT | TokenFlag.SAMPLED] * len(completion)) if not token_ids: raise ValueError("Trajectory contains no trainable choices") - return TokenizedTrajectory( + return TokenizedHistory( + model=model, token_ids=token_ids, logprobs=logprobs, flags=flags, - underlying=trajectory, ) -def tokenize_one( +def _tokenize_exchange_trajectory( trajectory: Trajectory, base_model: str | None, *, model: str | None, chat_template: str | None, chat_template_kwargs: Mapping[str, object] | None, - tokenizer_instance: _Tokenizer | None = None, -) -> TokenizedTrajectory: + tokenizer_instance: Tokenizer | None = None, +) -> TokenizedHistory: if trajectory.exchanges and ( trajectory.messages_and_choices or trajectory.tools is not None @@ -958,11 +1121,16 @@ def tokenize_one( "A trajectory cannot contain both exchanges and legacy histories" ) if not trajectory.exchanges: + if trajectory.additional_histories: + raise ValueError("Tokenization requires one history") + if model is None: + raise ValueError("Legacy trajectory tokenization requires model=") return _legacy_tokenize( - trajectory, - base_model, - chat_template=chat_template, - chat_template_kwargs=chat_template_kwargs, + LegacyHistory( + messages_and_choices=trajectory.messages_and_choices, + tools=trajectory.tools, + ), + model=model, ) exchanges = _exchange_list(trajectory, model) selected_model = exchanges[0].model @@ -986,6 +1154,8 @@ def tokenize_one( response_histories: dict[ str, tuple[list[dict[str, Any]] | None, ResponsesExchange] ] = {} + sampled_outputs: list[tuple[str, list[int]]] = [] + previous_render_state: tuple[Exchange, list[dict[str, Any]] | None] | None = None def fallback_config() -> _TokenizerConfig: nonlocal config @@ -1016,6 +1186,10 @@ def fallback_config() -> _TokenizerConfig: messages_override: list[dict[str, Any]] | None = None if isinstance(exchange, ResponsesExchange): request = exchange.request + if request.get("conversation") is not None and prompt is None: + raise ValueError( + "Responses conversation history requires exact prompt tokens" + ) try: messages_override = _responses_messages(request) except ValueError: @@ -1023,21 +1197,25 @@ def fallback_config() -> _TokenizerConfig: raise previous = request.get("previous_response_id") if previous is not None: - if not isinstance(previous, str) or previous not in response_histories: + if not isinstance(previous, str): + raise ValueError("Responses previous_response_id must be text") + if previous not in response_histories and prompt is None: raise ValueError( - "Responses exchange refers to a previous response outside this trajectory" + "Responses exchange refers to a previous response outside this " + "trajectory without exact prompt tokens" ) - previous_messages, previous_exchange = response_histories[previous] - if prompt is None: - if previous_messages is None or messages_override is None: - raise ValueError( - "Responses history cannot be rendered without exact prompt tokens" - ) - messages_override = [ - *previous_messages, - _response_message(previous_exchange), - *messages_override, - ] + if previous in response_histories: + previous_messages, previous_exchange = response_histories[previous] + if prompt is None: + if previous_messages is None or messages_override is None: + raise ValueError( + "Responses history cannot be rendered without exact prompt tokens" + ) + messages_override = [ + *previous_messages, + _response_message(previous_exchange), + *messages_override, + ] response_histories[exchange.response.id] = (messages_override, exchange) if prompt is None: resolved_config = fallback_config() @@ -1080,9 +1258,57 @@ def fallback_config() -> _TokenizerConfig: [TokenFlag.EXACT if prompt_is_exact else TokenFlag(0)] * len(prompt) ) elif len(prompt) < len(token_ids) or prompt[: len(token_ids)] != token_ids: - raise ValueError( - "Exchanges do not resolve to one append-only token history" + resolved_config = fallback_config() + if tokenizer is None: + tokenizer = _load_tokenizer(resolved_config) + repaired = _preserve_sampled_prefix( + prompt, + token_ids, + sampled_outputs, + tokenizer, ) + if repaired is None: + current_render = _template_ids( + tokenizer, + exchange, + completed=False, + config=resolved_config, + chat_template=chat_template, + chat_template_kwargs=chat_template_kwargs, + messages_override=messages_override, + ) + if previous_render_state is None: + repaired = [*token_ids, *current_render] + else: + previous_exchange, previous_messages = previous_render_state + previous_render = _template_ids( + tokenizer, + previous_exchange, + completed=True, + config=resolved_config, + chat_template=chat_template, + chat_template_kwargs=chat_template_kwargs, + messages_override=previous_messages, + ) + previous_canonical = _preserve_sampled_prefix( + previous_render, + token_ids, + sampled_outputs, + tokenizer, + ) + suffix = ( + current_render[len(previous_render) :] + if current_render[: len(previous_render)] == previous_render + else current_render + ) + repaired = [*(previous_canonical or token_ids), *suffix] + prompt = repaired + prompt_is_exact = False + _warn_prefix_retokenization() + suffix = prompt[len(token_ids) :] + token_ids.extend(suffix) + logprobs.extend([math.nan] * len(suffix)) + flags.extend([TokenFlag(0)] * len(suffix)) else: suffix = prompt[len(token_ids) :] token_ids.extend(suffix) @@ -1102,10 +1328,975 @@ def fallback_config() -> _TokenizerConfig: if completion_is_exact: completion_flag |= TokenFlag.EXACT flags.extend([completion_flag] * len(completion)) + if completion_is_exact and (text := _sampled_text(exchange)) is not None: + sampled_outputs.append((text, list(completion))) + previous_render_state = (exchange, messages_override) - return TokenizedTrajectory( + return TokenizedHistory( + model=selected_model, + token_ids=token_ids, + logprobs=logprobs, + flags=flags, + ) + + +def _unique_exchanges(history: History) -> list[Exchange]: + sources: Sequence[object] + if isinstance(history, ChatCompletionsHistory): + if len(history.messages) != len(history.message_sources): + raise ValueError("messages and message_sources differ in length") + sources = history.message_sources + elif isinstance(history, AnthropicMessagesHistory): + if len(history.messages) != len(history.message_sources): + raise ValueError("messages and message_sources differ in length") + sources = history.message_sources + elif isinstance(history, ResponsesHistory): + if len(history.input) != len(history.input_sources): + raise ValueError("input and input_sources differ in length") + sources = history.input_sources + else: + raise TypeError(f"Unsupported history type: {type(history).__name__}") + + exchanges: list[Exchange] = [] + seen: set[int] = set() + for source in sources: + exchange = getattr(source, "exchange", None) + if not isinstance( + exchange, + ( + ChatCompletionsExchange, + CompletionsExchange, + ResponsesExchange, + MessagesExchange, + ), + ): + continue + identity = id(exchange) + if identity in seen: + continue + seen.add(identity) + if isinstance(exchange, ChatCompletionsExchange): + choice_indexes = { + getattr(item, "choice_index", None) + for item in sources + if getattr(item, "exchange", None) is exchange + and getattr(item, "choice_index", None) is not None + } + if choice_indexes: + exchange = exchange.model_copy( + update={ + "response": exchange.response.model_copy( + update={ + "choices": [ + choice + for choice in exchange.response.choices + if choice.index in choice_indexes + ] + } + ) + } + ) + exchanges.append(exchange) + return sorted(exchanges, key=lambda item: (item.start_time, item.end_time)) + + +def _validate_history_sources(history: History) -> None: + if isinstance(history, ChatCompletionsHistory): + if len(history.messages) != len(history.message_sources): + raise ValueError("messages and message_sources differ in length") + for message, source in zip( + history.messages, history.message_sources, strict=True + ): + if source is None: + continue + exchange = source.exchange + expected: list[dict[str, Any]] = [] + if isinstance(exchange, ChatCompletionsExchange): + if source.choice_index is not None: + choice = next( + ( + item + for item in exchange.response.choices + if item.index == source.choice_index + ), + None, + ) + if choice is not None: + expected.append( + choice.message.model_dump(mode="python", exclude_none=True) + ) + elif source.request_index is not None: + request_messages = exchange.request.get("messages", []) + if source.request_index < len(request_messages): + expected.append(dict(request_messages[source.request_index])) + elif isinstance(exchange, MessagesExchange): + expected.extend(_anthropic_messages(_dump(exchange.request))) + expected.append(_response_message(exchange)) + visible_blocks = [ + block.model_dump(mode="python", exclude_none=True) + for block in exchange.response.content + if getattr(block, "type", None) + not in {"thinking", "redacted_thinking"} + ] + expected.extend( + _anthropic_messages( + {"messages": [{"role": "assistant", "content": visible_blocks}]} + ) + ) + elif isinstance(exchange, ResponsesExchange): + expected.extend(_responses_messages(_dump(exchange.request))) + expected.append(_response_message(exchange)) + if not any(dict(message) == candidate for candidate in expected): + raise ValueError( + "Chat Completions history no longer matches its source exchange" + ) + return + if isinstance(history, AnthropicMessagesHistory): + from ._history import _anthropic_message_key + + if len(history.messages) != len(history.message_sources): + raise ValueError("messages and message_sources differ in length") + for message, source in zip( + history.messages, history.message_sources, strict=True + ): + if source is None: + continue + if source.request_index is None: + expected = { + "role": "assistant", + "content": [ + block.model_dump(mode="json", exclude_none=True) + for block in source.exchange.response.content + ], + } + else: + request_messages = source.exchange.request.get("messages", []) + expected = ( + request_messages[source.request_index] + if source.request_index < len(request_messages) + else None + ) + matches = expected is not None and message == expected + if not matches and source.request_index is None and expected is not None: + matches = _anthropic_message_key(message) == _anthropic_message_key( + cast(MessageParam, expected), visible_only=True + ) + if not matches: + raise ValueError( + "Anthropic Messages history no longer matches its source exchange" + ) + return + if isinstance(history, ResponsesHistory): + from ._history import _responses_input + + if len(history.input) != len(history.input_sources): + raise ValueError("input and input_sources differ in length") + + for item, source in zip(history.input, history.input_sources, strict=True): + if source is None: + continue + if source.output_index is not None: + output = source.exchange.response.output + expected = ( + output[source.output_index].model_dump( + mode="json", exclude_none=True + ) + if source.output_index < len(output) + else None + ) + elif source.request_index is not None: + request_input = _responses_input(source.exchange.request.get("input")) + expected = ( + request_input[source.request_index] + if source.request_index < len(request_input) + else None + ) + else: + expected = None + if expected is None or item != expected: + raise ValueError( + "Responses history no longer matches its source exchange" + ) + + +def _trajectory_from_history(history: History) -> Trajectory: + _validate_history_sources(history) + exchanges = _unique_exchanges(history) + if not exchanges: + raise ValueError( + "History has no source exchanges; local history-only rendering is not yet possible" + ) + from . import TrajectoryExchanges + + return Trajectory( + exchanges=TrajectoryExchanges( + chat_completions=[ + item for item in exchanges if isinstance(item, ChatCompletionsExchange) + ], + completions=[ + item for item in exchanges if isinstance(item, CompletionsExchange) + ], + responses=[ + item for item in exchanges if isinstance(item, ResponsesExchange) + ], + messages=[item for item in exchanges if isinstance(item, MessagesExchange)], + ) + ) + + +def _last_source_exchange(sources: Sequence[object]) -> Exchange | None: + for source in reversed(sources): + exchange = getattr(source, "exchange", None) + if isinstance( + exchange, + ( + ChatCompletionsExchange, + CompletionsExchange, + ResponsesExchange, + MessagesExchange, + ), + ): + return exchange + return None + + +def _history_needs_render(history: History) -> bool: + if isinstance(history, ChatCompletionsHistory): + if any(source is None for source in history.message_sources): + return True + exchange = _last_source_exchange(history.message_sources) + if not isinstance(exchange, ChatCompletionsExchange): + return False + return ( + history.tools != exchange.request.get("tools") + or history.chat_template != exchange.request.get("chat_template") + or history.chat_template_kwargs + != exchange.request.get("chat_template_kwargs") + ) + if isinstance(history, AnthropicMessagesHistory): + if any(source is None for source in history.message_sources): + return True + for message, source in zip( + history.messages, history.message_sources, strict=True + ): + if source is None or source.request_index is not None: + continue + expected = { + "role": "assistant", + "content": [ + block.model_dump(mode="json", exclude_none=True) + for block in source.exchange.response.content + ], + } + if message != expected: + return True + exchange = _last_source_exchange(history.message_sources) + if not isinstance(exchange, MessagesExchange): + return False + return ( + history.system != exchange.request.get("system") + or history.tools != exchange.request.get("tools") + or history.chat_template != exchange.request.get("chat_template") + or history.chat_template_kwargs + != exchange.request.get("chat_template_kwargs") + ) + if isinstance(history, ResponsesHistory): + if any(source is None for source in history.input_sources): + return True + exchange = _last_source_exchange(history.input_sources) + if not isinstance(exchange, ResponsesExchange): + return False + return ( + history.instructions != exchange.request.get("instructions") + or history.tools != exchange.request.get("tools") + or history.chat_template != exchange.request.get("chat_template") + or history.chat_template_kwargs + != exchange.request.get("chat_template_kwargs") + ) + return False + + +def _chat_source_full_tokens( + source: object, +) -> tuple[list[int] | None, list[float]]: + exchange = getattr(source, "exchange", None) + if isinstance(exchange, ChatCompletionsExchange): + choice_index = getattr(source, "choice_index", None) + if choice_index is None: + return None, [] + choice = next( + item for item in exchange.response.choices if item.index == choice_index + ) + _, tokens, logprobs = _chat_choice_tokens(choice, _dump(exchange.response)) + return tokens, logprobs + if isinstance(exchange, (MessagesExchange, ResponsesExchange)): + _, tokens, logprobs = _exchange_tokens(exchange) + return tokens, logprobs + return None, [] + + +def _chat_source_tokens( + source: object, + text: str, + *, + part: str, +) -> tuple[list[int] | None, list[float]]: + exchange = getattr(source, "exchange", None) + if isinstance(exchange, ChatCompletionsExchange): + tokens, logprobs = _chat_source_full_tokens(source) + if tokens is None: + return None, [] + sampled_text = _sampled_text(exchange) + message = _dump( + next( + item + for item in exchange.response.choices + if item.index == getattr(source, "choice_index", None) + ).message + ) + if sampled_text == text or ( + part == "content" + and not any(message.get(key) for key in ("reasoning", "tool_calls")) + ): + return tokens, logprobs + return None, [] + if ( + isinstance(exchange, MessagesExchange) + and getattr(source, "request_index", None) is None + ): + block_type = "thinking" if part == "reasoning" else "text" + blocks = [ + block + for block in exchange.response.content + if getattr(block, "type", None) == block_type + ] + block_text = "".join( + str(getattr(block, "thinking" if part == "reasoning" else "text", "")) + for block in blocks + ) + if block_text == text: + token_ids: list[int] = [] + logprobs: list[float] = [] + for block in blocks: + extra = block.model_extra or {} + block_ids = _exact_token_ids( + extra.get("token_ids"), field="Messages content token_ids" + ) + block_logprobs = extra.get("logprobs") + if block_ids is None or not isinstance(block_logprobs, list): + break + if len(block_ids) != len(block_logprobs): + raise ValueError( + "Messages content token IDs and logprobs differ in length" + ) + token_ids.extend(block_ids) + logprobs.extend(float(value) for value in block_logprobs) + else: + return token_ids, logprobs + if part != "content" or any( + getattr(block, "type", None) in {"thinking", "redacted_thinking"} + for block in exchange.response.content + ): + return None, [] + _, tokens, logprobs = _messages_tokens(exchange.response) + return tokens, logprobs + if ( + isinstance(exchange, ResponsesExchange) + and getattr(source, "request_index", None) is None + ): + output_index = getattr(source, "output_index", None) + if isinstance(output_index, int) and output_index < len( + exchange.response.output + ): + item = _dump(exchange.response.output[output_index]) + if item.get("type") == "message" and part == "content": + item_text = _responses_output_text(item.get("content")) + if item_text == text: + token_ids: list[int] = [] + logprobs: list[float] = [] + for content in item.get("content") or []: + pairs, pair_logprobs = _pairs( + _dump(content).get("logprobs"), + field="Responses content logprobs", + ) + if not pairs: + break + token_ids.extend(pairs) + logprobs.extend(pair_logprobs) + else: + return token_ids, logprobs + if any( + getattr(item, "type", None) == "reasoning" + for item in exchange.response.output + ): + return None, [] + _, tokens, logprobs = _responses_tokens(exchange.response) + return tokens, logprobs + return None, [] + + +def _tokenize_chat_view( + history: ChatCompletionsHistory, + *, + base_model: str | None, + tokenizer: Tokenizer | None, + chat_template: str | None, + chat_template_kwargs: Mapping[str, object] | None, +) -> TokenizedHistory: + _validate_history_sources(history) + config = ( + _TokenizerConfig(base_model or history.model or "") + if tokenizer is not None + or (base_model is not None and chat_template is not None) + else _tokenizer_config(history.model or "", base_model) + ) + if tokenizer is None: + if not history.model and base_model is None: + raise ValueError("History tokenization requires a model or base_model") + tokenizer = _load_tokenizer(config) + messages = [dict(message) for message in history.messages] + kwargs = { + **(config.chat_template_kwargs or {}), + **(history.chat_template_kwargs or {}), + **(chat_template_kwargs or {}), + } + last_exchange = _last_source_exchange(history.message_sources) + if isinstance(last_exchange, MessagesExchange) and isinstance( + thinking := last_exchange.request.get("thinking"), dict + ): + kwargs.setdefault("enable_thinking", thinking.get("type") == "enabled") + if budget := thinking.get("budget_tokens"): + kwargs.setdefault("thinking_budget", budget) + template = chat_template or history.chat_template or config.chat_template + ends_with_assistant = bool(messages) and messages[-1].get("role") == "assistant" + rendered = _ids( + tokenizer.apply_chat_template( + messages, + tools=history.tools, + tokenize=True, + add_generation_prompt=not ends_with_assistant, + **({"chat_template": template} if template is not None else {}), + **kwargs, + ) + ) + replacements: list[tuple[int, int, list[int], list[float], bool]] = [] + for index, (message, source) in enumerate( + zip(history.messages, history.message_sources, strict=True) + ): + if message.get("role") != "assistant" or source is None: + continue + full_exact, full_logprobs = _chat_source_full_tokens(source) + content = _content_text(message.get("content")) + reasoning = message.get("reasoning") + if ( + full_exact + and content + and not reasoning + and not message.get("tool_calls") + and _ids(tokenizer(content, add_special_tokens=False)) == full_exact + ): + exact_starts = [ + start + for start in range(len(rendered) - len(full_exact) + 1) + if rendered[start : start + len(full_exact)] == full_exact + ] + if len(exact_starts) == 1: + start = exact_starts[0] + replacements.append( + ( + start, + start + len(full_exact), + full_exact, + full_logprobs + if len(full_logprobs) == len(full_exact) + else [math.nan] * len(full_exact), + True, + ) + ) + continue + stable_length = min(index + 2, len(history.messages)) + completed_prefix = ( + rendered + if stable_length == len(history.messages) + else _ids( + tokenizer.apply_chat_template( + [dict(item) for item in history.messages[:stable_length]], + tools=history.tools, + tokenize=True, + add_generation_prompt=False, + **({"chat_template": template} if template is not None else {}), + **kwargs, + ) + ) + ) + if rendered[: len(completed_prefix)] != completed_prefix: + raise ValueError( + "Chat template does not preserve completed message prefixes" + ) + message_end = len(completed_prefix) + if stable_length > index + 1: + next_text = _content_text(history.messages[index + 1].get("content")) + next_ids = _ids(tokenizer(next_text, add_special_tokens=False)) + if next_ids: + next_starts = [ + start + for start in range(len(completed_prefix) - len(next_ids) + 1) + if completed_prefix[start : start + len(next_ids)] == next_ids + ] + if next_starts: + message_end = next_starts[-1] + if full_exact: + exact_starts = [ + start + for start in range(message_end - len(full_exact) + 1) + if completed_prefix[start : start + len(full_exact)] == full_exact + ] + if exact_starts: + start = exact_starts[-1] + replacements.append( + ( + start, + start + len(full_exact), + full_exact, + full_logprobs + if len(full_logprobs) == len(full_exact) + else [math.nan] * len(full_exact), + True, + ) + ) + continue + + parts: list[tuple[str, str]] = [] + if isinstance(reasoning, str) and reasoning: + parts.append(("reasoning", reasoning)) + if content: + parts.append(("content", content)) + if not parts and isinstance( + exchange := getattr(source, "exchange", None), + (ChatCompletionsExchange, ResponsesExchange, MessagesExchange), + ): + if sampled_text := _sampled_text(exchange): + parts.append(("raw", sampled_text)) + + part_replacements: list[tuple[int, int, list[int], list[float], bool]] = [] + search_end = message_end + for part, text in reversed(parts): + local = _ids(tokenizer(text, add_special_tokens=False)) + if not local: + continue + starts = [ + start + for start in range(search_end - len(local) + 1) + if completed_prefix[start : start + len(local)] == local + ] + if not starts: + raise ValueError( + "Could not locate a sourced assistant message in the rendered history" + ) + start = starts[-1] + exact, logprobs = _chat_source_tokens(source, text, part=part) + replacement = exact if exact is not None else local + if exact is None and not logprobs: + exchange = getattr(source, "exchange", None) + if isinstance( + exchange, + (ChatCompletionsExchange, ResponsesExchange, MessagesExchange), + ): + logprobs = _align_visible_logprobs(tokenizer, local, exchange) or [] + part_replacements.append( + ( + start, + start + len(local), + replacement, + logprobs + if len(logprobs) == len(replacement) + else [math.nan] * len(replacement), + exact is not None, + ) + ) + search_end = start + replacements.extend(reversed(part_replacements)) + + token_ids: list[int] = [] + logprobs: list[float] = [] + flags: list[TokenFlag] = [] + cursor = 0 + for start, end, replacement, replacement_logprobs, exact in sorted(replacements): + if start < cursor: + raise ValueError("Rendered assistant source spans overlap") + token_ids.extend(rendered[cursor:start]) + logprobs.extend([math.nan] * (start - cursor)) + flags.extend([TokenFlag(0)] * (start - cursor)) + token_ids.extend(replacement) + logprobs.extend(replacement_logprobs) + flag = TokenFlag.SAMPLED | (TokenFlag.EXACT if exact else TokenFlag(0)) + flags.extend([flag] * len(replacement)) + cursor = end + token_ids.extend(rendered[cursor:]) + logprobs.extend([math.nan] * (len(rendered) - cursor)) + flags.extend([TokenFlag(0)] * (len(rendered) - cursor)) + if history.model is None: + raise ValueError("History tokenization requires a model") + return TokenizedHistory( + model=history.model, token_ids=token_ids, logprobs=logprobs, flags=flags, - underlying=trajectory, + ) + + +def _tokenize_completions_token_history( + history: CompletionsTokenHistory, +) -> TokenizedHistory: + if any( + span.start < 0 or span.end <= span.start or span.end > len(history.prompt) + for span in history.prompt_sources + ): + raise ValueError("Completions token source spans are out of bounds") + if not _spans_are_exhaustive(len(history.prompt), history.prompt_sources): + raise ValueError( + "Completions token source spans must exhaustively cover prompt" + ) + + flags = [TokenFlag(0)] * len(history.prompt) + logprobs = [math.nan] * len(history.prompt) + for start, end in history.sampled_spans: + if start < 0 or end <= start or end > len(history.prompt): + raise ValueError("Completions sampled spans are out of bounds") + flags[start:end] = [TokenFlag.SAMPLED] * (end - start) + for span in history.prompt_sources: + if span.source is None: + continue + flags[span.start : span.end] = [ + flag | TokenFlag.EXACT for flag in flags[span.start : span.end] + ] + prompt, completion, prompt_logprobs, completion_logprobs = ( + _completion_source_evidence(span.source) + ) + selected = prompt if span.source.choice_index is None else completion + selected_logprobs = ( + prompt_logprobs if span.source.choice_index is None else completion_logprobs + ) + if ( + selected == history.prompt[span.start : span.end] + and len(selected_logprobs) == span.end - span.start + ): + logprobs[span.start : span.end] = selected_logprobs + return TokenizedHistory( + model=history.model, + token_ids=list(history.prompt), + logprobs=logprobs, + flags=flags, + ) + + +def _tokenize_completions_string_history( + history: CompletionsStringHistory, + *, + base_model: str | None, + tokenizer: Tokenizer | None, +) -> TokenizedHistory: + if not _spans_are_exhaustive(len(history.prompt), history.prompt_sources): + raise ValueError( + "Completions string source spans must exhaustively cover prompt" + ) + sampled = [False] * len(history.prompt) + for start, end in history.sampled_spans: + if start < 0 or end <= start or end > len(history.prompt): + raise ValueError("Completions sampled spans are out of bounds") + sampled[start:end] = [True] * (end - start) + + config: _TokenizerConfig | None = None + + def resolved_tokenizer() -> Tokenizer: + nonlocal config, tokenizer + if tokenizer is None: + config = config or _tokenizer_config(history.model, base_model) + tokenizer = _load_tokenizer(config) + return tokenizer + + token_ids: list[int] = [] + logprobs: list[float] = [] + flags: list[TokenFlag] = [] + for span in history.prompt_sources: + text = history.prompt[span.start : span.end] + source = span.source + exact: list[int] | None = None + source_logprobs: list[float] = [] + is_sampled = any(sampled[span.start : span.end]) + if source is not None: + prompt, completion, prompt_logprobs, completion_logprobs = ( + _completion_source_evidence(source) + ) + if source.choice_index is None: + exact = prompt + source_logprobs = prompt_logprobs + else: + choice = next( + item + for item in source.exchange.response.choices + if item.index == source.choice_index + ) + expected = choice.text + request_prompt = source.exchange.request.get("prompt") + if source.exchange.request.get("echo") is True and isinstance( + request_prompt, str + ): + expected = expected.removeprefix(request_prompt) + if text != expected: + raise ValueError( + "Completions history text no longer matches its source exchange" + ) + exact = completion + source_logprobs = completion_logprobs + ids = ( + exact + if exact is not None + else _ids(resolved_tokenizer()(text, add_special_tokens=False)) + ) + token_ids.extend(ids) + if exact is not None and len(source_logprobs) == len(ids): + logprobs.extend(source_logprobs) + else: + logprobs.extend([math.nan] * len(ids)) + flag = TokenFlag.SAMPLED if is_sampled else TokenFlag(0) + if exact is not None: + flag |= TokenFlag.EXACT + flags.extend([flag] * len(ids)) + return TokenizedHistory( + model=history.model, + token_ids=token_ids, + logprobs=logprobs, + flags=flags, + ) + + +def _completion_source_evidence( + source: CompletionsSource, +) -> tuple[list[int] | None, list[int] | None, list[float], list[float]]: + choices = source.exchange.response.choices + selected = ( + next(choice for choice in choices if choice.index == source.choice_index) + if source.choice_index is not None + else choices[0] + ) + return _completion_evidence( + source.exchange.response.model_copy(update={"choices": [selected]}), + echo=source.exchange.request.get("echo") is True, + ) + + +def _spans_are_exhaustive(length: int, spans: Sequence[object]) -> bool: + cursor = 0 + for span in spans: + start = getattr(span, "start", None) + end = getattr(span, "end", None) + if ( + not isinstance(start, int) + or not isinstance(end, int) + or start != cursor + or end <= start + or end > length + ): + return False + cursor = end + return cursor == length + + +def tokenize_history( + history: History | LegacyHistory, + *, + model: str | None, + base_model: str | None, + tokenizer: Tokenizer | None, + chat_template: str | None, + chat_template_kwargs: Mapping[str, object] | None, +) -> TokenizedHistory: + if isinstance(history, LegacyHistory): + if model is None: + raise ValueError("Legacy history tokenization requires model=") + return _legacy_tokenize(history, model=model) + if model is None: + raise ValueError("History tokenization requires a model") + if isinstance(history, CompletionsTokenHistory): + return _tokenize_completions_token_history(history) + if isinstance(history, CompletionsStringHistory): + return _tokenize_completions_string_history( + history, + base_model=base_model, + tokenizer=tokenizer, + ) + override_requires_render = ( + chat_template is not None + and chat_template != getattr(history, "chat_template", None) + ) or ( + chat_template_kwargs is not None + and dict(chat_template_kwargs) + != (getattr(history, "chat_template_kwargs", None) or {}) + ) + if isinstance(history, ChatCompletionsHistory) and ( + _history_needs_render(history) or override_requires_render + ): + return _tokenize_chat_view( + history, + base_model=base_model, + tokenizer=tokenizer, + chat_template=chat_template, + chat_template_kwargs=chat_template_kwargs, + ) + if ( + isinstance(history, AnthropicMessagesHistory) + and (_history_needs_render(history) or override_requires_render) + ) or ( + isinstance(history, ResponsesHistory) + and (_history_needs_render(history) or override_requires_render) + ): + _validate_history_sources(history) + return _tokenize_chat_view( + history.as_chat_completions_history(), + base_model=base_model, + tokenizer=tokenizer, + chat_template=chat_template, + chat_template_kwargs=chat_template_kwargs, + ) + trajectory = _trajectory_from_history(history) + history_template = getattr(history, "chat_template", None) + history_kwargs = getattr(history, "chat_template_kwargs", None) + return _tokenize_exchange_trajectory( + trajectory, + base_model, + model=model, + chat_template=( + chat_template if chat_template is not None else history_template + ), + chat_template_kwargs={ + **(history_kwargs or {}), + **(chat_template_kwargs or {}), + } + or None, + tokenizer_instance=tokenizer, + ) + + +def _materialize_trajectory( + tokenized: TokenizedHistory, trajectory: Trajectory +) -> TokenizedTrajectory: + return TokenizedTrajectory( + **tokenized.model_dump(), + reward=trajectory.reward, + metrics=dict(trajectory.metrics), + metadata=dict(trajectory.metadata), + ) + + +def tokenize_trajectory( + trajectory: Trajectory, + *, + multi_history: bool, + model: str | None, + base_model: str | None, + tokenizer: Tokenizer | None, + chat_template: str | None, + chat_template_kwargs: Mapping[str, object] | None, +) -> TokenizedTrajectory | TokenizedMultiHistoryTrajectory: + histories = trajectory.histories(model=model) + if not multi_history: + if len(histories) != 1: + selected_models = { + history.model + for history in histories + if not isinstance(history, LegacyHistory) + } + if model is None and len(selected_models) > 1: + raise ValueError( + "Trajectory tokenization requires exactly one model; pass model= to select one" + ) + raise ValueError( + f"Trajectory tokenization requires exactly one history; found {len(histories)}" + ) + return _materialize_trajectory( + tokenize_history( + histories[0], + model=model + if isinstance(histories[0], LegacyHistory) + else histories[0].model, + base_model=base_model, + tokenizer=tokenizer, + chat_template=chat_template, + chat_template_kwargs=chat_template_kwargs, + ), + trajectory, + ) + tokenized = [ + tokenize_history( + history, + model=model if isinstance(history, LegacyHistory) else history.model, + base_model=base_model, + tokenizer=tokenizer, + chat_template=chat_template, + chat_template_kwargs=chat_template_kwargs, + ) + for history in histories + ] + return TokenizedMultiHistoryTrajectory( + histories=tokenized, + reward=trajectory.reward, + metrics=dict(trajectory.metrics), + metadata=dict(trajectory.metadata), + ) + + +def tokenize_group( + group: TrajectoryGroup, + *, + multi_history: bool, + model: str | None, + base_model: str | None, + tokenizer: Tokenizer | None, + chat_template: str | None, + chat_template_kwargs: Mapping[str, object] | None, +) -> ( + TokenizedTrajectoryGroup[TokenizedTrajectory] + | TokenizedTrajectoryGroup[TokenizedMultiHistoryTrajectory] +): + if multi_history: + trajectories = [ + cast( + TokenizedMultiHistoryTrajectory, + tokenize_trajectory( + trajectory, + multi_history=True, + model=model, + base_model=base_model, + tokenizer=tokenizer, + chat_template=chat_template, + chat_template_kwargs=chat_template_kwargs, + ), + ) + for trajectory in group.trajectories + ] + return TokenizedTrajectoryGroup[TokenizedMultiHistoryTrajectory]( + trajectories=trajectories, + metrics=dict(group.metrics), + metadata=dict(group.metadata), + ) + single_trajectories = [ + cast( + TokenizedTrajectory, + tokenize_trajectory( + trajectory, + multi_history=False, + model=model, + base_model=base_model, + tokenizer=tokenizer, + chat_template=chat_template, + chat_template_kwargs=chat_template_kwargs, + ), + ) + for trajectory in group.trajectories + ] + return TokenizedTrajectoryGroup[TokenizedTrajectory]( + trajectories=single_trajectories, + metrics=dict(group.metrics), + metadata=dict(group.metadata), ) diff --git a/src/art/utils/trajectory_logging.py b/src/art/utils/trajectory_logging.py index 325f576d6..eb697229a 100644 --- a/src/art/utils/trajectory_logging.py +++ b/src/art/utils/trajectory_logging.py @@ -9,7 +9,7 @@ import pydantic from art.openai import ART_MOE_ROUTING_METADATA_KEY -from art.trajectories import History, Trajectory, TrajectoryGroup +from art.trajectories import LegacyHistory, Trajectory, TrajectoryGroup if TYPE_CHECKING: import pyarrow as pa @@ -88,7 +88,7 @@ def _choice_data(item: object) -> dict[str, Any] | None: ) -def _history_data(history: History) -> dict[str, Any]: +def _history_data(history: LegacyHistory) -> dict[str, Any]: data = history.model_dump( mode="json", exclude={"messages_and_choices"}, warnings="error" ) @@ -123,7 +123,7 @@ def _trajectory_data(trajectory: Trajectory) -> dict[str, Any]: return data -def _restore_history(data: object) -> History: +def _restore_history(data: object) -> LegacyHistory: if not isinstance(data, dict): raise ValueError("Parquet additional history must be a JSON object") restored = dict(data) @@ -136,7 +136,7 @@ def _restore_history(data: object) -> History: else item for item in messages ] - return History.model_validate(restored) + return LegacyHistory.model_validate(restored) def _restore_trajectory(payload: object) -> Trajectory: diff --git a/src/art/utils/trajectory_migration.py b/src/art/utils/trajectory_migration.py index 7a2121893..10c40b1b8 100644 --- a/src/art/utils/trajectory_migration.py +++ b/src/art/utils/trajectory_migration.py @@ -19,7 +19,7 @@ import pydantic import yaml -from art.trajectories import History, Trajectory, TrajectoryGroup +from art.trajectories import LegacyHistory, Trajectory, TrajectoryGroup from art.types import Choice, Message, MessageOrChoice from art.utils.trajectory_logging import write_trajectory_groups_parquet @@ -49,7 +49,7 @@ def trajectory_group_to_dict(trajectory_group: TrajectoryGroup) -> dict[str, Any } -def history_to_dict(history: History) -> dict[str, Any]: +def history_to_dict(history: LegacyHistory) -> dict[str, Any]: messages_and_choices = [ message_or_choice_to_dict(message_or_choice) for message_or_choice in history.messages_and_choices @@ -146,8 +146,8 @@ def dict_to_trajectory(d: dict[str, Any]) -> Trajectory: ) -def dict_to_history(d: dict[str, Any]) -> History: - return History.model_validate( +def dict_to_history(d: dict[str, Any]) -> LegacyHistory: + return LegacyHistory.model_validate( { **d, "messages_and_choices": [ diff --git a/tests/integration/megatron/model_support/chat_template_rollout.py b/tests/integration/megatron/model_support/chat_template_rollout.py index 30a42148e..3ea2d897f 100644 --- a/tests/integration/megatron/model_support/chat_template_rollout.py +++ b/tests/integration/megatron/model_support/chat_template_rollout.py @@ -16,7 +16,7 @@ tokenize_trajectory, tokenize_trajectory_groups, ) -from art.trajectories import History +from art.trajectories import LegacyHistory from tests.support.chat_template_conformance_cases import ( build_chat_template_conformance_inputs, ) @@ -33,8 +33,8 @@ def _artifact_dir(base_model: str) -> Path: return path -def _history(trajectory: art.Trajectory) -> History: - return History( +def _history(trajectory: art.Trajectory) -> LegacyHistory: + return LegacyHistory( messages_and_choices=trajectory.messages_and_choices, tools=trajectory.tools, ) diff --git a/tests/support/chat_template_conformance_cases.py b/tests/support/chat_template_conformance_cases.py index 960bc6599..a3b15d03e 100644 --- a/tests/support/chat_template_conformance_cases.py +++ b/tests/support/chat_template_conformance_cases.py @@ -11,7 +11,7 @@ _apply_chat_template_token_ids, _messages_for_chat_template, ) -from art.trajectories import History, Trajectory, TrajectoryGroup +from art.trajectories import LegacyHistory, Trajectory, TrajectoryGroup from art.types import MessagesAndChoices, Tools @@ -133,7 +133,7 @@ def _rendered_ids( def _attach_token_metadata_to_history( tokenizer: PreTrainedTokenizerBase, - history: Trajectory | History, + history: Trajectory | LegacyHistory, ) -> None: items = history.messages_and_choices for index, item in enumerate(items): @@ -291,7 +291,7 @@ def build_chat_template_conformance_inputs( _choice_for_text("maybe", maybe_ids), ), additional_histories=[ - History( + LegacyHistory( messages_and_choices=_messages_and_choices( {"role": "user", "content": "Previous turn."}, _choice_for_text("prior yes", prior_yes_ids), @@ -306,7 +306,7 @@ def build_chat_template_conformance_inputs( _choice_for_text("yes", yes_ids), ), additional_histories=[ - History( + LegacyHistory( messages_and_choices=_messages_and_choices( {"role": "user", "content": "Previous turn."}, _choice_for_text("prior yes", prior_yes_ids), diff --git a/tests/unit/test_tinker_native_exchanges.py b/tests/unit/test_tinker_native_exchanges.py index d1bdcb8a4..3e7dd4359 100644 --- a/tests/unit/test_tinker_native_exchanges.py +++ b/tests/unit/test_tinker_native_exchanges.py @@ -1,7 +1,7 @@ from datetime import datetime from typing import Any -from openai.types.chat import ChatCompletion +from openai.types.chat import ChatCompletion, ChatCompletionMessageParam from openai.types.chat.chat_completion import Choice import pytest @@ -16,6 +16,16 @@ def _exchange( prompt: list[int], output: list[int], *, logprobs: bool = True ) -> ChatCompletionsExchange: + messages: list[ChatCompletionMessageParam] = [ + {"role": "user", "content": "question"} + ] + if len(prompt) > 1: + messages.extend( + [ + {"role": "assistant", "content": "answer"}, + {"role": "user", "content": "next"}, + ] + ) response = ChatCompletion.model_validate( { "id": "chat-1", @@ -48,7 +58,7 @@ def _exchange( } ) return ChatCompletionsExchange( - request=ChatCompletionsRequest(model="test/model", messages=[]), + request=ChatCompletionsRequest(model="test/model", messages=messages), response=response, start_time=datetime.now(), end_time=datetime.now(), diff --git a/tests/unit/test_trajectory_parquet.py b/tests/unit/test_trajectory_parquet.py index e766239b2..24ab9688a 100644 --- a/tests/unit/test_trajectory_parquet.py +++ b/tests/unit/test_trajectory_parquet.py @@ -35,7 +35,7 @@ ) import pytest -from art import History, Trajectory, TrajectoryGroup +from art import LegacyHistory, Trajectory, TrajectoryGroup from art.types import MessageOrChoice from art.utils.trajectory_logging import ( read_trajectory_groups_parquet, @@ -295,7 +295,7 @@ def test_legacy_serializer_round_trips_complete_legacy_trajectory() -> None: [ Trajectory( messages_and_choices=[choice], - additional_histories=[History(messages_and_choices=[choice])], + additional_histories=[LegacyHistory(messages_and_choices=[choice])], reward=1.0, initial_policy_version=3, final_policy_version=4, @@ -513,7 +513,9 @@ def test_complete_models_round_trip(self, tmp_path: Path) -> None: } ], additional_histories=[ - History(messages_and_choices=[{"role": "user", "content": "alternate"}]) + LegacyHistory( + messages_and_choices=[{"role": "user", "content": "alternate"}] + ) ], initial_policy_version=7, final_policy_version=8, diff --git a/tests/unit/trajectories/test_capture.py b/tests/unit/trajectories/test_capture.py index c28c96157..e13792af9 100644 --- a/tests/unit/trajectories/test_capture.py +++ b/tests/unit/trajectories/test_capture.py @@ -123,6 +123,7 @@ "stop_reason": "end_turn", "stop_sequence": None, "usage": {"input_tokens": 1, "output_tokens": 1}, + "prompt_token_ids": [1], "token_ids": [2], "logprobs": [-0.2], } @@ -172,7 +173,7 @@ async def handler(request: web.Request) -> web.StreamResponse: async def test_contexts_are_nested_and_task_local() -> None: assert art.current_trajectory() is None with art.Trajectory() as outer: - assert art.current_trajectory(required=True) is outer + assert art.current_trajectory(require=True) is outer with art.Trajectory() as inner: assert art.current_trajectory() is inner assert art.current_trajectory() is outer @@ -187,7 +188,7 @@ async def child() -> art.Trajectory: assert first is not second assert art.current_trajectory() is None with pytest.raises(RuntimeError, match="No trajectory"): - art.current_trajectory(required=True) + art.current_trajectory(require=True) async def test_group_context_and_async_helpers() -> None: @@ -643,6 +644,7 @@ def test_all_streaming_protocols_reconstruct_final_responses() -> None: "content": [], "stop_reason": None, "usage": {"input_tokens": 1, "output_tokens": 0}, + "prompt_token_ids": [1], }, }, ), @@ -660,6 +662,8 @@ def test_all_streaming_protocols_reconstruct_final_responses() -> None: "type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "hello"}, + "token_ids": [2], + "logprobs": [-0.2], }, ), ("content_block_stop", {"type": "content_block_stop", "index": 0}), @@ -669,6 +673,7 @@ def test_all_streaming_protocols_reconstruct_final_responses() -> None: "type": "message_delta", "delta": {"stop_reason": "end_turn", "stop_sequence": None}, "usage": {"output_tokens": 1}, + "prompt_token_ids": [1], "token_ids": [2], "logprobs": [-0.2], }, @@ -710,8 +715,12 @@ def test_all_streaming_protocols_reconstruct_final_responses() -> None: content = exchange.response.content[0] assert isinstance(content, TextBlock) assert content.text == "hello" + assert content.model_extra is not None + assert content.model_extra["token_ids"] == [2] + assert content.model_extra["logprobs"] == [-0.2] assert getattr(exchange.response, "token_ids") == [2] assert getattr(exchange.response, "logprobs") == [-0.2] + assert getattr(exchange.response, "prompt_token_ids") == [1] def test_streaming_chat_choices_are_accumulated_by_index() -> None: diff --git a/tests/unit/trajectories/test_history.py b/tests/unit/trajectories/test_history.py index 809fa8196..e46685fbb 100644 --- a/tests/unit/trajectories/test_history.py +++ b/tests/unit/trajectories/test_history.py @@ -1,5 +1,6 @@ from datetime import datetime, timedelta import importlib +from typing import Any from anthropic.types import Message from openai.types import Completion @@ -199,30 +200,113 @@ def test_chat_history_resolves_one_model_and_append_only_sequence() -> None: "user", "assistant", ] + assert [source.exchange for source in history.message_sources if source] == [ + first, + first, + second, + second, + ] + assert history.messages is not first.request["messages"] assert trajectory.chat_completions_history(model="test/model") == history second.request["cache_salt"] = "new-cache" - with pytest.raises(ValueError, match="different cache_salt"): - trajectory.chat_completions_history(model="test/model") + assert trajectory.chat_completions_history(model="test/model") == history second.request.pop("cache_salt") second.request["messages"] = [{"role": "user", "content": "branch"}] - with pytest.raises(ValueError, match="append-only"): + assert len(trajectory.chat_completions_histories(model="test/model")) == 2 + with pytest.raises(ValueError, match="exactly one history"): trajectory.chat_completions_history(model="test/model") +def test_chat_choices_branch_and_identical_continuation_uses_first_choice() -> None: + first = _chat([{"role": "user", "content": "one"}], "same") + response = first.response.model_dump(mode="python") + second_choice = dict(response["choices"][0]) + second_choice["index"] = 1 + response["choices"].append(second_choice) + first.response = ChatCompletion.model_validate(response) + continuation = _chat( + [ + {"role": "user", "content": "one"}, + {"role": "assistant", "content": "same"}, + {"role": "user", "content": "two"}, + ], + "continued", + offset=1, + ) + trajectory = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[first, continuation]) + ) + + histories = trajectory.chat_completions_histories() + + assert len(histories) == 2 + assert [len(history.messages) for history in histories] == [4, 2] + assert histories[0].message_sources[1] is not None + assert histories[0].message_sources[1].choice_index == 0 + assert histories[1].message_sources[1] is not None + assert histories[1].message_sources[1].choice_index == 1 + + +def test_history_mutation_must_keep_source_sidecar_consistent() -> None: + trajectory = art.Trajectory( + exchanges=TrajectoryExchanges( + chat_completions=[_chat([{"role": "user", "content": "one"}], "first")] + ) + ) + history = trajectory.chat_completions_history() + history.messages.append({"role": "user", "content": "next"}) + with pytest.raises(ValueError, match="differ in length"): + history.tokenize() + + history.message_sources.append(None) + history.messages[0] = {"role": "user", "content": "edited"} + with pytest.raises(ValueError, match="no longer matches"): + history.tokenize() + + +def test_history_accepts_user_authored_messages_with_none_source() -> None: + trajectory = art.Trajectory( + exchanges=TrajectoryExchanges( + chat_completions=[_chat([{"role": "user", "content": "one"}], "first")] + ) + ) + history = trajectory.chat_completions_history() + history.messages.append({"role": "user", "content": "next"}) + history.message_sources.append(None) + + class Tokenizer: + def __call__(self, text: str, *, add_special_tokens: bool = False) -> list[int]: + assert not add_special_tokens + return {"first": [20], "next": [30]}[text] + + def apply_chat_template( + self, + messages: list[dict[str, Any]], + *, + tools: object, + tokenize: bool, + add_generation_prompt: bool, + chat_template: str | None = None, + **kwargs: object, + ) -> list[int]: + del tools, tokenize, add_generation_prompt, chat_template, kwargs + return [10] if len(messages) == 1 else [10, 20, 30] + + tokenized = history.tokenize(tokenizer=Tokenizer()) + + assert tokenized.token_ids == [10, 20, 30] + assert tokenized.flags[1] == art.TokenFlag.SAMPLED + + def test_protocol_histories_convert_to_chat_and_history_rejects_ambiguity() -> None: message_trajectory = art.Trajectory( exchanges=TrajectoryExchanges(messages=[_message()]) ) messages_history = message_trajectory.anthropic_messages_history() assert messages_history.system == "Be concise" - assert ( - art.AnthropicMessagesHistory.model_validate_json( - messages_history.model_dump_json() - ) - == messages_history - ) + assert not hasattr(messages_history, "model_dump") assert [message["role"] for message in messages_history.messages] == [ "user", "assistant", @@ -244,6 +328,43 @@ def test_protocol_histories_convert_to_chat_and_history_rejects_ambiguity() -> N assert isinstance(mixed.anthropic_messages_history(), art.AnthropicMessagesHistory) +def test_anthropic_chat_conversion_preserves_sources_for_expanded_messages() -> None: + exchange = _message() + exchange.request["messages"] = [ + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "call-1", + "content": "result", + }, + {"type": "text", "text": "continue"}, + ], + } + ] + history = art.Trajectory( + exchanges=TrajectoryExchanges(messages=[exchange]) + ).anthropic_messages_history() + + converted = history.as_chat_completions_history() + + assert [message["role"] for message in converted.messages] == [ + "system", + "tool", + "user", + "assistant", + ] + for source in converted.message_sources[1:3]: + assert source is not None + assert source.exchange is exchange + assert source.request_index == 0 + + converted.messages[2] = {"role": "user", "content": "changed"} + with pytest.raises(ValueError, match="no longer matches"): + converted.tokenize() + + def test_responses_history_expands_previous_response_chain() -> None: trajectory = art.Trajectory( exchanges=TrajectoryExchanges( @@ -261,9 +382,6 @@ def test_responses_history_expands_previous_response_chain() -> None: history = trajectory.responses_history() assert len(history.input) == 5 - assert ( - art.ResponsesHistory.model_validate_json(history.model_dump_json()) == history - ) chat_history = history.as_chat_completions_history() assert [message["role"] for message in chat_history.messages] == [ "user", @@ -272,16 +390,14 @@ def test_responses_history_expands_previous_response_chain() -> None: "assistant", ] assert dict(chat_history.messages[1]).get("reasoning") == "think" - assert ( - art.ChatCompletionsHistory.model_validate_json( - chat_history.model_dump_json(warnings="error") - ) - == chat_history - ) + assert all(source is not None for source in chat_history.message_sources) + assert chat_history.message_sources[1] is not None + assert chat_history.message_sources[1].output_index == 0 trajectory.exchanges.responses[1].request["previous_response_id"] = "missing" - with pytest.raises(ValueError, match="outside this history"): - trajectory.responses_history() + external = trajectory.responses_histories() + assert len(external) == 2 + assert external[1].previous_response_id == "missing" def test_completions_history_preserves_exact_tokens_and_sampled_spans() -> None: @@ -294,25 +410,177 @@ def test_completions_history_preserves_exact_tokens_and_sampled_spans() -> None: ) ) - history = trajectory.completions_history() - assert history.token_ids == [1, 2, 3, 4] + history = trajectory.completions_token_history() + assert history.prompt == [1, 2, 3, 4] assert history.sampled_spans == [(1, 2), (3, 4)] with pytest.raises(ValueError, match="no chat-message structure"): history.as_chat_completions_history() -def test_completions_history_uses_request_token_ids_and_rejects_echo() -> None: +def test_completions_history_uses_request_token_ids() -> None: exchange = _completion([1], [2]) response = exchange.response.model_dump(mode="python") response["choices"][0].pop("prompt_token_ids") exchange.response = Completion.model_validate(response) trajectory = art.Trajectory(exchanges=TrajectoryExchanges(completions=[exchange])) - assert trajectory.completions_history().token_ids == [1, 2] + assert trajectory.completions_token_history().prompt == [1, 2] - exchange.request["echo"] = True - with pytest.raises(ValueError, match="echo=True"): - trajectory.completions_history() + +def test_batched_completions_create_every_prompt_choice_history() -> None: + exchange = _completion([1], [10]) + exchange.request["prompt"] = ["first", "second"] + response = exchange.response.model_dump(mode="python") + response["choices"] = [ + { + "index": index, + "finish_reason": "stop", + "text": f"answer-{index}", + "prompt_token_ids": [prompt_id], + "token_ids": [100 + index], + } + for index, prompt_id in enumerate((1, 1, 2, 2)) + ] + exchange.response = Completion.model_validate(response) + trajectory = art.Trajectory(exchanges=TrajectoryExchanges(completions=[exchange])) + + histories = trajectory.completions_token_histories() + + assert [history.prompt for history in histories] == [ + [1, 100], + [1, 101], + [2, 102], + [2, 103], + ] + with pytest.raises(ValueError, match="exactly one history"): + trajectory.history() + + +def test_completions_reject_ambiguous_batches_and_suffix() -> None: + ambiguous = _completion([1], [10]) + ambiguous.request["prompt"] = ["first", "second"] + with pytest.raises(ValueError, match="associate Completions choices"): + art.Trajectory( + exchanges=TrajectoryExchanges(completions=[ambiguous]) + ).histories() + + insertion = _completion([1], [10]) + insertion.request["suffix"] = "tail" + with pytest.raises(ValueError, match="suffix is not supported"): + art.Trajectory( + exchanges=TrajectoryExchanges(completions=[insertion]) + ).histories() + + +def test_reasoning_stripping_produces_truthful_history_per_generation() -> None: + def exchange( + offset: int, + request_messages: list[dict[str, object]], + answer: str, + ) -> MessagesExchange: + start, end = _times(offset) + return MessagesExchange( + request={ + "model": "test/model", + "messages": request_messages, + "max_tokens": 16, + "thinking": {"type": "enabled", "budget_tokens": 8}, + }, + response=Message.model_validate( + { + "id": f"message-{offset}", + "type": "message", + "role": "assistant", + "model": "test/model", + "content": [ + { + "type": "thinking", + "thinking": f"thought-{offset}", + "signature": "sig", + }, + {"type": "text", "text": answer}, + ], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 1, "output_tokens": 1}, + } + ), + start_time=start, + end_time=end, + ) + + first = exchange(0, [{"role": "user", "content": "one"}], "first") + second = exchange( + 1, + [ + {"role": "user", "content": "one"}, + {"role": "assistant", "content": "first"}, + {"role": "user", "content": "two"}, + ], + "second", + ) + third = exchange( + 2, + [ + {"role": "user", "content": "one"}, + {"role": "assistant", "content": "first"}, + {"role": "user", "content": "two"}, + {"role": "assistant", "content": "second"}, + {"role": "user", "content": "three"}, + ], + "third", + ) + trajectory = art.Trajectory( + exchanges=TrajectoryExchanges(messages=[first, second, third]) + ) + + histories = trajectory.anthropic_messages_histories() + + assert len(histories) == 3 + assert [len(history.messages) for history in histories] == [2, 4, 6] + assert histories[1].message_sources[1] is not None + assert histories[1].message_sources[1].exchange is first + assert histories[1].message_sources[1].request_index is None + with pytest.raises(ValueError, match="exactly one history"): + trajectory.tokenize() + + +def test_chat_template_stripped_reasoning_splits_exact_histories() -> None: + first = _chat([{"role": "user", "content": "one"}], "first") + first_data = first.response.model_dump(mode="python") + first_data["prompt_token_ids"] = [1] + first_data["choices"][0]["message"]["reasoning"] = "thought-one" + first_data["choices"][0]["token_ids"] = [2, 3] + first.response = ChatCompletion.model_validate(first_data) + + second = _chat( + [ + {"role": "user", "content": "one"}, + { + "role": "assistant", + "content": "first", + "reasoning": "thought-one", + }, + {"role": "user", "content": "two"}, + ], + "second", + offset=1, + ) + second_data = second.response.model_dump(mode="python") + second_data["prompt_token_ids"] = [1, 3, 4] + second_data["choices"][0]["message"]["reasoning"] = "thought-two" + second_data["choices"][0]["token_ids"] = [5, 6] + second.response = ChatCompletion.model_validate(second_data) + trajectory = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[first, second]) + ) + + histories = trajectory.chat_completions_histories() + + assert len(histories) == 2 + assert [len(history.messages) for history in histories] == [2, 4] + with pytest.raises(ValueError, match="exactly one history"): + trajectory.tokenize() def test_history_rejects_mutated_mixed_representation() -> None: @@ -332,22 +600,27 @@ def test_legacy_messages_delegate_through_history() -> None: messages_and_choices=[{"role": "user", "content": "hello"}] ) - assert isinstance(trajectory.history(), art.History) + assert isinstance(trajectory.history(), art.LegacyHistory) assert trajectory.messages() == [{"role": "user", "content": "hello"}] - with pytest.raises(ValueError, match="do not identify a model"): - trajectory.history(model="test/model") + assert isinstance(trajectory.history(model="test/model"), art.LegacyHistory) + with pytest.raises(ValueError, match="requires model="): + trajectory.tokenize() def test_legacy_messages_preserve_primary_history_with_additional_histories() -> None: trajectory = art.Trajectory( messages_and_choices=[{"role": "user", "content": "primary"}], additional_histories=[ - art.History(messages_and_choices=[{"role": "user", "content": "alternate"}]) + art.LegacyHistory( + messages_and_choices=[{"role": "user", "content": "alternate"}] + ) ], ) assert trajectory.messages() == [{"role": "user", "content": "primary"}] - with pytest.raises(ValueError, match="multiple legacy histories"): + assert len(trajectory.histories()) == 2 + assert len(trajectory.chat_completions_histories()) == 2 + with pytest.raises(ValueError, match="exactly one history"): trajectory.history() diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index 7dc2d80d8..4c102e11d 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -3,12 +3,13 @@ import builtins from datetime import datetime, timedelta import math -from types import SimpleNamespace +import sys +from types import ModuleType, SimpleNamespace from typing import Any from anthropic.types import ImageBlockParam, Message, MessageParam from openai.types import Completion -from openai.types.chat import ChatCompletion +from openai.types.chat import ChatCompletion, ChatCompletionMessageParam from openai.types.chat.chat_completion_token_logprob import ChatCompletionTokenLogprob from openai.types.responses import Response import pytest @@ -34,6 +35,11 @@ def _chat_exchange( model: str = "test/model", offset: int = 0, ) -> ChatCompletionsExchange: + messages: list[ChatCompletionMessageParam] = [] + for turn in range(offset + 1): + messages.append({"role": "user", "content": f"turn {turn}"}) + if turn < offset: + messages.append({"role": "assistant", "content": "answer"}) response = ChatCompletion.model_validate( { "id": f"chat-{offset}", @@ -66,7 +72,7 @@ def _chat_exchange( return ChatCompletionsExchange( request=ChatCompletionsRequest( model=model, - messages=[{"role": "user", "content": f"turn {offset}"}], + messages=messages, ), response=response, start_time=start, @@ -89,7 +95,7 @@ def _completion_exchange( { "index": 0, "finish_reason": "stop", - "text": "answer", + "text": f"{'question' if echo else ''}answer", "prompt_token_ids": [1], "token_ids": [2], "logprobs": { @@ -138,7 +144,7 @@ def import_without_tokenizer_dependencies(name: str, *args: Any, **kwargs: Any): monkeypatch.delenv("WANDB_API_KEY", raising=False) monkeypatch.setattr(builtins, "__import__", import_without_tokenizer_dependencies) - tokenized = art.tokenize_trajectory(trajectory) + tokenized = trajectory.tokenize() assert tokenized.token_ids == [1, 2, 3, 4] assert tokenized.flags == [ @@ -153,6 +159,58 @@ def import_without_tokenizer_dependencies(name: str, *args: Any, **kwargs: Any): assert tokenized.logprobs[3] == -0.4 +def test_messages_exact_prompt_and_output_do_not_load_a_tokenizer( + monkeypatch: pytest.MonkeyPatch, +) -> None: + response = Message.model_validate( + { + "id": "message-exact", + "type": "message", + "role": "assistant", + "model": "test/model", + "content": [{"type": "text", "text": "answer"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 2, "output_tokens": 1}, + "prompt_token_ids": [1, 2], + "token_ids": [3], + "logprobs": [-0.3], + } + ) + start = datetime(2026, 1, 1) + trajectory = art.Trajectory( + exchanges=TrajectoryExchanges( + messages=[ + MessagesExchange( + request=MessagesRequest( + model="test/model", + messages=[{"role": "user", "content": "question"}], + max_tokens=16, + ), + response=response, + start_time=start, + end_time=start + timedelta(milliseconds=1), + ) + ] + ) + ) + monkeypatch.setattr( + "art.trajectories._tokenize._load_tokenizer", + lambda _config: pytest.fail("exact Messages evidence loaded a tokenizer"), + ) + + tokenized = trajectory.tokenize() + + assert tokenized.token_ids == [1, 2, 3] + assert all(math.isnan(value) for value in tokenized.logprobs[:2]) + assert tokenized.logprobs[2] == -0.3 + assert tokenized.flags == [ + art.TokenFlag.EXACT, + art.TokenFlag.EXACT, + art.TokenFlag.EXACT | art.TokenFlag.SAMPLED, + ] + + def test_malformed_explicit_exact_token_metadata_fails_closed() -> None: chat = _chat_exchange([1], [2]) chat_extra = chat.response.choices[0].model_extra @@ -202,7 +260,7 @@ def test_malformed_explicit_exact_token_metadata_fails_closed() -> None: ] for trajectory in trajectories: with pytest.raises(ValueError, match="exact token"): - art.tokenize_trajectory(trajectory, base_model="base/model") + trajectory.tokenize(base_model="base/model") @pytest.mark.parametrize( @@ -213,26 +271,56 @@ def test_malformed_explicit_exact_token_metadata_fails_closed() -> None: _completion_exchange(echo=True), ], ) -def test_completions_reject_batch_prompts_and_echo( +def test_completions_support_single_item_batches_and_echo( exchange: CompletionsExchange, ) -> None: - with pytest.raises(ValueError, match="batched Completions|echo=True"): - art.tokenize_trajectory( - art.Trajectory(exchanges=TrajectoryExchanges(completions=[exchange])) - ) + tokenized = art.Trajectory( + exchanges=TrajectoryExchanges(completions=[exchange]) + ).tokenize() + assert tokenized.token_ids == [1, 2] + + +@pytest.mark.parametrize("response_token_ids", [[1, 2], [2]]) +def test_completions_echo_preserves_prompt_logprobs_without_sampling_them( + response_token_ids: list[int], +) -> None: + exchange = _completion_exchange(echo=True) + payload = exchange.response.model_dump(mode="python") + payload["choices"][0]["token_ids"] = response_token_ids + payload["choices"][0]["logprobs"] = { + "tokens": ["token_id:1", "token_id:2"], + "token_logprobs": [-0.1, -0.2], + "top_logprobs": [{}, {}], + "text_offset": [0, 8], + } + exchange.response = Completion.model_validate(payload) + + tokenized = art.Trajectory( + exchanges=TrajectoryExchanges(completions=[exchange]) + ).tokenize() + + assert tokenized.token_ids == [1, 2] + assert tokenized.logprobs == [-0.1, -0.2] + assert tokenized.flags == [ + art.TokenFlag.EXACT, + art.TokenFlag.EXACT | art.TokenFlag.SAMPLED, + ] def test_branching_and_multiple_models_require_explicit_resolution() -> None: + alternate = _chat_exchange([9], [3], offset=1) + alternate.request["messages"] = [{"role": "user", "content": "alternate"}] branching = art.Trajectory( exchanges=TrajectoryExchanges( chat_completions=[ _chat_exchange([1], [2], offset=0), - _chat_exchange([9], [3], offset=1), + alternate, ] ) ) - with pytest.raises(ValueError, match="append-only"): - art.tokenize_trajectory(branching) + with pytest.raises(ValueError, match="exactly one history"): + branching.tokenize() + assert len(branching.tokenize(multi_history=True).histories) == 2 mixed = art.Trajectory( exchanges=TrajectoryExchanges( @@ -243,14 +331,39 @@ def test_branching_and_multiple_models_require_explicit_resolution() -> None: ) ) with pytest.raises(ValueError, match="exactly one model"): - art.tokenize_trajectory(mixed) - assert art.tokenize_trajectory(mixed, model="two").token_ids == [3, 4] + mixed.tokenize() + assert mixed.tokenize(model="two").token_ids == [3, 4] + assert [ + history.model for history in mixed.tokenize(multi_history=True).histories + ] == ["one", "two"] + + +def test_legacy_additional_histories_require_multi_history_and_model() -> None: + first = _chat_exchange([1], [2]).response.choices[0] + second = _chat_exchange([3], [4]).response.choices[0] + trajectory = art.Trajectory( + messages_and_choices=[first], + additional_histories=[art.LegacyHistory(messages_and_choices=[second])], + ) + + with pytest.raises(ValueError, match="exactly one history"): + trajectory.tokenize(model="test/model") + with pytest.raises(ValueError, match="requires model="): + trajectory.tokenize(multi_history=True) + + tokenized = trajectory.tokenize(multi_history=True, model="test/model") + + assert [history.token_ids for history in tokenized.histories] == [[1, 2], [3, 4]] class _FakeTokenizer: def __init__(self) -> None: self.calls: list[dict[str, object]] = [] + def __call__(self, text: str, *, add_special_tokens: bool = False) -> list[int]: + del text, add_special_tokens + return [11] + def apply_chat_template( self, messages: list[dict[str, Any]], **kwargs: object ) -> list[int]: @@ -297,8 +410,9 @@ def test_fallback_uses_template_overrides_and_nan_logprobs( lambda _model: pytest.fail("explicit base_model should bypass W&B"), ) - result = art.tokenize_trajectory( - art.Trajectory(exchanges=TrajectoryExchanges(messages=[exchange])), + result = art.Trajectory( + exchanges=TrajectoryExchanges(messages=[exchange]) + ).tokenize( base_model="base/model", chat_template="explicit-template", chat_template_kwargs={"explicit": True}, @@ -309,16 +423,6 @@ def test_fallback_uses_template_overrides_and_nan_logprobs( assert result.flags == [art.TokenFlag(0), art.TokenFlag.SAMPLED] assert math.isnan(result.logprobs[1]) assert tokenizer.calls == [ - { - "tools": None, - "tokenize": True, - "add_generation_prompt": True, - "chat_template": "explicit-template", - "request": True, - "explicit": True, - "enable_thinking": True, - "thinking_budget": 128, - }, { "tools": None, "tokenize": True, @@ -360,7 +464,15 @@ def artifact(self, name: str) -> SimpleNamespace: } ) - monkeypatch.setattr("wandb.apis.public.Api", Api) + wandb = ModuleType("wandb") + apis = ModuleType("wandb.apis") + public = ModuleType("wandb.apis.public") + setattr(public, "Api", Api) + setattr(apis, "public", public) + setattr(wandb, "apis", apis) + monkeypatch.setitem(sys.modules, "wandb", wandb) + monkeypatch.setitem(sys.modules, "wandb.apis", apis) + monkeypatch.setitem(sys.modules, "wandb.apis.public", public) exchange = _chat_exchange([], [], model=model) extra = exchange.response.choices[0].model_extra assert extra is not None @@ -374,9 +486,9 @@ def artifact(self, name: str) -> SimpleNamespace: lambda config: configs.append(config) or tokenizer, ) - art.tokenize_trajectory( - art.Trajectory(exchanges=TrajectoryExchanges(chat_completions=[exchange])) - ) + art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[exchange]) + ).tokenize() config = configs[0] assert artifact_names == [artifact_name] @@ -388,6 +500,66 @@ def artifact(self, name: str) -> SimpleNamespace: assert tokenizer.calls[0]["thinking"] is True +def test_loaded_tokenizers_are_cached_by_model_and_revision( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from art.trajectories._tokenize import ( + _cached_tokenizer, + _load_tokenizer, + _TokenizerConfig, + ) + + loaded: list[tuple[str, str | None]] = [] + + class AutoTokenizer: + @staticmethod + def from_pretrained(model: str, *, revision: str | None) -> object: + loaded.append((model, revision)) + return object() + + transformers = ModuleType("transformers") + setattr(transformers, "AutoTokenizer", AutoTokenizer) + monkeypatch.setitem(sys.modules, "transformers", transformers) + _cached_tokenizer.cache_clear() + try: + config = _TokenizerConfig("test/model", revision="revision") + assert _load_tokenizer(config) is _load_tokenizer(config) + assert loaded == [("test/model", "revision")] + finally: + _cached_tokenizer.cache_clear() + + +def test_deepseek_v4_uses_arts_protocol_renderer( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from art.trajectories._tokenize import _cached_tokenizer + + raw = object() + wrapped = object() + + class AutoTokenizer: + @staticmethod + def from_pretrained(model: str, *, revision: str | None) -> object: + assert model == "deepseek-ai/DeepSeek-V4-Flash" + assert revision is None + return raw + + transformers = ModuleType("transformers") + setattr(transformers, "AutoTokenizer", AutoTokenizer) + monkeypatch.setitem(sys.modules, "transformers", transformers) + monkeypatch.setattr( + "art.megatron.dsv4.tokenizer.get_dsv4_tokenizer", + lambda tokenizer: ( + wrapped if tokenizer is raw else pytest.fail("wrong tokenizer") + ), + ) + _cached_tokenizer.cache_clear() + try: + assert _cached_tokenizer("deepseek-ai/DeepSeek-V4-Flash", None) is wrapped + finally: + _cached_tokenizer.cache_clear() + + def test_anthropic_fallback_preserves_thinking_and_tool_history() -> None: from art.trajectories._tokenize import _anthropic_messages @@ -447,6 +619,111 @@ def test_anthropic_fallback_preserves_thinking_and_tool_history() -> None: ] +def test_reasoning_stripped_history_uses_stream_block_tokens() -> None: + def exchange( + offset: int, + request_messages: list[MessageParam], + answer: str, + token_ids: list[int], + ) -> MessagesExchange: + start = datetime(2026, 1, 1) + timedelta(seconds=offset) + return MessagesExchange( + request=MessagesRequest( + model="test/model", + messages=request_messages, + max_tokens=16, + thinking={"type": "enabled", "budget_tokens": 8}, + ), + response=Message.model_validate( + { + "id": f"message-{offset}", + "type": "message", + "role": "assistant", + "model": "test/model", + "content": [ + { + "type": "thinking", + "thinking": f"thought-{offset}", + "signature": "signature", + "token_ids": [90 + offset], + "logprobs": [-9.0 - offset], + }, + { + "type": "text", + "text": answer, + "token_ids": token_ids, + "logprobs": [-0.1 * token for token in token_ids], + }, + ], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 1, "output_tokens": len(token_ids)}, + } + ), + start_time=start, + end_time=start + timedelta(milliseconds=1), + ) + + first = exchange( + 0, + [{"role": "user", "content": "one"}], + "first", + [101, 102], + ) + second = exchange( + 1, + [ + {"role": "user", "content": "one"}, + {"role": "assistant", "content": "first"}, + {"role": "user", "content": "two"}, + ], + "second", + [201], + ) + + class Tokenizer: + def __call__(self, text: str, *, add_special_tokens: bool = False) -> list[int]: + assert not add_special_tokens + return { + "first": [50], + "second": [60], + "thought-1": [70], + "two": [11], + }[text] + + def apply_chat_template( + self, + messages: list[dict[str, Any]], + **kwargs: object, + ) -> list[int]: + del kwargs + by_length = { + 1: [10], + 2: [10, 50], + 3: [10, 50, 11], + 4: [10, 50, 11, 70, 60], + } + return by_length[len(messages)] + + history = art.Trajectory( + exchanges=TrajectoryExchanges(messages=[first, second]) + ).anthropic_messages_histories()[1] + tokenized = history.tokenize(tokenizer=Tokenizer()) + + assert tokenized.token_ids == [10, 101, 102, 11, 91, 201] + assert tokenized.logprobs[1:3] == pytest.approx([-10.1, -10.2]) + assert tokenized.logprobs[-2] == pytest.approx(-10.0) + assert tokenized.logprobs[-1] == pytest.approx(-20.1) + assert tokenized.flags[1:3] == [ + art.TokenFlag.EXACT | art.TokenFlag.SAMPLED, + art.TokenFlag.EXACT | art.TokenFlag.SAMPLED, + ] + assert tokenized.flags[-2:] == [ + art.TokenFlag.EXACT | art.TokenFlag.SAMPLED, + art.TokenFlag.EXACT | art.TokenFlag.SAMPLED, + ] + + def test_choice_logprobs_survive_tokenizer_fallback( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -480,8 +757,9 @@ def __call__(self, text: str, **kwargs: object) -> SimpleNamespace: monkeypatch.setattr( "art.trajectories._tokenize._load_tokenizer", lambda _config: Tokenizer() ) - result = art.tokenize_trajectory( - art.Trajectory(exchanges=TrajectoryExchanges(chat_completions=[exchange])), + result = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[exchange]) + ).tokenize( base_model="base/model", ) assert result.token_ids == [10, 11, 12] @@ -522,8 +800,9 @@ def __call__(self, text: str, **kwargs: object) -> SimpleNamespace: monkeypatch.setattr( "art.trajectories._tokenize._load_tokenizer", lambda _config: Tokenizer() ) - result = art.tokenize_trajectory( - art.Trajectory(exchanges=TrajectoryExchanges(chat_completions=[exchange])), + result = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[exchange]) + ).tokenize( base_model="base/model", ) @@ -531,7 +810,7 @@ def __call__(self, text: str, **kwargs: object) -> SimpleNamespace: assert all(math.isnan(logprob) for logprob in result.logprobs[1:]) -def test_legacy_logprob_mismatch_fails_closed() -> None: +def test_legacy_token_and_logprob_length_mismatch_raises() -> None: exchange = _chat_exchange([1], [2, 3]) choice = exchange.response.choices[0] assert choice.logprobs is not None @@ -547,13 +826,8 @@ def test_legacy_logprob_mismatch_fails_closed() -> None: } ) - result = art.tokenize_trajectory( - art.Trajectory(messages_and_choices=[choice]), - ) - - assert result.token_ids == [1, 2, 3] - assert len(result.logprobs) == len(result.token_ids) - assert all(math.isnan(logprob) for logprob in result.logprobs) + with pytest.raises(ValueError, match="differ in length"): + art.Trajectory(messages_and_choices=[choice]).tokenize(model="test/model") def test_anthropic_fallback_rejects_unknown_content_blocks( @@ -595,8 +869,7 @@ def test_anthropic_fallback_rejects_unknown_content_blocks( ) with pytest.raises(ValueError, match="Unsupported Anthropic content block"): - art.tokenize_trajectory( - art.Trajectory(exchanges=TrajectoryExchanges(messages=[exchange])), + art.Trajectory(exchanges=TrajectoryExchanges(messages=[exchange])).tokenize( base_model="base/model", ) @@ -630,8 +903,9 @@ def apply_chat_template( monkeypatch.setattr( "art.trajectories._tokenize._load_tokenizer", lambda _config: Tokenizer() ) - result = art.tokenize_trajectory( - art.Trajectory(exchanges=TrajectoryExchanges(chat_completions=[exchange])), + result = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[exchange]) + ).tokenize( base_model="base/model", ) @@ -756,12 +1030,11 @@ def apply_chat_template( monkeypatch.setattr( "art.trajectories._tokenize._load_tokenizer", lambda _config: Tokenizer() ) - result = art.tokenize_trajectory( - art.Trajectory( - exchanges=TrajectoryExchanges( - responses=[_response_with_content_logprobs(exact_second=True)] - ) - ), + result = art.Trajectory( + exchanges=TrajectoryExchanges( + responses=[_response_with_content_logprobs(exact_second=True)] + ) + ).tokenize( base_model="base/model", ) @@ -780,8 +1053,9 @@ def test_responses_empty_raw_tokens_fall_back_for_visible_output( "art.trajectories._tokenize._load_tokenizer", lambda _config: _FakeTokenizer() ) - result = art.tokenize_trajectory( - art.Trajectory(exchanges=TrajectoryExchanges(responses=[exchange])), + result = art.Trajectory( + exchanges=TrajectoryExchanges(responses=[exchange]) + ).tokenize( base_model="base/model", chat_template="template", chat_template_kwargs={}, @@ -809,12 +1083,11 @@ def __call__(self, text: str, **kwargs: object) -> SimpleNamespace: monkeypatch.setattr( "art.trajectories._tokenize._load_tokenizer", lambda _config: Tokenizer() ) - result = art.tokenize_trajectory( - art.Trajectory( - exchanges=TrajectoryExchanges( - responses=[_response_with_content_logprobs(exact_second=False)] - ) - ), + result = art.Trajectory( + exchanges=TrajectoryExchanges( + responses=[_response_with_content_logprobs(exact_second=False)] + ) + ).tokenize( base_model="base/model", ) @@ -859,15 +1132,16 @@ def apply_chat_template( data.pop("raw_output_tokens", None) response_reasoning.response = Response.model_validate(data) - art.tokenize_trajectory( - art.Trajectory(exchanges=TrajectoryExchanges(responses=[request_reasoning])), + art.Trajectory( + exchanges=TrajectoryExchanges(responses=[request_reasoning]) + ).tokenize( base_model="base/model", ) single = art.Trajectory( exchanges=TrajectoryExchanges(responses=[response_reasoning]) ) - assert art.tokenize_trajectory(single, base_model="base/model").token_ids == [ + assert single.tokenize(base_model="base/model").token_ids == [ 10, 2, ] @@ -878,10 +1152,9 @@ def apply_chat_template( previous_response_id=response_reasoning.response.id, offset=1, ) - assert art.tokenize_trajectory( - art.Trajectory( - exchanges=TrajectoryExchanges(responses=[response_reasoning, continuation]) - ), + assert art.Trajectory( + exchanges=TrajectoryExchanges(responses=[response_reasoning, continuation]) + ).tokenize( base_model="base/model", ).token_ids == [10, 2, 3] @@ -904,8 +1177,7 @@ def test_responses_opaque_reasoning_requires_exact_tokens( "art.trajectories._tokenize._load_tokenizer", lambda _config: _FakeTokenizer() ) - assert art.tokenize_trajectory( - art.Trajectory(exchanges=TrajectoryExchanges(responses=[exchange])), + assert art.Trajectory(exchanges=TrajectoryExchanges(responses=[exchange])).tokenize( base_model="base/model", ).token_ids == [10, 2] @@ -913,8 +1185,7 @@ def test_responses_opaque_reasoning_requires_exact_tokens( response.pop("raw_output_tokens", None) exchange.response = Response.model_validate(response) with pytest.raises(ValueError, match="no renderable text"): - art.tokenize_trajectory( - art.Trajectory(exchanges=TrajectoryExchanges(responses=[exchange])), + art.Trajectory(exchanges=TrajectoryExchanges(responses=[exchange])).tokenize( base_model="base/model", ) @@ -949,8 +1220,7 @@ def apply_chat_template( monkeypatch.setattr( "art.trajectories._tokenize._load_tokenizer", lambda _config: Tokenizer() ) - art.tokenize_trajectory( - art.Trajectory(exchanges=TrajectoryExchanges(responses=[exchange])), + art.Trajectory(exchanges=TrajectoryExchanges(responses=[exchange])).tokenize( base_model="base/model", ) @@ -969,7 +1239,7 @@ def test_tokenization_rejects_mutated_mixed_representation() -> None: trajectory.exchanges.chat_completions.append(_chat_exchange([1], [2])) with pytest.raises(ValueError, match="both exchanges and legacy histories"): - art.tokenize_trajectory(trajectory) + trajectory.tokenize() def test_responses_previous_response_id_resolves_local_history( @@ -991,7 +1261,7 @@ def apply_chat_template( exchanges=TrajectoryExchanges(responses=[first, second]) ) - assert art.tokenize_trajectory(trajectory, base_model="base/model").token_ids == [ + assert trajectory.tokenize(base_model="base/model").token_ids == [ 10, 20, 11, @@ -999,8 +1269,185 @@ def apply_chat_template( ] second.request["previous_response_id"] = "missing" + assert len(trajectory.responses_histories()) == 2 + with pytest.raises(ValueError, match="exactly one history"): + trajectory.tokenize(base_model="base/model") with pytest.raises(ValueError, match="outside this trajectory"): - art.tokenize_trajectory(trajectory, base_model="base/model") + trajectory.responses_histories()[1].tokenize(base_model="base/model") + + +def test_prefix_retokenization_preserves_sampled_ids_and_logprobs( + monkeypatch: pytest.MonkeyPatch, +) -> None: + first = _chat_exchange([1], [101, 102]) + first.response.choices[0].message.content = "cat" + second = _chat_exchange([1, 500, 3], [4], offset=1) + second.request["messages"] = [ + {"role": "user", "content": "turn 0"}, + {"role": "assistant", "content": "cat"}, + {"role": "user", "content": "turn 1"}, + ] + + class Tokenizer: + def __call__(self, text: str, *, add_special_tokens: bool = False) -> list[int]: + assert text == "cat" + assert not add_special_tokens + return [500] + + monkeypatch.setattr( + "art.trajectories._tokenize._load_tokenizer", lambda _config: Tokenizer() + ) + monkeypatch.setattr( + "art.trajectories._tokenize._WARNED_PREFIX_RETOKENIZATION", False + ) + trajectory = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[first, second]) + ) + + with pytest.warns(UserWarning, match="preserved the original sampled token IDs"): + tokenized = trajectory.tokenize(base_model="base/model") + + assert tokenized.token_ids == [1, 101, 102, 3, 4] + assert tokenized.logprobs[1:3] == [-10.1, -10.2] + assert all(tokenized.flags[index] & art.TokenFlag.EXACT for index in (1, 2, 4)) + + +def test_template_change_rerenders_scaffold_but_preserves_sampled_output() -> None: + history = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[_chat_exchange([1], [2])]) + ).chat_completions_history() + history.chat_template = "custom" + + class Tokenizer: + def __call__(self, text: str, *, add_special_tokens: bool = False) -> list[int]: + assert text == "answer" + assert not add_special_tokens + return [20] + + def apply_chat_template( + self, + messages: list[dict[str, Any]], + *, + tools: object, + tokenize: bool, + add_generation_prompt: bool, + chat_template: str | None = None, + **kwargs: object, + ) -> list[int]: + del tools, tokenize, add_generation_prompt, kwargs + assert chat_template == "custom" + return [10] if len(messages) == 1 else [10, 20, 30] + + tokenized = history.tokenize(tokenizer=Tokenizer()) + + assert tokenized.token_ids == [10, 2, 30] + assert tokenized.logprobs[1] == -0.2 + assert tokenized.flags[1] == art.TokenFlag.EXACT | art.TokenFlag.SAMPLED + + +def test_explicit_template_override_rerenders_exact_exchange_scaffold() -> None: + trajectory = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[_chat_exchange([1], [2])]) + ) + + class Tokenizer: + def __call__(self, text: str, *, add_special_tokens: bool = False) -> list[int]: + assert text == "answer" + assert not add_special_tokens + return [20] + + def apply_chat_template( + self, + messages: list[dict[str, Any]], + *, + chat_template: str | None = None, + **kwargs: object, + ) -> list[int]: + del kwargs + assert chat_template == "custom" + return [10, 20, 30] + + tokenized = trajectory.tokenize( + tokenizer=Tokenizer(), + chat_template="custom", + ) + + assert tokenized.token_ids == [10, 2, 30] + assert tokenized.logprobs[1] == -0.2 + + +def test_responses_external_context_requires_or_uses_exact_prompt_tokens() -> None: + exchange = _response_exchange( + "external", 2, previous_response_id="outside-trajectory" + ) + history = art.Trajectory( + exchanges=TrajectoryExchanges(responses=[exchange]) + ).responses_history() + with pytest.raises(ValueError, match="without exact prompt tokens"): + history.tokenize(base_model="base/model") + + response = exchange.response.model_dump(mode="python") + response["prompt_token_ids"] = [7, 8] + exchange.response = Response.model_validate(response) + tokenized = ( + art.Trajectory(exchanges=TrajectoryExchanges(responses=[exchange])) + .responses_history() + .tokenize() + ) + + assert tokenized.token_ids == [7, 8, 2] + assert tokenized.flags == [ + art.TokenFlag.EXACT, + art.TokenFlag.EXACT, + art.TokenFlag.EXACT | art.TokenFlag.SAMPLED, + ] + + +def test_responses_conversation_requires_exact_prompt_tokens() -> None: + exchange = _response_exchange("conversation", 2) + exchange.request["conversation"] = "conversation-1" + trajectory = art.Trajectory(exchanges=TrajectoryExchanges(responses=[exchange])) + + with pytest.raises(ValueError, match="conversation history requires exact"): + trajectory.tokenize(base_model="base/model") + + response = exchange.response.model_dump(mode="python") + response["prompt_token_ids"] = [5] + exchange.response = Response.model_validate(response) + assert trajectory.tokenize().token_ids == [5, 2] + + +def test_tokenized_results_materialize_metadata_and_group_shape() -> None: + trajectory = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[_chat_exchange([1], [2])]), + reward=0.75, + metrics={"correct": True}, + metadata={"source": {"name": "unit"}}, + ) + group = art.TrajectoryGroup( + [trajectory], metrics={"batch": 1}, metadata={"split": "test"} + ) + + tokenized = group.tokenize() + + assert tokenized.trajectories[0].model == "test/model" + assert tokenized.trajectories[0].reward == 0.75 + assert tokenized.trajectories[0].metadata == {"source": {"name": "unit"}} + assert tokenized.metrics == {"batch": 1} + assert tokenized.metadata == {"split": "test"} + assert "underlying" not in tokenized.model_dump() + + +def test_completions_history_requires_exhaustive_source_spans() -> None: + history = art.CompletionsTokenHistory( + model="test/model", + prompt=[1], + prompt_sources=[], + sampled_spans=[], + ) + + with pytest.raises(ValueError, match="exhaustively cover"): + history.tokenize() def test_exchange_trajectories_feed_existing_training_tokenizer( @@ -1028,6 +1475,10 @@ def apply_chat_template( self.calls.append(kwargs) return [1, 2] if messages[-1]["role"] == "assistant" else [1] + def __call__(self, text: str, *, add_special_tokens: bool = False) -> list[int]: + del text, add_special_tokens + return [2] + def decode(self, token_id: int) -> str: return str(token_id) From aef9d4471196b21efb992243fe124da30fc0d1c8 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Thu, 23 Jul 2026 23:00:29 +0000 Subject: [PATCH 02/58] Harden protocol-native trajectory histories --- docs/features/additional-histories.mdx | 22 +- src/art/trajectories/__init__.py | 2 + src/art/trajectories/_history.py | 438 ++++++-- src/art/trajectories/_tokenize.py | 1210 +++++++++++++++++----- tests/unit/trajectories/test_history.py | 266 ++++- tests/unit/trajectories/test_tokenize.py | 889 +++++++++++++++- 6 files changed, 2493 insertions(+), 334 deletions(-) diff --git a/docs/features/additional-histories.mdx b/docs/features/additional-histories.mdx index 8eddd1669..620bac48e 100644 --- a/docs/features/additional-histories.mdx +++ b/docs/features/additional-histories.mdx @@ -30,7 +30,7 @@ Some models, like Qwen 3, use chat templates that remove special tokens (such as By splitting each turn into a separate history, you can preserve these tokens for training: ```python -from art.trajectories import Trajectory, History +from art.trajectories import LegacyHistory, Trajectory # Instead of a single multi-turn conversation that loses tokens # Train as separate histories to preserve them @@ -41,7 +41,7 @@ trajectory = Trajectory( {"role": "assistant", "content": "I need to add 2 and 24"} ], additional_histories=[ - History( + LegacyHistory( messages_and_choices=[ # The Qwen 3 chat template removes tokens from previous turns {"role": "user", "content": "What is 2+2?"}, @@ -82,7 +82,7 @@ trajectory = Trajectory( ], additional_histories=[ # Sub-agent 1: Code analysis - History( + LegacyHistory( messages_and_choices=[ {"role": "system", "content": "You are a code analysis expert"}, {"role": "user", "content": "Find potential bugs in main.py"}, @@ -90,7 +90,7 @@ trajectory = Trajectory( ] ), # Sub-agent 2: Bug fixing - History( + LegacyHistory( messages_and_choices=[ {"role": "system", "content": "You are a bug fixing expert"}, {"role": "user", "content": "Fix the null pointer issue on line 42"}, @@ -116,7 +116,7 @@ trajectory = Trajectory( ], additional_histories=[ # Previous conversation segment before compaction - History( + LegacyHistory( messages_and_choices=[ {"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "Compacted conversation history: the user asked about quantum entanglement, and the assistant explained..."}, @@ -152,11 +152,11 @@ for history in histories: ### Data Structure -The `History` class structure: +The legacy `LegacyHistory` payload structure: ```python @dataclass -class History: +class LegacyHistory: messages_and_choices: list[dict[str, Any]] tools: list[Tool] | None = None ``` @@ -168,7 +168,7 @@ The `Trajectory` class with additional histories: class Trajectory: messages_and_choices: list[dict[str, Any]] tools: list[Tool] | None = None - additional_histories: list[History] = field(default_factory=list) + additional_histories: list[LegacyHistory] = field(default_factory=list) reward: float | None = None metrics: dict[str, Any] = field(default_factory=dict) ``` @@ -178,7 +178,7 @@ class Trajectory: ### Creating a Trajectory with Additional Histories ```python -from art.trajectories import Trajectory, History +from art.trajectories import LegacyHistory, Trajectory # Create the main conversation main_messages = [ @@ -188,14 +188,14 @@ main_messages = [ ] # Create additional histories -history1 = History( +history1 = LegacyHistory( messages_and_choices=[ {"role": "user", "content": "First subtask"}, {"role": "assistant", "content": "Completing first subtask..."} ] ) -history2 = History( +history2 = LegacyHistory( messages_and_choices=[ {"role": "user", "content": "Second subtask"}, {"role": "assistant", "content": "Completing second subtask..."} diff --git a/src/art/trajectories/__init__.py b/src/art/trajectories/__init__.py index 87067648d..ffea450e7 100644 --- a/src/art/trajectories/__init__.py +++ b/src/art/trajectories/__init__.py @@ -323,6 +323,7 @@ class ChatCompletionsMessageSource: request_index: int | None = None choice_index: int | None = None output_index: int | None = None + generation_index: int | None = None @dataclass(frozen=True, slots=True) @@ -336,6 +337,7 @@ class ResponsesItemSource: exchange: ResponsesExchange request_index: int | None = None output_index: int | None = None + generation_index: int | None = None @dataclass(frozen=True, slots=True) diff --git a/src/art/trajectories/_history.py b/src/art/trajectories/_history.py index 112d555df..52e9feb20 100644 --- a/src/art/trajectories/_history.py +++ b/src/art/trajectories/_history.py @@ -2,7 +2,7 @@ from collections.abc import Callable, Iterable, Mapping, Sequence import copy -from dataclasses import dataclass +from dataclasses import dataclass, replace from datetime import datetime import json from typing import Generic, Protocol, TypeVar, cast @@ -95,6 +95,21 @@ class _ResponsesContext: kwargs: dict[str, object] | None +type _ResponsesRecord = tuple[ + ResponseInputParam, + list[ResponsesItemSource | None], + ResponsesConversation | None, + str | None, +] + + +@dataclass(frozen=True) +class _ResponsesGeneration: + prompt_token_ids: list[int] + output_token_ids: list[int] + output_indices: list[int] + + @dataclass class _Branch(Generic[_ItemT, _SourceT, _ContextT]): items: list[_ItemT] @@ -135,6 +150,28 @@ def _is_prefix(prefix: Sequence[object], value: Sequence[object]) -> bool: return len(prefix) <= len(value) and list(value[: len(prefix)]) == list(prefix) +def _lineage_prompt_sources( + branches: Sequence[_Branch[_ItemT, _SourceT, _ContextT]], + *, + prompt: Sequence[_ItemT], + defaults: Sequence[_SourceT | None], + equivalent: Callable[[_ItemT, _ItemT], bool], +) -> list[_SourceT | None]: + candidates = [ + branch + for branch in branches + if len(branch.items) < len(prompt) + and all( + equivalent(existing, current) + for existing, current in zip(branch.items, prompt, strict=False) + ) + ] + if not candidates: + return list(defaults) + parent = min(candidates, key=lambda branch: (-len(branch.items), branch.order)) + return [*parent.sources, *defaults[len(parent.items) :]] + + def _extend_branches( branches: list[_Branch[_ItemT, _SourceT, _ContextT]], *, @@ -172,7 +209,9 @@ def _extend_branches( regenerations = [ branch for branch in branches - if branch.context == context and _is_prefix(prompt, branch.items) + if branch.context == context + and _is_prefix(prompt, branch.items) + and (continuation is None or continuation(branch)) ] if regenerations: parent = min(regenerations, key=lambda branch: branch.order) @@ -249,6 +288,14 @@ def _chat_retains_sampled_reasoning( return True +def _chat_message_key(message: Message, *, visible_only: bool = False) -> str: + data = copy.deepcopy(dict(message)) + if visible_only: + data.pop("reasoning", None) + data.pop("reasoning_content", None) + return json.dumps(data, sort_keys=True, default=str) + + def _require_unmixed(trajectory: Trajectory) -> None: if trajectory.exchanges and ( trajectory.messages_and_choices @@ -297,10 +344,18 @@ def chat_completions_histories( ] = [] for sequence, exchange in enumerate(exchanges): prompt = _MESSAGES.validate_python(exchange.request.get("messages", [])) - prompt_sources = [ - ChatCompletionsMessageSource(exchange=exchange, request_index=index) - for index in range(len(prompt)) - ] + prompt_sources = _lineage_prompt_sources( + branches, + prompt=prompt, + defaults=[ + ChatCompletionsMessageSource(exchange=exchange, request_index=index) + for index in range(len(prompt)) + ], + equivalent=lambda left, right: ( + _chat_message_key(left, visible_only=True) + == _chat_message_key(right, visible_only=True) + ), + ) outputs: list[ tuple[ int, @@ -314,15 +369,14 @@ def chat_completions_histories( response = _MESSAGE.validate_python( choice.message.model_dump(mode="python", exclude_none=True) ) + response_source = ChatCompletionsMessageSource( + exchange=exchange, choice_index=choice.index + ) outputs.append( ( choice.index, [response], - [ - ChatCompletionsMessageSource( - exchange=exchange, choice_index=choice.index - ) - ], + [response_source], ) ) _extend_branches( @@ -376,7 +430,6 @@ def anthropic_messages_histories( for selected_model, exchanges in _selected_models( trajectory.exchanges.messages, model, "Anthropic Messages" ): - introduced: dict[tuple[object, ...], AnthropicMessageSource] = {} branches: list[ _Branch[AnthropicMessageParam, AnthropicMessageSource, _AnthropicContext] ] = [] @@ -385,13 +438,18 @@ def anthropic_messages_histories( copy.deepcopy(message) for message in exchange.request.get("messages", []) ] - prompt_sources = [ - introduced.get( - _anthropic_message_key(message), - AnthropicMessageSource(exchange=exchange, request_index=index), - ) - for index, message in enumerate(prompt) - ] + prompt_sources = _lineage_prompt_sources( + branches, + prompt=prompt, + defaults=[ + AnthropicMessageSource(exchange=exchange, request_index=index) + for index in range(len(prompt)) + ], + equivalent=lambda left, right: ( + _anthropic_message_key(left, visible_only=True) + == _anthropic_message_key(right, visible_only=True) + ), + ) response = cast( AnthropicMessageParam, { @@ -403,11 +461,6 @@ def anthropic_messages_histories( }, ) response_source = AnthropicMessageSource(exchange=exchange) - introduced.setdefault(_anthropic_message_key(response), response_source) - introduced.setdefault( - _anthropic_message_key(response, visible_only=True), - response_source, - ) _extend_branches( branches, prompt=prompt, @@ -434,7 +487,7 @@ def anthropic_messages_histories( sequence=sequence, start_time=exchange.start_time, ) - for branch in branches: + for branch in sorted(branches, key=lambda item: (item.first_time, item.order)): first_source = next( (source for source in branch.sources if source is not None), None ) @@ -521,9 +574,8 @@ def responses_histories( branches: list[ _Branch[ResponseInputItemParam, ResponsesItemSource, _ResponsesContext] ] = [] - responses: dict[ - str, tuple[ResponseInputParam, list[ResponsesItemSource | None]] - ] = {} + responses: dict[str, _ResponsesRecord] = {} + conversations: dict[str, _ResponsesRecord] = {} for sequence, exchange in enumerate(exchanges): request = exchange.request prompt = _responses_input(request.get("input")) @@ -531,11 +583,36 @@ def responses_histories( ResponsesItemSource(exchange=exchange, request_index=index) for index in range(len(prompt)) ] + requested_conversation = _RESPONSE_CONVERSATION.validate_python( + request.get("conversation") + ) previous = request.get("previous_response_id") + inherited_conversation: ResponsesConversation | None = None + inherited_previous: str | None = None + prior: _ResponsesRecord | None = None if isinstance(previous, str) and previous in responses: - prior_items, prior_sources = responses[previous] + prior = responses[previous] + elif requested_conversation is not None: + prior = conversations.get( + json.dumps(requested_conversation, sort_keys=True, default=str) + ) + if prior is not None: + ( + prior_items, + prior_sources, + inherited_conversation, + inherited_previous, + ) = prior prompt = [*copy.deepcopy(prior_items), *prompt] prompt_sources = [*copy.copy(prior_sources), *prompt_sources] + prompt_sources = _lineage_prompt_sources( + branches, + prompt=prompt, + defaults=prompt_sources, + equivalent=lambda left, right: ( + _response_item_key(left) == _response_item_key(right) + ), + ) output = [ cast( ResponseInputItemParam, @@ -543,8 +620,15 @@ def responses_histories( ) for item in exchange.response.output ] + generations, generation_indices = _responses_generations( + exchange, output=output + ) output_sources: list[ResponsesItemSource | None] = [ - ResponsesItemSource(exchange=exchange, output_index=index) + ResponsesItemSource( + exchange=exchange, + output_index=index, + generation_index=generation_indices[index], + ) for index in range(len(output)) ] instructions = request.get("instructions") @@ -553,33 +637,132 @@ def responses_histories( external_previous = ( previous if isinstance(previous, str) and previous not in responses - else None + else inherited_previous ) - _extend_branches( - branches, - prompt=prompt, - prompt_sources=prompt_sources, - outputs=[(0, output, output_sources)], - context=_ResponsesContext( - instructions=instructions, - tools=_RESPONSE_TOOLS.validate_python(request.get("tools")), - conversation=_RESPONSE_CONVERSATION.validate_python( - request.get("conversation") - ), - previous_response_id=external_previous, - template=request.get("chat_template"), - kwargs=_CHAT_KWARGS.validate_python( - request.get("chat_template_kwargs") - ), + conversation = ( + requested_conversation + if requested_conversation is not None + else inherited_conversation + ) + context = _ResponsesContext( + instructions=instructions, + tools=_RESPONSE_TOOLS.validate_python(request.get("tools")), + conversation=conversation, + previous_response_id=external_previous, + template=request.get("chat_template"), + kwargs=_CHAT_KWARGS.validate_python( + request.get("chat_template_kwargs") ), - sequence=sequence, - start_time=exchange.start_time, ) - responses[exchange.response.id] = ( - [*copy.deepcopy(prompt), *copy.deepcopy(output)], - [*copy.copy(prompt_sources), *copy.copy(output_sources)], + final_items = [*copy.deepcopy(prompt), *copy.deepcopy(output)] + final_sources = [*copy.copy(prompt_sources), *copy.copy(output_sources)] + if generations: + for generation_index, generation in enumerate(generations): + if not generation.output_indices: + raise ValueError( + "A Responses token generation without native output " + "items cannot yet be projected as a history" + ) + output_start = generation.output_indices[0] + output_end = generation.output_indices[-1] + 1 + if generation_index == len(generations) - 1: + output_end = len(output) + generation_prompt = [ + *copy.deepcopy(prompt), + *copy.deepcopy(output[:output_start]), + ] + generation_prompt_sources = [ + *copy.copy(prompt_sources), + *copy.copy(output_sources[:output_start]), + ] + continuation = lambda branch, generation=generation: ( + _responses_generation_extends( + branch, + exchange=exchange, + generation=generation, + generations=generations, + ) + ) + extends = any( + branch.context == context + and _is_prefix(branch.items, generation_prompt) + and continuation(branch) + for branch in branches + ) + if not extends and generation_index: + retained = [ + (item, source) + for item, source in zip( + generation_prompt, + generation_prompt_sources, + strict=True, + ) + if not ( + item.get("type") == "reasoning" + and source is not None + and source.exchange is exchange + and source.output_index is not None + ) + ] + generation_prompt = [item for item, _ in retained] + generation_prompt_sources = [ + ( + replace(source, generation_index=None) + if source is not None + and source.exchange is exchange + and source.generation_index is not None + else source + ) + for _, source in retained + ] + generation_output = output[output_start:output_end] + generation_output_sources = output_sources[output_start:output_end] + _extend_branches( + branches, + prompt=generation_prompt, + prompt_sources=generation_prompt_sources, + outputs=[ + ( + generation_index, + generation_output, + generation_output_sources, + ) + ], + context=context, + sequence=sequence, + start_time=exchange.start_time, + continuation=continuation, + ) + final_items = [ + *copy.deepcopy(generation_prompt), + *copy.deepcopy(generation_output), + ] + final_sources = [ + *copy.copy(generation_prompt_sources), + *copy.copy(generation_output_sources), + ] + else: + _extend_branches( + branches, + prompt=prompt, + prompt_sources=prompt_sources, + outputs=[(0, output, output_sources)], + context=context, + sequence=sequence, + start_time=exchange.start_time, + ) + record = ( + final_items, + final_sources, + copy.deepcopy(conversation), + external_previous, ) - for branch in branches: + responses[exchange.response.id] = record + if conversation is not None: + conversations[json.dumps(conversation, sort_keys=True, default=str)] = ( + record + ) + for branch in sorted(branches, key=lambda item: (item.first_time, item.order)): first_source = next( (source for source in branch.sources if source is not None), None ) @@ -602,6 +785,64 @@ def responses_histories( return histories +def _response_item_key(item: ResponseInputItemParam) -> str: + return json.dumps(item, sort_keys=True, default=str) + + +def _responses_generations( + exchange: ResponsesExchange, *, output: Sequence[ResponseInputItemParam] +) -> tuple[list[_ResponsesGeneration], list[int | None]]: + from ._tokenize import _response_generations + + generations = _response_generations(exchange.response) + result: list[int | None] = [None] * len(output) + parsed: list[_ResponsesGeneration] = [] + for generation_index, generation in enumerate(generations): + if generation.prompt_token_ids is None or generation.output_token_ids is None: + raise AssertionError("Validated Responses generation omitted exact tokens") + for output_index in generation.output_indices: + result[output_index] = generation_index + parsed.append( + _ResponsesGeneration( + prompt_token_ids=generation.prompt_token_ids, + output_token_ids=generation.output_token_ids, + output_indices=generation.output_indices, + ) + ) + return parsed, result + + +def _responses_generation_extends( + branch: _Branch[ResponseInputItemParam, ResponsesItemSource, _ResponsesContext], + *, + exchange: ResponsesExchange, + generation: _ResponsesGeneration, + generations: Sequence[_ResponsesGeneration], +) -> bool: + prior_generation = next( + ( + source.generation_index + for source in reversed(branch.sources) + if source is not None + and source.exchange is exchange + and source.generation_index is not None + ), + None, + ) + if prior_generation is None: + return True + prior = generations[prior_generation] + if _is_prefix( + [*prior.prompt_token_ids, *prior.output_token_ids], + generation.prompt_token_ids, + ): + return True + return not any( + getattr(exchange.response.output[index], "type", None) == "reasoning" + for index in prior.output_indices + ) + + def responses_history(trajectory: Trajectory, *, model: str | None) -> ResponsesHistory: histories = responses_histories(trajectory, model=model) if model is None and len({history.model for history in histories}) > 1: @@ -643,22 +884,26 @@ def completions_token_histories( trajectory.exchanges.completions, model, "Completions" ): branches: list[_Branch[int, CompletionsSource, tuple[()]]] = [] + complete = True for sequence, exchange in enumerate(exchanges): if exchange.request.get("suffix") is not None: raise ValueError("Completions suffix is not supported") + choice_groups = _completion_choice_groups(exchange) for prompt_index, request_prompt in enumerate( _completion_prompts(exchange.request.get("prompt")) ): - choices = _completion_choices_for_prompt(exchange, prompt_index) + choices = choice_groups[prompt_index] for choice in choices: prompt, completion = _completion_exact_tokens( exchange, choice.index ) if prompt is None: if not isinstance(request_prompt, list): + complete = False continue prompt = copy.deepcopy(request_prompt) if completion is None: + complete = False continue prompt_source = CompletionsSource( exchange=exchange, prompt_index=prompt_index @@ -683,6 +928,10 @@ def completions_token_histories( sequence=sequence, start_time=exchange.start_time, ) + if not complete: + raise ValueError( + "Completions token history requires exact token IDs for every choice" + ) for branch in branches: histories.append( CompletionsTokenHistory( @@ -729,28 +978,75 @@ def _completion_prompts(value: object) -> list[str | list[int]]: def _checked_token_ids(value: Sequence[object]) -> list[int]: result: list[int] = [] for token in value: - if not isinstance(token, int) or isinstance(token, bool): - raise ValueError("Completions token prompts must contain integers") + if not isinstance(token, int) or isinstance(token, bool) or token < 0: + raise ValueError( + "Completions token prompts must contain non-negative integers" + ) result.append(token) return result -def _completion_choices_for_prompt( - exchange: CompletionsExchange, prompt_index: int -) -> list[CompletionChoice]: - count = len(_completion_prompts(exchange.request.get("prompt"))) +def _completion_choice_groups( + exchange: CompletionsExchange, +) -> list[list[CompletionChoice]]: + prompts = _completion_prompts(exchange.request.get("prompt")) + count = len(prompts) choices = sorted(exchange.response.choices, key=lambda item: item.index) + if not choices: + raise ValueError("Completions response contains no choices") if count == 1: - return choices + return [choices] + requested_n = exchange.request.get("n") + if requested_n is not None and ( + not isinstance(requested_n, int) + or isinstance(requested_n, bool) + or requested_n < 1 + ): + raise ValueError("Completions n must be a positive integer") + + exact_matches: dict[int, int] = {} + for choice in choices: + prompt_ids, _ = _completion_exact_tokens(exchange, choice.index) + if prompt_ids is None: + continue + matches = [ + prompt_index + for prompt_index, prompt in enumerate(prompts) + if isinstance(prompt, list) and prompt == prompt_ids + ] + if len(matches) == 1: + exact_matches[choice.index] = matches[0] + + if len(exact_matches) == len(choices): + groups = [[] for _ in prompts] + for choice in choices: + groups[exact_matches[choice.index]].append(choice) + if all(groups) and ( + requested_n is None or all(len(group) == requested_n for group in groups) + ): + return groups + if len(choices) % count: raise ValueError("Cannot associate Completions choices with batched prompts") per_prompt = len(choices) // count - start = prompt_index * per_prompt - expected = list(range(start, start + per_prompt)) - selected = choices[start : start + per_prompt] - if [choice.index for choice in selected] != expected: + if requested_n is not None and requested_n != per_prompt: + raise ValueError("Cannot associate Completions choices with batched prompts") + if [choice.index for choice in choices] != list(range(len(choices))): raise ValueError("Ambiguous Completions choice-to-prompt association") - return selected + groups = [ + choices[prompt_index * per_prompt : (prompt_index + 1) * per_prompt] + for prompt_index in range(count) + ] + for prompt_index, group in enumerate(groups): + if any( + choice.index in exact_matches + and exact_matches[choice.index] != prompt_index + for choice in group + ): + raise ValueError( + "Completions exact prompt evidence contradicts choice indices" + ) + return groups def completions_string_histories( @@ -762,15 +1058,18 @@ def completions_string_histories( trajectory.exchanges.completions, model, "Completions" ): branches: list[_Branch[str, CompletionsSource, tuple[()]]] = [] + complete = True for sequence, exchange in enumerate(exchanges): if exchange.request.get("suffix") is not None: raise ValueError("Completions suffix is not supported") + choice_groups = _completion_choice_groups(exchange) for prompt_index, prompt in enumerate( _completion_prompts(exchange.request.get("prompt")) ): if not isinstance(prompt, str): + complete = False continue - for choice in _completion_choices_for_prompt(exchange, prompt_index): + for choice in choice_groups[prompt_index]: text = choice.text if exchange.request.get("echo") is True: if not text.startswith(prompt): @@ -801,6 +1100,10 @@ def completions_string_histories( sequence=sequence, start_time=exchange.start_time, ) + if not complete: + raise ValueError( + "Completions string history requires text prompts for every choice" + ) histories.extend( CompletionsStringHistory( model=selected_model, @@ -931,6 +1234,7 @@ def converted( exchange=source.exchange, request_index=source.request_index, output_index=source.output_index, + generation_index=source.generation_index, ) if source is not None else None diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index 0d2f0584b..25aaf1184 100644 --- a/src/art/trajectories/_tokenize.py +++ b/src/art/trajectories/_tokenize.py @@ -55,6 +55,13 @@ class _TokenizerConfig: chat_template_kwargs: Mapping[str, object] | None = None +@dataclass(frozen=True) +class _SampledOutput: + text: str | None + token_ids: list[int] + start: int + + def _as_tokenizer(tokenizer: object) -> Tokenizer: # Transformers' annotation permits only string-valued message dictionaries, # although its runtime API supports the structured content ART must tokenize. @@ -90,14 +97,16 @@ def _dump(value: object) -> dict[str, Any]: def _token_id(value: object) -> int | None: - if isinstance(value, int) and not isinstance(value, bool): + if isinstance(value, int) and not isinstance(value, bool) and value >= 0: return value if isinstance(value, str) and (match := _TOKEN_ID.fullmatch(value)): return int(match.group(1)) return None -def _exact_token_ids(values: object, *, field: str) -> list[int] | None: +def _exact_token_ids( + values: object, *, field: str, empty_is_missing: bool = False +) -> list[int] | None: if values is None: return None if not isinstance(values, list): @@ -108,7 +117,7 @@ def _exact_token_ids(values: object, *, field: str) -> list[int] | None: if token_id is None: raise ValueError(f"{field} contains an invalid exact token ID") token_ids.append(token_id) - return token_ids + return None if empty_is_missing and not token_ids else token_ids def _pair_token_id(data: dict[str, Any], *, required: bool, field: str) -> int | None: @@ -173,7 +182,11 @@ def _chat_choice_tokens( prompt = choice_data.get("prompt_token_ids") if prompt is None: prompt = response_data.get("prompt_token_ids") - prompt_ids = _exact_token_ids(prompt, field="Chat Completions prompt_token_ids") + prompt_ids = _exact_token_ids( + prompt, + field="Chat Completions prompt_token_ids", + empty_is_missing=True, + ) token_ids = _exact_token_ids( choice_data.get("token_ids"), field="Chat Completions token_ids", @@ -187,16 +200,18 @@ def _chat_choice_tokens( values or any(message.get(key) for key in ("content", "refusal", "tool_calls")) ): token_ids = None - pair_ids, logprobs = _pairs(values, field="Chat Completions logprobs") + pair_ids, pair_logprobs = _pairs(values, field="Chat Completions logprobs") + positional_logprobs = _logprob_values(values) if token_ids is not None and pair_ids and token_ids != pair_ids: raise ValueError("Response token IDs disagree with choice logprobs") selected = token_ids if token_ids is not None else pair_ids or None + logprobs = pair_logprobs or positional_logprobs if selected is not None and values and len(logprobs) != len(selected): raise ValueError("Chat Completions token IDs and logprobs differ in length") return ( prompt_ids, selected, - logprobs or _logprob_values(values) or [math.nan] * len(selected or []), + logprobs or [math.nan] * len(selected or []), ) @@ -212,6 +227,7 @@ def _completion_evidence( response: Completion, *, echo: bool = False, + empty_prompt_is_exact: bool = False, ) -> tuple[list[int] | None, list[int] | None, list[float], list[float]]: if len(response.choices) != 1: raise ValueError("Trajectory tokenization requires exactly one response choice") @@ -221,7 +237,11 @@ def _completion_evidence( prompt = choice_data.get("prompt_token_ids") if prompt is None: prompt = response_data.get("prompt_token_ids") - prompt_ids = _exact_token_ids(prompt, field="Completions prompt_token_ids") + prompt_ids = _exact_token_ids( + prompt, + field="Completions prompt_token_ids", + empty_is_missing=not empty_prompt_is_exact, + ) token_ids = _exact_token_ids( choice_data.get("token_ids"), field="Completions token_ids" ) @@ -244,7 +264,9 @@ def _completion_evidence( if not complete_pairs: pair_ids = [] pair_logprobs = [ - float(value) if isinstance(value, (int, float)) else math.nan + float(value) + if isinstance(value, (int, float)) and not isinstance(value, bool) + else math.nan for value in logprobs.get("token_logprobs") or [] ] if token_ids is not None and pair_ids and token_ids != pair_ids: @@ -278,28 +300,167 @@ def _completion_tokens( response: Completion, *, echo: bool = False, + empty_prompt_is_exact: bool = False, ) -> tuple[list[int] | None, list[int] | None, list[float]]: prompt, completion, _, completion_logprobs = _completion_evidence( - response, echo=echo + response, + echo=echo, + empty_prompt_is_exact=empty_prompt_is_exact, ) return prompt, completion, completion_logprobs +@dataclass(frozen=True) +class _ResponseGeneration: + prompt_token_ids: list[int] | None + output_token_ids: list[int] | None + output_logprobs: list[float] + output_indices: list[int] + output_text: str | None + + +def _response_generations(response: Response) -> list[_ResponseGeneration]: + data = _dump(response) + raw_generations = data.get("token_generations") + if raw_generations is None: + return [] + if not isinstance(raw_generations, list): + raise ValueError("Responses token_generations exact metadata must be a list") + if not raw_generations: + raise ValueError( + "Responses token_generations must be omitted when exact evidence " + "is unavailable" + ) + generations: list[_ResponseGeneration] = [] + assigned_outputs: set[int] = set() + last_output_index = -1 + for index, raw_generation in enumerate(raw_generations): + generation = _string_dict(raw_generation) + if generation is None: + raise ValueError( + f"Responses token_generations[{index}] must be a JSON object" + ) + prompt = _exact_token_ids( + generation.get("prompt_token_ids"), + field=f"Responses token_generations[{index}].prompt_token_ids", + empty_is_missing=True, + ) + if prompt is None: + raise ValueError( + f"Responses token_generations[{index}].prompt_token_ids must " + "contain the full non-empty generation prompt" + ) + raw_indices = generation.get("output_indices") + if not isinstance(raw_indices, list) or any( + not isinstance(value, int) or isinstance(value, bool) or value < 0 + for value in raw_indices + ): + raise ValueError( + f"Responses token_generations[{index}].output_indices must be " + "a list of non-negative integers" + ) + output_indices = [ + value + for value in raw_indices + if isinstance(value, int) and not isinstance(value, bool) + ] + if len(set(output_indices)) != len(output_indices): + raise ValueError( + f"Responses token_generations[{index}].output_indices contains duplicates" + ) + if output_indices != sorted(output_indices): + raise ValueError( + f"Responses token_generations[{index}].output_indices must be ordered" + ) + if output_indices and output_indices[0] <= last_output_index: + raise ValueError( + "Responses token_generations output_indices must be ordered " + "and nonoverlapping" + ) + if output_indices: + last_output_index = output_indices[-1] + if any(value >= len(response.output) for value in output_indices): + raise ValueError( + f"Responses token_generations[{index}].output_indices is out of bounds" + ) + if assigned_outputs.intersection(output_indices): + raise ValueError( + "Responses token_generations output_indices overlap between generations" + ) + assigned_outputs.update(output_indices) + output = generation.get("output_tokens") + output_ids: list[int] | None + output_logprobs: list[float] + if output is None: + output_ids, output_logprobs = None, [] + output_text = None + else: + output_ids, output_logprobs = _pairs( + output, + require_token_ids=True, + field=f"Responses token_generations[{index}].output_tokens", + ) + if not output_ids: + output_ids = None + output_logprobs = [] + output_texts = [ + item.get("text") + for item in (_string_dict(value) for value in output) + if item is not None + ] + output_text = ( + "".join(cast(str, text) for text in output_texts) + if len(output_texts) == len(output) + and all(isinstance(text, str) for text in output_texts) + else None + ) + if output_ids is None and any( + not str(getattr(response.output[value], "type", "")).endswith("_output") + for value in output_indices + ): + raise ValueError( + f"Responses token_generations[{index}].output_tokens must contain " + "exact tokens for its sampled output items" + ) + generations.append( + _ResponseGeneration( + prompt_token_ids=prompt, + output_token_ids=output_ids, + output_logprobs=output_logprobs, + output_indices=output_indices, + output_text=output_text, + ) + ) + required_outputs = { + index + for index, item in enumerate(response.output) + if not str(getattr(item, "type", "")).endswith("_output") + } + if not required_outputs.issubset(assigned_outputs): + raise ValueError( + "Responses token_generations output_indices must cover every sampled " + "output item" + ) + return generations + + def _responses_tokens( response: Response, ) -> tuple[list[int] | None, list[int] | None, list[float]]: data = _dump(response) - prompt_ids = _exact_token_ids( - data.get("prompt_token_ids"), field="Responses prompt_token_ids" - ) - if "raw_output_tokens" in data: - token_ids, logprobs = _pairs( - data["raw_output_tokens"], - require_token_ids=True, - field="Responses raw_output_tokens", + generations = _response_generations(response) + if generations: + if len(generations) != 1: + raise ValueError( + "A multi-generation Responses exchange must be tokenized through " + "its protocol history" + ) + generation = generations[0] + return ( + generation.prompt_token_ids, + generation.output_token_ids, + generation.output_logprobs, ) - if token_ids or not data.get("output"): - return prompt_ids, token_ids, logprobs token_ids: list[int] = [] logprobs: list[float] = [] saw_rendered_output = False @@ -324,8 +485,8 @@ def _responses_tokens( token_ids.extend(pair_ids) logprobs.extend(pair_logprobs) if saw_rendered_output and complete: - return prompt_ids, token_ids, logprobs - return prompt_ids, None, [] + return None, token_ids, logprobs + return None, None, [] def _messages_tokens( @@ -333,13 +494,17 @@ def _messages_tokens( ) -> tuple[list[int] | None, list[int] | None, list[float]]: data = _dump(response) prompt_ids = _exact_token_ids( - data.get("prompt_token_ids"), field="Messages prompt_token_ids" + data.get("prompt_token_ids"), + field="Messages prompt_token_ids", + empty_is_missing=True, ) token_ids = _exact_token_ids(data.get("token_ids"), field="Messages token_ids") if token_ids == [] and data.get("content"): token_ids = None logprobs = [ - float(value) if isinstance(value, (int, float)) else math.nan + float(value) + if isinstance(value, (int, float)) and not isinstance(value, bool) + else math.nan for value in data.get("logprobs") or [] ] if token_ids is not None and logprobs and len(logprobs) != len(token_ids): @@ -372,6 +537,31 @@ def _exchange_list(trajectory: Trajectory, model: str | None) -> list[Exchange]: ) +def _require_training_model(trajectory: Trajectory, model: str | None) -> str: + """Resolve one exchange model without guessing which policy should train.""" + + models = { + exchange.model + for exchange in ( + *trajectory.exchanges.chat_completions, + *trajectory.exchanges.completions, + *trajectory.exchanges.responses, + *trajectory.exchanges.messages, + ) + } + if None in models: + raise ValueError("Every training exchange must identify its model") + if model is not None: + if model not in models: + raise ValueError(f"Trajectory contains no exchanges for model {model!r}") + return model + if len(models) != 1: + raise ValueError( + "Exchange training requires exactly one model; pass model= to select one" + ) + return cast(str, next(iter(models))) + + def _artifact_name(model: str) -> str: return model.removeprefix("wandb-artifact:///") @@ -896,7 +1086,9 @@ def _exchange_tokens( return _chat_tokens(exchange.response) if isinstance(exchange, CompletionsExchange): return _completion_tokens( - exchange.response, echo=exchange.request.get("echo") is True + exchange.response, + echo=exchange.request.get("echo") is True, + empty_prompt_is_exact=exchange.request.get("prompt") in ("", []), ) if isinstance(exchange, ResponsesExchange): return _responses_tokens(exchange.response) @@ -969,37 +1161,30 @@ def _sampled_text(exchange: Exchange) -> str | None: return None -def _replace_unique( - values: list[int], old: list[int], new: list[int] -) -> list[int] | None: - if not old: - return None - starts = [ - index - for index in range(len(values) - len(old) + 1) - if values[index : index + len(old)] == old - ] - if len(starts) != 1: - return None - start = starts[0] - return [*values[:start], *new, *values[start + len(old) :]] - - def _preserve_sampled_prefix( prompt: list[int], canonical_prefix: list[int], - sampled_outputs: list[tuple[str, list[int]]], + sampled_outputs: list[_SampledOutput], tokenizer: Tokenizer, ) -> list[int] | None: repaired = prompt - for text, exact_ids in sampled_outputs: - rendered_ids = _ids(tokenizer(text, add_special_tokens=False)) - if rendered_ids == exact_ids: + for sampled in sampled_outputs: + if sampled.text is None: + return None + rendered_ids = _ids(tokenizer(sampled.text, add_special_tokens=False)) + if rendered_ids == sampled.token_ids: continue - replaced = _replace_unique(repaired, rendered_ids, exact_ids) - if replaced is None: + start = sampled.start + if ( + repaired[:start] != canonical_prefix[:start] + or repaired[start : start + len(rendered_ids)] != rendered_ids + ): return None - repaired = replaced + repaired = [ + *repaired[:start], + *sampled.token_ids, + *repaired[start + len(rendered_ids) :], + ] return repaired if repaired[: len(canonical_prefix)] == canonical_prefix else None @@ -1025,6 +1210,8 @@ def _align_visible_logprobs( token_ids: list[int] = [] logprobs: list[float] = [] for text, logprob in values: + if not text: + return None encoded = _ids(tokenizer(text, add_special_tokens=False)) if len(encoded) != 1: return None @@ -1154,7 +1341,7 @@ def _tokenize_exchange_trajectory( response_histories: dict[ str, tuple[list[dict[str, Any]] | None, ResponsesExchange] ] = {} - sampled_outputs: list[tuple[str, list[int]]] = [] + sampled_outputs: list[_SampledOutput] = [] previous_render_state: tuple[Exchange, list[dict[str, Any]] | None] | None = None def fallback_config() -> _TokenizerConfig: @@ -1234,6 +1421,19 @@ def fallback_config() -> _TokenizerConfig: resolved_config = fallback_config() if tokenizer is None: tokenizer = _load_tokenizer(resolved_config) + rendered_prompt = ( + _template_ids( + tokenizer, + exchange, + completed=False, + config=resolved_config, + chat_template=chat_template, + chat_template_kwargs=chat_template_kwargs, + messages_override=messages_override, + ) + if prompt_is_exact + else prompt + ) completed = _template_ids( tokenizer, exchange, @@ -1243,11 +1443,11 @@ def fallback_config() -> _TokenizerConfig: chat_template_kwargs=chat_template_kwargs, messages_override=messages_override, ) - if completed[: len(prompt)] != prompt: + if completed[: len(rendered_prompt)] != rendered_prompt: raise ValueError( "Completed response does not extend its generation prompt" ) - completion = completed[len(prompt) :] + completion = completed[len(rendered_prompt) :] completion_logprobs = _align_visible_logprobs( tokenizer, completion, exchange ) or [math.nan] * len(completion) @@ -1328,8 +1528,14 @@ def fallback_config() -> _TokenizerConfig: if completion_is_exact: completion_flag |= TokenFlag.EXACT flags.extend([completion_flag] * len(completion)) - if completion_is_exact and (text := _sampled_text(exchange)) is not None: - sampled_outputs.append((text, list(completion))) + if completion_is_exact: + sampled_outputs.append( + _SampledOutput( + text=_sampled_text(exchange), + token_ids=list(completion), + start=len(token_ids) - len(completion), + ) + ) previous_render_state = (exchange, messages_override) return TokenizedHistory( @@ -1400,6 +1606,57 @@ def _unique_exchanges(history: History) -> list[Exchange]: return sorted(exchanges, key=lambda item: (item.start_time, item.end_time)) +def _without_reasoning(message: Mapping[str, object]) -> dict[str, object]: + visible = dict(message) + visible.pop("reasoning", None) + visible.pop("reasoning_content", None) + return visible + + +def _chat_choice_message(source: object) -> dict[str, Any] | None: + exchange = getattr(source, "exchange", None) + choice_index = getattr(source, "choice_index", None) + if not isinstance(exchange, ChatCompletionsExchange) or not isinstance( + choice_index, int + ): + return None + choice = next( + (item for item in exchange.response.choices if item.index == choice_index), + None, + ) + return ( + choice.message.model_dump(mode="python", exclude_none=True) + if choice is not None + else None + ) + + +def _source_index(value: object, *, length: int, field: str) -> int: + if not isinstance(value, int) or isinstance(value, bool) or not 0 <= value < length: + raise ValueError(f"{field} is out of bounds") + return value + + +def _chat_choice(source: object) -> Choice: + exchange = getattr(source, "exchange", None) + choice_index = getattr(source, "choice_index", None) + if not isinstance(exchange, ChatCompletionsExchange): + raise ValueError("Chat choice source has the wrong exchange type") + if ( + not isinstance(choice_index, int) + or isinstance(choice_index, bool) + or choice_index < 0 + ): + raise ValueError("Chat choice source index is invalid") + choice = next( + (item for item in exchange.response.choices if item.index == choice_index), + None, + ) + if choice is None: + raise ValueError("Chat choice source index is out of bounds") + return choice + + def _validate_history_sources(history: History) -> None: if isinstance(history, ChatCompletionsHistory): if len(history.messages) != len(history.message_sources): @@ -1413,39 +1670,161 @@ def _validate_history_sources(history: History) -> None: expected: list[dict[str, Any]] = [] if isinstance(exchange, ChatCompletionsExchange): if source.choice_index is not None: - choice = next( - ( - item - for item in exchange.response.choices - if item.index == source.choice_index - ), - None, + if ( + source.request_index is not None + or source.output_index is not None + or source.generation_index is not None + ): + raise ValueError("Chat source has conflicting indices") + choice_message = _chat_choice(source).message.model_dump( + mode="python", exclude_none=True ) - if choice is not None: - expected.append( - choice.message.model_dump(mode="python", exclude_none=True) - ) + expected.append(choice_message) + visible = _without_reasoning(choice_message) + if visible != choice_message: + expected.append(visible) elif source.request_index is not None: + if ( + source.output_index is not None + or source.generation_index is not None + ): + raise ValueError("Chat source has conflicting indices") request_messages = exchange.request.get("messages", []) - if source.request_index < len(request_messages): - expected.append(dict(request_messages[source.request_index])) + request_index = _source_index( + source.request_index, + length=len(request_messages), + field="Chat request source index", + ) + expected.append(dict(request_messages[request_index])) + else: + raise ValueError("Chat source has no request or choice index") elif isinstance(exchange, MessagesExchange): - expected.extend(_anthropic_messages(_dump(exchange.request))) - expected.append(_response_message(exchange)) - visible_blocks = [ - block.model_dump(mode="python", exclude_none=True) - for block in exchange.response.content - if getattr(block, "type", None) - not in {"thinking", "redacted_thinking"} - ] - expected.extend( - _anthropic_messages( - {"messages": [{"role": "assistant", "content": visible_blocks}]} + if ( + source.choice_index is not None + or source.output_index is not None + or source.generation_index is not None + ): + raise ValueError("Anthropic-to-Chat source has invalid indices") + if source.request_index is not None: + request_messages = exchange.request.get("messages", []) + request_index = _source_index( + source.request_index, + length=len(request_messages), + field="Anthropic request source index", + ) + expected.extend( + _anthropic_messages( + {"messages": [request_messages[request_index]]} + ) + ) + else: + if message.get("role") == "system": + expected.extend( + _anthropic_messages( + { + "system": exchange.request.get("system"), + "messages": [], + } + ) + ) + expected.append(_response_message(exchange)) + visible_blocks = [ + block.model_dump(mode="python", exclude_none=True) + for block in exchange.response.content + if getattr(block, "type", None) + not in {"thinking", "redacted_thinking"} + ] + expected.extend( + _anthropic_messages( + { + "messages": [ + {"role": "assistant", "content": visible_blocks} + ] + } + ) ) - ) elif isinstance(exchange, ResponsesExchange): - expected.extend(_responses_messages(_dump(exchange.request))) - expected.append(_response_message(exchange)) + if source.request_index is not None: + if ( + source.choice_index is not None + or source.output_index is not None + or source.generation_index is not None + ): + raise ValueError("Responses-to-Chat source has invalid indices") + request_input = exchange.request.get("input") + if isinstance(request_input, list): + request_index = _source_index( + source.request_index, + length=len(request_input), + field="Responses request source index", + ) + expected.extend( + _responses_messages( + {"input": [request_input[request_index]]} + ) + ) + elif isinstance(request_input, str): + _source_index( + source.request_index, + length=1, + field="Responses request source index", + ) + expected.extend(_responses_messages({"input": request_input})) + else: + raise ValueError( + "Responses request source index is out of bounds" + ) + else: + if source.choice_index is not None: + raise ValueError("Responses-to-Chat source has invalid indices") + if source.output_index is not None: + output_index = _source_index( + source.output_index, + length=len(exchange.response.output), + field="Responses output source index", + ) + if source.generation_index is not None: + generations = _response_generations(exchange.response) + generation_index = _source_index( + source.generation_index, + length=len(generations), + field="Responses generation source index", + ) + if ( + output_index + not in generations[generation_index].output_indices + ): + raise ValueError( + "Responses output source does not belong to " + "its generation" + ) + expected.extend( + _responses_messages( + { + "input": [ + exchange.response.output[ + output_index + ].model_dump(mode="python", exclude_none=True) + ] + } + ) + ) + elif message.get("role") == "system": + expected.extend( + _responses_messages( + {"instructions": exchange.request.get("instructions")} + ) + ) + expected.extend( + _responses_messages( + { + "input": [ + item.model_dump(mode="python", exclude_none=True) + for item in exchange.response.output + ] + } + ) + ) if not any(dict(message) == candidate for candidate in expected): raise ValueError( "Chat Completions history no longer matches its source exchange" @@ -1471,11 +1850,12 @@ def _validate_history_sources(history: History) -> None: } else: request_messages = source.exchange.request.get("messages", []) - expected = ( - request_messages[source.request_index] - if source.request_index < len(request_messages) - else None + request_index = _source_index( + source.request_index, + length=len(request_messages), + field="Anthropic request source index", ) + expected = request_messages[request_index] matches = expected is not None and message == expected if not matches and source.request_index is None and expected is not None: matches = _anthropic_message_key(message) == _anthropic_message_key( @@ -1495,24 +1875,41 @@ def _validate_history_sources(history: History) -> None: for item, source in zip(history.input, history.input_sources, strict=True): if source is None: continue + if source.request_index is not None and source.output_index is not None: + raise ValueError("Responses source has conflicting indices") if source.output_index is not None: output = source.exchange.response.output - expected = ( - output[source.output_index].model_dump( - mode="json", exclude_none=True - ) - if source.output_index < len(output) - else None + output_index = _source_index( + source.output_index, + length=len(output), + field="Responses output source index", + ) + expected = output[output_index].model_dump( + mode="json", exclude_none=True ) + if source.generation_index is not None: + generations = _response_generations(source.exchange.response) + generation_index = _source_index( + source.generation_index, + length=len(generations), + field="Responses generation source index", + ) + if output_index not in generations[generation_index].output_indices: + raise ValueError( + "Responses output source does not belong to its generation" + ) elif source.request_index is not None: request_input = _responses_input(source.exchange.request.get("input")) - expected = ( - request_input[source.request_index] - if source.request_index < len(request_input) - else None + request_index = _source_index( + source.request_index, + length=len(request_input), + field="Responses request source index", ) + if source.generation_index is not None: + raise ValueError("Responses request source has a generation index") + expected = request_input[request_index] else: - expected = None + raise ValueError("Responses source has no request or output index") if expected is None or item != expected: raise ValueError( "Responses history no longer matches its source exchange" @@ -1564,15 +1961,25 @@ def _history_needs_render(history: History) -> bool: if isinstance(history, ChatCompletionsHistory): if any(source is None for source in history.message_sources): return True + for message, source in zip( + history.messages, history.message_sources, strict=True + ): + if source is None or (original := _chat_choice_message(source)) is None: + continue + if dict(message) != original and dict(message) == _without_reasoning( + original + ): + return True exchange = _last_source_exchange(history.message_sources) if not isinstance(exchange, ChatCompletionsExchange): - return False - return ( + return True + context_changed = ( history.tools != exchange.request.get("tools") or history.chat_template != exchange.request.get("chat_template") or history.chat_template_kwargs != exchange.request.get("chat_template_kwargs") ) + return context_changed or not _history_matches_projection(history) if isinstance(history, AnthropicMessagesHistory): if any(source is None for source in history.message_sources): return True @@ -1593,29 +2000,228 @@ def _history_needs_render(history: History) -> bool: exchange = _last_source_exchange(history.message_sources) if not isinstance(exchange, MessagesExchange): return False - return ( + context_changed = ( history.system != exchange.request.get("system") or history.tools != exchange.request.get("tools") or history.chat_template != exchange.request.get("chat_template") or history.chat_template_kwargs != exchange.request.get("chat_template_kwargs") ) + return context_changed or not _history_matches_projection(history) if isinstance(history, ResponsesHistory): if any(source is None for source in history.input_sources): return True exchange = _last_source_exchange(history.input_sources) if not isinstance(exchange, ResponsesExchange): return False - return ( + context_changed = ( history.instructions != exchange.request.get("instructions") or history.tools != exchange.request.get("tools") or history.chat_template != exchange.request.get("chat_template") or history.chat_template_kwargs != exchange.request.get("chat_template_kwargs") ) + return context_changed or not _history_matches_projection(history) + return False + + +def _source_signature(source: object) -> tuple[object, ...] | None: + if source is None: + return None + exchange = getattr(source, "exchange", None) + if not isinstance( + exchange, + ( + ChatCompletionsExchange, + CompletionsExchange, + ResponsesExchange, + MessagesExchange, + ), + ): + return None + response_id = getattr(exchange.response, "id", None) + return ( + type(source), + type(exchange), + exchange.start_time, + exchange.end_time, + response_id, + getattr(source, "request_index", None), + getattr(source, "choice_index", None), + getattr(source, "output_index", None), + getattr(source, "generation_index", None), + getattr(source, "prompt_index", None), + ) + + +def _sources_match(left: Sequence[object], right: Sequence[object]) -> bool: + return [_source_signature(item) for item in left] == [ + _source_signature(item) for item in right + ] + + +def _history_matches_projection(history: History) -> bool: + exchanges = _unique_exchanges(history) + if not exchanges: + return False + from . import TrajectoryExchanges + from ._history import ( + anthropic_messages_histories, + chat_completions_histories, + responses_histories, + ) + + trajectory = Trajectory( + exchanges=TrajectoryExchanges( + chat_completions=[ + item for item in exchanges if isinstance(item, ChatCompletionsExchange) + ], + completions=[ + item for item in exchanges if isinstance(item, CompletionsExchange) + ], + responses=[ + item for item in exchanges if isinstance(item, ResponsesExchange) + ], + messages=[item for item in exchanges if isinstance(item, MessagesExchange)], + ) + ) + if isinstance(history, ChatCompletionsHistory): + candidates: Sequence[History] = chat_completions_histories( + trajectory, model=history.model + ) + return any( + isinstance(candidate, ChatCompletionsHistory) + and candidate.messages == history.messages + and candidate.tools == history.tools + and candidate.chat_template == history.chat_template + and candidate.chat_template_kwargs == history.chat_template_kwargs + and _sources_match(candidate.message_sources, history.message_sources) + for candidate in candidates + ) + if isinstance(history, AnthropicMessagesHistory): + candidates = anthropic_messages_histories(trajectory, model=history.model) + return any( + isinstance(candidate, AnthropicMessagesHistory) + and candidate.messages == history.messages + and candidate.system == history.system + and candidate.tools == history.tools + and candidate.chat_template == history.chat_template + and candidate.chat_template_kwargs == history.chat_template_kwargs + and _sources_match(candidate.message_sources, history.message_sources) + for candidate in candidates + ) + if isinstance(history, ResponsesHistory): + candidates = responses_histories(trajectory, model=history.model) + return any( + isinstance(candidate, ResponsesHistory) + and candidate.input == history.input + and candidate.instructions == history.instructions + and candidate.tools == history.tools + and candidate.conversation == history.conversation + and candidate.previous_response_id == history.previous_response_id + and candidate.chat_template == history.chat_template + and candidate.chat_template_kwargs == history.chat_template_kwargs + and _sources_match(candidate.input_sources, history.input_sources) + for candidate in candidates + ) return False +def _response_generation_text( + response: Response, generation: _ResponseGeneration +) -> str | None: + parts: list[str] = [] + for output_index in generation.output_indices: + item = _dump(response.output[output_index]) + kind = item.get("type") + if kind == "message": + parts.append(_responses_output_text(item.get("content"))) + elif kind == "reasoning": + parts.append(_responses_reasoning_text(item)) + else: + return None + return "".join(parts) or None + + +def _tokenize_exact_responses_history( + history: ResponsesHistory, + *, + base_model: str | None, + tokenizer: Tokenizer | None, +) -> TokenizedHistory | None: + generation_keys: list[tuple[ResponsesExchange, int]] = [] + seen: set[tuple[int, int]] = set() + for source in history.input_sources: + if source is None or source.generation_index is None: + continue + key = (id(source.exchange), source.generation_index) + if key not in seen: + seen.add(key) + generation_keys.append((source.exchange, source.generation_index)) + if not generation_keys: + return None + + token_ids: list[int] = [] + logprobs: list[float] = [] + flags: list[TokenFlag] = [] + sampled_outputs: list[_SampledOutput] = [] + for exchange, generation_index in generation_keys: + generations = _response_generations(exchange.response) + if not 0 <= generation_index < len(generations): + raise ValueError("Responses source generation index is out of bounds") + generation = generations[generation_index] + prompt = generation.prompt_token_ids + output = generation.output_token_ids + if prompt is None or output is None: + return None + if not token_ids: + token_ids.extend(prompt) + logprobs.extend([math.nan] * len(prompt)) + flags.extend([TokenFlag.EXACT] * len(prompt)) + elif prompt[: len(token_ids)] == token_ids: + suffix = prompt[len(token_ids) :] + token_ids.extend(suffix) + logprobs.extend([math.nan] * len(suffix)) + flags = [flag | TokenFlag.EXACT for flag in flags] + flags.extend([TokenFlag.EXACT] * len(suffix)) + else: + if tokenizer is None and sampled_outputs: + tokenizer = _load_tokenizer( + _tokenizer_config(history.model, base_model) + ) + repaired = ( + _preserve_sampled_prefix(prompt, token_ids, sampled_outputs, tokenizer) + if tokenizer is not None + else None + ) + if repaired is None: + raise ValueError( + "Responses token generations do not form one append-only history" + ) + _warn_prefix_retokenization() + suffix = repaired[len(token_ids) :] + token_ids.extend(suffix) + logprobs.extend([math.nan] * len(suffix)) + flags.extend([TokenFlag.EXACT] * len(suffix)) + token_ids.extend(output) + logprobs.extend(generation.output_logprobs) + flags.extend([TokenFlag.EXACT | TokenFlag.SAMPLED] * len(output)) + sampled_outputs.append( + _SampledOutput( + text=generation.output_text + or _response_generation_text(exchange.response, generation), + token_ids=list(output), + start=len(token_ids) - len(output), + ) + ) + return TokenizedHistory( + model=history.model, + token_ids=token_ids, + logprobs=logprobs, + flags=flags, + ) + + def _chat_source_full_tokens( source: object, ) -> tuple[list[int] | None, list[float]]: @@ -1629,12 +2235,68 @@ def _chat_source_full_tokens( ) _, tokens, logprobs = _chat_choice_tokens(choice, _dump(exchange.response)) return tokens, logprobs - if isinstance(exchange, (MessagesExchange, ResponsesExchange)): + if isinstance(exchange, ResponsesExchange): + generation_index = getattr(source, "generation_index", None) + generations = _response_generations(exchange.response) + if isinstance(generation_index, int): + if not 0 <= generation_index < len(generations): + raise ValueError("Responses source generation index is out of bounds") + generation = generations[generation_index] + output_index = getattr(source, "output_index", None) + if ( + isinstance(output_index, int) + and output_index not in generation.output_indices + ): + raise ValueError( + "Responses source output index does not belong to its generation" + ) + return generation.output_token_ids, generation.output_logprobs + if len(generations) > 1: + return None, [] + return _responses_tokens(exchange.response)[1:] + if isinstance(exchange, MessagesExchange): _, tokens, logprobs = _exchange_tokens(exchange) return tokens, logprobs return None, [] +def _source_is_sampled(source: object) -> bool: + exchange = getattr(source, "exchange", None) + if isinstance(exchange, ChatCompletionsExchange): + return getattr(source, "choice_index", None) is not None + if isinstance(exchange, MessagesExchange): + return getattr(source, "request_index", None) is None + if isinstance(exchange, ResponsesExchange): + return ( + getattr(source, "output_index", None) is not None + or getattr(source, "generation_index", None) is not None + ) + return False + + +def _chat_message_parts(message: Mapping[str, object]) -> list[tuple[str, str]]: + parts: list[tuple[str, str]] = [] + reasoning = message.get("reasoning") or message.get("reasoning_content") + if isinstance(reasoning, str) and reasoning: + parts.append(("reasoning", reasoning)) + if content := _content_text(message.get("content")): + parts.append(("content", content)) + refusal = message.get("refusal") + if isinstance(refusal, str) and refusal: + parts.append(("content", refusal)) + tool_calls = message.get("tool_calls") + if isinstance(tool_calls, list): + for tool_call in tool_calls: + call = _string_dict(tool_call) + function = _string_dict(call.get("function")) if call is not None else None + if function is None: + continue + for value in (function.get("name"), function.get("arguments")): + if isinstance(value, str) and value: + parts.append(("tool_call", value)) + return parts + + def _chat_source_tokens( source: object, text: str, @@ -1654,8 +2316,9 @@ def _chat_source_tokens( if item.index == getattr(source, "choice_index", None) ).message ) - if sampled_text == text or ( + if ( part == "content" + and sampled_text == text and not any(message.get(key) for key in ("reasoning", "tool_calls")) ): return tokens, logprobs @@ -1700,10 +2363,22 @@ def _chat_source_tokens( return None, [] _, tokens, logprobs = _messages_tokens(exchange.response) return tokens, logprobs - if ( - isinstance(exchange, ResponsesExchange) - and getattr(source, "request_index", None) is None - ): + if isinstance(exchange, ResponsesExchange) and _source_is_sampled(source): + generation_index = getattr(source, "generation_index", None) + if isinstance(generation_index, int): + generations = _response_generations(exchange.response) + if not 0 <= generation_index < len(generations): + raise ValueError("Responses source generation index is out of bounds") + generation = generations[generation_index] + output_index = getattr(source, "output_index", None) + if ( + isinstance(output_index, int) + and output_index not in generation.output_indices + ): + raise ValueError( + "Responses source output index does not belong to its generation" + ) + return generation.output_token_ids, generation.output_logprobs output_index = getattr(source, "output_index", None) if isinstance(output_index, int) and output_index < len( exchange.response.output @@ -1779,121 +2454,72 @@ def _tokenize_chat_view( **kwargs, ) ) + + def locate(needle: Sequence[int], start: int) -> tuple[int, int] | None: + if not needle: + return None + for index in range(start, len(rendered) - len(needle) + 1): + if rendered[index : index + len(needle)] == list(needle): + return index, index + len(needle) + return None + replacements: list[tuple[int, int, list[int], list[float], bool]] = [] - for index, (message, source) in enumerate( - zip(history.messages, history.message_sources, strict=True) - ): - if message.get("role") != "assistant" or source is None: - continue - full_exact, full_logprobs = _chat_source_full_tokens(source) - content = _content_text(message.get("content")) - reasoning = message.get("reasoning") - if ( - full_exact - and content - and not reasoning - and not message.get("tool_calls") - and _ids(tokenizer(content, add_special_tokens=False)) == full_exact - ): - exact_starts = [ - start - for start in range(len(rendered) - len(full_exact) + 1) - if rendered[start : start + len(full_exact)] == full_exact - ] - if len(exact_starts) == 1: - start = exact_starts[0] - replacements.append( - ( - start, - start + len(full_exact), - full_exact, - full_logprobs - if len(full_logprobs) == len(full_exact) - else [math.nan] * len(full_exact), - True, - ) - ) - continue - stable_length = min(index + 2, len(history.messages)) - completed_prefix = ( - rendered - if stable_length == len(history.messages) - else _ids( - tokenizer.apply_chat_template( - [dict(item) for item in history.messages[:stable_length]], - tools=history.tools, - tokenize=True, - add_generation_prompt=False, - **({"chat_template": template} if template is not None else {}), - **kwargs, - ) - ) + search_cursor = 0 + sampled_texts = { + text + for message, source in zip( + history.messages, history.message_sources, strict=True ) - if rendered[: len(completed_prefix)] != completed_prefix: - raise ValueError( - "Chat template does not preserve completed message prefixes" - ) - message_end = len(completed_prefix) - if stable_length > index + 1: - next_text = _content_text(history.messages[index + 1].get("content")) - next_ids = _ids(tokenizer(next_text, add_special_tokens=False)) - if next_ids: - next_starts = [ - start - for start in range(len(completed_prefix) - len(next_ids) + 1) - if completed_prefix[start : start + len(next_ids)] == next_ids - ] - if next_starts: - message_end = next_starts[-1] - if full_exact: - exact_starts = [ - start - for start in range(message_end - len(full_exact) + 1) - if completed_prefix[start : start + len(full_exact)] == full_exact - ] - if exact_starts: - start = exact_starts[-1] - replacements.append( - ( - start, - start + len(full_exact), - full_exact, - full_logprobs - if len(full_logprobs) == len(full_exact) - else [math.nan] * len(full_exact), - True, - ) + if message.get("role") == "assistant" + and source is not None + and _source_is_sampled(source) + for _, text in _chat_message_parts(message) + } + for message, source in zip(history.messages, history.message_sources, strict=True): + parts = _chat_message_parts(message) + sampled = ( + message.get("role") == "assistant" + and source is not None + and _source_is_sampled(source) + ) + full_exact, full_logprobs = ( + _chat_source_full_tokens(source) if sampled else (None, []) + ) + if sampled and full_exact and (span := locate(full_exact, search_cursor)): + start, end = span + replacements.append( + ( + start, + end, + full_exact, + full_logprobs + if len(full_logprobs) == len(full_exact) + else [math.nan] * len(full_exact), + True, ) - continue - - parts: list[tuple[str, str]] = [] - if isinstance(reasoning, str) and reasoning: - parts.append(("reasoning", reasoning)) - if content: - parts.append(("content", content)) - if not parts and isinstance( - exchange := getattr(source, "exchange", None), - (ChatCompletionsExchange, ResponsesExchange, MessagesExchange), - ): - if sampled_text := _sampled_text(exchange): - parts.append(("raw", sampled_text)) + ) + search_cursor = end + continue - part_replacements: list[tuple[int, int, list[int], list[float], bool]] = [] - search_end = message_end - for part, text in reversed(parts): + replacement_start = len(replacements) + for part, text in parts: + if not sampled and text not in sampled_texts: + continue local = _ids(tokenizer(text, add_special_tokens=False)) if not local: continue - starts = [ - start - for start in range(search_end - len(local) + 1) - if completed_prefix[start : start + len(local)] == local - ] - if not starts: + span = locate(local, search_cursor) + if span is None: + if not sampled: + continue raise ValueError( - "Could not locate a sourced assistant message in the rendered history" + "Could not locate a history message in the rendered history" ) - start = starts[-1] + start, end = span + search_cursor = end + if not sampled: + continue + assert source is not None exact, logprobs = _chat_source_tokens(source, text, part=part) replacement = exact if exact is not None else local if exact is None and not logprobs: @@ -1903,10 +2529,10 @@ def _tokenize_chat_view( (ChatCompletionsExchange, ResponsesExchange, MessagesExchange), ): logprobs = _align_visible_logprobs(tokenizer, local, exchange) or [] - part_replacements.append( + replacements.append( ( start, - start + len(local), + end, replacement, logprobs if len(logprobs) == len(replacement) @@ -1914,8 +2540,28 @@ def _tokenize_chat_view( exact is not None, ) ) - search_end = start - replacements.extend(reversed(part_replacements)) + message_replacements = replacements[replacement_start:] + if ( + sampled + and message_replacements + and all(part == "tool_call" for part, _ in parts) + ): + start = message_replacements[0][0] + end = message_replacements[-1][1] + del replacements[replacement_start:] + replacements.append( + ( + start, + end, + rendered[start:end], + [math.nan] * (end - start), + False, + ) + ) + if sampled and not parts and full_exact is not None: + raise ValueError( + "Could not locate exact sampled output in the rendered history" + ) token_ids: list[int] = [] logprobs: list[float] = [] @@ -1977,10 +2623,12 @@ def _tokenize_completions_token_history( selected_logprobs = ( prompt_logprobs if span.source.choice_index is None else completion_logprobs ) - if ( - selected == history.prompt[span.start : span.end] - and len(selected_logprobs) == span.end - span.start - ): + if selected != history.prompt[span.start : span.end]: + flags[span.start : span.end] = [ + flag & ~TokenFlag.EXACT for flag in flags[span.start : span.end] + ] + continue + if len(selected_logprobs) == span.end - span.start: logprobs[span.start : span.end] = selected_logprobs return TokenizedHistory( model=history.model, @@ -1990,6 +2638,58 @@ def _tokenize_completions_token_history( ) +def _completion_visible_logprobs( + source: CompletionsSource, + text: str, + tokenizer: Tokenizer, + token_ids: list[int], +) -> list[float] | None: + choice = next( + item + for item in source.exchange.response.choices + if item.index == source.choice_index + ) + data = _dump(choice.logprobs) + raw_tokens = data.get("tokens") + raw_logprobs = data.get("token_logprobs") + if not isinstance(raw_tokens, list) or not isinstance(raw_logprobs, list): + return None + values = list(zip(raw_tokens, raw_logprobs, strict=False)) + request_prompt = source.exchange.request.get("prompt") + if source.exchange.request.get("echo") is True and isinstance(request_prompt, str): + consumed = "" + cursor = 0 + while cursor < len(values) and len(consumed) < len(request_prompt): + token_text = values[cursor][0] + if not isinstance(token_text, str): + return None + consumed += token_text + cursor += 1 + if consumed != request_prompt: + return None + values = values[cursor:] + if ( + "".join(token_text for token_text, _ in values if isinstance(token_text, str)) + != text + ): + return None + aligned_ids: list[int] = [] + aligned_logprobs: list[float] = [] + for token_text, logprob in values: + if ( + not isinstance(token_text, str) + or not isinstance(logprob, (int, float)) + or isinstance(logprob, bool) + ): + return None + encoded = _ids(tokenizer(token_text, add_special_tokens=False)) + if len(encoded) != 1: + return None + aligned_ids.append(encoded[0]) + aligned_logprobs.append(float(logprob)) + return aligned_logprobs if aligned_ids == token_ids else None + + def _tokenize_completions_string_history( history: CompletionsStringHistory, *, @@ -2029,8 +2729,19 @@ def resolved_tokenizer() -> Tokenizer: _completion_source_evidence(source) ) if source.choice_index is None: - exact = prompt - source_logprobs = prompt_logprobs + from ._history import _completion_prompts + + request_prompts = _completion_prompts( + source.exchange.request.get("prompt") + ) + original = request_prompts[source.prompt_index] + if not isinstance(original, str): + raise ValueError( + "A string Completions history cannot reference a token prompt" + ) + if text == original: + exact = prompt + source_logprobs = prompt_logprobs else: choice = next( item @@ -2058,7 +2769,12 @@ def resolved_tokenizer() -> Tokenizer: if exact is not None and len(source_logprobs) == len(ids): logprobs.extend(source_logprobs) else: - logprobs.extend([math.nan] * len(ids)) + visible = ( + _completion_visible_logprobs(source, text, resolved_tokenizer(), ids) + if source is not None and source.choice_index is not None + else None + ) + logprobs.extend(visible or [math.nan] * len(ids)) flag = TokenFlag.SAMPLED if is_sampled else TokenFlag(0) if exact is not None: flag |= TokenFlag.EXACT @@ -2074,15 +2790,37 @@ def resolved_tokenizer() -> Tokenizer: def _completion_source_evidence( source: CompletionsSource, ) -> tuple[list[int] | None, list[int] | None, list[float], list[float]]: - choices = source.exchange.response.choices - selected = ( - next(choice for choice in choices if choice.index == source.choice_index) - if source.choice_index is not None - else choices[0] + from ._history import _completion_choice_groups + + prompt_groups = _completion_choice_groups(source.exchange) + prompt_index = _source_index( + source.prompt_index, + length=len(prompt_groups), + field="Completions prompt source index", ) + if source.choice_index is None: + selected = prompt_groups[prompt_index][0] + else: + if ( + not isinstance(source.choice_index, int) + or isinstance(source.choice_index, bool) + or source.choice_index < 0 + ): + raise ValueError("Completions choice source index is invalid") + selected = next( + ( + choice + for choice in prompt_groups[prompt_index] + if choice.index == source.choice_index + ), + None, + ) + if selected is None: + raise ValueError("Completions choice source does not belong to its prompt") return _completion_evidence( source.exchange.response.model_copy(update={"choices": [selected]}), echo=source.exchange.request.get("echo") is True, + empty_prompt_is_exact=source.exchange.request.get("prompt") in ("", []), ) @@ -2126,6 +2864,7 @@ def tokenize_history( base_model=base_model, tokenizer=tokenizer, ) + _validate_history_sources(history) override_requires_render = ( chat_template is not None and chat_template != getattr(history, "chat_template", None) @@ -2134,9 +2873,13 @@ def tokenize_history( and dict(chat_template_kwargs) != (getattr(history, "chat_template_kwargs", None) or {}) ) - if isinstance(history, ChatCompletionsHistory) and ( - _history_needs_render(history) or override_requires_render - ): + needs_render = _history_needs_render(history) or override_requires_render + if isinstance(history, ResponsesHistory) and not needs_render: + if exact := _tokenize_exact_responses_history( + history, base_model=base_model, tokenizer=tokenizer + ): + return exact + if isinstance(history, ChatCompletionsHistory) and (needs_render): return _tokenize_chat_view( history, base_model=base_model, @@ -2144,14 +2887,9 @@ def tokenize_history( chat_template=chat_template, chat_template_kwargs=chat_template_kwargs, ) - if ( - isinstance(history, AnthropicMessagesHistory) - and (_history_needs_render(history) or override_requires_render) - ) or ( - isinstance(history, ResponsesHistory) - and (_history_needs_render(history) or override_requires_render) + if (isinstance(history, AnthropicMessagesHistory) and needs_render) or ( + isinstance(history, ResponsesHistory) and needs_render ): - _validate_history_sources(history) return _tokenize_chat_view( history.as_chat_completions_history(), base_model=base_model, diff --git a/tests/unit/trajectories/test_history.py b/tests/unit/trajectories/test_history.py index e46685fbb..da6066981 100644 --- a/tests/unit/trajectories/test_history.py +++ b/tests/unit/trajectories/test_history.py @@ -1,6 +1,6 @@ from datetime import datetime, timedelta import importlib -from typing import Any +from typing import Any, cast from anthropic.types import Message from openai.types import Completion @@ -219,6 +219,37 @@ def test_chat_history_resolves_one_model_and_append_only_sequence() -> None: trajectory.chat_completions_history(model="test/model") +def test_chat_history_preserves_provider_specific_nested_fields() -> None: + exchange = _chat( + [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "one", + "cache_control": {"type": "ephemeral"}, + } + ], + } + ], + "first", + ) + + trajectory = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[exchange]) + ) + history = trajectory.chat_completions_history() + + content = cast(list[dict[str, object]], history.messages[0]["content"]) + assert content[0]["cache_control"] == {"type": "ephemeral"} + dumped = trajectory.model_dump(mode="json", warnings="error") + dumped_content = dumped["exchanges"]["chat_completions"][0]["request"]["messages"][ + 0 + ]["content"] + assert dumped_content[0]["cache_control"] == {"type": "ephemeral"} + + def test_chat_choices_branch_and_identical_continuation_uses_first_choice() -> None: first = _chat([{"role": "user", "content": "one"}], "same") response = first.response.model_dump(mode="python") @@ -249,6 +280,43 @@ def test_chat_choices_branch_and_identical_continuation_uses_first_choice() -> N assert histories[1].message_sources[1].choice_index == 1 +def test_same_content_seeded_inputs_remain_request_sourced() -> None: + first_chat = _chat([{"role": "user", "content": "prompt"}], "same") + seeded_chat = _chat([{"role": "assistant", "content": "same"}], "next", offset=1) + chat_histories = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[first_chat, seeded_chat]) + ).chat_completions_histories() + seeded_chat_source = chat_histories[1].message_sources[0] + assert seeded_chat_source is not None + assert seeded_chat_source.exchange is seeded_chat + assert seeded_chat_source.request_index == 0 + + first_message = _message() + seeded_message = _message() + seeded_message.start_time, seeded_message.end_time = _times(1) + seeded_message.request["messages"] = [{"role": "assistant", "content": "Hi"}] + message_histories = art.Trajectory( + exchanges=TrajectoryExchanges(messages=[first_message, seeded_message]) + ).anthropic_messages_histories() + seeded_message_source = message_histories[1].message_sources[0] + assert seeded_message_source is not None + assert seeded_message_source.exchange is seeded_message + assert seeded_message_source.request_index == 0 + + first_response = _response("response-1", "same") + seeded_response = _response("response-2", "next", offset=1) + seeded_response.request["input"] = [ + first_response.response.output[0].model_dump(mode="json", exclude_none=True) + ] + response_histories = art.Trajectory( + exchanges=TrajectoryExchanges(responses=[first_response, seeded_response]) + ).responses_histories() + seeded_response_source = response_histories[1].input_sources[0] + assert seeded_response_source is not None + assert seeded_response_source.exchange is seeded_response + assert seeded_response_source.request_index == 0 + + def test_history_mutation_must_keep_source_sidecar_consistent() -> None: trajectory = art.Trajectory( exchanges=TrajectoryExchanges( @@ -279,7 +347,7 @@ def test_history_accepts_user_authored_messages_with_none_source() -> None: class Tokenizer: def __call__(self, text: str, *, add_special_tokens: bool = False) -> list[int]: assert not add_special_tokens - return {"first": [20], "next": [30]}[text] + return {"one": [10], "first": [20], "next": [30]}[text] def apply_chat_template( self, @@ -400,6 +468,147 @@ def test_responses_history_expands_previous_response_chain() -> None: assert external[1].previous_response_id == "missing" +def test_responses_history_propagates_opaque_context_and_first_sources() -> None: + first = _response( + "response-1", + "first", + previous_response_id="outside-trajectory", + ) + first.request["conversation"] = "conversation-1" + second = _response( + "response-2", + "second", + previous_response_id="response-1", + offset=1, + ) + trajectory = art.Trajectory( + exchanges=TrajectoryExchanges(responses=[first, second]) + ) + + history = trajectory.responses_history() + + assert history.previous_response_id == "outside-trajectory" + assert history.conversation == "conversation-1" + assert history.input_sources[1] is not None + assert history.input_sources[1].exchange is first + assert history.input_sources[1].output_index == 0 + + +def test_responses_history_maps_and_validates_generation_sources() -> None: + exchange = _response("response-1", "first", reasoning="think") + assert exchange.response.__pydantic_extra__ is not None + exchange.response.__pydantic_extra__["token_generations"] = [ + { + "prompt_token_ids": [1], + "output_tokens": [{"token_id": 2, "logprob": -0.1}], + "output_indices": [0, 1], + } + ] + + trajectory = art.Trajectory(exchanges=TrajectoryExchanges(responses=[exchange])) + history = trajectory.responses_history() + + assert { + source.generation_index + for source in history.input_sources + if source is not None and source.output_index is not None + } == {0} + + exchange.response.__pydantic_extra__["token_generations"][0]["output_indices"] = [0] + with pytest.raises(ValueError, match="every sampled output item"): + trajectory.responses_history() + + exchange.response.__pydantic_extra__["token_generations"] = [ + { + "prompt_token_ids": [1], + "output_tokens": [{"token_id": 2}], + "output_indices": [0], + }, + { + "prompt_token_ids": [1, 2], + "output_tokens": [{"token_id": 3}], + "output_indices": [0], + }, + ] + with pytest.raises(ValueError, match="nonoverlapping"): + trajectory.responses_history() + + +@pytest.mark.parametrize( + ("second_prompt", "reasoning", "expected_histories"), + [([1, 2, 3], False, 1), ([1, 3], True, 2)], +) +def test_responses_multi_generation_history_leaves_match_tokenization( + second_prompt: list[int], reasoning: bool, expected_histories: int +) -> None: + exchange = _response( + "response-1", "first", reasoning="think" if reasoning else None + ) + data = exchange.response.model_dump(mode="python") + first_output_indices = list(range(len(data["output"]))) + tool_output_index = len(data["output"]) + second_output_index = tool_output_index + 1 + data["output"].extend( + [ + { + "id": "tool-output", + "type": "function_call_output", + "call_id": "call-1", + "output": "result", + "status": "completed", + }, + { + "id": "message-second", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [ + { + "type": "output_text", + "text": "second", + "annotations": [], + "logprobs": [], + } + ], + }, + ] + ) + data["token_generations"] = [ + { + "prompt_token_ids": [1], + "output_tokens": [{"token_id": 2, "logprob": -0.2}], + "output_indices": first_output_indices, + }, + { + "prompt_token_ids": second_prompt, + "output_tokens": [{"token_id": 4, "logprob": -0.4}], + "output_indices": [second_output_index], + }, + ] + exchange.response = Response.model_validate(data) + trajectory = art.Trajectory(exchanges=TrajectoryExchanges(responses=[exchange])) + + histories = trajectory.responses_histories() + tokenized = trajectory.tokenize(multi_history=True) + + assert len(histories) == expected_histories + assert len(tokenized.histories) == expected_histories + final_sources = histories[-1].input_sources + assert final_sources[-2] is not None + assert final_sources[-2].generation_index is None + assert final_sources[-1] is not None + assert final_sources[-1].generation_index == 1 + if expected_histories == 2: + assert final_sources[1] is not None + assert final_sources[1].generation_index is None + assert all(item.get("type") != "reasoning" for item in histories[-1].input) + converted = histories[-1].as_chat_completions_history() + assert any( + source is not None and source.generation_index == 1 + for source in converted.message_sources + ) + + def test_completions_history_preserves_exact_tokens_and_sampled_spans() -> None: trajectory = art.Trajectory( exchanges=TrajectoryExchanges( @@ -456,6 +665,50 @@ def test_batched_completions_create_every_prompt_choice_history() -> None: trajectory.history() +def test_batched_completions_prefer_exact_prompt_association() -> None: + exchange = _completion([1], [10]) + exchange.request["prompt"] = [[1], [2]] + response = exchange.response.model_dump(mode="python") + response["choices"] = [ + { + "index": 0, + "finish_reason": "stop", + "text": "second", + "prompt_token_ids": [2], + "token_ids": [20], + }, + { + "index": 1, + "finish_reason": "stop", + "text": "first", + "prompt_token_ids": [1], + "token_ids": [10], + }, + ] + exchange.response = Completion.model_validate(response) + trajectory = art.Trajectory(exchanges=TrajectoryExchanges(completions=[exchange])) + + histories = trajectory.completions_token_histories() + + assert [history.prompt for history in histories] == [[1, 10], [2, 20]] + + +def test_completions_histories_never_silently_omit_mixed_evidence() -> None: + exact = _completion([1], [2]) + missing = _completion([3], [4], offset=1) + missing.request["prompt"] = "question" + response = missing.response.model_dump(mode="python") + response["choices"][0].pop("prompt_token_ids") + response["choices"][0].pop("token_ids") + missing.response = Completion.model_validate(response) + trajectory = art.Trajectory( + exchanges=TrajectoryExchanges(completions=[exact, missing]) + ) + + with pytest.raises(ValueError, match="text prompts for every choice"): + trajectory.histories() + + def test_completions_reject_ambiguous_batches_and_suffix() -> None: ambiguous = _completion([1], [10]) ambiguous.request["prompt"] = ["first", "second"] @@ -556,11 +809,7 @@ def test_chat_template_stripped_reasoning_splits_exact_histories() -> None: second = _chat( [ {"role": "user", "content": "one"}, - { - "role": "assistant", - "content": "first", - "reasoning": "thought-one", - }, + {"role": "assistant", "content": "first"}, {"role": "user", "content": "two"}, ], "second", @@ -579,6 +828,9 @@ def test_chat_template_stripped_reasoning_splits_exact_histories() -> None: assert len(histories) == 2 assert [len(history.messages) for history in histories] == [2, 4] + assert histories[1].message_sources[1] is not None + assert histories[1].message_sources[1].exchange is first + assert histories[1].message_sources[1].choice_index == 0 with pytest.raises(ValueError, match="exactly one history"): trajectory.tokenize() diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index 4c102e11d..af81b6558 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -3,6 +3,7 @@ import builtins from datetime import datetime, timedelta import math +import random import sys from types import ModuleType, SimpleNamespace from typing import Any @@ -23,6 +24,7 @@ MessagesExchange, MessagesRequest, ResponsesExchange, + ResponsesItemSource, ResponsesRequest, TrajectoryExchanges, ) @@ -225,7 +227,7 @@ def test_malformed_explicit_exact_token_metadata_fails_closed() -> None: response = _response_exchange("response-invalid", 2) response_extra = response.response.model_extra assert response_extra is not None - response_extra["raw_output_tokens"] = [{"token_id": "invalid"}] + response_extra["token_generations"][0]["output_tokens"] = [{"token_id": "invalid"}] message_response = Message.model_validate( { @@ -263,6 +265,19 @@ def test_malformed_explicit_exact_token_metadata_fails_closed() -> None: trajectory.tokenize(base_model="base/model") +@pytest.mark.parametrize("token_id", [-1, True]) +def test_exact_token_metadata_rejects_negative_ids_and_booleans( + token_id: object, +) -> None: + exchange = _response_exchange("response-invalid-token", 2) + extra = exchange.response.model_extra + assert extra is not None + extra["token_generations"][0]["output_tokens"] = [{"token_id": token_id}] + + with pytest.raises(ValueError, match="invalid exact token ID"): + art.Trajectory(exchanges=TrajectoryExchanges(responses=[exchange])).tokenize() + + @pytest.mark.parametrize( "exchange", [ @@ -307,6 +322,226 @@ def test_completions_echo_preserves_prompt_logprobs_without_sampling_them( ] +def test_completions_exact_ids_accept_textual_logprobs() -> None: + exchange = _completion_exchange() + payload = exchange.response.model_dump(mode="python") + payload["choices"][0]["logprobs"]["tokens"] = ["answer"] + payload["choices"][0]["logprobs"]["token_logprobs"] = [-0.75] + exchange.response = Completion.model_validate(payload) + + tokenized = art.Trajectory( + exchanges=TrajectoryExchanges(completions=[exchange]) + ).tokenize() + + assert tokenized.token_ids == [1, 2] + assert tokenized.logprobs[1] == -0.75 + + +def test_mutated_completions_token_prompt_drops_stale_exact_evidence() -> None: + history = art.Trajectory( + exchanges=TrajectoryExchanges(completions=[_completion_exchange()]) + ).completions_token_history() + history.prompt[0] = 99 + + tokenized = history.tokenize() + + assert tokenized.token_ids == [99, 2] + assert tokenized.flags == [ + art.TokenFlag(0), + art.TokenFlag.EXACT | art.TokenFlag.SAMPLED, + ] + + +def test_completions_string_history_preserves_textual_logprobs() -> None: + exchange = _completion_exchange() + payload = exchange.response.model_dump(mode="python") + payload["choices"][0].pop("prompt_token_ids", None) + payload["choices"][0].pop("token_ids", None) + payload["choices"][0]["logprobs"]["tokens"] = ["answer"] + payload["choices"][0]["logprobs"]["token_logprobs"] = [-0.8] + exchange.response = Completion.model_validate(payload) + history = art.Trajectory( + exchanges=TrajectoryExchanges(completions=[exchange]) + ).completions_string_history() + + class Tokenizer: + def __call__(self, text: str, **kwargs: object) -> list[int]: + del kwargs + return {"question": [1], "answer": [2]}[text] + + def apply_chat_template(self, *args: object, **kwargs: object) -> list[int]: + raise AssertionError("Completions tokenization does not render chat") + + tokenized = history.tokenize(tokenizer=Tokenizer()) + + assert tokenized.token_ids == [1, 2] + assert tokenized.logprobs[1] == -0.8 + assert tokenized.flags[1] == art.TokenFlag.SAMPLED + + +def test_mutated_completions_string_prompt_retokens_without_stale_exact() -> None: + history = art.Trajectory( + exchanges=TrajectoryExchanges(completions=[_completion_exchange()]) + ).completions_string_history() + history.prompt = "changed" + history.prompt[len("question") :] + first = history.prompt_sources[0] + history.prompt_sources[0] = type(first)( + start=0, + end=len("changed"), + source=first.source, + ) + shift = len("changed") - len("question") + second = history.prompt_sources[1] + history.prompt_sources[1] = type(second)( + start=second.start + shift, + end=second.end + shift, + source=second.source, + ) + history.sampled_spans = [ + (start + shift, end + shift) for start, end in history.sampled_spans + ] + + class Tokenizer: + def __call__(self, text: str, **kwargs: object) -> list[int]: + del kwargs + return {"changed": [9], "answer": [2]}[text] + + def apply_chat_template(self, *args: object, **kwargs: object) -> list[int]: + raise AssertionError("Completions tokenization does not render chat") + + tokenized = history.tokenize(tokenizer=Tokenizer()) + + assert tokenized.token_ids == [9, 2] + assert tokenized.flags[0] == art.TokenFlag(0) + + +def test_batched_completions_map_each_choice_to_its_prompt() -> None: + exchange = _completion_exchange(prompt=["first", "second"]) + payload = exchange.response.model_dump(mode="python") + payload["choices"] = [ + { + **payload["choices"][0], + "index": 0, + "text": "one", + "prompt_token_ids": [10], + "token_ids": [11], + "logprobs": { + "tokens": ["token_id:11"], + "token_logprobs": [-0.1], + "top_logprobs": [{}], + "text_offset": [0], + }, + }, + { + **payload["choices"][0], + "index": 1, + "text": "two", + "prompt_token_ids": [20], + "token_ids": [21], + "logprobs": { + "tokens": ["token_id:21"], + "token_logprobs": [-0.2], + "top_logprobs": [{}], + "text_offset": [0], + }, + }, + ] + exchange.response = Completion.model_validate(payload) + + tokenized = art.Trajectory( + exchanges=TrajectoryExchanges(completions=[exchange]) + ).tokenize(multi_history=True) + + assert [history.token_ids for history in tokenized.histories] == [ + [10, 11], + [20, 21], + ] + assert [history.flags for history in tokenized.histories] == [ + [art.TokenFlag.EXACT, art.TokenFlag.EXACT | art.TokenFlag.SAMPLED], + [art.TokenFlag.EXACT, art.TokenFlag.EXACT | art.TokenFlag.SAMPLED], + ] + + +def test_history_tokenization_rejects_negative_source_indices() -> None: + exchange = _response_exchange("response-negative-source", 2) + history = art.Trajectory( + exchanges=TrajectoryExchanges(responses=[exchange]) + ).responses_history() + history.input_sources[-1] = ResponsesItemSource( + exchange=exchange, + output_index=-1, + generation_index=0, + ) + + with pytest.raises(ValueError, match="out of bounds"): + history.tokenize() + + +def test_batched_completions_reject_ambiguous_choice_indices() -> None: + exchange = _completion_exchange(prompt=["first", "second"]) + payload = exchange.response.model_dump(mode="python") + payload["choices"] = [ + {**payload["choices"][0], "index": 0}, + {**payload["choices"][0], "index": 2}, + ] + exchange.response = Completion.model_validate(payload) + + with pytest.raises(ValueError, match="Ambiguous"): + art.Trajectory(exchanges=TrajectoryExchanges(completions=[exchange])).tokenize( + multi_history=True + ) + + +def test_randomized_completions_projection_preserves_every_choice_once() -> None: + rng = random.Random(0) + for case in range(20): + prompt_count = rng.randint(1, 5) + choices_per_prompt = rng.randint(1, 4) + exchange = _completion_exchange( + prompt=[f"prompt-{case}-{index}" for index in range(prompt_count)] + ) + exchange.request["n"] = choices_per_prompt + template = exchange.response.model_dump(mode="python")["choices"][0] + choices: list[dict[str, object]] = [] + expected: list[list[int]] = [] + for prompt_index in range(prompt_count): + prompt_id = 10_000 + case * 100 + prompt_index + for local_choice in range(choices_per_prompt): + choice_index = prompt_index * choices_per_prompt + local_choice + output_id = 20_000 + case * 100 + choice_index + expected.append([prompt_id, output_id]) + choices.append( + { + **template, + "index": choice_index, + "text": f"answer-{output_id}", + "prompt_token_ids": [prompt_id], + "token_ids": [output_id], + "logprobs": { + "tokens": [f"token_id:{output_id}"], + "token_logprobs": [-0.1], + "top_logprobs": [{}], + "text_offset": [0], + }, + } + ) + rng.shuffle(choices) + response = exchange.response.model_dump(mode="python") + response["choices"] = choices + exchange.response = Completion.model_validate(response) + + tokenized = art.Trajectory( + exchanges=TrajectoryExchanges(completions=[exchange]) + ).tokenize(multi_history=True) + + assert [history.token_ids for history in tokenized.histories] == expected + assert all( + history.flags + == [art.TokenFlag.EXACT, art.TokenFlag.EXACT | art.TokenFlag.SAMPLED] + for history in tokenized.histories + ) + + def test_branching_and_multiple_models_require_explicit_resolution() -> None: alternate = _chat_exchange([9], [3], offset=1) alternate.request["messages"] = [{"role": "user", "content": "alternate"}] @@ -767,6 +1002,106 @@ def __call__(self, text: str, **kwargs: object) -> SimpleNamespace: assert math.isnan(result.logprobs[2]) +def test_exact_chat_ids_accept_ordinary_positional_logprobs() -> None: + exchange = _chat_exchange([1], [2]) + logprobs = exchange.response.choices[0].logprobs + assert logprobs is not None and logprobs.content + exchange.response.choices[0].logprobs = logprobs.model_copy( + update={ + "content": [ + logprobs.content[0].model_copy( + update={"token": "answer", "logprob": -0.75} + ) + ] + } + ) + + tokenized = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[exchange]) + ).tokenize() + + assert tokenized.token_ids == [1, 2] + assert tokenized.logprobs[1] == -0.75 + + +def test_empty_chat_prompt_ids_are_missing_evidence( + monkeypatch: pytest.MonkeyPatch, +) -> None: + exchange = _chat_exchange([], [2]) + + class Tokenizer: + def apply_chat_template( + self, messages: list[dict[str, Any]], **kwargs: object + ) -> list[int]: + del kwargs + return [10, 20] if messages[-1]["role"] == "assistant" else [10] + + def __call__(self, text: str, **kwargs: object) -> list[int]: + del text, kwargs + return [20] + + monkeypatch.setattr( + "art.trajectories._tokenize._load_tokenizer", lambda _config: Tokenizer() + ) + tokenized = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[exchange]) + ).tokenize(base_model="base/model") + + assert tokenized.token_ids == [10, 2] + assert tokenized.flags == [ + art.TokenFlag(0), + art.TokenFlag.EXACT | art.TokenFlag.SAMPLED, + ] + + +def test_missing_completion_renders_only_missing_region_when_prompt_is_exact( + monkeypatch: pytest.MonkeyPatch, +) -> None: + exchange = _chat_exchange([99], []) + extra = exchange.response.choices[0].model_extra + assert extra is not None + extra.pop("token_ids", None) + logprobs = exchange.response.choices[0].logprobs + assert logprobs is not None + exchange.response.choices[0].logprobs = logprobs.model_copy( + update={ + "content": [ + ChatCompletionTokenLogprob( + token="answer", + logprob=-0.5, + bytes=list(b"answer"), + top_logprobs=[], + ) + ] + } + ) + + class Tokenizer: + def apply_chat_template( + self, messages: list[dict[str, Any]], **kwargs: object + ) -> list[int]: + del kwargs + return [10, 20] if messages[-1]["role"] == "assistant" else [10] + + def __call__(self, text: str, **kwargs: object) -> list[int]: + del text, kwargs + return [20] + + monkeypatch.setattr( + "art.trajectories._tokenize._load_tokenizer", lambda _config: Tokenizer() + ) + tokenized = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[exchange]) + ).tokenize(base_model="base/model") + + assert tokenized.token_ids == [99, 20] + assert tokenized.logprobs[1] == -0.5 + assert tokenized.flags == [ + art.TokenFlag.EXACT, + art.TokenFlag.SAMPLED, + ] + + def test_ambiguous_visible_logprobs_fail_closed( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -938,6 +1273,7 @@ def _response_exchange( *, previous_response_id: str | None = None, offset: int = 0, + prompt_token_ids: list[int] | None = None, ) -> ResponsesExchange: response = Response.model_validate( { @@ -964,7 +1300,13 @@ def _response_exchange( "parallel_tool_calls": True, "tool_choice": "auto", "tools": [], - "raw_output_tokens": [{"token_id": output_id, "logprob": -0.1}], + "token_generations": [ + { + "prompt_token_ids": prompt_token_ids or [10], + "output_tokens": [{"token_id": output_id, "logprob": -0.1}], + "output_indices": [0], + } + ], } ) request = ResponsesRequest(model="test/model", input=f"turn {offset}") @@ -982,7 +1324,7 @@ def _response_exchange( def _response_with_content_logprobs(*, exact_second: bool) -> ResponsesExchange: exchange = _response_exchange("response-content-logprobs", 0) data = exchange.response.model_dump(mode="python") - data.pop("raw_output_tokens", None) + data.pop("token_generations", None) def entry(token: str, token_id: int | None, logprob: float) -> dict[str, Any]: return { @@ -1042,12 +1384,12 @@ def apply_chat_template( assert result.logprobs[1:] == [-0.1, -0.2] -def test_responses_empty_raw_tokens_fall_back_for_visible_output( +def test_responses_missing_token_generations_falls_back_for_visible_output( monkeypatch: pytest.MonkeyPatch, ) -> None: exchange = _response_exchange("response-empty-raw", 0) data = exchange.response.model_dump(mode="python") - data["raw_output_tokens"] = [] + data.pop("token_generations", None) exchange.response = Response.model_validate(data) monkeypatch.setattr( "art.trajectories._tokenize._load_tokenizer", lambda _config: _FakeTokenizer() @@ -1112,6 +1454,9 @@ def apply_chat_template( "art.trajectories._tokenize._load_tokenizer", lambda _config: Tokenizer() ) request_reasoning = _response_exchange("request-reasoning", 2) + request_data = request_reasoning.response.model_dump(mode="python") + request_data.pop("token_generations", None) + request_reasoning.response = Response.model_validate(request_data) request_reasoning.request["input"] = [ { "id": "reasoning-1", @@ -1129,7 +1474,7 @@ def apply_chat_template( "type": "reasoning", } ] - data.pop("raw_output_tokens", None) + data.pop("token_generations", None) response_reasoning.response = Response.model_validate(data) art.Trajectory( @@ -1152,6 +1497,9 @@ def apply_chat_template( previous_response_id=response_reasoning.response.id, offset=1, ) + continuation_data = continuation.response.model_dump(mode="python") + continuation_data.pop("token_generations", None) + continuation.response = Response.model_validate(continuation_data) assert art.Trajectory( exchanges=TrajectoryExchanges(responses=[response_reasoning, continuation]) ).tokenize( @@ -1182,7 +1530,7 @@ def test_responses_opaque_reasoning_requires_exact_tokens( ).token_ids == [10, 2] response = exchange.response.model_dump(mode="python") - response.pop("raw_output_tokens", None) + response.pop("token_generations", None) exchange.response = Response.model_validate(response) with pytest.raises(ValueError, match="no renderable text"): art.Trajectory(exchanges=TrajectoryExchanges(responses=[exchange])).tokenize( @@ -1194,6 +1542,9 @@ def test_responses_parallel_function_calls_form_one_assistant_turn( monkeypatch: pytest.MonkeyPatch, ) -> None: exchange = _response_exchange("parallel-tools", 2) + response_data = exchange.response.model_dump(mode="python") + response_data.pop("token_generations", None) + exchange.response = Response.model_validate(response_data) exchange.request["input"] = [ { "id": "reasoning-1", @@ -1232,6 +1583,189 @@ def apply_chat_template( ] +def test_responses_token_generations_preserve_every_generation() -> None: + exchange = _response_exchange("multi-generation", 2) + data = exchange.response.model_dump(mode="python") + data["output"].append( + { + "id": "message-second", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [ + { + "type": "output_text", + "text": "second", + "annotations": [], + "logprobs": [], + } + ], + } + ) + data["token_generations"] = [ + { + "prompt_token_ids": [1], + "output_tokens": [{"token_id": 2, "logprob": -0.2}], + "output_indices": [0], + }, + { + "prompt_token_ids": [1, 2, 3], + "output_tokens": [{"token_id": 4, "logprob": -0.4}], + "output_indices": [1], + }, + ] + exchange.response = Response.model_validate(data) + history = art.Trajectory( + exchanges=TrajectoryExchanges(responses=[exchange]) + ).responses_history() + + tokenized = history.tokenize() + + assert tokenized.token_ids == [1, 2, 3, 4] + assert tokenized.logprobs[1] == -0.2 + assert tokenized.logprobs[3] == -0.4 + assert tokenized.flags == [ + art.TokenFlag.EXACT, + art.TokenFlag.EXACT | art.TokenFlag.SAMPLED, + art.TokenFlag.EXACT, + art.TokenFlag.EXACT | art.TokenFlag.SAMPLED, + ] + + +def test_responses_prompt_disagreement_preserves_sampled_token_identity( + monkeypatch: pytest.MonkeyPatch, +) -> None: + exchange = _response_exchange("retokenized-generation", 101) + data = exchange.response.model_dump(mode="python") + data["output"].append( + { + "id": "message-second", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [ + { + "type": "output_text", + "text": "dog", + "annotations": [], + "logprobs": [], + } + ], + } + ) + data["token_generations"] = [ + { + "prompt_token_ids": [1], + "output_tokens": [{"token_id": 101, "logprob": -0.1, "text": "cat"}], + "output_indices": [0], + }, + { + "prompt_token_ids": [1, 500, 3], + "output_tokens": [{"token_id": 4, "logprob": -0.4, "text": "dog"}], + "output_indices": [1], + }, + ] + exchange.response = Response.model_validate(data) + history = art.Trajectory( + exchanges=TrajectoryExchanges(responses=[exchange]) + ).responses_history() + + class Tokenizer: + def __call__(self, text: str, **kwargs: object) -> list[int]: + del kwargs + return {"cat": [500], "dog": [4]}[text] + + def apply_chat_template(self, *args: object, **kwargs: object) -> list[int]: + raise AssertionError("Exact Responses tokenization does not render chat") + + monkeypatch.setattr( + "art.trajectories._tokenize._WARNED_PREFIX_RETOKENIZATION", False + ) + with pytest.warns(UserWarning, match="preserved the original sampled token IDs"): + tokenized = history.tokenize(tokenizer=Tokenizer()) + + assert tokenized.token_ids == [1, 101, 3, 4] + assert tokenized.logprobs[1] == -0.1 + + +@pytest.mark.parametrize( + "token_generations, match", + [ + ([], "must be omitted"), + ( + [ + { + "output_tokens": [{"token_id": 2}], + "output_indices": [0], + } + ], + "prompt_token_ids", + ), + ( + [ + { + "prompt_token_ids": [1], + "output_tokens": [{"token_id": True}], + "output_indices": [0], + } + ], + "exact token ID", + ), + ( + [ + { + "prompt_token_ids": [1], + "output_tokens": [{"token_id": 2}], + "output_indices": [True], + } + ], + "integers", + ), + ( + [ + { + "prompt_token_ids": [1], + "output_tokens": [{"token_id": 2}], + "output_indices": [1], + } + ], + "out of bounds", + ), + ], +) +def test_responses_token_generations_fail_closed( + token_generations: list[dict[str, Any]], match: str +) -> None: + exchange = _response_exchange("invalid-generation", 2) + data = exchange.response.model_dump(mode="python") + data["token_generations"] = token_generations + exchange.response = Response.model_validate(data) + + with pytest.raises(ValueError, match=match): + art.Trajectory( + exchanges=TrajectoryExchanges(responses=[exchange]) + ).responses_history().tokenize() + + +def test_responses_generation_without_output_items_is_not_silently_lost() -> None: + exchange = _response_exchange("unprojectable-generation", 2) + data = exchange.response.model_dump(mode="python") + data["output"] = [] + data["token_generations"] = [ + { + "prompt_token_ids": [1], + "output_tokens": [{"token_id": 2, "logprob": -0.2}], + "output_indices": [], + } + ] + exchange.response = Response.model_validate(data) + + with pytest.raises(ValueError, match="cannot yet be projected"): + art.Trajectory( + exchanges=TrajectoryExchanges(responses=[exchange]) + ).responses_history() + + def test_tokenization_rejects_mutated_mixed_representation() -> None: trajectory = art.Trajectory( messages_and_choices=[{"role": "user", "content": "hi"}] @@ -1255,8 +1789,14 @@ def apply_chat_template( monkeypatch.setattr( "art.trajectories._tokenize._load_tokenizer", lambda _config: Tokenizer() ) - first = _response_exchange("resp-1", 20) - second = _response_exchange("resp-2", 30, previous_response_id="resp-1", offset=1) + first = _response_exchange("resp-1", 20, prompt_token_ids=[10]) + second = _response_exchange( + "resp-2", + 30, + previous_response_id="resp-1", + offset=1, + prompt_token_ids=[10, 20, 11], + ) trajectory = art.Trajectory( exchanges=TrajectoryExchanges(responses=[first, second]) ) @@ -1272,8 +1812,9 @@ def apply_chat_template( assert len(trajectory.responses_histories()) == 2 with pytest.raises(ValueError, match="exactly one history"): trajectory.tokenize(base_model="base/model") - with pytest.raises(ValueError, match="outside this trajectory"): - trajectory.responses_histories()[1].tokenize(base_model="base/model") + assert trajectory.responses_histories()[1].tokenize( + base_model="base/model" + ).token_ids == [10, 20, 11, 30] def test_prefix_retokenization_preserves_sampled_ids_and_logprobs( @@ -1345,6 +1886,310 @@ def apply_chat_template( assert tokenized.flags[1] == art.TokenFlag.EXACT | art.TokenFlag.SAMPLED +def test_mutable_chat_history_is_authoritative_and_does_not_replay_removed_turns() -> ( + None +): + first = _chat_exchange([10], [20]) + first.request["messages"] = [{"role": "user", "content": "first"}] + second = _chat_exchange([10, 20, 30], [40], offset=1) + second.request["messages"] = [ + {"role": "user", "content": "first"}, + {"role": "assistant", "content": "answer"}, + {"role": "user", "content": "second"}, + ] + history = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[first, second]) + ).chat_completions_history() + del history.messages[:2] + del history.message_sources[:2] + + class Tokenizer: + def __call__(self, text: str, **kwargs: object) -> list[int]: + del kwargs + return {"second": [31], "answer": [41]}[text] + + def apply_chat_template( + self, messages: list[dict[str, Any]], **kwargs: object + ) -> list[int]: + del kwargs + assert [message["content"] for message in messages] == [ + "second", + "answer", + ] + return [30, 41, 50] + + tokenized = history.tokenize(tokenizer=Tokenizer()) + + assert tokenized.token_ids == [30, 40, 50] + assert 20 not in tokenized.token_ids + assert tokenized.flags == [ + art.TokenFlag(0), + art.TokenFlag.EXACT | art.TokenFlag.SAMPLED, + art.TokenFlag(0), + ] + + +def test_request_assistant_messages_are_not_marked_sampled() -> None: + exchange = _chat_exchange([10], [40]) + exchange.request["messages"] = [ + {"role": "assistant", "content": "seed"}, + {"role": "user", "content": "question"}, + ] + history = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[exchange]) + ).chat_completions_history() + history.chat_template = "rerender" + + class Tokenizer: + def __call__(self, text: str, **kwargs: object) -> list[int]: + del kwargs + return {"seed": [20], "question": [30], "answer": [41]}[text] + + def apply_chat_template( + self, messages: list[dict[str, Any]], **kwargs: object + ) -> list[int]: + del messages, kwargs + return [5, 20, 6, 30, 7, 41, 8] + + tokenized = history.tokenize(tokenizer=Tokenizer()) + + assert tokenized.token_ids == [5, 20, 6, 30, 7, 40, 8] + assert not tokenized.flags[1] & art.TokenFlag.SAMPLED + assert tokenized.flags[5] == art.TokenFlag.EXACT | art.TokenFlag.SAMPLED + + +def test_rerender_constrains_exact_output_to_its_message_region() -> None: + exchange = _chat_exchange([7], [7]) + exchange.request["messages"] = [{"role": "user", "content": "same"}] + exchange.response.choices[0].message.content = "same" + history = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[exchange]) + ).chat_completions_history() + history.chat_template = "rerender" + + class Tokenizer: + def __call__(self, text: str, **kwargs: object) -> list[int]: + del kwargs + assert text == "same" + return [7] + + def apply_chat_template( + self, messages: list[dict[str, Any]], **kwargs: object + ) -> list[int]: + del messages, kwargs + return [7, 99, 7] + + tokenized = history.tokenize(tokenizer=Tokenizer()) + + assert tokenized.flags == [ + art.TokenFlag(0), + art.TokenFlag(0), + art.TokenFlag.EXACT | art.TokenFlag.SAMPLED, + ] + assert math.isnan(tokenized.logprobs[0]) + assert tokenized.logprobs[2] == -0.7 + + +def test_reasoning_stripped_chat_histories_tokenize_authoritative_views() -> None: + first = _chat_exchange([1], [2, 3]) + first.request["messages"] = [{"role": "user", "content": "one"}] + first_data = first.response.model_dump(mode="python") + first_data["choices"][0]["message"] = { + "role": "assistant", + "content": "first", + "reasoning": "thought-one", + } + first.response = ChatCompletion.model_validate(first_data) + + second = _chat_exchange([1, 3, 4], [5, 6], offset=1) + second.request["messages"] = [ + {"role": "user", "content": "one"}, + {"role": "assistant", "content": "first"}, + {"role": "user", "content": "two"}, + ] + second_data = second.response.model_dump(mode="python") + second_data["choices"][0]["message"] = { + "role": "assistant", + "content": "second", + "reasoning": "thought-two", + } + second.response = ChatCompletion.model_validate(second_data) + trajectory = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[first, second]) + ) + + class Tokenizer: + def __call__(self, text: str, **kwargs: object) -> list[int]: + del kwargs + return { + "one": [1], + "first": [3], + "two": [4], + "thought-two": [5], + "second": [6], + }[text] + + def apply_chat_template( + self, messages: list[dict[str, Any]], **kwargs: object + ) -> list[int]: + del kwargs + assert [message.get("content") for message in messages] == [ + "one", + "first", + "two", + "second", + ] + return [1, 3, 4, 5, 6] + + tokenized = trajectory.tokenize(multi_history=True, tokenizer=Tokenizer()) + + assert [history.token_ids for history in tokenized.histories] == [ + [1, 2, 3], + [1, 3, 4, 5, 6], + ] + assert tokenized.histories[1].flags[1] & art.TokenFlag.SAMPLED + assert 2 not in tokenized.histories[1].token_ids + + +def test_responses_prompt_repair_uses_native_text_and_source_position() -> None: + exchange = _response_exchange("repeated-retokenization", 101) + data = exchange.response.model_dump(mode="python") + data["output"].append( + { + "id": "message-second", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [ + { + "type": "output_text", + "text": "dog", + "annotations": [], + "logprobs": [], + } + ], + } + ) + data["token_generations"] = [ + { + "prompt_token_ids": [500], + "output_tokens": [{"token_id": 101, "logprob": -0.1}], + "output_indices": [0], + }, + { + "prompt_token_ids": [500, 500, 3], + "output_tokens": [{"token_id": 4, "logprob": -0.4}], + "output_indices": [1], + }, + ] + exchange.response = Response.model_validate(data) + history = art.Trajectory( + exchanges=TrajectoryExchanges(responses=[exchange]) + ).responses_history() + + class Tokenizer: + def __call__(self, text: str, **kwargs: object) -> list[int]: + del kwargs + return {"answer": [500], "dog": [4]}[text] + + def apply_chat_template(self, *args: object, **kwargs: object) -> list[int]: + raise AssertionError("Exact Responses tokenization must not render chat") + + with pytest.warns(UserWarning, match="preserved the original sampled token IDs"): + tokenized = history.tokenize(tokenizer=Tokenizer()) + + assert tokenized.token_ids == [500, 101, 3, 4] + assert tokenized.logprobs[1] == -0.1 + + +def test_rerender_marks_tool_call_only_generated_region_sampled() -> None: + exchange = _chat_exchange([1], [2]) + data = exchange.response.model_dump(mode="python") + choice = data["choices"][0] + choice.pop("token_ids") + choice["logprobs"] = None + choice["message"] = { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call-1", + "type": "function", + "function": {"name": "lookup", "arguments": '{"x":1}'}, + } + ], + } + exchange.response = ChatCompletion.model_validate(data) + history = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[exchange]) + ).chat_completions_history() + history.chat_template = "rerender" + + class Tokenizer: + def __call__(self, text: str, **kwargs: object) -> list[int]: + del kwargs + return { + "turn 0": [1], + "lookup": [20], + '{"x":1}': [30, 31], + }[text] + + def apply_chat_template( + self, messages: list[dict[str, Any]], **kwargs: object + ) -> list[int]: + del messages, kwargs + return [1, 10, 20, 25, 30, 31, 26] + + tokenized = history.tokenize(tokenizer=Tokenizer()) + + assert tokenized.token_ids == [1, 10, 20, 25, 30, 31, 26] + assert tokenized.flags[2:6] == [art.TokenFlag.SAMPLED] * 4 + assert not tokenized.flags[0] & art.TokenFlag.SAMPLED + + +def test_rerender_calls_chat_template_once_for_many_turns() -> None: + exchanges: list[ChatCompletionsExchange] = [] + prompt: list[int] = [] + messages: list[ChatCompletionMessageParam] = [] + for index in range(32): + prompt.extend([index * 2]) + messages.append({"role": "user", "content": f"u{index}"}) + exchange = _chat_exchange(list(prompt), [index * 2 + 1], offset=index) + exchange.request["messages"] = list(messages) + exchanges.append(exchange) + prompt.append(index * 2 + 1) + messages.append({"role": "assistant", "content": "answer"}) + history = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=exchanges) + ).chat_completions_history() + history.chat_template = "rerender" + + class Tokenizer: + apply_calls = 0 + + def __call__(self, text: str, **kwargs: object) -> list[int]: + del kwargs + return [1000] if text == "answer" else [2000 + int(text[1:])] + + def apply_chat_template( + self, messages: list[dict[str, Any]], **kwargs: object + ) -> list[int]: + del kwargs + self.apply_calls += 1 + result: list[int] = [] + for message in messages: + content = str(message["content"]) + result.extend( + [1000] if content == "answer" else [2000 + int(content[1:])] + ) + return result + + tokenizer = Tokenizer() + history.tokenize(tokenizer=tokenizer) + + assert tokenizer.apply_calls == 1 + + def test_explicit_template_override_rerenders_exact_exchange_scaffold() -> None: trajectory = art.Trajectory( exchanges=TrajectoryExchanges(chat_completions=[_chat_exchange([1], [2])]) @@ -1380,6 +2225,9 @@ def test_responses_external_context_requires_or_uses_exact_prompt_tokens() -> No exchange = _response_exchange( "external", 2, previous_response_id="outside-trajectory" ) + response = exchange.response.model_dump(mode="python") + response.pop("token_generations", None) + exchange.response = Response.model_validate(response) history = art.Trajectory( exchanges=TrajectoryExchanges(responses=[exchange]) ).responses_history() @@ -1387,7 +2235,13 @@ def test_responses_external_context_requires_or_uses_exact_prompt_tokens() -> No history.tokenize(base_model="base/model") response = exchange.response.model_dump(mode="python") - response["prompt_token_ids"] = [7, 8] + response["token_generations"] = [ + { + "prompt_token_ids": [7, 8], + "output_tokens": [{"token_id": 2, "logprob": -0.1}], + "output_indices": [0], + } + ] exchange.response = Response.model_validate(response) tokenized = ( art.Trajectory(exchanges=TrajectoryExchanges(responses=[exchange])) @@ -1406,13 +2260,22 @@ def test_responses_external_context_requires_or_uses_exact_prompt_tokens() -> No def test_responses_conversation_requires_exact_prompt_tokens() -> None: exchange = _response_exchange("conversation", 2) exchange.request["conversation"] = "conversation-1" + response = exchange.response.model_dump(mode="python") + response.pop("token_generations", None) + exchange.response = Response.model_validate(response) trajectory = art.Trajectory(exchanges=TrajectoryExchanges(responses=[exchange])) with pytest.raises(ValueError, match="conversation history requires exact"): trajectory.tokenize(base_model="base/model") response = exchange.response.model_dump(mode="python") - response["prompt_token_ids"] = [5] + response["token_generations"] = [ + { + "prompt_token_ids": [5], + "output_tokens": [{"token_id": 2, "logprob": -0.1}], + "output_indices": [0], + } + ] exchange.response = Response.model_validate(response) assert trajectory.tokenize().token_ids == [5, 2] From 4bbf144a410ad8437c2d5ec6b628bf08ddb2af26 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Thu, 23 Jul 2026 23:00:35 +0000 Subject: [PATCH 03/58] Select exchange models explicitly during training --- src/art/local/backend.py | 1 + src/art/preprocessing/tokenize.py | 141 ++++++++--- src/art/tinker_native/backend.py | 1 + src/art/tinker_native/data.py | 8 +- .../test_exchange_training_model_selection.py | 226 ++++++++++++++++++ 5 files changed, 342 insertions(+), 35 deletions(-) create mode 100644 tests/unit/test_exchange_training_model_selection.py diff --git a/src/art/local/backend.py b/src/art/local/backend.py index 64657d0d3..63bbe2192 100644 --- a/src/art/local/backend.py +++ b/src/art/local/backend.py @@ -958,6 +958,7 @@ def _get_packed_tensors( image_processor=self._image_processors[model.base_model], chat_template_kwargs=chat_template_kwargs, chat_template_tool_schema_format=chat_template_tool_schema_format, + model=self._model_inference_name(model), ) ) if not tokenized_results: diff --git a/src/art/preprocessing/tokenize.py b/src/art/preprocessing/tokenize.py index 03208901a..0d4732c0d 100644 --- a/src/art/preprocessing/tokenize.py +++ b/src/art/preprocessing/tokenize.py @@ -63,6 +63,87 @@ def _flag_spans(flags: list[TokenFlag], flag: TokenFlag) -> list[tuple[int, int] return spans +@dataclass(frozen=True) +class _ChatChoiceTrace: + choices: list[Choice] + offsets: list[int] + lengths: list[int] + + +def _chat_choice_trace( + history: ChatCompletionsHistory, + token_ids: list[int], + flags: list[TokenFlag], +) -> _ChatChoiceTrace | None: + """Recover per-choice boundaries that adjacent SAMPLED flags cannot represent.""" + + sourced_choices: list[Choice] = [] + seen: set[tuple[int, int]] = set() + for source in history.message_sources: + if ( + source is None + or source.choice_index is None + or not isinstance(source.exchange, ChatCompletionsExchange) + ): + continue + key = (id(source.exchange), source.choice_index) + if key in seen: + continue + seen.add(key) + sourced_choices.append( + next( + item + for item in source.exchange.response.choices + if item.index == source.choice_index + ) + ) + if not sourced_choices: + return None + has_routing = any( + choice_moe_routing_metadata(choice) is not None for choice in sourced_choices + ) + if any(choice_vllm_token_metadata(choice) is None for choice in sourced_choices): + if has_routing: + raise RuntimeError( + "MoE routing replay requires exact token IDs for every sourced choice" + ) + return None + + choices: list[Choice] = [] + offsets: list[int] = [] + lengths: list[int] = [] + cursor = 0 + for choice in sourced_choices: + metadata = choice_vllm_token_metadata(choice) + assert metadata is not None + completion_ids = metadata[1] + matches = [ + index + for index in range(cursor, len(token_ids) - len(completion_ids) + 1) + if token_ids[index : index + len(completion_ids)] == completion_ids + and all( + flag & TokenFlag.SAMPLED + for flag in flags[index : index + len(completion_ids)] + ) + ] + if not matches: + if has_routing: + raise RuntimeError( + "MoE routed completion tokens are absent from tokenized history" + ) + return None + offset = matches[0] + choices.append(choice) + offsets.append(offset) + lengths.append(len(completion_ids)) + cursor = offset + len(completion_ids) + return ( + _ChatChoiceTrace(choices=choices, offsets=offsets, lengths=lengths) + if choices + else None + ) + + def _slice_moe_routes( routes: MoeRouteArray | MoeRouteSegments | None, start: int ) -> MoeRouteArray | MoeRouteSegments | None: @@ -549,6 +630,7 @@ def tokenize_trajectory_groups( image_processor: BaseImageProcessor | None = None, chat_template_kwargs: dict[str, Any] | None = None, chat_template_tool_schema_format: ChatTemplateToolSchemaFormat = "default", + model: str | None = None, ) -> Generator["TokenizedResult", None, None]: for group in trajectory_groups: if not group: @@ -568,16 +650,21 @@ def tokenize_trajectory_groups( if advantage == 0 and drop_zero_advantage_trajectories: continue if trajectory.exchanges: - from ..trajectories._tokenize import _as_tokenizer + from ..trajectories._tokenize import ( + _as_tokenizer, + _require_training_model, + ) + selected_model = _require_training_model(trajectory, model) exchange_results = trajectory.tokenize( multi_history=True, + model=selected_model, base_model=tokenizer.name_or_path, tokenizer=_as_tokenizer(tokenizer), chat_template=None, chat_template_kwargs=chat_template_kwargs, ) - histories = trajectory.histories() + histories = trajectory.histories(model=selected_model) trajectory_results = [] for exchange_result, history in zip( exchange_results.histories, histories, strict=True @@ -597,43 +684,29 @@ def tokenize_trajectory_groups( raise RuntimeError( "Exchange trajectory is missing logprobs for trainable tokens" ) - chat_choices = [] + chat_trace = None if isinstance(history, ChatCompletionsHistory): - seen_choices: set[tuple[int, int]] = set() - for source in history.message_sources: - if source is None or source.choice_index is None: - continue - key = (id(source.exchange), source.choice_index) - if key in seen_choices: - continue - seen_choices.add(key) - if not isinstance(source.exchange, ChatCompletionsExchange): - continue - chat_choices.append( - next( - choice - for choice in source.exchange.response.choices - if choice.index == source.choice_index - ) - ) - if len(chat_choices) == len(choice_spans): + chat_trace = _chat_choice_trace( + history, + exchange_result.token_ids, + exchange_result.flags, + ) + if chat_trace is not None: moe_routes, moe_stats = align_choice_routes_to_tokenized_result( token_ids=exchange_result.token_ids, - choices=chat_choices, - choice_offsets=[start for start, _ in choice_spans], - choice_token_lengths=[ - end - start for start, end in choice_spans - ], + choices=chat_trace.choices, + choice_offsets=chat_trace.offsets, + choice_token_lengths=chat_trace.lengths, ) - else: - if any( - choice_moe_routing_metadata(choice) is not None - for choice in chat_choices - ): - raise RuntimeError( - "MoE routing replay requires complete Chat Completions " - "history sources" + choice_spans = [ + (start, start + length) + for start, length in zip( + chat_trace.offsets, + chat_trace.lengths, + strict=True, ) + ] + else: moe_routes = None moe_stats = MoeRoutingAlignmentStats() trajectory_results.append( diff --git a/src/art/tinker_native/backend.py b/src/art/tinker_native/backend.py index d397adc85..6a0998606 100644 --- a/src/art/tinker_native/backend.py +++ b/src/art/tinker_native/backend.py @@ -358,6 +358,7 @@ async def train( state.tokenizer, normalize_advantages, base_model=model.base_model, + model=self._model_inference_name(model), ) metrics: dict[str, float] = { diff --git a/src/art/tinker_native/data.py b/src/art/tinker_native/data.py index 9c1428e9b..500b69e94 100644 --- a/src/art/tinker_native/data.py +++ b/src/art/tinker_native/data.py @@ -142,6 +142,7 @@ def trajectory_groups_to_datums( normalize_advantages: bool = True, *, base_model: str | None = None, + model: str | None = None, ) -> list[tinker.Datum]: datums: list[tinker.Datum] = [] @@ -162,8 +163,13 @@ def trajectory_groups_to_datums( continue for trajectory, advantage in zip(group.trajectories, advantages): if trajectory.exchanges: + from ..trajectories._tokenize import _require_training_model + + selected_model = _require_training_model(trajectory, model) tokenized = trajectory.tokenize( - multi_history=True, base_model=base_model + multi_history=True, + model=selected_model, + base_model=base_model, ) for history in tokenized.histories: datum = _tokenized_trajectory_to_datum(history, advantage) diff --git a/tests/unit/test_exchange_training_model_selection.py b/tests/unit/test_exchange_training_model_selection.py new file mode 100644 index 000000000..176fa2bfb --- /dev/null +++ b/tests/unit/test_exchange_training_model_selection.py @@ -0,0 +1,226 @@ +from datetime import datetime, timedelta +from typing import cast + +import numpy as np +from openai.types.chat import ChatCompletion, ChatCompletionMessageParam +import pytest +from transformers import PreTrainedTokenizerBase + +import art +from art.openai import ART_MOE_ROUTING_METADATA_KEY +from art.preprocessing.tokenize import tokenize_trajectory_groups +from art.tinker_native.data import trajectory_groups_to_datums +from art.trajectories import ChatCompletionsExchange, ChatCompletionsRequest + + +def _exchange(model: str, output_token: int) -> ChatCompletionsExchange: + response = ChatCompletion.model_validate( + { + "id": f"chatcmpl-{output_token}", + "object": "chat.completion", + "created": 1, + "model": model, + "choices": [ + { + "index": 0, + "finish_reason": "stop", + "message": {"role": "assistant", "content": "answer"}, + "prompt_token_ids": [1], + "token_ids": [output_token], + "logprobs": { + "content": [ + { + "token": f"token_id:{output_token}", + "logprob": -0.1, + "bytes": [], + "top_logprobs": [], + } + ] + }, + } + ], + } + ) + start = datetime(2026, 1, 1) + return ChatCompletionsExchange( + request=ChatCompletionsRequest( + model=model, + messages=[{"role": "user", "content": "question"}], + ), + response=response, + start_time=start, + end_time=start + timedelta(seconds=1), + ) + + +def _routed_exchange( + *, + prompt_token_ids: list[int], + output_token: int, + messages: list[ChatCompletionMessageParam], + content: str, +) -> ChatCompletionsExchange: + exchange = _exchange("policy", output_token) + exchange.request["messages"] = messages + choice = exchange.response.choices[0] + choice.message.content = content + extra = choice.model_extra + assert extra is not None + extra["prompt_token_ids"] = prompt_token_ids + extra[ART_MOE_ROUTING_METADATA_KEY] = { + "prompt_token_ids": prompt_token_ids, + "completion_token_ids": [output_token], + "routed_experts": np.asarray( + [[[10]]] * len(prompt_token_ids) + [[[output_token * 10]]], + dtype=np.int32, + ), + } + return exchange + + +def _group() -> art.TrajectoryGroup: + trajectories = [ + art.Trajectory( + exchanges=art.TrajectoryExchanges( + chat_completions=[ + _exchange("policy", 2), + _exchange("judge", 3), + ] + ), + reward=reward, + ) + for reward in (1.0, 0.0) + ] + return art.TrajectoryGroup(trajectories=trajectories) + + +class _Tokenizer: + name_or_path = "base/model" + + +def test_preprocessing_requires_model_selection() -> None: + tokenizer = cast(PreTrainedTokenizerBase, _Tokenizer()) + with pytest.raises(ValueError, match="exactly one model"): + list( + tokenize_trajectory_groups( + tokenizer, + [_group()], + allow_training_without_logprobs=False, + scale_rewards=False, + ) + ) + + results = list( + tokenize_trajectory_groups( + tokenizer, + [_group()], + allow_training_without_logprobs=False, + scale_rewards=False, + model="policy", + ) + ) + assert len(results) == 2 + assert all(result.token_ids == [1, 2] for result in results) + + +def test_tinker_requires_model_selection() -> None: + with pytest.raises(ValueError, match="exactly one model"): + trajectory_groups_to_datums([_group()], renderer=None, tokenizer=None) + + datums = trajectory_groups_to_datums( + [_group()], + renderer=None, + tokenizer=None, + model="policy", + ) + assert len(datums) == 2 + assert all(datum.model_input.to_ints() == [1] for datum in datums) + + +def test_preprocessing_preserves_adjacent_choice_boundaries_for_moe() -> None: + exchanges = [ + _routed_exchange( + prompt_token_ids=[1], + output_token=2, + messages=[{"role": "user", "content": "question"}], + content="first", + ), + _routed_exchange( + prompt_token_ids=[1, 2], + output_token=3, + messages=[ + {"role": "user", "content": "question"}, + {"role": "assistant", "content": "first"}, + {"role": "user", "content": "again"}, + ], + content="second", + ), + ] + group = art.TrajectoryGroup( + trajectories=[ + art.Trajectory( + exchanges=art.TrajectoryExchanges( + chat_completions=exchanges, + ), + reward=reward, + ) + for reward in (1.0, 0.0) + ] + ) + + results = list( + tokenize_trajectory_groups( + cast(PreTrainedTokenizerBase, _Tokenizer()), + [group], + allow_training_without_logprobs=False, + scale_rewards=False, + model="policy", + ) + ) + + assert len(results) == 2 + assert all(result.choice_offsets == [1, 2] for result in results) + assert all(result.moe_routed_experts is not None for result in results) + + +def test_preprocessing_rejects_partial_choice_evidence_before_moe_routes() -> None: + first = _routed_exchange( + prompt_token_ids=[1], + output_token=2, + messages=[{"role": "user", "content": "question"}], + content="first", + ) + first_extra = first.response.choices[0].model_extra + assert first_extra is not None + first_extra.pop("token_ids") + first_extra.pop(ART_MOE_ROUTING_METADATA_KEY) + second = _routed_exchange( + prompt_token_ids=[1, 2], + output_token=3, + messages=[ + {"role": "user", "content": "question"}, + {"role": "assistant", "content": "first"}, + {"role": "user", "content": "again"}, + ], + content="second", + ) + group = art.TrajectoryGroup( + trajectories=[ + art.Trajectory( + exchanges=art.TrajectoryExchanges(chat_completions=[first, second]), + reward=reward, + ) + for reward in (1.0, 0.0) + ] + ) + + with pytest.raises(RuntimeError, match="every sourced choice"): + list( + tokenize_trajectory_groups( + cast(PreTrainedTokenizerBase, _Tokenizer()), + [group], + allow_training_without_logprobs=False, + scale_rewards=False, + model="policy", + ) + ) From 9bfa78e9c49fc08dc5ca36d8a01b0de9baa5c36a Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Fri, 24 Jul 2026 15:31:07 +0000 Subject: [PATCH 04/58] Capture tau-bench rollouts as exchanges --- src/art/tau_bench/rollout.py | 130 ++++++++-------- tests/unit/test_tau_bench_client.py | 222 +++++++++++++++++++++++++--- 2 files changed, 272 insertions(+), 80 deletions(-) diff --git a/src/art/tau_bench/rollout.py b/src/art/tau_bench/rollout.py index 093bd6396..8e058428f 100644 --- a/src/art/tau_bench/rollout.py +++ b/src/art/tau_bench/rollout.py @@ -3,9 +3,10 @@ from collections.abc import Mapping import json import os -from typing import Any, overload +from typing import Any, cast, overload from openai import AsyncOpenAI, BadRequestError +from openai.types.chat import ChatCompletionMessageParam from openai.types.completion_usage import CompletionUsage from art.costs import get_model_pricing, tokens_to_cost @@ -94,12 +95,12 @@ async def rollout( model=model, base_model=base_model, ) + messages: list[ChatCompletionMessageParam] = [ + {"role": "system", "content": env.info["policy"]}, + {"role": "user", "content": env.observation.removeprefix("user: ")}, + ] + tools = env.info.get("tools") or [] trajectory = Trajectory( - messages_and_choices=[ - {"role": "system", "content": env.info["policy"]}, - {"role": "user", "content": env.observation.removeprefix("user: ")}, - ], - tools=env.info.get("tools"), reward=0, metrics={ "cost/tinker/prefill": 0.0, @@ -110,65 +111,76 @@ async def rollout( ) terminated = False num_turns = 0 - while not terminated: - if max_turns is not None and num_turns >= max_turns: - break - try: - chat_completion = await openai_client.chat.completions.create( - messages=trajectory.messages(), - model=model_name, - stream=False, - tool_choice="auto", - tools=trajectory.tools or [], - **chat_completion_kwargs, - ) - except BadRequestError as exc: - if _is_max_tokens_error(exc): + with trajectory: + while not terminated: + if max_turns is not None and num_turns >= max_turns: break - raise - _record_tinker_costs( - trajectory, - cost_model, - chat_completion.usage, - assert_costs=assert_costs, - ) - choice = chat_completion.choices[0] - trajectory.messages_and_choices.append(choice) - tool_calls = getattr(choice.message, "tool_calls", None) - if tool_calls: - for tool_call in tool_calls: - action = _tool_call_action(tool_call) - step = await client.step_environment(env.id, action) - trajectory.messages_and_choices.append( + try: + chat_completion = await openai_client.chat.completions.create( + messages=messages, + model=model_name, + stream=False, + tool_choice="auto", + tools=tools, + **chat_completion_kwargs, + ) + except BadRequestError as exc: + if _is_max_tokens_error(exc): + break + raise + _record_tinker_costs( + trajectory, + cost_model, + chat_completion.usage, + assert_costs=assert_costs, + ) + choice = chat_completion.choices[0] + messages.append( + cast( + ChatCompletionMessageParam, + choice.message.model_dump(exclude_none=True), + ) + ) + tool_calls = getattr(choice.message, "tool_calls", None) + if tool_calls: + for tool_call in tool_calls: + action = _tool_call_action(tool_call) + step = await client.step_environment(env.id, action) + messages.append( + { + "role": "tool", + "content": step.observation.removeprefix("tool: "), + "tool_call_id": tool_call.id, + } + ) + trajectory.reward += step.reward + terminated = step.terminated + else: + step = await client.step_environment( + env.id, + choice.message.content or "", + ) + if "user_message_cost" in step.info: + trajectory.metrics["cost/user"] += step.info[ + "user_message_cost" + ] + elif assert_costs: + raise ValueError("Costs are not supported for the user model") + messages.append( { - "role": "tool", - "content": step.observation.removeprefix("tool: "), - "tool_call_id": tool_call.id, + "role": "user", + "content": step.observation.removeprefix("user: "), } ) trajectory.reward += step.reward terminated = step.terminated - else: - step = await client.step_environment( - env.id, - choice.message.content or "", - ) - if "user_message_cost" in step.info: - trajectory.metrics["cost/user"] += step.info["user_message_cost"] - elif assert_costs: - raise ValueError("Costs are not supported for the user model") - trajectory.messages_and_choices.append( - {"role": "user", "content": step.observation.removeprefix("user: ")} - ) - trajectory.reward += step.reward - terminated = step.terminated - num_turns += 1 - usage = chat_completion.usage - if usage is not None and _would_exceed_context_limit( - usage.total_tokens, - _requested_completion_tokens(chat_completion_kwargs), - ): - break + num_turns += 1 + usage = chat_completion.usage + if usage is not None and _would_exceed_context_limit( + usage.total_tokens, + _requested_completion_tokens(chat_completion_kwargs), + ): + break trajectory.metrics["num_turns"] = num_turns return trajectory diff --git a/tests/unit/test_tau_bench_client.py b/tests/unit/test_tau_bench_client.py index fd8a06783..ae7fb87e0 100644 --- a/tests/unit/test_tau_bench_client.py +++ b/tests/unit/test_tau_bench_client.py @@ -6,7 +6,8 @@ from typing import Any import httpx -from openai.types.completion_usage import CompletionUsage +from openai import AsyncOpenAI +from openai.types.chat import ChatCompletion import pytest import art @@ -251,16 +252,25 @@ async def delete_environment(self, env_id: str) -> DeleteEnvironmentResponse: class FakeCompletions: async def create(self, **kwargs: Any) -> Any: self.kwargs = kwargs - choice = SimpleNamespace( - message=SimpleNamespace(content="hello", tool_calls=None) - ) - return SimpleNamespace( - choices=[choice], - usage=CompletionUsage( - prompt_tokens=10, - completion_tokens=5, - total_tokens=15, - ), + return ChatCompletion.model_validate( + { + "id": "chat-1", + "object": "chat.completion", + "created": 0, + "model": "default", + "choices": [ + { + "index": 0, + "finish_reason": "stop", + "message": {"role": "assistant", "content": "hello"}, + } + ], + "usage": { + "prompt_tokens": 10, + "completion_tokens": 5, + "total_tokens": 15, + }, + } ) @@ -320,6 +330,167 @@ async def test_rollout_supports_art_model_like_args() -> None: assert trajectory.metrics["num_turns"] == 1 +@pytest.mark.asyncio +async def test_rollout_captures_two_turn_tool_exchange_with_exact_tokens() -> None: + rollout_module = importlib.import_module("art.tau_bench.rollout") + rollout_module.openai_clients.clear() + request_bodies: list[dict[str, Any]] = [] + + class ToolTauBenchClient(FakeTauBenchClient): + async def step_environment( + self, env_id: str, action: str + ) -> StepEnvironmentResponse: + if action == "lookup(key='x')": + return StepEnvironmentResponse( + id=env_id, + observation="tool: result", + reward=0.25, + terminated=False, + truncated=False, + info={}, + ) + assert action == "hello" + return StepEnvironmentResponse( + id=env_id, + observation="user: done", + reward=0.75, + terminated=True, + truncated=False, + info={"user_message_cost": 0.25}, + ) + + async def handler(request: httpx.Request) -> httpx.Response: + request_bodies.append(json.loads(request.content)) + if len(request_bodies) == 1: + return httpx.Response( + 200, + json={ + "id": "chat-tool", + "object": "chat.completion", + "created": 0, + "model": "default", + "prompt_token_ids": [10, 11], + "choices": [ + { + "index": 0, + "finish_reason": "tool_calls", + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call-1", + "type": "function", + "function": { + "name": "lookup", + "arguments": '{"key":"x"}', + }, + } + ], + }, + "token_ids": [12], + "logprobs": { + "content": [ + { + "token": "token_id:12", + "logprob": -0.25, + "bytes": None, + "top_logprobs": [], + } + ] + }, + } + ], + "usage": { + "prompt_tokens": 2, + "completion_tokens": 1, + "total_tokens": 3, + }, + }, + ) + return httpx.Response( + 200, + json={ + "id": "chat-exact", + "object": "chat.completion", + "created": 0, + "model": "default", + "prompt_token_ids": [10, 11, 12, 13], + "choices": [ + { + "index": 0, + "finish_reason": "stop", + "message": {"role": "assistant", "content": "hello"}, + "token_ids": [14], + "logprobs": { + "content": [ + { + "token": "token_id:14", + "logprob": -0.5, + "bytes": [104, 101, 108, 108, 111], + "top_logprobs": [], + } + ] + }, + } + ], + "usage": { + "prompt_tokens": 4, + "completion_tokens": 1, + "total_tokens": 5, + }, + }, + ) + + http_client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) + openai_client = AsyncOpenAI( + api_key="model-key", + base_url="http://model.test/v1", + http_client=http_client, + ) + rollout_module.openai_clients[("http://model.test/v1", "model-key")] = openai_client + try: + trajectory = await rollout_module.rollout( + Scenario(domain="banking_knowledge", task=Task(id="task_001")), + "http://model.test/v1", + "model-key", + "default", + client=ToolTauBenchClient(), + max_turns=2, + ) + finally: + await openai_client.close() + await http_client.aclose() + rollout_module.openai_clients.clear() + + assert request_bodies[0]["messages"] == [ + {"role": "system", "content": "policy"}, + {"role": "user", "content": "hello"}, + ] + assert request_bodies[1]["messages"][2]["tool_calls"][0]["id"] == "call-1" + assert request_bodies[1]["messages"][3] == { + "role": "tool", + "content": "result", + "tool_call_id": "call-1", + } + assert len(trajectory.exchanges.chat_completions) == 2 + assert not trajectory.messages_and_choices + assert trajectory.tools is None + restored = art.Trajectory.model_validate_json(trajectory.model_dump_json()) + tokenized = restored.tokenize() + assert tokenized.token_ids == [10, 11, 12, 13, 14] + assert tokenized.logprobs[2] == -0.25 + assert tokenized.logprobs[3] != tokenized.logprobs[3] + assert tokenized.logprobs[4] == -0.5 + assert tokenized.flags == [ + art.TokenFlag.EXACT, + art.TokenFlag.EXACT, + art.TokenFlag.EXACT | art.TokenFlag.SAMPLED, + art.TokenFlag.EXACT, + art.TokenFlag.EXACT | art.TokenFlag.SAMPLED, + ] + + class FakeBadRequestError(Exception): def __init__(self, message: str) -> None: self.message = message @@ -382,16 +553,25 @@ async def test_rollout_stops_on_max_tokens_bad_request( class NearContextLimitCompletions: async def create(self, **kwargs: Any) -> Any: - choice = SimpleNamespace( - message=SimpleNamespace(content="hello", tool_calls=None) - ) - return SimpleNamespace( - choices=[choice], - usage=CompletionUsage( - prompt_tokens=32_000, - completion_tokens=700, - total_tokens=32_700, - ), + return ChatCompletion.model_validate( + { + "id": "chat-near-limit", + "object": "chat.completion", + "created": 0, + "model": "default", + "choices": [ + { + "index": 0, + "finish_reason": "stop", + "message": {"role": "assistant", "content": "hello"}, + } + ], + "usage": { + "prompt_tokens": 32_000, + "completion_tokens": 700, + "total_tokens": 32_700, + }, + } ) From ed09e19dcc8cbfeaad6feac40a5b9ff99298eecd Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Fri, 24 Jul 2026 15:31:13 +0000 Subject: [PATCH 05/58] Harden trajectory history projection --- src/art/trajectories/_history.py | 148 ++++++++++++----- tests/unit/trajectories/test_history.py | 212 ++++++++++++++++++++++++ 2 files changed, 322 insertions(+), 38 deletions(-) diff --git a/src/art/trajectories/_history.py b/src/art/trajectories/_history.py index 52e9feb20..38cfa8eeb 100644 --- a/src/art/trajectories/_history.py +++ b/src/art/trajectories/_history.py @@ -56,7 +56,12 @@ class _ModelledExchange(Protocol): def model(self) -> str | None: ... +class _Indexed(Protocol): + index: int + + _ExchangeT = TypeVar("_ExchangeT", bound=_ModelledExchange) +_IndexedT = TypeVar("_IndexedT", bound=_Indexed) _ItemT = TypeVar("_ItemT") _SourceT = TypeVar("_SourceT") _ContextT = TypeVar("_ContextT") @@ -117,6 +122,7 @@ class _Branch(Generic[_ItemT, _SourceT, _ContextT]): context: _ContextT order: tuple[int, ...] first_time: datetime + context_source: _ModelledExchange | None def _selected_models( @@ -146,6 +152,20 @@ def _one(history: Sequence[_ItemT], protocol: str) -> _ItemT: return history[0] +def _ordered_choices(choices: Sequence[_IndexedT], *, protocol: str) -> list[_IndexedT]: + if not choices: + raise ValueError(f"{protocol} response contains no choices") + indices = [choice.index for choice in choices] + if any( + not isinstance(index, int) or isinstance(index, bool) or index < 0 + for index in indices + ): + raise ValueError(f"{protocol} choice indices must be non-negative integers") + if len(set(indices)) != len(indices): + raise ValueError(f"{protocol} response contains duplicate choice indices") + return sorted(choices, key=lambda choice: choice.index) + + def _is_prefix(prefix: Sequence[object], value: Sequence[object]) -> bool: return len(prefix) <= len(value) and list(value[: len(prefix)]) == list(prefix) @@ -181,6 +201,7 @@ def _extend_branches( context: _ContextT, sequence: int, start_time: datetime, + context_source: _ModelledExchange | None = None, continuation: Callable[[_Branch[_ItemT, _SourceT, _ContextT]], bool] | None = None, ) -> None: if len(prompt) != len(prompt_sources): @@ -228,6 +249,7 @@ def _extend_branches( context=copy.deepcopy(context), order=(*base_order, choice_index), first_time=first_time, + context_source=context_source, ) for choice_index, output, output_sources in outputs ] @@ -289,13 +311,24 @@ def _chat_retains_sampled_reasoning( def _chat_message_key(message: Message, *, visible_only: bool = False) -> str: - data = copy.deepcopy(dict(message)) + data = normalize_chat_message(message) if visible_only: data.pop("reasoning", None) data.pop("reasoning_content", None) return json.dumps(data, sort_keys=True, default=str) +def normalize_chat_message(message: Mapping[str, object]) -> dict[str, object]: + """Return the canonical history form of an OpenAI chat message.""" + + data = copy.deepcopy(dict(message)) + if data.get("content") is None: + data.pop("content", None) + if data.get("tool_calls") == []: + data.pop("tool_calls") + return data + + def _require_unmixed(trajectory: Trajectory) -> None: if trajectory.exchanges and ( trajectory.messages_and_choices @@ -343,7 +376,12 @@ def chat_completions_histories( _Branch[Message, ChatCompletionsMessageSource, _ChatContext] ] = [] for sequence, exchange in enumerate(exchanges): - prompt = _MESSAGES.validate_python(exchange.request.get("messages", [])) + prompt = [ + cast(Message, normalize_chat_message(message)) + for message in _MESSAGES.validate_python( + exchange.request.get("messages", []) + ) + ] prompt_sources = _lineage_prompt_sources( branches, prompt=prompt, @@ -363,11 +401,13 @@ def chat_completions_histories( list[ChatCompletionsMessageSource | None], ] ] = [] - for choice in sorted( - exchange.response.choices, key=lambda item: item.index + for choice in _ordered_choices( + exchange.response.choices, protocol="Chat Completions" ): response = _MESSAGE.validate_python( - choice.message.model_dump(mode="python", exclude_none=True) + normalize_chat_message( + choice.message.model_dump(mode="python", exclude_none=True) + ) ) response_source = ChatCompletionsMessageSource( exchange=exchange, choice_index=choice.index @@ -434,10 +474,7 @@ def anthropic_messages_histories( _Branch[AnthropicMessageParam, AnthropicMessageSource, _AnthropicContext] ] = [] for sequence, exchange in enumerate(exchanges): - prompt = [ - copy.deepcopy(message) - for message in exchange.request.get("messages", []) - ] + prompt = _anthropic_prompt(exchange.request.get("messages", [])) prompt_sources = _lineage_prompt_sources( branches, prompt=prompt, @@ -486,11 +523,11 @@ def anthropic_messages_histories( ), sequence=sequence, start_time=exchange.start_time, + context_source=( + exchange if exchange.request.get("system") is not None else None + ), ) for branch in sorted(branches, key=lambda item: (item.first_time, item.order)): - first_source = next( - (source for source in branch.sources if source is not None), None - ) histories.append( AnthropicMessagesHistory( model=selected_model, @@ -498,7 +535,9 @@ def anthropic_messages_histories( message_sources=copy.copy(branch.sources), system=copy.deepcopy(branch.context.system), system_source=( - first_source.exchange if first_source is not None else None + cast(MessagesExchange, branch.context_source) + if branch.context_source is not None + else None ), tools=copy.deepcopy(branch.context.tools), chat_template=branch.context.template, @@ -508,13 +547,31 @@ def anthropic_messages_histories( return histories +def _anthropic_prompt(value: object) -> list[AnthropicMessageParam]: + if not isinstance(value, list): + raise ValueError("Anthropic messages must be a list") + messages: list[AnthropicMessageParam] = [] + for message in value: + if not isinstance(message, Mapping): + raise ValueError("Anthropic messages must be JSON objects") + content = message.get("content") + if not isinstance(content, (str, list)): + raise ValueError("Anthropic message content must be text or a list") + if isinstance(content, list) and any( + not isinstance(block, (Mapping, pydantic.BaseModel)) for block in content + ): + raise ValueError("Anthropic message content blocks must be JSON objects") + messages.append(cast(AnthropicMessageParam, copy.deepcopy(dict(message)))) + return messages + + def _anthropic_message_key( message: AnthropicMessageParam, *, visible_only: bool = False ) -> tuple[object, ...]: content = message.get("content") if isinstance(content, str): blocks: tuple[object, ...] = (("text", content),) - else: + elif isinstance(content, list): normalized: list[object] = [] for block in content: if isinstance(block, pydantic.BaseModel): @@ -522,7 +579,9 @@ def _anthropic_message_key( elif isinstance(block, Mapping): data = dict(block) else: - data = {"type": type(block).__name__, "value": str(block)} + raise ValueError( + "Anthropic message content blocks must be JSON objects" + ) kind = data.get("type") if visible_only and kind in {"thinking", "redacted_thinking"}: continue @@ -531,6 +590,8 @@ def _anthropic_message_key( else: normalized.append(json.dumps(data, sort_keys=True, default=str)) blocks = tuple(normalized) + else: + raise ValueError("Anthropic message content must be text or a list") return (message.get("role"), *blocks) @@ -678,9 +739,7 @@ def responses_histories( continuation = lambda branch, generation=generation: ( _responses_generation_extends( branch, - exchange=exchange, generation=generation, - generations=generations, ) ) extends = any( @@ -689,7 +748,7 @@ def responses_histories( and continuation(branch) for branch in branches ) - if not extends and generation_index: + if not extends: retained = [ (item, source) for item, source in zip( @@ -700,8 +759,8 @@ def responses_histories( if not ( item.get("type") == "reasoning" and source is not None - and source.exchange is exchange and source.output_index is not None + and source.generation_index is not None ) ] generation_prompt = [item for item, _ in retained] @@ -731,6 +790,7 @@ def responses_histories( context=context, sequence=sequence, start_time=exchange.start_time, + context_source=(exchange if instructions is not None else None), continuation=continuation, ) final_items = [ @@ -750,6 +810,7 @@ def responses_histories( context=context, sequence=sequence, start_time=exchange.start_time, + context_source=exchange if instructions is not None else None, ) record = ( final_items, @@ -763,9 +824,6 @@ def responses_histories( record ) for branch in sorted(branches, key=lambda item: (item.first_time, item.order)): - first_source = next( - (source for source in branch.sources if source is not None), None - ) histories.append( ResponsesHistory( model=selected_model, @@ -773,7 +831,9 @@ def responses_histories( input_sources=copy.copy(branch.sources), instructions=branch.context.instructions, instructions_source=( - first_source.exchange if first_source is not None else None + cast(ResponsesExchange, branch.context_source) + if branch.context_source is not None + else None ), tools=copy.deepcopy(branch.context.tools), conversation=copy.deepcopy(branch.context.conversation), @@ -798,8 +858,15 @@ def _responses_generations( result: list[int | None] = [None] * len(output) parsed: list[_ResponsesGeneration] = [] for generation_index, generation in enumerate(generations): - if generation.prompt_token_ids is None or generation.output_token_ids is None: - raise AssertionError("Validated Responses generation omitted exact tokens") + if generation.prompt_token_ids is None: + raise ValueError( + "Responses generation without an exact prompt cannot be projected" + ) + if generation.output_token_ids is None: + raise ValueError( + "Responses generation without exact output tokens cannot yet be " + "projected as a history" + ) for output_index in generation.output_indices: result[output_index] = generation_index parsed.append( @@ -815,30 +882,37 @@ def _responses_generations( def _responses_generation_extends( branch: _Branch[ResponseInputItemParam, ResponsesItemSource, _ResponsesContext], *, - exchange: ResponsesExchange, generation: _ResponsesGeneration, - generations: Sequence[_ResponsesGeneration], ) -> bool: - prior_generation = next( + prior_source = next( ( - source.generation_index + source for source in reversed(branch.sources) - if source is not None - and source.exchange is exchange - and source.generation_index is not None + if source is not None and source.generation_index is not None ), None, ) - if prior_generation is None: + if prior_source is None or prior_source.generation_index is None: return True - prior = generations[prior_generation] + prior_exchange = prior_source.exchange + prior_output = [ + cast( + ResponseInputItemParam, + item.model_dump(mode="json", exclude_none=True), + ) + for item in prior_exchange.response.output + ] + prior_generations, _ = _responses_generations(prior_exchange, output=prior_output) + if not 0 <= prior_source.generation_index < len(prior_generations): + raise ValueError("Responses generation source index is out of bounds") + prior = prior_generations[prior_source.generation_index] if _is_prefix( [*prior.prompt_token_ids, *prior.output_token_ids], generation.prompt_token_ids, ): return True return not any( - getattr(exchange.response.output[index], "type", None) == "reasoning" + getattr(prior_exchange.response.output[index], "type", None) == "reasoning" for index in prior.output_indices ) @@ -991,9 +1065,7 @@ def _completion_choice_groups( ) -> list[list[CompletionChoice]]: prompts = _completion_prompts(exchange.request.get("prompt")) count = len(prompts) - choices = sorted(exchange.response.choices, key=lambda item: item.index) - if not choices: - raise ValueError("Completions response contains no choices") + choices = _ordered_choices(exchange.response.choices, protocol="Completions") if count == 1: return [choices] requested_n = exchange.request.get("n") diff --git a/tests/unit/trajectories/test_history.py b/tests/unit/trajectories/test_history.py index da6066981..7a3794afa 100644 --- a/tests/unit/trajectories/test_history.py +++ b/tests/unit/trajectories/test_history.py @@ -13,6 +13,7 @@ ChatCompletionsExchange, CompletionsExchange, MessagesExchange, + MessagesRequest, ResponsesExchange, TrajectoryExchanges, ) @@ -280,6 +281,47 @@ def test_chat_choices_branch_and_identical_continuation_uses_first_choice() -> N assert histories[1].message_sources[1].choice_index == 1 +def test_chat_history_normalizes_empty_response_only_fields() -> None: + first = _chat([{"role": "user", "content": "one"}], "first") + data = first.response.model_dump(mode="python") + data["choices"][0]["message"]["tool_calls"] = [] + first.response = ChatCompletion.model_validate(data) + second = _chat( + [ + {"role": "user", "content": "one"}, + {"role": "assistant", "content": "first"}, + {"role": "user", "content": "two"}, + ], + "second", + offset=1, + ) + + history = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[first, second]) + ).chat_completions_history() + + assert len(history.messages) == 4 + assert "tool_calls" not in history.messages[1] + assert history.message_sources[1] is not None + assert history.message_sources[1].exchange is first + + +@pytest.mark.parametrize("indices", [[], [0, 0]]) +def test_chat_history_rejects_missing_or_duplicate_choice_indices( + indices: list[int], +) -> None: + exchange = _chat([{"role": "user", "content": "one"}], "first") + data = exchange.response.model_dump(mode="python") + choice = data["choices"][0] + data["choices"] = [{**choice, "index": index} for index in indices] + exchange.response = ChatCompletion.model_validate(data) + + with pytest.raises(ValueError, match="choices|choice indices"): + art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[exchange]) + ).chat_completions_histories() + + def test_same_content_seeded_inputs_remain_request_sourced() -> None: first_chat = _chat([{"role": "user", "content": "prompt"}], "same") seeded_chat = _chat([{"role": "assistant", "content": "same"}], "next", offset=1) @@ -494,6 +536,40 @@ def test_responses_history_propagates_opaque_context_and_first_sources() -> None assert history.input_sources[1].output_index == 0 +def test_branch_context_sources_follow_the_request_that_supplied_the_context() -> None: + first_message = _message() + second_message = _message() + second_message.start_time, second_message.end_time = _times(1) + second_message.request["system"] = "New instructions" + second_message.request["messages"] = [ + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Hi"}, + {"role": "user", "content": "Again"}, + ] + message_histories = art.Trajectory( + exchanges=TrajectoryExchanges(messages=[first_message, second_message]) + ).anthropic_messages_histories() + + assert message_histories[-1].system == "New instructions" + assert message_histories[-1].system_source is second_message + + first_response = _response("response-1", "first") + first_response.request["instructions"] = "Old instructions" + second_response = _response( + "response-2", + "second", + previous_response_id="response-1", + offset=1, + ) + second_response.request["instructions"] = "New instructions" + response_histories = art.Trajectory( + exchanges=TrajectoryExchanges(responses=[first_response, second_response]) + ).responses_histories() + + assert response_histories[-1].instructions == "New instructions" + assert response_histories[-1].instructions_source is second_response + + def test_responses_history_maps_and_validates_generation_sources() -> None: exchange = _response("response-1", "first", reasoning="think") assert exchange.response.__pydantic_extra__ is not None @@ -534,6 +610,49 @@ def test_responses_history_maps_and_validates_generation_sources() -> None: trajectory.responses_history() +def test_cross_exchange_responses_reasoning_stripping_splits_histories() -> None: + first = _response("response-1", "first", reasoning="think") + first_data = first.response.model_dump(mode="python") + first_data["token_generations"] = [ + { + "prompt_token_ids": [1], + "output_tokens": [ + {"token_id": 2, "logprob": -0.2}, + {"token_id": 3, "logprob": -0.3}, + ], + "output_indices": [0, 1], + } + ] + first.response = Response.model_validate(first_data) + second = _response( + "response-2", + "second", + previous_response_id="response-1", + offset=1, + ) + second_data = second.response.model_dump(mode="python") + second_data["token_generations"] = [ + { + "prompt_token_ids": [1, 3, 4], + "output_tokens": [{"token_id": 5, "logprob": -0.5}], + "output_indices": [0], + } + ] + second.response = Response.model_validate(second_data) + + histories = art.Trajectory( + exchanges=TrajectoryExchanges(responses=[first, second]) + ).responses_histories() + + assert len(histories) == 2 + assert any(item.get("type") == "reasoning" for item in histories[0].input) + assert all(item.get("type") != "reasoning" for item in histories[1].input) + first_answer_source = histories[1].input_sources[1] + assert first_answer_source is not None + assert first_answer_source.exchange is first + assert first_answer_source.generation_index == 0 + + @pytest.mark.parametrize( ("second_prompt", "reasoning", "expected_histories"), [([1, 2, 3], False, 1), ([1, 3], True, 2)], @@ -725,6 +844,57 @@ def test_completions_reject_ambiguous_batches_and_suffix() -> None: ).histories() +def test_completions_reject_duplicate_choice_indices() -> None: + exchange = _completion([1], [10]) + data = exchange.response.model_dump(mode="python") + data["choices"].append(dict(data["choices"][0])) + exchange.response = Completion.model_validate(data) + + with pytest.raises(ValueError, match="choice indices"): + art.Trajectory( + exchanges=TrajectoryExchanges(completions=[exchange]) + ).completions_token_histories() + + +def test_malformed_anthropic_content_raises_value_error() -> None: + exchange = _message() + exchange.request = cast( + MessagesRequest, + {"messages": [{"role": "user", "content": None}]}, + ) + + with pytest.raises(ValueError, match="Anthropic message content"): + art.Trajectory( + exchanges=TrajectoryExchanges(messages=[exchange]) + ).anthropic_messages_histories() + + +def test_tokenless_responses_generation_raises_value_error() -> None: + exchange = _response("response-1", "first") + data = exchange.response.model_dump(mode="python") + data["output"] = [ + { + "id": "tool-output", + "type": "function_call_output", + "call_id": "call-1", + "output": "result", + "status": "completed", + } + ] + data["token_generations"] = [ + { + "prompt_token_ids": [1], + "output_indices": [0], + } + ] + exchange.response = Response.model_validate(data) + + with pytest.raises(ValueError, match="without exact output tokens"): + art.Trajectory( + exchanges=TrajectoryExchanges(responses=[exchange]) + ).responses_histories() + + def test_reasoning_stripping_produces_truthful_history_per_generation() -> None: def exchange( offset: int, @@ -835,6 +1005,48 @@ def test_chat_template_stripped_reasoning_splits_exact_histories() -> None: trajectory.tokenize() +def test_reasoning_stripped_tool_call_keeps_first_sampled_source() -> None: + first = _chat([{"role": "user", "content": "one"}], "") + first_data = first.response.model_dump(mode="python") + first_data["choices"][0]["message"] = { + "role": "assistant", + "content": None, + "reasoning": "thought", + "tool_calls": [ + { + "id": "call-1", + "type": "function", + "function": {"name": "lookup", "arguments": "{}"}, + } + ], + } + first.response = ChatCompletion.model_validate(first_data) + second = _chat( + [ + {"role": "user", "content": "one"}, + { + "role": "assistant", + "content": None, + "tool_calls": first_data["choices"][0]["message"]["tool_calls"], + }, + {"role": "user", "content": "two"}, + ], + "second", + offset=1, + ) + + histories = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[first, second]) + ).chat_completions_histories() + + assert len(histories) == 2 + source = histories[1].message_sources[1] + assert source is not None + assert source.exchange is first + assert source.choice_index == 0 + assert source.request_index is None + + def test_history_rejects_mutated_mixed_representation() -> None: trajectory = art.Trajectory( messages_and_choices=[{"role": "user", "content": "hi"}] From 4643d77d2f7df03f96d03984272b4395499d7882 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Fri, 24 Jul 2026 15:31:16 +0000 Subject: [PATCH 06/58] Preserve exact evidence across history splits --- src/art/trajectories/_tokenize.py | 360 ++++++++++++++++--- tests/unit/trajectories/test_tokenize.py | 434 ++++++++++++++++++++++- 2 files changed, 741 insertions(+), 53 deletions(-) diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index 25aaf1184..066466248 100644 --- a/src/art/trajectories/_tokenize.py +++ b/src/art/trajectories/_tokenize.py @@ -278,15 +278,22 @@ def _completion_evidence( prompt_logprobs: list[float] = [] completion_logprobs = pair_logprobs if echo and prompt_ids is not None and selected is not None: - if selected[: len(prompt_ids)] == prompt_ids: - if pair_logprobs and len(pair_logprobs) != len(selected): - raise ValueError("Completions token IDs and logprobs differ in length") + pair_includes_prompt = token_ids is not None and pair_ids == [ + *prompt_ids, + *token_ids, + ] + if pair_includes_prompt: prompt_logprobs = pair_logprobs[: len(prompt_ids)] - selected = selected[len(prompt_ids) :] completion_logprobs = pair_logprobs[len(prompt_ids) :] elif len(pair_logprobs) == len(prompt_ids) + len(selected): prompt_logprobs = pair_logprobs[: len(prompt_ids)] completion_logprobs = pair_logprobs[len(prompt_ids) :] + elif selected[: len(prompt_ids)] == prompt_ids: + if pair_logprobs and len(pair_logprobs) != len(selected): + raise ValueError("Completions token IDs and logprobs differ in length") + prompt_logprobs = pair_logprobs[: len(prompt_ids)] + selected = selected[len(prompt_ids) :] + completion_logprobs = pair_logprobs[len(prompt_ids) :] elif pair_logprobs and len(pair_logprobs) != len(selected): raise ValueError("Completions token IDs and logprobs differ in length") elif selected is not None and pair_logprobs and len(pair_logprobs) != len(selected): @@ -1201,6 +1208,28 @@ def _warn_prefix_retokenization() -> None: ) +def _retained_output_suffix( + *, + prompt: Sequence[int], + output: Sequence[int], + logprobs: Sequence[float], + later_prompt: Sequence[int], +) -> tuple[list[int], list[float]] | None: + if list(later_prompt[: len(prompt)]) != list(prompt): + return None + continuation = later_prompt[len(prompt) :] + for start in range(len(output)): + suffix = list(output[start:]) + if list(continuation[: len(suffix)]) == suffix: + return ( + suffix, + list(logprobs[start:]) + if len(logprobs) == len(output) + else [math.nan] * len(suffix), + ) + return None + + def _align_visible_logprobs( tokenizer: Tokenizer | None, completion: list[int], exchange: Exchange ) -> list[float] | None: @@ -1468,6 +1497,10 @@ def fallback_config() -> _TokenizerConfig: tokenizer, ) if repaired is None: + if previous_render_state is None: + raise ValueError( + "Inference prompts do not form one append-only history" + ) current_render = _template_ids( tokenizer, exchange, @@ -1477,31 +1510,33 @@ def fallback_config() -> _TokenizerConfig: chat_template_kwargs=chat_template_kwargs, messages_override=messages_override, ) - if previous_render_state is None: - repaired = [*token_ids, *current_render] - else: - previous_exchange, previous_messages = previous_render_state - previous_render = _template_ids( - tokenizer, - previous_exchange, - completed=True, - config=resolved_config, - chat_template=chat_template, - chat_template_kwargs=chat_template_kwargs, - messages_override=previous_messages, - ) - previous_canonical = _preserve_sampled_prefix( - previous_render, - token_ids, - sampled_outputs, - tokenizer, - ) - suffix = ( - current_render[len(previous_render) :] - if current_render[: len(previous_render)] == previous_render - else current_render + previous_exchange, previous_messages = previous_render_state + previous_render = _template_ids( + tokenizer, + previous_exchange, + completed=True, + config=resolved_config, + chat_template=chat_template, + chat_template_kwargs=chat_template_kwargs, + messages_override=previous_messages, + ) + previous_canonical = _preserve_sampled_prefix( + previous_render, + token_ids, + sampled_outputs, + tokenizer, + ) + if ( + previous_canonical is None + or current_render[: len(previous_render)] != previous_render + ): + raise ValueError( + "Rendered inference prompts do not form one append-only history" ) - repaired = [*(previous_canonical or token_ids), *suffix] + repaired = [ + *previous_canonical, + *current_render[len(previous_render) :], + ] prompt = repaired prompt_is_exact = False _warn_prefix_retokenization() @@ -1659,6 +1694,8 @@ def _chat_choice(source: object) -> Choice: def _validate_history_sources(history: History) -> None: if isinstance(history, ChatCompletionsHistory): + from ._history import normalize_chat_message + if len(history.messages) != len(history.message_sources): raise ValueError("messages and message_sources differ in length") for message, source in zip( @@ -1825,7 +1862,10 @@ def _validate_history_sources(history: History) -> None: } ) ) - if not any(dict(message) == candidate for candidate in expected): + actual = normalize_chat_message(message) + if not any( + actual == normalize_chat_message(candidate) for candidate in expected + ): raise ValueError( "Chat Completions history no longer matches its source exchange" ) @@ -2150,11 +2190,14 @@ def _tokenize_exact_responses_history( tokenizer: Tokenizer | None, ) -> TokenizedHistory | None: generation_keys: list[tuple[ResponsesExchange, int]] = [] + retained_output_indices: dict[tuple[int, int], set[int]] = {} seen: set[tuple[int, int]] = set() for source in history.input_sources: if source is None or source.generation_index is None: continue key = (id(source.exchange), source.generation_index) + if source.output_index is not None: + retained_output_indices.setdefault(key, set()).add(source.output_index) if key not in seen: seen.add(key) generation_keys.append((source.exchange, source.generation_index)) @@ -2165,7 +2208,7 @@ def _tokenize_exact_responses_history( logprobs: list[float] = [] flags: list[TokenFlag] = [] sampled_outputs: list[_SampledOutput] = [] - for exchange, generation_index in generation_keys: + for position, (exchange, generation_index) in enumerate(generation_keys): generations = _response_generations(exchange.response) if not 0 <= generation_index < len(generations): raise ValueError("Responses source generation index is out of bounds") @@ -2174,6 +2217,32 @@ def _tokenize_exact_responses_history( output = generation.output_token_ids if prompt is None or output is None: return None + retained = retained_output_indices.get((id(exchange), generation_index), set()) + if retained != set(generation.output_indices): + if position + 1 >= len(generation_keys): + return None + next_exchange, next_generation_index = generation_keys[position + 1] + next_generations = _response_generations(next_exchange.response) + if not 0 <= next_generation_index < len(next_generations): + raise ValueError("Responses source generation index is out of bounds") + next_prompt = next_generations[next_generation_index].prompt_token_ids + if next_prompt is None: + return None + retained_suffix = _retained_output_suffix( + prompt=prompt, + output=output, + logprobs=generation.output_logprobs, + later_prompt=next_prompt, + ) + if retained_suffix is None: + return None + output, output_logprobs = retained_suffix + output_text = None + else: + output_logprobs = generation.output_logprobs + output_text = generation.output_text or _response_generation_text( + exchange.response, generation + ) if not token_ids: token_ids.extend(prompt) logprobs.extend([math.nan] * len(prompt)) @@ -2204,12 +2273,11 @@ def _tokenize_exact_responses_history( logprobs.extend([math.nan] * len(suffix)) flags.extend([TokenFlag.EXACT] * len(suffix)) token_ids.extend(output) - logprobs.extend(generation.output_logprobs) + logprobs.extend(output_logprobs) flags.extend([TokenFlag.EXACT | TokenFlag.SAMPLED] * len(output)) sampled_outputs.append( _SampledOutput( - text=generation.output_text - or _response_generation_text(exchange.response, generation), + text=output_text, token_ids=list(output), start=len(token_ids) - len(output), ) @@ -2260,6 +2328,30 @@ def _chat_source_full_tokens( return None, [] +def _chat_source_prompt_tokens(source: object) -> list[int] | None: + exchange = getattr(source, "exchange", None) + if isinstance(exchange, ChatCompletionsExchange): + choice_index = getattr(source, "choice_index", None) + if choice_index is None: + return None + choice = next( + item for item in exchange.response.choices if item.index == choice_index + ) + prompt, _, _ = _chat_choice_tokens(choice, _dump(exchange.response)) + return prompt + if isinstance(exchange, ResponsesExchange): + generation_index = getattr(source, "generation_index", None) + generations = _response_generations(exchange.response) + if isinstance(generation_index, int): + if not 0 <= generation_index < len(generations): + raise ValueError("Responses source generation index is out of bounds") + return generations[generation_index].prompt_token_ids + return _responses_tokens(exchange.response)[0] + if isinstance(exchange, MessagesExchange): + return _messages_tokens(exchange.response)[0] + return None + + def _source_is_sampled(source: object) -> bool: exchange = getattr(source, "exchange", None) if isinstance(exchange, ChatCompletionsExchange): @@ -2274,6 +2366,91 @@ def _source_is_sampled(source: object) -> bool: return False +def _tokenize_exact_projected_chat_history( + history: ChatCompletionsHistory, +) -> TokenizedHistory | None: + if not _history_matches_projection(history): + return None + sampled_sources: list[object] = [] + seen: set[tuple[object, ...]] = set() + for source in history.message_sources: + signature = _source_signature(source) + if ( + source is not None + and signature is not None + and signature not in seen + and _source_is_sampled(source) + ): + seen.add(signature) + sampled_sources.append(source) + if not sampled_sources: + return None + + final_source = sampled_sources[-1] + final_prompt = _chat_source_prompt_tokens(final_source) + final_output, final_logprobs = _chat_source_full_tokens(final_source) + if final_prompt is None or final_output is None: + return None + + token_ids = [*final_prompt, *final_output] + logprobs = [ + *([math.nan] * len(final_prompt)), + *( + final_logprobs + if len(final_logprobs) == len(final_output) + else [math.nan] * len(final_output) + ), + ] + flags = [ + *([TokenFlag.EXACT] * len(final_prompt)), + *([TokenFlag.EXACT | TokenFlag.SAMPLED] * len(final_output)), + ] + for index, source in enumerate(sampled_sources[:-1]): + prompt = _chat_source_prompt_tokens(source) + output, output_logprobs = _chat_source_full_tokens(source) + if ( + prompt is None + or output is None + or list(final_prompt[: len(prompt)]) != prompt + ): + return None + retained = next( + ( + evidence + for later_source in sampled_sources[index + 1 :] + if (later_prompt := _chat_source_prompt_tokens(later_source)) + is not None + and ( + evidence := _retained_output_suffix( + prompt=prompt, + output=output, + logprobs=output_logprobs, + later_prompt=later_prompt, + ) + ) + is not None + ), + None, + ) + if retained is None: + return None + retained_ids, retained_logprobs = retained + start = len(prompt) + end = start + len(retained_ids) + if final_prompt[start:end] != retained_ids: + return None + flags[start:end] = [TokenFlag.EXACT | TokenFlag.SAMPLED] * len(retained_ids) + logprobs[start:end] = retained_logprobs + if history.model is None: + raise ValueError("History tokenization requires a model") + return TokenizedHistory( + model=history.model, + token_ids=token_ids, + logprobs=logprobs, + flags=flags, + ) + + def _chat_message_parts(message: Mapping[str, object]) -> list[tuple[str, str]]: parts: list[tuple[str, str]] = [] reasoning = message.get("reasoning") or message.get("reasoning_content") @@ -2429,6 +2606,8 @@ def _tokenize_chat_view( if not history.model and base_model is None: raise ValueError("History tokenization requires a model or base_model") tokenizer = _load_tokenizer(config) + assert tokenizer is not None + resolved_tokenizer = tokenizer messages = [dict(message) for message in history.messages] kwargs = { **(config.chat_template_kwargs or {}), @@ -2455,13 +2634,14 @@ def _tokenize_chat_view( ) ) - def locate(needle: Sequence[int], start: int) -> tuple[int, int] | None: + def locations(needle: Sequence[int], start: int) -> list[tuple[int, int]]: if not needle: - return None - for index in range(start, len(rendered) - len(needle) + 1): - if rendered[index : index + len(needle)] == list(needle): - return index, index + len(needle) - return None + return [] + return [ + (index, index + len(needle)) + for index in range(start, len(rendered) - len(needle) + 1) + if rendered[index : index + len(needle)] == list(needle) + ] replacements: list[tuple[int, int, list[int], list[float], bool]] = [] search_cursor = 0 @@ -2475,6 +2655,26 @@ def locate(needle: Sequence[int], start: int) -> tuple[int, int] | None: and _source_is_sampled(source) for _, text in _chat_message_parts(message) } + part_ids_cache: dict[str, list[int]] = {} + + def part_ids(text: str) -> list[int]: + if text not in part_ids_cache: + part_ids_cache[text] = _ids( + resolved_tokenizer(text, add_special_tokens=False) + ) + return part_ids_cache[text] + + sampled_part_ids = [ + part_ids(text) + for message, source in zip( + history.messages, history.message_sources, strict=True + ) + if message.get("role") == "assistant" + and source is not None + and _source_is_sampled(source) + for _, text in _chat_message_parts(message) + ] + sampled_part_cursor = 0 for message, source in zip(history.messages, history.message_sources, strict=True): parts = _chat_message_parts(message) sampled = ( @@ -2485,7 +2685,35 @@ def locate(needle: Sequence[int], start: int) -> tuple[int, int] | None: full_exact, full_logprobs = ( _chat_source_full_tokens(source) if sampled else (None, []) ) - if sampled and full_exact and (span := locate(full_exact, search_cursor)): + source_boundary = False + if sampled and source is not None: + source_prompt = _chat_source_prompt_tokens(source) + source_exchange = getattr(source, "exchange", None) + source_context_matches = ( + isinstance(source_exchange, ChatCompletionsExchange) + and history.tools == source_exchange.request.get("tools") + and history.chat_template + == source_exchange.request.get("chat_template") + and history.chat_template_kwargs + == source_exchange.request.get("chat_template_kwargs") + ) + if ( + source_context_matches + and source_prompt + and rendered[: len(source_prompt)] == source_prompt + ): + search_cursor = max(search_cursor, len(source_prompt)) + source_boundary = True + full_matches = ( + locations(full_exact, search_cursor) if sampled and full_exact else [] + ) + if ( + full_exact is not None + and full_matches + and source_boundary + and full_matches[0][0] == search_cursor + ): + span = full_matches[0] start, end = span replacements.append( ( @@ -2499,28 +2727,59 @@ def locate(needle: Sequence[int], start: int) -> tuple[int, int] | None: ) ) search_cursor = end + sampled_part_cursor += len(parts) continue replacement_start = len(replacements) for part, text in parts: if not sampled and text not in sampled_texts: continue - local = _ids(tokenizer(text, add_special_tokens=False)) + local = part_ids(text) + same_sampled_parts = ( + sum( + candidate == local + for candidate in sampled_part_ids[sampled_part_cursor:] + ) + if sampled + else 0 + ) + if sampled: + sampled_part_cursor += 1 if not local: continue - span = locate(local, search_cursor) + local_matches = locations(local, search_cursor) + span = local_matches[0] if local_matches else None if span is None: if not sampled: continue raise ValueError( "Could not locate a history message in the rendered history" ) + if ( + sampled + and not source_boundary + and len(local_matches) != same_sampled_parts + ): + raise ValueError( + "Could not uniquely locate a sampled history message in the " + "rendered history" + ) start, end = span search_cursor = end if not sampled: continue assert source is not None - exact, logprobs = _chat_source_tokens(source, text, part=part) + exact, logprobs = _chat_source_tokens( + source, + text, + part=part, + ) + if exact is not None and rendered[start : start + len(exact)] == exact: + end = start + len(exact) + search_cursor = end + elif exact is not None and exact[: len(local)] == local: + exact = exact[: len(local)] + logprobs = logprobs[: len(local)] replacement = exact if exact is not None else local if exact is None and not logprobs: exchange = getattr(source, "exchange", None) @@ -2545,6 +2804,17 @@ def locate(needle: Sequence[int], start: int) -> tuple[int, int] | None: sampled and message_replacements and all(part == "tool_call" for part, _ in parts) + and not ( + all(replacement[4] for replacement in message_replacements) + and all( + left[1] == right[0] + for left, right in zip( + message_replacements, + message_replacements[1:], + strict=False, + ) + ) + ) ): start = message_replacements[0][0] end = message_replacements[-1][1] @@ -2880,6 +3150,10 @@ def tokenize_history( ): return exact if isinstance(history, ChatCompletionsHistory) and (needs_render): + if not override_requires_render and ( + exact := _tokenize_exact_projected_chat_history(history) + ): + return exact return _tokenize_chat_view( history, base_model=base_model, diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index af81b6558..c85f128b7 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -161,6 +161,21 @@ def import_without_tokenizer_dependencies(name: str, *args: Any, **kwargs: Any): assert tokenized.logprobs[3] == -0.4 +def test_empty_tool_calls_normalization_preserves_exact_continuation() -> None: + first = _chat_exchange([1], [2]) + first_data = first.response.model_dump(mode="python") + first_data["choices"][0]["message"]["tool_calls"] = [] + first.response = ChatCompletion.model_validate(first_data) + second = _chat_exchange([1, 2, 3], [4], offset=1) + + tokenized = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[first, second]) + ).tokenize() + + assert tokenized.token_ids == [1, 2, 3, 4] + assert tokenized.logprobs[1::2] == [-0.2, -0.4] + + def test_messages_exact_prompt_and_output_do_not_load_a_tokenizer( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -322,6 +337,31 @@ def test_completions_echo_preserves_prompt_logprobs_without_sampling_them( ] +def test_completions_echo_prefers_full_logprob_carrier_to_id_prefix_heuristic() -> None: + exchange = _completion_exchange(echo=True) + payload = exchange.response.model_dump(mode="python") + payload["choices"][0]["token_ids"] = [1, 2] + payload["choices"][0]["logprobs"] = { + "tokens": ["token_id:1", "token_id:1", "token_id:2"], + "token_logprobs": [-0.1, -0.2, -0.3], + "top_logprobs": [{}, {}, {}], + "text_offset": [0, 8, 9], + } + exchange.response = Completion.model_validate(payload) + + tokenized = art.Trajectory( + exchanges=TrajectoryExchanges(completions=[exchange]) + ).tokenize() + + assert tokenized.token_ids == [1, 1, 2] + assert tokenized.logprobs == [-0.1, -0.2, -0.3] + assert tokenized.flags == [ + art.TokenFlag.EXACT, + art.TokenFlag.EXACT | art.TokenFlag.SAMPLED, + art.TokenFlag.EXACT | art.TokenFlag.SAMPLED, + ] + + def test_completions_exact_ids_accept_textual_logprobs() -> None: exchange = _completion_exchange() payload = exchange.response.model_dump(mode="python") @@ -920,6 +960,7 @@ class Tokenizer: def __call__(self, text: str, *, add_special_tokens: bool = False) -> list[int]: assert not add_special_tokens return { + "one": [10], "first": [50], "second": [60], "thought-1": [70], @@ -1321,6 +1362,57 @@ def _response_exchange( ) +def test_cross_exchange_responses_reasoning_split_uses_later_prompt_backbone() -> None: + first = _response_exchange("response-1", 3, prompt_token_ids=[1]) + first_data = first.response.model_dump(mode="python") + first_data["output"] = [ + { + "id": "reasoning-response-1", + "type": "reasoning", + "summary": [{"type": "summary_text", "text": "think"}], + }, + first_data["output"][0], + ] + first_data["token_generations"] = [ + { + "prompt_token_ids": [1], + "output_tokens": [ + {"token_id": 2, "logprob": -0.2}, + {"token_id": 3, "logprob": -0.3}, + ], + "output_indices": [0, 1], + } + ] + first.response = Response.model_validate(first_data) + second = _response_exchange( + "response-2", + 5, + previous_response_id="response-1", + offset=1, + prompt_token_ids=[1, 3, 4], + ) + trajectory = art.Trajectory( + exchanges=TrajectoryExchanges(responses=[first, second]) + ) + + tokenized = trajectory.tokenize(multi_history=True) + + assert [history.token_ids for history in tokenized.histories] == [ + [1, 2, 3], + [1, 3, 4, 5], + ] + assert math.isnan(tokenized.histories[1].logprobs[0]) + assert tokenized.histories[1].logprobs[1] == -0.3 + assert math.isnan(tokenized.histories[1].logprobs[2]) + assert tokenized.histories[1].logprobs[3] == -0.1 + assert tokenized.histories[1].flags == [ + art.TokenFlag.EXACT, + art.TokenFlag.EXACT | art.TokenFlag.SAMPLED, + art.TokenFlag.EXACT, + art.TokenFlag.EXACT | art.TokenFlag.SAMPLED, + ] + + def _response_with_content_logprobs(*, exact_second: bool) -> ResponsesExchange: exchange = _response_exchange("response-content-logprobs", 0) data = exchange.response.model_dump(mode="python") @@ -1861,9 +1953,8 @@ def test_template_change_rerenders_scaffold_but_preserves_sampled_output() -> No class Tokenizer: def __call__(self, text: str, *, add_special_tokens: bool = False) -> list[int]: - assert text == "answer" assert not add_special_tokens - return [20] + return {"turn 0": [10], "answer": [20]}[text] def apply_chat_template( self, @@ -1990,8 +2081,138 @@ def apply_chat_template( assert tokenized.logprobs[2] == -0.7 +def test_rerender_does_not_bind_sampled_ids_to_token_equivalent_user_text() -> None: + exchange = _chat_exchange([7], [7]) + exchange.request["messages"] = [{"role": "user", "content": "cat"}] + exchange.response.choices[0].message.content = "dog" + history = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[exchange]) + ).chat_completions_history() + history.chat_template = "rerender" + + class Tokenizer: + def __call__(self, text: str, **kwargs: object) -> list[int]: + del text, kwargs + return [7] + + def apply_chat_template( + self, messages: list[dict[str, Any]], **kwargs: object + ) -> list[int]: + del messages, kwargs + return [7, 99, 7] + + with pytest.raises(ValueError, match="uniquely locate"): + history.tokenize(tokenizer=Tokenizer()) + + +def test_rerender_does_not_bind_unique_exact_id_outside_sampled_message() -> None: + exchange = _chat_exchange([42], [7]) + exchange.request["messages"] = [{"role": "user", "content": "cat"}] + exchange.response.choices[0].message.content = "dog" + history = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[exchange]) + ).chat_completions_history() + history.chat_template = "rerender" + + class Tokenizer: + def __call__(self, text: str, **kwargs: object) -> list[int]: + del kwargs + return {"cat": [7], "dog": [500]}[text] + + def apply_chat_template( + self, messages: list[dict[str, Any]], **kwargs: object + ) -> list[int]: + del messages, kwargs + return [42, 7, 99, 500] + + tokenized = history.tokenize(tokenizer=Tokenizer()) + + assert tokenized.token_ids == [42, 7, 99, 7] + assert not tokenized.flags[1] & art.TokenFlag.SAMPLED + assert tokenized.flags[-1] == (art.TokenFlag.EXACT | art.TokenFlag.SAMPLED) + assert math.isnan(tokenized.logprobs[1]) + assert tokenized.logprobs[-1] == -0.7 + + +def test_rerender_rejects_sampled_text_ambiguous_with_trailing_scaffold() -> None: + exchange = _chat_exchange([42], [7]) + exchange.request["messages"] = [{"role": "user", "content": "cat"}] + exchange.response.choices[0].message.content = "dog" + history = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[exchange]) + ).chat_completions_history() + history.chat_template = "rerender" + + class Tokenizer: + def __call__(self, text: str, **kwargs: object) -> list[int]: + del kwargs + return {"cat": [1], "dog": [500]}[text] + + def apply_chat_template( + self, messages: list[dict[str, Any]], **kwargs: object + ) -> list[int]: + del messages, kwargs + return [100, 500, 99, 500] + + with pytest.raises(ValueError, match="uniquely locate"): + history.tokenize(tokenizer=Tokenizer()) + + +def test_rerender_does_not_duplicate_sampled_trailing_eos() -> None: + exchange = _chat_exchange([1], [7, 2]) + exchange.request["messages"] = [{"role": "user", "content": "question"}] + exchange.response.choices[0].message.content = "answer" + history = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[exchange]) + ).chat_completions_history() + history.chat_template = "rerender" + + class Tokenizer: + def __call__(self, text: str, **kwargs: object) -> list[int]: + del kwargs + return {"question": [1], "answer": [7]}[text] + + def apply_chat_template( + self, messages: list[dict[str, Any]], **kwargs: object + ) -> list[int]: + del messages, kwargs + return [1, 99, 7, 2] + + tokenized = history.tokenize(tokenizer=Tokenizer()) + + assert tokenized.token_ids == [1, 99, 7, 2] + assert tokenized.token_ids.count(2) == 1 + assert tokenized.flags[-2:] == [ + art.TokenFlag.EXACT | art.TokenFlag.SAMPLED, + art.TokenFlag.EXACT | art.TokenFlag.SAMPLED, + ] + assert tokenized.logprobs[-2:] == [-0.7, -0.2] + + +def test_full_render_repair_rejects_non_append_only_templates() -> None: + first = _chat_exchange([1], [2]) + second = _chat_exchange([9, 8, 7], [3], offset=1) + trajectory = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[first, second]) + ) + + class Tokenizer: + def __call__(self, text: str, **kwargs: object) -> list[int]: + del kwargs + return {"answer": [2]}[text] + + def apply_chat_template( + self, messages: list[dict[str, Any]], **kwargs: object + ) -> list[int]: + del kwargs + return [100, len(messages)] + + with pytest.raises(ValueError, match="do not form one append-only history"): + trajectory.tokenize(tokenizer=Tokenizer()) + + def test_reasoning_stripped_chat_histories_tokenize_authoritative_views() -> None: - first = _chat_exchange([1], [2, 3]) + first = _chat_exchange([1], [2, 101, 102, 9]) first.request["messages"] = [{"role": "user", "content": "one"}] first_data = first.response.model_dump(mode="python") first_data["choices"][0]["message"] = { @@ -2001,7 +2222,7 @@ def test_reasoning_stripped_chat_histories_tokenize_authoritative_views() -> Non } first.response = ChatCompletion.model_validate(first_data) - second = _chat_exchange([1, 3, 4], [5, 6], offset=1) + second = _chat_exchange([1, 101, 102, 9, 4], [5, 6], offset=1) second.request["messages"] = [ {"role": "user", "content": "one"}, {"role": "assistant", "content": "first"}, @@ -2019,11 +2240,13 @@ def test_reasoning_stripped_chat_histories_tokenize_authoritative_views() -> Non ) class Tokenizer: + name_or_path = "test/model" + def __call__(self, text: str, **kwargs: object) -> list[int]: del kwargs return { "one": [1], - "first": [3], + "first": [500], "two": [4], "thought-two": [5], "second": [6], @@ -2039,16 +2262,207 @@ def apply_chat_template( "two", "second", ] - return [1, 3, 4, 5, 6] + return [1, 500, 9, 4, 5, 6] tokenized = trajectory.tokenize(multi_history=True, tokenizer=Tokenizer()) assert [history.token_ids for history in tokenized.histories] == [ - [1, 2, 3], - [1, 3, 4, 5, 6], + [1, 2, 101, 102, 9], + [1, 101, 102, 9, 4, 5, 6], ] assert tokenized.histories[1].flags[1] & art.TokenFlag.SAMPLED + assert tokenized.histories[1].flags[1] & art.TokenFlag.EXACT + assert tokenized.histories[1].logprobs[1:3] == [-10.1, -10.2] + assert tokenized.histories[1].flags[3] == ( + art.TokenFlag.EXACT | art.TokenFlag.SAMPLED + ) + assert tokenized.histories[1].logprobs[3] == -0.9 assert 2 not in tokenized.histories[1].token_ids + assert 500 not in tokenized.histories[1].token_ids + + +def test_reasoning_stripped_histories_remain_trainable_end_to_end( + monkeypatch: pytest.MonkeyPatch, +) -> None: + first = _chat_exchange([1], [2, 101, 102, 9]) + first.request["messages"] = [{"role": "user", "content": "one"}] + first_data = first.response.model_dump(mode="python") + first_data["choices"][0]["message"] = { + "role": "assistant", + "content": "first", + "reasoning": "thought-one", + } + first.response = ChatCompletion.model_validate(first_data) + second = _chat_exchange([1, 101, 102, 9, 4], [5, 6], offset=1) + second.request["messages"] = [ + {"role": "user", "content": "one"}, + {"role": "assistant", "content": "first"}, + {"role": "user", "content": "two"}, + ] + second_data = second.response.model_dump(mode="python") + second_data["choices"][0]["message"] = { + "role": "assistant", + "content": "second", + "reasoning": "thought-two", + } + second.response = ChatCompletion.model_validate(second_data) + + class Tokenizer: + name_or_path = "test/model" + + def __call__(self, text: str, **kwargs: object) -> list[int]: + del kwargs + return { + "one": [1], + "first": [500], + "two": [4], + "thought-two": [5], + "second": [6], + }[text] + + def apply_chat_template( + self, messages: list[dict[str, Any]], **kwargs: object + ) -> list[int]: + del kwargs + content = [message.get("content") for message in messages] + return ( + [1, 2, 101, 102, 9] + if content == ["one", "first"] + else [1, 500, 9, 4, 5, 6] + ) + + monkeypatch.setattr( + "art.trajectories._tokenize._load_tokenizer", lambda _: Tokenizer() + ) + trajectories = [ + art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[first, second]), + reward=reward, + ) + for reward in (1.0, 0.0) + ] + group = art.TrajectoryGroup(trajectories=trajectories) + + from art.preprocessing.tokenize import tokenize_trajectory_groups + from art.tinker_native.data import trajectory_groups_to_datums + + preprocessing = list( + tokenize_trajectory_groups( + Tokenizer(), # type: ignore[arg-type, ty:invalid-argument-type] + [group], + allow_training_without_logprobs=False, + scale_rewards=False, + shuffle_group_trajectories=False, + drop_zero_advantage_trajectories=False, + ) + ) + datums = trajectory_groups_to_datums( + [group], + renderer=None, + tokenizer=None, + normalize_advantages=False, + ) + + assert len(preprocessing) == 4 + assert len(datums) == 4 + assert all( + not math.isnan(logprob) + for result in preprocessing + for logprob, sampled in zip(result.logprobs, result.assistant_mask, strict=True) + if sampled + ) + + +def test_reasoning_stripped_tool_call_keeps_exact_evidence_for_strict_training( + monkeypatch: pytest.MonkeyPatch, +) -> None: + first = _chat_exchange([1], [2, 7, 8]) + first.request["messages"] = [{"role": "user", "content": "one"}] + first_data = first.response.model_dump(mode="python") + first_data["choices"][0]["message"] = { + "role": "assistant", + "content": None, + "reasoning": "thought", + "tool_calls": [ + { + "id": "call-1", + "type": "function", + "function": {"name": "lookup", "arguments": "{}"}, + } + ], + } + first.response = ChatCompletion.model_validate(first_data) + second = _chat_exchange([1, 7, 8, 4], [5], offset=1) + second.request["messages"] = [ + {"role": "user", "content": "one"}, + { + "role": "assistant", + "content": None, + "tool_calls": first_data["choices"][0]["message"]["tool_calls"], + }, + {"role": "user", "content": "two"}, + ] + + class Tokenizer: + name_or_path = "test/model" + + def __call__(self, text: str, **kwargs: object) -> list[int]: + del kwargs + return {"lookup": [7], "{}": [8]}[text] + + def apply_chat_template( + self, messages: list[dict[str, Any]], **kwargs: object + ) -> list[int]: + del kwargs + assert len(messages) == 4 + return [1, 7, 8, 4, 5] + + monkeypatch.setattr( + "art.trajectories._tokenize._load_tokenizer", lambda _: Tokenizer() + ) + trajectories = [ + art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[first, second]), + reward=reward, + ) + for reward in (1.0, 0.0) + ] + group = art.TrajectoryGroup(trajectories=trajectories) + + from art.preprocessing.tokenize import tokenize_trajectory_groups + from art.tinker_native.data import trajectory_groups_to_datums + + tokenized = trajectories[0].tokenize( + multi_history=True, + tokenizer=Tokenizer(), + ) + second_history = tokenized.histories[1] + assert second_history.token_ids == [1, 7, 8, 4, 5] + assert second_history.logprobs[1:3] == [-0.7, -0.8] + assert second_history.flags[1:3] == [ + art.TokenFlag.EXACT | art.TokenFlag.SAMPLED, + art.TokenFlag.EXACT | art.TokenFlag.SAMPLED, + ] + + preprocessing = list( + tokenize_trajectory_groups( + Tokenizer(), # type: ignore[arg-type, ty:invalid-argument-type] + [group], + allow_training_without_logprobs=False, + scale_rewards=False, + shuffle_group_trajectories=False, + drop_zero_advantage_trajectories=False, + ) + ) + datums = trajectory_groups_to_datums( + [group], + renderer=None, + tokenizer=None, + normalize_advantages=False, + ) + + assert len(preprocessing) == 4 + assert len(datums) == 4 def test_responses_prompt_repair_uses_native_text_and_source_position() -> None: @@ -2144,6 +2558,7 @@ def apply_chat_template( assert tokenized.token_ids == [1, 10, 20, 25, 30, 31, 26] assert tokenized.flags[2:6] == [art.TokenFlag.SAMPLED] * 4 + assert all(math.isnan(value) for value in tokenized.logprobs[2:6]) assert not tokenized.flags[0] & art.TokenFlag.SAMPLED @@ -2197,9 +2612,8 @@ def test_explicit_template_override_rerenders_exact_exchange_scaffold() -> None: class Tokenizer: def __call__(self, text: str, *, add_special_tokens: bool = False) -> list[int]: - assert text == "answer" assert not add_special_tokens - return [20] + return {"turn 0": [10], "answer": [20]}[text] def apply_chat_template( self, From 2cb0c4e8cc11108e32ca4281d3bf565fd043a5ab Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Fri, 24 Jul 2026 16:10:41 +0000 Subject: [PATCH 07/58] Represent terminal EOS as an empty assistant turn --- src/art/trajectories/_history.py | 36 ++++++++++++++----- src/art/trajectories/_tokenize.py | 20 ++++++++++- tests/unit/trajectories/test_tokenize.py | 45 ++++++++++++++++++++++-- 3 files changed, 89 insertions(+), 12 deletions(-) diff --git a/src/art/trajectories/_history.py b/src/art/trajectories/_history.py index 38cfa8eeb..9cb676a79 100644 --- a/src/art/trajectories/_history.py +++ b/src/art/trajectories/_history.py @@ -719,14 +719,18 @@ def responses_histories( final_sources = [*copy.copy(prompt_sources), *copy.copy(output_sources)] if generations: for generation_index, generation in enumerate(generations): - if not generation.output_indices: + outputless = not generation.output_indices + if outputless and generation_index != len(generations) - 1: raise ValueError( - "A Responses token generation without native output " - "items cannot yet be projected as a history" + "A nonterminal Responses token generation without " + "native output items cannot be projected as a history" ) - output_start = generation.output_indices[0] - output_end = generation.output_indices[-1] + 1 - if generation_index == len(generations) - 1: + if outputless: + output_start = output_end = len(output) + else: + output_start = generation.output_indices[0] + output_end = generation.output_indices[-1] + 1 + if not outputless and generation_index == len(generations) - 1: output_end = len(output) generation_prompt = [ *copy.deepcopy(prompt), @@ -774,8 +778,24 @@ def responses_histories( ) for _, source in retained ] - generation_output = output[output_start:output_end] - generation_output_sources = output_sources[output_start:output_end] + if outputless: + generation_output = [ + cast( + ResponseInputItemParam, + {"role": "assistant", "content": ""}, + ) + ] + generation_output_sources = [ + ResponsesItemSource( + exchange=exchange, + generation_index=generation_index, + ) + ] + else: + generation_output = output[output_start:output_end] + generation_output_sources = output_sources[ + output_start:output_end + ] _extend_branches( branches, prompt=generation_prompt, diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index 066466248..85626fae0 100644 --- a/src/art/trajectories/_tokenize.py +++ b/src/art/trajectories/_tokenize.py @@ -1949,7 +1949,25 @@ def _validate_history_sources(history: History) -> None: raise ValueError("Responses request source has a generation index") expected = request_input[request_index] else: - raise ValueError("Responses source has no request or output index") + generations = _response_generations(source.exchange.response) + if source.generation_index is None: + raise ValueError( + "Responses source has no request, output, or generation index" + ) + generation_index = _source_index( + source.generation_index, + length=len(generations), + field="Responses generation source index", + ) + if ( + generation_index != len(generations) - 1 + or generations[generation_index].output_indices + ): + raise ValueError( + "Responses generation-only source must refer to a terminal " + "generation without native output items" + ) + expected = {"role": "assistant", "content": ""} if expected is None or item != expected: raise ValueError( "Responses history no longer matches its source exchange" diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index c85f128b7..521ea7913 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -1839,8 +1839,8 @@ def test_responses_token_generations_fail_closed( ).responses_history().tokenize() -def test_responses_generation_without_output_items_is_not_silently_lost() -> None: - exchange = _response_exchange("unprojectable-generation", 2) +def test_responses_terminal_generation_without_output_items_is_tokenized() -> None: + exchange = _response_exchange("terminal-eos", 2) data = exchange.response.model_dump(mode="python") data["output"] = [] data["token_generations"] = [ @@ -1852,7 +1852,46 @@ def test_responses_generation_without_output_items_is_not_silently_lost() -> Non ] exchange.response = Response.model_validate(data) - with pytest.raises(ValueError, match="cannot yet be projected"): + history = art.Trajectory( + exchanges=TrajectoryExchanges(responses=[exchange]) + ).responses_history() + tokenized = history.tokenize() + + assert history.input[-1] == {"role": "assistant", "content": ""} + assert history.input_sources[-1] == ResponsesItemSource( + exchange=exchange, generation_index=0 + ) + assert tokenized.token_ids == [1, 2] + assert math.isnan(tokenized.logprobs[0]) + assert tokenized.logprobs[1] == pytest.approx(-0.2) + assert tokenized.flags == [ + art.TokenFlag.EXACT, + art.TokenFlag.EXACT | art.TokenFlag.SAMPLED, + ] + assert history.as_chat_completions_history().messages[-1] == { + "role": "assistant", + "content": "", + } + + +def test_responses_nonterminal_generation_without_output_items_raises() -> None: + exchange = _response_exchange("hidden-control", 2) + data = exchange.response.model_dump(mode="python") + data["token_generations"] = [ + { + "prompt_token_ids": [1], + "output_tokens": [{"token_id": 2, "logprob": -0.2}], + "output_indices": [], + }, + { + "prompt_token_ids": [1, 2], + "output_tokens": [{"token_id": 3, "logprob": -0.3}], + "output_indices": [0], + }, + ] + exchange.response = Response.model_validate(data) + + with pytest.raises(ValueError, match="nonterminal"): art.Trajectory( exchanges=TrajectoryExchanges(responses=[exchange]) ).responses_history() From 54e1fb36b1239b8b9e0c5bb50f549c5d8485d7b7 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Fri, 24 Jul 2026 17:28:06 +0000 Subject: [PATCH 08/58] Harden trajectory evidence and model selection --- src/art/local/backend.py | 3 +- src/art/preprocessing/tokenize.py | 81 +++++- src/art/tinker_native/backend.py | 3 +- src/art/trajectories/__init__.py | 1 + src/art/trajectories/_history.py | 17 +- src/art/trajectories/_protocols.py | 85 +++++-- src/art/trajectories/_tokenize.py | 45 +++- .../test_exchange_training_model_selection.py | 236 +++++++++++++++++- .../test_pipeline_trainer_local_backend.py | 3 +- tests/unit/trajectories/test_capture.py | 70 ++++++ tests/unit/trajectories/test_history.py | 37 +++ tests/unit/trajectories/test_tokenize.py | 49 ++-- 12 files changed, 561 insertions(+), 69 deletions(-) diff --git a/src/art/local/backend.py b/src/art/local/backend.py index 63bbe2192..670471259 100644 --- a/src/art/local/backend.py +++ b/src/art/local/backend.py @@ -91,6 +91,7 @@ ) from ..serving_capabilities import ServingCapabilities from ..trajectories import Trajectory, TrajectoryGroup +from ..trajectories._tokenize import _training_model_pattern from ..types import ( Choice, LocalTrainResult, @@ -958,7 +959,7 @@ def _get_packed_tensors( image_processor=self._image_processors[model.base_model], chat_template_kwargs=chat_template_kwargs, chat_template_tool_schema_format=chat_template_tool_schema_format, - model=self._model_inference_name(model), + model=_training_model_pattern(self._model_inference_name(model)), ) ) if not tokenized_results: diff --git a/src/art/preprocessing/tokenize.py b/src/art/preprocessing/tokenize.py index 0d4732c0d..d1b32ffbd 100644 --- a/src/art/preprocessing/tokenize.py +++ b/src/art/preprocessing/tokenize.py @@ -17,6 +17,7 @@ if TYPE_CHECKING: from transformers.image_processing_utils import BaseImageProcessor +from ..openai import ART_MOE_ROUTING_METADATA_KEY from ..trajectories import ( ChatCompletionsExchange, ChatCompletionsHistory, @@ -117,26 +118,47 @@ def _chat_choice_trace( metadata = choice_vllm_token_metadata(choice) assert metadata is not None completion_ids = metadata[1] - matches = [ - index - for index in range(cursor, len(token_ids) - len(completion_ids) + 1) - if token_ids[index : index + len(completion_ids)] == completion_ids - and all( - flag & TokenFlag.SAMPLED - for flag in flags[index : index + len(completion_ids)] - ) - ] + retained_start = 0 + + def matching_offsets(values: list[int]) -> list[int]: + return [ + index + for index in range(cursor, len(token_ids) - len(values) + 1) + if token_ids[index : index + len(values)] == values + and all( + flag & TokenFlag.SAMPLED + for flag in flags[index : index + len(values)] + ) + ] + + matches = matching_offsets(completion_ids) + if not matches and completion_ids: + for retained_start in range(1, len(completion_ids)): + matches = matching_offsets(completion_ids[retained_start:]) + if matches: + break if not matches: if has_routing: raise RuntimeError( "MoE routed completion tokens are absent from tokenized history" ) return None + if retained_start and len(matches) != 1: + if has_routing: + raise RuntimeError( + "MoE routed completion suffix is ambiguous in tokenized history" + ) + return None + retained_ids = completion_ids[retained_start:] offset = matches[0] - choices.append(choice) + choices.append( + _choice_with_retained_routing(choice, retained_start) + if retained_start + else choice + ) offsets.append(offset) - lengths.append(len(completion_ids)) - cursor = offset + len(completion_ids) + lengths.append(len(retained_ids)) + cursor = offset + len(retained_ids) return ( _ChatChoiceTrace(choices=choices, offsets=offsets, lengths=lengths) if choices @@ -144,6 +166,41 @@ def _chat_choice_trace( ) +def _choice_with_retained_routing(choice: Choice, start: int) -> Choice: + """Copy one routed choice while retaining only a completion suffix.""" + + routing = choice_moe_routing_metadata(choice) + if routing is None: + return choice + token_metadata = choice_vllm_token_metadata(choice) + assert token_metadata is not None + prompt_ids, completion_ids = token_metadata + routes = routing.get("routed_experts") + if not isinstance(routes, np.ndarray): + raise RuntimeError("Missing binary routed experts") + completion_route_count = len(routes) - len(prompt_ids) + if completion_route_count not in { + len(completion_ids), + max(len(completion_ids) - 1, 0), + }: + raise RuntimeError( + "routed_experts length does not match prompt/completion token ids" + ) + retained = choice.model_copy() + extra = retained.model_extra + if extra is None: + raise RuntimeError("OpenAI Choice.model_extra is unavailable for route replay") + extra["token_ids"] = completion_ids[start:] + extra[ART_MOE_ROUTING_METADATA_KEY] = { + **routing, + "completion_token_ids": completion_ids[start:], + "routed_experts": np.concatenate( + (routes[: len(prompt_ids)], routes[len(prompt_ids) + start :]) + ), + } + return retained + + def _slice_moe_routes( routes: MoeRouteArray | MoeRouteSegments | None, start: int ) -> MoeRouteArray | MoeRouteSegments | None: diff --git a/src/art/tinker_native/backend.py b/src/art/tinker_native/backend.py index 6a0998606..e6a28cd3c 100644 --- a/src/art/tinker_native/backend.py +++ b/src/art/tinker_native/backend.py @@ -41,6 +41,7 @@ from ..tinker.backend import get_renderer_name from ..tinker.server import get_free_port from ..trajectories import Trajectory, TrajectoryGroup +from ..trajectories._tokenize import _training_model_pattern from ..types import TrainResult, TrainSFTConfig from ..utils.lifecycle import process_shutdown_timeout from ..utils.output_dirs import get_model_dir @@ -358,7 +359,7 @@ async def train( state.tokenizer, normalize_advantages, base_model=model.base_model, - model=self._model_inference_name(model), + model=_training_model_pattern(self._model_inference_name(model)), ) metrics: dict[str, float] = { diff --git a/src/art/trajectories/__init__.py b/src/art/trajectories/__init__.py index ffea450e7..b12c9a280 100644 --- a/src/art/trajectories/__init__.py +++ b/src/art/trajectories/__init__.py @@ -511,6 +511,7 @@ async def track_duration(self, metric_name: str) -> AsyncGenerator[None, None]: def __str__(self) -> str: return f"Trajectory(reward={self.reward}, metrics={self.metrics}, metadata={self.metadata})" + # Every model selector accepts an exact identity or a shell-style pattern. def chat_completions_history( self, *, model: str | None = None ) -> ChatCompletionsHistory: diff --git a/src/art/trajectories/_history.py b/src/art/trajectories/_history.py index 9cb676a79..b948b8273 100644 --- a/src/art/trajectories/_history.py +++ b/src/art/trajectories/_history.py @@ -4,6 +4,7 @@ import copy from dataclasses import dataclass, replace from datetime import datetime +from fnmatch import fnmatchcase import json from typing import Generic, Protocol, TypeVar, cast @@ -129,7 +130,9 @@ def _selected_models( exchanges: Sequence[_ExchangeT], model: str | None, protocol: str ) -> list[tuple[str, list[_ExchangeT]]]: selected = [ - exchange for exchange in exchanges if model is None or exchange.model == model + exchange + for exchange in exchanges + if model is None or _model_matches(exchange.model, model) ] if not selected: suffix = f" for model {model!r}" if model is not None else "" @@ -144,6 +147,10 @@ def _selected_models( return list(grouped.items()) +def _model_matches(candidate: str | None, pattern: str) -> bool: + return candidate is not None and fnmatchcase(candidate, pattern) + + def _one(history: Sequence[_ItemT], protocol: str) -> _ItemT: if len(history) != 1: raise ValueError( @@ -1394,12 +1401,12 @@ def trajectory_histories( candidates: list[TrajectoryHistory] = [] if any( - model is None or exchange.model == model + model is None or _model_matches(exchange.model, model) for exchange in trajectory.exchanges.chat_completions ): candidates.extend(chat_completions_histories(trajectory, model=model)) if any( - model is None or exchange.model == model + model is None or _model_matches(exchange.model, model) for exchange in trajectory.exchanges.completions ): try: @@ -1409,12 +1416,12 @@ def trajectory_histories( raise candidates.extend(completions_string_histories(trajectory, model=model)) if any( - model is None or exchange.model == model + model is None or _model_matches(exchange.model, model) for exchange in trajectory.exchanges.responses ): candidates.extend(responses_histories(trajectory, model=model)) if any( - model is None or exchange.model == model + model is None or _model_matches(exchange.model, model) for exchange in trajectory.exchanges.messages ): candidates.extend(anthropic_messages_histories(trajectory, model=model)) diff --git a/src/art/trajectories/_protocols.py b/src/art/trajectories/_protocols.py index 704df4bbb..7d49b77b9 100644 --- a/src/art/trajectories/_protocols.py +++ b/src/art/trajectories/_protocols.py @@ -2,6 +2,7 @@ from datetime import datetime import json +import math from typing import Any, Literal from urllib.parse import urlsplit @@ -224,6 +225,35 @@ def _messages_response(body: bytes, *, stream: bool) -> Message: prompt_token_ids: list[int] = [] block_token_ids: dict[int, list[int]] = {} block_logprobs: dict[int, list[float]] = {} + + def parsed_token_ids(value: object, field: str) -> list[int] | None: + if value is None: + return None + if not isinstance(value, list): + raise ValueError(f"{field} must contain integer token IDs") + result: list[int] = [] + for item in value: + if not isinstance(item, int) or isinstance(item, bool) or item < 0: + raise ValueError(f"{field} must contain integer token IDs") + result.append(item) + return result + + def token_logprobs(value: object, field: str) -> list[float] | None: + if value is None: + return None + if not isinstance(value, list): + raise ValueError(f"{field} must contain numeric logprobs") + result: list[float] = [] + for item in value: + if ( + not isinstance(item, (int, float)) + or isinstance(item, bool) + or not math.isfinite(item) + ): + raise ValueError(f"{field} must contain numeric logprobs") + result.append(float(item)) + return result + for event_name, payload in _sse_events(body): if not isinstance(payload, dict): continue @@ -240,44 +270,49 @@ def _messages_response(body: bytes, *, stream: bool) -> Message: values = ( message.get("prompt_token_ids") if isinstance(message, dict) else None ) - if isinstance(values, list) and all( - isinstance(value, int) for value in values - ): + if ( + values := parsed_token_ids(values, "Messages prompt_token_ids") + ) is not None: prompt_token_ids = values elif event.type == "message_delta": - values = payload.get("prompt_token_ids") - if isinstance(values, list) and all( - isinstance(value, int) for value in values - ): + if ( + values := parsed_token_ids( + payload.get("prompt_token_ids"), "Messages prompt_token_ids" + ) + ) is not None: prompt_token_ids = values - event_token_ids = payload.get("token_ids") - event_logprobs = payload.get("logprobs") - if isinstance(event_token_ids, list) and all( - isinstance(value, int) for value in event_token_ids - ): + if ( + event_token_ids := parsed_token_ids( + payload.get("token_ids"), "Messages token_ids" + ) + ) is not None: token_ids = event_token_ids - if isinstance(event_logprobs, list) and all( - isinstance(value, (int, float)) for value in event_logprobs - ): - logprobs = [float(value) for value in event_logprobs] + if ( + event_logprobs := token_logprobs( + payload.get("logprobs"), "Messages logprobs" + ) + ) is not None: + logprobs = event_logprobs elif event.type in {"content_block_start", "content_block_delta"}: index = payload.get("index") - event_token_ids = payload.get("token_ids") - event_logprobs = payload.get("logprobs") + event_token_ids = parsed_token_ids( + payload.get("token_ids"), "Messages content token_ids" + ) + event_logprobs = token_logprobs( + payload.get("logprobs"), "Messages content logprobs" + ) if ( isinstance(index, int) - and isinstance(event_token_ids, list) - and all(isinstance(value, int) for value in event_token_ids) + and not isinstance(index, bool) + and event_token_ids ): block_token_ids.setdefault(index, []).extend(event_token_ids) if ( isinstance(index, int) - and isinstance(event_logprobs, list) - and all(isinstance(value, (int, float)) for value in event_logprobs) + and not isinstance(index, bool) + and event_logprobs ): - block_logprobs.setdefault(index, []).extend( - float(value) for value in event_logprobs - ) + block_logprobs.setdefault(index, []).extend(event_logprobs) complete = complete or event.type == "message_stop" if snapshot is None or not complete: raise ValueError("Incomplete Messages stream") diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index 85626fae0..8a9a047cd 100644 --- a/src/art/trajectories/_tokenize.py +++ b/src/art/trajectories/_tokenize.py @@ -38,6 +38,7 @@ TrajectoryGroup, TrajectoryHistory, ) +from ._history import _model_matches from ._protocols import Exchange if TYPE_CHECKING: @@ -529,7 +530,9 @@ def _exchange_list(trajectory: Trajectory, model: str | None) -> list[Exchange]: *trajectory.exchanges.messages, ] if model is not None: - exchanges = [exchange for exchange in exchanges if exchange.model == model] + exchanges = [ + exchange for exchange in exchanges if _model_matches(exchange.model, model) + ] if not exchanges: raise ValueError(f"Trajectory contains no exchanges for model {model!r}") models = {exchange.model for exchange in exchanges} @@ -559,7 +562,7 @@ def _require_training_model(trajectory: Trajectory, model: str | None) -> str: if None in models: raise ValueError("Every training exchange must identify its model") if model is not None: - if model not in models: + if not any(_model_matches(candidate, model) for candidate in models): raise ValueError(f"Trajectory contains no exchanges for model {model!r}") return model if len(models) != 1: @@ -569,6 +572,16 @@ def _require_training_model(trajectory: Trajectory, model: str | None) -> str: return cast(str, next(iter(models))) +def _training_model_pattern(model: str) -> str: + """Select every immutable checkpoint belonging to one automatic policy.""" + + if re.search(r"@\d+$", model): + return re.sub(r"\d+$", "*", model) + if re.search(r":step\d+$", model): + return re.sub(r"\d+$", "*", model) + return model + + def _artifact_name(model: str) -> str: return model.removeprefix("wandb-artifact:///") @@ -2386,8 +2399,10 @@ def _source_is_sampled(source: object) -> bool: def _tokenize_exact_projected_chat_history( history: ChatCompletionsHistory, + *, + projection_validated: bool = False, ) -> TokenizedHistory | None: - if not _history_matches_projection(history): + if not projection_validated and not _history_matches_projection(history): return None sampled_sources: list[object] = [] seen: set[tuple[object, ...]] = set() @@ -3179,9 +3194,27 @@ def tokenize_history( chat_template=chat_template, chat_template_kwargs=chat_template_kwargs, ) - if (isinstance(history, AnthropicMessagesHistory) and needs_render) or ( - isinstance(history, ResponsesHistory) and needs_render - ): + if isinstance(history, AnthropicMessagesHistory) and needs_render: + converted = history.as_chat_completions_history() + if ( + not override_requires_render + and _history_matches_projection(history) + and ( + exact := _tokenize_exact_projected_chat_history( + converted, + projection_validated=True, + ) + ) + ): + return exact + return _tokenize_chat_view( + converted, + base_model=base_model, + tokenizer=tokenizer, + chat_template=chat_template, + chat_template_kwargs=chat_template_kwargs, + ) + if isinstance(history, ResponsesHistory) and needs_render: return _tokenize_chat_view( history.as_chat_completions_history(), base_model=base_model, diff --git a/tests/unit/test_exchange_training_model_selection.py b/tests/unit/test_exchange_training_model_selection.py index 176fa2bfb..068426ac6 100644 --- a/tests/unit/test_exchange_training_model_selection.py +++ b/tests/unit/test_exchange_training_model_selection.py @@ -8,9 +8,15 @@ import art from art.openai import ART_MOE_ROUTING_METADATA_KEY -from art.preprocessing.tokenize import tokenize_trajectory_groups +from art.preprocessing.tokenize import _chat_choice_trace, tokenize_trajectory_groups from art.tinker_native.data import trajectory_groups_to_datums -from art.trajectories import ChatCompletionsExchange, ChatCompletionsRequest +from art.trajectories import ( + ChatCompletionsExchange, + ChatCompletionsHistory, + ChatCompletionsMessageSource, + ChatCompletionsRequest, +) +from art.trajectories._tokenize import _training_model_pattern def _exchange(model: str, output_token: int) -> ChatCompletionsExchange: @@ -94,6 +100,24 @@ def _group() -> art.TrajectoryGroup: return art.TrajectoryGroup(trajectories=trajectories) +def _versioned_group() -> art.TrajectoryGroup: + return art.TrajectoryGroup( + trajectories=[ + art.Trajectory( + exchanges=art.TrajectoryExchanges( + chat_completions=[ + _exchange("policy@12", 2), + _exchange("judge@4", 3), + _exchange("policy@13", 4), + ] + ), + reward=reward, + ) + for reward in (1.0, 0.0) + ] + ) + + class _Tokenizer: name_or_path = "base/model" @@ -137,6 +161,54 @@ def test_tinker_requires_model_selection() -> None: assert all(datum.model_input.to_ints() == [1] for datum in datums) +def test_training_model_patterns_select_all_policy_versions() -> None: + group = _versioned_group() + results = list( + tokenize_trajectory_groups( + cast(PreTrainedTokenizerBase, _Tokenizer()), + [group], + allow_training_without_logprobs=False, + scale_rewards=False, + model="policy@*", + ) + ) + assert len(results) == 4 + assert {result.token_ids[-1] for result in results} == {2, 4} + + datums = trajectory_groups_to_datums( + [group], + renderer=None, + tokenizer=None, + model="policy@*", + ) + assert len(datums) == 4 + + trajectory = group.trajectories[0] + with pytest.raises(ValueError, match="exactly one history"): + trajectory.tokenize(model="policy@*") + tokenized = trajectory.tokenize(model="policy@*", multi_history=True) + assert [history.model for history in tokenized.histories] == [ + "policy@12", + "policy@13", + ] + + +@pytest.mark.parametrize( + ("model", "pattern"), + [ + ("policy@12", "policy@*"), + ( + "wandb-artifact:///entity/project/run:step12", + "wandb-artifact:///entity/project/run:step*", + ), + ("policy:active", "policy:active"), + ("base/model", "base/model"), + ], +) +def test_automatic_training_model_pattern(model: str, pattern: str) -> None: + assert _training_model_pattern(model) == pattern + + def test_preprocessing_preserves_adjacent_choice_boundaries_for_moe() -> None: exchanges = [ _routed_exchange( @@ -183,6 +255,166 @@ def test_preprocessing_preserves_adjacent_choice_boundaries_for_moe() -> None: assert all(result.moe_routed_experts is not None for result in results) +def test_preprocessing_preserves_moe_routes_for_reasoning_stripped_suffix() -> None: + first = _routed_exchange( + prompt_token_ids=[1], + output_token=2, + messages=[{"role": "user", "content": "one"}], + content="first", + ) + first_data = first.response.model_dump(mode="python") + first_data["choices"][0]["message"] = { + "role": "assistant", + "content": "first", + "reasoning": "thought-one", + } + first_data["choices"][0]["token_ids"] = [2, 101, 102, 9] + first_data["choices"][0]["logprobs"]["content"] = [ + { + "token": f"token_id:{token}", + "logprob": -token / 10, + "bytes": [], + "top_logprobs": [], + } + for token in [2, 101, 102, 9] + ] + first.response = ChatCompletion.model_validate(first_data) + first_extra = first.response.choices[0].model_extra + assert first_extra is not None + first_extra[ART_MOE_ROUTING_METADATA_KEY] = { + "prompt_token_ids": [1], + "completion_token_ids": [2, 101, 102, 9], + "routed_experts": np.asarray( + [[[10]], [[20]], [[1010]], [[1020]], [[90]]], dtype=np.int32 + ), + } + + second = _routed_exchange( + prompt_token_ids=[1, 101, 102, 9, 4], + output_token=5, + messages=[ + {"role": "user", "content": "one"}, + {"role": "assistant", "content": "first"}, + {"role": "user", "content": "two"}, + ], + content="second", + ) + second_data = second.response.model_dump(mode="python") + second_data["choices"][0]["message"] = { + "role": "assistant", + "content": "second", + "reasoning": "thought-two", + } + second_data["choices"][0]["token_ids"] = [5, 6] + second_data["choices"][0]["logprobs"]["content"] = [ + { + "token": f"token_id:{token}", + "logprob": -token / 10, + "bytes": [], + "top_logprobs": [], + } + for token in [5, 6] + ] + second.response = ChatCompletion.model_validate(second_data) + second_extra = second.response.choices[0].model_extra + assert second_extra is not None + second_extra[ART_MOE_ROUTING_METADATA_KEY] = { + "prompt_token_ids": [1, 101, 102, 9, 4], + "completion_token_ids": [5, 6], + "routed_experts": np.asarray( + [[[10]], [[1010]], [[1020]], [[90]], [[40]], [[50]], [[60]]], + dtype=np.int32, + ), + } + + class Tokenizer(_Tokenizer): + def __call__(self, text: str, **kwargs: object) -> list[int]: + del kwargs + return { + "one": [1], + "first": [500], + "two": [4], + "thought-two": [5], + "second": [6], + }[text] + + def apply_chat_template( + self, messages: list[dict[str, object]], **kwargs: object + ) -> list[int]: + del kwargs + content = [message.get("content") for message in messages] + return ( + [1, 2, 101, 102, 9] + if content == ["one", "first"] + else [1, 500, 9, 4, 5, 6] + ) + + group = art.TrajectoryGroup( + trajectories=[ + art.Trajectory( + exchanges=art.TrajectoryExchanges( + chat_completions=[first, second], + ), + reward=reward, + ) + for reward in (1.0, 0.0) + ] + ) + + results = list( + tokenize_trajectory_groups( + cast(PreTrainedTokenizerBase, Tokenizer()), + [group], + allow_training_without_logprobs=False, + scale_rewards=False, + shuffle_group_trajectories=False, + drop_zero_advantage_trajectories=False, + model="policy", + ) + ) + + stripped = [result for result in results if result.token_ids[1] == 101] + assert len(stripped) == 2 + assert all(result.choice_offsets == [1, 5] for result in stripped) + expected_routes = np.asarray( + [[[10]], [[1010]], [[1020]], [[90]], [[40]], [[50]], [[60]]], + dtype=np.int32, + ) + for result in stripped: + assert isinstance(result.moe_routed_experts, np.ndarray) + assert np.array_equal(result.moe_routed_experts, expected_routes) + + +def test_ambiguous_non_moe_suffix_falls_back_to_sampled_spans() -> None: + first = _exchange("policy", 2) + first_extra = first.response.choices[0].model_extra + assert first_extra is not None + first_extra["token_ids"] = [2, 11] + second = _exchange("policy", 11) + history = ChatCompletionsHistory( + model="policy", + messages=[], + message_sources=[ + ChatCompletionsMessageSource(exchange=first, choice_index=0), + ChatCompletionsMessageSource(exchange=second, choice_index=0), + ], + ) + + assert ( + _chat_choice_trace( + history, + [1, 11, 2, 11], + [ + art.TokenFlag.EXACT, + art.TokenFlag.EXACT | art.TokenFlag.SAMPLED, + art.TokenFlag.EXACT, + art.TokenFlag.EXACT | art.TokenFlag.SAMPLED, + ], + ) + is None + ) + + def test_preprocessing_rejects_partial_choice_evidence_before_moe_routes() -> None: first = _routed_exchange( prompt_token_ids=[1], diff --git a/tests/unit/test_pipeline_trainer_local_backend.py b/tests/unit/test_pipeline_trainer_local_backend.py index 741b3bed4..340217b16 100644 --- a/tests/unit/test_pipeline_trainer_local_backend.py +++ b/tests/unit/test_pipeline_trainer_local_backend.py @@ -753,7 +753,7 @@ def test_local_backend_get_packed_tensors_warns_and_drops_overlong_results( patch( "art.local.backend.tokenize_trajectory_groups", return_value=iter([short_result, long_result]), - ), + ) as tokenize, pytest.warns(UserWarning, match="Dropping 1 tokenized results"), ): packed_tensors = backend._get_packed_tensors( @@ -769,6 +769,7 @@ def test_local_backend_get_packed_tensors_warns_and_drops_overlong_results( assert packed_tensors is not None assert packed_tensors["tokens"].shape == (1, 4) + assert tokenize.call_args.kwargs["model"] == f"{model.name}@*" @pytest.mark.asyncio diff --git a/tests/unit/trajectories/test_capture.py b/tests/unit/trajectories/test_capture.py index e13792af9..9353eed4a 100644 --- a/tests/unit/trajectories/test_capture.py +++ b/tests/unit/trajectories/test_capture.py @@ -723,6 +723,76 @@ def test_all_streaming_protocols_reconstruct_final_responses() -> None: assert getattr(exchange.response, "prompt_token_ids") == [1] +@pytest.mark.parametrize( + ("field", "value"), + [ + ("prompt_token_ids", [True]), + ("prompt_token_ids", [-1]), + ("token_ids", [False]), + ("token_ids", [-1]), + ("logprobs", [True]), + ("logprobs", [float("nan")]), + ("logprobs", [float("inf")]), + ], +) +def test_streaming_messages_reject_malformed_token_metadata( + field: str, value: list[object] +) -> None: + delta = { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 1}, + "prompt_token_ids": [1], + "token_ids": [2], + "logprobs": [-0.2], + field: value, + } + body = _sse( + [ + ( + "message_start", + { + "type": "message_start", + "message": { + **MESSAGE, + "content": [], + "stop_reason": None, + "usage": {"input_tokens": 1, "output_tokens": 0}, + }, + }, + ), + ( + "content_block_start", + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "text", "text": "", "citations": None}, + }, + ), + ( + "content_block_delta", + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "hello"}, + }, + ), + ("content_block_stop", {"type": "content_block_stop", "index": 0}), + ("message_delta", delta), + ("message_stop", {"type": "message_stop"}), + ] + ) + + with pytest.raises(ValueError, match="must contain"): + build_exchange( + "messages", + {"model": "test/model", "messages": [], "stream": True}, + body, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + def test_streaming_chat_choices_are_accumulated_by_index() -> None: now = datetime.now() diff --git a/tests/unit/trajectories/test_history.py b/tests/unit/trajectories/test_history.py index 7a3794afa..a4e671e57 100644 --- a/tests/unit/trajectories/test_history.py +++ b/tests/unit/trajectories/test_history.py @@ -220,6 +220,43 @@ def test_chat_history_resolves_one_model_and_append_only_sequence() -> None: trajectory.chat_completions_history(model="test/model") +def test_model_patterns_select_matching_histories_only() -> None: + policy_12 = _chat( + [{"role": "user", "content": "one"}], + "first", + model="policy@12", + ) + judge = _chat( + [{"role": "user", "content": "judge"}], + "score", + model="judge@4", + offset=1, + ) + policy_13 = _chat( + [{"role": "user", "content": "two"}], + "second", + model="policy@13", + offset=2, + ) + trajectory = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[policy_12, judge, policy_13]) + ) + + histories = trajectory.chat_completions_histories(model="policy@*") + assert [history.model for history in histories] == ["policy@12", "policy@13"] + generic_histories = trajectory.histories(model="policy@*") + assert all(isinstance(history, art.History) for history in generic_histories) + assert [cast(art.History, history).model for history in generic_histories] == [ + "policy@12", + "policy@13", + ] + assert trajectory.chat_completions_history(model="policy@12").model == "policy@12" + with pytest.raises(ValueError, match="exactly one history"): + trajectory.history(model="policy@*") + with pytest.raises(ValueError, match="no Chat Completions exchanges"): + trajectory.chat_completions_histories(model="foreign@*") + + def test_chat_history_preserves_provider_specific_nested_fields() -> None: exchange = _chat( [ diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index 521ea7913..ea78d79c1 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -894,14 +894,42 @@ def test_anthropic_fallback_preserves_thinking_and_tool_history() -> None: ] -def test_reasoning_stripped_history_uses_stream_block_tokens() -> None: +@pytest.mark.parametrize("top_level_only", [False, True]) +def test_reasoning_stripped_messages_history_preserves_exact_tokens( + top_level_only: bool, +) -> None: def exchange( offset: int, request_messages: list[MessageParam], answer: str, + prompt_token_ids: list[int], token_ids: list[int], ) -> MessagesExchange: start = datetime(2026, 1, 1) + timedelta(seconds=offset) + thinking: dict[str, Any] = { + "type": "thinking", + "thinking": f"thought-{offset}", + "signature": "signature", + } + text: dict[str, Any] = {"type": "text", "text": answer} + response_extra: dict[str, Any] = {} + if top_level_only: + response_extra = { + "prompt_token_ids": prompt_token_ids, + "token_ids": [90 + offset, *token_ids], + "logprobs": [ + -9.0 - offset, + *[-0.1 * token for token in token_ids], + ], + } + else: + thinking.update({"token_ids": [90 + offset], "logprobs": [-9.0 - offset]}) + text.update( + { + "token_ids": token_ids, + "logprobs": [-0.1 * token for token in token_ids], + } + ) return MessagesExchange( request=MessagesRequest( model="test/model", @@ -915,24 +943,11 @@ def exchange( "type": "message", "role": "assistant", "model": "test/model", - "content": [ - { - "type": "thinking", - "thinking": f"thought-{offset}", - "signature": "signature", - "token_ids": [90 + offset], - "logprobs": [-9.0 - offset], - }, - { - "type": "text", - "text": answer, - "token_ids": token_ids, - "logprobs": [-0.1 * token for token in token_ids], - }, - ], + "content": [thinking, text], "stop_reason": "end_turn", "stop_sequence": None, "usage": {"input_tokens": 1, "output_tokens": len(token_ids)}, + **response_extra, } ), start_time=start, @@ -943,6 +958,7 @@ def exchange( 0, [{"role": "user", "content": "one"}], "first", + [10], [101, 102], ) second = exchange( @@ -953,6 +969,7 @@ def exchange( {"role": "user", "content": "two"}, ], "second", + [10, 101, 102, 11], [201], ) From ef75a5285b36a0aae5156e958d7462778e2db63e Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Fri, 24 Jul 2026 19:04:19 +0000 Subject: [PATCH 09/58] Harden exchange training selection and serialization --- src/art/local/backend.py | 6 +- src/art/preprocessing/tokenize.py | 70 +++--- src/art/tinker_native/backend.py | 4 +- src/art/tinker_native/data.py | 7 +- src/art/trajectories/__init__.py | 2 + src/art/trajectories/_selection.py | 107 +++++++++ .../test_exchange_training_model_selection.py | 210 +++++++++++++++--- .../trajectories/test_tokenized_models.py | 66 ++++++ 8 files changed, 403 insertions(+), 69 deletions(-) create mode 100644 src/art/trajectories/_selection.py create mode 100644 tests/unit/trajectories/test_tokenized_models.py diff --git a/src/art/local/backend.py b/src/art/local/backend.py index 670471259..24638f364 100644 --- a/src/art/local/backend.py +++ b/src/art/local/backend.py @@ -91,7 +91,7 @@ ) from ..serving_capabilities import ServingCapabilities from ..trajectories import Trajectory, TrajectoryGroup -from ..trajectories._tokenize import _training_model_pattern +from ..trajectories._selection import automatic_training_model_selector from ..types import ( Choice, LocalTrainResult, @@ -959,7 +959,9 @@ def _get_packed_tensors( image_processor=self._image_processors[model.base_model], chat_template_kwargs=chat_template_kwargs, chat_template_tool_schema_format=chat_template_tool_schema_format, - model=_training_model_pattern(self._model_inference_name(model)), + model=automatic_training_model_selector( + self._model_inference_name(model) + ), ) ) if not tokenized_results: diff --git a/src/art/preprocessing/tokenize.py b/src/art/preprocessing/tokenize.py index d1b32ffbd..64343217d 100644 --- a/src/art/preprocessing/tokenize.py +++ b/src/art/preprocessing/tokenize.py @@ -27,6 +27,7 @@ TrajectoryGroup, get_messages, ) +from ..trajectories._selection import ModelSelector, resolve_training_model from ..types import MessagesAndChoices from ..utils.chat_template import ( default_chat_template_kwargs_for_tokenizer, @@ -110,47 +111,53 @@ def _chat_choice_trace( ) return None + metadata = [choice_vllm_token_metadata(choice) for choice in sourced_choices] + if any(value is None for value in metadata): + raise AssertionError("choice metadata was checked above") + exact_metadata = cast(list[tuple[list[int], list[int]]], metadata) + choices: list[Choice] = [] offsets: list[int] = [] lengths: list[int] = [] cursor = 0 - for choice in sourced_choices: - metadata = choice_vllm_token_metadata(choice) - assert metadata is not None - completion_ids = metadata[1] - retained_start = 0 - - def matching_offsets(values: list[int]) -> list[int]: - return [ - index - for index in range(cursor, len(token_ids) - len(values) + 1) - if token_ids[index : index + len(values)] == values - and all( - flag & TokenFlag.SAMPLED - for flag in flags[index : index + len(values)] + for position, (choice, (prompt_ids, completion_ids)) in enumerate( + zip(sourced_choices, exact_metadata, strict=True) + ): + offset = len(prompt_ids) + if ( + offset < cursor + or offset > len(token_ids) + or token_ids[:offset] != prompt_ids + ): + if has_routing: + raise RuntimeError( + "MoE routed prompt tokens are absent from tokenized history" ) - ] - - matches = matching_offsets(completion_ids) - if not matches and completion_ids: - for retained_start in range(1, len(completion_ids)): - matches = matching_offsets(completion_ids[retained_start:]) - if matches: - break - if not matches: + return None + next_boundary = ( + len(exact_metadata[position + 1][0]) + if position + 1 < len(exact_metadata) + else len(token_ids) + ) + end = offset + while ( + end < min(next_boundary, len(token_ids)) and flags[end] & TokenFlag.SAMPLED + ): + end += 1 + retained_ids = token_ids[offset:end] + if not retained_ids or not completion_ids: if has_routing: raise RuntimeError( "MoE routed completion tokens are absent from tokenized history" ) return None - if retained_start and len(matches) != 1: + retained_start = len(completion_ids) - len(retained_ids) + if retained_start < 0 or completion_ids[retained_start:] != retained_ids: if has_routing: raise RuntimeError( - "MoE routed completion suffix is ambiguous in tokenized history" + "MoE routed completion suffix disagrees with tokenized history" ) return None - retained_ids = completion_ids[retained_start:] - offset = matches[0] choices.append( _choice_with_retained_routing(choice, retained_start) if retained_start @@ -687,7 +694,7 @@ def tokenize_trajectory_groups( image_processor: BaseImageProcessor | None = None, chat_template_kwargs: dict[str, Any] | None = None, chat_template_tool_schema_format: ChatTemplateToolSchemaFormat = "default", - model: str | None = None, + model: ModelSelector | str | None = None, ) -> Generator["TokenizedResult", None, None]: for group in trajectory_groups: if not group: @@ -707,12 +714,9 @@ def tokenize_trajectory_groups( if advantage == 0 and drop_zero_advantage_trajectories: continue if trajectory.exchanges: - from ..trajectories._tokenize import ( - _as_tokenizer, - _require_training_model, - ) + from ..trajectories._tokenize import _as_tokenizer - selected_model = _require_training_model(trajectory, model) + selected_model = resolve_training_model(trajectory, model) exchange_results = trajectory.tokenize( multi_history=True, model=selected_model, diff --git a/src/art/tinker_native/backend.py b/src/art/tinker_native/backend.py index e6a28cd3c..2aba95947 100644 --- a/src/art/tinker_native/backend.py +++ b/src/art/tinker_native/backend.py @@ -41,7 +41,7 @@ from ..tinker.backend import get_renderer_name from ..tinker.server import get_free_port from ..trajectories import Trajectory, TrajectoryGroup -from ..trajectories._tokenize import _training_model_pattern +from ..trajectories._selection import automatic_training_model_selector from ..types import TrainResult, TrainSFTConfig from ..utils.lifecycle import process_shutdown_timeout from ..utils.output_dirs import get_model_dir @@ -359,7 +359,7 @@ async def train( state.tokenizer, normalize_advantages, base_model=model.base_model, - model=_training_model_pattern(self._model_inference_name(model)), + model=automatic_training_model_selector(self._model_inference_name(model)), ) metrics: dict[str, float] = { diff --git a/src/art/tinker_native/data.py b/src/art/tinker_native/data.py index 500b69e94..d500e0fd9 100644 --- a/src/art/tinker_native/data.py +++ b/src/art/tinker_native/data.py @@ -17,6 +17,7 @@ TrajectoryGroup, get_messages, ) +from ..trajectories._selection import ModelSelector, resolve_training_model from ..types import MessagesAndChoices @@ -142,7 +143,7 @@ def trajectory_groups_to_datums( normalize_advantages: bool = True, *, base_model: str | None = None, - model: str | None = None, + model: ModelSelector | str | None = None, ) -> list[tinker.Datum]: datums: list[tinker.Datum] = [] @@ -163,9 +164,7 @@ def trajectory_groups_to_datums( continue for trajectory, advantage in zip(group.trajectories, advantages): if trajectory.exchanges: - from ..trajectories._tokenize import _require_training_model - - selected_model = _require_training_model(trajectory, model) + selected_model = resolve_training_model(trajectory, model) tokenized = trajectory.tokenize( multi_history=True, model=selected_model, diff --git a/src/art/trajectories/__init__.py b/src/art/trajectories/__init__.py index b12c9a280..54a6b0de5 100644 --- a/src/art/trajectories/__init__.py +++ b/src/art/trajectories/__init__.py @@ -807,6 +807,8 @@ def tokenize( class TokenizedHistory(pydantic.BaseModel): + model_config = pydantic.ConfigDict(ser_json_inf_nan="strings") + model: str token_ids: list[int] logprobs: list[float] diff --git a/src/art/trajectories/_selection.py b/src/art/trajectories/_selection.py new file mode 100644 index 000000000..6858dfe20 --- /dev/null +++ b/src/art/trajectories/_selection.py @@ -0,0 +1,107 @@ +from __future__ import annotations + +from dataclasses import dataclass +from fnmatch import fnmatchcase +import re +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from . import Trajectory + + +@dataclass(frozen=True, slots=True) +class ModelSelector: + """Private model selection policy used at training boundaries.""" + + value: str + automatic_family: tuple[str, str] | None = None + + def matches(self, candidate: str) -> bool: + if self.automatic_family is not None: + prefix, separator = self.automatic_family + return bool( + re.fullmatch( + f"{re.escape(prefix)}{re.escape(separator)}[0-9]+", + candidate, + ) + ) + return fnmatchcase(candidate, self.value) + + +def public_model_selector(value: str) -> ModelSelector: + return ModelSelector(value) + + +def automatic_training_model_selector(value: str) -> ModelSelector: + if match := re.fullmatch(r"(.*)@([0-9]+)", value): + return ModelSelector(value, (match.group(1), "@")) + if match := re.fullmatch(r"(.*):step([0-9]+)", value): + return ModelSelector(value, (match.group(1), ":step")) + return ModelSelector(value) + + +def resolve_training_model( + trajectory: Trajectory, + selector: ModelSelector | str | None, +) -> str: + """Resolve one concrete, single-protocol captured model for training.""" + + exchanges_by_protocol = { + "Chat Completions": trajectory.exchanges.chat_completions, + "Completions": trajectory.exchanges.completions, + "Responses": trajectory.exchanges.responses, + "Anthropic Messages": trajectory.exchanges.messages, + } + exchanges = [ + (protocol, exchange) + for protocol, protocol_exchanges in exchanges_by_protocol.items() + for exchange in protocol_exchanges + ] + if not exchanges: + raise ValueError("Exchange training requires at least one captured exchange") + if any(exchange.model is None for _, exchange in exchanges): + raise ValueError("Every training exchange must identify its model") + + concrete_models = {exchange.model for _, exchange in exchanges} + if selector is None: + matches = concrete_models + else: + selector = ( + public_model_selector(selector) if isinstance(selector, str) else selector + ) + exact = { + candidate for candidate in concrete_models if candidate == selector.value + } + matches = exact or { + candidate + for candidate in concrete_models + if candidate is not None and selector.matches(candidate) + } + if not matches: + value = selector.value if isinstance(selector, ModelSelector) else selector + raise ValueError(f"Trajectory contains no exchanges for model {value!r}") + if len(matches) != 1: + raise ValueError( + "Exchange training requires exactly one concrete model; matched " + f"{sorted(matches)}" + ) + selected_model = next(iter(matches)) + if selected_model is None: + raise AssertionError("model identity was checked above") + protocols = { + protocol for protocol, exchange in exchanges if exchange.model == selected_model + } + if len(protocols) != 1: + raise ValueError( + "Exchange training does not support mixed protocols for one model; found " + f"{sorted(protocols)}" + ) + return selected_model + + +__all__ = [ + "ModelSelector", + "automatic_training_model_selector", + "public_model_selector", + "resolve_training_model", +] diff --git a/tests/unit/test_exchange_training_model_selection.py b/tests/unit/test_exchange_training_model_selection.py index 068426ac6..903c541a1 100644 --- a/tests/unit/test_exchange_training_model_selection.py +++ b/tests/unit/test_exchange_training_model_selection.py @@ -1,7 +1,8 @@ from datetime import datetime, timedelta -from typing import cast +from typing import SupportsIndex, cast, overload import numpy as np +from openai.types import Completion from openai.types.chat import ChatCompletion, ChatCompletionMessageParam import pytest from transformers import PreTrainedTokenizerBase @@ -15,8 +16,13 @@ ChatCompletionsHistory, ChatCompletionsMessageSource, ChatCompletionsRequest, + CompletionsExchange, + CompletionsRequest, +) +from art.trajectories._selection import ( + automatic_training_model_selector, + resolve_training_model, ) -from art.trajectories._tokenize import _training_model_pattern def _exchange(model: str, output_token: int) -> ChatCompletionsExchange: @@ -124,7 +130,7 @@ class _Tokenizer: def test_preprocessing_requires_model_selection() -> None: tokenizer = cast(PreTrainedTokenizerBase, _Tokenizer()) - with pytest.raises(ValueError, match="exactly one model"): + with pytest.raises(ValueError, match="exactly one concrete model"): list( tokenize_trajectory_groups( tokenizer, @@ -148,7 +154,7 @@ def test_preprocessing_requires_model_selection() -> None: def test_tinker_requires_model_selection() -> None: - with pytest.raises(ValueError, match="exactly one model"): + with pytest.raises(ValueError, match="exactly one concrete model"): trajectory_groups_to_datums([_group()], renderer=None, tokenizer=None) datums = trajectory_groups_to_datums( @@ -161,27 +167,26 @@ def test_tinker_requires_model_selection() -> None: assert all(datum.model_input.to_ints() == [1] for datum in datums) -def test_training_model_patterns_select_all_policy_versions() -> None: +def test_training_rejects_multiple_concrete_policy_versions() -> None: group = _versioned_group() - results = list( - tokenize_trajectory_groups( - cast(PreTrainedTokenizerBase, _Tokenizer()), + with pytest.raises(ValueError, match="exactly one concrete model"): + list( + tokenize_trajectory_groups( + cast(PreTrainedTokenizerBase, _Tokenizer()), + [group], + allow_training_without_logprobs=False, + scale_rewards=False, + model="policy@*", + ) + ) + + with pytest.raises(ValueError, match="exactly one concrete model"): + trajectory_groups_to_datums( [group], - allow_training_without_logprobs=False, - scale_rewards=False, + renderer=None, + tokenizer=None, model="policy@*", ) - ) - assert len(results) == 4 - assert {result.token_ids[-1] for result in results} == {2, 4} - - datums = trajectory_groups_to_datums( - [group], - renderer=None, - tokenizer=None, - model="policy@*", - ) - assert len(datums) == 4 trajectory = group.trajectories[0] with pytest.raises(ValueError, match="exactly one history"): @@ -194,19 +199,87 @@ def test_training_model_patterns_select_all_policy_versions() -> None: @pytest.mark.parametrize( - ("model", "pattern"), + ("model", "matches", "misses"), [ - ("policy@12", "policy@*"), + ("policy@12", ("policy@0", "policy@12"), ("policy@x", "policy@12x")), ( "wandb-artifact:///entity/project/run:step12", - "wandb-artifact:///entity/project/run:step*", + ( + "wandb-artifact:///entity/project/run:step0", + "wandb-artifact:///entity/project/run:step12", + ), + ("wandb-artifact:///entity/project/run:stepx",), ), - ("policy:active", "policy:active"), - ("base/model", "base/model"), + ("policy:active", ("policy:active",), ("policy:active2",)), + ("base/model", ("base/model",), ("base/model@1",)), ], ) -def test_automatic_training_model_pattern(model: str, pattern: str) -> None: - assert _training_model_pattern(model) == pattern +def test_automatic_training_model_selector( + model: str, matches: tuple[str, ...], misses: tuple[str, ...] +) -> None: + selector = automatic_training_model_selector(model) + assert all(selector.matches(candidate) for candidate in matches) + assert not any(selector.matches(candidate) for candidate in misses) + + +def test_automatic_training_model_selector_treats_family_metacharacters_literally() -> ( + None +): + selector = automatic_training_model_selector("policy[blue]*@12") + assert selector.matches("policy[blue]*@13") + assert not selector.matches("policyb@13") + assert not selector.matches("policy[blue]anything@13") + + +def test_public_training_selector_prefers_exact_model_over_glob_interpretation() -> ( + None +): + trajectory = art.Trajectory( + exchanges=art.TrajectoryExchanges( + chat_completions=[ + _exchange("policy[*]", 2), + _exchange("policyx", 3), + ] + ) + ) + assert resolve_training_model(trajectory, "policy[*]") == "policy[*]" + + +def test_training_selector_rejects_zero_matches() -> None: + with pytest.raises(ValueError, match="no exchanges"): + resolve_training_model(_group().trajectories[0], "missing*") + + +def test_training_selector_rejects_mixed_protocols() -> None: + start = datetime(2026, 1, 1) + completion = CompletionsExchange( + request=CompletionsRequest(model="policy", prompt="question"), + response=Completion.model_validate( + { + "id": "cmpl", + "object": "text_completion", + "created": 1, + "model": "policy", + "choices": [ + { + "index": 0, + "finish_reason": "stop", + "text": "answer", + } + ], + } + ), + start_time=start, + end_time=start + timedelta(seconds=1), + ) + trajectory = art.Trajectory( + exchanges=art.TrajectoryExchanges( + chat_completions=[_exchange("policy", 2)], + completions=[completion], + ) + ) + with pytest.raises(ValueError, match="mixed protocols"): + resolve_training_model(trajectory, "policy") def test_preprocessing_preserves_adjacent_choice_boundaries_for_moe() -> None: @@ -415,6 +488,87 @@ def test_ambiguous_non_moe_suffix_falls_back_to_sampled_spans() -> None: ) +def test_chat_choice_trace_anchors_retained_suffix_at_its_prompt_boundary() -> None: + first = _exchange("policy", 8) + first_extra = first.response.choices[0].model_extra + assert first_extra is not None + first_extra["prompt_token_ids"] = [1] + first_extra["token_ids"] = [7, 8] + second = _exchange("policy", 8) + second_extra = second.response.choices[0].model_extra + assert second_extra is not None + second_extra["prompt_token_ids"] = [1, 8, 9] + second_extra["token_ids"] = [7, 8] + history = ChatCompletionsHistory( + model="policy", + messages=[], + message_sources=[ + ChatCompletionsMessageSource(exchange=first, choice_index=0), + ChatCompletionsMessageSource(exchange=second, choice_index=0), + ], + ) + + trace = _chat_choice_trace( + history, + [1, 8, 9, 7, 8], + [ + art.TokenFlag.EXACT, + art.TokenFlag.EXACT | art.TokenFlag.SAMPLED, + art.TokenFlag.EXACT, + art.TokenFlag.EXACT | art.TokenFlag.SAMPLED, + art.TokenFlag.EXACT | art.TokenFlag.SAMPLED, + ], + ) + + assert trace is not None + assert trace.offsets == [1, 3] + assert trace.lengths == [1, 2] + + +def test_chat_choice_trace_does_bounded_work_for_a_retained_suffix() -> None: + class CountingList(list[int]): + slice_reads = 0 + + @overload + def __getitem__(self, key: SupportsIndex, /) -> int: ... + + @overload + def __getitem__(self, key: slice[SupportsIndex | None], /) -> list[int]: ... + + def __getitem__( + self, key: SupportsIndex | slice[SupportsIndex | None], / + ) -> int | list[int]: + if isinstance(key, slice): + self.slice_reads += 1 + return super().__getitem__(key) + + exchange = _exchange("policy", 7) + extra = exchange.response.choices[0].model_extra + assert extra is not None + extra["prompt_token_ids"] = [1] + extra["token_ids"] = [*([8] * 510), 7] + history = ChatCompletionsHistory( + model="policy", + messages=[], + message_sources=[ + ChatCompletionsMessageSource(exchange=exchange, choice_index=0) + ], + ) + token_ids = CountingList([1, 7, *([0] * 510)]) + flags = [ + art.TokenFlag.EXACT, + art.TokenFlag.EXACT | art.TokenFlag.SAMPLED, + *([art.TokenFlag.EXACT] * 510), + ] + + trace = _chat_choice_trace(history, token_ids, flags) + + assert trace is not None + assert trace.offsets == [1] + assert trace.lengths == [1] + assert token_ids.slice_reads < 10 + + def test_preprocessing_rejects_partial_choice_evidence_before_moe_routes() -> None: first = _routed_exchange( prompt_token_ids=[1], diff --git a/tests/unit/trajectories/test_tokenized_models.py b/tests/unit/trajectories/test_tokenized_models.py new file mode 100644 index 000000000..fc6e58b75 --- /dev/null +++ b/tests/unit/trajectories/test_tokenized_models.py @@ -0,0 +1,66 @@ +import math + +import art + + +def _history() -> art.TokenizedHistory: + return art.TokenizedHistory( + model="policy", + token_ids=[1, 2], + logprobs=[math.nan, -0.25], + flags=[art.TokenFlag.EXACT, art.TokenFlag.EXACT | art.TokenFlag.SAMPLED], + ) + + +def _assert_history_round_trip( + restored: art.TokenizedHistory, expected: art.TokenizedHistory +) -> None: + assert restored.model == expected.model + assert restored.token_ids == expected.token_ids + assert restored.flags == expected.flags + assert math.isnan(restored.logprobs[0]) + assert restored.logprobs[1:] == expected.logprobs[1:] + + +def test_tokenized_history_nan_json_round_trip() -> None: + value = _history() + payload = value.model_dump_json() + assert '"NaN"' in payload + restored = art.TokenizedHistory.model_validate_json(payload) + _assert_history_round_trip(restored, value) + + +def test_tokenized_trajectory_nan_json_round_trip() -> None: + value = art.TokenizedTrajectory( + **_history().model_dump(), + reward=1.0, + metrics={"count": 1}, + metadata={"source": "test"}, + ) + restored = art.TokenizedTrajectory.model_validate_json(value.model_dump_json()) + _assert_history_round_trip(restored, value) + assert restored.reward == value.reward + assert restored.metrics == value.metrics + assert restored.metadata == value.metadata + + +def test_nested_tokenized_models_nan_json_round_trip() -> None: + trajectory = art.TokenizedMultiHistoryTrajectory( + histories=[_history()], + reward=1.0, + metrics={}, + metadata={}, + ) + group = art.TokenizedTrajectoryGroup[art.TokenizedMultiHistoryTrajectory]( + trajectories=[trajectory], + metrics={}, + metadata={}, + ) + restored = art.TokenizedTrajectoryGroup[ + art.TokenizedMultiHistoryTrajectory + ].model_validate_json(group.model_dump_json()) + restored_trajectory = restored.trajectories[0] + _assert_history_round_trip(restored_trajectory.histories[0], _history()) + assert restored_trajectory.reward == trajectory.reward + assert restored.metrics == group.metrics + assert restored.metadata == group.metadata From 3eb8021137fcdba80f3a77761b14a3464f956862 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Fri, 24 Jul 2026 19:06:50 +0000 Subject: [PATCH 10/58] Reject ambiguous automatic policy families --- src/art/trajectories/_selection.py | 8 +++++--- tests/unit/test_exchange_training_model_selection.py | 9 +++++++++ 2 files changed, 14 insertions(+), 3 deletions(-) diff --git a/src/art/trajectories/_selection.py b/src/art/trajectories/_selection.py index 6858dfe20..c7f4c5c3a 100644 --- a/src/art/trajectories/_selection.py +++ b/src/art/trajectories/_selection.py @@ -69,9 +69,11 @@ def resolve_training_model( selector = ( public_model_selector(selector) if isinstance(selector, str) else selector ) - exact = { - candidate for candidate in concrete_models if candidate == selector.value - } + exact = ( + {candidate for candidate in concrete_models if candidate == selector.value} + if selector.automatic_family is None + else set() + ) matches = exact or { candidate for candidate in concrete_models diff --git a/tests/unit/test_exchange_training_model_selection.py b/tests/unit/test_exchange_training_model_selection.py index 903c541a1..d39638eb0 100644 --- a/tests/unit/test_exchange_training_model_selection.py +++ b/tests/unit/test_exchange_training_model_selection.py @@ -1,3 +1,5 @@ +from __future__ import annotations + from datetime import datetime, timedelta from typing import SupportsIndex, cast, overload @@ -231,6 +233,13 @@ def test_automatic_training_model_selector_treats_family_metacharacters_literall assert not selector.matches("policy[blue]anything@13") +def test_automatic_training_selector_rejects_multiple_numeric_steps() -> None: + trajectory = _versioned_group().trajectories[0] + selector = automatic_training_model_selector("policy@12") + with pytest.raises(ValueError, match="exactly one concrete model"): + resolve_training_model(trajectory, selector) + + def test_public_training_selector_prefers_exact_model_over_glob_interpretation() -> ( None ): From 0ab1d5433b7538f52e773f2ce6d20cfa56a1455b Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Fri, 24 Jul 2026 19:08:11 +0000 Subject: [PATCH 11/58] Harden streaming trajectory capture --- src/art/openai.py | 12 +- src/art/trajectories/_capture/core.py | 4 + src/art/trajectories/_capture/httpx.py | 93 ++++++- src/art/trajectories/_protocols.py | 5 + tests/unit/trajectories/test_capture.py | 344 ++++++++++++++++++++++++ 5 files changed, 438 insertions(+), 20 deletions(-) diff --git a/src/art/openai.py b/src/art/openai.py index 2c99fb044..de3a5e07c 100644 --- a/src/art/openai.py +++ b/src/art/openai.py @@ -163,14 +163,14 @@ def update_chat_completion( tool_call.function.arguments += ( tool_call_delta.function.arguments ) - if getattr(chunk_choice.delta, "reasoning", None): - if not hasattr(choice.message, "reasoning"): - setattr(choice.message, "reasoning", "") + for field in ("reasoning", "reasoning_content"): + value = getattr(chunk_choice.delta, field, None) + if not value: + continue setattr( choice.message, - "reasoning", - getattr(choice.message, "reasoning") - + getattr(chunk_choice.delta, "reasoning"), + field, + (getattr(choice.message, field, None) or "") + value, ) chat_completion.service_tier = chunk.service_tier chat_completion.system_fingerprint = chunk.system_fingerprint diff --git a/src/art/trajectories/_capture/core.py b/src/art/trajectories/_capture/core.py index 8de4f154b..b6d3f1cb1 100644 --- a/src/art/trajectories/_capture/core.py +++ b/src/art/trajectories/_capture/core.py @@ -38,6 +38,10 @@ def add(self, chunk: bytes) -> None: if not self.captured: self.body.extend(chunk) + def discard(self) -> None: + self.body.clear() + self.captured = True + def finish(self) -> None: if self.captured: return diff --git a/src/art/trajectories/_capture/httpx.py b/src/art/trajectories/_capture/httpx.py index 31293405c..c599b127e 100644 --- a/src/art/trajectories/_capture/httpx.py +++ b/src/art/trajectories/_capture/httpx.py @@ -4,12 +4,14 @@ import httpx from httpx._client import UseClientDefault +from httpx._decoders import ContentDecoder from httpx._types import AuthTypes from typing_extensions import TypedDict, Unpack from .core import CaptureState, begin, reset _STATE = "_art_trajectory_capture" +_RAW_ACTIVE = "_art_trajectory_capture_raw_active" class _SendOptions(TypedDict, total=False): @@ -18,13 +20,22 @@ class _SendOptions(TypedDict, total=False): follow_redirects: bool | UseClientDefault +def _capture_decoder(response: httpx.Response) -> ContentDecoder: + shadow = httpx.Response( + response.status_code, + headers=response.headers, + stream=httpx.ByteStream(b""), + ) + return shadow._get_content_decoder() + + def install() -> None: if getattr(httpx.Client.send, "_art_capture", False): return original_send = httpx.Client.send original_async_send = httpx.AsyncClient.send - original_iter = httpx.Response.iter_bytes - original_aiter = httpx.Response.aiter_bytes + original_iter = httpx.Response.iter_raw + original_aiter = httpx.Response.aiter_raw original_close = httpx.Response.close original_aclose = httpx.Response.aclose @@ -72,53 +83,107 @@ async def async_send( state.finish() return response - def iter_bytes( + def iter_raw( self: httpx.Response, chunk_size: int | None = None ) -> Iterator[bytes]: state: CaptureState | None = getattr(self, _STATE, None) + if state is None: + yield from original_iter(self, chunk_size) + return + try: + decoder = _capture_decoder(self) + except Exception: + state.discard() + yield from original_iter(self, chunk_size) + return completed = False + usable = True + setattr(self, _RAW_ACTIVE, True) try: for chunk in original_iter(self, chunk_size): - if state is not None: - state.add(chunk) + if usable: + try: + state.add(decoder.decode(chunk)) + except Exception: + state.discard() + usable = False yield chunk completed = True finally: - if state is not None and (completed or state.request.get("stream") is True): + if usable and (completed or state.request.get("stream") is True): + try: + state.add(decoder.flush()) + except Exception: + state.discard() + usable = False + setattr(self, _RAW_ACTIVE, False) + if usable and (completed or state.request.get("stream") is True): state.finish() - async def aiter_bytes( + async def aiter_raw( self: httpx.Response, chunk_size: int | None = None ) -> AsyncIterator[bytes]: state: CaptureState | None = getattr(self, _STATE, None) + if state is None: + async for chunk in original_aiter(self, chunk_size): + yield chunk + return + try: + decoder = _capture_decoder(self) + except Exception: + state.discard() + async for chunk in original_aiter(self, chunk_size): + yield chunk + return completed = False + usable = True + setattr(self, _RAW_ACTIVE, True) try: async for chunk in original_aiter(self, chunk_size): - if state is not None: - state.add(chunk) + if usable: + try: + state.add(decoder.decode(chunk)) + except Exception: + state.discard() + usable = False yield chunk completed = True finally: - if state is not None and (completed or state.request.get("stream") is True): + if usable and (completed or state.request.get("stream") is True): + try: + state.add(decoder.flush()) + except Exception: + state.discard() + usable = False + setattr(self, _RAW_ACTIVE, False) + if usable and (completed or state.request.get("stream") is True): state.finish() def close(self: httpx.Response) -> None: original_close(self) state: CaptureState | None = getattr(self, _STATE, None) - if state is not None and state.request.get("stream") is True: + if ( + state is not None + and not getattr(self, _RAW_ACTIVE, False) + and state.request.get("stream") is True + ): state.finish() async def aclose(self: httpx.Response) -> None: await original_aclose(self) state: CaptureState | None = getattr(self, _STATE, None) - if state is not None and state.request.get("stream") is True: + if ( + state is not None + and not getattr(self, _RAW_ACTIVE, False) + and state.request.get("stream") is True + ): state.finish() setattr(send, "_art_capture", True) setattr(async_send, "_art_capture", True) setattr(httpx.Client, "send", send) setattr(httpx.AsyncClient, "send", async_send) - setattr(httpx.Response, "iter_bytes", iter_bytes) - setattr(httpx.Response, "aiter_bytes", aiter_bytes) + setattr(httpx.Response, "iter_raw", iter_raw) + setattr(httpx.Response, "aiter_raw", aiter_raw) setattr(httpx.Response, "close", close) setattr(httpx.Response, "aclose", aclose) diff --git a/src/art/trajectories/_protocols.py b/src/art/trajectories/_protocols.py index 7d49b77b9..1c420e9df 100644 --- a/src/art/trajectories/_protocols.py +++ b/src/art/trajectories/_protocols.py @@ -257,6 +257,11 @@ def token_logprobs(value: object, field: str) -> list[float] | None: for event_name, payload in _sse_events(body): if not isinstance(payload, dict): continue + event_type = payload.get("type") or event_name + if event_name == "ping" or event_type == "ping": + continue + if event_name == "error" or event_type == "error": + raise ValueError("Anthropic Messages stream returned an error event") if event_name and "type" not in payload: payload = {**payload, "type": event_name} event = adapter.validate_python(payload) diff --git a/tests/unit/trajectories/test_capture.py b/tests/unit/trajectories/test_capture.py index 9353eed4a..f7221adcc 100644 --- a/tests/unit/trajectories/test_capture.py +++ b/tests/unit/trajectories/test_capture.py @@ -4,9 +4,11 @@ from collections.abc import AsyncGenerator, AsyncIterator, Generator import copy from datetime import datetime, timedelta +import gzip import json from typing import Any, cast from unittest.mock import Mock +import zlib import aiohttp from aiohttp import web @@ -129,6 +131,33 @@ } +class _SyncChunks(httpx.SyncByteStream): + def __init__(self, body: bytes) -> None: + self.body = body + + def __iter__(self) -> Generator[bytes, None, None]: + for index in range(0, len(self.body), 3): + yield self.body[index : index + 3] + + +class _AsyncChunks(httpx.AsyncByteStream): + def __init__(self, body: bytes) -> None: + self.body = body + + async def __aiter__(self) -> AsyncIterator[bytes]: + for index in range(0, len(self.body), 3): + yield self.body[index : index + 3] + + +def _encoded(body: bytes, encoding: str) -> bytes: + if encoding == "gzip": + return gzip.compress(body) + if encoding == "deflate": + return zlib.compress(body) + brotli = pytest.importorskip("brotli") + return cast(bytes, brotli.compress(body)) + + @pytest_asyncio.fixture async def endpoint_server(unused_tcp_port: int) -> AsyncIterator[str]: async def handler(request: web.Request) -> web.StreamResponse: @@ -304,6 +333,234 @@ def requests_stream() -> None: ) +@pytest.mark.parametrize("encoding", ["gzip", "deflate", "br"]) +@pytest.mark.parametrize("mode", ["raw", "bytes", "lines"]) +def test_httpx_sync_stream_consumption_captures_decoded_body_once( + encoding: str, mode: str +) -> None: + body = _sse( + [ + ( + None, + { + "id": "chatcmpl-1", + "object": "chat.completion.chunk", + "created": 1, + "model": "test/model", + "choices": [ + { + "index": 0, + "delta": {"role": "assistant", "content": "hello"}, + "finish_reason": "stop", + } + ], + }, + ), + (None, "[DONE]"), + ] + ) + compressed = _encoded(body, encoding) + + def response(_: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + headers={"content-encoding": encoding}, + stream=_SyncChunks(compressed), + ) + + with art.Trajectory() as trajectory: + with httpx.Client(transport=httpx.MockTransport(response)) as client: + with client.stream( + "POST", + "https://example.test/v1/chat/completions", + json={"model": "test/model", "messages": [], "stream": True}, + ) as result: + if mode == "raw": + assert b"".join(result.iter_raw()) == compressed + elif mode == "bytes": + assert b"".join(result.iter_bytes()) == body + else: + list(result.iter_lines()) + + assert len(trajectory.exchanges.chat_completions) == 1 + + +@pytest.mark.parametrize("encoding", ["gzip", "deflate", "br"]) +@pytest.mark.parametrize("mode", ["raw", "bytes", "lines"]) +async def test_httpx_async_stream_consumption_captures_decoded_body_once( + encoding: str, mode: str +) -> None: + body = _sse( + [ + ( + None, + { + "id": "chatcmpl-1", + "object": "chat.completion.chunk", + "created": 1, + "model": "test/model", + "choices": [ + { + "index": 0, + "delta": {"role": "assistant", "content": "hello"}, + "finish_reason": "stop", + } + ], + }, + ), + (None, "[DONE]"), + ] + ) + compressed = _encoded(body, encoding) + + async def response(_: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + headers={"content-encoding": encoding}, + stream=_AsyncChunks(compressed), + ) + + with art.Trajectory() as trajectory: + async with httpx.AsyncClient(transport=httpx.MockTransport(response)) as client: + async with client.stream( + "POST", + "https://example.test/v1/chat/completions", + json={"model": "test/model", "messages": [], "stream": True}, + ) as result: + if mode == "raw": + assert ( + b"".join([chunk async for chunk in result.aiter_raw()]) + == compressed + ) + elif mode == "bytes": + assert ( + b"".join([chunk async for chunk in result.aiter_bytes()]) + == body + ) + else: + _ = [line async for line in result.aiter_lines()] + + assert len(trajectory.exchanges.chat_completions) == 1 + + +def test_httpx_raw_capture_failure_does_not_change_user_stream() -> None: + malformed = b"not a gzip stream" + + def response(_: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + headers={"content-encoding": "gzip"}, + stream=_SyncChunks(malformed), + ) + + with art.Trajectory() as trajectory: + with httpx.Client(transport=httpx.MockTransport(response)) as client: + with client.stream( + "POST", + "https://example.test/v1/chat/completions", + json={"model": "test/model", "messages": [], "stream": True}, + ) as result: + assert b"".join(result.iter_raw()) == malformed + + assert not trajectory.exchanges + + +async def test_httpx_async_raw_capture_failure_does_not_change_user_stream() -> None: + malformed = b"not a gzip stream" + + async def response(_: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + headers={"content-encoding": "gzip"}, + stream=_AsyncChunks(malformed), + ) + + with art.Trajectory() as trajectory: + async with httpx.AsyncClient(transport=httpx.MockTransport(response)) as client: + async with client.stream( + "POST", + "https://example.test/v1/chat/completions", + json={"model": "test/model", "messages": [], "stream": True}, + ) as result: + assert ( + b"".join([chunk async for chunk in result.aiter_raw()]) == malformed + ) + + assert not trajectory.exchanges + + +def _streaming_chat_body() -> bytes: + return _sse( + [ + ( + None, + { + "id": "chatcmpl-1", + "object": "chat.completion.chunk", + "created": 1, + "model": "test/model", + "choices": [ + { + "index": 0, + "delta": {"role": "assistant", "content": "hello"}, + "finish_reason": "stop", + } + ], + }, + ), + (None, "[DONE]"), + ] + ) + + +def test_httpx_abandoned_compressed_raw_stream_is_excluded() -> None: + compressed = gzip.compress(_streaming_chat_body()) + + def response(_: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + headers={"content-encoding": "gzip"}, + stream=_SyncChunks(compressed), + ) + + with art.Trajectory() as trajectory: + with httpx.Client(transport=httpx.MockTransport(response)) as client: + with client.stream( + "POST", + "https://example.test/v1/chat/completions", + json={"model": "test/model", "messages": [], "stream": True}, + ) as result: + iterator = cast(Generator[bytes, None, None], result.iter_raw()) + next(iterator) + iterator.close() + + assert not trajectory.exchanges + + +async def test_httpx_abandoned_compressed_async_raw_stream_is_excluded() -> None: + compressed = gzip.compress(_streaming_chat_body()) + + async def response(_: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + headers={"content-encoding": "gzip"}, + stream=_AsyncChunks(compressed), + ) + + with art.Trajectory() as trajectory: + async with httpx.AsyncClient(transport=httpx.MockTransport(response)) as client: + async with client.stream( + "POST", + "https://example.test/v1/chat/completions", + json={"model": "test/model", "messages": [], "stream": True}, + ) as result: + iterator = cast(AsyncGenerator[bytes, None], result.aiter_raw()) + await anext(iterator) + await iterator.aclose() + + assert not trajectory.exchanges + + async def test_aiohttp_capture_covers_stream_reader_consumption_methods( endpoint_server: str, ) -> None: @@ -635,6 +892,7 @@ def test_all_streaming_protocols_reconstruct_final_responses() -> None: } response_event = {"type": "response.completed", "response": RESPONSE} message_events = [ + ("ping", {"type": "ping"}), ( "message_start", { @@ -722,6 +980,53 @@ def test_all_streaming_protocols_reconstruct_final_responses() -> None: assert getattr(exchange.response, "logprobs") == [-0.2] assert getattr(exchange.response, "prompt_token_ids") == [1] + with art.Trajectory() as trajectory: + state, token = begin( + "POST", + "https://example.test/v1/messages", + {"model": "test/model", "messages": [], "stream": True}, + ) + reset(token) + assert state is not None + state.status_code = 200 + state.add(_sse(message_events)) + state.finish() + + assert len(trajectory.exchanges.messages) == 1 + + +def test_streaming_messages_error_event_is_rejected_and_not_captured() -> None: + body = _sse( + [ + ( + "error", + { + "type": "error", + "error": {"type": "api_error", "message": "failed"}, + }, + ) + ] + ) + request = {"model": "test/model", "messages": [], "stream": True} + with pytest.raises(ValueError, match="returned an error event"): + build_exchange( + "messages", + request, + body, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + with art.Trajectory() as trajectory: + state, token = begin("POST", "https://example.test/v1/messages", request) + reset(token) + assert state is not None + state.status_code = 200 + state.add(body) + state.finish() + + assert not trajectory.exchanges + @pytest.mark.parametrize( ("field", "value"), @@ -836,6 +1141,45 @@ def chunk(index: int, content: str) -> dict[str, Any]: ] +def test_streaming_chat_preserves_reasoning_fields() -> None: + now = datetime.now() + + def chunk(**delta: str) -> dict[str, Any]: + return { + "id": "chatcmpl-1", + "object": "chat.completion.chunk", + "created": 1, + "model": "test/model", + "choices": [ + { + "index": 0, + "delta": {"role": "assistant", **delta}, + "finish_reason": None, + "logprobs": None, + } + ], + } + + exchange = build_exchange( + "chat_completions", + {"model": "test/model", "messages": [], "stream": True}, + _sse( + [ + (None, chunk(reasoning="r1", reasoning_content="c1")), + (None, chunk(reasoning="r2", reasoning_content="c2")), + (None, "[DONE]"), + ] + ), + start_time=now, + end_time=now + timedelta(seconds=1), + ) + + assert isinstance(exchange, ChatCompletionsExchange) + message = exchange.response.choices[0].message + assert getattr(message, "reasoning") == "r1r2" + assert getattr(message, "reasoning_content") == "c1c2" + + def test_streaming_chat_ignores_keepalives_and_azure_prologue() -> None: chunk = { "id": "chatcmpl-1", From df889ce9bfb3692c706a9746084ecb6884a39007 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Fri, 24 Jul 2026 19:08:13 +0000 Subject: [PATCH 12/58] Preserve explicit vLLM chat metadata options --- .../test_runtime_project_isolation.py | 48 +++++++++++++++---- vllm_runtime/src/art_vllm_runtime/patches.py | 9 ++-- 2 files changed, 41 insertions(+), 16 deletions(-) diff --git a/tests/integration/megatron/runtime_isolation/test_runtime_project_isolation.py b/tests/integration/megatron/runtime_isolation/test_runtime_project_isolation.py index eb7501217..ac8baf57d 100644 --- a/tests/integration/megatron/runtime_isolation/test_runtime_project_isolation.py +++ b/tests/integration/megatron/runtime_isolation/test_runtime_project_isolation.py @@ -58,24 +58,52 @@ def test_runtime_patch_always_returns_token_ids( ) -> None: payload = _runtime_python( "import json; " - "from art_vllm_runtime.patches import apply_vllm_runtime_patches; " - "apply_vllm_runtime_patches(); " + "from art_vllm_runtime.patches import subclass_chat_completion_request; " + "subclass_chat_completion_request(); " "from vllm.entrypoints.openai.chat_completion import protocol; " - "request = protocol.ChatCompletionRequest(" + "default_request = protocol.ChatCompletionRequest(" "model='m', messages=[{'role': 'user', 'content': 'x'}]" "); " + "explicit_request = protocol.ChatCompletionRequest(" + "model='m', messages=[{'role': 'user', 'content': 'x'}], " + "logprobs=False, top_logprobs=None, return_token_ids=None" + "); " "print(json.dumps({" - "'logprobs': request.logprobs, " - "'top_logprobs': request.top_logprobs, " - "'return_token_ids': request.return_token_ids" + "'default': {" + "'logprobs': default_request.logprobs, " + "'top_logprobs': default_request.top_logprobs, " + "'return_token_ids': default_request.return_token_ids" + "}, " + "'default_fields_set': sorted(default_request.model_fields_set), " + "'explicit': {" + "'logprobs': explicit_request.logprobs, " + "'top_logprobs': explicit_request.top_logprobs, " + "'return_token_ids': explicit_request.return_token_ids" + "}, " + "'explicit_fields_set': sorted(explicit_request.model_fields_set)" "}))", artifact_dir, "route_token_ids", ) - assert json.loads(payload) == { - "logprobs": True, - "top_logprobs": 0, - "return_token_ids": True, + assert json.loads(payload.splitlines()[-1]) == { + "default": { + "logprobs": True, + "top_logprobs": 0, + "return_token_ids": True, + }, + "default_fields_set": ["messages", "model"], + "explicit": { + "logprobs": False, + "top_logprobs": None, + "return_token_ids": None, + }, + "explicit_fields_set": [ + "logprobs", + "messages", + "model", + "return_token_ids", + "top_logprobs", + ], } diff --git a/vllm_runtime/src/art_vllm_runtime/patches.py b/vllm_runtime/src/art_vllm_runtime/patches.py index 97bc4fdb7..cef798784 100644 --- a/vllm_runtime/src/art_vllm_runtime/patches.py +++ b/vllm_runtime/src/art_vllm_runtime/patches.py @@ -110,12 +110,9 @@ def subclass_chat_completion_request() -> None: return class ChatCompletionRequest(protocol.ChatCompletionRequest): - def __init__(self, *args: object, **kwargs: object) -> None: - super().__init__(*args, **kwargs) # ty:ignore[invalid-argument-type] - self.logprobs = True - if self.top_logprobs is None: - self.top_logprobs = 0 - self.return_token_ids = True + logprobs: bool | None = True + top_logprobs: int | None = 0 + return_token_ids: bool | None = True protocol.ChatCompletionRequest = ChatCompletionRequest # ty:ignore[invalid-assignment] setattr(protocol, "_art_chat_completion_request_patched", True) From efbfa1250ba883959b4d3e6672ec862c8060509a Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Fri, 24 Jul 2026 19:12:29 +0000 Subject: [PATCH 13/58] Test explicit vLLM token metadata opt-out --- .../test_runtime_project_isolation.py | 28 ++++++++++++------- 1 file changed, 18 insertions(+), 10 deletions(-) diff --git a/tests/integration/megatron/runtime_isolation/test_runtime_project_isolation.py b/tests/integration/megatron/runtime_isolation/test_runtime_project_isolation.py index ac8baf57d..6a82c36cf 100644 --- a/tests/integration/megatron/runtime_isolation/test_runtime_project_isolation.py +++ b/tests/integration/megatron/runtime_isolation/test_runtime_project_isolation.py @@ -64,9 +64,13 @@ def test_runtime_patch_always_returns_token_ids( "default_request = protocol.ChatCompletionRequest(" "model='m', messages=[{'role': 'user', 'content': 'x'}]" "); " - "explicit_request = protocol.ChatCompletionRequest(" + "explicit_false = protocol.ChatCompletionRequest(" "model='m', messages=[{'role': 'user', 'content': 'x'}], " - "logprobs=False, top_logprobs=None, return_token_ids=None" + "logprobs=False, top_logprobs=None, return_token_ids=False" + "); " + "explicit_none = protocol.ChatCompletionRequest(" + "model='m', messages=[{'role': 'user', 'content': 'x'}], " + "return_token_ids=None" "); " "print(json.dumps({" "'default': {" @@ -75,12 +79,14 @@ def test_runtime_patch_always_returns_token_ids( "'return_token_ids': default_request.return_token_ids" "}, " "'default_fields_set': sorted(default_request.model_fields_set), " - "'explicit': {" - "'logprobs': explicit_request.logprobs, " - "'top_logprobs': explicit_request.top_logprobs, " - "'return_token_ids': explicit_request.return_token_ids" + "'explicit_false': {" + "'logprobs': explicit_false.logprobs, " + "'top_logprobs': explicit_false.top_logprobs, " + "'return_token_ids': explicit_false.return_token_ids" "}, " - "'explicit_fields_set': sorted(explicit_request.model_fields_set)" + "'explicit_false_fields_set': sorted(explicit_false.model_fields_set), " + "'explicit_none_return_token_ids': explicit_none.return_token_ids, " + "'explicit_none_fields_set': sorted(explicit_none.model_fields_set)" "}))", artifact_dir, "route_token_ids", @@ -92,18 +98,20 @@ def test_runtime_patch_always_returns_token_ids( "return_token_ids": True, }, "default_fields_set": ["messages", "model"], - "explicit": { + "explicit_false": { "logprobs": False, "top_logprobs": None, - "return_token_ids": None, + "return_token_ids": False, }, - "explicit_fields_set": [ + "explicit_false_fields_set": [ "logprobs", "messages", "model", "return_token_ids", "top_logprobs", ], + "explicit_none_return_token_ids": None, + "explicit_none_fields_set": ["messages", "model", "return_token_ids"], } From eeb249c4719fd5382cd681fbaffe3526905acefd Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Fri, 24 Jul 2026 19:14:06 +0000 Subject: [PATCH 14/58] Fix trajectory tokenization evidence --- src/art/trajectories/_history.py | 14 +- src/art/trajectories/_tokenize.py | 639 ++++++++++++++++++++--- tests/unit/trajectories/test_tokenize.py | 283 +++++++++- 3 files changed, 862 insertions(+), 74 deletions(-) diff --git a/src/art/trajectories/_history.py b/src/art/trajectories/_history.py index b948b8273..a1aaeb5db 100644 --- a/src/art/trajectories/_history.py +++ b/src/art/trajectories/_history.py @@ -129,10 +129,18 @@ class _Branch(Generic[_ItemT, _SourceT, _ContextT]): def _selected_models( exchanges: Sequence[_ExchangeT], model: str | None, protocol: str ) -> list[tuple[str, list[_ExchangeT]]]: + has_exact_match = model is not None and any( + exchange.model == model for exchange in exchanges + ) selected = [ exchange for exchange in exchanges - if model is None or _model_matches(exchange.model, model) + if model is None + or ( + exchange.model == model + if has_exact_match + else _model_matches(exchange.model, model) + ) ] if not selected: suffix = f" for model {model!r}" if model is not None else "" @@ -148,7 +156,9 @@ def _selected_models( def _model_matches(candidate: str | None, pattern: str) -> bool: - return candidate is not None and fnmatchcase(candidate, pattern) + return candidate is not None and ( + candidate == pattern or fnmatchcase(candidate, pattern) + ) def _one(history: Sequence[_ItemT], protocol: str) -> _ItemT: diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index 8a9a047cd..59c01e478 100644 --- a/src/art/trajectories/_tokenize.py +++ b/src/art/trajectories/_tokenize.py @@ -1,11 +1,13 @@ from __future__ import annotations +from bisect import bisect_left from collections.abc import Mapping, Sequence from dataclasses import dataclass +from datetime import datetime from functools import lru_cache import math import re -from typing import TYPE_CHECKING, Any, cast +from typing import TYPE_CHECKING, Any, Literal, cast import warnings from anthropic.types import Message, MessageParam, TextBlock @@ -22,7 +24,9 @@ CompletionsExchange, CompletionsSource, CompletionsStringHistory, + CompletionsStringSourceSpan, CompletionsTokenHistory, + CompletionsTokenSourceSpan, History, LegacyHistory, MessagesExchange, @@ -63,6 +67,125 @@ class _SampledOutput: start: int +@dataclass(frozen=True) +class _SampledSourceKey: + protocol: Literal["chat_completions", "responses", "messages", "completions"] + response_id: str + start_time: datetime + end_time: datetime + index: int + prompt_index: int | None + exchange_identity: int + + +@dataclass +class _HistoryTokenizationTrace: + source_keys: list[_SampledSourceKey | None] + sources: dict[_SampledSourceKey, object] + + def validate(self, tokenized: TokenizedHistory) -> None: + if len(self.source_keys) != len(tokenized.token_ids): + raise AssertionError( + "Tokenization trace differs in length from tokenized data" + ) + for flag, source_key in zip(tokenized.flags, self.source_keys, strict=True): + if bool(flag & TokenFlag.SAMPLED) != (source_key is not None): + raise AssertionError( + "Tokenization trace must identify every sampled token exactly once" + ) + if source_key is not None and source_key not in self.sources: + raise AssertionError("Tokenization trace source key is unresolved") + + +@dataclass +class _TraceBuilder: + trace: _HistoryTokenizationTrace | None = None + + def set( + self, + tokenized: TokenizedHistory, + source_keys: list[_SampledSourceKey | None], + sources: dict[_SampledSourceKey, object], + ) -> None: + trace = _HistoryTokenizationTrace(source_keys=source_keys, sources=sources) + trace.validate(tokenized) + self.trace = trace + + +def _source_key( + exchange: Exchange, + *, + protocol: Literal["chat_completions", "responses", "messages", "completions"], + index: int, + prompt_index: int | None = None, +) -> _SampledSourceKey: + return _SampledSourceKey( + protocol=protocol, + response_id=str(getattr(exchange.response, "id", "")), + start_time=exchange.start_time, + end_time=exchange.end_time, + index=index, + prompt_index=prompt_index, + # The trace is private and consumed during this tokenization call. Object + # identity makes repeated empty provider IDs collision-free while keeping + # all histories projected from one exchange on the same source key. + exchange_identity=id(exchange), + ) + + +def _sampled_source_key(source: object) -> _SampledSourceKey: + exchange = getattr(source, "exchange", None) + if isinstance(exchange, ChatCompletionsExchange): + index = getattr(source, "choice_index", None) + if not isinstance(index, int) or isinstance(index, bool): + raise ValueError("Sampled Chat source has no choice index") + return _source_key(exchange, protocol="chat_completions", index=index) + if isinstance(exchange, ResponsesExchange): + index = getattr(source, "generation_index", None) + if index is None and getattr(source, "output_index", None) is not None: + index = 0 + if not isinstance(index, int) or isinstance(index, bool): + raise ValueError("Sampled Responses source has no generation identity") + return _source_key(exchange, protocol="responses", index=index) + if isinstance(exchange, MessagesExchange): + return _source_key(exchange, protocol="messages", index=0) + if isinstance(exchange, CompletionsExchange): + index = getattr(source, "choice_index", None) + prompt_index = getattr(source, "prompt_index", None) + if not isinstance(index, int) or isinstance(index, bool): + raise ValueError("Sampled Completions source has no choice index") + if not isinstance(prompt_index, int) or isinstance(prompt_index, bool): + raise ValueError("Sampled Completions source has no prompt index") + return _source_key( + exchange, + protocol="completions", + index=index, + prompt_index=prompt_index, + ) + raise ValueError("Sampled token source has an unsupported exchange") + + +def _exchange_sampled_source_key(exchange: Exchange) -> _SampledSourceKey: + if isinstance(exchange, ChatCompletionsExchange): + return _source_key( + exchange, + protocol="chat_completions", + index=exchange.response.choices[0].index, + ) + if isinstance(exchange, CompletionsExchange): + return _source_key( + exchange, + protocol="completions", + index=exchange.response.choices[0].index, + prompt_index=0, + ) + if isinstance(exchange, ResponsesExchange): + return _source_key(exchange, protocol="responses", index=0) + if isinstance(exchange, MessagesExchange): + return _source_key(exchange, protocol="messages", index=0) + raise TypeError(f"Unsupported sampled exchange: {type(exchange).__name__}") + + def _as_tokenizer(tokenizer: object) -> Tokenizer: # Transformers' annotation permits only string-valued message dictionaries, # although its runtime API supports the structured content ART must tokenize. @@ -289,12 +412,22 @@ def _completion_evidence( elif len(pair_logprobs) == len(prompt_ids) + len(selected): prompt_logprobs = pair_logprobs[: len(prompt_ids)] completion_logprobs = pair_logprobs[len(prompt_ids) :] - elif selected[: len(prompt_ids)] == prompt_ids: + elif selected[: len(prompt_ids)] == prompt_ids and ( + pair_ids == selected + or ( + len(tokens) == len(selected) + and all(isinstance(value, str) for value in tokens) + and "".join(cast(str, value) for value in tokens) == choice.text + ) + ): if pair_logprobs and len(pair_logprobs) != len(selected): raise ValueError("Completions token IDs and logprobs differ in length") prompt_logprobs = pair_logprobs[: len(prompt_ids)] selected = selected[len(prompt_ids) :] completion_logprobs = pair_logprobs[len(prompt_ids) :] + elif selected[: len(prompt_ids)] == prompt_ids: + selected = None + completion_logprobs = [] elif pair_logprobs and len(pair_logprobs) != len(selected): raise ValueError("Completions token IDs and logprobs differ in length") elif selected is not None and pair_logprobs and len(pair_logprobs) != len(selected): @@ -1340,6 +1473,7 @@ def _tokenize_exchange_trajectory( chat_template: str | None, chat_template_kwargs: Mapping[str, object] | None, tokenizer_instance: Tokenizer | None = None, + _trace: _TraceBuilder | None = None, ) -> TokenizedHistory: if trajectory.exchanges and ( trajectory.messages_and_choices @@ -1380,6 +1514,8 @@ def _tokenize_exchange_trajectory( token_ids: list[int] = [] logprobs: list[float] = [] flags: list[TokenFlag] = [] + source_keys: list[_SampledSourceKey | None] = [] + sources: dict[_SampledSourceKey, object] = {} response_histories: dict[ str, tuple[list[dict[str, Any]] | None, ResponsesExchange] ] = {} @@ -1499,6 +1635,7 @@ def fallback_config() -> _TokenizerConfig: flags.extend( [TokenFlag.EXACT if prompt_is_exact else TokenFlag(0)] * len(prompt) ) + source_keys.extend([None] * len(prompt)) elif len(prompt) < len(token_ids) or prompt[: len(token_ids)] != token_ids: resolved_config = fallback_config() if tokenizer is None: @@ -1557,6 +1694,7 @@ def fallback_config() -> _TokenizerConfig: token_ids.extend(suffix) logprobs.extend([math.nan] * len(suffix)) flags.extend([TokenFlag(0)] * len(suffix)) + source_keys.extend([None] * len(suffix)) else: suffix = prompt[len(token_ids) :] token_ids.extend(suffix) @@ -1566,6 +1704,7 @@ def fallback_config() -> _TokenizerConfig: flags.extend( [TokenFlag.EXACT if prompt_is_exact else TokenFlag(0)] * len(suffix) ) + source_keys.extend([None] * len(suffix)) if len(completion_logprobs) != len(completion): completion_logprobs = _align_visible_logprobs( tokenizer, completion, exchange @@ -1576,6 +1715,9 @@ def fallback_config() -> _TokenizerConfig: if completion_is_exact: completion_flag |= TokenFlag.EXACT flags.extend([completion_flag] * len(completion)) + source_key = _exchange_sampled_source_key(exchange) + source_keys.extend([source_key] * len(completion)) + sources[source_key] = exchange if completion_is_exact: sampled_outputs.append( _SampledOutput( @@ -1586,12 +1728,15 @@ def fallback_config() -> _TokenizerConfig: ) previous_render_state = (exchange, messages_override) - return TokenizedHistory( + tokenized = TokenizedHistory( model=selected_model, token_ids=token_ids, logprobs=logprobs, flags=flags, ) + if _trace is not None: + _trace.set(tokenized, source_keys, sources) + return tokenized def _unique_exchanges(history: History) -> list[Exchange]: @@ -1848,6 +1993,22 @@ def _validate_history_sources(history: History) -> None: "Responses output source does not belong to " "its generation" ) + expected.extend( + _responses_messages( + { + "input": [ + exchange.response.output[ + generation_output_index + ].model_dump( + mode="python", exclude_none=True + ) + for generation_output_index in generations[ + generation_index + ].output_indices + ] + } + ) + ) expected.extend( _responses_messages( { @@ -1865,16 +2026,8 @@ def _validate_history_sources(history: History) -> None: {"instructions": exchange.request.get("instructions")} ) ) - expected.extend( - _responses_messages( - { - "input": [ - item.model_dump(mode="python", exclude_none=True) - for item in exchange.response.output - ] - } - ) - ) + elif source.generation_index is not None: + expected.append({"role": "assistant", "content": ""}) actual = normalize_chat_message(message) if not any( actual == normalize_chat_message(candidate) for candidate in expected @@ -2028,10 +2181,17 @@ def _last_source_exchange(sources: Sequence[object]) -> Exchange | None: return None -def _history_needs_render(history: History) -> bool: +@dataclass(frozen=True) +class _HistoryRenderState: + needs_render: bool + context_changed: bool = False + projection_matches: bool | None = None + + +def _history_render_state(history: History) -> _HistoryRenderState: if isinstance(history, ChatCompletionsHistory): if any(source is None for source in history.message_sources): - return True + return _HistoryRenderState(needs_render=True, projection_matches=False) for message, source in zip( history.messages, history.message_sources, strict=True ): @@ -2040,20 +2200,26 @@ def _history_needs_render(history: History) -> bool: if dict(message) != original and dict(message) == _without_reasoning( original ): - return True + return _HistoryRenderState(needs_render=True) exchange = _last_source_exchange(history.message_sources) if not isinstance(exchange, ChatCompletionsExchange): - return True + return _HistoryRenderState(needs_render=True) context_changed = ( history.tools != exchange.request.get("tools") or history.chat_template != exchange.request.get("chat_template") or history.chat_template_kwargs != exchange.request.get("chat_template_kwargs") ) - return context_changed or not _history_matches_projection(history) + if context_changed: + return _HistoryRenderState(needs_render=True, context_changed=True) + projection_matches = _history_matches_projection(history) + return _HistoryRenderState( + needs_render=not projection_matches, + projection_matches=projection_matches, + ) if isinstance(history, AnthropicMessagesHistory): if any(source is None for source in history.message_sources): - return True + return _HistoryRenderState(needs_render=True, projection_matches=False) for message, source in zip( history.messages, history.message_sources, strict=True ): @@ -2067,10 +2233,10 @@ def _history_needs_render(history: History) -> bool: ], } if message != expected: - return True + return _HistoryRenderState(needs_render=True) exchange = _last_source_exchange(history.message_sources) if not isinstance(exchange, MessagesExchange): - return False + return _HistoryRenderState(needs_render=False) context_changed = ( history.system != exchange.request.get("system") or history.tools != exchange.request.get("tools") @@ -2078,13 +2244,19 @@ def _history_needs_render(history: History) -> bool: or history.chat_template_kwargs != exchange.request.get("chat_template_kwargs") ) - return context_changed or not _history_matches_projection(history) + if context_changed: + return _HistoryRenderState(needs_render=True, context_changed=True) + projection_matches = _history_matches_projection(history) + return _HistoryRenderState( + needs_render=not projection_matches, + projection_matches=projection_matches, + ) if isinstance(history, ResponsesHistory): if any(source is None for source in history.input_sources): - return True + return _HistoryRenderState(needs_render=True, projection_matches=False) exchange = _last_source_exchange(history.input_sources) if not isinstance(exchange, ResponsesExchange): - return False + return _HistoryRenderState(needs_render=False) context_changed = ( history.instructions != exchange.request.get("instructions") or history.tools != exchange.request.get("tools") @@ -2092,8 +2264,14 @@ def _history_needs_render(history: History) -> bool: or history.chat_template_kwargs != exchange.request.get("chat_template_kwargs") ) - return context_changed or not _history_matches_projection(history) - return False + if context_changed: + return _HistoryRenderState(needs_render=True, context_changed=True) + projection_matches = _history_matches_projection(history) + return _HistoryRenderState( + needs_render=not projection_matches, + projection_matches=projection_matches, + ) + return _HistoryRenderState(needs_render=False) def _source_signature(source: object) -> tuple[object, ...] | None: @@ -2157,9 +2335,14 @@ def _history_matches_projection(history: History) -> bool: ) ) if isinstance(history, ChatCompletionsHistory): - candidates: Sequence[History] = chat_completions_histories( - trajectory, model=history.model - ) + try: + candidates: Sequence[History] = chat_completions_histories( + trajectory, model=history.model + ) + except ValueError as error: + if "no Chat Completions exchanges" not in str(error): + raise + return False return any( isinstance(candidate, ChatCompletionsHistory) and candidate.messages == history.messages @@ -2219,6 +2402,7 @@ def _tokenize_exact_responses_history( *, base_model: str | None, tokenizer: Tokenizer | None, + _trace: _TraceBuilder | None = None, ) -> TokenizedHistory | None: generation_keys: list[tuple[ResponsesExchange, int]] = [] retained_output_indices: dict[tuple[int, int], set[int]] = {} @@ -2238,6 +2422,8 @@ def _tokenize_exact_responses_history( token_ids: list[int] = [] logprobs: list[float] = [] flags: list[TokenFlag] = [] + source_keys: list[_SampledSourceKey | None] = [] + sources: dict[_SampledSourceKey, object] = {} sampled_outputs: list[_SampledOutput] = [] for position, (exchange, generation_index) in enumerate(generation_keys): generations = _response_generations(exchange.response) @@ -2278,12 +2464,14 @@ def _tokenize_exact_responses_history( token_ids.extend(prompt) logprobs.extend([math.nan] * len(prompt)) flags.extend([TokenFlag.EXACT] * len(prompt)) + source_keys.extend([None] * len(prompt)) elif prompt[: len(token_ids)] == token_ids: suffix = prompt[len(token_ids) :] token_ids.extend(suffix) logprobs.extend([math.nan] * len(suffix)) flags = [flag | TokenFlag.EXACT for flag in flags] flags.extend([TokenFlag.EXACT] * len(suffix)) + source_keys.extend([None] * len(suffix)) else: if tokenizer is None and sampled_outputs: tokenizer = _load_tokenizer( @@ -2303,9 +2491,25 @@ def _tokenize_exact_responses_history( token_ids.extend(suffix) logprobs.extend([math.nan] * len(suffix)) flags.extend([TokenFlag.EXACT] * len(suffix)) + source_keys.extend([None] * len(suffix)) token_ids.extend(output) logprobs.extend(output_logprobs) flags.extend([TokenFlag.EXACT | TokenFlag.SAMPLED] * len(output)) + source = next( + ( + item + for item in history.input_sources + if item is not None + and item.exchange is exchange + and item.generation_index == generation_index + ), + None, + ) + if source is None: + raise AssertionError("Responses generation has no history source") + source_key = _sampled_source_key(source) + source_keys.extend([source_key] * len(output)) + sources[source_key] = source sampled_outputs.append( _SampledOutput( text=output_text, @@ -2313,12 +2517,15 @@ def _tokenize_exact_responses_history( start=len(token_ids) - len(output), ) ) - return TokenizedHistory( + tokenized = TokenizedHistory( model=history.model, token_ids=token_ids, logprobs=logprobs, flags=flags, ) + if _trace is not None: + _trace.set(tokenized, source_keys, sources) + return tokenized def _chat_source_full_tokens( @@ -2401,15 +2608,17 @@ def _tokenize_exact_projected_chat_history( history: ChatCompletionsHistory, *, projection_validated: bool = False, + _trace: _TraceBuilder | None = None, ) -> TokenizedHistory | None: if not projection_validated and not _history_matches_projection(history): return None sampled_sources: list[object] = [] seen: set[tuple[object, ...]] = set() - for source in history.message_sources: + for message, source in zip(history.messages, history.message_sources, strict=True): signature = _source_signature(source) if ( - source is not None + message.get("role") == "assistant" + and source is not None and signature is not None and signature not in seen and _source_is_sampled(source) @@ -2438,6 +2647,12 @@ def _tokenize_exact_projected_chat_history( *([TokenFlag.EXACT] * len(final_prompt)), *([TokenFlag.EXACT | TokenFlag.SAMPLED] * len(final_output)), ] + final_key = _sampled_source_key(final_source) + source_keys: list[_SampledSourceKey | None] = [ + *([None] * len(final_prompt)), + *([final_key] * len(final_output)), + ] + sources: dict[_SampledSourceKey, object] = {final_key: final_source} for index, source in enumerate(sampled_sources[:-1]): prompt = _chat_source_prompt_tokens(source) output, output_logprobs = _chat_source_full_tokens(source) @@ -2474,14 +2689,20 @@ def _tokenize_exact_projected_chat_history( return None flags[start:end] = [TokenFlag.EXACT | TokenFlag.SAMPLED] * len(retained_ids) logprobs[start:end] = retained_logprobs + source_key = _sampled_source_key(source) + source_keys[start:end] = [source_key] * len(retained_ids) + sources[source_key] = source if history.model is None: raise ValueError("History tokenization requires a model") - return TokenizedHistory( + tokenized = TokenizedHistory( model=history.model, token_ids=token_ids, logprobs=logprobs, flags=flags, ) + if _trace is not None: + _trace.set(tokenized, source_keys, sources) + return tokenized def _chat_message_parts(message: Mapping[str, object]) -> list[tuple[str, str]]: @@ -2588,7 +2809,6 @@ def _chat_source_tokens( raise ValueError( "Responses source output index does not belong to its generation" ) - return generation.output_token_ids, generation.output_logprobs output_index = getattr(source, "output_index", None) if isinstance(output_index, int) and output_index < len( exchange.response.output @@ -2620,6 +2840,48 @@ def _chat_source_tokens( return None, [] +def _source_covers_complete_sampled_message( + message: Mapping[str, object], source: object +) -> bool: + from ._history import normalize_chat_message + + exchange = getattr(source, "exchange", None) + if isinstance(exchange, ChatCompletionsExchange): + expected = _chat_choice_message(source) + return expected is not None and normalize_chat_message( + message + ) == normalize_chat_message(expected) + if isinstance(exchange, MessagesExchange): + if getattr(source, "request_index", None) is not None: + return False + return normalize_chat_message(message) == normalize_chat_message( + _response_message(exchange) + ) + if not isinstance(exchange, ResponsesExchange): + return False + generation_index = getattr(source, "generation_index", None) + if not isinstance(generation_index, int): + return False + generations = _response_generations(exchange.response) + if not 0 <= generation_index < len(generations): + raise ValueError("Responses source generation index is out of bounds") + if not generations[generation_index].output_indices: + return message.get("role") == "assistant" and not _chat_message_parts(message) + projected = _responses_messages( + { + "input": [ + exchange.response.output[index].model_dump( + mode="python", exclude_none=True + ) + for index in generations[generation_index].output_indices + ] + } + ) + return len(projected) == 1 and normalize_chat_message( + message + ) == normalize_chat_message(projected[0]) + + def _tokenize_chat_view( history: ChatCompletionsHistory, *, @@ -2627,6 +2889,7 @@ def _tokenize_chat_view( tokenizer: Tokenizer | None, chat_template: str | None, chat_template_kwargs: Mapping[str, object] | None, + _trace: _TraceBuilder | None = None, ) -> TokenizedHistory: _validate_history_sources(history) config = ( @@ -2656,27 +2919,55 @@ def _tokenize_chat_view( kwargs.setdefault("thinking_budget", budget) template = chat_template or history.chat_template or config.chat_template ends_with_assistant = bool(messages) and messages[-1].get("role") == "assistant" - rendered = _ids( - tokenizer.apply_chat_template( - messages, - tools=history.tools, - tokenize=True, - add_generation_prompt=not ends_with_assistant, - **({"chat_template": template} if template is not None else {}), - **kwargs, + + def render( + selected_messages: list[dict[str, Any]], *, add_generation_prompt: bool + ) -> list[int]: + return _ids( + resolved_tokenizer.apply_chat_template( + selected_messages, + tools=history.tools, + tokenize=True, + add_generation_prompt=add_generation_prompt, + **({"chat_template": template} if template is not None else {}), + **kwargs, + ) ) + + rendered = render( + messages, + add_generation_prompt=not ends_with_assistant, ) + positions_by_first_token: dict[int, list[int]] = {} + for index, token_id in enumerate(rendered): + positions_by_first_token.setdefault(token_id, []).append(index) + locations_by_needle: dict[tuple[int, ...], list[tuple[int, int]]] = {} + def locations(needle: Sequence[int], start: int) -> list[tuple[int, int]]: if not needle: return [] - return [ - (index, index + len(needle)) - for index in range(start, len(rendered) - len(needle) + 1) - if rendered[index : index + len(needle)] == list(needle) + key = tuple(needle) + if key not in locations_by_needle: + locations_by_needle[key] = [ + (index, index + len(key)) + for index in positions_by_first_token.get(key[0], []) + if rendered[index : index + len(key)] == list(key) + ] + spans = locations_by_needle[key] + return spans[bisect_left(spans, (start, -1)) :] + + replacements: list[ + tuple[ + int, + int, + list[int], + list[float], + bool, + _SampledSourceKey, + object, ] - - replacements: list[tuple[int, int, list[int], list[float], bool]] = [] + ] = [] search_cursor = 0 sampled_texts = { text @@ -2708,7 +2999,9 @@ def part_ids(text: str) -> list[int]: for _, text in _chat_message_parts(message) ] sampled_part_cursor = 0 - for message, source in zip(history.messages, history.message_sources, strict=True): + for message_index, (message, source) in enumerate( + zip(history.messages, history.message_sources, strict=True) + ): parts = _chat_message_parts(message) sampled = ( message.get("role") == "assistant" @@ -2718,6 +3011,11 @@ def part_ids(text: str) -> list[int]: full_exact, full_logprobs = ( _chat_source_full_tokens(source) if sampled else (None, []) ) + complete_sampled_message = ( + sampled + and source is not None + and _source_covers_complete_sampled_message(message, source) + ) source_boundary = False if sampled and source is not None: source_prompt = _chat_source_prompt_tokens(source) @@ -2740,11 +3038,21 @@ def part_ids(text: str) -> list[int]: full_matches = ( locations(full_exact, search_cursor) if sampled and full_exact else [] ) + first_part_matches = ( + locations(part_ids(parts[0][1]), search_cursor) if parts else [] + ) if ( full_exact is not None and full_matches - and source_boundary - and full_matches[0][0] == search_cursor + and ( + (source_boundary and full_matches[0][0] == search_cursor) + or ( + complete_sampled_message + and len(full_matches) == 1 + and first_part_matches + and full_matches[0][0] == first_part_matches[0][0] + ) + ) ): span = full_matches[0] start, end = span @@ -2757,12 +3065,48 @@ def part_ids(text: str) -> list[int]: if len(full_logprobs) == len(full_exact) else [math.nan] * len(full_exact), True, + _sampled_source_key(source), + source, ) ) search_cursor = end sampled_part_cursor += len(parts) continue + if ( + complete_sampled_message + and full_exact is not None + and (len(parts) != 1 or len(full_exact) != len(part_ids(parts[0][1]))) + ): + prompt_render = render(messages[:message_index], add_generation_prompt=True) + completed_render = render( + messages[: message_index + 1], add_generation_prompt=False + ) + if ( + len(completed_render) <= len(prompt_render) + or completed_render[: len(prompt_render)] != prompt_render + or rendered[: len(completed_render)] != completed_render + ): + raise ValueError( + "Could not locate a complete sampled message in the rendered history" + ) + replacements.append( + ( + len(prompt_render), + len(completed_render), + full_exact, + full_logprobs + if len(full_logprobs) == len(full_exact) + else [math.nan] * len(full_exact), + True, + _sampled_source_key(source), + source, + ) + ) + search_cursor = len(completed_render) + sampled_part_cursor += len(parts) + continue + replacement_start = len(replacements) for part, text in parts: if not sampled and text not in sampled_texts: @@ -2810,9 +3154,6 @@ def part_ids(text: str) -> list[int]: if exact is not None and rendered[start : start + len(exact)] == exact: end = start + len(exact) search_cursor = end - elif exact is not None and exact[: len(local)] == local: - exact = exact[: len(local)] - logprobs = logprobs[: len(local)] replacement = exact if exact is not None else local if exact is None and not logprobs: exchange = getattr(source, "exchange", None) @@ -2830,6 +3171,8 @@ def part_ids(text: str) -> list[int]: if len(logprobs) == len(replacement) else [math.nan] * len(replacement), exact is not None, + _sampled_source_key(source), + source, ) ) message_replacements = replacements[replacement_start:] @@ -2859,6 +3202,8 @@ def part_ids(text: str) -> list[int]: rendered[start:end], [math.nan] * (end - start), False, + message_replacements[0][5], + message_replacements[0][6], ) ) if sampled and not parts and full_exact is not None: @@ -2869,33 +3214,52 @@ def part_ids(text: str) -> list[int]: token_ids: list[int] = [] logprobs: list[float] = [] flags: list[TokenFlag] = [] + source_keys: list[_SampledSourceKey | None] = [] + sources: dict[_SampledSourceKey, object] = {} cursor = 0 - for start, end, replacement, replacement_logprobs, exact in sorted(replacements): + for ( + start, + end, + replacement, + replacement_logprobs, + exact, + source_key, + source, + ) in sorted(replacements, key=lambda item: (item[0], item[1])): if start < cursor: raise ValueError("Rendered assistant source spans overlap") token_ids.extend(rendered[cursor:start]) logprobs.extend([math.nan] * (start - cursor)) flags.extend([TokenFlag(0)] * (start - cursor)) + source_keys.extend([None] * (start - cursor)) token_ids.extend(replacement) logprobs.extend(replacement_logprobs) flag = TokenFlag.SAMPLED | (TokenFlag.EXACT if exact else TokenFlag(0)) flags.extend([flag] * len(replacement)) + source_keys.extend([source_key] * len(replacement)) + sources[source_key] = source cursor = end token_ids.extend(rendered[cursor:]) logprobs.extend([math.nan] * (len(rendered) - cursor)) flags.extend([TokenFlag(0)] * (len(rendered) - cursor)) + source_keys.extend([None] * (len(rendered) - cursor)) if history.model is None: raise ValueError("History tokenization requires a model") - return TokenizedHistory( + tokenized = TokenizedHistory( model=history.model, token_ids=token_ids, logprobs=logprobs, flags=flags, ) + if _trace is not None: + _trace.set(tokenized, source_keys, sources) + return tokenized def _tokenize_completions_token_history( history: CompletionsTokenHistory, + *, + _trace: _TraceBuilder | None = None, ) -> TokenizedHistory: if any( span.start < 0 or span.end <= span.start or span.end > len(history.prompt) @@ -2906,9 +3270,16 @@ def _tokenize_completions_token_history( raise ValueError( "Completions token source spans must exhaustively cover prompt" ) + _validate_completions_sources( + model=history.model, + source_spans=history.prompt_sources, + sampled_spans=history.sampled_spans, + ) flags = [TokenFlag(0)] * len(history.prompt) logprobs = [math.nan] * len(history.prompt) + source_keys: list[_SampledSourceKey | None] = [None] * len(history.prompt) + sources: dict[_SampledSourceKey, object] = {} for start, end in history.sampled_spans: if start < 0 or end <= start or end > len(history.prompt): raise ValueError("Completions sampled spans are out of bounds") @@ -2916,6 +3287,10 @@ def _tokenize_completions_token_history( for span in history.prompt_sources: if span.source is None: continue + if span.source.choice_index is not None: + source_key = _sampled_source_key(span.source) + source_keys[span.start : span.end] = [source_key] * (span.end - span.start) + sources[source_key] = span.source flags[span.start : span.end] = [ flag | TokenFlag.EXACT for flag in flags[span.start : span.end] ] @@ -2933,12 +3308,15 @@ def _tokenize_completions_token_history( continue if len(selected_logprobs) == span.end - span.start: logprobs[span.start : span.end] = selected_logprobs - return TokenizedHistory( + tokenized = TokenizedHistory( model=history.model, token_ids=list(history.prompt), logprobs=logprobs, flags=flags, ) + if _trace is not None: + _trace.set(tokenized, source_keys, sources) + return tokenized def _completion_visible_logprobs( @@ -2958,7 +3336,10 @@ def _completion_visible_logprobs( if not isinstance(raw_tokens, list) or not isinstance(raw_logprobs, list): return None values = list(zip(raw_tokens, raw_logprobs, strict=False)) - request_prompt = source.exchange.request.get("prompt") + from ._history import _completion_prompts + + request_prompts = _completion_prompts(source.exchange.request.get("prompt")) + request_prompt = request_prompts[source.prompt_index] if source.exchange.request.get("echo") is True and isinstance(request_prompt, str): consumed = "" cursor = 0 @@ -2998,11 +3379,17 @@ def _tokenize_completions_string_history( *, base_model: str | None, tokenizer: Tokenizer | None, + _trace: _TraceBuilder | None = None, ) -> TokenizedHistory: if not _spans_are_exhaustive(len(history.prompt), history.prompt_sources): raise ValueError( "Completions string source spans must exhaustively cover prompt" ) + _validate_completions_sources( + model=history.model, + source_spans=history.prompt_sources, + sampled_spans=history.sampled_spans, + ) sampled = [False] * len(history.prompt) for start, end in history.sampled_spans: if start < 0 or end <= start or end > len(history.prompt): @@ -3021,6 +3408,8 @@ def resolved_tokenizer() -> Tokenizer: token_ids: list[int] = [] logprobs: list[float] = [] flags: list[TokenFlag] = [] + source_keys: list[_SampledSourceKey | None] = [] + sources: dict[_SampledSourceKey, object] = {} for span in history.prompt_sources: text = history.prompt[span.start : span.end] source = span.source @@ -3052,11 +3441,20 @@ def resolved_tokenizer() -> Tokenizer: if item.index == source.choice_index ) expected = choice.text - request_prompt = source.exchange.request.get("prompt") - if source.exchange.request.get("echo") is True and isinstance( - request_prompt, str - ): - expected = expected.removeprefix(request_prompt) + from ._history import _completion_prompts + + request_prompts = _completion_prompts( + source.exchange.request.get("prompt") + ) + request_prompt = request_prompts[source.prompt_index] + if source.exchange.request.get("echo") is True: + if not isinstance(request_prompt, str) or not expected.startswith( + request_prompt + ): + raise ValueError( + "Cannot locate echoed Completions prompt boundary" + ) + expected = expected[len(request_prompt) :] if text != expected: raise ValueError( "Completions history text no longer matches its source exchange" @@ -3082,12 +3480,23 @@ def resolved_tokenizer() -> Tokenizer: if exact is not None: flag |= TokenFlag.EXACT flags.extend([flag] * len(ids)) - return TokenizedHistory( + if is_sampled: + if source is None: + raise AssertionError("Sampled Completions span has no source") + source_key = _sampled_source_key(source) + source_keys.extend([source_key] * len(ids)) + sources[source_key] = source + else: + source_keys.extend([None] * len(ids)) + tokenized = TokenizedHistory( model=history.model, token_ids=token_ids, logprobs=logprobs, flags=flags, ) + if _trace is not None: + _trace.set(tokenized, source_keys, sources) + return tokenized def _completion_source_evidence( @@ -3127,6 +3536,31 @@ def _completion_source_evidence( ) +def _validate_completions_sources( + *, + model: str, + source_spans: Sequence[CompletionsTokenSourceSpan | CompletionsStringSourceSpan], + sampled_spans: Sequence[tuple[int, int]], +) -> None: + expected_sampled: list[tuple[int, int]] = [] + for span in source_spans: + source = getattr(span, "source", None) + if source is None: + continue + if not isinstance(source, CompletionsSource): + raise ValueError("Completions history has an invalid source") + if source.exchange.model != model: + raise ValueError( + "Completions history model no longer matches its source exchange" + ) + if source.choice_index is not None: + expected_sampled.append((span.start, span.end)) + if list(sampled_spans) != expected_sampled: + raise ValueError( + "Completions sampled spans must exactly match choice-backed source spans" + ) + + def _spans_are_exhaustive(length: int, spans: Sequence[object]) -> bool: cursor = 0 for span in spans: @@ -3152,6 +3586,7 @@ def tokenize_history( tokenizer: Tokenizer | None, chat_template: str | None, chat_template_kwargs: Mapping[str, object] | None, + _trace: _TraceBuilder | None = None, ) -> TokenizedHistory: if isinstance(history, LegacyHistory): if model is None: @@ -3160,12 +3595,13 @@ def tokenize_history( if model is None: raise ValueError("History tokenization requires a model") if isinstance(history, CompletionsTokenHistory): - return _tokenize_completions_token_history(history) + return _tokenize_completions_token_history(history, _trace=_trace) if isinstance(history, CompletionsStringHistory): return _tokenize_completions_string_history( history, base_model=base_model, tokenizer=tokenizer, + _trace=_trace, ) _validate_history_sources(history) override_requires_render = ( @@ -3176,15 +3612,21 @@ def tokenize_history( and dict(chat_template_kwargs) != (getattr(history, "chat_template_kwargs", None) or {}) ) - needs_render = _history_needs_render(history) or override_requires_render + render_state = _history_render_state(history) + needs_render = render_state.needs_render or override_requires_render if isinstance(history, ResponsesHistory) and not needs_render: if exact := _tokenize_exact_responses_history( - history, base_model=base_model, tokenizer=tokenizer + history, base_model=base_model, tokenizer=tokenizer, _trace=_trace ): return exact if isinstance(history, ChatCompletionsHistory) and (needs_render): - if not override_requires_render and ( - exact := _tokenize_exact_projected_chat_history(history) + if ( + not override_requires_render + and not render_state.context_changed + and render_state.projection_matches is not False + and ( + exact := _tokenize_exact_projected_chat_history(history, _trace=_trace) + ) ): return exact return _tokenize_chat_view( @@ -3193,16 +3635,23 @@ def tokenize_history( tokenizer=tokenizer, chat_template=chat_template, chat_template_kwargs=chat_template_kwargs, + _trace=_trace, ) if isinstance(history, AnthropicMessagesHistory) and needs_render: converted = history.as_chat_completions_history() if ( not override_requires_render - and _history_matches_projection(history) + and not render_state.context_changed + and ( + render_state.projection_matches + if render_state.projection_matches is not None + else _history_matches_projection(history) + ) and ( exact := _tokenize_exact_projected_chat_history( converted, projection_validated=True, + _trace=_trace, ) ) ): @@ -3213,6 +3662,7 @@ def tokenize_history( tokenizer=tokenizer, chat_template=chat_template, chat_template_kwargs=chat_template_kwargs, + _trace=_trace, ) if isinstance(history, ResponsesHistory) and needs_render: return _tokenize_chat_view( @@ -3221,6 +3671,7 @@ def tokenize_history( tokenizer=tokenizer, chat_template=chat_template, chat_template_kwargs=chat_template_kwargs, + _trace=_trace, ) trajectory = _trajectory_from_history(history) history_template = getattr(history, "chat_template", None) @@ -3238,6 +3689,7 @@ def tokenize_history( } or None, tokenizer_instance=tokenizer, + _trace=_trace, ) @@ -3309,6 +3761,53 @@ def tokenize_trajectory( ) +def _tokenize_trajectory_with_trace( + trajectory: Trajectory, + *, + model: str | None = None, + base_model: str | None = None, + tokenizer: Tokenizer | None = None, + chat_template: str | None = None, + chat_template_kwargs: Mapping[str, object] | None = None, +) -> tuple[ + TokenizedMultiHistoryTrajectory, + list[_HistoryTokenizationTrace], +]: + if not trajectory.exchanges: + raise ValueError("Private exchange tokenization trace requires exchanges") + histories = trajectory.histories(model=model) + tokenized_histories: list[TokenizedHistory] = [] + traces: list[_HistoryTokenizationTrace] = [] + for history in histories: + if isinstance(history, LegacyHistory): + raise AssertionError( + "Exchange trajectories cannot produce legacy histories" + ) + trace_builder = _TraceBuilder() + tokenized = tokenize_history( + history, + model=history.model, + base_model=base_model, + tokenizer=tokenizer, + chat_template=chat_template, + chat_template_kwargs=chat_template_kwargs, + _trace=trace_builder, + ) + if trace_builder.trace is None: + raise AssertionError("Exchange tokenization did not produce a source trace") + tokenized_histories.append(tokenized) + traces.append(trace_builder.trace) + return ( + TokenizedMultiHistoryTrajectory( + histories=tokenized_histories, + reward=trajectory.reward, + metrics=dict(trajectory.metrics), + metadata=dict(trajectory.metadata), + ), + traces, + ) + + def tokenize_group( group: TrajectoryGroup, *, diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index ea78d79c1..487cf7cd6 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -362,6 +362,61 @@ def test_completions_echo_prefers_full_logprob_carrier_to_id_prefix_heuristic() ] +def test_batched_completions_echo_uses_selected_prompt_boundary() -> None: + exchange = _completion_exchange(prompt=["p0", "p1"], echo=True) + payload = exchange.response.model_dump(mode="python") + payload["choices"] = [ + { + "index": 0, + "finish_reason": "stop", + "text": "p0a", + "logprobs": { + "tokens": ["p0", "a"], + "token_logprobs": [-9.0, -0.1], + "top_logprobs": [{}, {}], + "text_offset": [0, 2], + }, + }, + { + "index": 1, + "finish_reason": "stop", + "text": "p1b", + "logprobs": { + "tokens": ["p1", "b"], + "token_logprobs": [-9.0, -0.2], + "top_logprobs": [{}, {}], + "text_offset": [0, 2], + }, + }, + ] + exchange.request["n"] = 1 + exchange.response = Completion.model_validate(payload) + + class Tokenizer: + def __call__(self, text: str, **kwargs: object) -> list[int]: + del kwargs + return {"p0": [1], "p1": [2], "a": [3], "b": [4]}[text] + + def apply_chat_template( + self, messages: list[dict[str, Any]], **kwargs: object + ) -> list[int]: + raise AssertionError((messages, kwargs)) + + histories = art.Trajectory( + exchanges=TrajectoryExchanges(completions=[exchange]) + ).completions_string_histories() + + assert [ + history.tokenize(tokenizer=Tokenizer()).token_ids for history in histories + ] == [ + [1, 3], + [2, 4], + ] + assert [ + history.tokenize(tokenizer=Tokenizer()).logprobs[-1] for history in histories + ] == pytest.approx([-0.1, -0.2]) + + def test_completions_exact_ids_accept_textual_logprobs() -> None: exchange = _completion_exchange() payload = exchange.response.model_dump(mode="python") @@ -392,6 +447,20 @@ def test_mutated_completions_token_prompt_drops_stale_exact_evidence() -> None: ] +def test_completions_history_rejects_model_and_sampled_span_mutation() -> None: + history = art.Trajectory( + exchanges=TrajectoryExchanges(completions=[_completion_exchange()]) + ).completions_token_history() + history.model = "other/model" + with pytest.raises(ValueError, match="model no longer matches"): + history.tokenize() + + history.model = "test/model" + history.sampled_spans = [(0, len(history.prompt))] + with pytest.raises(ValueError, match="exactly match choice-backed"): + history.tokenize() + + def test_completions_string_history_preserves_textual_logprobs() -> None: exchange = _completion_exchange() payload = exchange.response.model_dump(mode="python") @@ -533,6 +602,8 @@ def test_batched_completions_reject_ambiguous_choice_indices() -> None: def test_randomized_completions_projection_preserves_every_choice_once() -> None: + from art.trajectories._tokenize import _tokenize_trajectory_with_trace + rng = random.Random(0) for case in range(20): prompt_count = rng.randint(1, 5) @@ -570,9 +641,10 @@ def test_randomized_completions_projection_preserves_every_choice_once() -> None response["choices"] = choices exchange.response = Completion.model_validate(response) - tokenized = art.Trajectory( + trajectory = art.Trajectory( exchanges=TrajectoryExchanges(completions=[exchange]) - ).tokenize(multi_history=True) + ) + tokenized = trajectory.tokenize(multi_history=True) assert [history.token_ids for history in tokenized.histories] == expected assert all( @@ -580,6 +652,17 @@ def test_randomized_completions_projection_preserves_every_choice_once() -> None == [art.TokenFlag.EXACT, art.TokenFlag.EXACT | art.TokenFlag.SAMPLED] for history in tokenized.histories ) + traced, traces = _tokenize_trajectory_with_trace(trajectory) + assert [history.token_ids for history in traced.histories] == expected + assert [ + next(key for key in trace.source_keys if key is not None).prompt_index + for trace in traces + ] == [ + prompt_index + for prompt_index in range(prompt_count) + for _ in range(choices_per_prompt) + ] + assert all(len(trace.sources) == 1 for trace in traces) def test_branching_and_multiple_models_require_explicit_resolution() -> None: @@ -613,6 +696,20 @@ def test_branching_and_multiple_models_require_explicit_resolution() -> None: ] == ["one", "two"] +def test_model_selection_prefers_exact_identity_over_glob_interpretation() -> None: + literal_model = "org/model[1]" + wildcard_match = _chat_exchange([3], [4], model="org/model1", offset=1) + exact_match = _chat_exchange([1], [2], model=literal_model) + trajectory = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[wildcard_match, exact_match]) + ) + + tokenized = trajectory.tokenize(model=literal_model) + + assert tokenized.model == literal_model + assert tokenized.token_ids == [1, 2] + + def test_legacy_additional_histories_require_multi_history_and_model() -> None: first = _chat_exchange([1], [2]).response.choices[0] second = _chat_exchange([3], [4]).response.choices[0] @@ -2033,6 +2130,106 @@ def apply_chat_template( assert tokenized.flags[1] == art.TokenFlag.EXACT | art.TokenFlag.SAMPLED +def test_template_change_preserves_complete_exact_sampled_suffix() -> None: + exchange = _chat_exchange([1], [2, 3]) + history = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[exchange]) + ).chat_completions_history() + history.chat_template = "custom" + + class Tokenizer: + def __call__(self, text: str, **kwargs: object) -> list[int]: + del kwargs + return {"turn 0": [1], "answer": [2]}[text] + + def apply_chat_template( + self, messages: list[dict[str, Any]], **kwargs: object + ) -> list[int]: + del kwargs + if messages[-1]["role"] != "assistant": + return [1] + return [1, 2, 9] + + tokenized = history.tokenize(tokenizer=Tokenizer()) + + assert tokenized.token_ids == [1, 2, 3] + assert tokenized.logprobs[1:] == pytest.approx([-0.2, -0.3]) + assert tokenized.flags[1:] == [ + art.TokenFlag.EXACT | art.TokenFlag.SAMPLED, + art.TokenFlag.EXACT | art.TokenFlag.SAMPLED, + ] + + +def test_responses_generation_evidence_is_atomic_and_partial_edits_do_not_replay() -> ( + None +): + exchange = _response_exchange("reasoning-and-answer", 0) + data = exchange.response.model_dump(mode="python") + data["output"] = [ + { + "id": "reasoning", + "type": "reasoning", + "summary": [{"type": "summary_text", "text": "think"}], + }, + { + "id": "message", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [ + { + "type": "output_text", + "text": "answer", + "annotations": [], + "logprobs": [], + } + ], + }, + ] + data["token_generations"] = [ + { + "prompt_token_ids": [1], + "output_tokens": [ + {"token_id": 2, "logprob": -0.2, "text": "think"}, + {"token_id": 3, "logprob": -0.3, "text": "answer"}, + ], + "output_indices": [0, 1], + } + ] + exchange.response = Response.model_validate(data) + + class Tokenizer: + def __call__(self, text: str, **kwargs: object) -> list[int]: + del kwargs + return {"turn 0": [1], "think": [20], "answer": [30]}[text] + + def apply_chat_template( + self, messages: list[dict[str, Any]], **kwargs: object + ) -> list[int]: + del kwargs + if messages[-1]["role"] != "assistant": + return [1] + assistant = messages[-1] + if assistant.get("reasoning"): + return [1, 20, 30, 9] + return [1, 30, 9] + + history = art.Trajectory( + exchanges=TrajectoryExchanges(responses=[exchange]) + ).responses_history() + exact = history.tokenize(tokenizer=Tokenizer(), chat_template="custom") + assert exact.token_ids == [1, 2, 3] + assert exact.logprobs[1:] == pytest.approx([-0.2, -0.3]) + + del history.input[1] + del history.input_sources[1] + partial = history.tokenize(tokenizer=Tokenizer(), chat_template="custom") + assert partial.token_ids == [1, 30, 9] + assert 2 not in partial.token_ids + assert partial.flags[1] == art.TokenFlag.SAMPLED + assert math.isnan(partial.logprobs[1]) + + def test_mutable_chat_history_is_authoritative_and_does_not_replay_removed_turns() -> ( None ): @@ -2771,6 +2968,88 @@ def test_tokenized_results_materialize_metadata_and_group_shape() -> None: assert "underlying" not in tokenized.model_dump() +def test_private_trace_covers_sampled_tokens_for_every_protocol() -> None: + from art.trajectories._tokenize import _tokenize_trajectory_with_trace + + message_response = Message.model_validate( + { + "id": "message-trace", + "type": "message", + "role": "assistant", + "model": "test/model", + "content": [{"type": "text", "text": "answer"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 1, "output_tokens": 1}, + "prompt_token_ids": [1], + "token_ids": [2], + "logprobs": [-0.2], + } + ) + start = datetime(2026, 1, 1) + trajectories = [ + art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[_chat_exchange([1], [2])]) + ), + art.Trajectory( + exchanges=TrajectoryExchanges(completions=[_completion_exchange()]) + ), + art.Trajectory( + exchanges=TrajectoryExchanges( + responses=[_response_exchange("response-trace", 2)] + ) + ), + art.Trajectory( + exchanges=TrajectoryExchanges( + messages=[ + MessagesExchange( + request=MessagesRequest( + model="test/model", + messages=[{"role": "user", "content": "question"}], + max_tokens=16, + ), + response=message_response, + start_time=start, + end_time=start + timedelta(milliseconds=1), + ) + ] + ) + ), + ] + + for trajectory in trajectories: + tokenized, traces = _tokenize_trajectory_with_trace(trajectory) + + assert len(tokenized.histories) == len(traces) == 1 + history = tokenized.histories[0] + trace = traces[0] + trace.validate(history) + assert sum(key is not None for key in trace.source_keys) == sum( + bool(flag & art.TokenFlag.SAMPLED) for flag in history.flags + ) + assert len(trace.sources) == 1 + + +def test_private_trace_keys_do_not_collide_for_repeated_empty_response_ids() -> None: + from art.trajectories._tokenize import _tokenize_trajectory_with_trace + + first = _chat_exchange([1], [2]) + second = _chat_exchange([1, 2, 3], [4], offset=1) + first.response.id = second.response.id = "" + second.start_time = first.start_time + second.end_time = first.end_time + trajectory = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[first, second]) + ) + + tokenized, [trace] = _tokenize_trajectory_with_trace(trajectory) + sampled_keys = [key for key in trace.source_keys if key is not None] + + assert tokenized.histories[0].token_ids == [1, 2, 3, 4] + assert len(set(sampled_keys)) == 2 + assert len(trace.sources) == 2 + + def test_completions_history_requires_exhaustive_source_spans() -> None: history = art.CompletionsTokenHistory( model="test/model", From 2ed95a43793a5ac88b99a7681a37446db02780ac Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Fri, 24 Jul 2026 19:15:24 +0000 Subject: [PATCH 15/58] Handle ambiguous echoed completion evidence --- src/art/trajectories/_tokenize.py | 8 +++++- tests/unit/trajectories/test_tokenize.py | 34 ++++++++++++++++++++++++ 2 files changed, 41 insertions(+), 1 deletion(-) diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index 59c01e478..1c60bafdf 100644 --- a/src/art/trajectories/_tokenize.py +++ b/src/art/trajectories/_tokenize.py @@ -401,7 +401,13 @@ def _completion_evidence( selected = token_ids if token_ids is not None else pair_ids or None prompt_logprobs: list[float] = [] completion_logprobs = pair_logprobs - if echo and prompt_ids is not None and selected is not None: + if echo and prompt_ids is None and token_ids is None: + # A combined logprobs.tokens carrier does not reveal where an echoed + # prompt ends. Let the string history tokenize prompt and completion + # independently rather than mislabeling the prompt as sampled output. + selected = None + completion_logprobs = [] + elif echo and prompt_ids is not None and selected is not None: pair_includes_prompt = token_ids is not None and pair_ids == [ *prompt_ids, *token_ids, diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index 487cf7cd6..f13ee8789 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -362,6 +362,40 @@ def test_completions_echo_prefers_full_logprob_carrier_to_id_prefix_heuristic() ] +def test_completions_echo_without_prompt_ids_falls_back_without_sampling_prompt() -> ( + None +): + exchange = _completion_exchange(echo=True) + payload = exchange.response.model_dump(mode="python") + payload["choices"][0].pop("prompt_token_ids") + payload["choices"][0].pop("token_ids") + payload["choices"][0]["logprobs"] = { + "tokens": ["token_id:1", "token_id:2"], + "token_logprobs": [-0.1, -0.2], + "top_logprobs": [{}, {}], + "text_offset": [0, 8], + } + exchange.response = Completion.model_validate(payload) + + class Tokenizer: + def __call__(self, text: str, **kwargs: object) -> list[int]: + del kwargs + return {"question": [1], "answer": [2]}[text] + + def apply_chat_template( + self, messages: list[dict[str, Any]], **kwargs: object + ) -> list[int]: + raise AssertionError((messages, kwargs)) + + tokenized = art.Trajectory( + exchanges=TrajectoryExchanges(completions=[exchange]) + ).tokenize(tokenizer=Tokenizer()) + + assert tokenized.token_ids == [1, 2] + assert tokenized.flags == [art.TokenFlag(0), art.TokenFlag.SAMPLED] + assert all(math.isnan(logprob) for logprob in tokenized.logprobs) + + def test_batched_completions_echo_uses_selected_prompt_boundary() -> None: exchange = _completion_exchange(prompt=["p0", "p1"], echo=True) payload = exchange.response.model_dump(mode="python") From 7ce39a0e75efb0ce30f63f260a46d37750b8732d Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Fri, 24 Jul 2026 19:18:03 +0000 Subject: [PATCH 16/58] Use generation-scoped Responses fixtures --- tests/unit/test_trajectory_parquet.py | 14 +++++++++++++- tests/unit/trajectories/test_capture.py | 10 +++++++++- 2 files changed, 22 insertions(+), 2 deletions(-) diff --git a/tests/unit/test_trajectory_parquet.py b/tests/unit/test_trajectory_parquet.py index 24ab9688a..a0898a14f 100644 --- a/tests/unit/test_trajectory_parquet.py +++ b/tests/unit/test_trajectory_parquet.py @@ -154,7 +154,19 @@ def _exchange_trajectory() -> Trajectory: "parallel_tool_calls": True, "tool_choice": "auto", "tools": [], - "raw_output_tokens": [{"token_id": 20, "logprob": -0.3}], + "token_generations": [ + { + "prompt_token_ids": [10], + "output_tokens": [ + { + "token_id": 20, + "logprob": -0.3, + "text": "done", + } + ], + "output_indices": [0], + } + ], }, "start_time": "2026-01-01T00:00:04", "end_time": "2026-01-01T00:00:05", diff --git a/tests/unit/trajectories/test_capture.py b/tests/unit/trajectories/test_capture.py index f7221adcc..0e2bda5e5 100644 --- a/tests/unit/trajectories/test_capture.py +++ b/tests/unit/trajectories/test_capture.py @@ -113,7 +113,15 @@ "output_tokens_details": {"reasoning_tokens": 0}, "total_tokens": 5, }, - "raw_output_tokens": [{"token_id": 2, "logprob": -0.2}], + "token_generations": [ + { + "prompt_token_ids": [1], + "output_tokens": [ + {"token_id": 2, "logprob": -0.2, "text": "hello"} + ], + "output_indices": [0], + } + ], } MESSAGE: dict[str, Any] = { From 1693a2fa1fcfc19a3846f3b1505c7abd62ca05d5 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Fri, 24 Jul 2026 19:27:05 +0000 Subject: [PATCH 17/58] Harden projected history tokenization --- src/art/trajectories/_tokenize.py | 157 ++++++++++++-- tests/unit/trajectories/test_tokenize.py | 252 +++++++++++++++++++++-- 2 files changed, 370 insertions(+), 39 deletions(-) diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index 1c60bafdf..b9e77b038 100644 --- a/src/art/trajectories/_tokenize.py +++ b/src/art/trajectories/_tokenize.py @@ -5,6 +5,8 @@ from dataclasses import dataclass from datetime import datetime from functools import lru_cache +from hashlib import sha256 +import json import math import re from typing import TYPE_CHECKING, Any, Literal, cast @@ -75,7 +77,7 @@ class _SampledSourceKey: end_time: datetime index: int prompt_index: int | None - exchange_identity: int + evidence_fingerprint: str @dataclass @@ -112,6 +114,72 @@ def set( self.trace = trace +def _fingerprint(value: object) -> str: + serialized = json.dumps( + value, + ensure_ascii=False, + separators=(",", ":"), + sort_keys=True, + ) + return sha256(serialized.encode()).hexdigest() + + +def _sampled_evidence_fingerprint( + exchange: Exchange, + *, + protocol: Literal["chat_completions", "responses", "messages", "completions"], + index: int, +) -> str: + if protocol == "chat_completions": + if not isinstance(exchange, ChatCompletionsExchange): + raise TypeError("Chat source has the wrong exchange type") + evidence = next( + choice.model_dump(mode="json") + for choice in exchange.response.choices + if choice.index == index + ) + elif protocol == "completions": + if not isinstance(exchange, CompletionsExchange): + raise TypeError("Completions source has the wrong exchange type") + evidence = next( + choice.model_dump(mode="json") + for choice in exchange.response.choices + if choice.index == index + ) + elif protocol == "responses": + if not isinstance(exchange, ResponsesExchange): + raise TypeError("Responses source has the wrong exchange type") + response = exchange.response.model_dump(mode="json") + generations = response.get("token_generations") + if isinstance(generations, list) and 0 <= index < len(generations): + generation = generations[index] + output_indices = ( + generation.get("output_indices") + if isinstance(generation, dict) + else None + ) + outputs = response.get("output") + evidence = { + "generation": generation, + "outputs": [ + outputs[output_index] + for output_index in output_indices + if isinstance(output_index, int) + and isinstance(outputs, list) + and 0 <= output_index < len(outputs) + ] + if isinstance(output_indices, list) + else [], + } + else: + evidence = response.get("output") + else: + if not isinstance(exchange, MessagesExchange): + raise TypeError("Messages source has the wrong exchange type") + evidence = exchange.response.model_dump(mode="json") + return _fingerprint(evidence) + + def _source_key( exchange: Exchange, *, @@ -126,10 +194,12 @@ def _source_key( end_time=exchange.end_time, index=index, prompt_index=prompt_index, - # The trace is private and consumed during this tokenization call. Object - # identity makes repeated empty provider IDs collision-free while keeping - # all histories projected from one exchange on the same source key. - exchange_identity=id(exchange), + # Internal projection may copy an exchange to isolate one choice. A + # source-specific identity remains stable across those copies without + # hashing a growing request or unrelated choices. + evidence_fingerprint=_sampled_evidence_fingerprint( + exchange, protocol=protocol, index=index + ), ) @@ -2295,16 +2365,35 @@ def _source_signature(source: object) -> tuple[object, ...] | None: ): return None response_id = getattr(exchange.response, "id", None) + choice_index = getattr(source, "choice_index", None) + generation_index = getattr(source, "generation_index", None) + evidence_fingerprint: str | None = None + if isinstance(exchange, ChatCompletionsExchange) and isinstance(choice_index, int): + evidence_fingerprint = _sampled_evidence_fingerprint( + exchange, protocol="chat_completions", index=choice_index + ) + elif isinstance(exchange, ResponsesExchange) and isinstance(generation_index, int): + evidence_fingerprint = _sampled_evidence_fingerprint( + exchange, protocol="responses", index=generation_index + ) + elif ( + isinstance(exchange, MessagesExchange) + and getattr(source, "request_index", None) is None + ): + evidence_fingerprint = _sampled_evidence_fingerprint( + exchange, protocol="messages", index=0 + ) return ( type(source), type(exchange), exchange.start_time, exchange.end_time, response_id, + evidence_fingerprint, getattr(source, "request_index", None), - getattr(source, "choice_index", None), + choice_index, getattr(source, "output_index", None), - getattr(source, "generation_index", None), + generation_index, getattr(source, "prompt_index", None), ) @@ -2348,7 +2437,22 @@ def _history_matches_projection(history: History) -> bool: except ValueError as error: if "no Chat Completions exchanges" not in str(error): raise - return False + if trajectory.exchanges.messages and not trajectory.exchanges.responses: + candidates = [ + candidate.as_chat_completions_history() + for candidate in anthropic_messages_histories( + trajectory, model=history.model + ) + ] + elif trajectory.exchanges.responses and not trajectory.exchanges.messages: + candidates = [ + candidate.as_chat_completions_history() + for candidate in responses_histories( + trajectory, model=history.model + ) + ] + else: + return False return any( isinstance(candidate, ChatCompletionsHistory) and candidate.messages == history.messages @@ -3593,6 +3697,7 @@ def tokenize_history( chat_template: str | None, chat_template_kwargs: Mapping[str, object] | None, _trace: _TraceBuilder | None = None, + _projection_validated: bool = False, ) -> TokenizedHistory: if isinstance(history, LegacyHistory): if model is None: @@ -3618,31 +3723,42 @@ def tokenize_history( and dict(chat_template_kwargs) != (getattr(history, "chat_template_kwargs", None) or {}) ) - render_state = _history_render_state(history) + render_state = ( + _HistoryRenderState(needs_render=False, projection_matches=True) + if _projection_validated + else _history_render_state(history) + ) needs_render = render_state.needs_render or override_requires_render if isinstance(history, ResponsesHistory) and not needs_render: if exact := _tokenize_exact_responses_history( history, base_model=base_model, tokenizer=tokenizer, _trace=_trace ): return exact - if isinstance(history, ChatCompletionsHistory) and (needs_render): + if isinstance(history, ChatCompletionsHistory): if ( not override_requires_render and not render_state.context_changed and render_state.projection_matches is not False and ( - exact := _tokenize_exact_projected_chat_history(history, _trace=_trace) + exact := _tokenize_exact_projected_chat_history( + history, + projection_validated=( + _projection_validated or render_state.projection_matches is True + ), + _trace=_trace, + ) ) ): return exact - return _tokenize_chat_view( - history, - base_model=base_model, - tokenizer=tokenizer, - chat_template=chat_template, - chat_template_kwargs=chat_template_kwargs, - _trace=_trace, - ) + if needs_render: + return _tokenize_chat_view( + history, + base_model=base_model, + tokenizer=tokenizer, + chat_template=chat_template, + chat_template_kwargs=chat_template_kwargs, + _trace=_trace, + ) if isinstance(history, AnthropicMessagesHistory) and needs_render: converted = history.as_chat_completions_history() if ( @@ -3745,6 +3861,7 @@ def tokenize_trajectory( tokenizer=tokenizer, chat_template=chat_template, chat_template_kwargs=chat_template_kwargs, + _projection_validated=not isinstance(histories[0], LegacyHistory), ), trajectory, ) @@ -3756,6 +3873,7 @@ def tokenize_trajectory( tokenizer=tokenizer, chat_template=chat_template, chat_template_kwargs=chat_template_kwargs, + _projection_validated=not isinstance(history, LegacyHistory), ) for history in histories ] @@ -3798,6 +3916,7 @@ def _tokenize_trajectory_with_trace( chat_template=chat_template, chat_template_kwargs=chat_template_kwargs, _trace=trace_builder, + _projection_validated=True, ) if trace_builder.trace is None: raise AssertionError("Exchange tokenization did not produce a source trace") diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index f13ee8789..7d15c5b13 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -4,7 +4,9 @@ from datetime import datetime, timedelta import math import random +from statistics import median import sys +from time import perf_counter from types import ModuleType, SimpleNamespace from typing import Any @@ -228,6 +230,78 @@ def test_messages_exact_prompt_and_output_do_not_load_a_tokenizer( ] +def test_converted_anthropic_system_history_preserves_exact_assistant_evidence( + monkeypatch: pytest.MonkeyPatch, +) -> None: + response = Message.model_validate( + { + "id": "message-system-exact", + "type": "message", + "role": "assistant", + "model": "test/model", + "content": [{"type": "text", "text": "answer"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 2, "output_tokens": 1}, + "prompt_token_ids": [10, 11], + "token_ids": [12], + "logprobs": [-0.12], + } + ) + start = datetime(2026, 1, 1) + exchange = MessagesExchange( + request=MessagesRequest( + model="test/model", + system="system", + messages=[{"role": "user", "content": "question"}], + max_tokens=16, + ), + response=response, + start_time=start, + end_time=start + timedelta(milliseconds=1), + ) + trajectory = art.Trajectory(exchanges=TrajectoryExchanges(messages=[exchange])) + monkeypatch.setattr( + "art.trajectories._tokenize._load_tokenizer", + lambda _config: pytest.fail("exact converted history loaded a tokenizer"), + ) + + history = trajectory.anthropic_messages_history().as_chat_completions_history() + assert all( + source is None or source.exchange is exchange + for source in history.message_sources + ) + tokenized = history.tokenize() + + assert tokenized.token_ids == [10, 11, 12] + assert tokenized.logprobs[-1] == pytest.approx(-0.12) + assert tokenized.flags[-1] == art.TokenFlag.EXACT | art.TokenFlag.SAMPLED + + +def test_converted_responses_history_tokenizes_without_native_chat_exchange( + monkeypatch: pytest.MonkeyPatch, +) -> None: + exchange = _response_exchange( + "response-converted-exact", 12, prompt_token_ids=[10, 11] + ) + trajectory = art.Trajectory(exchanges=TrajectoryExchanges(responses=[exchange])) + monkeypatch.setattr( + "art.trajectories._tokenize._load_tokenizer", + lambda _config: pytest.fail("exact converted history loaded a tokenizer"), + ) + + history = trajectory.responses_history().as_chat_completions_history() + assert all( + source is None or source.exchange is exchange + for source in history.message_sources + ) + tokenized = history.tokenize() + + assert tokenized.token_ids == [10, 11, 12] + assert tokenized.logprobs[-1] == pytest.approx(-0.1) + assert tokenized.flags[-1] == art.TokenFlag.EXACT | art.TokenFlag.SAMPLED + + def test_malformed_explicit_exact_token_metadata_fails_closed() -> None: chat = _chat_exchange([1], [2]) chat_extra = chat.response.choices[0].model_extra @@ -2264,6 +2338,59 @@ def apply_chat_template( assert math.isnan(partial.logprobs[1]) +def test_responses_generation_source_rejects_content_from_another_generation() -> None: + exchange = _response_exchange("generation-provenance", 2) + data = exchange.response.model_dump(mode="python") + data["output"].append( + { + "id": "message-second-generation", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [ + { + "type": "output_text", + "text": "second", + "annotations": [], + "logprobs": [], + } + ], + } + ) + data["token_generations"] = [ + { + "prompt_token_ids": [1], + "output_tokens": [{"token_id": 2, "logprob": -0.2}], + "output_indices": [0], + }, + { + "prompt_token_ids": [1, 2, 3], + "output_tokens": [{"token_id": 4, "logprob": -0.4}], + "output_indices": [1], + }, + ] + exchange.response = Response.model_validate(data) + history = ( + art.Trajectory(exchanges=TrajectoryExchanges(responses=[exchange])) + .responses_history() + .as_chat_completions_history() + ) + first_generation = next( + index + for index, source in enumerate(history.message_sources) + if source is not None + and source.generation_index == 0 + and history.messages[index].get("role") == "assistant" + ) + history.messages[first_generation] = { + "role": "assistant", + "content": "second", + } + + with pytest.raises(ValueError, match="no longer matches its source exchange"): + history.tokenize() + + def test_mutable_chat_history_is_authoritative_and_does_not_replay_removed_turns() -> ( None ): @@ -2849,11 +2976,11 @@ def apply_chat_template( assert not tokenized.flags[0] & art.TokenFlag.SAMPLED -def test_rerender_calls_chat_template_once_for_many_turns() -> None: +def _repeated_text_rerender_history(turn_count: int) -> art.ChatCompletionsHistory: exchanges: list[ChatCompletionsExchange] = [] prompt: list[int] = [] messages: list[ChatCompletionMessageParam] = [] - for index in range(32): + for index in range(turn_count): prompt.extend([index * 2]) messages.append({"role": "user", "content": f"u{index}"}) exchange = _chat_exchange(list(prompt), [index * 2 + 1], offset=index) @@ -2865,33 +2992,118 @@ def test_rerender_calls_chat_template_once_for_many_turns() -> None: exchanges=TrajectoryExchanges(chat_completions=exchanges) ).chat_completions_history() history.chat_template = "rerender" + return history - class Tokenizer: - apply_calls = 0 - def __call__(self, text: str, **kwargs: object) -> list[int]: - del kwargs - return [1000] if text == "answer" else [2000 + int(text[1:])] +class _RepeatedTextTokenizer: + def __init__(self) -> None: + self.apply_calls = 0 - def apply_chat_template( - self, messages: list[dict[str, Any]], **kwargs: object - ) -> list[int]: - del kwargs - self.apply_calls += 1 - result: list[int] = [] - for message in messages: - content = str(message["content"]) - result.extend( - [1000] if content == "answer" else [2000 + int(content[1:])] - ) - return result + def __call__(self, text: str, **kwargs: object) -> list[int]: + del kwargs + return [1000] if text == "answer" else [2000 + int(text[1:])] - tokenizer = Tokenizer() + def apply_chat_template( + self, messages: list[dict[str, Any]], **kwargs: object + ) -> list[int]: + del kwargs + self.apply_calls += 1 + result: list[int] = [] + for message in messages: + content = str(message["content"]) + result.extend([1000] if content == "answer" else [2000 + int(content[1:])]) + return result + + +def test_rerender_calls_chat_template_once_for_many_turns() -> None: + history = _repeated_text_rerender_history(32) + + tokenizer = _RepeatedTextTokenizer() history.tokenize(tokenizer=tokenizer) assert tokenizer.apply_calls == 1 +def test_repeated_text_rerender_scaling_is_near_linear() -> None: + medians: list[float] = [] + for turn_count in (32, 64, 128): + history = _repeated_text_rerender_history(turn_count) + samples: list[float] = [] + for _ in range(5): + tokenizer = _RepeatedTextTokenizer() + started = perf_counter() + history.tokenize(tokenizer=tokenizer) + samples.append(perf_counter() - started) + assert tokenizer.apply_calls == 1 + medians.append(median(samples)) + + assert medians[1] < medians[0] * 3 + assert medians[2] < medians[1] * 3 + + +def test_reasoning_split_trajectory_reuses_prevalidated_projections( + monkeypatch: pytest.MonkeyPatch, +) -> None: + exchanges: list[ChatCompletionsExchange] = [] + request_messages: list[ChatCompletionMessageParam] = [] + prompt: list[int] = [] + for index in range(40): + request_messages.append({"role": "user", "content": f"u{index}"}) + prompt.append(3000 + index) + exchange = _chat_exchange( + list(prompt), [1000 + index, 2000 + index], offset=index + ) + exchange.request["messages"] = list(request_messages) + payload = exchange.response.model_dump(mode="python") + payload["choices"][0]["message"] = { + "role": "assistant", + "reasoning": f"r{index}", + "content": f"a{index}", + } + exchange.response = ChatCompletion.model_validate(payload) + exchanges.append(exchange) + request_messages.append({"role": "assistant", "content": f"a{index}"}) + prompt.append(2000 + index) + trajectory = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=exchanges) + ) + projected = trajectory.chat_completions_histories() + assert len(projected) == 40 + for history in projected: + for source in history.message_sources: + if source is not None: + assert any(source.exchange is exchange for exchange in exchanges) + assert all( + any( + source is not None + and source.choice_index == 0 + and source.exchange is exchanges[0] + for source in history.message_sources + ) + for history in projected + ) + + from art.trajectories import _tokenize + + original = _tokenize._history_matches_projection + calls = 0 + + def counted(history: art.History) -> bool: + nonlocal calls + calls += 1 + return original(history) + + monkeypatch.setattr(_tokenize, "_history_matches_projection", counted) + + tokenized = trajectory.tokenize(multi_history=True) + _, traces = _tokenize._tokenize_trajectory_with_trace(trajectory) + first_key = next(key for key in traces[0].source_keys if key is not None) + + assert len(tokenized.histories) == 40 + assert all(first_key in trace.sources for trace in traces) + assert calls == 0 + + def test_explicit_template_override_rerenders_exact_exchange_scaffold() -> None: trajectory = art.Trajectory( exchanges=TrajectoryExchanges(chat_completions=[_chat_exchange([1], [2])]) From eb785adfdb6ba73cf596c7c8c9c7d88dac2b49c2 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Fri, 24 Jul 2026 19:29:58 +0000 Subject: [PATCH 18/58] Exclude routing extras from source fingerprints --- src/art/trajectories/_tokenize.py | 111 +++++++++++++++++++++--------- 1 file changed, 80 insertions(+), 31 deletions(-) diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index b9e77b038..c9a20efcd 100644 --- a/src/art/trajectories/_tokenize.py +++ b/src/art/trajectories/_tokenize.py @@ -124,6 +124,29 @@ def _fingerprint(value: object) -> str: return sha256(serialized.encode()).hexdigest() +def _chat_logprob_fingerprint_evidence(choice: Choice) -> dict[str, object] | None: + if choice.logprobs is None: + return None + + def values(items: Sequence[object] | None) -> list[dict[str, object]]: + result: list[dict[str, object]] = [] + for item in items or []: + data = _dump(item) + result.append( + { + key: data[key] + for key in ("token", "token_id", "logprob", "bytes") + if key in data + } + ) + return result + + return { + "content": values(choice.logprobs.content), + "refusal": values(choice.logprobs.refusal), + } + + def _sampled_evidence_fingerprint( exchange: Exchange, *, @@ -133,50 +156,76 @@ def _sampled_evidence_fingerprint( if protocol == "chat_completions": if not isinstance(exchange, ChatCompletionsExchange): raise TypeError("Chat source has the wrong exchange type") - evidence = next( - choice.model_dump(mode="json") - for choice in exchange.response.choices - if choice.index == index - ) + choice = next( + choice for choice in exchange.response.choices if choice.index == index + ) + choice_extra = choice.model_extra or {} + evidence = { + "message": choice.message.model_dump( + mode="json", + include={ + "role", + "content", + "refusal", + "reasoning", + "reasoning_content", + "tool_calls", + "function_call", + "audio", + }, + exclude_none=True, + ), + "token_ids": choice_extra.get("token_ids"), + "logprobs": _chat_logprob_fingerprint_evidence(choice), + "finish_reason": choice.finish_reason, + } elif protocol == "completions": if not isinstance(exchange, CompletionsExchange): raise TypeError("Completions source has the wrong exchange type") - evidence = next( - choice.model_dump(mode="json") - for choice in exchange.response.choices - if choice.index == index - ) + choice = next( + choice for choice in exchange.response.choices if choice.index == index + ) + choice_extra = choice.model_extra or {} + logprobs = _dump(choice.logprobs) + evidence = { + "text": choice.text, + "token_ids": choice_extra.get("token_ids"), + "logprobs": { + key: logprobs[key] + for key in ("tokens", "token_logprobs") + if key in logprobs + }, + "finish_reason": choice.finish_reason, + } elif protocol == "responses": if not isinstance(exchange, ResponsesExchange): raise TypeError("Responses source has the wrong exchange type") - response = exchange.response.model_dump(mode="json") - generations = response.get("token_generations") + generations = (exchange.response.model_extra or {}).get("token_generations") if isinstance(generations, list) and 0 <= index < len(generations): - generation = generations[index] - output_indices = ( - generation.get("output_indices") - if isinstance(generation, dict) - else None - ) - outputs = response.get("output") + generation = _string_dict(generations[index]) or {} evidence = { - "generation": generation, - "outputs": [ - outputs[output_index] - for output_index in output_indices - if isinstance(output_index, int) - and isinstance(outputs, list) - and 0 <= output_index < len(outputs) - ] - if isinstance(output_indices, list) - else [], + key: generation[key] + for key in ("output_tokens", "output_indices") + if key in generation } else: - evidence = response.get("output") + evidence = [ + output.model_dump(mode="json", exclude_none=True) + for output in exchange.response.output + ] else: if not isinstance(exchange, MessagesExchange): raise TypeError("Messages source has the wrong exchange type") - evidence = exchange.response.model_dump(mode="json") + response_extra = exchange.response.model_extra or {} + evidence = { + "content": [ + block.model_dump(mode="json", exclude_none=True) + for block in exchange.response.content + ], + "token_ids": response_extra.get("token_ids"), + "logprobs": response_extra.get("logprobs"), + "stop_reason": exchange.response.stop_reason, + } return _fingerprint(evidence) From 3a70b1ab723f260b83bc53bd601e69f5f151c72a Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Fri, 24 Jul 2026 19:20:33 +0000 Subject: [PATCH 19/58] Prefer exact public model selectors --- src/art/trajectories/_history.py | 29 +++++++++++++++++++++-------- src/art/trajectories/_tokenize.py | 9 ++++++++- 2 files changed, 29 insertions(+), 9 deletions(-) diff --git a/src/art/trajectories/_history.py b/src/art/trajectories/_history.py index a1aaeb5db..3765e3f16 100644 --- a/src/art/trajectories/_history.py +++ b/src/art/trajectories/_history.py @@ -1409,15 +1409,30 @@ def trajectory_histories( *copy.deepcopy(trajectory.additional_histories), ] + all_exchanges = [ + *trajectory.exchanges.chat_completions, + *trajectory.exchanges.completions, + *trajectory.exchanges.responses, + *trajectory.exchanges.messages, + ] + has_exact_match = model is not None and any( + exchange.model == model for exchange in all_exchanges + ) + + def is_selected(exchange: _ModelledExchange) -> bool: + if model is None: + return True + if has_exact_match: + return exchange.model == model + return _model_matches(exchange.model, model) + candidates: list[TrajectoryHistory] = [] if any( - model is None or _model_matches(exchange.model, model) - for exchange in trajectory.exchanges.chat_completions + is_selected(exchange) for exchange in trajectory.exchanges.chat_completions ): candidates.extend(chat_completions_histories(trajectory, model=model)) if any( - model is None or _model_matches(exchange.model, model) - for exchange in trajectory.exchanges.completions + is_selected(exchange) for exchange in trajectory.exchanges.completions ): try: candidates.extend(completions_token_histories(trajectory, model=model)) @@ -1426,13 +1441,11 @@ def trajectory_histories( raise candidates.extend(completions_string_histories(trajectory, model=model)) if any( - model is None or _model_matches(exchange.model, model) - for exchange in trajectory.exchanges.responses + is_selected(exchange) for exchange in trajectory.exchanges.responses ): candidates.extend(responses_histories(trajectory, model=model)) if any( - model is None or _model_matches(exchange.model, model) - for exchange in trajectory.exchanges.messages + is_selected(exchange) for exchange in trajectory.exchanges.messages ): candidates.extend(anthropic_messages_histories(trajectory, model=model)) if not candidates: diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index c9a20efcd..817d45c26 100644 --- a/src/art/trajectories/_tokenize.py +++ b/src/art/trajectories/_tokenize.py @@ -788,8 +788,15 @@ def _exchange_list(trajectory: Trajectory, model: str | None) -> list[Exchange]: *trajectory.exchanges.messages, ] if model is not None: + has_exact_match = any(exchange.model == model for exchange in exchanges) exchanges = [ - exchange for exchange in exchanges if _model_matches(exchange.model, model) + exchange + for exchange in exchanges + if ( + exchange.model == model + if has_exact_match + else _model_matches(exchange.model, model) + ) ] if not exchanges: raise ValueError(f"Trajectory contains no exchanges for model {model!r}") From 97ce26a7c45ac765ac8f35734c0791377a76043d Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Fri, 24 Jul 2026 19:32:33 +0000 Subject: [PATCH 20/58] Train projected exchange sources once --- src/art/preprocessing/tokenize.py | 130 ++++++++++++------ src/art/tinker_native/data.py | 50 +++++-- src/art/trajectories/_tokenize.py | 10 ++ .../test_exchange_training_model_selection.py | 114 ++++++++++++++- 4 files changed, 248 insertions(+), 56 deletions(-) diff --git a/src/art/preprocessing/tokenize.py b/src/art/preprocessing/tokenize.py index 64343217d..e98a2161a 100644 --- a/src/art/preprocessing/tokenize.py +++ b/src/art/preprocessing/tokenize.py @@ -1,6 +1,6 @@ from __future__ import annotations -from collections.abc import Callable +from collections.abc import Callable, Iterable from dataclasses import dataclass, field from functools import cached_property from itertools import takewhile @@ -21,6 +21,7 @@ from ..trajectories import ( ChatCompletionsExchange, ChatCompletionsHistory, + ChatCompletionsMessageSource, LegacyHistory, TokenFlag, Trajectory, @@ -65,6 +66,13 @@ def _flag_spans(flags: list[TokenFlag], flag: TokenFlag) -> list[tuple[int, int] return spans +def _true_spans(mask: list[bool]) -> list[tuple[int, int]]: + return _flag_spans( + [TokenFlag.SAMPLED if value else TokenFlag(0) for value in mask], + TokenFlag.SAMPLED, + ) + + @dataclass(frozen=True) class _ChatChoiceTrace: choices: list[Choice] @@ -76,27 +84,42 @@ def _chat_choice_trace( history: ChatCompletionsHistory, token_ids: list[int], flags: list[TokenFlag], +) -> _ChatChoiceTrace | None: + return _chat_source_choice_trace( + ( + (source, source.choice_index) + for source in history.message_sources + if source is not None + ), + token_ids, + flags, + ) + + +def _chat_source_choice_trace( + sources: Iterable[tuple[object, int | None]], + token_ids: list[int], + flags: list[TokenFlag], ) -> _ChatChoiceTrace | None: """Recover per-choice boundaries that adjacent SAMPLED flags cannot represent.""" sourced_choices: list[Choice] = [] seen: set[tuple[int, int]] = set() - for source in history.message_sources: - if ( - source is None - or source.choice_index is None - or not isinstance(source.exchange, ChatCompletionsExchange) - ): + for source, choice_index in sources: + exchange = ( + source.exchange + if isinstance(source, ChatCompletionsMessageSource) + else source + ) + if not isinstance(exchange, ChatCompletionsExchange) or choice_index is None: continue - key = (id(source.exchange), source.choice_index) + key = (id(exchange), choice_index) if key in seen: continue seen.add(key) sourced_choices.append( next( - item - for item in source.exchange.response.choices - if item.index == source.choice_index + item for item in exchange.response.choices if item.index == choice_index ) ) if not sourced_choices: @@ -714,30 +737,37 @@ def tokenize_trajectory_groups( if advantage == 0 and drop_zero_advantage_trajectories: continue if trajectory.exchanges: - from ..trajectories._tokenize import _as_tokenizer + from ..trajectories._tokenize import ( + _as_tokenizer, + _first_introduction_mask, + _SampledSourceKey, + _tokenize_trajectory_with_trace, + ) selected_model = resolve_training_model(trajectory, model) - exchange_results = trajectory.tokenize( - multi_history=True, + exchange_results, traces = _tokenize_trajectory_with_trace( + trajectory, model=selected_model, base_model=tokenizer.name_or_path, tokenizer=_as_tokenizer(tokenizer), chat_template=None, chat_template_kwargs=chat_template_kwargs, ) - histories = trajectory.histories(model=selected_model) trajectory_results = [] - for exchange_result, history in zip( - exchange_results.histories, histories, strict=True + seen_source_keys: set[_SampledSourceKey] = set() + for exchange_result, trace in zip( + exchange_results.histories, traces, strict=True ): - sampled = [ - bool(flag & TokenFlag.SAMPLED) for flag in exchange_result.flags - ] - choice_spans = _flag_spans(exchange_result.flags, TokenFlag.SAMPLED) + trainable = _first_introduction_mask( + trace.source_keys, seen_source_keys + ) + if not any(trainable): + continue + choice_spans = _true_spans(trainable) if not allow_training_without_logprobs and any( - trainable and math.isnan(logprob) - for trainable, logprob in zip( - sampled, + selected and math.isnan(logprob) + for selected, logprob in zip( + trainable, exchange_result.logprobs, strict=True, ) @@ -745,27 +775,47 @@ def tokenize_trajectory_groups( raise RuntimeError( "Exchange trajectory is missing logprobs for trainable tokens" ) - chat_trace = None - if isinstance(history, ChatCompletionsHistory): - chat_trace = _chat_choice_trace( - history, - exchange_result.token_ids, - exchange_result.flags, - ) + ordered_source_keys = dict.fromkeys( + key for key in trace.source_keys if key is not None + ) + chat_trace = _chat_source_choice_trace( + ( + (trace.sources[key], key.index) + for key in ordered_source_keys + ), + exchange_result.token_ids, + exchange_result.flags, + ) if chat_trace is not None: + selected_choices = [ + index + for index, (start, length) in enumerate( + zip( + chat_trace.offsets, + chat_trace.lengths, + strict=True, + ) + ) + if any(trainable[start : start + length]) + ] moe_routes, moe_stats = align_choice_routes_to_tokenized_result( token_ids=exchange_result.token_ids, - choices=chat_trace.choices, - choice_offsets=chat_trace.offsets, - choice_token_lengths=chat_trace.lengths, + choices=[ + chat_trace.choices[index] for index in selected_choices + ], + choice_offsets=[ + chat_trace.offsets[index] for index in selected_choices + ], + choice_token_lengths=[ + chat_trace.lengths[index] for index in selected_choices + ], ) choice_spans = [ - (start, start + length) - for start, length in zip( - chat_trace.offsets, - chat_trace.lengths, - strict=True, + ( + chat_trace.offsets[index], + chat_trace.offsets[index] + chat_trace.lengths[index], ) + for index in selected_choices ] else: moe_routes = None @@ -776,7 +826,7 @@ def tokenize_trajectory_groups( chat="", token_ids=exchange_result.token_ids, input_pos=list(range(len(exchange_result.token_ids))), - assistant_mask=[int(value) for value in sampled], + assistant_mask=[int(value) for value in trainable], logprobs=exchange_result.logprobs, pixel_values=None, image_grid_thw=None, diff --git a/src/art/tinker_native/data.py b/src/art/tinker_native/data.py index d500e0fd9..0bc205452 100644 --- a/src/art/tinker_native/data.py +++ b/src/art/tinker_native/data.py @@ -164,14 +164,30 @@ def trajectory_groups_to_datums( continue for trajectory, advantage in zip(group.trajectories, advantages): if trajectory.exchanges: + from ..trajectories._tokenize import ( + _as_tokenizer, + _first_introduction_mask, + _SampledSourceKey, + _tokenize_trajectory_with_trace, + ) + selected_model = resolve_training_model(trajectory, model) - tokenized = trajectory.tokenize( - multi_history=True, + tokenized, traces = _tokenize_trajectory_with_trace( + trajectory, model=selected_model, base_model=base_model, + tokenizer=_as_tokenizer(tokenizer) + if tokenizer is not None + else None, ) - for history in tokenized.histories: - datum = _tokenized_trajectory_to_datum(history, advantage) + seen_source_keys: set[_SampledSourceKey] = set() + for history, trace in zip(tokenized.histories, traces, strict=True): + trainable = _first_introduction_mask( + trace.source_keys, seen_source_keys + ) + datum = _tokenized_trajectory_to_datum( + history, advantage, trainable=trainable + ) if datum is not None: datums.append(datum) continue @@ -289,17 +305,31 @@ def build_datum( def _tokenized_trajectory_to_datum( - tokenized: TokenizedHistory, advantage: float + tokenized: TokenizedHistory, + advantage: float, + *, + trainable: list[bool] | None = None, ) -> tinker.Datum | None: - sampled = [bool(flag & TokenFlag.SAMPLED) for flag in tokenized.flags] - if not (len(tokenized.token_ids) == len(tokenized.logprobs) == len(sampled)): + if trainable is None: + trainable = [bool(flag & TokenFlag.SAMPLED) for flag in tokenized.flags] + if not ( + len(tokenized.token_ids) + == len(tokenized.logprobs) + == len(tokenized.flags) + == len(trainable) + ): raise ValueError("Tokenized trajectory fields differ in length") - if len(tokenized.token_ids) < 2 or not any(sampled): + if any( + selected and not flag & TokenFlag.SAMPLED + for selected, flag in zip(trainable, tokenized.flags, strict=True) + ): + raise ValueError("Only sampled tokens can be selected for Tinker training") + if len(tokenized.token_ids) < 2 or not any(trainable): return None - if sampled[0]: + if trainable[0]: raise ValueError("A trainable trajectory cannot start with a sampled token") - action_mask = sampled[1:] + action_mask = trainable[1:] if any( trainable and math.isnan(logprob) for trainable, logprob in zip(action_mask, tokenized.logprobs[1:], strict=True) diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index 817d45c26..822284a86 100644 --- a/src/art/trajectories/_tokenize.py +++ b/src/art/trajectories/_tokenize.py @@ -99,6 +99,16 @@ def validate(self, tokenized: TokenizedHistory) -> None: raise AssertionError("Tokenization trace source key is unresolved") +def _first_introduction_mask( + source_keys: Sequence[_SampledSourceKey | None], + seen: set[_SampledSourceKey], +) -> list[bool]: + keys = {key for key in source_keys if key is not None} + new_keys = keys - seen + seen.update(keys) + return [key in new_keys if key is not None else False for key in source_keys] + + @dataclass class _TraceBuilder: trace: _HistoryTokenizationTrace | None = None diff --git a/tests/unit/test_exchange_training_model_selection.py b/tests/unit/test_exchange_training_model_selection.py index d39638eb0..6267435e6 100644 --- a/tests/unit/test_exchange_training_model_selection.py +++ b/tests/unit/test_exchange_training_model_selection.py @@ -11,7 +11,11 @@ import art from art.openai import ART_MOE_ROUTING_METADATA_KEY -from art.preprocessing.tokenize import _chat_choice_trace, tokenize_trajectory_groups +from art.preprocessing.moe_routing import MoeRouteSegments +from art.preprocessing.tokenize import ( + _chat_choice_trace, + tokenize_trajectory_groups, +) from art.tinker_native.data import trajectory_groups_to_datums from art.trajectories import ( ChatCompletionsExchange, @@ -21,10 +25,12 @@ CompletionsExchange, CompletionsRequest, ) +from art.trajectories import _tokenize as trajectory_tokenization from art.trajectories._selection import ( automatic_training_model_selector, resolve_training_model, ) +from art.trajectories._tokenize import _first_introduction_mask def _exchange(model: str, output_token: int) -> ChatCompletionsExchange: @@ -130,6 +136,36 @@ class _Tokenizer: name_or_path = "base/model" +def test_first_introduction_mask_trains_repeated_sources_once_across_histories() -> ( + None +): + seen: set[str] = set() + assert _first_introduction_mask([None, "a", "a"], seen) == [ + False, + True, + True, + ] + assert _first_introduction_mask([None, "a", "a", None, "b"], seen) == [ + False, + False, + False, + False, + True, + ] + assert _first_introduction_mask( + [None, "a", "a", None, "b", None, "c", "c"], seen + ) == [ + False, + False, + False, + False, + False, + False, + True, + True, + ] + + def test_preprocessing_requires_model_selection() -> None: tokenizer = cast(PreTrainedTokenizerBase, _Tokenizer()) with pytest.raises(ValueError, match="exactly one concrete model"): @@ -169,6 +205,42 @@ def test_tinker_requires_model_selection() -> None: assert all(datum.model_input.to_ints() == [1] for datum in datums) +def test_training_tokenizes_each_exchange_trajectory_once( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls = 0 + original = trajectory_tokenization._tokenize_trajectory_with_trace + + def counted(*args: object, **kwargs: object) -> object: + nonlocal calls + calls += 1 + return original(*args, **kwargs) + + monkeypatch.setattr( + trajectory_tokenization, "_tokenize_trajectory_with_trace", counted + ) + group = _group() + list( + tokenize_trajectory_groups( + cast(PreTrainedTokenizerBase, _Tokenizer()), + [group], + allow_training_without_logprobs=False, + scale_rewards=False, + model="policy", + ) + ) + assert calls == len(group.trajectories) + + calls = 0 + trajectory_groups_to_datums( + [group], + renderer=None, + tokenizer=None, + model="policy", + ) + assert calls == len(group.trajectories) + + def test_training_rejects_multiple_concrete_policy_versions() -> None: group = _versioned_group() with pytest.raises(ValueError, match="exactly one concrete model"): @@ -246,12 +318,21 @@ def test_public_training_selector_prefers_exact_model_over_glob_interpretation() trajectory = art.Trajectory( exchanges=art.TrajectoryExchanges( chat_completions=[ - _exchange("policy[*]", 2), + _exchange("policy*", 2), _exchange("policyx", 3), ] ) ) - assert resolve_training_model(trajectory, "policy[*]") == "policy[*]" + assert resolve_training_model(trajectory, "policy*") == "policy*" + assert [history.model for history in trajectory.histories(model="policy*")] == [ + "policy*" + ] + assert [ + history.model + for history in trajectory.tokenize( + model="policy*", multi_history=True + ).histories + ] == ["policy*"] def test_training_selector_rejects_zero_matches() -> None: @@ -455,16 +536,37 @@ def apply_chat_template( ) ) + initial = [result for result in results if result.token_ids[1] == 2] stripped = [result for result in results if result.token_ids[1] == 101] + assert len(initial) == 2 assert len(stripped) == 2 - assert all(result.choice_offsets == [1, 5] for result in stripped) + assert all(result.choice_offsets == [1] for result in initial) + assert all(result.choice_offsets == [5] for result in stripped) + assert all(result.assistant_mask == [0, 1, 1, 1, 1] for result in initial) + assert all(result.assistant_mask == [0, 0, 0, 0, 0, 1, 1] for result in stripped) + assert all(result.weight == pytest.approx(1 / 6) for result in results) expected_routes = np.asarray( [[[10]], [[1010]], [[1020]], [[90]], [[40]], [[50]], [[60]]], dtype=np.int32, ) for result in stripped: - assert isinstance(result.moe_routed_experts, np.ndarray) - assert np.array_equal(result.moe_routed_experts, expected_routes) + assert isinstance(result.moe_routed_experts, MoeRouteSegments) + assert np.array_equal( + np.concatenate(result.moe_routed_experts.segments), + expected_routes, + ) + + datums = trajectory_groups_to_datums( + [group], + renderer=None, + tokenizer=Tokenizer(), + normalize_advantages=False, + base_model="base/model", + model="policy", + ) + masks = [datum.loss_fn_inputs["mask"].to_torch().tolist() for datum in datums] + assert masks.count([1, 1, 1, 1]) == 2 + assert masks.count([0, 0, 0, 0, 1, 1]) == 2 def test_ambiguous_non_moe_suffix_falls_back_to_sampled_spans() -> None: From 64daac2740d0d00e94b30131f81f1b122905498d Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Fri, 24 Jul 2026 19:44:52 +0000 Subject: [PATCH 21/58] Type training trace regressions --- src/art/trajectories/_history.py | 16 ++------ src/art/trajectories/_tokenize.py | 11 ++++-- .../test_exchange_training_model_selection.py | 39 ++++++++++++++++--- tests/unit/trajectories/test_capture.py | 4 +- 4 files changed, 45 insertions(+), 25 deletions(-) diff --git a/src/art/trajectories/_history.py b/src/art/trajectories/_history.py index 3765e3f16..da961e3e2 100644 --- a/src/art/trajectories/_history.py +++ b/src/art/trajectories/_history.py @@ -1427,26 +1427,18 @@ def is_selected(exchange: _ModelledExchange) -> bool: return _model_matches(exchange.model, model) candidates: list[TrajectoryHistory] = [] - if any( - is_selected(exchange) for exchange in trajectory.exchanges.chat_completions - ): + if any(is_selected(exchange) for exchange in trajectory.exchanges.chat_completions): candidates.extend(chat_completions_histories(trajectory, model=model)) - if any( - is_selected(exchange) for exchange in trajectory.exchanges.completions - ): + if any(is_selected(exchange) for exchange in trajectory.exchanges.completions): try: candidates.extend(completions_token_histories(trajectory, model=model)) except ValueError as error: if "requires exact token IDs" not in str(error): raise candidates.extend(completions_string_histories(trajectory, model=model)) - if any( - is_selected(exchange) for exchange in trajectory.exchanges.responses - ): + if any(is_selected(exchange) for exchange in trajectory.exchanges.responses): candidates.extend(responses_histories(trajectory, model=model)) - if any( - is_selected(exchange) for exchange in trajectory.exchanges.messages - ): + if any(is_selected(exchange) for exchange in trajectory.exchanges.messages): candidates.extend(anthropic_messages_histories(trajectory, model=model)) if not candidates: suffix = f" for model {model!r}" if model is not None else "" diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index 822284a86..500f79f90 100644 --- a/src/art/trajectories/_tokenize.py +++ b/src/art/trajectories/_tokenize.py @@ -1,7 +1,7 @@ from __future__ import annotations from bisect import bisect_left -from collections.abc import Mapping, Sequence +from collections.abc import Hashable, Mapping, Sequence from dataclasses import dataclass from datetime import datetime from functools import lru_cache @@ -9,7 +9,7 @@ import json import math import re -from typing import TYPE_CHECKING, Any, Literal, cast +from typing import TYPE_CHECKING, Any, Literal, TypeVar, cast import warnings from anthropic.types import Message, MessageParam, TextBlock @@ -99,9 +99,12 @@ def validate(self, tokenized: TokenizedHistory) -> None: raise AssertionError("Tokenization trace source key is unresolved") +_SourceKeyT = TypeVar("_SourceKeyT", bound=Hashable) + + def _first_introduction_mask( - source_keys: Sequence[_SampledSourceKey | None], - seen: set[_SampledSourceKey], + source_keys: Sequence[_SourceKeyT | None], + seen: set[_SourceKeyT], ) -> list[bool]: keys = {key for key in source_keys if key is not None} new_keys = keys - seen diff --git a/tests/unit/test_exchange_training_model_selection.py b/tests/unit/test_exchange_training_model_selection.py index 6267435e6..c5f325a1a 100644 --- a/tests/unit/test_exchange_training_model_selection.py +++ b/tests/unit/test_exchange_training_model_selection.py @@ -1,6 +1,7 @@ from __future__ import annotations from datetime import datetime, timedelta +from collections.abc import Mapping from typing import SupportsIndex, cast, overload import numpy as np @@ -24,13 +25,18 @@ ChatCompletionsRequest, CompletionsExchange, CompletionsRequest, + TokenizedMultiHistoryTrajectory, + Tokenizer, ) from art.trajectories import _tokenize as trajectory_tokenization from art.trajectories._selection import ( automatic_training_model_selector, resolve_training_model, ) -from art.trajectories._tokenize import _first_introduction_mask +from art.trajectories._tokenize import ( + _first_introduction_mask, + _HistoryTokenizationTrace, +) def _exchange(model: str, output_token: int) -> ChatCompletionsExchange: @@ -211,10 +217,28 @@ def test_training_tokenizes_each_exchange_trajectory_once( calls = 0 original = trajectory_tokenization._tokenize_trajectory_with_trace - def counted(*args: object, **kwargs: object) -> object: + def counted( + trajectory: art.Trajectory, + *, + model: str | None = None, + base_model: str | None = None, + tokenizer: Tokenizer | None = None, + chat_template: str | None = None, + chat_template_kwargs: Mapping[str, object] | None = None, + ) -> tuple[ + TokenizedMultiHistoryTrajectory, + list[_HistoryTokenizationTrace], + ]: nonlocal calls calls += 1 - return original(*args, **kwargs) + return original( + trajectory, + model=model, + base_model=base_model, + tokenizer=tokenizer, + chat_template=chat_template, + chat_template_kwargs=chat_template_kwargs, + ) monkeypatch.setattr( trajectory_tokenization, "_tokenize_trajectory_with_trace", counted @@ -324,9 +348,12 @@ def test_public_training_selector_prefers_exact_model_over_glob_interpretation() ) ) assert resolve_training_model(trajectory, "policy*") == "policy*" - assert [history.model for history in trajectory.histories(model="policy*")] == [ - "policy*" - ] + histories = trajectory.histories(model="policy*") + assert [ + history.model + for history in histories + if not isinstance(history, art.LegacyHistory) + ] == ["policy*"] assert [ history.model for history in trajectory.tokenize( diff --git a/tests/unit/trajectories/test_capture.py b/tests/unit/trajectories/test_capture.py index 0e2bda5e5..d1e6a84a3 100644 --- a/tests/unit/trajectories/test_capture.py +++ b/tests/unit/trajectories/test_capture.py @@ -116,9 +116,7 @@ "token_generations": [ { "prompt_token_ids": [1], - "output_tokens": [ - {"token_id": 2, "logprob": -0.2, "text": "hello"} - ], + "output_tokens": [{"token_id": 2, "logprob": -0.2, "text": "hello"}], "output_indices": [0], } ], From 6d2bef999057e314c5ce329552e3274ce259702c Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Fri, 24 Jul 2026 19:44:20 +0000 Subject: [PATCH 22/58] Reject training without a causal predecessor --- src/art/preprocessing/tokenize.py | 2 + src/art/tinker_native/data.py | 4 +- src/art/trajectories/_tokenize.py | 5 ++ .../test_exchange_training_model_selection.py | 88 +++++++++++++++++++ 4 files changed, 97 insertions(+), 2 deletions(-) diff --git a/src/art/preprocessing/tokenize.py b/src/art/preprocessing/tokenize.py index e98a2161a..853add2e9 100644 --- a/src/art/preprocessing/tokenize.py +++ b/src/art/preprocessing/tokenize.py @@ -740,6 +740,7 @@ def tokenize_trajectory_groups( from ..trajectories._tokenize import ( _as_tokenizer, _first_introduction_mask, + _require_causal_predecessor, _SampledSourceKey, _tokenize_trajectory_with_trace, ) @@ -761,6 +762,7 @@ def tokenize_trajectory_groups( trainable = _first_introduction_mask( trace.source_keys, seen_source_keys ) + _require_causal_predecessor(trainable) if not any(trainable): continue choice_spans = _true_spans(trainable) diff --git a/src/art/tinker_native/data.py b/src/art/tinker_native/data.py index 0bc205452..e0051e57f 100644 --- a/src/art/tinker_native/data.py +++ b/src/art/tinker_native/data.py @@ -324,10 +324,10 @@ def _tokenized_trajectory_to_datum( for selected, flag in zip(trainable, tokenized.flags, strict=True) ): raise ValueError("Only sampled tokens can be selected for Tinker training") + if trainable and trainable[0]: + raise ValueError("A trainable trajectory cannot start with a sampled token") if len(tokenized.token_ids) < 2 or not any(trainable): return None - if trainable[0]: - raise ValueError("A trainable trajectory cannot start with a sampled token") action_mask = trainable[1:] if any( diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index 500f79f90..73cbd56dc 100644 --- a/src/art/trajectories/_tokenize.py +++ b/src/art/trajectories/_tokenize.py @@ -112,6 +112,11 @@ def _first_introduction_mask( return [key in new_keys if key is not None else False for key in source_keys] +def _require_causal_predecessor(trainable: Sequence[bool]) -> None: + if trainable and trainable[0]: + raise ValueError("A trainable trajectory cannot start with a sampled token") + + @dataclass class _TraceBuilder: trace: _HistoryTokenizationTrace | None = None diff --git a/tests/unit/test_exchange_training_model_selection.py b/tests/unit/test_exchange_training_model_selection.py index c5f325a1a..016f49611 100644 --- a/tests/unit/test_exchange_training_model_selection.py +++ b/tests/unit/test_exchange_training_model_selection.py @@ -14,6 +14,7 @@ from art.openai import ART_MOE_ROUTING_METADATA_KEY from art.preprocessing.moe_routing import MoeRouteSegments from art.preprocessing.tokenize import ( + TokenizedResult, _chat_choice_trace, tokenize_trajectory_groups, ) @@ -79,6 +80,42 @@ def _exchange(model: str, output_token: int) -> ChatCompletionsExchange: ) +def _empty_prompt_completion_exchange( + output_token_ids: list[int], +) -> CompletionsExchange: + start = datetime(2026, 1, 1) + return CompletionsExchange( + request=CompletionsRequest(model="policy", prompt=""), + response=Completion.model_validate( + { + "id": "cmpl-empty", + "object": "text_completion", + "created": 1, + "model": "policy", + "choices": [ + { + "index": 0, + "finish_reason": "stop", + "text": "answer", + "prompt_token_ids": [], + "token_ids": output_token_ids, + "logprobs": { + "tokens": [ + f"token_id:{token_id}" for token_id in output_token_ids + ], + "token_logprobs": [-0.1] * len(output_token_ids), + "top_logprobs": [{}] * len(output_token_ids), + "text_offset": list(range(len(output_token_ids))), + }, + } + ], + } + ), + start_time=start, + end_time=start + timedelta(seconds=1), + ) + + def _routed_exchange( *, prompt_token_ids: list[int], @@ -399,6 +436,57 @@ def test_training_selector_rejects_mixed_protocols() -> None: resolve_training_model(trajectory, "policy") +@pytest.mark.parametrize("output_token_ids", [[2], [2, 3]]) +def test_training_rejects_sampled_token_without_causal_predecessor( + output_token_ids: list[int], + monkeypatch: pytest.MonkeyPatch, +) -> None: + exchange = _empty_prompt_completion_exchange(output_token_ids) + group = art.TrajectoryGroup( + [ + art.Trajectory( + exchanges=art.TrajectoryExchanges(completions=[exchange]), + reward=reward, + ) + for reward in (1.0, 0.0) + ] + ) + tokenized = group.trajectories[0].tokenize() + assert tokenized.token_ids == output_token_ids + assert all(flag & art.TokenFlag.SAMPLED for flag in tokenized.flags) + + weight_writes = 0 + original_setattr = TokenizedResult.__setattr__ + + def tracked_setattr(result: TokenizedResult, name: str, value: object) -> None: + nonlocal weight_writes + if name == "weight": + weight_writes += 1 + original_setattr(result, name, value) + + monkeypatch.setattr(TokenizedResult, "__setattr__", tracked_setattr) + with pytest.raises(ValueError, match="cannot start with a sampled token"): + list( + tokenize_trajectory_groups( + cast(PreTrainedTokenizerBase, _Tokenizer()), + [group], + allow_training_without_logprobs=False, + scale_rewards=False, + model="policy", + ) + ) + assert weight_writes == 0 + + with pytest.raises(ValueError, match="cannot start with a sampled token"): + trajectory_groups_to_datums( + [group], + renderer=None, + tokenizer=None, + normalize_advantages=False, + model="policy", + ) + + def test_preprocessing_preserves_adjacent_choice_boundaries_for_moe() -> None: exchanges = [ _routed_exchange( From 2e7a6d23962f542ccaf9976c4bd4f6d9590d9508 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Fri, 24 Jul 2026 19:46:00 +0000 Subject: [PATCH 23/58] Preserve Anthropic prompt lineage metadata --- src/art/trajectories/_history.py | 20 +++--- tests/unit/trajectories/test_tokenize.py | 85 ++++++++++++++++++++++++ 2 files changed, 95 insertions(+), 10 deletions(-) diff --git a/src/art/trajectories/_history.py b/src/art/trajectories/_history.py index da961e3e2..df4c5584d 100644 --- a/src/art/trajectories/_history.py +++ b/src/art/trajectories/_history.py @@ -584,17 +584,18 @@ def _anthropic_prompt(value: object) -> list[AnthropicMessageParam]: def _anthropic_message_key( message: AnthropicMessageParam, *, visible_only: bool = False -) -> tuple[object, ...]: +) -> str: + normalized = copy.deepcopy(dict(message)) content = message.get("content") if isinstance(content, str): - blocks: tuple[object, ...] = (("text", content),) + blocks: list[dict[str, object]] = [{"type": "text", "text": content}] elif isinstance(content, list): - normalized: list[object] = [] + blocks = [] for block in content: if isinstance(block, pydantic.BaseModel): data = block.model_dump(mode="json", exclude_none=True) elif isinstance(block, Mapping): - data = dict(block) + data = copy.deepcopy(dict(block)) else: raise ValueError( "Anthropic message content blocks must be JSON objects" @@ -602,14 +603,13 @@ def _anthropic_message_key( kind = data.get("type") if visible_only and kind in {"thinking", "redacted_thinking"}: continue - if kind == "text": - normalized.append(("text", data.get("text"))) - else: - normalized.append(json.dumps(data, sort_keys=True, default=str)) - blocks = tuple(normalized) + for field in ("token_ids", "logprobs"): + data.pop(field, None) + blocks.append(data) else: raise ValueError("Anthropic message content must be text or a list") - return (message.get("role"), *blocks) + normalized["content"] = blocks + return json.dumps(normalized, sort_keys=True, default=str) def anthropic_messages_history( diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index 7d15c5b13..ae5bced7b 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -230,6 +230,91 @@ def test_messages_exact_prompt_and_output_do_not_load_a_tokenizer( ] +def test_anthropic_cache_control_change_starts_new_source_lineage() -> None: + start = datetime(2026, 1, 1) + + def exchange( + *, + identifier: str, + messages: list[MessageParam], + prompt_token_ids: list[int], + output_token_id: int, + offset: int, + ) -> MessagesExchange: + response = Message.model_validate( + { + "id": identifier, + "type": "message", + "role": "assistant", + "model": "test/model", + "content": [{"type": "text", "text": f"answer {offset}"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": { + "input_tokens": len(prompt_token_ids), + "output_tokens": 1, + }, + "prompt_token_ids": prompt_token_ids, + "token_ids": [output_token_id], + "logprobs": [-0.1], + } + ) + timestamp = start + timedelta(seconds=offset) + return MessagesExchange( + request=MessagesRequest( + model="test/model", + messages=messages, + max_tokens=16, + ), + response=response, + start_time=timestamp, + end_time=timestamp + timedelta(milliseconds=1), + ) + + first = exchange( + identifier="message-1", + messages=[{"role": "user", "content": [{"type": "text", "text": "question"}]}], + prompt_token_ids=[1], + output_token_id=2, + offset=0, + ) + second = exchange( + identifier="message-2", + messages=[ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "question", + "cache_control": {"type": "ephemeral"}, + } + ], + }, + { + "role": "assistant", + "content": [{"type": "text", "text": "answer 0"}], + }, + {"role": "user", "content": "follow up"}, + ], + prompt_token_ids=[1, 2, 3], + output_token_id=4, + offset=1, + ) + + histories = art.Trajectory( + exchanges=TrajectoryExchanges(messages=[first, second]) + ).anthropic_messages_histories() + + assert len(histories) == 2 + updated = histories[1] + source = updated.message_sources[0] + assert source is not None + assert source.exchange is second + assert source.request_index == 0 + assert updated.tokenize().token_ids == [1, 2, 3, 4] + + def test_converted_anthropic_system_history_preserves_exact_assistant_evidence( monkeypatch: pytest.MonkeyPatch, ) -> None: From 1fffa1ff2f78350862d1e8431ae2887c48cbd8c2 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Fri, 24 Jul 2026 19:47:04 +0000 Subject: [PATCH 24/58] Fix converted history source attribution --- src/art/trajectories/_history.py | 1 + src/art/trajectories/_tokenize.py | 208 ++++++++++++++++++----- tests/unit/trajectories/test_tokenize.py | 121 +++++++++++++ 3 files changed, 287 insertions(+), 43 deletions(-) diff --git a/src/art/trajectories/_history.py b/src/art/trajectories/_history.py index df4c5584d..3c1c226c4 100644 --- a/src/art/trajectories/_history.py +++ b/src/art/trajectories/_history.py @@ -1299,6 +1299,7 @@ def anthropic_as_chat_completions_history( ChatCompletionsMessageSource( exchange=source.exchange, request_index=source.request_index, + output_index=0 if source.request_index is None else None, ) if source is not None else None diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index 73cbd56dc..bad6691b6 100644 --- a/src/art/trajectories/_tokenize.py +++ b/src/art/trajectories/_tokenize.py @@ -2046,11 +2046,12 @@ def _validate_history_sources(history: History) -> None: elif isinstance(exchange, MessagesExchange): if ( source.choice_index is not None - or source.output_index is not None or source.generation_index is not None ): raise ValueError("Anthropic-to-Chat source has invalid indices") if source.request_index is not None: + if source.output_index is not None: + raise ValueError("Anthropic-to-Chat source has invalid indices") request_messages = exchange.request.get("messages", []) request_index = _source_index( source.request_index, @@ -2062,16 +2063,16 @@ def _validate_history_sources(history: History) -> None: {"messages": [request_messages[request_index]]} ) ) - else: - if message.get("role") == "system": - expected.extend( - _anthropic_messages( - { - "system": exchange.request.get("system"), - "messages": [], - } - ) + elif source.output_index is None: + expected.extend( + _anthropic_messages( + { + "system": exchange.request.get("system"), + "messages": [], + } ) + ) + elif source.output_index == 0: expected.append(_response_message(exchange)) visible_blocks = [ block.model_dump(mode="python", exclude_none=True) @@ -2088,6 +2089,8 @@ def _validate_history_sources(history: History) -> None: } ) ) + else: + raise ValueError("Anthropic response source index is out of bounds") elif isinstance(exchange, ResponsesExchange): if source.request_index is not None: if ( @@ -2143,22 +2146,20 @@ def _validate_history_sources(history: History) -> None: "Responses output source does not belong to " "its generation" ) - expected.extend( - _responses_messages( - { - "input": [ - exchange.response.output[ - generation_output_index - ].model_dump( - mode="python", exclude_none=True - ) - for generation_output_index in generations[ - generation_index - ].output_indices - ] - } - ) + generation_messages = _responses_messages( + { + "input": [ + exchange.response.output[ + generation_output_index + ].model_dump(mode="python", exclude_none=True) + for generation_output_index in generations[ + generation_index + ].output_indices + ] + } ) + if len(generation_messages) == 1: + expected.extend(generation_messages) expected.extend( _responses_messages( { @@ -2452,7 +2453,7 @@ def _source_signature(source: object) -> tuple[object, ...] | None: ) elif ( isinstance(exchange, MessagesExchange) - and getattr(source, "request_index", None) is None + and getattr(source, "output_index", None) == 0 ): evidence_fingerprint = _sampled_evidence_fingerprint( exchange, protocol="messages", index=0 @@ -2712,6 +2713,54 @@ def _tokenize_exact_responses_history( return tokenized +def _responses_source_generation( + source: object, +) -> tuple[ResponsesExchange, _ResponseGeneration] | None: + exchange = getattr(source, "exchange", None) + generation_index = getattr(source, "generation_index", None) + if not isinstance(exchange, ResponsesExchange) or not isinstance( + generation_index, int + ): + return None + generations = _response_generations(exchange.response) + if not 0 <= generation_index < len(generations): + raise ValueError("Responses source generation index is out of bounds") + generation = generations[generation_index] + output_index = getattr(source, "output_index", None) + if isinstance(output_index, int) and output_index not in generation.output_indices: + raise ValueError( + "Responses source output index does not belong to its generation" + ) + return exchange, generation + + +def _responses_generation_messages(source: object) -> list[dict[str, Any]] | None: + selected = _responses_source_generation(source) + if selected is None: + return None + exchange, generation = selected + return _responses_messages( + { + "input": [ + exchange.response.output[index].model_dump( + mode="python", exclude_none=True + ) + for index in generation.output_indices + ] + } + ) + + +def _responses_generation_full_tokens( + source: object, +) -> tuple[list[int] | None, list[float]]: + selected = _responses_source_generation(source) + if selected is None: + return None, [] + _, generation = selected + return generation.output_token_ids, generation.output_logprobs + + def _chat_source_full_tokens( source: object, ) -> tuple[list[int] | None, list[float]]: @@ -2726,25 +2775,18 @@ def _chat_source_full_tokens( _, tokens, logprobs = _chat_choice_tokens(choice, _dump(exchange.response)) return tokens, logprobs if isinstance(exchange, ResponsesExchange): - generation_index = getattr(source, "generation_index", None) + generation_messages = _responses_generation_messages(source) + if generation_messages is not None: + if len(generation_messages) != 1: + return None, [] + return _responses_generation_full_tokens(source) generations = _response_generations(exchange.response) - if isinstance(generation_index, int): - if not 0 <= generation_index < len(generations): - raise ValueError("Responses source generation index is out of bounds") - generation = generations[generation_index] - output_index = getattr(source, "output_index", None) - if ( - isinstance(output_index, int) - and output_index not in generation.output_indices - ): - raise ValueError( - "Responses source output index does not belong to its generation" - ) - return generation.output_token_ids, generation.output_logprobs if len(generations) > 1: return None, [] return _responses_tokens(exchange.response)[1:] if isinstance(exchange, MessagesExchange): + if getattr(source, "output_index", None) != 0: + return None, [] _, tokens, logprobs = _exchange_tokens(exchange) return tokens, logprobs return None, [] @@ -2770,6 +2812,8 @@ def _chat_source_prompt_tokens(source: object) -> list[int] | None: return generations[generation_index].prompt_token_ids return _responses_tokens(exchange.response)[0] if isinstance(exchange, MessagesExchange): + if getattr(source, "output_index", None) != 0: + return None return _messages_tokens(exchange.response)[0] return None @@ -2779,7 +2823,7 @@ def _source_is_sampled(source: object) -> bool: if isinstance(exchange, ChatCompletionsExchange): return getattr(source, "choice_index", None) is not None if isinstance(exchange, MessagesExchange): - return getattr(source, "request_index", None) is None + return getattr(source, "output_index", None) == 0 if isinstance(exchange, ResponsesExchange): return ( getattr(source, "output_index", None) is not None @@ -2940,7 +2984,7 @@ def _chat_source_tokens( return None, [] if ( isinstance(exchange, MessagesExchange) - and getattr(source, "request_index", None) is None + and getattr(source, "output_index", None) == 0 ): block_type = "thinking" if part == "reasoning" else "text" blocks = [ @@ -3036,7 +3080,7 @@ def _source_covers_complete_sampled_message( message ) == normalize_chat_message(expected) if isinstance(exchange, MessagesExchange): - if getattr(source, "request_index", None) is not None: + if getattr(source, "output_index", None) != 0: return False return normalize_chat_message(message) == normalize_chat_message( _response_message(exchange) @@ -3182,10 +3226,88 @@ def part_ids(text: str) -> list[int]: and _source_is_sampled(source) for _, text in _chat_message_parts(message) ] + from ._history import normalize_chat_message + + atomic_generations: dict[int, tuple[int, object, list[int], list[float]]] = {} + seen_generations: set[tuple[int, int]] = set() + for message_index, source in enumerate(history.message_sources): + if source is None or not isinstance( + getattr(source, "exchange", None), ResponsesExchange + ): + continue + generation_index = getattr(source, "generation_index", None) + output_index = getattr(source, "output_index", None) + if not isinstance(generation_index, int) or not isinstance(output_index, int): + continue + key = (id(source.exchange), generation_index) + if key in seen_generations: + continue + seen_generations.add(key) + projected = _responses_generation_messages(source) + if projected is None or len(projected) <= 1: + continue + end = message_index + len(projected) + if end > len(messages) or [ + normalize_chat_message(item) for item in messages[message_index:end] + ] != [normalize_chat_message(item) for item in projected]: + continue + group_sources = history.message_sources[message_index:end] + if any( + item is None + or item.exchange is not source.exchange + or item.generation_index != generation_index + for item in group_sources + ): + continue + exact, exact_logprobs = _responses_generation_full_tokens(source) + if exact is not None: + atomic_generations[message_index] = ( + end, + source, + exact, + exact_logprobs, + ) + sampled_part_cursor = 0 + atomic_end = 0 for message_index, (message, source) in enumerate( zip(history.messages, history.message_sources, strict=True) ): + if message_index < atomic_end: + continue + atomic = atomic_generations.get(message_index) + if atomic is not None: + end, atomic_source, exact, exact_logprobs = atomic + prompt_render = render(messages[:message_index], add_generation_prompt=True) + completed_render = render(messages[:end], add_generation_prompt=False) + if ( + len(completed_render) <= len(prompt_render) + or completed_render[: len(prompt_render)] != prompt_render + or rendered[: len(completed_render)] != completed_render + ): + raise ValueError( + "Could not locate a complete sampled generation in the " + "rendered history" + ) + replacements.append( + ( + len(prompt_render), + len(completed_render), + exact, + exact_logprobs + if len(exact_logprobs) == len(exact) + else [math.nan] * len(exact), + True, + _sampled_source_key(atomic_source), + atomic_source, + ) + ) + search_cursor = len(completed_render) + sampled_part_cursor += sum( + len(_chat_message_parts(item)) for item in messages[message_index:end] + ) + atomic_end = end + continue parts = _chat_message_parts(message) sampled = ( message.get("role") == "assistant" diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index ae5bced7b..cb4619a26 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -363,6 +363,45 @@ def test_converted_anthropic_system_history_preserves_exact_assistant_evidence( assert tokenized.flags[-1] == art.TokenFlag.EXACT | art.TokenFlag.SAMPLED +def test_converted_anthropic_system_source_rejects_sampled_response_mutation() -> None: + response = Message.model_validate( + { + "id": "message-system-source", + "type": "message", + "role": "assistant", + "model": "test/model", + "content": [{"type": "text", "text": "answer"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 2, "output_tokens": 1}, + "prompt_token_ids": [10, 11], + "token_ids": [12], + "logprobs": [-0.12], + } + ) + start = datetime(2026, 1, 1) + exchange = MessagesExchange( + request=MessagesRequest( + model="test/model", + system="system", + messages=[{"role": "user", "content": "question"}], + max_tokens=16, + ), + response=response, + start_time=start, + end_time=start + timedelta(milliseconds=1), + ) + history = ( + art.Trajectory(exchanges=TrajectoryExchanges(messages=[exchange])) + .anthropic_messages_history() + .as_chat_completions_history() + ) + history.messages[0] = {"role": "assistant", "content": "answer"} + + with pytest.raises(ValueError, match="no longer matches its source exchange"): + history.tokenize() + + def test_converted_responses_history_tokenizes_without_native_chat_exchange( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -2476,6 +2515,88 @@ def test_responses_generation_source_rejects_content_from_another_generation() - history.tokenize() +def _multi_output_responses_chat_history() -> art.ChatCompletionsHistory: + exchange = _response_exchange("multi-output-generation", 2) + data = exchange.response.model_dump(mode="python") + data["output"][0]["content"][0]["text"] = "first" + data["output"].append( + { + "id": "message-second-output", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [ + { + "type": "output_text", + "text": "second", + "annotations": [], + "logprobs": [], + } + ], + } + ) + data["token_generations"] = [ + { + "prompt_token_ids": [1], + "output_tokens": [ + {"token_id": 2, "logprob": -0.2}, + {"token_id": 3, "logprob": -0.3}, + ], + "output_indices": [0, 1], + } + ] + exchange.response = Response.model_validate(data) + return ( + art.Trajectory(exchanges=TrajectoryExchanges(responses=[exchange])) + .responses_history() + .as_chat_completions_history() + ) + + +def test_responses_output_source_rejects_sibling_generation_output() -> None: + history = _multi_output_responses_chat_history() + first_output = next( + index + for index, source in enumerate(history.message_sources) + if source is not None and source.output_index == 0 + ) + history.messages[first_output] = { + "role": "assistant", + "content": "second", + } + + with pytest.raises(ValueError, match="no longer matches its source exchange"): + history.tokenize() + + +def test_responses_multi_output_generation_exact_evidence_is_consumed_once() -> None: + history = _multi_output_responses_chat_history() + + class Tokenizer: + def __call__(self, text: str, **kwargs: object) -> list[int]: + del kwargs + return {"turn 0": [1], "first": [20], "second": [30]}[text] + + def apply_chat_template( + self, messages: list[dict[str, Any]], **kwargs: object + ) -> list[int]: + del kwargs + result: list[int] = [] + for message in messages: + result.extend(self(str(message["content"]))) + return result + + tokenized = history.tokenize(tokenizer=Tokenizer(), chat_template="custom") + + assert tokenized.token_ids == [1, 2, 3] + assert tokenized.logprobs[1:] == pytest.approx([-0.2, -0.3]) + assert tokenized.flags == [ + art.TokenFlag(0), + art.TokenFlag.EXACT | art.TokenFlag.SAMPLED, + art.TokenFlag.EXACT | art.TokenFlag.SAMPLED, + ] + + def test_mutable_chat_history_is_authoritative_and_does_not_replay_removed_turns() -> ( None ): From 60721cce7f1e6a996954cd9659e079da457c2393 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Fri, 24 Jul 2026 19:53:09 +0000 Subject: [PATCH 25/58] Scale Chat projection with captured evidence --- src/art/trajectories/_history.py | 60 ++++++++++++++++++++--- tests/unit/trajectories/test_history.py | 65 +++++++++++++++++++++++++ 2 files changed, 118 insertions(+), 7 deletions(-) diff --git a/src/art/trajectories/_history.py b/src/art/trajectories/_history.py index 3c1c226c4..bf8ac1467 100644 --- a/src/art/trajectories/_history.py +++ b/src/art/trajectories/_history.py @@ -124,6 +124,7 @@ class _Branch(Generic[_ItemT, _SourceT, _ContextT]): order: tuple[int, ...] first_time: datetime context_source: _ModelledExchange | None + lineage_keys: list[object] | None = None def _selected_models( @@ -184,7 +185,9 @@ def _ordered_choices(choices: Sequence[_IndexedT], *, protocol: str) -> list[_In def _is_prefix(prefix: Sequence[object], value: Sequence[object]) -> bool: - return len(prefix) <= len(value) and list(value[: len(prefix)]) == list(prefix) + return len(prefix) <= len(value) and all( + left == right for left, right in zip(prefix, value[: len(prefix)], strict=True) + ) def _lineage_prompt_sources( @@ -193,14 +196,19 @@ def _lineage_prompt_sources( prompt: Sequence[_ItemT], defaults: Sequence[_SourceT | None], equivalent: Callable[[_ItemT, _ItemT], bool], + prompt_keys: Sequence[object] | None = None, ) -> list[_SourceT | None]: candidates = [ branch for branch in branches if len(branch.items) < len(prompt) - and all( - equivalent(existing, current) - for existing, current in zip(branch.items, prompt, strict=False) + and ( + list(prompt_keys[: len(branch.items)]) == branch.lineage_keys + if prompt_keys is not None and branch.lineage_keys is not None + else all( + equivalent(existing, current) + for existing, current in zip(branch.items, prompt, strict=False) + ) ) ] if not candidates: @@ -220,9 +228,22 @@ def _extend_branches( start_time: datetime, context_source: _ModelledExchange | None = None, continuation: Callable[[_Branch[_ItemT, _SourceT, _ContextT]], bool] | None = None, + prompt_lineage_keys: Sequence[object] | None = None, + output_lineage_keys: Sequence[Sequence[object]] | None = None, ) -> None: if len(prompt) != len(prompt_sources): raise AssertionError("prompt sources must parallel prompt items") + if (prompt_lineage_keys is None) != (output_lineage_keys is None): + raise AssertionError("prompt and output lineage keys must be provided together") + if prompt_lineage_keys is not None and len(prompt_lineage_keys) != len(prompt): + raise AssertionError("prompt lineage keys must parallel prompt items") + if output_lineage_keys is not None and len(output_lineage_keys) != len(outputs): + raise AssertionError("output lineage keys must parallel outputs") + if output_lineage_keys is not None and any( + len(keys) != len(output) + for keys, (_, output, _) in zip(output_lineage_keys, outputs, strict=True) + ): + raise AssertionError("output lineage keys must parallel output items") candidates = [ (index, branch) for index, branch in enumerate(branches) @@ -261,14 +282,23 @@ def _extend_branches( first_time = parent.first_time if parent is not None else start_time created = [ _Branch( - items=[*copy.deepcopy(prompt), *copy.deepcopy(output)], + items=( + [*prompt, *output] + if prompt_lineage_keys is not None + else [*copy.deepcopy(prompt), *copy.deepcopy(output)] + ), sources=[*sources, *output_sources], context=copy.deepcopy(context), order=(*base_order, choice_index), first_time=first_time, context_source=context_source, + lineage_keys=( + [*prompt_lineage_keys, *output_lineage_keys[position]] + if prompt_lineage_keys is not None and output_lineage_keys is not None + else None + ), ) - for choice_index, output, output_sources in outputs + for position, (choice_index, output, output_sources) in enumerate(outputs) ] if not created and remove_parent and parent is not None: branches.append(parent) @@ -328,7 +358,13 @@ def _chat_retains_sampled_reasoning( def _chat_message_key(message: Message, *, visible_only: bool = False) -> str: - data = normalize_chat_message(message) + # History messages are already normalized and detached from the exchange. + # Copy only the top-level mapping because keying never mutates nested values. + data = dict(message) + if data.get("content") is None: + data.pop("content", None) + if data.get("tool_calls") == []: + data.pop("tool_calls") if visible_only: data.pop("reasoning", None) data.pop("reasoning_content", None) @@ -399,6 +435,9 @@ def chat_completions_histories( exchange.request.get("messages", []) ) ] + prompt_lineage_keys = [ + _chat_message_key(message, visible_only=True) for message in prompt + ] prompt_sources = _lineage_prompt_sources( branches, prompt=prompt, @@ -410,6 +449,7 @@ def chat_completions_histories( _chat_message_key(left, visible_only=True) == _chat_message_key(right, visible_only=True) ), + prompt_keys=prompt_lineage_keys, ) outputs: list[ tuple[ @@ -436,6 +476,10 @@ def chat_completions_histories( [response_source], ) ) + output_lineage_keys = [ + [_chat_message_key(output[0], visible_only=True)] + for _, output, _ in outputs + ] _extend_branches( branches, prompt=prompt, @@ -453,6 +497,8 @@ def chat_completions_histories( continuation=lambda branch: _chat_retains_sampled_reasoning( branch, exchange, len(prompt) ), + prompt_lineage_keys=prompt_lineage_keys, + output_lineage_keys=output_lineage_keys, ) for branch in sorted(branches, key=lambda item: (item.first_time, item.order)): histories.append( diff --git a/tests/unit/trajectories/test_history.py b/tests/unit/trajectories/test_history.py index a4e671e57..1e28c64a6 100644 --- a/tests/unit/trajectories/test_history.py +++ b/tests/unit/trajectories/test_history.py @@ -1,5 +1,7 @@ from datetime import datetime, timedelta import importlib +from statistics import median +from time import perf_counter from typing import Any, cast from anthropic.types import Message @@ -17,6 +19,7 @@ ResponsesExchange, TrajectoryExchanges, ) +from art.types import Message as ChatMessage def _times(offset: int = 0) -> tuple[datetime, datetime]: @@ -54,6 +57,16 @@ def _chat( ) +def _growing_chat_trajectory(turn_count: int) -> art.Trajectory: + exchanges: list[ChatCompletionsExchange] = [] + messages: list[dict[str, object]] = [] + for index in range(turn_count): + messages.append({"role": "user", "content": f"question {index}"}) + exchanges.append(_chat(list(messages), f"answer {index}", offset=index)) + messages.append({"role": "assistant", "content": f"answer {index}"}) + return art.Trajectory(exchanges=TrajectoryExchanges(chat_completions=exchanges)) + + def _completion( prompt: list[int], output: list[int], *, offset: int = 0 ) -> CompletionsExchange: @@ -220,6 +233,58 @@ def test_chat_history_resolves_one_model_and_append_only_sequence() -> None: trajectory.chat_completions_history(model="test/model") +def test_chat_projection_keys_each_captured_message_once( + monkeypatch: pytest.MonkeyPatch, +) -> None: + history_module = importlib.import_module("art.trajectories._history") + original = history_module._chat_message_key + calls = 0 + + def count(message: ChatMessage, *, visible_only: bool = False) -> str: + nonlocal calls + calls += 1 + return original(message, visible_only=visible_only) + + monkeypatch.setattr(history_module, "_chat_message_key", count) + trajectory = _growing_chat_trajectory(16) + + trajectory.chat_completions_history() + + expected = sum( + len(exchange.request["messages"]) + len(exchange.response.choices) + for exchange in trajectory.exchanges.chat_completions + ) + assert calls == expected + + +def test_chat_projection_scales_with_captured_messages() -> None: + normalized_medians: list[float] = [] + raw_medians: list[float] = [] + for turn_count in (32, 64, 128): + trajectory = _growing_chat_trajectory(turn_count) + captured_messages = sum( + len(exchange.request["messages"]) + len(exchange.response.choices) + for exchange in trajectory.exchanges.chat_completions + ) + trajectory.chat_completions_history() + samples: list[float] = [] + for _ in range(5): + started = perf_counter() + trajectory.chat_completions_history() + samples.append(perf_counter() - started) + elapsed = median(samples) + raw_medians.append(elapsed) + normalized_medians.append(elapsed / captured_messages) + + # A growing transcript contains O(turns²) captured messages: doubling the + # turn count quadruples the input that projection must validate. Projection + # should remain linear in that actual input, without hidden superlinear work. + assert raw_medians[1] < raw_medians[0] * 5 + assert raw_medians[2] < raw_medians[1] * 5 + assert normalized_medians[1] < normalized_medians[0] * 2 + assert normalized_medians[2] < normalized_medians[1] * 2 + + def test_model_patterns_select_matching_histories_only() -> None: policy_12 = _chat( [{"role": "user", "content": "one"}], From 0229f30085a0211173814de7552f3bf304b47d34 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Fri, 24 Jul 2026 19:53:45 +0000 Subject: [PATCH 26/58] Format training trace tests --- tests/unit/test_exchange_training_model_selection.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/test_exchange_training_model_selection.py b/tests/unit/test_exchange_training_model_selection.py index 016f49611..6b9bbb932 100644 --- a/tests/unit/test_exchange_training_model_selection.py +++ b/tests/unit/test_exchange_training_model_selection.py @@ -1,7 +1,7 @@ from __future__ import annotations -from datetime import datetime, timedelta from collections.abc import Mapping +from datetime import datetime, timedelta from typing import SupportsIndex, cast, overload import numpy as np From 310dc839544b1dd92a7ea9935bda223a95530c4a Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Fri, 24 Jul 2026 19:53:59 +0000 Subject: [PATCH 27/58] Gate Chat projection by captured bytes --- tests/unit/trajectories/test_history.py | 28 ++++++++++++------------- 1 file changed, 14 insertions(+), 14 deletions(-) diff --git a/tests/unit/trajectories/test_history.py b/tests/unit/trajectories/test_history.py index 1e28c64a6..2fd5d7e17 100644 --- a/tests/unit/trajectories/test_history.py +++ b/tests/unit/trajectories/test_history.py @@ -258,12 +258,11 @@ def count(message: ChatMessage, *, visible_only: bool = False) -> str: def test_chat_projection_scales_with_captured_messages() -> None: - normalized_medians: list[float] = [] - raw_medians: list[float] = [] + measurements: list[tuple[int, float, int]] = [] for turn_count in (32, 64, 128): trajectory = _growing_chat_trajectory(turn_count) - captured_messages = sum( - len(exchange.request["messages"]) + len(exchange.response.choices) + captured_bytes = sum( + len(exchange.model_dump_json().encode()) for exchange in trajectory.exchanges.chat_completions ) trajectory.chat_completions_history() @@ -273,16 +272,17 @@ def test_chat_projection_scales_with_captured_messages() -> None: trajectory.chat_completions_history() samples.append(perf_counter() - started) elapsed = median(samples) - raw_medians.append(elapsed) - normalized_medians.append(elapsed / captured_messages) - - # A growing transcript contains O(turns²) captured messages: doubling the - # turn count quadruples the input that projection must validate. Projection - # should remain linear in that actual input, without hidden superlinear work. - assert raw_medians[1] < raw_medians[0] * 5 - assert raw_medians[2] < raw_medians[1] * 5 - assert normalized_medians[1] < normalized_medians[0] * 2 - assert normalized_medians[2] < normalized_medians[1] * 2 + measurements.append((turn_count, elapsed, captured_bytes)) + + # A growing transcript contains O(turns²) serialized evidence: doubling the + # turn count roughly quadruples the bytes that projection must validate. + # Preserve the turn-count curve as diagnostics, but gate near-linear cost in + # the actual input size rather than imposing an impossible per-turn ratio. + normalized = [ + elapsed / captured_bytes for _, elapsed, captured_bytes in measurements + ] + assert normalized[1] < normalized[0] * 2, measurements + assert normalized[2] < normalized[1] * 2, measurements def test_model_patterns_select_matching_histories_only() -> None: From 09a649207e7ae25652bb52af745c278862c8b69c Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Fri, 24 Jul 2026 20:08:01 +0000 Subject: [PATCH 28/58] Disambiguate echoed completion evidence --- src/art/trajectories/_tokenize.py | 33 +++++---- tests/unit/trajectories/test_tokenize.py | 90 ++++++++++++++++++++++-- 2 files changed, 105 insertions(+), 18 deletions(-) diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index bad6691b6..1682643be 100644 --- a/src/art/trajectories/_tokenize.py +++ b/src/art/trajectories/_tokenize.py @@ -530,10 +530,21 @@ def _completion_evidence( else math.nan for value in logprobs.get("token_logprobs") or [] ] + pair_includes_prompt = ( + echo + and prompt_ids is not None + and token_ids is not None + and pair_ids == [*prompt_ids, *token_ids] + ) + token_ids_include_prompt = ( + echo + and prompt_ids is not None + and token_ids is not None + and bool(pair_ids) + and token_ids == [*prompt_ids, *pair_ids] + ) if token_ids is not None and pair_ids and token_ids != pair_ids: - if not ( - echo and prompt_ids is not None and pair_ids == [*prompt_ids, *token_ids] - ): + if not pair_includes_prompt and not token_ids_include_prompt: raise ValueError("Response token IDs disagree with completion logprobs") selected = token_ids if token_ids is not None else pair_ids or None prompt_logprobs: list[float] = [] @@ -545,23 +556,19 @@ def _completion_evidence( selected = None completion_logprobs = [] elif echo and prompt_ids is not None and selected is not None: - pair_includes_prompt = token_ids is not None and pair_ids == [ - *prompt_ids, - *token_ids, - ] if pair_includes_prompt: prompt_logprobs = pair_logprobs[: len(prompt_ids)] completion_logprobs = pair_logprobs[len(prompt_ids) :] elif len(pair_logprobs) == len(prompt_ids) + len(selected): prompt_logprobs = pair_logprobs[: len(prompt_ids)] completion_logprobs = pair_logprobs[len(prompt_ids) :] + elif token_ids_include_prompt: + selected = selected[len(prompt_ids) :] + completion_logprobs = pair_logprobs elif selected[: len(prompt_ids)] == prompt_ids and ( - pair_ids == selected - or ( - len(tokens) == len(selected) - and all(isinstance(value, str) for value in tokens) - and "".join(cast(str, value) for value in tokens) == choice.text - ) + len(tokens) == len(selected) + and all(isinstance(value, str) for value in tokens) + and "".join(cast(str, value) for value in tokens) == choice.text ): if pair_logprobs and len(pair_logprobs) != len(selected): raise ValueError("Completions token IDs and logprobs differ in length") diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index cb4619a26..affcf793a 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -508,13 +508,10 @@ def test_completions_support_single_item_batches_and_echo( assert tokenized.token_ids == [1, 2] -@pytest.mark.parametrize("response_token_ids", [[1, 2], [2]]) -def test_completions_echo_preserves_prompt_logprobs_without_sampling_them( - response_token_ids: list[int], -) -> None: +def test_completions_echo_preserves_prompt_logprobs_without_sampling_them() -> None: exchange = _completion_exchange(echo=True) payload = exchange.response.model_dump(mode="python") - payload["choices"][0]["token_ids"] = response_token_ids + payload["choices"][0]["token_ids"] = [2] payload["choices"][0]["logprobs"] = { "tokens": ["token_id:1", "token_id:2"], "token_logprobs": [-0.1, -0.2], @@ -535,6 +532,89 @@ def test_completions_echo_preserves_prompt_logprobs_without_sampling_them( ] +def test_completions_echo_does_not_strip_repeated_prompt_token_from_completion() -> ( + None +): + exchange = _completion_exchange(echo=True) + payload = exchange.response.model_dump(mode="python") + payload["choices"][0]["text"] = "questionquestionanswer" + payload["choices"][0]["token_ids"] = [1, 2] + payload["choices"][0]["logprobs"] = { + "tokens": ["token_id:1", "token_id:2"], + "token_logprobs": [-0.2, -0.3], + "top_logprobs": [{}, {}], + "text_offset": [8, 16], + } + exchange.response = Completion.model_validate(payload) + + class Tokenizer: + def __call__(self, text: str, **kwargs: object) -> list[int]: + del kwargs + return {"question": [1], "questionanswer": [1, 2]}[text] + + tokenized = art.Trajectory( + exchanges=TrajectoryExchanges(completions=[exchange]) + ).tokenize(tokenizer=Tokenizer()) + + assert tokenized.token_ids == [1, 1, 2] + assert tokenized.flags == [ + art.TokenFlag.EXACT, + art.TokenFlag.SAMPLED, + art.TokenFlag.SAMPLED, + ] + assert all(math.isnan(logprob) for logprob in tokenized.logprobs) + + +def test_completions_echo_strips_prompt_from_proven_combined_token_carrier() -> None: + exchange = _completion_exchange(echo=True) + payload = exchange.response.model_dump(mode="python") + payload["choices"][0]["text"] = "questionquestionanswer" + payload["choices"][0]["token_ids"] = [1, 1, 2] + payload["choices"][0]["logprobs"] = { + "tokens": ["token_id:1", "token_id:2"], + "token_logprobs": [-0.2, -0.3], + "top_logprobs": [{}, {}], + "text_offset": [8, 16], + } + exchange.response = Completion.model_validate(payload) + + tokenized = art.Trajectory( + exchanges=TrajectoryExchanges(completions=[exchange]) + ).tokenize() + + assert tokenized.token_ids == [1, 1, 2] + assert tokenized.logprobs[1:] == pytest.approx([-0.2, -0.3]) + assert tokenized.flags == [ + art.TokenFlag.EXACT, + art.TokenFlag.EXACT | art.TokenFlag.SAMPLED, + art.TokenFlag.EXACT | art.TokenFlag.SAMPLED, + ] + + +def test_completions_echo_strips_prompt_from_proven_textual_carrier() -> None: + exchange = _completion_exchange(echo=True) + payload = exchange.response.model_dump(mode="python") + payload["choices"][0]["token_ids"] = [1, 2] + payload["choices"][0]["logprobs"] = { + "tokens": ["question", "answer"], + "token_logprobs": [-0.1, -0.2], + "top_logprobs": [{}, {}], + "text_offset": [0, 8], + } + exchange.response = Completion.model_validate(payload) + + tokenized = art.Trajectory( + exchanges=TrajectoryExchanges(completions=[exchange]) + ).tokenize() + + assert tokenized.token_ids == [1, 2] + assert tokenized.logprobs == pytest.approx([-0.1, -0.2]) + assert tokenized.flags == [ + art.TokenFlag.EXACT, + art.TokenFlag.EXACT | art.TokenFlag.SAMPLED, + ] + + def test_completions_echo_prefers_full_logprob_carrier_to_id_prefix_heuristic() -> None: exchange = _completion_exchange(echo=True) payload = exchange.response.model_dump(mode="python") From 402d2386c5930b3675f2d744bbe310f95ec1cb34 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Fri, 24 Jul 2026 20:11:45 +0000 Subject: [PATCH 29/58] Finalize captured streams at terminal events --- src/art/trajectories/_capture/core.py | 58 +++++++++++ src/art/trajectories/_protocols.py | 7 +- tests/unit/trajectories/test_capture.py | 124 +++++++++++++++++++++--- 3 files changed, 175 insertions(+), 14 deletions(-) diff --git a/src/art/trajectories/_capture/core.py b/src/art/trajectories/_capture/core.py index b6d3f1cb1..3e4cac320 100644 --- a/src/art/trajectories/_capture/core.py +++ b/src/art/trajectories/_capture/core.py @@ -22,6 +22,43 @@ _adapter_active: contextvars.ContextVar[bool] = contextvars.ContextVar( "art_capture_adapter_active", default=False ) +_SSE_DELIMITERS = (b"\r\n\r\n", b"\n\n", b"\r\r") + + +def _terminal_sse_event(endpoint: Endpoint, block: bytes) -> bool: + try: + lines = block.decode("utf-8").splitlines() + except UnicodeDecodeError: + return False + event_name: str | None = None + data_lines: list[str] = [] + for line in lines: + field, separator, value = line.partition(":") + if not separator: + continue + if value.startswith(" "): + value = value[1:] + if field == "event": + event_name = value + elif field == "data": + data_lines.append(value) + data = "\n".join(data_lines) + if endpoint in {"chat_completions", "completions"}: + return data == "[DONE]" + if endpoint == "responses" and event_name == "response.completed": + return True + if endpoint == "messages" and event_name == "message_stop": + return True + try: + payload = json.loads(data) + except json.JSONDecodeError: + return False + if not isinstance(payload, dict): + return False + event_type = payload.get("type") + return (endpoint == "responses" and event_type == "response.completed") or ( + endpoint == "messages" and event_type == "message_stop" + ) @dataclass @@ -33,10 +70,31 @@ class CaptureState: status_code: int | None = None body: bytearray = field(default_factory=bytearray) captured: bool = False + _event_start: int = field(default=0, init=False, repr=False) + _scan_start: int = field(default=0, init=False, repr=False) def add(self, chunk: bytes) -> None: if not self.captured: self.body.extend(chunk) + if self.request.get("stream") is True and self._reached_terminal_event(): + self.finish() + + def _reached_terminal_event(self) -> bool: + while True: + boundaries = [ + (index, len(delimiter)) + for delimiter in _SSE_DELIMITERS + if (index := self.body.find(delimiter, self._scan_start)) >= 0 + ] + if not boundaries: + self._scan_start = max(self._event_start, len(self.body) - 3) + return False + index, delimiter_length = min(boundaries) + block = bytes(self.body[self._event_start : index]) + self._event_start = index + delimiter_length + self._scan_start = self._event_start + if _terminal_sse_event(self.endpoint, block): + return True def discard(self) -> None: self.body.clear() diff --git a/src/art/trajectories/_protocols.py b/src/art/trajectories/_protocols.py index 1c420e9df..594d5a0e3 100644 --- a/src/art/trajectories/_protocols.py +++ b/src/art/trajectories/_protocols.py @@ -53,9 +53,12 @@ def endpoint_for_url(url: str) -> Endpoint | None: def _sse_events(body: bytes) -> list[tuple[str | None, SSEPayload]]: - text = body.decode("utf-8").replace("\r\n", "\n") + text = body.decode("utf-8").replace("\r\n", "\n").replace("\r", "\n") events: list[tuple[str | None, SSEPayload]] = [] - for block in text.split("\n\n"): + blocks = text.split("\n\n") + if not text.endswith("\n\n"): + blocks.pop() + for block in blocks: event_name: str | None = None data_lines: list[str] = [] for line in block.splitlines(): diff --git a/tests/unit/trajectories/test_capture.py b/tests/unit/trajectories/test_capture.py index d1e6a84a3..aef8b1b9c 100644 --- a/tests/unit/trajectories/test_capture.py +++ b/tests/unit/trajectories/test_capture.py @@ -449,6 +449,45 @@ async def response(_: httpx.Request) -> httpx.Response: assert len(trajectory.exchanges.chat_completions) == 1 +@pytest.mark.parametrize("encoding", ["gzip", "deflate", "br"]) +async def test_httpx_terminal_event_captures_before_response_close( + encoding: str, +) -> None: + body = _streaming_chat_body() + compressed = _encoded(body, encoding) + + async def response(_: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + headers={ + "content-encoding": encoding, + "content-type": "text/event-stream", + }, + stream=_AsyncChunks(compressed), + ) + + client = httpx.AsyncClient(transport=httpx.MockTransport(response)) + request = client.build_request( + "POST", + "https://example.test/v1/chat/completions", + json={"model": "test/model", "messages": [], "stream": True}, + ) + with art.Trajectory() as trajectory: + result = await client.send(request, stream=True) + iterator = cast(AsyncGenerator[bytes, None], result.aiter_bytes()) + received = bytearray() + async for chunk in iterator: + received.extend(chunk) + if b"data: [DONE]\n\n" in received: + break + assert len(trajectory.exchanges.chat_completions) == 1 + + assert not result.is_closed + await iterator.aclose() + await result.aclose() + await client.aclose() + + def test_httpx_raw_capture_failure_does_not_change_user_stream() -> None: malformed = b"not a gzip stream" @@ -638,6 +677,53 @@ async def test_native_openai_and_anthropic_sdks(endpoint_server: str) -> None: assert len(trajectory.exchanges.messages) == 1 +async def test_native_openai_chat_stream_captures_at_done_event() -> None: + async def response(_: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + stream=_AsyncChunks(_streaming_chat_body()), + ) + + http_client = httpx.AsyncClient(transport=httpx.MockTransport(response)) + client = AsyncOpenAI( + base_url="https://example.test/v1", + api_key="test", + http_client=http_client, + ) + with art.Trajectory() as trajectory: + stream = await client.chat.completions.create( + model="test/model", + messages=[], + stream=True, + ) + chunks = [chunk async for chunk in stream] + assert len(trajectory.exchanges.chat_completions) == 1 + + assert len(chunks) == 1 + assert chunks[0].choices[0].delta.content == "hello" + await stream.close() + await client.close() + + +def test_stream_terminal_without_sse_boundary_is_excluded() -> None: + body = _streaming_chat_body()[:-1] + with art.Trajectory() as trajectory: + state, token = begin( + "POST", + "https://example.test/v1/chat/completions", + {"model": "test/model", "messages": [], "stream": True}, + ) + reset(token) + assert state is not None + state.status_code = 200 + state.add(body) + assert not state.captured + state.finish() + + assert not trajectory.exchanges + + async def test_failed_and_incomplete_calls_are_excluded(endpoint_server: str) -> None: async with httpx.AsyncClient() as client: with art.Trajectory() as trajectory: @@ -986,19 +1072,33 @@ def test_all_streaming_protocols_reconstruct_final_responses() -> None: assert getattr(exchange.response, "logprobs") == [-0.2] assert getattr(exchange.response, "prompt_token_ids") == [1] - with art.Trajectory() as trajectory: - state, token = begin( - "POST", - "https://example.test/v1/messages", - {"model": "test/model", "messages": [], "stream": True}, + with art.Trajectory() as trajectory: + state, token = begin( + "POST", + f"https://example.test/v1/{endpoint.replace('_', '/')}", + request, + ) + reset(token) + assert state is not None + state.status_code = 200 + for byte in body[:-1]: + state.add(bytes([byte])) + assert not state.captured + state.add(body[-1:]) + assert state.captured + + assert ( + sum( + len(exchanges) + for exchanges in ( + trajectory.exchanges.chat_completions, + trajectory.exchanges.completions, + trajectory.exchanges.responses, + trajectory.exchanges.messages, + ) + ) + == 1 ) - reset(token) - assert state is not None - state.status_code = 200 - state.add(_sse(message_events)) - state.finish() - - assert len(trajectory.exchanges.messages) == 1 def test_streaming_messages_error_event_is_rejected_and_not_captured() -> None: From 6c388a7fe4baa86021cfcd5f7a3fdb909b30015a Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Fri, 24 Jul 2026 20:14:36 +0000 Subject: [PATCH 30/58] Keep automatic model selectors literal --- src/art/trajectories/_selection.py | 13 +++++-- src/art/trajectories/_tokenize.py | 35 ------------------- .../test_exchange_training_model_selection.py | 17 +++++++++ 3 files changed, 28 insertions(+), 37 deletions(-) diff --git a/src/art/trajectories/_selection.py b/src/art/trajectories/_selection.py index c7f4c5c3a..69ccc9d2b 100644 --- a/src/art/trajectories/_selection.py +++ b/src/art/trajectories/_selection.py @@ -15,6 +15,11 @@ class ModelSelector: value: str automatic_family: tuple[str, str] | None = None + allow_glob: bool = False + + def __post_init__(self) -> None: + if not self.value: + raise ValueError("A model selector cannot be empty") def matches(self, candidate: str) -> bool: if self.automatic_family is not None: @@ -25,11 +30,15 @@ def matches(self, candidate: str) -> bool: candidate, ) ) - return fnmatchcase(candidate, self.value) + return ( + fnmatchcase(candidate, self.value) + if self.allow_glob + else candidate == self.value + ) def public_model_selector(value: str) -> ModelSelector: - return ModelSelector(value) + return ModelSelector(value, allow_glob=True) def automatic_training_model_selector(value: str) -> ModelSelector: diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index 1682643be..0f115b5c5 100644 --- a/src/art/trajectories/_tokenize.py +++ b/src/art/trajectories/_tokenize.py @@ -837,41 +837,6 @@ def _exchange_list(trajectory: Trajectory, model: str | None) -> list[Exchange]: ) -def _require_training_model(trajectory: Trajectory, model: str | None) -> str: - """Resolve one exchange model without guessing which policy should train.""" - - models = { - exchange.model - for exchange in ( - *trajectory.exchanges.chat_completions, - *trajectory.exchanges.completions, - *trajectory.exchanges.responses, - *trajectory.exchanges.messages, - ) - } - if None in models: - raise ValueError("Every training exchange must identify its model") - if model is not None: - if not any(_model_matches(candidate, model) for candidate in models): - raise ValueError(f"Trajectory contains no exchanges for model {model!r}") - return model - if len(models) != 1: - raise ValueError( - "Exchange training requires exactly one model; pass model= to select one" - ) - return cast(str, next(iter(models))) - - -def _training_model_pattern(model: str) -> str: - """Select every immutable checkpoint belonging to one automatic policy.""" - - if re.search(r"@\d+$", model): - return re.sub(r"\d+$", "*", model) - if re.search(r":step\d+$", model): - return re.sub(r"\d+$", "*", model) - return model - - def _artifact_name(model: str) -> str: return model.removeprefix("wandb-artifact:///") diff --git a/tests/unit/test_exchange_training_model_selection.py b/tests/unit/test_exchange_training_model_selection.py index 6b9bbb932..f265d03ed 100644 --- a/tests/unit/test_exchange_training_model_selection.py +++ b/tests/unit/test_exchange_training_model_selection.py @@ -366,6 +366,23 @@ def test_automatic_training_model_selector_treats_family_metacharacters_literall assert not selector.matches("policy[blue]anything@13") +def test_automatic_training_model_selector_treats_non_family_metacharacters_literally() -> ( + None +): + selector = automatic_training_model_selector("policy*") + assert selector.matches("policy*") + assert not selector.matches("policy-judge") + + +@pytest.mark.parametrize("automatic", [False, True]) +def test_training_model_selector_rejects_empty_value(automatic: bool) -> None: + with pytest.raises(ValueError, match="cannot be empty"): + if automatic: + automatic_training_model_selector("") + else: + resolve_training_model(_group().trajectories[0], "") + + def test_automatic_training_selector_rejects_multiple_numeric_steps() -> None: trajectory = _versioned_group().trajectories[0] selector = automatic_training_model_selector("policy@12") From a7f0e287c431789143d1e8058bdba02411ad8ef1 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Fri, 24 Jul 2026 20:15:07 +0000 Subject: [PATCH 31/58] Preserve sources across overlength histories --- src/art/local/backend.py | 8 +- src/art/preprocessing/tokenize.py | 31 ++++ .../test_exchange_training_model_selection.py | 141 ++++++++++++++++++ 3 files changed, 179 insertions(+), 1 deletion(-) diff --git a/src/art/local/backend.py b/src/art/local/backend.py index 24638f364..9f93d4832 100644 --- a/src/art/local/backend.py +++ b/src/art/local/backend.py @@ -950,6 +950,12 @@ def _get_packed_tensors( chat_template_tool_schema_format = self._chat_template_tool_schema_format( internal_config ) + model_max_sequence_length = self._model_max_sequence_length(model) + training_max_sequence_length = ( + min(model_max_sequence_length, packed_sequence_length) + if packed_sequence_length is not None + else model_max_sequence_length + ) tokenized_results = list( tokenize_trajectory_groups( tokenizer, @@ -962,11 +968,11 @@ def _get_packed_tensors( model=automatic_training_model_selector( self._model_inference_name(model) ), + _max_sequence_length=training_max_sequence_length, ) ) if not tokenized_results: return None - model_max_sequence_length = self._model_max_sequence_length(model) too_long_for_model = [ result for result in tokenized_results diff --git a/src/art/preprocessing/tokenize.py b/src/art/preprocessing/tokenize.py index 853add2e9..cca1b2474 100644 --- a/src/art/preprocessing/tokenize.py +++ b/src/art/preprocessing/tokenize.py @@ -718,6 +718,7 @@ def tokenize_trajectory_groups( chat_template_kwargs: dict[str, Any] | None = None, chat_template_tool_schema_format: ChatTemplateToolSchemaFormat = "default", model: ModelSelector | str | None = None, + _max_sequence_length: int | None = None, ) -> Generator["TokenizedResult", None, None]: for group in trajectory_groups: if not group: @@ -759,6 +760,36 @@ def tokenize_trajectory_groups( for exchange_result, trace in zip( exchange_results.histories, traces, strict=True ): + if ( + _max_sequence_length is not None + and len(exchange_result.token_ids) > _max_sequence_length + ): + preview_seen = set(seen_source_keys) + would_train = _first_introduction_mask( + trace.source_keys, preview_seen + ) + if any(would_train): + trajectory_results.append( + TokenizedResult( + advantage=advantage, + chat="", + token_ids=exchange_result.token_ids, + input_pos=list( + range(len(exchange_result.token_ids)) + ), + assistant_mask=[0] * len(exchange_result.token_ids), + logprobs=exchange_result.logprobs, + pixel_values=None, + image_grid_thw=None, + trajectory=trajectory, + choice_offsets=[], + extra_logprobs={}, + moe_routed_experts=None, + moe_routing_alignment_stats=MoeRoutingAlignmentStats(), + _tokenizer=tokenizer, + ) + ) + continue trainable = _first_introduction_mask( trace.source_keys, seen_source_keys ) diff --git a/tests/unit/test_exchange_training_model_selection.py b/tests/unit/test_exchange_training_model_selection.py index f265d03ed..f402407bf 100644 --- a/tests/unit/test_exchange_training_model_selection.py +++ b/tests/unit/test_exchange_training_model_selection.py @@ -2,7 +2,10 @@ from collections.abc import Mapping from datetime import datetime, timedelta +from pathlib import Path +from types import SimpleNamespace from typing import SupportsIndex, cast, overload +from unittest.mock import patch import numpy as np from openai.types import Completion @@ -11,6 +14,9 @@ from transformers import PreTrainedTokenizerBase import art +from art import TrainableModel +from art.dev.model import InternalModelConfig +from art.local import LocalBackend from art.openai import ART_MOE_ROUTING_METADATA_KEY from art.preprocessing.moe_routing import MoeRouteSegments from art.preprocessing.tokenize import ( @@ -141,6 +147,75 @@ def _routed_exchange( return exchange +def _reasoning_stripped_group() -> art.TrajectoryGroup: + def set_choice( + exchange: ChatCompletionsExchange, + token_ids: list[int], + *, + content: str, + reasoning: str, + ) -> None: + data = exchange.response.model_dump(mode="python") + choice = data["choices"][0] + choice["message"] = { + "role": "assistant", + "content": content, + "reasoning": reasoning, + } + choice["token_ids"] = token_ids + choice["logprobs"]["content"] = [ + { + "token": f"token_id:{token_id}", + "logprob": -0.1, + "bytes": [], + "top_logprobs": [], + } + for token_id in token_ids + ] + exchange.response = ChatCompletion.model_validate(data) + extra = exchange.response.choices[0].model_extra + assert extra is not None + extra.pop(ART_MOE_ROUTING_METADATA_KEY, None) + + first = _routed_exchange( + prompt_token_ids=[1], + output_token=2, + messages=[{"role": "user", "content": "one"}], + content="first", + ) + set_choice( + first, + [2, 101, 102, 103, 104, 9], + content="first", + reasoning="long reasoning", + ) + second = _routed_exchange( + prompt_token_ids=[1, 9, 4], + output_token=5, + messages=[ + {"role": "user", "content": "one"}, + {"role": "assistant", "content": "first"}, + {"role": "user", "content": "two"}, + ], + content="second", + ) + set_choice( + second, + [5, 6], + content="second", + reasoning="short reasoning", + ) + return art.TrajectoryGroup( + [ + art.Trajectory( + exchanges=art.TrajectoryExchanges(chat_completions=[first, second]), + reward=reward, + ) + for reward in (1.0, 0.0) + ] + ) + + def _group() -> art.TrajectoryGroup: trajectories = [ art.Trajectory( @@ -302,6 +377,72 @@ def counted( assert calls == len(group.trajectories) +def test_overlength_history_does_not_claim_sources_from_fitting_history() -> None: + results = list( + tokenize_trajectory_groups( + cast(PreTrainedTokenizerBase, _Tokenizer()), + [_reasoning_stripped_group()], + allow_training_without_logprobs=False, + scale_rewards=False, + shuffle_group_trajectories=False, + drop_zero_advantage_trajectories=False, + model="policy", + _max_sequence_length=5, + ) + ) + + long = [result for result in results if len(result.token_ids) > 5] + fitting = [result for result in results if len(result.token_ids) <= 5] + assert len(long) == len(fitting) == 2 + assert all(result.assistant_mask == [0] * 7 for result in long) + assert all(result.token_ids == [1, 9, 4, 5, 6] for result in fitting) + assert all(result.assistant_mask == [0, 1, 0, 1, 1] for result in fitting) + assert all(result.weight == pytest.approx(1 / 3) for result in results) + + +def test_local_backend_trains_retained_source_after_overlength_history( + tmp_path: Path, +) -> None: + backend = LocalBackend(path=str(tmp_path)) + model = TrainableModel( + run_name="reasoning-stripped-overlength", + name="policy", + project="pipeline-tests", + base_model="test-model", + base_path=str(tmp_path), + _internal_config=InternalModelConfig(init_args={"max_seq_length": 5}), + ) + tokenizer = cast( + PreTrainedTokenizerBase, + SimpleNamespace( + name_or_path="test-model", + eos_token_id=0, + decode=lambda token_id: str(token_id), + ), + ) + backend._tokenizers[("test-model", None)] = tokenizer + backend._image_processors["test-model"] = None + + with ( + patch.object(backend, "_model_inference_name", return_value="policy"), + pytest.warns(UserWarning, match="Dropping 2 tokenized results"), + ): + packed = backend._get_packed_tensors( + model, + [_reasoning_stripped_group()], + advantage_balance=0.0, + allow_training_without_logprobs=False, + scale_rewards=False, + plot_tensors=False, + packed_sequence_length=5, + logprob_calculation_chunk_size=1, + ) + + assert packed is not None + assert packed["tokens"].tolist() == [[1, 9, 4, 5, 6]] * 2 + assert packed["assistant_mask"].tolist() == [[False, True, False, True, True]] * 2 + + def test_training_rejects_multiple_concrete_policy_versions() -> None: group = _versioned_group() with pytest.raises(ValueError, match="exactly one concrete model"): From 85fb0de92939e4e9765daf1685d591632fab889c Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Fri, 24 Jul 2026 20:15:49 +0000 Subject: [PATCH 32/58] Capture preloaded HTTPX streams --- src/art/trajectories/_capture/httpx.py | 23 +++++-- tests/unit/trajectories/test_capture.py | 85 ++++++++++++++++++++++++- 2 files changed, 101 insertions(+), 7 deletions(-) diff --git a/src/art/trajectories/_capture/httpx.py b/src/art/trajectories/_capture/httpx.py index c599b127e..01c2f7efc 100644 --- a/src/art/trajectories/_capture/httpx.py +++ b/src/art/trajectories/_capture/httpx.py @@ -29,6 +29,21 @@ def _capture_decoder(response: httpx.Response) -> ContentDecoder: return shadow._get_content_decoder() +def _capture_preloaded( + state: CaptureState, response: httpx.Response, *, stream: bool +) -> None: + if stream and not response.is_stream_consumed: + return + try: + content = response.content + except httpx.ResponseNotRead: + if not stream: + raise + return + state.add(content) + state.finish() + + def install() -> None: if getattr(httpx.Client.send, "_art_capture", False): return @@ -56,9 +71,7 @@ def send( if state is not None: state.status_code = response.status_code setattr(response, _STATE, state) - if not kwargs.get("stream", False): - state.add(response.content) - state.finish() + _capture_preloaded(state, response, stream=kwargs.get("stream", False)) return response async def async_send( @@ -78,9 +91,7 @@ async def async_send( if state is not None: state.status_code = response.status_code setattr(response, _STATE, state) - if not kwargs.get("stream", False): - state.add(response.content) - state.finish() + _capture_preloaded(state, response, stream=kwargs.get("stream", False)) return response def iter_raw( diff --git a/tests/unit/trajectories/test_capture.py b/tests/unit/trajectories/test_capture.py index aef8b1b9c..bf6430bc8 100644 --- a/tests/unit/trajectories/test_capture.py +++ b/tests/unit/trajectories/test_capture.py @@ -15,7 +15,7 @@ from anthropic import AsyncAnthropic from anthropic.types import TextBlock import httpx -from openai import AsyncOpenAI +from openai import AsyncOpenAI, OpenAI import pytest import pytest_asyncio import requests @@ -703,9 +703,92 @@ async def response(_: httpx.Request) -> httpx.Response: assert len(chunks) == 1 assert chunks[0].choices[0].delta.content == "hello" await stream.close() + assert len(trajectory.exchanges.chat_completions) == 1 + await client.close() + + +def test_native_openai_preloaded_chat_stream_captures_once() -> None: + def response(_: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content=_streaming_chat_body(), + ) + + http_client = httpx.Client(transport=httpx.MockTransport(response)) + client = OpenAI( + base_url="https://example.test/v1", + api_key="test", + http_client=http_client, + ) + with art.Trajectory() as trajectory: + stream = client.chat.completions.create( + model="test/model", + messages=[], + stream=True, + ) + assert len(trajectory.exchanges.chat_completions) == 1 + chunks = list(stream) + assert len(trajectory.exchanges.chat_completions) == 1 + + assert len(chunks) == 1 + assert chunks[0].choices[0].delta.content == "hello" + stream.close() + assert len(trajectory.exchanges.chat_completions) == 1 + client.close() + + +async def test_native_async_openai_preloaded_chat_stream_captures_once() -> None: + async def response(_: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content=_streaming_chat_body(), + ) + + http_client = httpx.AsyncClient(transport=httpx.MockTransport(response)) + client = AsyncOpenAI( + base_url="https://example.test/v1", + api_key="test", + http_client=http_client, + ) + with art.Trajectory() as trajectory: + stream = await client.chat.completions.create( + model="test/model", + messages=[], + stream=True, + ) + assert len(trajectory.exchanges.chat_completions) == 1 + chunks = [chunk async for chunk in stream] + assert len(trajectory.exchanges.chat_completions) == 1 + + assert len(chunks) == 1 + assert chunks[0].choices[0].delta.content == "hello" + await stream.close() + assert len(trajectory.exchanges.chat_completions) == 1 await client.close() +def test_preloaded_malformed_stream_is_excluded() -> None: + def response(_: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content=_sse([(None, "corrupt"), (None, "[DONE]")]), + ) + + with art.Trajectory() as trajectory: + with httpx.Client(transport=httpx.MockTransport(response)) as client: + with client.stream( + "POST", + "https://example.test/v1/chat/completions", + json={"model": "test/model", "messages": [], "stream": True}, + ) as result: + assert list(result.iter_bytes()) + + assert not trajectory.exchanges + + def test_stream_terminal_without_sse_boundary_is_excluded() -> None: body = _streaming_chat_body()[:-1] with art.Trajectory() as trajectory: From 8c51ff2e5c7b5be0b57a20073749cd5561c4b510 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Fri, 24 Jul 2026 20:18:25 +0000 Subject: [PATCH 33/58] Preserve exact multi-generation chat evidence --- src/art/trajectories/_tokenize.py | 15 +++- tests/unit/trajectories/test_tokenize.py | 109 +++++++++++++++++++++++ 2 files changed, 121 insertions(+), 3 deletions(-) diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index 0f115b5c5..863c040ed 100644 --- a/src/art/trajectories/_tokenize.py +++ b/src/art/trajectories/_tokenize.py @@ -279,8 +279,8 @@ def _sampled_source_key(source: object) -> _SampledSourceKey: return _source_key(exchange, protocol="chat_completions", index=index) if isinstance(exchange, ResponsesExchange): index = getattr(source, "generation_index", None) - if index is None and getattr(source, "output_index", None) is not None: - index = 0 + if index is None and not _response_generations(exchange.response): + index = getattr(source, "output_index", None) if not isinstance(index, int) or isinstance(index, bool): raise ValueError("Sampled Responses source has no generation identity") return _source_key(exchange, protocol="responses", index=index) @@ -3351,10 +3351,19 @@ def part_ids(text: str) -> list[int]: sampled_part_cursor += len(parts) continue + source_exchange = getattr(source, "exchange", None) + multi_generation_response = ( + isinstance(source_exchange, ResponsesExchange) + and len(_response_generations(source_exchange.response)) > 1 + ) if ( complete_sampled_message and full_exact is not None - and (len(parts) != 1 or len(full_exact) != len(part_ids(parts[0][1]))) + and ( + multi_generation_response + or len(parts) != 1 + or len(full_exact) != len(part_ids(parts[0][1])) + ) ): prompt_render = render(messages[:message_index], add_generation_prompt=True) completed_render = render( diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index affcf793a..7d51cb051 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -20,6 +20,7 @@ import art from art.trajectories import ( ChatCompletionsExchange, + ChatCompletionsMessageSource, ChatCompletionsRequest, CompletionsExchange, CompletionsRequest, @@ -2595,6 +2596,114 @@ def test_responses_generation_source_rejects_content_from_another_generation() - history.tokenize() +def test_responses_chat_rerender_preserves_equal_length_generation_evidence() -> None: + exchange = _response_exchange("equal-length-generations", 2) + data = exchange.response.model_dump(mode="python") + data["output"].append( + { + "id": "message-second-generation", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [ + { + "type": "output_text", + "text": "second", + "annotations": [], + "logprobs": [], + } + ], + } + ) + data["token_generations"] = [ + { + "prompt_token_ids": [1], + "output_tokens": [{"token_id": 2, "logprob": -0.2}], + "output_indices": [0], + }, + { + "prompt_token_ids": [1, 2, 3], + "output_tokens": [{"token_id": 4, "logprob": -0.4}], + "output_indices": [1], + }, + ] + exchange.response = Response.model_validate(data) + history = ( + art.Trajectory(exchanges=TrajectoryExchanges(responses=[exchange])) + .responses_history() + .as_chat_completions_history() + ) + + class Tokenizer: + def __call__(self, text: str, **kwargs: object) -> list[int]: + del kwargs + return {"turn 0": [10], "answer": [20], "second": [40]}[text] + + def apply_chat_template( + self, + messages: list[dict[str, Any]], + *, + add_generation_prompt: bool, + **kwargs: object, + ) -> list[int]: + del kwargs + result: list[int] = [] + for index, message in enumerate(messages): + result.extend(self(str(message["content"]))) + if index == 1 and (len(messages) > 2 or add_generation_prompt): + result.append(30) + return result + + tokenized = history.tokenize(tokenizer=Tokenizer(), chat_template="custom") + + assert tokenized.token_ids == [10, 2, 30, 4] + assert 20 not in tokenized.token_ids + assert 40 not in tokenized.token_ids + assert tokenized.logprobs[1::2] == pytest.approx([-0.2, -0.4]) + assert tokenized.flags == [ + art.TokenFlag(0), + art.TokenFlag.EXACT | art.TokenFlag.SAMPLED, + art.TokenFlag(0), + art.TokenFlag.EXACT | art.TokenFlag.SAMPLED, + ] + + +def test_responses_chat_sampled_source_requires_generation_identity() -> None: + exchange = _response_exchange("missing-generation-identity", 2) + history = ( + art.Trajectory(exchanges=TrajectoryExchanges(responses=[exchange])) + .responses_history() + .as_chat_completions_history() + ) + assistant_index = next( + index + for index, source in enumerate(history.message_sources) + if source is not None and source.output_index is not None + ) + source = history.message_sources[assistant_index] + assert source is not None + history.message_sources[assistant_index] = ChatCompletionsMessageSource( + exchange=source.exchange, + output_index=source.output_index, + ) + + class Tokenizer: + def __call__(self, text: str, **kwargs: object) -> list[int]: + del kwargs + return {"turn 0": [1], "answer": [2]}[text] + + def apply_chat_template( + self, messages: list[dict[str, Any]], **kwargs: object + ) -> list[int]: + del kwargs + return [ + token for message in messages for token in self(str(message["content"])) + ] + + with pytest.raises(ValueError, match="no generation identity"): + history.tokenize(tokenizer=Tokenizer()) + + def _multi_output_responses_chat_history() -> art.ChatCompletionsHistory: exchange = _response_exchange("multi-output-generation", 2) data = exchange.response.model_dump(mode="python") From cbd27e94cd5c5cfe472d73dba68aaaa3ebaac7ef Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Fri, 24 Jul 2026 20:24:43 +0000 Subject: [PATCH 34/58] Tighten remediation test typing --- .../runtime_isolation/test_runtime_project_isolation.py | 2 +- tests/unit/trajectories/test_tokenize.py | 7 +++++++ 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/tests/integration/megatron/runtime_isolation/test_runtime_project_isolation.py b/tests/integration/megatron/runtime_isolation/test_runtime_project_isolation.py index 6a82c36cf..4eddbb176 100644 --- a/tests/integration/megatron/runtime_isolation/test_runtime_project_isolation.py +++ b/tests/integration/megatron/runtime_isolation/test_runtime_project_isolation.py @@ -53,7 +53,7 @@ def test_runtime_server_source_contains_only_required_custom_routes() -> None: assert route in source -def test_runtime_patch_always_returns_token_ids( +def test_runtime_patch_defaults_evidence_on_and_honors_opt_out( artifact_dir: Path, ) -> None: payload = _runtime_python( diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index 7d51cb051..5ee0df714 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -553,6 +553,13 @@ def __call__(self, text: str, **kwargs: object) -> list[int]: del kwargs return {"question": [1], "questionanswer": [1, 2]}[text] + def apply_chat_template( + self, messages: list[dict[str, object]], **kwargs: object + ) -> list[int]: + raise AssertionError( + "Completions tokenization must not render chat messages" + ) + tokenized = art.Trajectory( exchanges=TrajectoryExchanges(completions=[exchange]) ).tokenize(tokenizer=Tokenizer()) From 76aaa2ca58636e3cd61b570285a2eccee70fef9d Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Fri, 24 Jul 2026 20:39:53 +0000 Subject: [PATCH 35/58] Unify Responses fallback trace identity --- src/art/trajectories/_tokenize.py | 2 +- tests/unit/trajectories/test_tokenize.py | 103 +++++++++++++++++++++++ 2 files changed, 104 insertions(+), 1 deletion(-) diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index 863c040ed..ce3924f8c 100644 --- a/src/art/trajectories/_tokenize.py +++ b/src/art/trajectories/_tokenize.py @@ -280,7 +280,7 @@ def _sampled_source_key(source: object) -> _SampledSourceKey: if isinstance(exchange, ResponsesExchange): index = getattr(source, "generation_index", None) if index is None and not _response_generations(exchange.response): - index = getattr(source, "output_index", None) + index = 0 if not isinstance(index, int) or isinstance(index, bool): raise ValueError("Sampled Responses source has no generation identity") return _source_key(exchange, protocol="responses", index=index) diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index 5ee0df714..037495ffd 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -3698,6 +3698,109 @@ def test_private_trace_keys_do_not_collide_for_repeated_empty_response_ids() -> assert len(trace.sources) == 2 +def test_responses_fallback_trace_does_not_retrain_echoed_output_items() -> None: + from art.trajectories._tokenize import ( + _first_introduction_mask, + _tokenize_trajectory_with_trace, + ) + + def output(item_id: str, text: str) -> dict[str, Any]: + return { + "id": item_id, + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": text, "annotations": []}], + } + + def response( + response_id: str, items: list[dict[str, Any]], offset: int + ) -> Response: + return Response.model_validate( + { + "id": response_id, + "created_at": float(offset), + "model": "test/model", + "object": "response", + "output": items, + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [], + } + ) + + user = {"role": "user", "content": "question"} + first_response = response( + "response-fallback", + [output("one", "first"), output("two", "second")], + 0, + ) + first_request = ResponsesRequest(model="test/model", input=[user]) + first_request["chat_template_kwargs"] = {"enable_thinking": False} + first = ResponsesExchange( + request=first_request, + response=first_response, + start_time=datetime(2026, 1, 1), + end_time=datetime(2026, 1, 1, 0, 0, 0, 1000), + ) + echoed = [ + item.model_dump(mode="json", exclude_none=True) + for item in first_response.output + ] + second = ResponsesExchange( + request=ResponsesRequest( + model="test/model", + input=[user, *echoed, {"role": "user", "content": "continue"}], + ), + response=response("response-final", [output("final", "final")], 1), + start_time=datetime(2026, 1, 1, 0, 0, 1), + end_time=datetime(2026, 1, 1, 0, 0, 1, 1000), + ) + + class Tokenizer: + def __call__(self, text: str, **kwargs: object) -> list[int]: + del kwargs + return { + "question": [1], + "first": [2], + "second": [3], + "firstsecond": [2, 3], + "continue": [4], + "final": [5], + }[text] + + def apply_chat_template( + self, messages: list[dict[str, Any]], **kwargs: object + ) -> list[int]: + del kwargs + return [ + token for message in messages for token in self(str(message["content"])) + ] + + trajectory = art.Trajectory( + exchanges=TrajectoryExchanges(responses=[first, second]) + ) + tokenized, traces = _tokenize_trajectory_with_trace( + trajectory, + tokenizer=Tokenizer(), + chat_template_kwargs={"enable_thinking": False}, + ) + + assert len(tokenized.histories) == len(traces) == 2 + seen: set[object] = set() + trained_first_response = 0 + first_response_indices: set[int] = set() + for trace in traces: + trainable = _first_introduction_mask(trace.source_keys, seen) + for selected, key in zip(trainable, trace.source_keys, strict=True): + if key is not None and key.response_id == first_response.id: + first_response_indices.add(key.index) + trained_first_response += selected + + assert first_response_indices == {0} + assert trained_first_response == 2 + + def test_completions_history_requires_exhaustive_source_spans() -> None: history = art.CompletionsTokenHistory( model="test/model", From b01e2a8c2ce57ea4f7a4dcd6f8676d9ae2175f64 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Fri, 24 Jul 2026 20:43:28 +0000 Subject: [PATCH 36/58] Type Responses fallback regression --- tests/unit/trajectories/test_tokenize.py | 29 ++++++++++++++++-------- 1 file changed, 19 insertions(+), 10 deletions(-) diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index 037495ffd..9e553625a 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -14,7 +14,12 @@ from openai.types import Completion from openai.types.chat import ChatCompletion, ChatCompletionMessageParam from openai.types.chat.chat_completion_token_logprob import ChatCompletionTokenLogprob -from openai.types.responses import Response +from openai.types.responses import ( + EasyInputMessageParam, + Response, + ResponseInputParam, + ResponseOutputMessageParam, +) import pytest import art @@ -3704,7 +3709,7 @@ def test_responses_fallback_trace_does_not_retrain_echoed_output_items() -> None _tokenize_trajectory_with_trace, ) - def output(item_id: str, text: str) -> dict[str, Any]: + def output(item_id: str, text: str) -> ResponseOutputMessageParam: return { "id": item_id, "type": "message", @@ -3714,7 +3719,7 @@ def output(item_id: str, text: str) -> dict[str, Any]: } def response( - response_id: str, items: list[dict[str, Any]], offset: int + response_id: str, items: list[ResponseOutputMessageParam], offset: int ) -> Response: return Response.model_validate( { @@ -3729,13 +3734,15 @@ def response( } ) - user = {"role": "user", "content": "question"} + user = EasyInputMessageParam(role="user", content="question") + first_outputs = [output("one", "first"), output("two", "second")] first_response = response( "response-fallback", - [output("one", "first"), output("two", "second")], + first_outputs, 0, ) - first_request = ResponsesRequest(model="test/model", input=[user]) + first_input: ResponseInputParam = [user] + first_request = ResponsesRequest(model="test/model", input=first_input) first_request["chat_template_kwargs"] = {"enable_thinking": False} first = ResponsesExchange( request=first_request, @@ -3743,14 +3750,16 @@ def response( start_time=datetime(2026, 1, 1), end_time=datetime(2026, 1, 1, 0, 0, 0, 1000), ) - echoed = [ - item.model_dump(mode="json", exclude_none=True) - for item in first_response.output + echoed: ResponseInputParam = [*first_outputs] + second_input: ResponseInputParam = [ + user, + *echoed, + EasyInputMessageParam(role="user", content="continue"), ] second = ResponsesExchange( request=ResponsesRequest( model="test/model", - input=[user, *echoed, {"role": "user", "content": "continue"}], + input=second_input, ), response=response("response-final", [output("final", "final")], 1), start_time=datetime(2026, 1, 1, 0, 0, 1), From 86dc78dd1acf21cb68dce922de631a679a674cf9 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Fri, 24 Jul 2026 21:10:47 +0000 Subject: [PATCH 37/58] Fix trajectory CI contract checks --- pyproject.toml | 2 +- tests/unit/test_pipeline_trainer_local_backend.py | 5 ++++- uv.lock | 2 +- 3 files changed, 6 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 20858ccc3..9be997359 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,7 +26,7 @@ plotting = ["matplotlib>=3.10.1", "seaborn>=0.13.2"] backend = [ "peft>=0.14.0", "hf-xet>=1.1.0", - "bitsandbytes>=0.45.2", + "bitsandbytes>=0.45.2,!=0.50.0", "unsloth==2026.3.3", "unsloth-zoo==2026.3.1", "torch==2.11.0", diff --git a/tests/unit/test_pipeline_trainer_local_backend.py b/tests/unit/test_pipeline_trainer_local_backend.py index 340217b16..e347f22e5 100644 --- a/tests/unit/test_pipeline_trainer_local_backend.py +++ b/tests/unit/test_pipeline_trainer_local_backend.py @@ -769,7 +769,10 @@ def test_local_backend_get_packed_tensors_warns_and_drops_overlong_results( assert packed_tensors is not None assert packed_tensors["tokens"].shape == (1, 4) - assert tokenize.call_args.kwargs["model"] == f"{model.name}@*" + selector = tokenize.call_args.kwargs["model"] + assert selector.value == f"{model.name}@0" + assert selector.automatic_family == (model.name, "@") + assert selector.allow_glob is False @pytest.mark.asyncio diff --git a/uv.lock b/uv.lock index 4709e1923..cf4aae62d 100644 --- a/uv.lock +++ b/uv.lock @@ -4882,7 +4882,7 @@ requires-dist = [ { name = "anthropic", specifier = ">=0.77.0" }, { name = "apex", marker = "extra == 'megatron'", git = "https://github.com/NVIDIA/apex.git?rev=25.09" }, { name = "awscli", marker = "extra == 'backend'", specifier = ">=1.38.1" }, - { name = "bitsandbytes", marker = "extra == 'backend'", specifier = ">=0.45.2" }, + { name = "bitsandbytes", marker = "extra == 'backend'", specifier = ">=0.45.2,!=0.50.0" }, { name = "causal-conv1d", marker = "python_full_version < '3.12' and platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'megatron'", specifier = "==1.6.1" }, { name = "datrie", marker = "extra == 'tinker'", specifier = ">=0.8.3" }, { name = "duckdb", marker = "extra == 'backend'", specifier = ">=1.0.0" }, From 05f21dfae91a8956577da33270e6522e5ef49f53 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Fri, 24 Jul 2026 22:30:36 +0000 Subject: [PATCH 38/58] refactor trajectories to reduce duplication --- src/art/trajectories/_tokenize.py | 70 ++----- tests/unit/trajectories/test_capture.py | 44 +---- tests/unit/trajectories/test_tokenize.py | 239 +++++++---------------- 3 files changed, 95 insertions(+), 258 deletions(-) diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index ce3924f8c..03c3813e8 100644 --- a/src/art/trajectories/_tokenize.py +++ b/src/art/trajectories/_tokenize.py @@ -4028,20 +4028,6 @@ def tokenize_trajectory( raise ValueError( f"Trajectory tokenization requires exactly one history; found {len(histories)}" ) - return _materialize_trajectory( - tokenize_history( - histories[0], - model=model - if isinstance(histories[0], LegacyHistory) - else histories[0].model, - base_model=base_model, - tokenizer=tokenizer, - chat_template=chat_template, - chat_template_kwargs=chat_template_kwargs, - _projection_validated=not isinstance(histories[0], LegacyHistory), - ), - trajectory, - ) tokenized = [ tokenize_history( history, @@ -4054,6 +4040,8 @@ def tokenize_trajectory( ) for history in histories ] + if not multi_history: + return _materialize_trajectory(tokenized[0], trajectory) return TokenizedMultiHistoryTrajectory( histories=tokenized, reward=trajectory.reward, @@ -4123,44 +4111,24 @@ def tokenize_group( TokenizedTrajectoryGroup[TokenizedTrajectory] | TokenizedTrajectoryGroup[TokenizedMultiHistoryTrajectory] ): - if multi_history: - trajectories = [ - cast( - TokenizedMultiHistoryTrajectory, - tokenize_trajectory( - trajectory, - multi_history=True, - model=model, - base_model=base_model, - tokenizer=tokenizer, - chat_template=chat_template, - chat_template_kwargs=chat_template_kwargs, - ), - ) - for trajectory in group.trajectories - ] - return TokenizedTrajectoryGroup[TokenizedMultiHistoryTrajectory]( - trajectories=trajectories, - metrics=dict(group.metrics), - metadata=dict(group.metadata), - ) - single_trajectories = [ - cast( - TokenizedTrajectory, - tokenize_trajectory( - trajectory, - multi_history=False, - model=model, - base_model=base_model, - tokenizer=tokenizer, - chat_template=chat_template, - chat_template_kwargs=chat_template_kwargs, - ), + trajectories = [ + tokenize_trajectory( + trajectory, + multi_history=multi_history, + model=model, + base_model=base_model, + tokenizer=tokenizer, + chat_template=chat_template, + chat_template_kwargs=chat_template_kwargs, ) for trajectory in group.trajectories ] - return TokenizedTrajectoryGroup[TokenizedTrajectory]( - trajectories=single_trajectories, - metrics=dict(group.metrics), - metadata=dict(group.metadata), + return cast( + TokenizedTrajectoryGroup[TokenizedTrajectory] + | TokenizedTrajectoryGroup[TokenizedMultiHistoryTrajectory], + TokenizedTrajectoryGroup( + trajectories=trajectories, + metrics=dict(group.metrics), + metadata=dict(group.metadata), + ), ) diff --git a/tests/unit/trajectories/test_capture.py b/tests/unit/trajectories/test_capture.py index bf6430bc8..4bd81758f 100644 --- a/tests/unit/trajectories/test_capture.py +++ b/tests/unit/trajectories/test_capture.py @@ -344,27 +344,7 @@ def requests_stream() -> None: def test_httpx_sync_stream_consumption_captures_decoded_body_once( encoding: str, mode: str ) -> None: - body = _sse( - [ - ( - None, - { - "id": "chatcmpl-1", - "object": "chat.completion.chunk", - "created": 1, - "model": "test/model", - "choices": [ - { - "index": 0, - "delta": {"role": "assistant", "content": "hello"}, - "finish_reason": "stop", - } - ], - }, - ), - (None, "[DONE]"), - ] - ) + body = _streaming_chat_body() compressed = _encoded(body, encoding) def response(_: httpx.Request) -> httpx.Response: @@ -396,27 +376,7 @@ def response(_: httpx.Request) -> httpx.Response: async def test_httpx_async_stream_consumption_captures_decoded_body_once( encoding: str, mode: str ) -> None: - body = _sse( - [ - ( - None, - { - "id": "chatcmpl-1", - "object": "chat.completion.chunk", - "created": 1, - "model": "test/model", - "choices": [ - { - "index": 0, - "delta": {"role": "assistant", "content": "hello"}, - "finish_reason": "stop", - } - ], - }, - ), - (None, "[DONE]"), - ] - ) + body = _streaming_chat_body() compressed = _encoded(body, encoding) async def response(_: httpx.Request) -> httpx.Response: diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index 9e553625a..f885a2186 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -129,6 +129,38 @@ def _completion_exchange( ) +def _message_exchange( + request: MessagesRequest, + *, + identifier: str = "message-1", + content: list[dict[str, object]] | None = None, + duration: timedelta = timedelta(milliseconds=1), + offset: int = 0, + response_model: str = "test/model", + **response_extra: object, +) -> MessagesExchange: + start = datetime(2026, 1, 1) + timedelta(seconds=offset) + response = Message.model_validate( + { + "id": identifier, + "type": "message", + "role": "assistant", + "model": response_model, + "content": content or [{"type": "text", "text": "answer"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 0, "output_tokens": 0}, + **response_extra, + } + ) + return MessagesExchange( + request=request, + response=response, + start_time=start, + end_time=start + duration, + ) + + def test_exact_tokens_form_one_append_only_history_without_tokenizer( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -187,34 +219,19 @@ def test_empty_tool_calls_normalization_preserves_exact_continuation() -> None: def test_messages_exact_prompt_and_output_do_not_load_a_tokenizer( monkeypatch: pytest.MonkeyPatch, ) -> None: - response = Message.model_validate( - { - "id": "message-exact", - "type": "message", - "role": "assistant", - "model": "test/model", - "content": [{"type": "text", "text": "answer"}], - "stop_reason": "end_turn", - "stop_sequence": None, - "usage": {"input_tokens": 2, "output_tokens": 1}, - "prompt_token_ids": [1, 2], - "token_ids": [3], - "logprobs": [-0.3], - } - ) - start = datetime(2026, 1, 1) trajectory = art.Trajectory( exchanges=TrajectoryExchanges( messages=[ - MessagesExchange( - request=MessagesRequest( + _message_exchange( + MessagesRequest( model="test/model", messages=[{"role": "user", "content": "question"}], max_tokens=16, ), - response=response, - start_time=start, - end_time=start + timedelta(milliseconds=1), + identifier="message-exact", + prompt_token_ids=[1, 2], + token_ids=[3], + logprobs=[-0.3], ) ] ) @@ -237,8 +254,6 @@ def test_messages_exact_prompt_and_output_do_not_load_a_tokenizer( def test_anthropic_cache_control_change_starts_new_source_lineage() -> None: - start = datetime(2026, 1, 1) - def exchange( *, identifier: str, @@ -247,34 +262,18 @@ def exchange( output_token_id: int, offset: int, ) -> MessagesExchange: - response = Message.model_validate( - { - "id": identifier, - "type": "message", - "role": "assistant", - "model": "test/model", - "content": [{"type": "text", "text": f"answer {offset}"}], - "stop_reason": "end_turn", - "stop_sequence": None, - "usage": { - "input_tokens": len(prompt_token_ids), - "output_tokens": 1, - }, - "prompt_token_ids": prompt_token_ids, - "token_ids": [output_token_id], - "logprobs": [-0.1], - } - ) - timestamp = start + timedelta(seconds=offset) - return MessagesExchange( - request=MessagesRequest( + return _message_exchange( + MessagesRequest( model="test/model", messages=messages, max_tokens=16, ), - response=response, - start_time=timestamp, - end_time=timestamp + timedelta(milliseconds=1), + identifier=identifier, + content=[{"type": "text", "text": f"answer {offset}"}], + offset=offset, + prompt_token_ids=prompt_token_ids, + token_ids=[output_token_id], + logprobs=[-0.1], ) first = exchange( @@ -324,32 +323,17 @@ def exchange( def test_converted_anthropic_system_history_preserves_exact_assistant_evidence( monkeypatch: pytest.MonkeyPatch, ) -> None: - response = Message.model_validate( - { - "id": "message-system-exact", - "type": "message", - "role": "assistant", - "model": "test/model", - "content": [{"type": "text", "text": "answer"}], - "stop_reason": "end_turn", - "stop_sequence": None, - "usage": {"input_tokens": 2, "output_tokens": 1}, - "prompt_token_ids": [10, 11], - "token_ids": [12], - "logprobs": [-0.12], - } - ) - start = datetime(2026, 1, 1) - exchange = MessagesExchange( - request=MessagesRequest( + exchange = _message_exchange( + MessagesRequest( model="test/model", system="system", messages=[{"role": "user", "content": "question"}], max_tokens=16, ), - response=response, - start_time=start, - end_time=start + timedelta(milliseconds=1), + identifier="message-system-exact", + prompt_token_ids=[10, 11], + token_ids=[12], + logprobs=[-0.12], ) trajectory = art.Trajectory(exchanges=TrajectoryExchanges(messages=[exchange])) monkeypatch.setattr( @@ -370,32 +354,17 @@ def test_converted_anthropic_system_history_preserves_exact_assistant_evidence( def test_converted_anthropic_system_source_rejects_sampled_response_mutation() -> None: - response = Message.model_validate( - { - "id": "message-system-source", - "type": "message", - "role": "assistant", - "model": "test/model", - "content": [{"type": "text", "text": "answer"}], - "stop_reason": "end_turn", - "stop_sequence": None, - "usage": {"input_tokens": 2, "output_tokens": 1}, - "prompt_token_ids": [10, 11], - "token_ids": [12], - "logprobs": [-0.12], - } - ) - start = datetime(2026, 1, 1) - exchange = MessagesExchange( - request=MessagesRequest( + exchange = _message_exchange( + MessagesRequest( model="test/model", system="system", messages=[{"role": "user", "content": "question"}], max_tokens=16, ), - response=response, - start_time=start, - end_time=start + timedelta(milliseconds=1), + identifier="message-system-source", + prompt_token_ids=[10, 11], + token_ids=[12], + logprobs=[-0.12], ) history = ( art.Trajectory(exchanges=TrajectoryExchanges(messages=[exchange])) @@ -448,29 +417,14 @@ def test_malformed_explicit_exact_token_metadata_fails_closed() -> None: assert response_extra is not None response_extra["token_generations"][0]["output_tokens"] = [{"token_id": "invalid"}] - message_response = Message.model_validate( - { - "id": "message-invalid", - "type": "message", - "role": "assistant", - "model": "test/model", - "content": [{"type": "text", "text": "answer"}], - "stop_reason": "end_turn", - "stop_sequence": None, - "usage": {"input_tokens": 1, "output_tokens": 1}, - "token_ids": [2, "invalid"], - } - ) - start = datetime(2026, 1, 1) - message = MessagesExchange( - request=MessagesRequest( + message = _message_exchange( + MessagesRequest( model="test/model", messages=[{"role": "user", "content": "question"}], max_tokens=16, ), - response=message_response, - start_time=start, - end_time=start + timedelta(milliseconds=1), + identifier="message-invalid", + token_ids=[2, "invalid"], ) trajectories = [ @@ -1071,30 +1025,15 @@ def apply_chat_template( def test_fallback_uses_template_overrides_and_nan_logprobs( monkeypatch: pytest.MonkeyPatch, ) -> None: - response = Message.model_validate( - { - "id": "msg_1", - "type": "message", - "role": "assistant", - "model": "test/model", - "content": [{"type": "text", "text": "answer"}], - "stop_reason": "end_turn", - "stop_sequence": None, - "usage": {"input_tokens": 1, "output_tokens": 1}, - } - ) - start = datetime(2026, 1, 1) - exchange = MessagesExchange( - request=MessagesRequest( + exchange = _message_exchange( + MessagesRequest( model="wandb-artifact:///entity/project/run:step0", messages=[{"role": "user", "content": "question"}], chat_template="request-template", chat_template_kwargs={"request": True}, thinking={"type": "enabled", "budget_tokens": 128}, ), - response=response, - start_time=start, - end_time=start + timedelta(seconds=1), + duration=timedelta(seconds=1), ) tokenizer = _FakeTokenizer() loaded_base_models: list[str] = [] @@ -1648,19 +1587,6 @@ def test_legacy_token_and_logprob_length_mismatch_raises() -> None: def test_anthropic_fallback_rejects_unknown_content_blocks( monkeypatch: pytest.MonkeyPatch, ) -> None: - response = Message.model_validate( - { - "id": "msg_1", - "type": "message", - "role": "assistant", - "model": "test/model", - "content": [{"type": "text", "text": "answer"}], - "stop_reason": "end_turn", - "stop_sequence": None, - "usage": {"input_tokens": 1, "output_tokens": 1}, - } - ) - start = datetime(2026, 1, 1) image: ImageBlockParam = { "type": "image", "source": { @@ -1670,14 +1596,12 @@ def test_anthropic_fallback_rejects_unknown_content_blocks( }, } message: MessageParam = {"role": "user", "content": [image]} - exchange = MessagesExchange( - request=MessagesRequest( + exchange = _message_exchange( + MessagesRequest( model="test/model", messages=[message], ), - response=response, - start_time=start, - end_time=start + timedelta(seconds=1), + duration=timedelta(seconds=1), ) monkeypatch.setattr( "art.trajectories._tokenize._load_tokenizer", lambda _config: _FakeTokenizer() @@ -3624,22 +3548,6 @@ def test_tokenized_results_materialize_metadata_and_group_shape() -> None: def test_private_trace_covers_sampled_tokens_for_every_protocol() -> None: from art.trajectories._tokenize import _tokenize_trajectory_with_trace - message_response = Message.model_validate( - { - "id": "message-trace", - "type": "message", - "role": "assistant", - "model": "test/model", - "content": [{"type": "text", "text": "answer"}], - "stop_reason": "end_turn", - "stop_sequence": None, - "usage": {"input_tokens": 1, "output_tokens": 1}, - "prompt_token_ids": [1], - "token_ids": [2], - "logprobs": [-0.2], - } - ) - start = datetime(2026, 1, 1) trajectories = [ art.Trajectory( exchanges=TrajectoryExchanges(chat_completions=[_chat_exchange([1], [2])]) @@ -3655,15 +3563,16 @@ def test_private_trace_covers_sampled_tokens_for_every_protocol() -> None: art.Trajectory( exchanges=TrajectoryExchanges( messages=[ - MessagesExchange( - request=MessagesRequest( + _message_exchange( + MessagesRequest( model="test/model", messages=[{"role": "user", "content": "question"}], max_tokens=16, ), - response=message_response, - start_time=start, - end_time=start + timedelta(milliseconds=1), + identifier="message-trace", + prompt_token_ids=[1], + token_ids=[2], + logprobs=[-0.2], ) ] ) From 1b437854422b7be76014125a66850cb8233dbf60 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Sat, 25 Jul 2026 00:04:41 +0000 Subject: [PATCH 39/58] change chat source to explicit output indices --- src/art/trajectories/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/art/trajectories/__init__.py b/src/art/trajectories/__init__.py index 54a6b0de5..0e1068cf9 100644 --- a/src/art/trajectories/__init__.py +++ b/src/art/trajectories/__init__.py @@ -322,7 +322,7 @@ class ChatCompletionsMessageSource: exchange: ChatCompletionsExchange | MessagesExchange | ResponsesExchange request_index: int | None = None choice_index: int | None = None - output_index: int | None = None + output_indices: tuple[int, ...] | None = None generation_index: int | None = None From 110c19d9cf1d446a2f59cb51d42478d37f1b5f51 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Sat, 25 Jul 2026 00:13:49 +0000 Subject: [PATCH 40/58] Fix trajectory history projection provenance --- src/art/trajectories/_history.py | 141 +++++++++++------ tests/unit/trajectories/test_history.py | 192 +++++++++++++++++++++++- 2 files changed, 285 insertions(+), 48 deletions(-) diff --git a/src/art/trajectories/_history.py b/src/art/trajectories/_history.py index bf8ac1467..348d77590 100644 --- a/src/art/trajectories/_history.py +++ b/src/art/trajectories/_history.py @@ -358,13 +358,7 @@ def _chat_retains_sampled_reasoning( def _chat_message_key(message: Message, *, visible_only: bool = False) -> str: - # History messages are already normalized and detached from the exchange. - # Copy only the top-level mapping because keying never mutates nested values. - data = dict(message) - if data.get("content") is None: - data.pop("content", None) - if data.get("tool_calls") == []: - data.pop("tool_calls") + data = normalize_chat_message(message) if visible_only: data.pop("reasoning", None) data.pop("reasoning_content", None) @@ -375,7 +369,10 @@ def normalize_chat_message(message: Mapping[str, object]) -> dict[str, object]: """Return the canonical history form of an OpenAI chat message.""" data = copy.deepcopy(dict(message)) - if data.get("content") is None: + data.pop("annotations", None) + if data.get("role") == "assistant" and data.get("content") is None: + data["content"] = "" + elif data.get("content") is None: data.pop("content", None) if data.get("tool_calls") == []: data.pop("tool_calls") @@ -1149,6 +1146,23 @@ def _completion_choice_groups( prompts = _completion_prompts(exchange.request.get("prompt")) count = len(prompts) choices = _ordered_choices(exchange.response.choices, protocol="Completions") + missing_prompt_index = object() + explicit_matches: dict[int, int] = {} + for choice in choices: + raw_prompt_index = (choice.model_extra or {}).get( + "prompt_index", missing_prompt_index + ) + if raw_prompt_index is missing_prompt_index: + continue + if ( + not isinstance(raw_prompt_index, int) + or isinstance(raw_prompt_index, bool) + or not 0 <= raw_prompt_index < count + ): + raise ValueError( + "Completions choice prompt_index must identify a batched prompt" + ) + explicit_matches[choice.index] = raw_prompt_index if count == 1: return [choices] requested_n = exchange.request.get("n") @@ -1169,17 +1183,32 @@ def _completion_choice_groups( for prompt_index, prompt in enumerate(prompts) if isinstance(prompt, list) and prompt == prompt_ids ] + explicit = explicit_matches.get(choice.index) + if ( + explicit is not None + and isinstance(prompts[explicit], list) + and prompts[explicit] != prompt_ids + ): + raise ValueError( + "Completions prompt_index contradicts exact prompt evidence" + ) if len(matches) == 1: exact_matches[choice.index] = matches[0] + if explicit is not None and explicit != matches[0]: + raise ValueError( + "Completions prompt_index contradicts exact prompt evidence" + ) - if len(exact_matches) == len(choices): + evidence_matches = {**exact_matches, **explicit_matches} + if len(evidence_matches) == len(choices): groups = [[] for _ in prompts] for choice in choices: - groups[exact_matches[choice.index]].append(choice) + groups[evidence_matches[choice.index]].append(choice) if all(groups) and ( requested_n is None or all(len(group) == requested_n for group in groups) ): return groups + raise ValueError("Cannot associate Completions choices with batched prompts") if len(choices) % count: raise ValueError("Cannot associate Completions choices with batched prompts") @@ -1194,13 +1223,11 @@ def _completion_choice_groups( ] for prompt_index, group in enumerate(groups): if any( - choice.index in exact_matches - and exact_matches[choice.index] != prompt_index + choice.index in evidence_matches + and evidence_matches[choice.index] != prompt_index for choice in group ): - raise ValueError( - "Completions exact prompt evidence contradicts choice indices" - ) + raise ValueError("Completions prompt evidence contradicts choice indices") return groups @@ -1345,7 +1372,7 @@ def anthropic_as_chat_completions_history( ChatCompletionsMessageSource( exchange=source.exchange, request_index=source.request_index, - output_index=0 if source.request_index is None else None, + output_indices=(0,) if source.request_index is None else None, ) if source is not None else None @@ -1383,53 +1410,75 @@ def responses_as_chat_completions_history( ) def converted( - source: ResponsesItemSource | None, + contributors: Sequence[ResponsesItemSource | None], ) -> ChatCompletionsMessageSource | None: - return ( - ChatCompletionsMessageSource( - exchange=source.exchange, - request_index=source.request_index, - output_index=source.output_index, - generation_index=source.generation_index, + present = [source for source in contributors if source is not None] + if not present: + return None + exchanges = {id(source.exchange): source.exchange for source in present} + if len(exchanges) != 1: + raise ValueError( + "One projected Chat message cannot span multiple Responses exchanges" + ) + generation_indices = { + source.generation_index + for source in present + if source.generation_index is not None + } + if len(generation_indices) > 1: + raise ValueError( + "One projected Chat message cannot span multiple Responses generations" + ) + output_indices = tuple( + dict.fromkeys( + source.output_index + for source in present + if source.output_index is not None ) - if source is not None - else None + ) + generation_index = next(iter(generation_indices), None) + first = present[0] + if output_indices or generation_index is not None: + return ChatCompletionsMessageSource( + exchange=first.exchange, + output_indices=output_indices, + generation_index=generation_index, + ) + return ChatCompletionsMessageSource( + exchange=first.exchange, + request_index=first.request_index, ) + source_groups: list[list[ResponsesItemSource | None]] = [] pending_reasoning: list[ResponsesItemSource | None] = [] - tool_message_open = False + tool_message_sources: list[ResponsesItemSource | None] | None = None for item, source in zip(history.input, history.input_sources, strict=True): kind = item.get("type") if kind == "reasoning": - tool_message_open = False + tool_message_sources = None pending_reasoning.append(source) continue if kind == "function_call": - if not tool_message_open: - first = next( - (item for item in pending_reasoning if item is not None), source - ) - sources.append(converted(first)) + if tool_message_sources is None: + tool_message_sources = [*pending_reasoning, source] + source_groups.append(tool_message_sources) pending_reasoning.clear() - tool_message_open = True + else: + tool_message_sources.append(source) continue - tool_message_open = False + tool_message_sources = None if pending_reasoning: - first = next( - (item for item in pending_reasoning if item is not None), source - ) - sources.append(converted(first)) - if not (kind in {None, "message"} and item.get("role") == "assistant"): - sources.append(converted(source)) + if kind in {None, "message"} and item.get("role") == "assistant": + source_groups.append([*pending_reasoning, source]) + else: + source_groups.append([*pending_reasoning]) + source_groups.append([source]) pending_reasoning.clear() else: - sources.append(converted(source)) + source_groups.append([source]) if pending_reasoning: - sources.append( - converted( - next((item for item in pending_reasoning if item is not None), None) - ) - ) + source_groups.append(pending_reasoning) + sources.extend(converted(group) for group in source_groups) if len(sources) != len(messages): raise AssertionError("Responses conversion sources must parallel messages") tools = _TOOLS.validate_python(_openai_tools(history.tools, dialect="responses")) diff --git a/tests/unit/trajectories/test_history.py b/tests/unit/trajectories/test_history.py index 2fd5d7e17..2aacd6c1b 100644 --- a/tests/unit/trajectories/test_history.py +++ b/tests/unit/trajectories/test_history.py @@ -387,6 +387,7 @@ def test_chat_history_normalizes_empty_response_only_fields() -> None: first = _chat([{"role": "user", "content": "one"}], "first") data = first.response.model_dump(mode="python") data["choices"][0]["message"]["tool_calls"] = [] + data["choices"][0]["message"]["annotations"] = [] first.response = ChatCompletion.model_validate(data) second = _chat( [ @@ -404,6 +405,41 @@ def test_chat_history_normalizes_empty_response_only_fields() -> None: assert len(history.messages) == 4 assert "tool_calls" not in history.messages[1] + assert "annotations" not in history.messages[1] + assert history.message_sources[1] is not None + assert history.message_sources[1].exchange is first + + +def test_chat_history_normalizes_assistant_missing_content_to_empty() -> None: + first = _chat([{"role": "user", "content": ""}], "") + data = first.response.model_dump(mode="python") + data["choices"][0]["message"] = { + "role": "assistant", + "content": None, + "annotations": [], + } + first.response = ChatCompletion.model_validate(data) + second = _chat( + [ + {"role": "user", "content": ""}, + {"role": "assistant", "content": ""}, + {"role": "user", "content": "next"}, + ], + "second", + offset=1, + ) + + history = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[first, second]) + ).chat_completions_history() + + assert [message["content"] for message in history.messages] == [ + "", + "", + "next", + "second", + ] + assert "annotations" not in history.messages[1] assert history.message_sources[1] is not None assert history.message_sources[1].exchange is first @@ -571,6 +607,11 @@ def test_anthropic_chat_conversion_preserves_sources_for_expanded_messages() -> assert source is not None assert source.exchange is exchange assert source.request_index == 0 + assert source.output_indices is None + response_source = converted.message_sources[-1] + assert response_source is not None + assert response_source.exchange is exchange + assert response_source.output_indices == (0,) converted.messages[2] = {"role": "user", "content": "changed"} with pytest.raises(ValueError, match="no longer matches"): @@ -604,7 +645,7 @@ def test_responses_history_expands_previous_response_chain() -> None: assert dict(chat_history.messages[1]).get("reasoning") == "think" assert all(source is not None for source in chat_history.message_sources) assert chat_history.message_sources[1] is not None - assert chat_history.message_sources[1].output_index == 0 + assert chat_history.message_sources[1].output_indices == (0, 1) trajectory.exchanges.responses[1].request["previous_response_id"] = "missing" external = trajectory.responses_histories() @@ -825,11 +866,39 @@ def test_responses_multi_generation_history_leaves_match_tokenization( assert all(item.get("type") != "reasoning" for item in histories[-1].input) converted = histories[-1].as_chat_completions_history() assert any( - source is not None and source.generation_index == 1 + source is not None + and source.generation_index == 1 + and source.output_indices == (second_output_index,) for source in converted.message_sources ) +def test_responses_generation_only_chat_source_has_empty_output_indices() -> None: + exchange = _response("response-empty-generation", "") + data = exchange.response.model_dump(mode="python") + data["output"] = [] + data["token_generations"] = [ + { + "prompt_token_ids": [1], + "output_tokens": [{"token_id": 2, "logprob": -0.2}], + "output_indices": [], + } + ] + exchange.response = Response.model_validate(data) + + converted = ( + art.Trajectory(exchanges=TrajectoryExchanges(responses=[exchange])) + .responses_history() + .as_chat_completions_history() + ) + + source = converted.message_sources[-1] + assert converted.messages[-1] == {"role": "assistant", "content": ""} + assert source is not None + assert source.output_indices == () + assert source.generation_index == 0 + + def test_completions_history_preserves_exact_tokens_and_sampled_spans() -> None: trajectory = art.Trajectory( exchanges=TrajectoryExchanges( @@ -914,6 +983,125 @@ def test_batched_completions_prefer_exact_prompt_association() -> None: assert [history.prompt for history in histories] == [[1, 10], [2, 20]] +def test_batched_completions_honor_interleaved_explicit_prompt_indices() -> None: + exchange = _completion([1], [10]) + exchange.request["prompt"] = ["first", "second"] + response = exchange.response.model_dump(mode="python") + response["choices"] = [ + { + "index": choice_index, + "finish_reason": "stop", + "text": text, + "prompt_index": prompt_index, + } + for choice_index, prompt_index, text in ( + (0, 1, "B0"), + (1, 0, "A0"), + (2, 1, "B1"), + (3, 0, "A1"), + ) + ] + exchange.response = Completion.model_validate(response) + + histories = art.Trajectory( + exchanges=TrajectoryExchanges(completions=[exchange]) + ).completions_string_histories() + + assert [history.prompt for history in histories] == [ + "firstA0", + "firstA1", + "secondB0", + "secondB1", + ] + + +@pytest.mark.parametrize(("prompt_index", "raises"), [(0, False), (1, True)]) +def test_batched_completions_validate_partial_prompt_index_fallback( + prompt_index: int, raises: bool +) -> None: + exchange = _completion([1], [10]) + exchange.request["prompt"] = ["first", "second"] + response = exchange.response.model_dump(mode="python") + response["choices"] = [ + { + "index": index, + "finish_reason": "stop", + "text": text, + **({"prompt_index": prompt_index} if index == 0 else {}), + } + for index, text in enumerate(("A0", "A1", "B0", "B1")) + ] + exchange.response = Completion.model_validate(response) + trajectory = art.Trajectory(exchanges=TrajectoryExchanges(completions=[exchange])) + + if raises: + with pytest.raises(ValueError, match="contradicts choice indices"): + trajectory.completions_string_histories() + return + assert [ + history.prompt for history in trajectory.completions_string_histories() + ] == ["firstA0", "firstA1", "secondB0", "secondB1"] + + +@pytest.mark.parametrize("prompt_index", [-1, 2, True, None]) +def test_batched_completions_reject_invalid_explicit_prompt_index( + prompt_index: object, +) -> None: + exchange = _completion([1], [10]) + exchange.request["prompt"] = ["first", "second"] + response = exchange.response.model_dump(mode="python") + response["choices"] = [ + { + "index": 0, + "finish_reason": "stop", + "text": "answer", + "prompt_index": prompt_index, + }, + { + "index": 1, + "finish_reason": "stop", + "text": "answer", + "prompt_index": 1, + }, + ] + exchange.response = Completion.model_validate(response) + + with pytest.raises(ValueError, match="prompt_index"): + art.Trajectory( + exchanges=TrajectoryExchanges(completions=[exchange]) + ).completions_string_histories() + + +def test_batched_completions_reject_prompt_index_exact_evidence_contradiction() -> None: + exchange = _completion([1], [10]) + exchange.request["prompt"] = [[1], [2]] + response = exchange.response.model_dump(mode="python") + response["choices"] = [ + { + "index": 0, + "finish_reason": "stop", + "text": "answer-0", + "prompt_index": 0, + "prompt_token_ids": [2], + "token_ids": [20], + }, + { + "index": 1, + "finish_reason": "stop", + "text": "answer-1", + "prompt_index": 1, + "prompt_token_ids": [1], + "token_ids": [10], + }, + ] + exchange.response = Completion.model_validate(response) + + with pytest.raises(ValueError, match="contradicts exact prompt evidence"): + art.Trajectory( + exchanges=TrajectoryExchanges(completions=[exchange]) + ).completions_token_histories() + + def test_completions_histories_never_silently_omit_mixed_evidence() -> None: exact = _completion([1], [2]) missing = _completion([3], [4], offset=1) From fbd51405c0515d98c6863303cfb3680822bb8bb6 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Sat, 25 Jul 2026 00:12:25 +0000 Subject: [PATCH 41/58] fix trajectory token evidence provenance --- src/art/trajectories/_tokenize.py | 189 ++++++++++++---------- tests/unit/trajectories/test_tokenize.py | 193 ++++++++++++++++++++++- 2 files changed, 294 insertions(+), 88 deletions(-) diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index 03c3813e8..a2e19fc43 100644 --- a/src/art/trajectories/_tokenize.py +++ b/src/art/trajectories/_tokenize.py @@ -436,6 +436,15 @@ def _logprob_values(values: object) -> list[float]: return result +def _chat_logprob_entries(choice: Choice) -> list[object]: + if choice.logprobs is None: + return [] + return [ + *(choice.logprobs.content or []), + *(choice.logprobs.refusal or []), + ] + + def _chat_choice_tokens( choice: Choice, response_data: dict[str, Any] ) -> tuple[list[int] | None, list[int] | None, list[float]]: @@ -452,10 +461,7 @@ def _chat_choice_tokens( choice_data.get("token_ids"), field="Chat Completions token_ids", ) - logprob_values = None - if choice.logprobs is not None: - logprob_values = choice.logprobs.content or choice.logprobs.refusal - values = list(logprob_values or []) + values = _chat_logprob_entries(choice) message = _dump(choice.message) if token_ids == [] and ( values or any(message.get(key) for key in ("content", "refusal", "tool_calls")) @@ -1375,9 +1381,7 @@ def _exchange_tokens( def _visible_logprobs(exchange: Exchange) -> list[tuple[str, float]]: values: list[tuple[str, float]] = [] if isinstance(exchange, ChatCompletionsExchange): - logprobs = exchange.response.choices[0].logprobs - entries = (logprobs.content or logprobs.refusal or []) if logprobs else [] - for entry in entries: + for entry in _chat_logprob_entries(exchange.response.choices[0]): data = _dump(entry) raw_bytes = data.get("bytes") if isinstance(raw_bytes, list): @@ -1952,6 +1956,19 @@ def _source_index(value: object, *, length: int, field: str) -> int: return value +def _chat_output_indices(source: object) -> tuple[int, ...] | None: + value = getattr(source, "output_indices", None) + if value is None: + return None + if not isinstance(value, tuple) or any( + not isinstance(index, int) or isinstance(index, bool) for index in value + ): + raise ValueError("Chat output source indices are invalid") + if tuple(sorted(set(value))) != value: + raise ValueError("Chat output source indices are not strictly ordered") + return value + + def _chat_choice(source: object) -> Choice: exchange = getattr(source, "exchange", None) choice_index = getattr(source, "choice_index", None) @@ -1984,12 +2001,18 @@ def _validate_history_sources(history: History) -> None: if source is None: continue exchange = source.exchange + if exchange.model != history.model: + raise ValueError( + "Chat Completions history model no longer matches its source " + "exchange" + ) + output_indices = _chat_output_indices(source) expected: list[dict[str, Any]] = [] if isinstance(exchange, ChatCompletionsExchange): if source.choice_index is not None: if ( source.request_index is not None - or source.output_index is not None + or output_indices is not None or source.generation_index is not None ): raise ValueError("Chat source has conflicting indices") @@ -2002,7 +2025,7 @@ def _validate_history_sources(history: History) -> None: expected.append(visible) elif source.request_index is not None: if ( - source.output_index is not None + output_indices is not None or source.generation_index is not None ): raise ValueError("Chat source has conflicting indices") @@ -2022,7 +2045,7 @@ def _validate_history_sources(history: History) -> None: ): raise ValueError("Anthropic-to-Chat source has invalid indices") if source.request_index is not None: - if source.output_index is not None: + if output_indices is not None: raise ValueError("Anthropic-to-Chat source has invalid indices") request_messages = exchange.request.get("messages", []) request_index = _source_index( @@ -2035,7 +2058,7 @@ def _validate_history_sources(history: History) -> None: {"messages": [request_messages[request_index]]} ) ) - elif source.output_index is None: + elif output_indices is None: expected.extend( _anthropic_messages( { @@ -2044,7 +2067,7 @@ def _validate_history_sources(history: History) -> None: } ) ) - elif source.output_index == 0: + elif output_indices == (0,): expected.append(_response_message(exchange)) visible_blocks = [ block.model_dump(mode="python", exclude_none=True) @@ -2067,7 +2090,7 @@ def _validate_history_sources(history: History) -> None: if source.request_index is not None: if ( source.choice_index is not None - or source.output_index is not None + or output_indices is not None or source.generation_index is not None ): raise ValueError("Responses-to-Chat source has invalid indices") @@ -2097,11 +2120,14 @@ def _validate_history_sources(history: History) -> None: else: if source.choice_index is not None: raise ValueError("Responses-to-Chat source has invalid indices") - if source.output_index is not None: - output_index = _source_index( - source.output_index, - length=len(exchange.response.output), - field="Responses output source index", + if output_indices is not None: + resolved_indices = tuple( + _source_index( + output_index, + length=len(exchange.response.output), + field="Responses output source index", + ) + for output_index in output_indices ) if source.generation_index is not None: generations = _response_generations(exchange.response) @@ -2110,28 +2136,19 @@ def _validate_history_sources(history: History) -> None: length=len(generations), field="Responses generation source index", ) - if ( - output_index - not in generations[generation_index].output_indices + generation_outputs = generations[ + generation_index + ].output_indices + if any( + output_index not in generation_outputs + for output_index in resolved_indices ): raise ValueError( "Responses output source does not belong to " "its generation" ) - generation_messages = _responses_messages( - { - "input": [ - exchange.response.output[ - generation_output_index - ].model_dump(mode="python", exclude_none=True) - for generation_output_index in generations[ - generation_index - ].output_indices - ] - } - ) - if len(generation_messages) == 1: - expected.extend(generation_messages) + if not resolved_indices and source.generation_index is not None: + expected.append({"role": "assistant", "content": ""}) expected.extend( _responses_messages( { @@ -2139,6 +2156,7 @@ def _validate_history_sources(history: History) -> None: exchange.response.output[ output_index ].model_dump(mode="python", exclude_none=True) + for output_index in resolved_indices ] } ) @@ -2169,6 +2187,11 @@ def _validate_history_sources(history: History) -> None: ): if source is None: continue + if source.exchange.model != history.model: + raise ValueError( + "Anthropic Messages history model no longer matches its source " + "exchange" + ) if source.request_index is None: expected = { "role": "assistant", @@ -2204,6 +2227,10 @@ def _validate_history_sources(history: History) -> None: for item, source in zip(history.input, history.input_sources, strict=True): if source is None: continue + if source.exchange.model != history.model: + raise ValueError( + "Responses history model no longer matches its source exchange" + ) if source.request_index is not None and source.output_index is not None: raise ValueError("Responses source has conflicting indices") if source.output_index is not None: @@ -2423,9 +2450,9 @@ def _source_signature(source: object) -> tuple[object, ...] | None: evidence_fingerprint = _sampled_evidence_fingerprint( exchange, protocol="responses", index=generation_index ) - elif ( - isinstance(exchange, MessagesExchange) - and getattr(source, "output_index", None) == 0 + elif isinstance(exchange, MessagesExchange) and ( + getattr(source, "output_index", None) == 0 + or _chat_output_indices(source) == (0,) ): evidence_fingerprint = _sampled_evidence_fingerprint( exchange, protocol="messages", index=0 @@ -2440,6 +2467,7 @@ def _source_signature(source: object) -> tuple[object, ...] | None: getattr(source, "request_index", None), choice_index, getattr(source, "output_index", None), + _chat_output_indices(source), generation_index, getattr(source, "prompt_index", None), ) @@ -2687,7 +2715,7 @@ def _tokenize_exact_responses_history( def _responses_source_generation( source: object, -) -> tuple[ResponsesExchange, _ResponseGeneration] | None: +) -> tuple[ResponsesExchange, _ResponseGeneration, tuple[int, ...]] | None: exchange = getattr(source, "exchange", None) generation_index = getattr(source, "generation_index", None) if not isinstance(exchange, ResponsesExchange) or not isinstance( @@ -2698,26 +2726,28 @@ def _responses_source_generation( if not 0 <= generation_index < len(generations): raise ValueError("Responses source generation index is out of bounds") generation = generations[generation_index] - output_index = getattr(source, "output_index", None) - if isinstance(output_index, int) and output_index not in generation.output_indices: + output_indices = _chat_output_indices(source) + if output_indices is None: + return None + if any(index not in generation.output_indices for index in output_indices): raise ValueError( "Responses source output index does not belong to its generation" ) - return exchange, generation + return exchange, generation, output_indices def _responses_generation_messages(source: object) -> list[dict[str, Any]] | None: selected = _responses_source_generation(source) if selected is None: return None - exchange, generation = selected + exchange, _, output_indices = selected return _responses_messages( { "input": [ exchange.response.output[index].model_dump( mode="python", exclude_none=True ) - for index in generation.output_indices + for index in output_indices ] } ) @@ -2729,7 +2759,12 @@ def _responses_generation_full_tokens( selected = _responses_source_generation(source) if selected is None: return None, [] - _, generation = selected + _, generation, output_indices = selected + if output_indices != tuple(generation.output_indices): + return None, [] + messages = _responses_generation_messages(source) + if messages is None or (messages and len(messages) != 1): + return None, [] return generation.output_token_ids, generation.output_logprobs @@ -2757,7 +2792,7 @@ def _chat_source_full_tokens( return None, [] return _responses_tokens(exchange.response)[1:] if isinstance(exchange, MessagesExchange): - if getattr(source, "output_index", None) != 0: + if _chat_output_indices(source) != (0,): return None, [] _, tokens, logprobs = _exchange_tokens(exchange) return tokens, logprobs @@ -2784,7 +2819,7 @@ def _chat_source_prompt_tokens(source: object) -> list[int] | None: return generations[generation_index].prompt_token_ids return _responses_tokens(exchange.response)[0] if isinstance(exchange, MessagesExchange): - if getattr(source, "output_index", None) != 0: + if _chat_output_indices(source) != (0,): return None return _messages_tokens(exchange.response)[0] return None @@ -2795,10 +2830,10 @@ def _source_is_sampled(source: object) -> bool: if isinstance(exchange, ChatCompletionsExchange): return getattr(source, "choice_index", None) is not None if isinstance(exchange, MessagesExchange): - return getattr(source, "output_index", None) == 0 + return _chat_output_indices(source) == (0,) if isinstance(exchange, ResponsesExchange): return ( - getattr(source, "output_index", None) is not None + _chat_output_indices(source) is not None or getattr(source, "generation_index", None) is not None ) return False @@ -2954,10 +2989,7 @@ def _chat_source_tokens( ): return tokens, logprobs return None, [] - if ( - isinstance(exchange, MessagesExchange) - and getattr(source, "output_index", None) == 0 - ): + if isinstance(exchange, MessagesExchange) and _chat_output_indices(source) == (0,): block_type = "thinking" if part == "reasoning" else "text" blocks = [ block @@ -2995,24 +3027,10 @@ def _chat_source_tokens( _, tokens, logprobs = _messages_tokens(exchange.response) return tokens, logprobs if isinstance(exchange, ResponsesExchange) and _source_is_sampled(source): - generation_index = getattr(source, "generation_index", None) - if isinstance(generation_index, int): - generations = _response_generations(exchange.response) - if not 0 <= generation_index < len(generations): - raise ValueError("Responses source generation index is out of bounds") - generation = generations[generation_index] - output_index = getattr(source, "output_index", None) - if ( - isinstance(output_index, int) - and output_index not in generation.output_indices - ): - raise ValueError( - "Responses source output index does not belong to its generation" - ) - output_index = getattr(source, "output_index", None) - if isinstance(output_index, int) and output_index < len( - exchange.response.output - ): + selected = _responses_source_generation(source) + output_indices = selected[2] if selected is not None else () + if len(output_indices) == 1: + output_index = output_indices[0] item = _dump(exchange.response.output[output_index]) if item.get("type") == "message" and part == "content": item_text = _responses_output_text(item.get("content")) @@ -3052,7 +3070,7 @@ def _source_covers_complete_sampled_message( message ) == normalize_chat_message(expected) if isinstance(exchange, MessagesExchange): - if getattr(source, "output_index", None) != 0: + if _chat_output_indices(source) != (0,): return False return normalize_chat_message(message) == normalize_chat_message( _response_message(exchange) @@ -3065,18 +3083,19 @@ def _source_covers_complete_sampled_message( generations = _response_generations(exchange.response) if not 0 <= generation_index < len(generations): raise ValueError("Responses source generation index is out of bounds") - if not generations[generation_index].output_indices: + output_indices = _chat_output_indices(source) + if output_indices is None: + return False + if not output_indices: return message.get("role") == "assistant" and not _chat_message_parts(message) - projected = _responses_messages( - { - "input": [ - exchange.response.output[index].model_dump( - mode="python", exclude_none=True - ) - for index in generations[generation_index].output_indices - ] - } - ) + generation = generations[generation_index] + if any(index not in generation.output_indices for index in output_indices): + raise ValueError( + "Responses source output index does not belong to its generation" + ) + projected = _responses_generation_messages(source) + if projected is None: + return False return len(projected) == 1 and normalize_chat_message( message ) == normalize_chat_message(projected[0]) @@ -3208,8 +3227,8 @@ def part_ids(text: str) -> list[int]: ): continue generation_index = getattr(source, "generation_index", None) - output_index = getattr(source, "output_index", None) - if not isinstance(generation_index, int) or not isinstance(output_index, int): + output_indices = _chat_output_indices(source) + if not isinstance(generation_index, int) or output_indices is None: continue key = (id(source.exchange), generation_index) if key in seen_generations: diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index f885a2186..38403acfd 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -740,6 +740,49 @@ def test_completions_history_rejects_model_and_sampled_span_mutation() -> None: history.tokenize() +@pytest.mark.parametrize( + "protocol", + ["chat_completions", "messages", "responses", "completions"], +) +def test_source_backed_histories_reject_model_mutation(protocol: str) -> None: + if protocol == "chat_completions": + history = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[_chat_exchange([1], [2])]) + ).chat_completions_history() + elif protocol == "messages": + history = art.Trajectory( + exchanges=TrajectoryExchanges( + messages=[ + _message_exchange( + MessagesRequest( + model="test/model", + messages=[{"role": "user", "content": "question"}], + max_tokens=16, + ), + prompt_token_ids=[1], + token_ids=[2], + logprobs=[-0.2], + ) + ] + ) + ).anthropic_messages_history() + elif protocol == "responses": + history = art.Trajectory( + exchanges=TrajectoryExchanges( + responses=[_response_exchange("response-model", 2)] + ) + ).responses_history() + else: + history = art.Trajectory( + exchanges=TrajectoryExchanges(completions=[_completion_exchange()]) + ).completions_token_history() + + history.model = "other/model" + + with pytest.raises(ValueError, match="model no longer matches"): + history.tokenize() + + def test_completions_string_history_preserves_textual_logprobs() -> None: exchange = _completion_exchange() payload = exchange.response.model_dump(mode="python") @@ -1443,6 +1486,103 @@ def test_exact_chat_ids_accept_ordinary_positional_logprobs() -> None: assert tokenized.logprobs[1] == -0.75 +def test_chat_content_and_refusal_logprobs_are_combined_in_protocol_order() -> None: + exchange = _chat_exchange([1], [2, 3]) + data = exchange.response.model_dump(mode="python") + choice = data["choices"][0] + choice["message"]["refusal"] = "refusal" + choice["logprobs"] = { + "content": [ + { + "token": "token_id:2", + "logprob": -0.2, + "bytes": [], + "top_logprobs": [], + } + ], + "refusal": [ + { + "token": "token_id:3", + "logprob": -0.3, + "bytes": [], + "top_logprobs": [], + } + ], + } + exchange.response = ChatCompletion.model_validate(data) + + tokenized = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[exchange]) + ).tokenize() + + assert tokenized.token_ids == [1, 2, 3] + assert tokenized.logprobs[1:] == pytest.approx([-0.2, -0.3]) + + +def test_chat_content_and_refusal_logprobs_must_match_exact_ids() -> None: + exchange = _chat_exchange([1], [2, 3]) + data = exchange.response.model_dump(mode="python") + choice = data["choices"][0] + choice["message"]["refusal"] = "refusal" + choice["logprobs"] = { + "content": [ + { + "token": "token_id:2", + "logprob": -0.2, + "bytes": [], + "top_logprobs": [], + } + ], + "refusal": [ + { + "token": "token_id:4", + "logprob": -0.4, + "bytes": [], + "top_logprobs": [], + } + ], + } + exchange.response = ChatCompletion.model_validate(data) + + with pytest.raises(ValueError, match="disagree with choice logprobs"): + art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[exchange]) + ).tokenize() + + +def test_chat_visible_logprobs_include_content_and_refusal() -> None: + from art.trajectories._tokenize import _visible_logprobs + + exchange = _chat_exchange([], []) + data = exchange.response.model_dump(mode="python") + choice = data["choices"][0] + choice["message"]["refusal"] = "refusal" + choice["logprobs"] = { + "content": [ + { + "token": "answer", + "logprob": -0.2, + "bytes": list(b"answer"), + "top_logprobs": [], + } + ], + "refusal": [ + { + "token": "refusal", + "logprob": -0.3, + "bytes": list(b"refusal"), + "top_logprobs": [], + } + ], + } + exchange.response = ChatCompletion.model_validate(data) + + assert _visible_logprobs(exchange) == [ + ("answer", -0.2), + ("refusal", -0.3), + ] + + def test_empty_chat_prompt_ids_are_missing_evidence( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -2614,13 +2754,13 @@ def test_responses_chat_sampled_source_requires_generation_identity() -> None: assistant_index = next( index for index, source in enumerate(history.message_sources) - if source is not None and source.output_index is not None + if source is not None and source.output_indices is not None ) source = history.message_sources[assistant_index] assert source is not None history.message_sources[assistant_index] = ChatCompletionsMessageSource( exchange=source.exchange, - output_index=source.output_index, + output_indices=source.output_indices, ) class Tokenizer: @@ -2640,6 +2780,53 @@ def apply_chat_template( history.tokenize(tokenizer=Tokenizer()) +def test_responses_chat_output_indices_reuse_one_complete_generation() -> None: + exchange = _response_exchange( + "response-output-indices", + 2, + prompt_token_ids=[1], + ) + history = art.ChatCompletionsHistory( + model="test/model", + messages=[ + {"role": "user", "content": "turn 0"}, + {"role": "assistant", "content": "answer"}, + ], + message_sources=[ + ChatCompletionsMessageSource(exchange=exchange, request_index=0), + ChatCompletionsMessageSource( + exchange=exchange, + output_indices=(0,), + generation_index=0, + ), + ], + ) + + tokenized = history.tokenize() + + assert tokenized.token_ids == [1, 2] + assert tokenized.logprobs[-1] == pytest.approx(-0.1) + assert tokenized.flags[-1] == art.TokenFlag.EXACT | art.TokenFlag.SAMPLED + + +def test_responses_chat_output_indices_are_bounds_checked() -> None: + exchange = _response_exchange("response-output-indices-invalid", 2) + history = art.ChatCompletionsHistory( + model="test/model", + messages=[{"role": "assistant", "content": "answer"}], + message_sources=[ + ChatCompletionsMessageSource( + exchange=exchange, + output_indices=(1,), + generation_index=0, + ) + ], + ) + + with pytest.raises(ValueError, match="out of bounds"): + history.tokenize() + + def _multi_output_responses_chat_history() -> art.ChatCompletionsHistory: exchange = _response_exchange("multi-output-generation", 2) data = exchange.response.model_dump(mode="python") @@ -2683,7 +2870,7 @@ def test_responses_output_source_rejects_sibling_generation_output() -> None: first_output = next( index for index, source in enumerate(history.message_sources) - if source is not None and source.output_index == 0 + if source is not None and source.output_indices == (0,) ) history.messages[first_output] = { "role": "assistant", From a8b77a8b2ffa3985b2da447afd0d92ba1574d1c4 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Sat, 25 Jul 2026 00:16:51 +0000 Subject: [PATCH 42/58] avoid duplicating aggregate response evidence --- src/art/trajectories/_tokenize.py | 11 ++++------- tests/unit/trajectories/test_tokenize.py | 12 +++++++----- 2 files changed, 11 insertions(+), 12 deletions(-) diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index a2e19fc43..5daba0842 100644 --- a/src/art/trajectories/_tokenize.py +++ b/src/art/trajectories/_tokenize.py @@ -3048,13 +3048,10 @@ def _chat_source_tokens( logprobs.extend(pair_logprobs) else: return token_ids, logprobs - if any( - getattr(item, "type", None) == "reasoning" - for item in exchange.response.output - ): - return None, [] - _, tokens, logprobs = _responses_tokens(exchange.response) - return tokens, logprobs + # Aggregate generation evidence cannot be partitioned safely across + # multiple projected messages. Item-local pairs above remain usable; + # otherwise render this message without claiming exact token identity. + return None, [] return None, [] diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index 38403acfd..73c78ae28 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -2881,7 +2881,9 @@ def test_responses_output_source_rejects_sibling_generation_output() -> None: history.tokenize() -def test_responses_multi_output_generation_exact_evidence_is_consumed_once() -> None: +def test_responses_multi_output_generation_is_rendered_without_duplicate_evidence() -> ( + None +): history = _multi_output_responses_chat_history() class Tokenizer: @@ -2900,12 +2902,12 @@ def apply_chat_template( tokenized = history.tokenize(tokenizer=Tokenizer(), chat_template="custom") - assert tokenized.token_ids == [1, 2, 3] - assert tokenized.logprobs[1:] == pytest.approx([-0.2, -0.3]) + assert tokenized.token_ids == [1, 20, 30] + assert all(math.isnan(value) for value in tokenized.logprobs) assert tokenized.flags == [ art.TokenFlag(0), - art.TokenFlag.EXACT | art.TokenFlag.SAMPLED, - art.TokenFlag.EXACT | art.TokenFlag.SAMPLED, + art.TokenFlag.SAMPLED, + art.TokenFlag.SAMPLED, ] From 13e7102f138595f8167c71bb087130ef7c738be1 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Sat, 25 Jul 2026 00:28:49 +0000 Subject: [PATCH 43/58] preserve outputless response evidence --- src/art/trajectories/_tokenize.py | 88 +++--------------------- tests/unit/trajectories/test_tokenize.py | 63 +++++++++++++++++ 2 files changed, 71 insertions(+), 80 deletions(-) diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index 5daba0842..1c0cb315a 100644 --- a/src/art/trajectories/_tokenize.py +++ b/src/art/trajectories/_tokenize.py @@ -2139,6 +2139,11 @@ def _validate_history_sources(history: History) -> None: generation_outputs = generations[ generation_index ].output_indices + if not resolved_indices and generation_outputs: + raise ValueError( + "Responses empty output source references a " + "generation with output items" + ) if any( output_index not in generation_outputs for output_index in resolved_indices @@ -2784,7 +2789,7 @@ def _chat_source_full_tokens( if isinstance(exchange, ResponsesExchange): generation_messages = _responses_generation_messages(source) if generation_messages is not None: - if len(generation_messages) != 1: + if len(generation_messages) > 1: return None, [] return _responses_generation_full_tokens(source) generations = _response_generations(exchange.response) @@ -3214,88 +3219,10 @@ def part_ids(text: str) -> list[int]: and _source_is_sampled(source) for _, text in _chat_message_parts(message) ] - from ._history import normalize_chat_message - - atomic_generations: dict[int, tuple[int, object, list[int], list[float]]] = {} - seen_generations: set[tuple[int, int]] = set() - for message_index, source in enumerate(history.message_sources): - if source is None or not isinstance( - getattr(source, "exchange", None), ResponsesExchange - ): - continue - generation_index = getattr(source, "generation_index", None) - output_indices = _chat_output_indices(source) - if not isinstance(generation_index, int) or output_indices is None: - continue - key = (id(source.exchange), generation_index) - if key in seen_generations: - continue - seen_generations.add(key) - projected = _responses_generation_messages(source) - if projected is None or len(projected) <= 1: - continue - end = message_index + len(projected) - if end > len(messages) or [ - normalize_chat_message(item) for item in messages[message_index:end] - ] != [normalize_chat_message(item) for item in projected]: - continue - group_sources = history.message_sources[message_index:end] - if any( - item is None - or item.exchange is not source.exchange - or item.generation_index != generation_index - for item in group_sources - ): - continue - exact, exact_logprobs = _responses_generation_full_tokens(source) - if exact is not None: - atomic_generations[message_index] = ( - end, - source, - exact, - exact_logprobs, - ) - sampled_part_cursor = 0 - atomic_end = 0 for message_index, (message, source) in enumerate( zip(history.messages, history.message_sources, strict=True) ): - if message_index < atomic_end: - continue - atomic = atomic_generations.get(message_index) - if atomic is not None: - end, atomic_source, exact, exact_logprobs = atomic - prompt_render = render(messages[:message_index], add_generation_prompt=True) - completed_render = render(messages[:end], add_generation_prompt=False) - if ( - len(completed_render) <= len(prompt_render) - or completed_render[: len(prompt_render)] != prompt_render - or rendered[: len(completed_render)] != completed_render - ): - raise ValueError( - "Could not locate a complete sampled generation in the " - "rendered history" - ) - replacements.append( - ( - len(prompt_render), - len(completed_render), - exact, - exact_logprobs - if len(exact_logprobs) == len(exact) - else [math.nan] * len(exact), - True, - _sampled_source_key(atomic_source), - atomic_source, - ) - ) - search_cursor = len(completed_render) - sampled_part_cursor += sum( - len(_chat_message_parts(item)) for item in messages[message_index:end] - ) - atomic_end = end - continue parts = _chat_message_parts(message) sampled = ( message.get("role") == "assistant" @@ -3386,7 +3313,8 @@ def part_ids(text: str) -> list[int]: messages[: message_index + 1], add_generation_prompt=False ) if ( - len(completed_render) <= len(prompt_render) + len(completed_render) < len(prompt_render) + or (parts and len(completed_render) == len(prompt_render)) or completed_render[: len(prompt_render)] != prompt_render or rendered[: len(completed_render)] != completed_render ): diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index 73c78ae28..34f623e42 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -2377,6 +2377,69 @@ def test_responses_terminal_generation_without_output_items_is_tokenized() -> No } +def test_responses_terminal_generation_without_output_items_survives_rerender() -> None: + exchange = _response_exchange("terminal-eos-rerender", 2) + data = exchange.response.model_dump(mode="python") + data["output"] = [] + data["token_generations"] = [ + { + "prompt_token_ids": [1], + "output_tokens": [{"token_id": 2, "logprob": -0.2}], + "output_indices": [], + } + ] + exchange.response = Response.model_validate(data) + history = ( + art.Trajectory(exchanges=TrajectoryExchanges(responses=[exchange])) + .responses_history() + .as_chat_completions_history() + ) + + class Tokenizer: + def __call__(self, text: str, **kwargs: object) -> list[int]: + del text, kwargs + return [1] + + def apply_chat_template( + self, messages: list[dict[str, Any]], **kwargs: object + ) -> list[int]: + del messages, kwargs + return [1] + + tokenized = history.tokenize(tokenizer=Tokenizer(), chat_template="custom") + + assert tokenized.token_ids == [1, 2] + assert math.isnan(tokenized.logprobs[0]) + assert tokenized.logprobs[1] == pytest.approx(-0.2) + assert tokenized.flags == [ + art.TokenFlag(0), + art.TokenFlag.EXACT | art.TokenFlag.SAMPLED, + ] + + +def test_responses_empty_chat_source_requires_outputless_generation() -> None: + exchange = _response_exchange("nonempty-generation-empty-source", 2) + history = ( + art.Trajectory(exchanges=TrajectoryExchanges(responses=[exchange])) + .responses_history() + .as_chat_completions_history() + ) + assistant_index = next( + index + for index, source in enumerate(history.message_sources) + if source is not None and source.output_indices is not None + ) + history.messages[assistant_index] = {"role": "assistant", "content": ""} + history.message_sources[assistant_index] = ChatCompletionsMessageSource( + exchange=exchange, + output_indices=(), + generation_index=0, + ) + + with pytest.raises(ValueError, match="empty output source"): + history.tokenize() + + def test_responses_nonterminal_generation_without_output_items_raises() -> None: exchange = _response_exchange("hidden-control", 2) data = exchange.response.model_dump(mode="python") From 8c653b790540a4a040adb9f9efd1c98afcbda475 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Sat, 25 Jul 2026 00:29:15 +0000 Subject: [PATCH 44/58] Preserve composite Responses chat provenance --- src/art/trajectories/_history.py | 93 ++++++++++++++++++++++--- tests/unit/trajectories/test_history.py | 86 +++++++++++++++++++++++ 2 files changed, 170 insertions(+), 9 deletions(-) diff --git a/src/art/trajectories/_history.py b/src/art/trajectories/_history.py index 348d77590..b81fc41c7 100644 --- a/src/art/trajectories/_history.py +++ b/src/art/trajectories/_history.py @@ -1398,9 +1398,7 @@ def responses_as_chat_completions_history( ) from ._tokenize import _openai_tools, _responses_messages - messages = _responses_messages( - {"instructions": history.instructions, "input": history.input} - ) + messages = _responses_messages({"instructions": history.instructions, "input": []}) sources: list[ChatCompletionsMessageSource | None] = [] if history.instructions: sources.append( @@ -1412,9 +1410,9 @@ def responses_as_chat_completions_history( def converted( contributors: Sequence[ResponsesItemSource | None], ) -> ChatCompletionsMessageSource | None: - present = [source for source in contributors if source is not None] - if not present: + if not contributors or any(source is None for source in contributors): return None + present = [source for source in contributors if source is not None] exchanges = {id(source.exchange): source.exchange for source in present} if len(exchanges) != 1: raise ValueError( @@ -1438,47 +1436,124 @@ def converted( ) generation_index = next(iter(generation_indices), None) first = present[0] + request_index = next( + ( + source.request_index + for source in present + if source.request_index is not None + ), + None, + ) if output_indices or generation_index is not None: return ChatCompletionsMessageSource( exchange=first.exchange, + request_index=request_index, output_indices=output_indices, generation_index=generation_index, ) return ChatCompletionsMessageSource( exchange=first.exchange, - request_index=first.request_index, + request_index=request_index, ) + no_output = object() + + def compatible( + contributors: Sequence[ResponsesItemSource | None], + source: ResponsesItemSource | None, + ) -> bool: + if source is None: + return True + present = [item for item in contributors if item is not None] + if present and source.exchange is not present[0].exchange: + return False + + def output_generation(item: ResponsesItemSource) -> object: + if item.output_index is None and item.generation_index is None: + return no_output + return item.generation_index + + generations = { + value + for item in present + if (value := output_generation(item)) is not no_output + } + candidate = output_generation(source) + return candidate is no_output or not generations or candidate in generations + + item_groups: list[list[ResponseInputItemParam]] = [] source_groups: list[list[ResponsesItemSource | None]] = [] + pending_reasoning_items: list[ResponseInputItemParam] = [] pending_reasoning: list[ResponsesItemSource | None] = [] + tool_message_items: list[ResponseInputItemParam] | None = None tool_message_sources: list[ResponsesItemSource | None] | None = None for item, source in zip(history.input, history.input_sources, strict=True): kind = item.get("type") if kind == "reasoning": + if pending_reasoning and not compatible(pending_reasoning, source): + item_groups.append([*pending_reasoning_items]) + source_groups.append([*pending_reasoning]) + pending_reasoning_items.clear() + pending_reasoning.clear() + tool_message_items = None tool_message_sources = None + pending_reasoning_items.append(item) pending_reasoning.append(source) continue if kind == "function_call": if tool_message_sources is None: + if pending_reasoning and not compatible(pending_reasoning, source): + item_groups.append([*pending_reasoning_items]) + source_groups.append([*pending_reasoning]) + pending_reasoning_items.clear() + pending_reasoning.clear() + tool_message_items = [*pending_reasoning_items, item] tool_message_sources = [*pending_reasoning, source] + item_groups.append(tool_message_items) source_groups.append(tool_message_sources) + pending_reasoning_items.clear() pending_reasoning.clear() - else: + elif compatible(tool_message_sources, source): + assert tool_message_items is not None + tool_message_items.append(item) tool_message_sources.append(source) + else: + tool_message_items = [item] + tool_message_sources = [source] + item_groups.append(tool_message_items) + source_groups.append(tool_message_sources) continue + tool_message_items = None tool_message_sources = None if pending_reasoning: - if kind in {None, "message"} and item.get("role") == "assistant": + if ( + kind in {None, "message"} + and item.get("role") == "assistant" + and compatible(pending_reasoning, source) + ): + item_groups.append([*pending_reasoning_items, item]) source_groups.append([*pending_reasoning, source]) else: + item_groups.append([*pending_reasoning_items]) source_groups.append([*pending_reasoning]) + item_groups.append([item]) source_groups.append([source]) + pending_reasoning_items.clear() pending_reasoning.clear() else: + item_groups.append([item]) source_groups.append([source]) if pending_reasoning: + item_groups.append(pending_reasoning_items) source_groups.append(pending_reasoning) - sources.extend(converted(group) for group in source_groups) + for items, group in zip(item_groups, source_groups, strict=True): + projected = _responses_messages({"input": items}) + if len(projected) != 1: + raise AssertionError( + "One Responses projection group must produce one Chat message" + ) + messages.extend(projected) + sources.append(converted(group)) if len(sources) != len(messages): raise AssertionError("Responses conversion sources must parallel messages") tools = _TOOLS.validate_python(_openai_tools(history.tools, dialect="responses")) diff --git a/tests/unit/trajectories/test_history.py b/tests/unit/trajectories/test_history.py index 2aacd6c1b..516e845b5 100644 --- a/tests/unit/trajectories/test_history.py +++ b/tests/unit/trajectories/test_history.py @@ -653,6 +653,92 @@ def test_responses_history_expands_previous_response_chain() -> None: assert external[1].previous_response_id == "missing" +def test_responses_chat_conversion_preserves_request_and_output_sources() -> None: + exchange = _response("response-mixed-source", "answer") + exchange.request["input"] = [ + { + "id": "request-reasoning", + "type": "reasoning", + "summary": [{"type": "summary_text", "text": "prior thought"}], + } + ] + + converted = ( + art.Trajectory(exchanges=TrajectoryExchanges(responses=[exchange])) + .responses_history() + .as_chat_completions_history() + ) + + assert converted.messages == [ + { + "role": "assistant", + "content": "answer", + "reasoning": "prior thought", + } + ] + source = converted.message_sources[0] + assert source is not None + assert source.exchange is exchange + assert source.request_index == 0 + assert source.output_indices == (0,) + + +def test_responses_chat_conversion_splits_cross_exchange_assistant_sources() -> None: + first = _response("response-reasoning", "", reasoning="think") + first_data = first.response.model_dump(mode="python") + first_data["output"] = first_data["output"][:1] + first.response = Response.model_validate(first_data) + second = _response( + "response-answer", + "answer", + previous_response_id="response-reasoning", + offset=1, + ) + second.request["input"] = [] + + converted = ( + art.Trajectory(exchanges=TrajectoryExchanges(responses=[first, second])) + .responses_history() + .as_chat_completions_history() + ) + + assert converted.messages[-2:] == [ + {"role": "assistant", "content": "", "reasoning": "think"}, + {"role": "assistant", "content": "answer"}, + ] + first_source, second_source = converted.message_sources[-2:] + assert first_source is not None and first_source.exchange is first + assert first_source.output_indices == (0,) + assert second_source is not None and second_source.exchange is second + assert second_source.output_indices == (0,) + + +def test_responses_chat_conversion_owns_request_tool_group_by_first_item() -> None: + exchange = _response("response-request-tools", "answer") + exchange.request["input"] = [ + { + "type": "function_call", + "call_id": f"call-{index}", + "name": f"tool_{index}", + "arguments": "{}", + } + for index in range(2) + ] + + converted = ( + art.Trajectory(exchanges=TrajectoryExchanges(responses=[exchange])) + .responses_history() + .as_chat_completions_history() + ) + + assert len(converted.messages[0].get("tool_calls", [])) == 2 + source = converted.message_sources[0] + assert source is not None + assert source.exchange is exchange + assert source.request_index == 0 + assert source.output_indices is None + + def test_responses_history_propagates_opaque_context_and_first_sources() -> None: first = _response( "response-1", From babbd5a0abb2707de5f17cec19eafd51c28ef68c Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Sat, 25 Jul 2026 00:31:32 +0000 Subject: [PATCH 45/58] Trust explicit string completion routing --- src/art/trajectories/_history.py | 4 ---- tests/unit/trajectories/test_history.py | 26 +++++++++++++++++++++++++ 2 files changed, 26 insertions(+), 4 deletions(-) diff --git a/src/art/trajectories/_history.py b/src/art/trajectories/_history.py index b81fc41c7..96b43e3d5 100644 --- a/src/art/trajectories/_history.py +++ b/src/art/trajectories/_history.py @@ -1194,10 +1194,6 @@ def _completion_choice_groups( ) if len(matches) == 1: exact_matches[choice.index] = matches[0] - if explicit is not None and explicit != matches[0]: - raise ValueError( - "Completions prompt_index contradicts exact prompt evidence" - ) evidence_matches = {**exact_matches, **explicit_matches} if len(evidence_matches) == len(choices): diff --git a/tests/unit/trajectories/test_history.py b/tests/unit/trajectories/test_history.py index 516e845b5..c9926c89d 100644 --- a/tests/unit/trajectories/test_history.py +++ b/tests/unit/trajectories/test_history.py @@ -1188,6 +1188,32 @@ def test_batched_completions_reject_prompt_index_exact_evidence_contradiction() ).completions_token_histories() +def test_batched_completions_trust_explicit_string_prompt_index() -> None: + from art.trajectories._history import _completion_choice_groups + + exchange = _completion([1], [10]) + # Captured provider payloads can be broader than the SDK's prompt union. + exchange.request["prompt"] = cast(Any, ["same tokenization", [42]]) + response = exchange.response.model_dump(mode="python") + response["choices"] = [ + { + "index": index, + "finish_reason": "stop", + "text": f"answer-{index}", + "prompt_index": index, + "prompt_token_ids": [42], + "token_ids": [index + 10], + } + for index in range(2) + ] + exchange.response = Completion.model_validate(response) + + assert [ + [choice.index for choice in group] + for group in _completion_choice_groups(exchange) + ] == [[0], [1]] + + def test_completions_histories_never_silently_omit_mixed_evidence() -> None: exact = _completion([1], [2]) missing = _completion([3], [4], offset=1) From 775599249f67ef58a760758adf274635e8a23b60 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Sat, 25 Jul 2026 00:35:27 +0000 Subject: [PATCH 46/58] Separate request and sampled response sources --- src/art/trajectories/_history.py | 34 +++++++++++++++---------- tests/unit/trajectories/test_history.py | 22 +++++++++++----- 2 files changed, 36 insertions(+), 20 deletions(-) diff --git a/src/art/trajectories/_history.py b/src/art/trajectories/_history.py index 96b43e3d5..41b246d89 100644 --- a/src/art/trajectories/_history.py +++ b/src/art/trajectories/_history.py @@ -1432,24 +1432,19 @@ def converted( ) generation_index = next(iter(generation_indices), None) first = present[0] - request_index = next( - ( - source.request_index - for source in present - if source.request_index is not None - ), - None, - ) if output_indices or generation_index is not None: return ChatCompletionsMessageSource( exchange=first.exchange, - request_index=request_index, output_indices=output_indices, generation_index=generation_index, ) return ChatCompletionsMessageSource( exchange=first.exchange, - request_index=request_index, + request_index=next( + source.request_index + for source in present + if source.request_index is not None + ), ) no_output = object() @@ -1458,11 +1453,24 @@ def compatible( contributors: Sequence[ResponsesItemSource | None], source: ResponsesItemSource | None, ) -> bool: - if source is None: - return True + def sampled(item: ResponsesItemSource | None) -> bool: + return item is not None and ( + item.output_index is not None or item.generation_index is not None + ) + present = [item for item in contributors if item is not None] - if present and source.exchange is not present[0].exchange: + if ( + source is not None + and present + and source.exchange is not present[0].exchange + ): + return False + if contributors and sampled(source) != any( + sampled(item) for item in contributors + ): return False + if source is None: + return True def output_generation(item: ResponsesItemSource) -> object: if item.output_index is None and item.generation_index is None: diff --git a/tests/unit/trajectories/test_history.py b/tests/unit/trajectories/test_history.py index c9926c89d..a79deffad 100644 --- a/tests/unit/trajectories/test_history.py +++ b/tests/unit/trajectories/test_history.py @@ -672,15 +672,23 @@ def test_responses_chat_conversion_preserves_request_and_output_sources() -> Non assert converted.messages == [ { "role": "assistant", - "content": "answer", + "content": "", "reasoning": "prior thought", - } + }, + { + "role": "assistant", + "content": "answer", + }, ] - source = converted.message_sources[0] - assert source is not None - assert source.exchange is exchange - assert source.request_index == 0 - assert source.output_indices == (0,) + request_source, output_source = converted.message_sources + assert request_source is not None + assert request_source.exchange is exchange + assert request_source.request_index == 0 + assert request_source.output_indices is None + assert output_source is not None + assert output_source.exchange is exchange + assert output_source.request_index is None + assert output_source.output_indices == (0,) def test_responses_chat_conversion_splits_cross_exchange_assistant_sources() -> None: From 76d7aa899e04fd40efcab6e7235dcd8f7efe188a Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Sat, 25 Jul 2026 00:36:21 +0000 Subject: [PATCH 47/58] validate split responses source evidence --- src/art/trajectories/_tokenize.py | 11 +++++---- tests/unit/trajectories/test_tokenize.py | 30 +++++++++++++++++++++++- 2 files changed, 36 insertions(+), 5 deletions(-) diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index 1c0cb315a..1308a6762 100644 --- a/src/art/trajectories/_tokenize.py +++ b/src/art/trajectories/_tokenize.py @@ -2101,11 +2101,12 @@ def _validate_history_sources(history: History) -> None: length=len(request_input), field="Responses request source index", ) - expected.extend( - _responses_messages( - {"input": [request_input[request_index]]} + for end in range(request_index + 1, len(request_input) + 1): + projected = _responses_messages( + {"input": request_input[request_index:end]} ) - ) + if len(projected) == 1: + expected.extend(projected) elif isinstance(request_input, str): _source_index( source.request_index, @@ -2734,6 +2735,8 @@ def _responses_source_generation( output_indices = _chat_output_indices(source) if output_indices is None: return None + if bool(output_indices) != bool(generation.output_indices): + raise ValueError("Responses empty output source does not match its generation") if any(index not in generation.output_indices for index in output_indices): raise ValueError( "Responses source output index does not belong to its generation" diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index 34f623e42..86ee03f72 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -2418,6 +2418,8 @@ def apply_chat_template( def test_responses_empty_chat_source_requires_outputless_generation() -> None: + from art.trajectories._tokenize import _responses_source_generation + exchange = _response_exchange("nonempty-generation-empty-source", 2) history = ( art.Trajectory(exchanges=TrajectoryExchanges(responses=[exchange])) @@ -2430,16 +2432,42 @@ def test_responses_empty_chat_source_requires_outputless_generation() -> None: if source is not None and source.output_indices is not None ) history.messages[assistant_index] = {"role": "assistant", "content": ""} - history.message_sources[assistant_index] = ChatCompletionsMessageSource( + invalid_source = ChatCompletionsMessageSource( exchange=exchange, output_indices=(), generation_index=0, ) + history.message_sources[assistant_index] = invalid_source + with pytest.raises(ValueError, match="empty output source"): + _responses_source_generation(invalid_source) with pytest.raises(ValueError, match="empty output source"): history.tokenize() +def test_responses_request_composite_source_validation_uses_contiguous_items() -> None: + from art.trajectories._tokenize import _validate_history_sources + + exchange = _response_exchange("request-composite", 2) + exchange.request["input"] = [ + { + "type": "function_call", + "call_id": f"call-{index}", + "name": f"tool_{index}", + "arguments": "{}", + } + for index in range(2) + ] + history = ( + art.Trajectory(exchanges=TrajectoryExchanges(responses=[exchange])) + .responses_history() + .as_chat_completions_history() + ) + + assert len(history.messages[0].get("tool_calls", [])) == 2 + _validate_history_sources(history) + + def test_responses_nonterminal_generation_without_output_items_raises() -> None: exchange = _response_exchange("hidden-control", 2) data = exchange.response.model_dump(mode="python") From 28777501e9197df9c6f146102399e22002f88b95 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Sat, 25 Jul 2026 00:51:50 +0000 Subject: [PATCH 48/58] fix fallback trajectory token evidence --- src/art/trajectories/_tokenize.py | 153 ++++++++++++------ tests/unit/trajectories/test_tokenize.py | 195 ++++++++++++++++++++++- 2 files changed, 296 insertions(+), 52 deletions(-) diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index 1308a6762..da8250c0b 100644 --- a/src/art/trajectories/_tokenize.py +++ b/src/art/trajectories/_tokenize.py @@ -42,7 +42,6 @@ Tokenizer, Trajectory, TrajectoryGroup, - TrajectoryHistory, ) from ._history import _model_matches from ._protocols import Exchange @@ -357,6 +356,14 @@ def _dump(value: object) -> dict[str, Any]: return _string_dict(value) or {} +def _field(value: object, name: str, default: object = None) -> object: + return ( + value.get(name, default) + if isinstance(value, Mapping) + else getattr(value, name, default) + ) + + def _token_id(value: object) -> int | None: if isinstance(value, int) and not isinstance(value, bool) and value >= 0: return value @@ -616,6 +623,10 @@ class _ResponseGeneration: output_text: str | None +def _responses_output_is_sampled(item: object) -> bool: + return not str(_field(item, "type", "")).endswith("_output") + + def _response_generations(response: Response) -> list[_ResponseGeneration]: data = _dump(response) raw_generations = data.get("token_generations") @@ -712,7 +723,7 @@ def _response_generations(response: Response) -> list[_ResponseGeneration]: else None ) if output_ids is None and any( - not str(getattr(response.output[value], "type", "")).endswith("_output") + _responses_output_is_sampled(response.output[value]) for value in output_indices ): raise ValueError( @@ -731,7 +742,7 @@ def _response_generations(response: Response) -> list[_ResponseGeneration]: required_outputs = { index for index, item in enumerate(response.output) - if not str(getattr(item, "type", "")).endswith("_output") + if _responses_output_is_sampled(item) } if not required_outputs.issubset(assigned_outputs): raise ValueError( @@ -2719,22 +2730,39 @@ def _tokenize_exact_responses_history( return tokenized +def _responses_source_outputs( + source: object, +) -> tuple[ResponsesExchange, tuple[int, ...]] | None: + exchange = getattr(source, "exchange", None) + if not isinstance(exchange, ResponsesExchange): + return None + output_indices = _chat_output_indices(source) + if output_indices is None: + return None + return exchange, tuple( + _source_index( + index, + length=len(exchange.response.output), + field="Responses source output index", + ) + for index in output_indices + ) + + def _responses_source_generation( source: object, ) -> tuple[ResponsesExchange, _ResponseGeneration, tuple[int, ...]] | None: - exchange = getattr(source, "exchange", None) + selected = _responses_source_outputs(source) generation_index = getattr(source, "generation_index", None) - if not isinstance(exchange, ResponsesExchange) or not isinstance( - generation_index, int - ): + if selected is None or generation_index is None: return None + if not isinstance(generation_index, int) or isinstance(generation_index, bool): + raise ValueError("Responses source generation index is invalid") + exchange, output_indices = selected generations = _response_generations(exchange.response) if not 0 <= generation_index < len(generations): raise ValueError("Responses source generation index is out of bounds") generation = generations[generation_index] - output_indices = _chat_output_indices(source) - if output_indices is None: - return None if bool(output_indices) != bool(generation.output_indices): raise ValueError("Responses empty output source does not match its generation") if any(index not in generation.output_indices for index in output_indices): @@ -2790,15 +2818,7 @@ def _chat_source_full_tokens( _, tokens, logprobs = _chat_choice_tokens(choice, _dump(exchange.response)) return tokens, logprobs if isinstance(exchange, ResponsesExchange): - generation_messages = _responses_generation_messages(source) - if generation_messages is not None: - if len(generation_messages) > 1: - return None, [] - return _responses_generation_full_tokens(source) - generations = _response_generations(exchange.response) - if len(generations) > 1: - return None, [] - return _responses_tokens(exchange.response)[1:] + return _responses_generation_full_tokens(source) if isinstance(exchange, MessagesExchange): if _chat_output_indices(source) != (0,): return None, [] @@ -2840,9 +2860,12 @@ def _source_is_sampled(source: object) -> bool: if isinstance(exchange, MessagesExchange): return _chat_output_indices(source) == (0,) if isinstance(exchange, ResponsesExchange): - return ( - _chat_output_indices(source) is not None - or getattr(source, "generation_index", None) is not None + if getattr(source, "generation_index", None) is not None: + return True + selected = _responses_source_outputs(source) + return selected is not None and any( + _responses_output_is_sampled(selected[0].response.output[index]) + for index in selected[1] ) return False @@ -3035,8 +3058,8 @@ def _chat_source_tokens( _, tokens, logprobs = _messages_tokens(exchange.response) return tokens, logprobs if isinstance(exchange, ResponsesExchange) and _source_is_sampled(source): - selected = _responses_source_generation(source) - output_indices = selected[2] if selected is not None else () + selected = _responses_source_outputs(source) + output_indices = selected[1] if selected is not None else () if len(output_indices) == 1: output_index = output_indices[0] item = _dump(exchange.response.output[output_index]) @@ -3113,6 +3136,7 @@ def _tokenize_chat_view( tokenizer: Tokenizer | None, chat_template: str | None, chat_template_kwargs: Mapping[str, object] | None, + _projection_matches: bool | None = None, _trace: _TraceBuilder | None = None, ) -> TokenizedHistory: _validate_history_sources(history) @@ -3163,6 +3187,40 @@ def render( add_generation_prompt=not ends_with_assistant, ) + def source_matches_context(source: object) -> bool: + exchange = getattr(source, "exchange", None) + return ( + isinstance(exchange, ChatCompletionsExchange) + and history.tools == exchange.request.get("tools") + and history.chat_template == exchange.request.get("chat_template") + and history.chat_template_kwargs + == exchange.request.get("chat_template_kwargs") + ) + + exact_prefix_length = 0 + if ( + chat_template is None + and chat_template_kwargs is None + and _projection_matches is True + ): + for message_index, (message, source) in enumerate( + zip(history.messages, history.message_sources, strict=True) + ): + if message.get("role") != "assistant" or source is None: + continue + source_prompt = _chat_source_prompt_tokens(source) + if source_prompt and source_matches_context(source): + rendered_prompt = render( + messages[:message_index], add_generation_prompt=True + ) + if rendered[: len(rendered_prompt)] == rendered_prompt: + rendered = [ + *source_prompt, + *rendered[len(rendered_prompt) :], + ] + exact_prefix_length = len(source_prompt) + break + positions_by_first_token: dict[int, list[int]] = {} for index, token_id in enumerate(rendered): positions_by_first_token.setdefault(token_id, []).append(index) @@ -3244,14 +3302,7 @@ def part_ids(text: str) -> list[int]: if sampled and source is not None: source_prompt = _chat_source_prompt_tokens(source) source_exchange = getattr(source, "exchange", None) - source_context_matches = ( - isinstance(source_exchange, ChatCompletionsExchange) - and history.tools == source_exchange.request.get("tools") - and history.chat_template - == source_exchange.request.get("chat_template") - and history.chat_template_kwargs - == source_exchange.request.get("chat_template_kwargs") - ) + source_context_matches = source_matches_context(source) if ( source_context_matches and source_prompt @@ -3337,6 +3388,12 @@ def part_ids(text: str) -> list[int]: source, ) ) + if ( + message_index < len(messages) - 1 + and isinstance(source_exchange, ChatCompletionsExchange) + and completed_render[len(prompt_render) :] != full_exact + ): + _warn_prefix_retokenization() search_cursor = len(completed_render) sampled_part_cursor += len(parts) continue @@ -3371,10 +3428,12 @@ def part_ids(text: str) -> list[int]: and not source_boundary and len(local_matches) != same_sampled_parts ): - raise ValueError( - "Could not uniquely locate a sampled history message in the " - "rendered history" - ) + if not source_context_matches: + raise ValueError( + "Could not uniquely locate a sampled history message in the " + "rendered history" + ) + continue start, end = span search_cursor = end if not sampled: @@ -3477,6 +3536,8 @@ def part_ids(text: str) -> list[int]: logprobs.extend([math.nan] * (len(rendered) - cursor)) flags.extend([TokenFlag(0)] * (len(rendered) - cursor)) source_keys.extend([None] * (len(rendered) - cursor)) + for index in range(min(exact_prefix_length, len(flags))): + flags[index] |= TokenFlag.EXACT if history.model is None: raise ValueError("History tokenization requires a model") tokenized = TokenizedHistory( @@ -3874,15 +3935,17 @@ def tokenize_history( ) ): return exact - if needs_render: - return _tokenize_chat_view( - history, - base_model=base_model, - tokenizer=tokenizer, - chat_template=chat_template, - chat_template_kwargs=chat_template_kwargs, - _trace=_trace, - ) + return _tokenize_chat_view( + history, + base_model=base_model, + tokenizer=tokenizer, + chat_template=chat_template, + chat_template_kwargs=chat_template_kwargs, + _projection_matches=( + True if _projection_validated else render_state.projection_matches + ), + _trace=_trace, + ) if isinstance(history, AnthropicMessagesHistory) and needs_render: converted = history.as_chat_completions_history() if ( diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index 86ee03f72..7b3304c98 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -1486,6 +1486,122 @@ def test_exact_chat_ids_accept_ordinary_positional_logprobs() -> None: assert tokenized.logprobs[1] == -0.75 +def test_chat_view_preserves_two_turn_textual_logprobs_without_exact_ids() -> None: + exchanges = [_chat_exchange([], [], offset=index) for index in range(2)] + for index, exchange in enumerate(exchanges): + choice = exchange.response.choices[0] + assert choice.model_extra is not None + choice.model_extra.pop("prompt_token_ids", None) + choice.model_extra.pop("token_ids", None) + assert choice.logprobs is not None + choice.logprobs = choice.logprobs.model_copy( + update={ + "content": [ + ChatCompletionTokenLogprob( + token="answer", + logprob=-0.4 - index / 10, + bytes=list(b"answer"), + top_logprobs=[], + ) + ] + } + ) + + class Tokenizer: + def apply_chat_template( + self, messages: list[dict[str, Any]], **kwargs: object + ) -> list[int]: + del kwargs + if len(messages) == 4: + return [10, 11, 20, 11] + if messages[-1]["role"] == "assistant": + return [10, 11] + return [10] + + def __call__(self, text: str, **kwargs: object) -> SimpleNamespace: + del text, kwargs + return SimpleNamespace(input_ids=[11]) + + history = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=exchanges) + ).chat_completions_history() + tokenized = history.tokenize(tokenizer=Tokenizer()) + + assert tokenized.token_ids == [10, 11, 20, 11] + assert tokenized.logprobs[1] == pytest.approx(-0.4) + assert tokenized.logprobs[3] == pytest.approx(-0.5) + assert tokenized.flags == [ + art.TokenFlag(0), + art.TokenFlag.SAMPLED, + art.TokenFlag(0), + art.TokenFlag.SAMPLED, + ] + + +def test_chat_view_uses_later_exact_prompt_when_first_is_missing() -> None: + first = _chat_exchange([], [11]) + second = _chat_exchange([10, 11, 20], [], offset=1) + second.response.choices[0].message.content = "final" + choice = second.response.choices[0] + assert choice.model_extra is not None + choice.model_extra.pop("token_ids", None) + assert choice.logprobs is not None + choice.logprobs = choice.logprobs.model_copy( + update={ + "content": [ + ChatCompletionTokenLogprob( + token="final", + logprob=-0.5, + bytes=list(b"final"), + top_logprobs=[], + ) + ] + } + ) + + class Tokenizer: + def apply_chat_template( + self, messages: list[dict[str, Any]], **kwargs: object + ) -> list[int]: + del kwargs + token_ids = { + "turn 0": [10], + "answer": [11], + "turn 1": [20], + "final": [21], + } + return [ + token + for message in messages + for token in token_ids[str(message["content"])] + ] + + def __call__(self, text: str, **kwargs: object) -> SimpleNamespace: + del kwargs + return SimpleNamespace( + input_ids={ + "turn 0": [10], + "answer": [11], + "turn 1": [20], + "final": [21], + }[text] + ) + + history = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[first, second]) + ).chat_completions_history() + tokenized = history.tokenize(tokenizer=Tokenizer()) + + assert tokenized.token_ids == [10, 11, 20, 21] + assert tokenized.logprobs[-1] == pytest.approx(-0.5) + assert tokenized.flags == [ + art.TokenFlag.EXACT, + art.TokenFlag.EXACT | art.TokenFlag.SAMPLED, + art.TokenFlag.EXACT, + art.TokenFlag.SAMPLED, + ] + + def test_chat_content_and_refusal_logprobs_are_combined_in_protocol_order() -> None: exchange = _chat_exchange([1], [2, 3]) data = exchange.response.model_dump(mode="python") @@ -1779,6 +1895,10 @@ def apply_chat_template( del kwargs return [10, 11] if messages[-1]["role"] == "assistant" else [10] + def __call__(self, text: str, **kwargs: object) -> list[int]: + del text, kwargs + return [11] + monkeypatch.setattr( "art.trajectories._tokenize._load_tokenizer", lambda _config: Tokenizer() ) @@ -1977,6 +2097,56 @@ def apply_chat_template( assert result.token_ids == [10, 11, 12] assert result.logprobs[1:] == [-0.1, -0.2] + assert result.flags[1:] == [ + art.TokenFlag.EXACT | art.TokenFlag.SAMPLED, + art.TokenFlag.EXACT | art.TokenFlag.SAMPLED, + ] + + +def test_responses_tool_output_source_is_not_sampled_without_generation() -> None: + from art.trajectories._tokenize import ( + _responses_output_is_sampled, + _source_is_sampled, + ) + + exchange = _response_exchange("response-tool-output", 0) + data = exchange.response.model_dump(mode="python") + data.pop("token_generations", None) + data["output"] = [ + { + "type": "function_call_output", + "id": "output-1", + "call_id": "call-1", + "output": "result", + "status": "completed", + } + ] + exchange.response = Response.model_validate(data) + + history = ( + art.Trajectory(exchanges=TrajectoryExchanges(responses=[exchange])) + .responses_history() + .as_chat_completions_history() + ) + assert history.messages[-1]["role"] == "tool" + source = history.message_sources[-1] + assert source is not None + assert not _source_is_sampled(source) + assert not _responses_output_is_sampled({"type": "function_call_output"}) + + +def test_responses_source_rejects_boolean_generation_index() -> None: + from art.trajectories._tokenize import _responses_source_generation + + exchange = _response_exchange("response-bool-generation", 2) + source = ChatCompletionsMessageSource( + exchange=exchange, + output_indices=(0,), + generation_index=True, + ) + + with pytest.raises(ValueError, match="generation index is invalid"): + _responses_source_generation(source) def test_responses_missing_token_generations_falls_back_for_visible_output( @@ -2556,9 +2726,17 @@ def test_prefix_retokenization_preserves_sampled_ids_and_logprobs( class Tokenizer: def __call__(self, text: str, *, add_special_tokens: bool = False) -> list[int]: - assert text == "cat" assert not add_special_tokens - return [500] + return {"cat": [500], "answer": [4], "turn 0": [1], "turn 1": [3]}[text] + + def apply_chat_template( + self, messages: list[dict[str, Any]], **kwargs: object + ) -> list[int]: + del kwargs + rendered: list[int] = [] + for message in messages: + rendered.extend(self(str(message["content"]))) + return rendered monkeypatch.setattr( "art.trajectories._tokenize._load_tokenizer", lambda _config: Tokenizer() @@ -3214,7 +3392,7 @@ def apply_chat_template( assert tokenized.logprobs[-2:] == [-0.7, -0.2] -def test_full_render_repair_rejects_non_append_only_templates() -> None: +def test_chat_view_preserves_initial_prompt_and_ignores_later_disagreement() -> None: first = _chat_exchange([1], [2]) second = _chat_exchange([9, 8, 7], [3], offset=1) trajectory = art.Trajectory( @@ -3224,16 +3402,19 @@ def test_full_render_repair_rejects_non_append_only_templates() -> None: class Tokenizer: def __call__(self, text: str, **kwargs: object) -> list[int]: del kwargs - return {"answer": [2]}[text] + return {"turn 0": [100], "answer": [2], "turn 1": [7]}[text] def apply_chat_template( self, messages: list[dict[str, Any]], **kwargs: object ) -> list[int]: del kwargs - return [100, len(messages)] + return [ + token for message in messages for token in self(str(message["content"])) + ] + + tokenized = trajectory.tokenize(tokenizer=Tokenizer()) - with pytest.raises(ValueError, match="do not form one append-only history"): - trajectory.tokenize(tokenizer=Tokenizer()) + assert tokenized.token_ids == [1, 2, 7, 3] def test_reasoning_stripped_chat_histories_tokenize_authoritative_views() -> None: From 7d868ad34c0737f30c2393e9a39e1d1d53f9a639 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Sat, 25 Jul 2026 01:58:49 +0000 Subject: [PATCH 49/58] Harden trajectory fallback token evidence --- src/art/trajectories/_tokenize.py | 76 ++++++--- tests/unit/trajectories/test_tokenize.py | 152 ++++++++++++++++-- .../trajectories/test_tokenized_models.py | 81 ++++++++++ 3 files changed, 282 insertions(+), 27 deletions(-) diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index da8250c0b..5fbd8c3fb 100644 --- a/src/art/trajectories/_tokenize.py +++ b/src/art/trajectories/_tokenize.py @@ -1389,10 +1389,15 @@ def _exchange_tokens( raise TypeError(f"Unknown exchange type: {type(exchange)!r}") -def _visible_logprobs(exchange: Exchange) -> list[tuple[str, float]]: +def _visible_logprobs( + exchange: Exchange, *, source: object | None = None +) -> list[tuple[str, float]]: values: list[tuple[str, float]] = [] if isinstance(exchange, ChatCompletionsExchange): - for entry in _chat_logprob_entries(exchange.response.choices[0]): + choice = ( + _chat_choice(source) if source is not None else exchange.response.choices[0] + ) + for entry in _chat_logprob_entries(choice): data = _dump(entry) raw_bytes = data.get("bytes") if isinstance(raw_bytes, list): @@ -1414,7 +1419,7 @@ def _visible_logprobs(exchange: Exchange) -> list[tuple[str, float]]: if logprob is not None: values.append((text, float(logprob))) elif isinstance(exchange, ResponsesExchange): - for output in _dump(exchange.response).get("output") or []: + for output in exchange.response.output: for content in _dump(output).get("content") or []: for entry in _dump(content).get("logprobs") or []: data = _dump(entry) @@ -1425,12 +1430,15 @@ def _visible_logprobs(exchange: Exchange) -> list[tuple[str, float]]: return values -def _sampled_text(exchange: Exchange) -> str | None: - visible = "".join(text for text, _ in _visible_logprobs(exchange)) +def _sampled_text(exchange: Exchange, *, source: object | None = None) -> str | None: + visible = "".join(text for text, _ in _visible_logprobs(exchange, source=source)) if visible: return visible if isinstance(exchange, ChatCompletionsExchange): - content = exchange.response.choices[0].message.content + choice = ( + _chat_choice(source) if source is not None else exchange.response.choices[0] + ) + content = choice.message.content return content if isinstance(content, str) else None if isinstance(exchange, CompletionsExchange): return exchange.response.choices[0].text @@ -1443,7 +1451,7 @@ def _sampled_text(exchange: Exchange) -> str | None: return "".join(parts) if parts else None if isinstance(exchange, ResponsesExchange): parts: list[str] = [] - for output in _dump(exchange.response).get("output") or []: + for output in exchange.response.output: data = _dump(output) if data.get("type") == "message": parts.append(_responses_output_text(data.get("content"))) @@ -1514,9 +1522,13 @@ def _retained_output_suffix( def _align_visible_logprobs( - tokenizer: Tokenizer | None, completion: list[int], exchange: Exchange + tokenizer: Tokenizer | None, + completion: list[int], + exchange: Exchange, + *, + source: object | None = None, ) -> list[float] | None: - values = _visible_logprobs(exchange) + values = _visible_logprobs(exchange, source=source) if not values or tokenizer is None: return None token_ids: list[int] = [] @@ -3005,7 +3017,7 @@ def _chat_source_tokens( tokens, logprobs = _chat_source_full_tokens(source) if tokens is None: return None, [] - sampled_text = _sampled_text(exchange) + sampled_text = _sampled_text(exchange, source=source) message = _dump( next( item @@ -3428,12 +3440,29 @@ def part_ids(text: str) -> list[int]: and not source_boundary and len(local_matches) != same_sampled_parts ): - if not source_context_matches: + prompt_render = render( + messages[:message_index], add_generation_prompt=True + ) + completed_render = render( + messages[: message_index + 1], add_generation_prompt=False + ) + bounded_matches = [ + match + for match in locations(local, len(prompt_render)) + if match[1] <= len(completed_render) + ] + if ( + len(completed_render) < len(prompt_render) + or completed_render[: len(prompt_render)] != prompt_render + or rendered[: len(completed_render)] != completed_render + or len(bounded_matches) != 1 + or bounded_matches[0][0] < search_cursor + ): raise ValueError( "Could not uniquely locate a sampled history message in the " "rendered history" ) - continue + span = bounded_matches[0] start, end = span search_cursor = end if not sampled: @@ -3454,7 +3483,12 @@ def part_ids(text: str) -> list[int]: exchange, (ChatCompletionsExchange, ResponsesExchange, MessagesExchange), ): - logprobs = _align_visible_logprobs(tokenizer, local, exchange) or [] + logprobs = ( + _align_visible_logprobs( + tokenizer, local, exchange, source=source + ) + or [] + ) replacements.append( ( start, @@ -3597,6 +3631,10 @@ def _tokenize_completions_token_history( prompt_logprobs if span.source.choice_index is None else completion_logprobs ) if selected != history.prompt[span.start : span.end]: + if span.source.choice_index is not None: + raise ValueError( + "Completions sampled output no longer matches its source exchange" + ) flags[span.start : span.end] = [ flag & ~TokenFlag.EXACT for flag in flags[span.start : span.end] ] @@ -4133,12 +4171,14 @@ def tokenize_group( ) for trajectory in group.trajectories ] - return cast( - TokenizedTrajectoryGroup[TokenizedTrajectory] - | TokenizedTrajectoryGroup[TokenizedMultiHistoryTrajectory], - TokenizedTrajectoryGroup( + if multi_history: + return TokenizedTrajectoryGroup[TokenizedMultiHistoryTrajectory]( trajectories=trajectories, metrics=dict(group.metrics), metadata=dict(group.metadata), - ), + ) + return TokenizedTrajectoryGroup[TokenizedTrajectory]( + trajectories=trajectories, + metrics=dict(group.metrics), + metadata=dict(group.metadata), ) diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index 7b3304c98..f53728ac1 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -726,6 +726,16 @@ def test_mutated_completions_token_prompt_drops_stale_exact_evidence() -> None: ] +def test_mutated_completions_token_choice_is_rejected() -> None: + history = art.Trajectory( + exchanges=TrajectoryExchanges(completions=[_completion_exchange()]) + ).completions_token_history() + history.prompt[-1] = 99 + + with pytest.raises(ValueError, match="sampled output"): + history.tokenize() + + def test_completions_history_rejects_model_and_sampled_span_mutation() -> None: history = art.Trajectory( exchanges=TrajectoryExchanges(completions=[_completion_exchange()]) @@ -1464,6 +1474,132 @@ def __call__(self, text: str, **kwargs: object) -> SimpleNamespace: assert math.isnan(result.logprobs[2]) +def test_each_chat_choice_preserves_its_visible_fallback_logprobs() -> None: + exchange = _chat_exchange([], []) + exchange.request["messages"] = [{"role": "user", "content": "question"}] + data = exchange.response.model_dump(mode="python") + choices = [] + for index, (text, logprob) in enumerate((("left", -0.1), ("right", -0.2))): + choice = data["choices"][0].copy() + choice.pop("prompt_token_ids", None) + choice.pop("token_ids", None) + choice["index"] = index + choice["message"] = {"role": "assistant", "content": text} + choice["logprobs"] = { + "content": [ + { + "token": text, + "logprob": logprob, + "bytes": list(text.encode()), + "top_logprobs": [], + } + ] + } + choices.append(choice) + data["choices"] = choices + exchange.response = ChatCompletion.model_validate(data) + trajectory = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[exchange]) + ) + + class Tokenizer: + def __call__(self, text: str, **kwargs: object) -> list[int]: + del kwargs + return {"question": [1], "left": [2], "right": [3]}[text] + + def apply_chat_template( + self, messages: list[dict[str, Any]], **kwargs: object + ) -> list[int]: + del kwargs + return [ + token for message in messages for token in self(str(message["content"])) + ] + + histories = trajectory.chat_completions_histories() + direct = [history.tokenize(tokenizer=Tokenizer()) for history in histories] + tokenized = trajectory.tokenize(multi_history=True, tokenizer=Tokenizer()) + + assert [history.logprobs[-1] for history in direct] == [-0.1, -0.2] + assert [history.logprobs[-1] for history in tokenized.histories] == [-0.1, -0.2] + + from art.trajectories._tokenize import _tokenize_trajectory_with_trace + + traced, traces = _tokenize_trajectory_with_trace(trajectory, tokenizer=Tokenizer()) + assert [history.logprobs[-1] for history in traced.histories] == [-0.1, -0.2] + assert [ + {key.index for key in trace.source_keys if key is not None} for trace in traces + ] == [{0}, {1}] + + +def test_chat_fallback_anchors_sampled_text_away_from_equal_user_text() -> None: + first = _chat_exchange([], [], offset=0) + first.request["messages"] = [{"role": "user", "content": "q"}] + first_data = first.response.model_dump(mode="python") + first_choice = first_data["choices"][0] + first_choice.pop("prompt_token_ids", None) + first_choice.pop("token_ids", None) + first_choice["message"]["content"] = "same" + first_choice["logprobs"]["content"] = [ + { + "token": "same", + "logprob": -0.1, + "bytes": list(b"same"), + "top_logprobs": [], + } + ] + first.response = ChatCompletion.model_validate(first_data) + + second = _chat_exchange([], [], offset=1) + second.request["messages"] = [ + {"role": "user", "content": "q"}, + {"role": "assistant", "content": "same"}, + {"role": "user", "content": "same"}, + ] + second_data = second.response.model_dump(mode="python") + second_choice = second_data["choices"][0] + second_choice.pop("prompt_token_ids", None) + second_choice.pop("token_ids", None) + second_choice["message"]["content"] = "other" + second_choice["logprobs"]["content"] = [ + { + "token": "other", + "logprob": -0.2, + "bytes": list(b"other"), + "top_logprobs": [], + } + ] + second.response = ChatCompletion.model_validate(second_data) + + class Tokenizer: + def __call__(self, text: str, **kwargs: object) -> list[int]: + del kwargs + return {"q": [1], "same": [2], "other": [3]}[text] + + def apply_chat_template( + self, messages: list[dict[str, Any]], **kwargs: object + ) -> list[int]: + del kwargs + return [ + token for message in messages for token in self(str(message["content"])) + ] + + history = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[first, second]) + ).chat_completions_history() + tokenized = history.tokenize(tokenizer=Tokenizer()) + + assert tokenized.token_ids == [1, 2, 2, 3] + assert tokenized.flags == [ + art.TokenFlag(0), + art.TokenFlag.SAMPLED, + art.TokenFlag(0), + art.TokenFlag.SAMPLED, + ] + assert tokenized.logprobs[1] == -0.1 + assert math.isnan(tokenized.logprobs[2]) + assert tokenized.logprobs[3] == -0.2 + + def test_exact_chat_ids_accept_ordinary_positional_logprobs() -> None: exchange = _chat_exchange([1], [2]) logprobs = exchange.response.choices[0].logprobs @@ -1777,7 +1913,7 @@ def __call__(self, text: str, **kwargs: object) -> list[int]: ] -def test_ambiguous_visible_logprobs_fail_closed( +def test_ambiguous_visible_logprobs_raise( monkeypatch: pytest.MonkeyPatch, ) -> None: exchange = _chat_exchange([], []) @@ -1810,14 +1946,12 @@ def __call__(self, text: str, **kwargs: object) -> SimpleNamespace: monkeypatch.setattr( "art.trajectories._tokenize._load_tokenizer", lambda _config: Tokenizer() ) - result = art.Trajectory( - exchanges=TrajectoryExchanges(chat_completions=[exchange]) - ).tokenize( - base_model="base/model", - ) - - assert result.token_ids == [10, 11, 12, 11] - assert all(math.isnan(logprob) for logprob in result.logprobs[1:]) + with pytest.raises(ValueError, match="uniquely locate"): + art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[exchange]) + ).tokenize( + base_model="base/model", + ) def test_legacy_token_and_logprob_length_mismatch_raises() -> None: diff --git a/tests/unit/trajectories/test_tokenized_models.py b/tests/unit/trajectories/test_tokenized_models.py index fc6e58b75..97c7eaca0 100644 --- a/tests/unit/trajectories/test_tokenized_models.py +++ b/tests/unit/trajectories/test_tokenized_models.py @@ -64,3 +64,84 @@ def test_nested_tokenized_models_nan_json_round_trip() -> None: assert restored_trajectory.reward == trajectory.reward assert restored.metrics == group.metrics assert restored.metadata == group.metadata + + +def test_public_group_tokenization_nan_json_round_trip() -> None: + from datetime import datetime + + from openai.types.chat import ChatCompletion + + from art.trajectories import ( + ChatCompletionsExchange, + ChatCompletionsRequest, + TrajectoryExchanges, + ) + + exchange = ChatCompletionsExchange( + request=ChatCompletionsRequest( + model="policy", + messages=[{"role": "user", "content": "question"}], + ), + response=ChatCompletion.model_validate( + { + "id": "chat", + "object": "chat.completion", + "created": 0, + "model": "policy", + "choices": [ + { + "index": index, + "finish_reason": "stop", + "message": {"role": "assistant", "content": text}, + "prompt_token_ids": [1], + "token_ids": [token_id], + "logprobs": { + "content": [ + { + "token": f"token_id:{token_id}", + "logprob": -0.1 * token_id, + "bytes": [], + "top_logprobs": [], + } + ] + }, + } + for index, (text, token_id) in enumerate( + (("left", 2), ("right", 3)) + ) + ], + } + ), + start_time=datetime(2026, 1, 1), + end_time=datetime(2026, 1, 1), + ) + group = art.TrajectoryGroup( + [art.Trajectory(exchanges=TrajectoryExchanges(chat_completions=[exchange]))] + ) + + single_exchange = exchange.model_copy( + update={ + "response": exchange.response.model_copy( + update={"choices": [exchange.response.choices[0]]} + ) + } + ) + single = art.TrajectoryGroup( + [ + art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[single_exchange]) + ) + ] + ).tokenize() + single_json = single.model_dump_json() + assert '"NaN"' in single_json + art.TokenizedTrajectoryGroup[art.TokenizedTrajectory].model_validate_json( + single_json + ) + + multi = group.tokenize(multi_history=True) + multi_json = multi.model_dump_json() + assert '"NaN"' in multi_json + art.TokenizedTrajectoryGroup[ + art.TokenizedMultiHistoryTrajectory + ].model_validate_json(multi_json) From 8e86252c31836bdface4650c5e01ed5dcc86b8d2 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Sat, 25 Jul 2026 01:55:42 +0000 Subject: [PATCH 50/58] fix trajectory requests stream transparency --- src/art/trajectories/_capture/requests.py | 4 ++-- tests/unit/trajectories/test_capture.py | 22 ++++++++++++++++++++++ 2 files changed, 24 insertions(+), 2 deletions(-) diff --git a/src/art/trajectories/_capture/requests.py b/src/art/trajectories/_capture/requests.py index 9af7bc94c..8ff248cb3 100644 --- a/src/art/trajectories/_capture/requests.py +++ b/src/art/trajectories/_capture/requests.py @@ -45,8 +45,8 @@ def iter_content( ): if state is not None: if isinstance(chunk, str): - chunk = chunk.encode(self.encoding or "utf-8") - if isinstance(chunk, bytes): + state.add(chunk.encode(self.encoding or "utf-8")) + elif isinstance(chunk, bytes): state.add(chunk) yield chunk completed = True diff --git a/tests/unit/trajectories/test_capture.py b/tests/unit/trajectories/test_capture.py index 4bd81758f..7a2d2a647 100644 --- a/tests/unit/trajectories/test_capture.py +++ b/tests/unit/trajectories/test_capture.py @@ -612,6 +612,28 @@ def consume(trajectory: art.Trajectory) -> None: assert len(trajectory.exchanges.chat_completions) == 1 +async def test_requests_decode_unicode_preserves_string_chunks( + endpoint_server: str, +) -> None: + body = {"model": "test/model", "messages": [{"role": "user", "content": "hi"}]} + + def consume() -> list[str | bytes]: + with requests.post( + f"{endpoint_server}/chat/completions", + json=body, + stream=True, + timeout=5, + ) as response: + return list(response.iter_content(chunk_size=5, decode_unicode=True)) + + with art.Trajectory() as trajectory: + chunks = await asyncio.to_thread(consume) + + assert chunks + assert all(isinstance(chunk, str) for chunk in chunks) + assert len(trajectory.exchanges.chat_completions) == 1 + + async def test_native_openai_and_anthropic_sdks(endpoint_server: str) -> None: openai = AsyncOpenAI(base_url=endpoint_server, api_key="test") anthropic = AsyncAnthropic( From 200304c1fec1b512c8639e1e2a087a38a13e3a2d Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Sat, 25 Jul 2026 02:12:29 +0000 Subject: [PATCH 51/58] Prove sampled chat evidence boundaries --- src/art/trajectories/_tokenize.py | 117 +++++++----- tests/unit/trajectories/test_tokenize.py | 231 ++++++++++++++++++++--- 2 files changed, 278 insertions(+), 70 deletions(-) diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index 5fbd8c3fb..2754b1865 100644 --- a/src/art/trajectories/_tokenize.py +++ b/src/art/trajectories/_tokenize.py @@ -1419,7 +1419,16 @@ def _visible_logprobs( if logprob is not None: values.append((text, float(logprob))) elif isinstance(exchange, ResponsesExchange): - for output in exchange.response.output: + outputs = exchange.response.output + if source is not None: + selected = _responses_source_outputs(source) + if selected is None: + return [] + selected_exchange, output_indices = selected + if selected_exchange is not exchange: + raise ValueError("Responses source belongs to a different exchange") + outputs = [exchange.response.output[index] for index in output_indices] + for output in outputs: for content in _dump(output).get("content") or []: for entry in _dump(content).get("logprobs") or []: data = _dump(entry) @@ -1450,8 +1459,17 @@ def _sampled_text(exchange: Exchange, *, source: object | None = None) -> str | ] return "".join(parts) if parts else None if isinstance(exchange, ResponsesExchange): + outputs = exchange.response.output + if source is not None: + selected = _responses_source_outputs(source) + if selected is None: + return None + selected_exchange, output_indices = selected + if selected_exchange is not exchange: + raise ValueError("Responses source belongs to a different exchange") + outputs = [exchange.response.output[index] for index in output_indices] parts: list[str] = [] - for output in exchange.response.output: + for output in outputs: data = _dump(output) if data.get("type") == "message": parts.append(_responses_output_text(data.get("content"))) @@ -3282,17 +3300,20 @@ def part_ids(text: str) -> list[int]: ) return part_ids_cache[text] - sampled_part_ids = [ - part_ids(text) - for message, source in zip( - history.messages, history.message_sources, strict=True - ) - if message.get("role") == "assistant" - and source is not None - and _source_is_sampled(source) - for _, text in _chat_message_parts(message) - ] - sampled_part_cursor = 0 + direct_render: list[int] = [] + direct_bounds: list[tuple[int, int]] = [] + for message in messages: + start = len(direct_render) + for _, text in _chat_message_parts(message): + direct_render.extend(part_ids(text)) + direct_bounds.append((start, len(direct_render))) + if direct_render != rendered and not ( + exact_prefix_length + and len(direct_render) == len(rendered) + and direct_render[exact_prefix_length:] == rendered[exact_prefix_length:] + ): + direct_bounds = [] + for message_index, (message, source) in enumerate( zip(history.messages, history.message_sources, strict=True) ): @@ -3322,12 +3343,45 @@ def part_ids(text: str) -> list[int]: ): search_cursor = max(search_cursor, len(source_prompt)) source_boundary = True + sampled_bounds: tuple[int, int] | None = None + if sampled and not source_boundary: + if direct_bounds: + sampled_bounds = direct_bounds[message_index] + else: + prompt_render = render( + messages[:message_index], add_generation_prompt=True + ) + completed_render = render( + messages[: message_index + 1], add_generation_prompt=False + ) + if ( + len(completed_render) < len(prompt_render) + or completed_render[: len(prompt_render)] != prompt_render + or rendered[: len(completed_render)] != completed_render + ): + raise ValueError( + "Could not locate a sampled history message in the " + "rendered history" + ) + sampled_bounds = (len(prompt_render), len(completed_render)) full_matches = ( locations(full_exact, search_cursor) if sampled and full_exact else [] ) first_part_matches = ( locations(part_ids(parts[0][1]), search_cursor) if parts else [] ) + if sampled_bounds is not None: + lower, upper = sampled_bounds + full_matches = [ + match + for match in full_matches + if match[0] >= lower and match[1] <= upper + ] + first_part_matches = [ + match + for match in first_part_matches + if match[0] >= lower and match[1] <= upper + ] if ( full_exact is not None and full_matches @@ -3357,7 +3411,6 @@ def part_ids(text: str) -> list[int]: ) ) search_cursor = end - sampled_part_cursor += len(parts) continue source_exchange = getattr(source, "exchange", None) @@ -3407,7 +3460,6 @@ def part_ids(text: str) -> list[int]: ): _warn_prefix_retokenization() search_cursor = len(completed_render) - sampled_part_cursor += len(parts) continue replacement_start = len(replacements) @@ -3415,16 +3467,6 @@ def part_ids(text: str) -> list[int]: if not sampled and text not in sampled_texts: continue local = part_ids(text) - same_sampled_parts = ( - sum( - candidate == local - for candidate in sampled_part_ids[sampled_part_cursor:] - ) - if sampled - else 0 - ) - if sampled: - sampled_part_cursor += 1 if not local: continue local_matches = locations(local, search_cursor) @@ -3435,29 +3477,14 @@ def part_ids(text: str) -> list[int]: raise ValueError( "Could not locate a history message in the rendered history" ) - if ( - sampled - and not source_boundary - and len(local_matches) != same_sampled_parts - ): - prompt_render = render( - messages[:message_index], add_generation_prompt=True - ) - completed_render = render( - messages[: message_index + 1], add_generation_prompt=False - ) + if sampled_bounds is not None: + lower, upper = sampled_bounds bounded_matches = [ match - for match in locations(local, len(prompt_render)) - if match[1] <= len(completed_render) + for match in locations(local, max(lower, search_cursor)) + if match[1] <= upper ] - if ( - len(completed_render) < len(prompt_render) - or completed_render[: len(prompt_render)] != prompt_render - or rendered[: len(completed_render)] != completed_render - or len(bounded_matches) != 1 - or bounded_matches[0][0] < search_cursor - ): + if len(bounded_matches) != 1: raise ValueError( "Could not uniquely locate a sampled history message in the " "rendered history" diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index f53728ac1..9deede862 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -1111,8 +1111,18 @@ def test_fallback_uses_template_overrides_and_nan_logprobs( assert loaded_base_models == ["base/model"] assert result.flags == [art.TokenFlag(0), art.TokenFlag.SAMPLED] assert math.isnan(result.logprobs[1]) - assert tokenizer.calls == [ + assert len(tokenizer.calls) == 3 + assert [call["add_generation_prompt"] for call in tokenizer.calls] == [ + False, + True, + False, + ] + assert all( { + **call, + "add_generation_prompt": False, + } + == { "tools": None, "tokenize": True, "add_generation_prompt": False, @@ -1121,8 +1131,9 @@ def test_fallback_uses_template_overrides_and_nan_logprobs( "explicit": True, "enable_thinking": True, "thinking_budget": 128, - }, - ] + } + for call in tokenizer.calls + ) @pytest.mark.parametrize( @@ -1474,6 +1485,81 @@ def __call__(self, text: str, **kwargs: object) -> SimpleNamespace: assert math.isnan(result.logprobs[2]) +def test_chat_fallback_rejects_unique_text_match_in_generation_scaffold() -> None: + exchange = _chat_exchange([], []) + exchange.request["messages"] = [{"role": "user", "content": "question"}] + choice = exchange.response.choices[0] + assert choice.model_extra is not None + choice.model_extra.pop("prompt_token_ids", None) + choice.model_extra.pop("token_ids", None) + assert choice.logprobs is not None + choice.logprobs = choice.logprobs.model_copy( + update={ + "content": [ + ChatCompletionTokenLogprob( + token="answer", + logprob=-0.7, + bytes=list(b"answer"), + top_logprobs=[], + ) + ] + } + ) + + class Tokenizer: + def apply_chat_template( + self, + messages: list[dict[str, Any]], + *, + add_generation_prompt: bool, + **kwargs: object, + ) -> list[int]: + del kwargs + if messages[-1]["role"] == "assistant": + return [1, 11, 12] + return [1, 11] if add_generation_prompt else [1] + + def __call__(self, text: str, **kwargs: object) -> SimpleNamespace: + del kwargs + return SimpleNamespace(input_ids={"question": [1], "answer": [11]}[text]) + + history = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[exchange]) + ).chat_completions_history() + + with pytest.raises(ValueError, match="uniquely locate"): + history.tokenize(tokenizer=Tokenizer()) + + +def test_chat_exact_ids_reject_unique_match_in_generation_scaffold() -> None: + exchange = _chat_exchange([], [11]) + exchange.request["messages"] = [{"role": "user", "content": "question"}] + + class Tokenizer: + def apply_chat_template( + self, + messages: list[dict[str, Any]], + *, + add_generation_prompt: bool, + **kwargs: object, + ) -> list[int]: + del kwargs + if messages[-1]["role"] == "assistant": + return [1, 11, 12] + return [1, 11] if add_generation_prompt else [1] + + def __call__(self, text: str, **kwargs: object) -> SimpleNamespace: + del kwargs + return SimpleNamespace(input_ids={"question": [1], "answer": [11]}[text]) + + history = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[exchange]) + ).chat_completions_history() + + with pytest.raises(ValueError, match="uniquely locate"): + history.tokenize(tokenizer=Tokenizer()) + + def test_each_chat_choice_preserves_its_visible_fallback_logprobs() -> None: exchange = _chat_exchange([], []) exchange.request["messages"] = [{"role": "user", "content": "question"}] @@ -3314,6 +3400,59 @@ def apply_chat_template( ] +def test_responses_multi_output_chat_conversion_preserves_item_logprobs() -> None: + projected = _multi_output_responses_chat_history() + exchange = next( + source.exchange + for source in projected.message_sources + if source is not None and source.output_indices == (0,) + ) + assert isinstance(exchange, ResponsesExchange) + data = exchange.response.model_dump(mode="python") + data.pop("token_generations") + for output, (text, logprob) in zip( + data["output"], (("first", -0.1), ("second", -0.2)), strict=True + ): + output["content"][0]["logprobs"] = [ + { + "token": text, + "logprob": logprob, + "bytes": list(text.encode()), + "top_logprobs": [], + } + ] + exchange.response = Response.model_validate(data) + history = ( + art.Trajectory(exchanges=TrajectoryExchanges(responses=[exchange])) + .responses_history() + .as_chat_completions_history() + ) + + class Tokenizer: + def __call__(self, text: str, **kwargs: object) -> list[int]: + del kwargs + return {"turn 0": [1], "first": [2], "second": [3]}[text] + + def apply_chat_template( + self, messages: list[dict[str, Any]], **kwargs: object + ) -> list[int]: + del kwargs + return [ + token for message in messages for token in self(str(message["content"])) + ] + + tokenized = history.tokenize(tokenizer=Tokenizer()) + + assert tokenized.token_ids == [1, 2, 3] + assert math.isnan(tokenized.logprobs[0]) + assert tokenized.logprobs[1:] == [-0.1, -0.2] + assert tokenized.flags == [ + art.TokenFlag(0), + art.TokenFlag.SAMPLED, + art.TokenFlag.SAMPLED, + ] + + def test_mutable_chat_history_is_authoritative_and_does_not_replay_removed_turns() -> ( None ): @@ -3337,14 +3476,18 @@ def __call__(self, text: str, **kwargs: object) -> list[int]: return {"second": [31], "answer": [41]}[text] def apply_chat_template( - self, messages: list[dict[str, Any]], **kwargs: object + self, + messages: list[dict[str, Any]], + *, + add_generation_prompt: bool, + **kwargs: object, ) -> list[int]: del kwargs - assert [message["content"] for message in messages] == [ - "second", - "answer", - ] - return [30, 41, 50] + contents = [message["content"] for message in messages] + if contents == ["second", "answer"]: + return [30, 41, 50] + assert contents == ["second"] + return [30] if add_generation_prompt else [30] tokenized = history.tokenize(tokenizer=Tokenizer()) @@ -3374,10 +3517,17 @@ def __call__(self, text: str, **kwargs: object) -> list[int]: return {"seed": [20], "question": [30], "answer": [41]}[text] def apply_chat_template( - self, messages: list[dict[str, Any]], **kwargs: object + self, + messages: list[dict[str, Any]], + *, + add_generation_prompt: bool, + **kwargs: object, ) -> list[int]: - del messages, kwargs - return [5, 20, 6, 30, 7, 41, 8] + del kwargs + if messages[-1]["role"] == "assistant" and len(messages) == 3: + return [5, 20, 6, 30, 7, 41, 8] + assert add_generation_prompt + return [5, 20, 6, 30, 7] tokenized = history.tokenize(tokenizer=Tokenizer()) @@ -3402,10 +3552,17 @@ def __call__(self, text: str, **kwargs: object) -> list[int]: return [7] def apply_chat_template( - self, messages: list[dict[str, Any]], **kwargs: object + self, + messages: list[dict[str, Any]], + *, + add_generation_prompt: bool, + **kwargs: object, ) -> list[int]: - del messages, kwargs - return [7, 99, 7] + del kwargs + if messages[-1]["role"] == "assistant": + return [7, 99, 7] + assert add_generation_prompt + return [7, 99] tokenized = history.tokenize(tokenizer=Tokenizer()) @@ -3457,10 +3614,16 @@ def __call__(self, text: str, **kwargs: object) -> list[int]: return {"cat": [7], "dog": [500]}[text] def apply_chat_template( - self, messages: list[dict[str, Any]], **kwargs: object + self, + messages: list[dict[str, Any]], + *, + add_generation_prompt: bool, + **kwargs: object, ) -> list[int]: - del messages, kwargs - return [42, 7, 99, 500] + del kwargs + if messages[-1]["role"] == "assistant": + return [42, 7, 99, 500] + return [42, 7, 99] if add_generation_prompt else [42, 7] tokenized = history.tokenize(tokenizer=Tokenizer()) @@ -3510,10 +3673,17 @@ def __call__(self, text: str, **kwargs: object) -> list[int]: return {"question": [1], "answer": [7]}[text] def apply_chat_template( - self, messages: list[dict[str, Any]], **kwargs: object + self, + messages: list[dict[str, Any]], + *, + add_generation_prompt: bool, + **kwargs: object, ) -> list[int]: - del messages, kwargs - return [1, 99, 7, 2] + del kwargs + if messages[-1]["role"] == "assistant": + return [1, 99, 7, 2] + assert add_generation_prompt + return [1, 99] tokenized = history.tokenize(tokenizer=Tokenizer()) @@ -3889,10 +4059,17 @@ def __call__(self, text: str, **kwargs: object) -> list[int]: }[text] def apply_chat_template( - self, messages: list[dict[str, Any]], **kwargs: object + self, + messages: list[dict[str, Any]], + *, + add_generation_prompt: bool, + **kwargs: object, ) -> list[int]: - del messages, kwargs - return [1, 10, 20, 25, 30, 31, 26] + del kwargs + if messages[-1]["role"] == "assistant": + return [1, 10, 20, 25, 30, 31, 26] + assert add_generation_prompt + return [1, 10] tokenized = history.tokenize(tokenizer=Tokenizer()) @@ -4044,12 +4221,16 @@ def apply_chat_template( self, messages: list[dict[str, Any]], *, + add_generation_prompt: bool, chat_template: str | None = None, **kwargs: object, ) -> list[int]: del kwargs assert chat_template == "custom" - return [10, 20, 30] + if messages[-1]["role"] == "assistant": + return [10, 20, 30] + assert add_generation_prompt + return [10] tokenized = trajectory.tokenize( tokenizer=Tokenizer(), From 32afa2dec46ba8400b2568895240c4bb02f7a026 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Sat, 25 Jul 2026 02:57:35 +0000 Subject: [PATCH 52/58] Harden sampled token attribution --- src/art/trajectories/_tokenize.py | 681 ++++++++++++++++++++--- tests/unit/trajectories/test_tokenize.py | 567 +++++++++++++++++-- 2 files changed, 1129 insertions(+), 119 deletions(-) diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index 2754b1865..5e2404236 100644 --- a/src/art/trajectories/_tokenize.py +++ b/src/art/trajectories/_tokenize.py @@ -2,6 +2,7 @@ from bisect import bisect_left from collections.abc import Hashable, Mapping, Sequence +from copy import deepcopy from dataclasses import dataclass from datetime import datetime from functools import lru_cache @@ -9,7 +10,7 @@ import json import math import re -from typing import TYPE_CHECKING, Any, Literal, TypeVar, cast +from typing import TYPE_CHECKING, Any, Literal, Protocol, TypeVar, cast import warnings from anthropic.types import Message, MessageParam, TextBlock @@ -61,6 +62,16 @@ class _TokenizerConfig: chat_template_kwargs: Mapping[str, object] | None = None +class _OffsetTokenizer(Protocol): + def __call__( + self, + text: str, + *, + add_special_tokens: bool, + return_offsets_mapping: Literal[True], + ) -> object: ... + + @dataclass(frozen=True) class _SampledOutput: text: str | None @@ -331,7 +342,7 @@ def _as_tokenizer(tokenizer: object) -> Tokenizer: def _string_dict(value: object) -> dict[str, Any] | None: - if not isinstance(value, dict) or not all(isinstance(key, str) for key in value): + if not isinstance(value, Mapping) or not all(isinstance(key, str) for key in value): return None return {key: item for key, item in value.items() if isinstance(key, str)} @@ -452,18 +463,10 @@ def _chat_logprob_entries(choice: Choice) -> list[object]: ] -def _chat_choice_tokens( - choice: Choice, response_data: dict[str, Any] -) -> tuple[list[int] | None, list[int] | None, list[float]]: +def _chat_choice_output_tokens( + choice: Choice, +) -> tuple[list[int] | None, list[float]]: choice_data = _dump(choice) - prompt = choice_data.get("prompt_token_ids") - if prompt is None: - prompt = response_data.get("prompt_token_ids") - prompt_ids = _exact_token_ids( - prompt, - field="Chat Completions prompt_token_ids", - empty_is_missing=True, - ) token_ids = _exact_token_ids( choice_data.get("token_ids"), field="Chat Completions token_ids", @@ -482,10 +485,26 @@ def _chat_choice_tokens( logprobs = pair_logprobs or positional_logprobs if selected is not None and values and len(logprobs) != len(selected): raise ValueError("Chat Completions token IDs and logprobs differ in length") + return selected, logprobs or [math.nan] * len(selected or []) + + +def _chat_choice_tokens( + choice: Choice, response_data: dict[str, Any] +) -> tuple[list[int] | None, list[int] | None, list[float]]: + choice_data = _dump(choice) + prompt = choice_data.get("prompt_token_ids") + if prompt is None: + prompt = response_data.get("prompt_token_ids") + prompt_ids = _exact_token_ids( + prompt, + field="Chat Completions prompt_token_ids", + empty_is_missing=True, + ) + selected, logprobs = _chat_choice_output_tokens(choice) return ( prompt_ids, selected, - logprobs or [math.nan] * len(selected or []), + logprobs, ) @@ -1545,10 +1564,25 @@ def _align_visible_logprobs( exchange: Exchange, *, source: object | None = None, + sampled_text: str | None = None, ) -> list[float] | None: values = _visible_logprobs(exchange, source=source) if not values or tokenizer is None: return None + if sampled_text is not None: + matches: list[list[tuple[str, float]]] = [] + for start in range(len(values)): + combined = "" + for end in range(start, len(values)): + combined += values[end][0] + if combined == sampled_text: + matches.append(values[start : end + 1]) + break + if len(combined) >= len(sampled_text): + break + if len(matches) != 1: + return None + values = matches[0] token_ids: list[int] = [] logprobs: list[float] = [] for text, logprob in values: @@ -2385,6 +2419,27 @@ class _HistoryRenderState: projection_matches: bool | None = None +def _matches_final_chat_exchange(history: ChatCompletionsHistory) -> bool | None: + from ._history import normalize_chat_message + + for source in reversed(history.message_sources): + if ( + source is None + or not isinstance(source.exchange, ChatCompletionsExchange) + or source.choice_index is None + ): + continue + choice = _chat_choice(source) + expected = [ + *source.exchange.request.get("messages", []), + choice.message.model_dump(mode="python", exclude_none=True), + ] + return [normalize_chat_message(message) for message in history.messages] == [ + normalize_chat_message(message) for message in expected + ] + return None + + def _history_render_state(history: History) -> _HistoryRenderState: if isinstance(history, ChatCompletionsHistory): if any(source is None for source in history.message_sources): @@ -2409,7 +2464,9 @@ def _history_render_state(history: History) -> _HistoryRenderState: ) if context_changed: return _HistoryRenderState(needs_render=True, context_changed=True) - projection_matches = _history_matches_projection(history) + projection_matches = _matches_final_chat_exchange(history) + if projection_matches is None: + projection_matches = _history_matches_projection(history) return _HistoryRenderState( needs_render=not projection_matches, projection_matches=projection_matches, @@ -2845,8 +2902,7 @@ def _chat_source_full_tokens( choice = next( item for item in exchange.response.choices if item.index == choice_index ) - _, tokens, logprobs = _chat_choice_tokens(choice, _dump(exchange.response)) - return tokens, logprobs + return _chat_choice_output_tokens(choice) if isinstance(exchange, ResponsesExchange): return _responses_generation_full_tokens(source) if isinstance(exchange, MessagesExchange): @@ -3024,15 +3080,61 @@ def _chat_message_parts(message: Mapping[str, object]) -> list[tuple[str, str]]: return parts +def _chat_message_text_slot_groups( + message: dict[str, Any], +) -> list[list[tuple[dict[str, Any], str]]]: + groups: list[list[tuple[dict[str, Any], str]]] = [] + reasoning_key = ( + "reasoning" + if isinstance(message.get("reasoning"), str) and message["reasoning"] + else "reasoning_content" + ) + if isinstance(message.get(reasoning_key), str) and message[reasoning_key]: + groups.append([(message, reasoning_key)]) + content = message.get("content") + if isinstance(content, str) and content: + groups.append([(message, "content")]) + elif isinstance(content, list): + content_slots = [ + (block, "text") + for block in content + if ( + isinstance(block, dict) + and block.get("type") in {"input_text", "output_text", "text"} + and isinstance(block.get("text"), str) + and block["text"] + ) + ] + if content_slots: + groups.append(content_slots) + if isinstance(message.get("refusal"), str) and message["refusal"]: + groups.append([(message, "refusal")]) + tool_calls = message.get("tool_calls") + if isinstance(tool_calls, list): + for call in tool_calls: + if not isinstance(call, dict): + continue + function = call.get("function") + if not isinstance(function, dict): + continue + for key in ("name", "arguments"): + if isinstance(function.get(key), str) and function[key]: + groups.append([(function, key)]) + return groups + + def _chat_source_tokens( source: object, text: str, *, part: str, + full_tokens: tuple[list[int] | None, list[float]] | None = None, ) -> tuple[list[int] | None, list[float]]: exchange = getattr(source, "exchange", None) if isinstance(exchange, ChatCompletionsExchange): - tokens, logprobs = _chat_source_full_tokens(source) + tokens, logprobs = ( + full_tokens if full_tokens is not None else _chat_source_full_tokens(source) + ) if tokens is None: return None, [] sampled_text = _sampled_text(exchange, source=source) @@ -3212,10 +3314,89 @@ def render( ) ) - rendered = render( - messages, - add_generation_prompt=not ends_with_assistant, - ) + rendered = render(messages, add_generation_prompt=not ends_with_assistant) + if any( + message.get("role") == "assistant" + and isinstance(message.get("reasoning"), str) + and message["reasoning"] + and not message.get("reasoning_content") + for message in messages + ): + without_reasoning = deepcopy(messages) + for message in without_reasoning: + message.pop("reasoning", None) + if ( + render( + without_reasoning, + add_generation_prompt=not ends_with_assistant, + ) + == rendered + ): + messages = deepcopy(messages) + for message in messages: + reasoning = message.pop("reasoning", None) + if isinstance(reasoning, str) and reasoning: + message.setdefault("reasoning_content", reasoning) + rendered = render( + messages, + add_generation_prompt=not ends_with_assistant, + ) + if any( + message.get("role") == "assistant" + and isinstance(message.get("refusal"), str) + and message["refusal"] + for message in messages + ): + without_refusals = deepcopy(messages) + for message in without_refusals: + message.pop("refusal", None) + if ( + render( + without_refusals, + add_generation_prompt=not ends_with_assistant, + ) + == rendered + ): + messages = deepcopy(messages) + for message in messages: + refusal = message.pop("refusal", None) + if not isinstance(refusal, str) or not refusal: + continue + content = message.get("content") + if isinstance(content, str): + message["content"] = content + refusal + elif isinstance(content, list): + message["content"] = [ + *content, + {"type": "text", "text": refusal}, + ] + elif content is None: + message["content"] = refusal + else: + raise ValueError( + "Cannot render an assistant refusal with this content shape" + ) + rendered = render( + messages, + add_generation_prompt=not ends_with_assistant, + ) + + prompt_cache: dict[int, list[int] | None] = {} + output_cache: dict[int, tuple[list[int] | None, list[float]]] = {} + + def source_prompt_tokens(source: object) -> list[int] | None: + key = id(source) + if key not in prompt_cache: + prompt_cache[key] = _chat_source_prompt_tokens(source) + return prompt_cache[key] + + def source_output_tokens( + source: object, + ) -> tuple[list[int] | None, list[float]]: + key = id(source) + if key not in output_cache: + output_cache[key] = _chat_source_full_tokens(source) + return output_cache[key] def source_matches_context(source: object) -> bool: exchange = getattr(source, "exchange", None) @@ -3238,7 +3419,7 @@ def source_matches_context(source: object) -> bool: ): if message.get("role") != "assistant" or source is None: continue - source_prompt = _chat_source_prompt_tokens(source) + source_prompt = source_prompt_tokens(source) if source_prompt and source_matches_context(source): rendered_prompt = render( messages[:message_index], add_generation_prompt=True @@ -3314,6 +3495,286 @@ def part_ids(text: str) -> list[int]: ): direct_bounds = [] + marked_bounds: dict[int, tuple[int, int]] = {} + marked_part_bounds: dict[int, list[tuple[int, int]]] = {} + if not direct_bounds: + marked_messages = deepcopy(messages) + marker_prefix = f"ART_TRAJECTORY_{id(marked_messages):x}_" + markers: dict[str, tuple[int, int, Literal["start", "end"]]] = {} + part_counts: dict[int, int] = {} + part_whitespace: dict[tuple[int, int], tuple[str, str]] = {} + for message_index, (message, source) in enumerate( + zip(marked_messages, history.message_sources, strict=True) + ): + if ( + message.get("role") != "assistant" + or source is None + or not _source_is_sampled(source) + ): + continue + slot_groups = _chat_message_text_slot_groups(message) + if not slot_groups: + continue + part_counts[message_index] = len(slot_groups) + for part_index, slots in enumerate(slot_groups): + start = f"{marker_prefix}{message_index}_{part_index}_START" + end = f"{marker_prefix}{message_index}_{part_index}_END" + first, first_key = slots[0] + last, last_key = slots[-1] + first_text = str(first[first_key]) + last_text = str(last[last_key]) + leading = first_text[: len(first_text) - len(first_text.lstrip())] + trailing = last_text[len(last_text.rstrip()) :] + if first is last and first_key == last_key: + core = first_text[ + len(leading) : len(first_text) - len(trailing) + if trailing + else len(first_text) + ] + first[first_key] = leading + start + core + end + trailing + else: + first[first_key] = leading + start + first_text[len(leading) :] + last[last_key] = ( + last_text[: len(last_text) - len(trailing)] + end + trailing + ) + part_whitespace[(message_index, part_index)] = (leading, trailing) + markers[start] = (message_index, part_index, "start") + markers[end] = (message_index, part_index, "end") + if markers: + try: + marked_text = resolved_tokenizer.apply_chat_template( + marked_messages, + tools=history.tools, + tokenize=False, + add_generation_prompt=not ends_with_assistant, + **({"chat_template": template} if template is not None else {}), + **kwargs, + ) + except (TypeError, NotImplementedError): + marked_text = None + if isinstance(marked_text, str): + marker_pattern = re.compile( + rf"{re.escape(marker_prefix)}\d+_\d+_(?:START|END)" + ) + matches = list(marker_pattern.finditer(marked_text)) + found_markers = [match.group(0) for match in matches] + else: + matches = [] + found_markers = [] + if ( + isinstance(marked_text, str) + and len(found_markers) == len(markers) + and set(found_markers) == set(markers) + ): + unmarked_parts: list[str] = [] + char_bounds: dict[tuple[int, int], list[int]] = {} + source_cursor = 0 + target_cursor = 0 + for match in matches: + position = match.start() + marker = match.group(0) + message_index, part_index, boundary = markers[marker] + chunk = marked_text[source_cursor:position] + unmarked_parts.append(chunk) + target_cursor += len(chunk) + char_bounds.setdefault((message_index, part_index), [0, 0])[ + 0 if boundary == "start" else 1 + ] = target_cursor + source_cursor = match.end() + unmarked_parts.append(marked_text[source_cursor:]) + unmarked_text = "".join(unmarked_parts) + for key, bounds in char_bounds.items(): + leading, trailing = part_whitespace[key] + if ( + leading + and unmarked_text[max(0, bounds[0] - len(leading)) : bounds[0]] + == leading + ): + bounds[0] -= len(leading) + if ( + trailing + and unmarked_text[bounds[1] : bounds[1] + len(trailing)] + == trailing + ): + bounds[1] += len(trailing) + try: + encoded = cast(_OffsetTokenizer, resolved_tokenizer)( + unmarked_text, + add_special_tokens=False, + return_offsets_mapping=True, + ) + except (TypeError, NotImplementedError): + encoded = None + encoded_data = _string_dict(encoded) + raw_offsets = ( + encoded_data.get("offset_mapping") + if encoded_data is not None + else None + ) + if ( + encoded is not None + and _ids(encoded) == rendered + and isinstance(raw_offsets, list) + and len(raw_offsets) == len(rendered) + ): + offsets: list[tuple[int, int]] = [] + for value in raw_offsets: + if ( + not isinstance(value, (list, tuple)) + or len(value) != 2 + or not all(isinstance(item, int) for item in value) + ): + break + offsets.append((value[0], value[1])) + if len(offsets) == len(rendered): + token_bounds: dict[tuple[int, int], tuple[int, int]] = {} + token_cursor = 0 + for key, (char_start, char_end) in sorted( + char_bounds.items(), key=lambda item: item[1] + ): + while ( + token_cursor < len(offsets) + and offsets[token_cursor][1] <= char_start + ): + token_cursor += 1 + token_end = token_cursor + while ( + token_end < len(offsets) + and offsets[token_end][0] < char_end + ): + token_end += 1 + if ( + token_end > token_cursor + and offsets[token_cursor][0] >= char_start + and offsets[token_end - 1][1] <= char_end + ): + token_bounds[key] = (token_cursor, token_end) + token_cursor = token_end + for message_index, part_count in part_counts.items(): + bounds = [ + token_bounds[(message_index, part_index)] + for part_index in range(part_count) + if (message_index, part_index) in token_bounds + ] + if len(bounds) == part_count: + marked_part_bounds[message_index] = bounds + marked_bounds[message_index] = ( + bounds[0][0], + bounds[-1][1], + ) + + def differing_span(probe: Sequence[int]) -> tuple[int, int] | None: + prefix = 0 + while ( + prefix < len(rendered) + and prefix < len(probe) + and rendered[prefix] == probe[prefix] + ): + prefix += 1 + suffix = 0 + while ( + suffix < len(rendered) - prefix + and suffix < len(probe) - prefix + and rendered[-suffix - 1] == probe[-suffix - 1] + ): + suffix += 1 + end = len(rendered) - suffix + return (prefix, end) if prefix < end else None + + probed_bounds: dict[int, tuple[int, int]] = {} + if not direct_bounds: + for message_index, (message, source) in enumerate( + zip(messages, history.message_sources, strict=True) + ): + if ( + message_index in marked_bounds + or message.get("role") != "assistant" + or source is None + or not _source_is_sampled(source) + ): + continue + parts = _chat_message_parts(message) + if parts and all(part == "tool_call" for part, _ in parts): + probe_messages = deepcopy(messages) + tool_calls = probe_messages[message_index].get("tool_calls") + if isinstance(tool_calls, list): + for call_index, call in enumerate(tool_calls): + if not isinstance(call, dict): + continue + function = call.get("function") + if not isinstance(function, dict): + continue + function["name"] = f"art_trajectory_probe_{call_index}" + function["arguments"] = '{"art_trajectory_probe":true}' + try: + probe = render( + probe_messages, + add_generation_prompt=not ends_with_assistant, + ) + except Exception: + probe = [] + if span := differing_span(probe): + probed_bounds[message_index] = span + continue + if len(parts) == 1: + try: + prefix = render( + messages[:message_index], add_generation_prompt=True + ) + except Exception: + prefix = [] + local = part_ids(parts[0][1]) + if rendered == [*prefix, *local]: + probed_bounds[message_index] = ( + len(prefix), + len(rendered), + ) + continue + try: + completed = render( + messages[: message_index + 1], + add_generation_prompt=False, + ) + except Exception: + completed = [] + if ( + completed == [*prefix, *local] + and rendered[: len(completed)] == completed + ): + probed_bounds[message_index] = ( + len(prefix), + len(completed), + ) + continue + probe_messages = deepcopy(messages) + slot_groups = _chat_message_text_slot_groups(probe_messages[message_index]) + if len(slot_groups) != 1 or len(slot_groups[0]) != 1: + continue + container, key = slot_groups[0][0] + original = str(container[key]) + leading = original[: len(original) - len(original.lstrip())] + trailing = original[len(original.rstrip()) :] + container[key] = ( + leading + f"ART_TRAJECTORY_{id(probe_messages):x}_PROBE" + trailing + ) + try: + probe = render( + probe_messages, + add_generation_prompt=not ends_with_assistant, + ) + except Exception: + continue + if span := differing_span(probe): + probed_bounds[message_index] = span + + sampled_message_count = sum( + message.get("role") == "assistant" + and source is not None + and _source_is_sampled(source) + for message, source in zip( + history.messages, history.message_sources, strict=True + ) + ) for message_index, (message, source) in enumerate( zip(history.messages, history.message_sources, strict=True) ): @@ -3324,46 +3785,58 @@ def part_ids(text: str) -> list[int]: and _source_is_sampled(source) ) full_exact, full_logprobs = ( - _chat_source_full_tokens(source) if sampled else (None, []) + source_output_tokens(source) if sampled else (None, []) ) + if sampled and not parts and not full_exact: + continue complete_sampled_message = ( sampled and source is not None and _source_covers_complete_sampled_message(message, source) ) source_boundary = False - if sampled and source is not None: - source_prompt = _chat_source_prompt_tokens(source) - source_exchange = getattr(source, "exchange", None) - source_context_matches = source_matches_context(source) - if ( - source_context_matches - and source_prompt - and rendered[: len(source_prompt)] == source_prompt - ): - search_cursor = max(search_cursor, len(source_prompt)) - source_boundary = True + generation_start: int | None = None sampled_bounds: tuple[int, int] | None = None - if sampled and not source_boundary: + content_bounds_proven = False + if sampled: if direct_bounds: sampled_bounds = direct_bounds[message_index] + content_bounds_proven = True + elif message_index in marked_bounds: + sampled_bounds = marked_bounds[message_index] + content_bounds_proven = True + elif message_index in probed_bounds: + sampled_bounds = probed_bounds[message_index] + content_bounds_proven = True else: - prompt_render = render( - messages[:message_index], add_generation_prompt=True - ) - completed_render = render( - messages[: message_index + 1], add_generation_prompt=False - ) + assert source is not None + source_prompt = source_prompt_tokens(source) + source_context_matches = source_matches_context(source) if ( - len(completed_render) < len(prompt_render) - or completed_render[: len(prompt_render)] != prompt_render - or rendered[: len(completed_render)] != completed_render + source_context_matches + and source_prompt + and rendered[: len(source_prompt)] == source_prompt ): + search_cursor = max(search_cursor, len(source_prompt)) + source_boundary = True + generation_start = len(source_prompt) + sampled_bounds = (generation_start, len(rendered)) + elif sampled_message_count == 1: + prompt_render = render( + messages[:message_index], add_generation_prompt=True + ) + if rendered[: len(prompt_render)] != prompt_render: + raise ValueError( + "Could not locate a sampled history message in the " + "rendered history" + ) + generation_start = len(prompt_render) + sampled_bounds = (generation_start, len(rendered)) + else: raise ValueError( - "Could not locate a sampled history message in the " - "rendered history" + "Could not prove a sampled history message boundary with this " + "tokenizer" ) - sampled_bounds = (len(prompt_render), len(completed_render)) full_matches = ( locations(full_exact, search_cursor) if sampled and full_exact else [] ) @@ -3418,6 +3891,33 @@ def part_ids(text: str) -> list[int]: isinstance(source_exchange, ResponsesExchange) and len(_response_generations(source_exchange.response)) > 1 ) + if sampled and not content_bounds_proven: + raise ValueError( + "Could not uniquely locate or prove the sampled content boundary " + "with this tokenizer" + ) + if ( + sampled + and full_exact is None + and message_index in probed_bounds + and parts + and all(part == "tool_call" for part, _ in parts) + ): + assert source is not None and sampled_bounds is not None + start, end = sampled_bounds + replacements.append( + ( + start, + end, + rendered[start:end], + [math.nan] * (end - start), + False, + _sampled_source_key(source), + source, + ) + ) + search_cursor = end + continue if ( complete_sampled_message and full_exact is not None @@ -3427,23 +3927,38 @@ def part_ids(text: str) -> list[int]: or len(full_exact) != len(part_ids(parts[0][1])) ) ): - prompt_render = render(messages[:message_index], add_generation_prompt=True) - completed_render = render( - messages[: message_index + 1], add_generation_prompt=False - ) - if ( - len(completed_render) < len(prompt_render) - or (parts and len(completed_render) == len(prompt_render)) - or completed_render[: len(prompt_render)] != prompt_render - or rendered[: len(completed_render)] != completed_render - ): + if not parts and sampled_bounds is not None: + start = end = sampled_bounds[0] + elif sampled_bounds is None or sampled_bounds[0] == sampled_bounds[1]: raise ValueError( "Could not locate a complete sampled message in the rendered history" ) + else: + start, end = sampled_bounds + if parts and len(parts) == 1 and parts[0][0] == "content": + visible_matches = [ + match + for match in locations(part_ids(parts[0][1]), start) + if match[1] <= end + ] + if len(visible_matches) != 1 or ( + not content_bounds_proven + and generation_start is not None + and visible_matches[0][0] != generation_start + ): + raise ValueError( + "Could not prove the sampled content boundary in the " + "rendered history" + ) + start, end = visible_matches[0] + elif generation_start is not None and ( + multi_generation_response or len(parts) != 1 or parts[0][0] != "content" + ): + start = generation_start replacements.append( ( - len(prompt_render), - len(completed_render), + start, + end, full_exact, full_logprobs if len(full_logprobs) == len(full_exact) @@ -3456,21 +3971,25 @@ def part_ids(text: str) -> list[int]: if ( message_index < len(messages) - 1 and isinstance(source_exchange, ChatCompletionsExchange) - and completed_render[len(prompt_render) :] != full_exact + and rendered[start:end] != full_exact ): _warn_prefix_retokenization() - search_cursor = len(completed_render) + search_cursor = end continue replacement_start = len(replacements) - for part, text in parts: + for part_index, (part, text) in enumerate(parts): if not sampled and text not in sampled_texts: continue local = part_ids(text) if not local: continue - local_matches = locations(local, search_cursor) - span = local_matches[0] if local_matches else None + proven_part_bounds = marked_part_bounds.get(message_index) + span = ( + proven_part_bounds[part_index] + if proven_part_bounds is not None + else next(iter(locations(local, search_cursor)), None) + ) if span is None: if not sampled: continue @@ -3479,17 +3998,18 @@ def part_ids(text: str) -> list[int]: ) if sampled_bounds is not None: lower, upper = sampled_bounds - bounded_matches = [ - match - for match in locations(local, max(lower, search_cursor)) - if match[1] <= upper - ] - if len(bounded_matches) != 1: - raise ValueError( - "Could not uniquely locate a sampled history message in the " - "rendered history" - ) - span = bounded_matches[0] + if proven_part_bounds is None: + bounded_matches = [ + match + for match in locations(local, max(lower, search_cursor)) + if match[1] <= upper + ] + if len(bounded_matches) != 1: + raise ValueError( + "Could not uniquely locate a sampled history message in " + "the rendered history" + ) + span = bounded_matches[0] start, end = span search_cursor = end if not sampled: @@ -3499,11 +4019,12 @@ def part_ids(text: str) -> list[int]: source, text, part=part, + full_tokens=(full_exact, full_logprobs), ) if exact is not None and rendered[start : start + len(exact)] == exact: end = start + len(exact) search_cursor = end - replacement = exact if exact is not None else local + replacement = exact if exact is not None else rendered[start:end] if exact is None and not logprobs: exchange = getattr(source, "exchange", None) if isinstance( @@ -3512,7 +4033,11 @@ def part_ids(text: str) -> list[int]: ): logprobs = ( _align_visible_logprobs( - tokenizer, local, exchange, source=source + tokenizer, + replacement, + exchange, + source=source, + sampled_text=text, ) or [] ) diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index 9deede862..e5fb8aaba 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -4,6 +4,7 @@ from datetime import datetime, timedelta import math import random +import re from statistics import median import sys from time import perf_counter @@ -1114,18 +1115,18 @@ def test_fallback_uses_template_overrides_and_nan_logprobs( assert len(tokenizer.calls) == 3 assert [call["add_generation_prompt"] for call in tokenizer.calls] == [ False, - True, False, + True, ] + assert [call["tokenize"] for call in tokenizer.calls] == [True, False, True] assert all( { - **call, - "add_generation_prompt": False, + key: value + for key, value in call.items() + if key not in {"add_generation_prompt", "tokenize"} } == { "tools": None, - "tokenize": True, - "add_generation_prompt": False, "chat_template": "explicit-template", "request": True, "explicit": True, @@ -1466,11 +1467,17 @@ def apply_chat_template( self, messages: list[dict[str, Any]], **kwargs: object ) -> list[int]: del kwargs - return [10, 11, 12] if messages[-1]["role"] == "assistant" else [10] + if messages[-1]["role"] != "assistant": + return [10] + if str(messages[-1]["content"]).startswith("ART_TRAJECTORY_"): + return [10, 99, 12] + return [10, 11, 12] def __call__(self, text: str, **kwargs: object) -> SimpleNamespace: - del text, kwargs - return SimpleNamespace(input_ids=[11]) + del kwargs + return SimpleNamespace( + input_ids={"turn 0": [10], "answer": [11], "turn 1": [20]}[text] + ) monkeypatch.setattr( "art.trajectories._tokenize._load_tokenizer", lambda _config: Tokenizer() @@ -1560,6 +1567,78 @@ def __call__(self, text: str, **kwargs: object) -> SimpleNamespace: history.tokenize(tokenizer=Tokenizer()) +def test_chat_prompt_ids_do_not_bind_output_to_trailing_scaffold() -> None: + exchange = _chat_exchange([1, 2], []) + choice = exchange.response.choices[0] + assert choice.logprobs is not None + choice.logprobs = choice.logprobs.model_copy( + update={ + "content": [ + ChatCompletionTokenLogprob( + token="answer", + logprob=-0.7, + bytes=list(b"answer"), + top_logprobs=[], + ) + ] + } + ) + + class Tokenizer: + def apply_chat_template( + self, messages: list[dict[str, Any]], **kwargs: object + ) -> list[int]: + del kwargs + return [1, 2, 12, 11] if messages[-1]["role"] == "assistant" else [1, 2] + + def __call__(self, text: str, **kwargs: object) -> SimpleNamespace: + del kwargs + return SimpleNamespace( + input_ids={"turn 0": [10], "answer": [11], "turn 1": [20]}[text] + ) + + history = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[exchange]) + ).chat_completions_history() + + with pytest.raises(ValueError, match="content boundary"): + history.tokenize(tokenizer=Tokenizer()) + + +def test_chat_exact_hidden_suffix_preserves_rendered_trailing_scaffold() -> None: + exchange = _chat_exchange([1, 99], [7, 8]) + history = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[exchange]) + ).chat_completions_history() + history.chat_template = "rerender" + + class Tokenizer: + def apply_chat_template( + self, messages: list[dict[str, Any]], **kwargs: object + ) -> list[int]: + del kwargs + if messages[-1]["role"] != "assistant": + return [1, 99] + if str(messages[-1]["content"]).startswith("ART_TRAJECTORY_"): + return [1, 99, 999, 9] + return [1, 99, 7, 9] + + def __call__(self, text: str, **kwargs: object) -> SimpleNamespace: + del text, kwargs + return SimpleNamespace(input_ids=[7]) + + tokenized = history.tokenize(tokenizer=Tokenizer()) + + assert tokenized.token_ids == [1, 99, 7, 8, 9] + assert tokenized.flags == [ + art.TokenFlag(0), + art.TokenFlag(0), + art.TokenFlag.EXACT | art.TokenFlag.SAMPLED, + art.TokenFlag.EXACT | art.TokenFlag.SAMPLED, + art.TokenFlag(0), + ] + + def test_each_chat_choice_preserves_its_visible_fallback_logprobs() -> None: exchange = _chat_exchange([], []) exchange.request["messages"] = [{"role": "user", "content": "question"}] @@ -1741,8 +1820,10 @@ def apply_chat_template( return [10] def __call__(self, text: str, **kwargs: object) -> SimpleNamespace: - del text, kwargs - return SimpleNamespace(input_ids=[11]) + del kwargs + return SimpleNamespace( + input_ids={"turn 0": [10], "answer": [11], "turn 1": [20]}[text] + ) history = art.Trajectory( exchanges=TrajectoryExchanges(chat_completions=exchanges) @@ -2999,7 +3080,11 @@ def apply_chat_template( ) -> list[int]: del tools, tokenize, add_generation_prompt, kwargs assert chat_template == "custom" - return [10] if len(messages) == 1 else [10, 20, 30] + if len(messages) == 1: + return [10] + if str(messages[-1]["content"]).startswith("ART_TRAJECTORY_"): + return [10, 999, 30] + return [10, 20, 30] tokenized = history.tokenize(tokenizer=Tokenizer()) @@ -3026,15 +3111,19 @@ def apply_chat_template( del kwargs if messages[-1]["role"] != "assistant": return [1] + if str(messages[-1]["content"]).startswith("ART_TRAJECTORY_"): + return [1, 999, 9] return [1, 2, 9] tokenized = history.tokenize(tokenizer=Tokenizer()) - assert tokenized.token_ids == [1, 2, 3] - assert tokenized.logprobs[1:] == pytest.approx([-0.2, -0.3]) + assert tokenized.token_ids == [1, 2, 3, 9] + assert tokenized.logprobs[1:3] == pytest.approx([-0.2, -0.3]) + assert math.isnan(tokenized.logprobs[3]) assert tokenized.flags[1:] == [ art.TokenFlag.EXACT | art.TokenFlag.SAMPLED, art.TokenFlag.EXACT | art.TokenFlag.SAMPLED, + art.TokenFlag(0), ] @@ -3077,17 +3166,51 @@ def test_responses_generation_evidence_is_atomic_and_partial_edits_do_not_replay exchange.response = Response.model_validate(data) class Tokenizer: - def __call__(self, text: str, **kwargs: object) -> list[int]: - del kwargs + def __call__(self, text: str, **kwargs: object) -> object: + if kwargs.get("return_offsets_mapping"): + answer_start = text.index("answer") + answer_end = answer_start + len("answer") + offsets: list[tuple[int, int]] + token_ids = [1] + if "think" in text: + think_start = text.index("think") + think_end = think_start + len("think") + offsets = [ + (0, think_start), + (think_start, think_end), + (answer_start, answer_end), + (answer_end, len(text)), + ] + token_ids.extend([20, 30, 9]) + else: + offsets = [ + (0, answer_start), + (answer_start, answer_end), + (answer_end, len(text)), + ] + token_ids.extend([30, 9]) + return {"input_ids": token_ids, "offset_mapping": offsets} return {"turn 0": [1], "think": [20], "answer": [30]}[text] def apply_chat_template( self, messages: list[dict[str, Any]], **kwargs: object - ) -> list[int]: + ) -> object: + tokenize = kwargs.pop("tokenize") del kwargs if messages[-1]["role"] != "assistant": return [1] assistant = messages[-1] + rendered = ( + f"{messages[0]['content']}" + + ( + f"{assistant['reasoning']}" + if assistant.get("reasoning") + else "" + ) + + f"{assistant['content']}" + ) + if not tokenize: + return rendered if assistant.get("reasoning"): return [1, 20, 30, 9] return [1, 30, 9] @@ -3096,8 +3219,9 @@ def apply_chat_template( exchanges=TrajectoryExchanges(responses=[exchange]) ).responses_history() exact = history.tokenize(tokenizer=Tokenizer(), chat_template="custom") - assert exact.token_ids == [1, 2, 3] - assert exact.logprobs[1:] == pytest.approx([-0.2, -0.3]) + assert exact.token_ids == [1, 2, 3, 9] + assert exact.logprobs[1:3] == pytest.approx([-0.2, -0.3]) + assert math.isnan(exact.logprobs[3]) del history.input[1] del history.input_sources[1] @@ -3200,24 +3324,36 @@ def test_responses_chat_rerender_preserves_equal_length_generation_evidence() -> ) class Tokenizer: - def __call__(self, text: str, **kwargs: object) -> list[int]: - del kwargs - return {"turn 0": [10], "answer": [20], "second": [40]}[text] + def __call__(self, text: str, **kwargs: object) -> object: + mapping = {"turn 0": 10, "answer": 20, "second": 40} + if not kwargs.get("return_offsets_mapping"): + return [mapping[text]] + token_ids: list[int] = [] + offsets: list[tuple[int, int]] = [] + for match in re.finditer(r"(.*?)|", text): + if match.group(1) is None: + token_ids.append(30) + offsets.append(match.span()) + else: + token_ids.append(mapping[match.group(1)]) + offsets.append(match.span(1)) + return {"input_ids": token_ids, "offset_mapping": offsets} def apply_chat_template( self, messages: list[dict[str, Any]], *, add_generation_prompt: bool, + tokenize: bool, **kwargs: object, - ) -> list[int]: + ) -> object: del kwargs - result: list[int] = [] + result = "" for index, message in enumerate(messages): - result.extend(self(str(message["content"]))) + result += f"{message['content']}" if index == 1 and (len(messages) > 2 or add_generation_prompt): - result.append(30) - return result + result += "" + return self(result, return_offsets_mapping=True) if tokenize else result tokenized = history.tokenize(tokenizer=Tokenizer(), chat_template="custom") @@ -3486,6 +3622,8 @@ def apply_chat_template( contents = [message["content"] for message in messages] if contents == ["second", "answer"]: return [30, 41, 50] + if len(contents) == 2 and str(contents[-1]).startswith("ART_TRAJECTORY_"): + return [30, 99, 50] assert contents == ["second"] return [30] if add_generation_prompt else [30] @@ -3525,6 +3663,8 @@ def apply_chat_template( ) -> list[int]: del kwargs if messages[-1]["role"] == "assistant" and len(messages) == 3: + if str(messages[-1]["content"]).startswith("ART_TRAJECTORY_"): + return [5, 20, 6, 30, 7, 99, 8] return [5, 20, 6, 30, 7, 41, 8] assert add_generation_prompt return [5, 20, 6, 30, 7] @@ -4050,8 +4190,25 @@ def test_rerender_marks_tool_call_only_generated_region_sampled() -> None: history.chat_template = "rerender" class Tokenizer: - def __call__(self, text: str, **kwargs: object) -> list[int]: - del kwargs + def __call__(self, text: str, **kwargs: object) -> object: + if kwargs.get("return_offsets_mapping"): + name_start = text.index("lookup") + name_end = name_start + len("lookup") + args_start = text.index('{"x":1}') + args_end = args_start + len('{"x":1}') + midpoint = args_start + 3 + return { + "input_ids": [1, 10, 20, 25, 30, 31, 26], + "offset_mapping": [ + (0, text.index("")), + (text.index(""), name_start), + (name_start, name_end), + (name_end, args_start), + (args_start, midpoint), + (midpoint, args_end), + (args_end, len(text)), + ], + } return { "turn 0": [1], "lookup": [20], @@ -4064,9 +4221,18 @@ def apply_chat_template( *, add_generation_prompt: bool, **kwargs: object, - ) -> list[int]: + ) -> object: + tokenize = kwargs.pop("tokenize") del kwargs if messages[-1]["role"] == "assistant": + function = messages[-1]["tool_calls"][0]["function"] + rendered = ( + f"{messages[0]['content']}" + f"{function['name']}" + f"{function['arguments']}" + ) + if not tokenize: + return rendered return [1, 10, 20, 25, 30, 31, 26] assert add_generation_prompt return [1, 10] @@ -4079,6 +4245,297 @@ def apply_chat_template( assert not tokenized.flags[0] & art.TokenFlag.SAMPLED +def test_tool_call_probe_handles_contextual_tokenization() -> None: + exchange = _chat_exchange([], []) + data = exchange.response.model_dump(mode="python") + choice = data["choices"][0] + choice.pop("prompt_token_ids") + choice.pop("token_ids") + choice["logprobs"] = None + choice["message"] = { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call-1", + "type": "function", + "function": {"name": "lookup", "arguments": "{}"}, + } + ], + } + exchange.response = ChatCompletion.model_validate(data) + history = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[exchange]) + ).chat_completions_history() + + class Tokenizer: + def __call__(self, text: str, **kwargs: object) -> list[int]: + del kwargs + return {"turn 0": [1], "lookup": [90], "{}": [91]}[text] + + def apply_chat_template( + self, messages: list[dict[str, Any]], **kwargs: object + ) -> list[int]: + del kwargs + if messages[-1]["role"] != "assistant": + return [1, 10] + function = messages[-1]["tool_calls"][0]["function"] + if str(function["name"]).startswith("art_trajectory_probe_"): + return [1, 10, 98, 25, 99, 26] + return [1, 10, 20, 25, 30, 26] + + tokenized = history.tokenize(tokenizer=Tokenizer()) + + assert tokenized.token_ids == [1, 10, 20, 25, 30, 26] + assert tokenized.flags[2:5] == [art.TokenFlag.SAMPLED] * 3 + assert all(math.isnan(value) for value in tokenized.logprobs[2:5]) + + +def test_rerender_proves_each_reasoning_and_content_part_separately() -> None: + exchange = _chat_exchange([], []) + data = exchange.response.model_dump(mode="python") + choice = data["choices"][0] + choice.pop("prompt_token_ids") + choice.pop("token_ids") + choice["message"] = { + "role": "assistant", + "reasoning": "think", + "content": "answer", + } + choice["logprobs"] = { + "content": [ + { + "token": "answer", + "logprob": -0.7, + "bytes": list(b"answer"), + "top_logprobs": [], + } + ], + "refusal": None, + } + exchange.response = ChatCompletion.model_validate(data) + history = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[exchange]) + ).chat_completions_history() + history.chat_template = "rerender" + + class Tokenizer: + def __call__(self, text: str, **kwargs: object) -> object: + if kwargs.get("return_offsets_mapping"): + think_start = text.index("think") + think_end = think_start + len("think") + answer_start = text.index("answer") + answer_end = answer_start + len("answer") + return { + "input_ids": [1, 10, 11, 12, 13, 14], + "offset_mapping": [ + (0, text.index("")), + (text.index(""), think_start), + (think_start, think_end), + (think_end, answer_start), + (answer_start, answer_end), + (answer_end, len(text)), + ], + } + return {"turn 0": [1], "think": [11], "answer": [13]}[text] + + def apply_chat_template( + self, messages: list[dict[str, Any]], **kwargs: object + ) -> object: + tokenize = kwargs.pop("tokenize") + del kwargs + assistant = messages[-1] + reasoning = assistant.get("reasoning") or assistant.get("reasoning_content") + rendered = ( + f"{messages[0]['content']}" + + (f"{reasoning}" if reasoning else "") + + f"{assistant['content']}" + ) + if not tokenize: + return rendered + if not reasoning: + return [1, 10, 13, 14] + return [1, 10, 11, 12, 13, 14] + + tokenized = history.tokenize(tokenizer=Tokenizer()) + + assert tokenized.flags == [ + art.TokenFlag(0), + art.TokenFlag(0), + art.TokenFlag.SAMPLED, + art.TokenFlag(0), + art.TokenFlag.SAMPLED, + art.TokenFlag(0), + ] + assert tokenized.logprobs[4] == -0.7 + + +def test_rerender_rejects_unproved_multi_part_boundaries() -> None: + exchange = _chat_exchange([], []) + data = exchange.response.model_dump(mode="python") + data["choices"][0].pop("prompt_token_ids") + data["choices"][0].pop("token_ids") + data["choices"][0]["logprobs"] = None + data["choices"][0]["message"] = { + "role": "assistant", + "reasoning": "think", + "content": "answer", + } + exchange.response = ChatCompletion.model_validate(data) + history = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[exchange]) + ).chat_completions_history() + history.chat_template = "rerender" + + class Tokenizer: + def __call__(self, text: str, **kwargs: object) -> list[int]: + del kwargs + return {"turn 0": [1], "think": [11], "answer": [12]}[text] + + def apply_chat_template( + self, messages: list[dict[str, Any]], **kwargs: object + ) -> list[int]: + del messages, kwargs + return [1, 10, 11, 12, 13, 14] + + with pytest.raises(ValueError, match="sampled content boundary"): + history.tokenize(tokenizer=Tokenizer()) + + +def test_empty_sampled_messages_need_no_content_boundary() -> None: + exchange = _chat_exchange([], []) + data = exchange.response.model_dump(mode="python") + data["choices"][0].pop("prompt_token_ids") + data["choices"][0].pop("token_ids") + data["choices"][0]["logprobs"] = None + data["choices"][0]["message"]["content"] = "" + exchange.response = ChatCompletion.model_validate(data) + history = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[exchange]) + ).chat_completions_history() + + class Tokenizer: + def __call__(self, text: str, **kwargs: object) -> list[int]: + del text, kwargs + return [1] + + def apply_chat_template( + self, messages: list[dict[str, Any]], **kwargs: object + ) -> list[int]: + del messages, kwargs + return [1, 2] + + tokenized = history.tokenize(tokenizer=Tokenizer()) + + assert tokenized.token_ids == [1, 2] + assert tokenized.flags == [art.TokenFlag(0), art.TokenFlag(0)] + + +def test_renderer_ignored_refusal_is_appended_for_tokenization() -> None: + exchange = _chat_exchange([], []) + data = exchange.response.model_dump(mode="python") + choice = data["choices"][0] + choice.pop("prompt_token_ids") + choice.pop("token_ids") + choice["message"] = { + "role": "assistant", + "content": "answer", + "refusal": "declined", + } + choice["logprobs"] = { + "content": [ + { + "token": "answer", + "logprob": -0.4, + "bytes": list(b"answer"), + "top_logprobs": [], + } + ], + "refusal": [ + { + "token": "declined", + "logprob": -0.5, + "bytes": list(b"declined"), + "top_logprobs": [], + } + ], + } + exchange.response = ChatCompletion.model_validate(data) + history = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[exchange]) + ).chat_completions_history() + + class Tokenizer: + def __call__(self, text: str, **kwargs: object) -> list[int]: + del kwargs + return { + "turn 0": [1], + "answer": [4], + "declined": [5], + "answerdeclined": [4, 5], + }[text] + + def apply_chat_template( + self, messages: list[dict[str, Any]], **kwargs: object + ) -> list[int]: + del kwargs + return [ + token + for message in messages + for token in self(str(message.get("content") or "")) + ] + + tokenized = history.tokenize(tokenizer=Tokenizer()) + + assert tokenized.token_ids == [1, 4, 5] + assert tokenized.logprobs[1:] == [-0.4, -0.5] + assert tokenized.flags[1:] == [art.TokenFlag.SAMPLED] * 2 + + +def test_renderer_reasoning_content_alias_preserves_reasoning() -> None: + exchange = _chat_exchange([], []) + data = exchange.response.model_dump(mode="python") + choice = data["choices"][0] + choice.pop("prompt_token_ids") + choice.pop("token_ids") + choice["logprobs"] = None + choice["message"] = { + "role": "assistant", + "reasoning": "think", + "content": "answer", + } + exchange.response = ChatCompletion.model_validate(data) + history = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[exchange]) + ).chat_completions_history() + + class Tokenizer: + def __call__(self, text: str, **kwargs: object) -> list[int]: + del kwargs + return {"turn 0": [1], "think": [2], "answer": [3]}[text] + + def apply_chat_template( + self, messages: list[dict[str, Any]], **kwargs: object + ) -> list[int]: + del kwargs + result: list[int] = [] + for message in messages: + if reasoning := message.get("reasoning_content"): + result.extend(self(str(reasoning))) + if content := message.get("content"): + result.extend(self(str(content))) + return result + + tokenized = history.tokenize(tokenizer=Tokenizer()) + + assert tokenized.token_ids == [1, 2, 3] + assert tokenized.flags == [ + art.TokenFlag(0), + art.TokenFlag.SAMPLED, + art.TokenFlag.SAMPLED, + ] + + def _repeated_text_rerender_history(turn_count: int) -> art.ChatCompletionsHistory: exchanges: list[ChatCompletionsExchange] = [] prompt: list[int] = [] @@ -4102,20 +4559,46 @@ class _RepeatedTextTokenizer: def __init__(self) -> None: self.apply_calls = 0 - def __call__(self, text: str, **kwargs: object) -> list[int]: - del kwargs - return [1000] if text == "answer" else [2000 + int(text[1:])] + def __call__(self, text: str, **kwargs: object) -> object: + if not kwargs.get("return_offsets_mapping"): + return [1000] if text == "answer" else [2000 + int(text[1:])] + token_ids: list[int] = [] + offsets: list[tuple[int, int]] = [] + for match in re.finditer(r"<([ua])>(.*?)", text): + content_start, content_end = match.span(2) + token_ids.extend( + [ + 3000 if match.group(1) == "u" else 3001, + 1000 + if match.group(2) == "answer" + else 2000 + int(match.group(2)[1:]), + 3002, + ] + ) + offsets.extend( + [ + (match.start(), content_start), + (content_start, content_end), + (content_end, match.end()), + ] + ) + return {"input_ids": token_ids, "offset_mapping": offsets} def apply_chat_template( - self, messages: list[dict[str, Any]], **kwargs: object - ) -> list[int]: + self, + messages: list[dict[str, Any]], + *, + tokenize: bool, + **kwargs: object, + ) -> object: del kwargs self.apply_calls += 1 - result: list[int] = [] - for message in messages: - content = str(message["content"]) - result.extend([1000] if content == "answer" else [2000 + int(content[1:])]) - return result + rendered = "".join( + f"<{'a' if message['role'] == 'assistant' else 'u'}>" + f"{message['content']}" + for message in messages + ) + return self(rendered, return_offsets_mapping=True) if tokenize else rendered def test_rerender_calls_chat_template_once_for_many_turns() -> None: @@ -4124,7 +4607,7 @@ def test_rerender_calls_chat_template_once_for_many_turns() -> None: tokenizer = _RepeatedTextTokenizer() history.tokenize(tokenizer=tokenizer) - assert tokenizer.apply_calls == 1 + assert tokenizer.apply_calls == 2 def test_repeated_text_rerender_scaling_is_near_linear() -> None: @@ -4137,7 +4620,7 @@ def test_repeated_text_rerender_scaling_is_near_linear() -> None: started = perf_counter() history.tokenize(tokenizer=tokenizer) samples.append(perf_counter() - started) - assert tokenizer.apply_calls == 1 + assert tokenizer.apply_calls == 2 medians.append(median(samples)) assert medians[1] < medians[0] * 3 @@ -4228,6 +4711,8 @@ def apply_chat_template( del kwargs assert chat_template == "custom" if messages[-1]["role"] == "assistant": + if str(messages[-1]["content"]).startswith("ART_TRAJECTORY_"): + return [10, 999, 30] return [10, 20, 30] assert add_generation_prompt return [10] From d99a480cff397f14e51ea58f17378bedb55dced3 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Sat, 25 Jul 2026 03:01:35 +0000 Subject: [PATCH 53/58] Preserve textual logprob token identity --- src/art/trajectories/_tokenize.py | 57 +++++++++++++++----- tests/unit/trajectories/test_tokenize.py | 68 ++++++++++++++++++++++++ 2 files changed, 111 insertions(+), 14 deletions(-) diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index 5e2404236..ec93d59b9 100644 --- a/src/art/trajectories/_tokenize.py +++ b/src/art/trajectories/_tokenize.py @@ -1558,14 +1558,13 @@ def _retained_output_suffix( return None -def _align_visible_logprobs( +def _visible_token_evidence( tokenizer: Tokenizer | None, - completion: list[int], exchange: Exchange, *, source: object | None = None, sampled_text: str | None = None, -) -> list[float] | None: +) -> tuple[list[int], list[float]] | None: values = _visible_logprobs(exchange, source=source) if not values or tokenizer is None: return None @@ -1593,6 +1592,26 @@ def _align_visible_logprobs( return None token_ids.append(encoded[0]) logprobs.append(logprob) + return token_ids, logprobs + + +def _align_visible_logprobs( + tokenizer: Tokenizer | None, + completion: list[int], + exchange: Exchange, + *, + source: object | None = None, + sampled_text: str | None = None, +) -> list[float] | None: + evidence = _visible_token_evidence( + tokenizer, + exchange, + source=source, + sampled_text=sampled_text, + ) + if evidence is None: + return None + token_ids, logprobs = evidence left: list[int] = [] cursor = 0 @@ -3701,9 +3720,10 @@ def differing_span(probe: Sequence[int]) -> tuple[int, int] | None: for call_index, call in enumerate(tool_calls): if not isinstance(call, dict): continue - function = call.get("function") - if not isinstance(function, dict): + raw_function = call.get("function") + if not isinstance(raw_function, dict): continue + function = cast(dict[str, Any], raw_function) function["name"] = f"art_trajectory_probe_{call_index}" function["arguments"] = '{"art_trajectory_probe":true}' try: @@ -4031,16 +4051,25 @@ def differing_span(probe: Sequence[int]) -> tuple[int, int] | None: exchange, (ChatCompletionsExchange, ResponsesExchange, MessagesExchange), ): - logprobs = ( - _align_visible_logprobs( - tokenizer, - replacement, - exchange, - source=source, - sampled_text=text, - ) - or [] + evidence = _visible_token_evidence( + tokenizer, + exchange, + source=source, + sampled_text=text, ) + if evidence is not None: + replacement, logprobs = evidence + else: + logprobs = ( + _align_visible_logprobs( + tokenizer, + replacement, + exchange, + source=source, + sampled_text=text, + ) + or [] + ) replacements.append( ( start, diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index e5fb8aaba..c43d4e637 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -4536,6 +4536,74 @@ def apply_chat_template( ] +def test_trimmed_render_preserves_authoritative_textual_logprob_tokens() -> None: + exchange = _chat_exchange([], []) + data = exchange.response.model_dump(mode="python") + choice = data["choices"][0] + choice.pop("prompt_token_ids") + choice.pop("token_ids") + choice["message"]["content"] = " first " + choice["logprobs"] = { + "content": [ + { + "token": " first", + "logprob": -0.4, + "bytes": list(b" first"), + "top_logprobs": [], + }, + { + "token": " ", + "logprob": -0.5, + "bytes": [32], + "top_logprobs": [], + }, + ], + "refusal": None, + } + exchange.response = ChatCompletion.model_validate(data) + history = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[exchange]) + ).chat_completions_history() + + class Tokenizer: + def __call__(self, text: str, **kwargs: object) -> object: + if kwargs.get("return_offsets_mapping"): + content_start = text.index("first") + content_end = content_start + len("first") + return { + "input_ids": [1, 3765], + "offset_mapping": [ + (0, content_start), + (content_start, content_end), + ], + } + return { + "turn 0": [1], + " first ": [1118, 220], + " first": [1118], + " ": [220], + }[text] + + def apply_chat_template( + self, messages: list[dict[str, Any]], **kwargs: object + ) -> object: + tokenize = kwargs.pop("tokenize") + del kwargs + rendered = ( + f"{messages[0]['content']}" + f"{str(messages[-1]['content']).strip()}" + ) + if not tokenize: + return rendered + return [1, 3765] + + tokenized = history.tokenize(tokenizer=Tokenizer()) + + assert tokenized.token_ids == [1, 1118, 220] + assert tokenized.logprobs[1:] == [-0.4, -0.5] + assert tokenized.flags[1:] == [art.TokenFlag.SAMPLED] * 2 + + def _repeated_text_rerender_history(turn_count: int) -> art.ChatCompletionsHistory: exchanges: list[ChatCompletionsExchange] = [] prompt: list[int] = [] From 6e067a39bae86d0d8e0e672e19dcfa69016e4aa2 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Sat, 25 Jul 2026 03:03:46 +0000 Subject: [PATCH 54/58] Harden fallback probes --- src/art/trajectories/_tokenize.py | 32 ++++++++++++----- tests/unit/trajectories/test_tokenize.py | 44 ++++++++++++++++++++---- 2 files changed, 60 insertions(+), 16 deletions(-) diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index ec93d59b9..599648c7d 100644 --- a/src/art/trajectories/_tokenize.py +++ b/src/art/trajectories/_tokenize.py @@ -3483,9 +3483,7 @@ def locations(needle: Sequence[int], start: int) -> list[tuple[int, int]]: search_cursor = 0 sampled_texts = { text - for message, source in zip( - history.messages, history.message_sources, strict=True - ) + for message, source in zip(messages, history.message_sources, strict=True) if message.get("role") == "assistant" and source is not None and _source_is_sampled(source) @@ -3717,6 +3715,14 @@ def differing_span(probe: Sequence[int]) -> tuple[int, int] | None: probe_messages = deepcopy(messages) tool_calls = probe_messages[message_index].get("tool_calls") if isinstance(tool_calls, list): + existing_values = { + value + for call in tool_calls + if isinstance(call, dict) + and isinstance(call.get("function"), dict) + for value in cast(dict[str, Any], call["function"]).values() + if isinstance(value, str) + } for call_index, call in enumerate(tool_calls): if not isinstance(call, dict): continue @@ -3724,8 +3730,18 @@ def differing_span(probe: Sequence[int]) -> tuple[int, int] | None: if not isinstance(raw_function, dict): continue function = cast(dict[str, Any], raw_function) - function["name"] = f"art_trajectory_probe_{call_index}" - function["arguments"] = '{"art_trajectory_probe":true}' + probe_name = ( + f"art_trajectory_probe_{id(probe_messages):x}_{call_index}" + ) + probe_arguments = json.dumps({probe_name: True}) + while ( + probe_name in existing_values + or probe_arguments in existing_values + ): + probe_name += "_" + probe_arguments = json.dumps({probe_name: True}) + function["name"] = probe_name + function["arguments"] = probe_arguments try: probe = render( probe_messages, @@ -3791,12 +3807,10 @@ def differing_span(probe: Sequence[int]) -> tuple[int, int] | None: message.get("role") == "assistant" and source is not None and _source_is_sampled(source) - for message, source in zip( - history.messages, history.message_sources, strict=True - ) + for message, source in zip(messages, history.message_sources, strict=True) ) for message_index, (message, source) in enumerate( - zip(history.messages, history.message_sources, strict=True) + zip(messages, history.message_sources, strict=True) ): parts = _chat_message_parts(message) sampled = ( diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index c43d4e637..45665117e 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -4245,7 +4245,16 @@ def apply_chat_template( assert not tokenized.flags[0] & art.TokenFlag.SAMPLED -def test_tool_call_probe_handles_contextual_tokenization() -> None: +@pytest.mark.parametrize( + ("name", "arguments"), + [ + ("lookup", "{}"), + ("art_trajectory_probe_0", '{"art_trajectory_probe":true}'), + ], +) +def test_tool_call_probe_handles_contextual_tokenization( + name: str, arguments: str +) -> None: exchange = _chat_exchange([], []) data = exchange.response.model_dump(mode="python") choice = data["choices"][0] @@ -4259,7 +4268,7 @@ def test_tool_call_probe_handles_contextual_tokenization() -> None: { "id": "call-1", "type": "function", - "function": {"name": "lookup", "arguments": "{}"}, + "function": {"name": name, "arguments": arguments}, } ], } @@ -4271,7 +4280,11 @@ def test_tool_call_probe_handles_contextual_tokenization() -> None: class Tokenizer: def __call__(self, text: str, **kwargs: object) -> list[int]: del kwargs - return {"turn 0": [1], "lookup": [90], "{}": [91]}[text] + return { + "turn 0": [1], + name: [90], + arguments: [91], + }[text] def apply_chat_template( self, messages: list[dict[str, Any]], **kwargs: object @@ -4280,7 +4293,7 @@ def apply_chat_template( if messages[-1]["role"] != "assistant": return [1, 10] function = messages[-1]["tool_calls"][0]["function"] - if str(function["name"]).startswith("art_trajectory_probe_"): + if function["name"] != name: return [1, 10, 98, 25, 99, 26] return [1, 10, 20, 25, 30, 26] @@ -4466,8 +4479,19 @@ def test_renderer_ignored_refusal_is_appended_for_tokenization() -> None: ).chat_completions_history() class Tokenizer: - def __call__(self, text: str, **kwargs: object) -> list[int]: - del kwargs + def __call__(self, text: str, **kwargs: object) -> object: + if kwargs.get("return_offsets_mapping"): + answer_start = text.index("answer") + declined_start = text.index("declined") + declined_end = declined_start + len("declined") + return { + "input_ids": [1, 4, 5], + "offset_mapping": [ + (0, answer_start), + (answer_start, declined_start), + (declined_start, declined_end), + ], + } return { "turn 0": [1], "answer": [4], @@ -4477,8 +4501,14 @@ def __call__(self, text: str, **kwargs: object) -> list[int]: def apply_chat_template( self, messages: list[dict[str, Any]], **kwargs: object - ) -> list[int]: + ) -> object: + tokenize = kwargs.pop("tokenize") del kwargs + if not tokenize: + return "".join( + f"{message.get('content') or ''}" + for message in messages + ) return [ token for message in messages From ac6ebb34c42912f363a684b793771f688ed5c035 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Sat, 25 Jul 2026 03:08:01 +0000 Subject: [PATCH 55/58] Preserve zero-width sampled outputs --- src/art/trajectories/_tokenize.py | 38 ++++++++++- tests/unit/trajectories/test_tokenize.py | 83 ++++++++++++++++++++++++ 2 files changed, 119 insertions(+), 2 deletions(-) diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index 599648c7d..57801385d 100644 --- a/src/art/trajectories/_tokenize.py +++ b/src/art/trajectories/_tokenize.py @@ -3542,6 +3542,8 @@ def part_ids(text: str) -> list[int]: last_text = str(last[last_key]) leading = first_text[: len(first_text) - len(first_text.lstrip())] trailing = last_text[len(last_text.rstrip()) :] + if first is last and first_key == last_key and not first_text.strip(): + trailing = "" if first is last and first_key == last_key: core = first_text[ len(leading) : len(first_text) - len(trailing) @@ -3660,7 +3662,9 @@ def part_ids(text: str) -> list[int]: and offsets[token_end][0] < char_end ): token_end += 1 - if ( + if char_start == char_end: + token_bounds[key] = (token_cursor, token_cursor) + elif ( token_end > token_cursor and offsets[token_cursor][0] >= char_start and offsets[token_end - 1][1] <= char_end @@ -3711,6 +3715,29 @@ def differing_span(probe: Sequence[int]) -> tuple[int, int] | None: ): continue parts = _chat_message_parts(message) + if not parts: + exact_output, _ = source_output_tokens(source) + if exact_output: + try: + prefix = render( + messages[:message_index], + add_generation_prompt=True, + ) + completed = render( + messages[: message_index + 1], + add_generation_prompt=False, + ) + except Exception: + continue + if ( + completed[: len(prefix)] == prefix + and rendered[: len(completed)] == completed + ): + probed_bounds[message_index] = ( + len(prefix), + len(prefix), + ) + continue if parts and all(part == "tool_call" for part, _ in parts): probe_messages = deepcopy(messages) tool_calls = probe_messages[message_index].get("tool_calls") @@ -3963,10 +3990,17 @@ def differing_span(probe: Sequence[int]) -> tuple[int, int] | None: ): if not parts and sampled_bounds is not None: start = end = sampled_bounds[0] - elif sampled_bounds is None or sampled_bounds[0] == sampled_bounds[1]: + elif sampled_bounds is None: raise ValueError( "Could not locate a complete sampled message in the rendered history" ) + elif sampled_bounds[0] == sampled_bounds[1]: + if not content_bounds_proven: + raise ValueError( + "Could not locate a complete sampled message in the rendered " + "history" + ) + start = end = sampled_bounds[0] else: start, end = sampled_bounds if parts and len(parts) == 1 and parts[0][0] == "content": diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index 45665117e..cc34171e7 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -4444,6 +4444,38 @@ def apply_chat_template( assert tokenized.flags == [art.TokenFlag(0), art.TokenFlag(0)] +def test_empty_sampled_message_inserts_exact_control_token() -> None: + exchange = _chat_exchange([], [2]) + exchange.response.choices[0].message.content = "" + history = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[exchange]) + ).chat_completions_history() + + class Tokenizer: + def __call__(self, text: str, **kwargs: object) -> list[int]: + del text, kwargs + return [] + + def apply_chat_template( + self, + messages: list[dict[str, Any]], + *, + add_generation_prompt: bool, + **kwargs: object, + ) -> list[int]: + del kwargs + if messages[-1]["role"] == "assistant": + return [1, 99, 9] + assert add_generation_prompt + return [1, 99] + + tokenized = history.tokenize(tokenizer=Tokenizer()) + + assert tokenized.token_ids == [1, 99, 2, 9] + assert tokenized.logprobs[2] == -0.2 + assert tokenized.flags[2] == art.TokenFlag.EXACT | art.TokenFlag.SAMPLED + + def test_renderer_ignored_refusal_is_appended_for_tokenization() -> None: exchange = _chat_exchange([], []) data = exchange.response.model_dump(mode="python") @@ -4634,6 +4666,57 @@ def apply_chat_template( assert tokenized.flags[1:] == [art.TokenFlag.SAMPLED] * 2 +def test_trimmed_whitespace_output_inserts_authoritative_logprob_token() -> None: + exchange = _chat_exchange([], []) + data = exchange.response.model_dump(mode="python") + choice = data["choices"][0] + choice.pop("prompt_token_ids") + choice.pop("token_ids") + choice["message"]["content"] = " " + choice["logprobs"] = { + "content": [ + { + "token": " ", + "logprob": -0.5, + "bytes": [32], + "top_logprobs": [], + } + ], + "refusal": None, + } + exchange.response = ChatCompletion.model_validate(data) + history = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[exchange]) + ).chat_completions_history() + + class Tokenizer: + def __call__(self, text: str, **kwargs: object) -> object: + if kwargs.get("return_offsets_mapping"): + boundary = text.index("") + return { + "input_ids": [1, 9], + "offset_mapping": [(0, boundary), (boundary, len(text))], + } + return {"turn 0": [1], " ": [220]}[text] + + def apply_chat_template( + self, messages: list[dict[str, Any]], **kwargs: object + ) -> object: + tokenize = kwargs.pop("tokenize") + del kwargs + rendered = ( + f"{messages[0]['content']}" + f"{str(messages[-1]['content']).strip()}" + ) + return [1, 9] if tokenize else rendered + + tokenized = history.tokenize(tokenizer=Tokenizer()) + + assert tokenized.token_ids == [1, 220, 9] + assert tokenized.logprobs[1] == -0.5 + assert tokenized.flags[1] == art.TokenFlag.SAMPLED + + def _repeated_text_rerender_history(turn_count: int) -> art.ChatCompletionsHistory: exchanges: list[ChatCompletionsExchange] = [] prompt: list[int] = [] From 0ea64ce5ed4ae3ed7e6ce76a32b8f7176888d45f Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Sat, 25 Jul 2026 03:14:25 +0000 Subject: [PATCH 56/58] Handle trimmed newline samples --- src/art/trajectories/_tokenize.py | 9 +++++++-- tests/unit/trajectories/test_tokenize.py | 15 +++++++++------ 2 files changed, 16 insertions(+), 8 deletions(-) diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index 57801385d..5cd887873 100644 --- a/src/art/trajectories/_tokenize.py +++ b/src/art/trajectories/_tokenize.py @@ -3542,7 +3542,10 @@ def part_ids(text: str) -> list[int]: last_text = str(last[last_key]) leading = first_text[: len(first_text) - len(first_text.lstrip())] trailing = last_text[len(last_text.rstrip()) :] - if first is last and first_key == last_key and not first_text.strip(): + whitespace_only = ( + first is last and first_key == last_key and not first_text.strip() + ) + if whitespace_only: trailing = "" if first is last and first_key == last_key: core = first_text[ @@ -3556,7 +3559,9 @@ def part_ids(text: str) -> list[int]: last[last_key] = ( last_text[: len(last_text) - len(trailing)] + end + trailing ) - part_whitespace[(message_index, part_index)] = (leading, trailing) + part_whitespace[(message_index, part_index)] = ( + ("", "") if whitespace_only else (leading, trailing) + ) markers[start] = (message_index, part_index, "start") markers[end] = (message_index, part_index, "end") if markers: diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index cc34171e7..ee4939eab 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -4666,19 +4666,22 @@ def apply_chat_template( assert tokenized.flags[1:] == [art.TokenFlag.SAMPLED] * 2 -def test_trimmed_whitespace_output_inserts_authoritative_logprob_token() -> None: +@pytest.mark.parametrize(("content", "token_id"), [(" ", 220), ("\n", 198)]) +def test_trimmed_whitespace_output_inserts_authoritative_logprob_token( + content: str, token_id: int +) -> None: exchange = _chat_exchange([], []) data = exchange.response.model_dump(mode="python") choice = data["choices"][0] choice.pop("prompt_token_ids") choice.pop("token_ids") - choice["message"]["content"] = " " + choice["message"]["content"] = content choice["logprobs"] = { "content": [ { - "token": " ", + "token": content, "logprob": -0.5, - "bytes": [32], + "bytes": list(content.encode()), "top_logprobs": [], } ], @@ -4697,7 +4700,7 @@ def __call__(self, text: str, **kwargs: object) -> object: "input_ids": [1, 9], "offset_mapping": [(0, boundary), (boundary, len(text))], } - return {"turn 0": [1], " ": [220]}[text] + return {"turn 0": [1], content: [token_id]}[text] def apply_chat_template( self, messages: list[dict[str, Any]], **kwargs: object @@ -4712,7 +4715,7 @@ def apply_chat_template( tokenized = history.tokenize(tokenizer=Tokenizer()) - assert tokenized.token_ids == [1, 220, 9] + assert tokenized.token_ids == [1, token_id, 9] assert tokenized.logprobs[1] == -0.5 assert tokenized.flags[1] == art.TokenFlag.SAMPLED From 4329366a24af7ae26e82fbcd460f0d9b15a0fb64 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Sat, 25 Jul 2026 03:16:20 +0000 Subject: [PATCH 57/58] Finish sampled evidence hardening --- src/art/trajectories/_tokenize.py | 19 +++++++++++-------- tests/unit/trajectories/test_tokenize.py | 13 ++++++++----- 2 files changed, 19 insertions(+), 13 deletions(-) diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index 5cd887873..109d5deca 100644 --- a/src/art/trajectories/_tokenize.py +++ b/src/art/trajectories/_tokenize.py @@ -3747,14 +3747,17 @@ def differing_span(probe: Sequence[int]) -> tuple[int, int] | None: probe_messages = deepcopy(messages) tool_calls = probe_messages[message_index].get("tool_calls") if isinstance(tool_calls, list): - existing_values = { - value - for call in tool_calls - if isinstance(call, dict) - and isinstance(call.get("function"), dict) - for value in cast(dict[str, Any], call["function"]).values() - if isinstance(value, str) - } + existing_values: set[str] = set() + for existing_call in tool_calls: + if not isinstance(existing_call, dict): + continue + existing_function = existing_call.get("function") + if isinstance(existing_function, dict): + existing_values.update( + value + for value in existing_function.values() + if isinstance(value, str) + ) for call_index, call in enumerate(tool_calls): if not isinstance(call, dict): continue diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index ee4939eab..4213daa30 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -4541,11 +4541,14 @@ def apply_chat_template( f"{message.get('content') or ''}" for message in messages ) - return [ - token - for message in messages - for token in self(str(message.get("content") or "")) - ] + result: list[int] = [] + for message in messages: + encoded = self(str(message.get("content") or "")) + assert isinstance(encoded, list) + for token in encoded: + assert isinstance(token, int) + result.append(token) + return result tokenized = history.tokenize(tokenizer=Tokenizer()) From cb7cd0db35b3b0e72e91e82210412eab59d3fad4 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Sat, 25 Jul 2026 03:22:12 +0000 Subject: [PATCH 58/58] Avoid duplicating preserved whitespace samples --- src/art/trajectories/_tokenize.py | 150 ++++++++++++----- tests/unit/trajectories/test_tokenize.py | 199 +++++++++++++++++++++-- 2 files changed, 294 insertions(+), 55 deletions(-) diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index 109d5deca..84c420764 100644 --- a/src/art/trajectories/_tokenize.py +++ b/src/art/trajectories/_tokenize.py @@ -1,6 +1,7 @@ from __future__ import annotations from bisect import bisect_left +import codecs from collections.abc import Hashable, Mapping, Sequence from copy import deepcopy from dataclasses import dataclass @@ -1416,15 +1417,24 @@ def _visible_logprobs( choice = ( _chat_choice(source) if source is not None else exchange.response.choices[0] ) - for entry in _chat_logprob_entries(choice): + entries = _chat_logprob_entries(choice) + decoder = codecs.getincrementaldecoder("utf-8")() + for index, entry in enumerate(entries): data = _dump(entry) raw_bytes = data.get("bytes") if isinstance(raw_bytes, list): try: - text = bytes(raw_bytes).decode("utf-8") + next_data = ( + _dump(entries[index + 1]) if index + 1 < len(entries) else {} + ) + text = decoder.decode( + bytes(raw_bytes), + final=not isinstance(next_data.get("bytes"), list), + ) except (TypeError, ValueError, UnicodeDecodeError): return [] else: + decoder = codecs.getincrementaldecoder("utf-8")() text = data.get("token") logprob = data.get("logprob") if isinstance(text, str) and isinstance(logprob, (int, float)): @@ -1569,29 +1579,53 @@ def _visible_token_evidence( if not values or tokenizer is None: return None if sampled_text is not None: - matches: list[list[tuple[str, float]]] = [] - for start in range(len(values)): - combined = "" - for end in range(start, len(values)): - combined += values[end][0] - if combined == sampled_text: - matches.append(values[start : end + 1]) - break - if len(combined) >= len(sampled_text): - break - if len(matches) != 1: - return None - values = matches[0] + if "".join(text for text, _ in values) != sampled_text: + matches: list[list[tuple[str, float]]] = [] + for start in range(len(values)): + combined = "" + for end in range(start, len(values)): + combined += values[end][0] + if combined == sampled_text: + matches.append(values[start : end + 1]) + break + if len(combined) >= len(sampled_text): + break + if len(matches) != 1: + return None + values = matches[0] + logprobs = [logprob for _, logprob in values] + text = "".join(text for text, _ in values) + if any(not value for value, _ in values): + token_ids = _ids(tokenizer(text, add_special_tokens=False)) + return (token_ids, logprobs) if len(token_ids) == len(values) else None + try: + contextual = cast(_OffsetTokenizer, tokenizer)( + text, + add_special_tokens=False, + return_offsets_mapping=True, + ) + except (TypeError, ValueError, NotImplementedError): + contextual = None + contextual_data = _string_dict(contextual) + offsets = ( + contextual_data.get("offset_mapping") if contextual_data is not None else None + ) + if isinstance(offsets, list) and len(offsets) == len(values): + boundaries: list[tuple[int, int]] = [] + cursor = 0 + for value, _ in values: + boundaries.append((cursor, cursor + len(value))) + cursor += len(value) + if offsets == boundaries: + token_ids = _ids(contextual) + if len(token_ids) == len(values): + return token_ids, logprobs token_ids: list[int] = [] - logprobs: list[float] = [] - for text, logprob in values: - if not text: - return None + for text, _ in values: encoded = _ids(tokenizer(text, add_special_tokens=False)) if len(encoded) != 1: return None token_ids.append(encoded[0]) - logprobs.append(logprob) return token_ids, logprobs @@ -3512,6 +3546,24 @@ def part_ids(text: str) -> list[int]: ): direct_bounds = [] + def differing_span(probe: Sequence[int]) -> tuple[int, int] | None: + prefix = 0 + while ( + prefix < len(rendered) + and prefix < len(probe) + and rendered[prefix] == probe[prefix] + ): + prefix += 1 + suffix = 0 + while ( + suffix < len(rendered) - prefix + and suffix < len(probe) - prefix + and rendered[-suffix - 1] == probe[-suffix - 1] + ): + suffix += 1 + end = len(rendered) - suffix + return (prefix, end) if prefix < end else None + marked_bounds: dict[int, tuple[int, int]] = {} marked_part_bounds: dict[int, list[tuple[int, int]]] = {} if not direct_bounds: @@ -3688,24 +3740,46 @@ def part_ids(text: str) -> list[int]: bounds[0][0], bounds[-1][1], ) - - def differing_span(probe: Sequence[int]) -> tuple[int, int] | None: - prefix = 0 - while ( - prefix < len(rendered) - and prefix < len(probe) - and rendered[prefix] == probe[prefix] - ): - prefix += 1 - suffix = 0 - while ( - suffix < len(rendered) - prefix - and suffix < len(probe) - prefix - and rendered[-suffix - 1] == probe[-suffix - 1] - ): - suffix += 1 - end = len(rendered) - suffix - return (prefix, end) if prefix < end else None + for message_index, bounds in list(marked_part_bounds.items()): + whitespace_parts = [ + (part_index, text) + for part_index, (_, text) in enumerate( + _chat_message_parts(messages[message_index]) + ) + if text and not text.strip() + ] + for part_index, _ in whitespace_parts: + empty_messages = deepcopy(messages) + groups = _chat_message_text_slot_groups(empty_messages[message_index]) + if part_index >= len(groups): + marked_bounds.pop(message_index, None) + marked_part_bounds.pop(message_index, None) + break + for container, key in groups[part_index]: + container[key] = "" + try: + empty_render = render( + empty_messages, + add_generation_prompt=not ends_with_assistant, + ) + except Exception: + marked_bounds.pop(message_index, None) + marked_part_bounds.pop(message_index, None) + break + if empty_render == rendered: + continue + anchor = bounds[part_index][0] + start = anchor - (len(rendered) - len(empty_render)) + span = (start, anchor) + if ( + start < 0 + or rendered[: span[0]] + rendered[span[1] :] != empty_render + ): + marked_bounds.pop(message_index, None) + marked_part_bounds.pop(message_index, None) + break + bounds[part_index] = span + marked_bounds[message_index] = (bounds[0][0], bounds[-1][1]) probed_bounds: dict[int, tuple[int, int]] = {} if not direct_bounds: diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index 4213daa30..914eadaa5 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -4607,13 +4607,19 @@ def test_trimmed_render_preserves_authoritative_textual_logprob_tokens() -> None choice = data["choices"][0] choice.pop("prompt_token_ids") choice.pop("token_ids") - choice["message"]["content"] = " first " + choice["message"]["content"] = " helloworld " choice["logprobs"] = { "content": [ { - "token": " first", + "token": " hello", "logprob": -0.4, - "bytes": list(b" first"), + "bytes": list(b" hello"), + "top_logprobs": [], + }, + { + "token": "world", + "logprob": -0.45, + "bytes": list(b"world"), "top_logprobs": [], }, { @@ -4633,8 +4639,13 @@ def test_trimmed_render_preserves_authoritative_textual_logprob_tokens() -> None class Tokenizer: def __call__(self, text: str, **kwargs: object) -> object: if kwargs.get("return_offsets_mapping"): - content_start = text.index("first") - content_end = content_start + len("first") + if text == " helloworld ": + return { + "input_ids": [1118, 2222, 220], + "offset_mapping": [(0, 6), (6, 11), (11, 12)], + } + content_start = text.index("helloworld") + content_end = content_start + len("helloworld") return { "input_ids": [1, 3765], "offset_mapping": [ @@ -4644,8 +4655,9 @@ def __call__(self, text: str, **kwargs: object) -> object: } return { "turn 0": [1], - " first ": [1118, 220], - " first": [1118], + " helloworld ": [1118, 2222, 220], + " hello": [1118], + "world": [3333], " ": [220], }[text] @@ -4664,14 +4676,129 @@ def apply_chat_template( tokenized = history.tokenize(tokenizer=Tokenizer()) - assert tokenized.token_ids == [1, 1118, 220] - assert tokenized.logprobs[1:] == [-0.4, -0.5] - assert tokenized.flags[1:] == [art.TokenFlag.SAMPLED] * 2 + assert tokenized.token_ids == [1, 1118, 2222, 220] + assert tokenized.logprobs[1:] == [-0.4, -0.45, -0.5] + assert tokenized.flags[1:] == [art.TokenFlag.SAMPLED] * 3 -@pytest.mark.parametrize(("content", "token_id"), [(" ", 220), ("\n", 198)]) +def test_textual_logprobs_reconstruct_split_utf8_bytes() -> None: + from art.trajectories._tokenize import _visible_token_evidence + + exchange = _chat_exchange([], []) + data = exchange.response.model_dump(mode="python") + choice = data["choices"][0] + choice.pop("prompt_token_ids") + choice.pop("token_ids") + choice["message"]["content"] = "😊" + choice["logprobs"]["content"] = [ + { + "token": "�", + "logprob": -0.4, + "bytes": [240, 159, 152], + "top_logprobs": [], + }, + { + "token": "�", + "logprob": -0.5, + "bytes": [138], + "top_logprobs": [], + }, + ] + exchange.response = ChatCompletion.model_validate(data) + + class Tokenizer: + def __call__(self, text: str, **kwargs: object) -> list[int]: + del kwargs + assert text == "😊" + return [11, 12] + + def apply_chat_template( + self, + messages: list[dict[str, Any]], + *, + tools: object, + tokenize: bool, + add_generation_prompt: bool, + chat_template: str | None = None, + **kwargs: object, + ) -> object: + del messages, tools, tokenize, add_generation_prompt, chat_template, kwargs + raise AssertionError("template rendering is not expected") + + assert _visible_token_evidence(Tokenizer(), exchange, sampled_text="😊") == ( + [11, 12], + [-0.4, -0.5], + ) + + +def test_textual_logprobs_reject_shifted_contextual_token_boundaries() -> None: + from art.trajectories._tokenize import _visible_token_evidence + + exchange = _chat_exchange([], []) + data = exchange.response.model_dump(mode="python") + choice = data["choices"][0] + choice.pop("prompt_token_ids") + choice.pop("token_ids") + choice["message"]["content"] = " penalates" + choice["logprobs"]["content"] = [ + { + "token": " pena", + "logprob": -0.4, + "bytes": list(b" pena"), + "top_logprobs": [], + }, + { + "token": "lates", + "logprob": -0.5, + "bytes": list(b"lates"), + "top_logprobs": [], + }, + ] + exchange.response = ChatCompletion.model_validate(data) + + class Tokenizer: + def __call__(self, text: str, **kwargs: object) -> object: + if kwargs.get("return_offsets_mapping"): + return { + "input_ids": [30, 31], + "offset_mapping": [(0, 6), (6, 10)], + } + return {" pena": [11], "lates": [12]}[text] + + def apply_chat_template( + self, + messages: list[dict[str, Any]], + *, + tools: object, + tokenize: bool, + add_generation_prompt: bool, + chat_template: str | None = None, + **kwargs: object, + ) -> object: + del messages, tools, tokenize, add_generation_prompt, chat_template, kwargs + raise AssertionError("template rendering is not expected") + + assert _visible_token_evidence( + Tokenizer(), exchange, sampled_text=" penalates" + ) == ([11, 12], [-0.4, -0.5]) + + +@pytest.mark.parametrize( + ("content", "token_id", "trim", "adjacent_scaffold", "reject_empty"), + [ + (" ", 220, True, False, False), + ("\n", 198, True, False, False), + (" ", 220, False, False, False), + (" ", 220, False, True, False), + (" ", 220, False, False, True), + ], +) def test_trimmed_whitespace_output_inserts_authoritative_logprob_token( - content: str, token_id: int + content: str, + token_id: int, + trim: bool, + adjacent_scaffold: bool, + reject_empty: bool, ) -> None: exchange = _chat_exchange([], []) data = exchange.response.model_dump(mode="python") @@ -4698,10 +4825,27 @@ def test_trimmed_whitespace_output_inserts_authoritative_logprob_token( class Tokenizer: def __call__(self, text: str, **kwargs: object) -> object: if kwargs.get("return_offsets_mapping"): - boundary = text.index("") + start = text.index("") + len("") + end = text.index("") + if start != end: + scaffold_start = end - int(adjacent_scaffold) + return { + "input_ids": [ + 1, + token_id, + *([token_id] if adjacent_scaffold else []), + 9, + ], + "offset_mapping": [ + (0, start), + (start, scaffold_start), + *([(scaffold_start, end)] if adjacent_scaffold else []), + (end, len(text)), + ], + } return { "input_ids": [1, 9], - "offset_mapping": [(0, boundary), (boundary, len(text))], + "offset_mapping": [(0, start), (start, len(text))], } return {"turn 0": [1], content: [token_id]}[text] @@ -4710,17 +4854,38 @@ def apply_chat_template( ) -> object: tokenize = kwargs.pop("tokenize") del kwargs + content_value = str(messages[-1]["content"]) + if trim: + content_value = content_value.strip() + scaffold = " " if adjacent_scaffold else "" rendered = ( f"{messages[0]['content']}" - f"{str(messages[-1]['content']).strip()}" + f"{content_value}{scaffold}" ) - return [1, 9] if tokenize else rendered + if not tokenize: + return rendered + if reject_empty and not content_value: + raise RuntimeError("template rejects empty assistant content") + rendered_id = 777 if "ART_TRAJECTORY" in content_value else token_id + return [ + 1, + *([rendered_id] if content_value else []), + *([token_id] if adjacent_scaffold else []), + 9, + ] tokenized = history.tokenize(tokenizer=Tokenizer()) - assert tokenized.token_ids == [1, token_id, 9] + assert tokenized.token_ids == [ + 1, + token_id, + *([token_id] if adjacent_scaffold else []), + 9, + ] assert tokenized.logprobs[1] == -0.5 assert tokenized.flags[1] == art.TokenFlag.SAMPLED + if adjacent_scaffold: + assert tokenized.flags[2] == art.TokenFlag(0) def _repeated_text_rerender_history(turn_count: int) -> art.ChatCompletionsHistory: