From 03fd13dc292dd5e4f399e8713b2c3f633164f62f Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Tue, 14 Jul 2026 22:38:27 +0000 Subject: [PATCH 01/17] feat: add protocol-aware trajectory capture --- .gitignore | 2 +- pyproject.toml | 3 + src/art/__init__.py | 38 +- src/art/auto_trajectory.py | 208 +----- src/art/preprocessing/tokenize.py | 62 +- src/art/trajectories.py | 319 --------- src/art/trajectories/__init__.py | 413 ++++++++++++ src/art/trajectories/_capture/__init__.py | 23 + src/art/trajectories/_capture/aiohttp.py | 81 +++ src/art/trajectories/_capture/core.py | 91 +++ src/art/trajectories/_capture/httpx.py | 106 +++ src/art/trajectories/_capture/requests.py | 53 ++ src/art/trajectories/_compat.py | 198 ++++++ src/art/trajectories/_protocols.py | 257 ++++++++ src/art/trajectories/_scope.py | 100 +++ src/art/trajectories/_tokenize.py | 749 ++++++++++++++++++++++ tests/unit/test_auto_trajectory.py | 55 +- tests/unit/trajectories/test_capture.py | 419 ++++++++++++ tests/unit/trajectories/test_tokenize.py | 445 +++++++++++++ uv.lock | 25 + 20 files changed, 3064 insertions(+), 583 deletions(-) delete mode 100644 src/art/trajectories.py create mode 100644 src/art/trajectories/__init__.py create mode 100644 src/art/trajectories/_capture/__init__.py create mode 100644 src/art/trajectories/_capture/aiohttp.py create mode 100644 src/art/trajectories/_capture/core.py create mode 100644 src/art/trajectories/_capture/httpx.py create mode 100644 src/art/trajectories/_capture/requests.py create mode 100644 src/art/trajectories/_compat.py create mode 100644 src/art/trajectories/_protocols.py create mode 100644 src/art/trajectories/_scope.py create mode 100644 src/art/trajectories/_tokenize.py create mode 100644 tests/unit/trajectories/test_capture.py create mode 100644 tests/unit/trajectories/test_tokenize.py diff --git a/.gitignore b/.gitignore index 0dfae3afe..ed198af87 100644 --- a/.gitignore +++ b/.gitignore @@ -14,7 +14,7 @@ wandb/ docs/node_modules/ dist/ replays/ -trajectories/ +/trajectories/ .DS_Store .local/ .claude/settings.local.json diff --git a/pyproject.toml b/pyproject.toml index 3fa87c527..371aeacb4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,7 +5,10 @@ description = "The OpenPipe Agent Reinforcement Training (ART) library" readme = "README.md" requires-python = ">=3.12" dependencies = [ + "aiohttp>=3.10.0", + "anthropic>=0.75.0", "openai>=2.14.0", + "requests>=2.32.0", "typer>=0.15.2", "litellm>=1.71.1,<=1.82.0", "weave>=0.52.24", diff --git a/src/art/__init__.py b/src/art/__init__.py index ddb0d8c1f..c715dafb2 100644 --- a/src/art/__init__.py +++ b/src/art/__init__.py @@ -59,7 +59,6 @@ from . import dev -from .auto_trajectory import auto_trajectory, capture_auto_trajectory from .backend import Backend from .batches import trajectory_group_batches from .dev import LoRAConfig @@ -70,7 +69,27 @@ ) from .model import Model, TrainableModel from .serverless import ServerlessBackend -from .trajectories import Trajectory, TrajectoryGroup +from .trajectories import ( + ChatCompletionsExchange, + CompletionsExchange, + MessagesExchange, + ResponsesExchange, + TokenizedTrajectory, + TokenizedTrajectoryGroup, + Trajectory, + TrajectoryExchanges, + TrajectoryGroup, + auto_trajectory, + capture_auto_trajectory, + current_trajectory, + current_trajectory_group, + tokenize_trajectories, + tokenize_trajectory, + tokenize_trajectory_group, + tokenize_trajectory_groups, + trajectory, + trajectory_group, +) from .types import ( LocalTrainResult, MegatronRuntimeConfig, @@ -90,6 +109,8 @@ "dev", "auto_trajectory", "capture_auto_trajectory", + "current_trajectory", + "current_trajectory_group", "gather_trajectories", "gather_trajectory_groups", "trajectory_group_batches", @@ -112,7 +133,20 @@ "TrainConfig", "TrainResult", "Trajectory", + "TrajectoryExchanges", "TrajectoryGroup", + "ChatCompletionsExchange", + "CompletionsExchange", + "ResponsesExchange", + "MessagesExchange", + "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/auto_trajectory.py b/src/art/auto_trajectory.py index 434bcc403..f4f2f64d3 100644 --- a/src/art/auto_trajectory.py +++ b/src/art/auto_trajectory.py @@ -1,207 +1,5 @@ -import contextvars -import json -import logging -from typing import Any, AsyncIterator, Coroutine, Iterator, Literal, overload +"""Deprecated compatibility names for automatic trajectory capture.""" -import httpx._models -from openai.types.chat.chat_completion import ChatCompletion, Choice -from openai.types.chat.chat_completion_chunk import ChatCompletionChunk +from .trajectories import auto_trajectory, capture_auto_trajectory -from .openai import init_chat_completion, update_chat_completion -from .preprocessing.moe_routing import attach_moe_routing_metadata_to_choice -from .preprocessing.vllm_tokens import attach_vllm_token_metadata_to_choice -from .trajectories import History, Trajectory - -logger = logging.getLogger(__name__) - - -def parse_sse_to_chat_completion(content: bytes) -> ChatCompletion: - """Parse SSE (Server-Sent Events) content and build a ChatCompletion. - - This handles the case where streaming responses have already been consumed - and we need to reconstruct the ChatCompletion from buffered bytes. - """ - chat_completion: ChatCompletion | None = None - - # Parse SSE format: each line starting with "data: " contains JSON - for line in content.decode("utf-8").split("\n"): - line = line.strip() - if not line.startswith("data: "): - continue - data = line[6:] # Remove "data: " prefix - if data == "[DONE]": - continue - - chunk_data = json.loads(data) - chunk = ChatCompletionChunk(**chunk_data) - if chat_completion is None: - chat_completion = init_chat_completion(chunk) - update_chat_completion(chat_completion, chunk) - - if chat_completion is None: - raise ValueError("No valid chat completion chunks found in SSE content") - - return chat_completion - - -@overload -def auto_trajectory(*, required: Literal[True]) -> Trajectory: ... - - -@overload -def auto_trajectory(*, required: Literal[False] = False) -> Trajectory | None: ... - - -def auto_trajectory(*, required: bool = False) -> Trajectory | None: - context = auto_trajectory_context_var.get(None) - if context is None: - if required: - raise RuntimeError( - "No auto trajectory in context. `auto_trajectory(required=True)` must be called in a `capture_auto_trajectory(...)` scope." - ) - return None - return context.trajectory - - -async def capture_auto_trajectory(coroutine: Coroutine[Any, Any, Any]) -> Trajectory: - with AutoTrajectoryContext() as trajectory: - await coroutine - return trajectory - - -class AutoTrajectoryContext: - def __init__(self) -> None: - self.trajectory = Trajectory() - - def __enter__(self) -> Trajectory: - self.token = auto_trajectory_context_var.set(self) - return self.trajectory - - def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None: - auto_trajectory_context_var.reset(self.token) - self.trajectory.finish() - - def handle_httpx_response(self, response: httpx._models.Response) -> None: - # Get buffered content (set by patched aiter_bytes/iter_bytes) - content = getattr(response, "_content_so_far", b"") - if not content: - # No content captured, nothing to process - return - - try: - request_content = json.loads(getattr(response.request, "_content", b"")) - except (json.JSONDecodeError, AttributeError): - # Not a JSON request body, skip - return - - messages = request_content.get("messages") - if messages is None: - # Not a chat completion request - return - - tools = request_content.get("tools", None) - - try: - if request_content.get("stream", False): - # Parse SSE content directly from buffered bytes - chat_completion = parse_sse_to_chat_completion(content) - choice = chat_completion.choices[0] - response_payload = chat_completion.model_dump(mode="python") - attach_vllm_token_metadata_to_choice( - choice=choice, - response_payload=response_payload, - choice_index=0, - ) - attach_moe_routing_metadata_to_choice( - choice=choice, - response_payload=response_payload, - choice_index=0, - ) - else: - response_payload = json.loads(content) - choice = Choice(**response_payload["choices"][0]) - attach_vllm_token_metadata_to_choice( - choice=choice, - response_payload=response_payload, - choice_index=0, - ) - attach_moe_routing_metadata_to_choice( - choice=choice, - response_payload=response_payload, - choice_index=0, - ) - except (json.JSONDecodeError, KeyError, ValueError) as e: - logger.debug(f"Failed to parse response content: {e}") - return - - # Find the appropriate history to add this response to - history: Trajectory | History = self.trajectory - history_index = -1 - while True: - history_messages = history.messages() - if history_messages == messages[: len(history_messages)] and ( - history.tools == tools - or (history_messages == [] and history.tools is None) - ): - break - history_index += 1 - try: - history = self.trajectory.additional_histories[history_index] - except IndexError: - history = History(messages_and_choices=[]) - self.trajectory.additional_histories.append(history) - break - - history.messages_and_choices.extend( - messages[len(history.messages_and_choices) :] - ) - history.messages_and_choices.append(choice) - history.tools = tools - - -auto_trajectory_context_var: contextvars.ContextVar[AutoTrajectoryContext] = ( - contextvars.ContextVar("auto_trajectory_context") -) - - -def patch_httpx() -> None: - original_iter_bytes = httpx._models.Response.iter_bytes - original_aiter_bytes = httpx._models.Response.aiter_bytes - original_close = httpx._models.Response.close - original_aclose = httpx._models.Response.aclose - - def patched_iter_bytes( - self: httpx._models.Response, chunk_size: int | None = None - ) -> Iterator[bytes]: - for chunk in original_iter_bytes(self, chunk_size): - setattr( - self, "_content_so_far", getattr(self, "_content_so_far", b"") + chunk - ) - yield chunk - - async def patched_aiter_bytes( - self: httpx._models.Response, chunk_size: int | None = None - ) -> AsyncIterator[bytes]: - async for chunk in original_aiter_bytes(self, chunk_size): - setattr( - self, "_content_so_far", getattr(self, "_content_so_far", b"") + chunk - ) - yield chunk - - def patched_close(self: httpx._models.Response) -> None: - original_close(self) - if context := auto_trajectory_context_var.get(None): - context.handle_httpx_response(self) - - async def patched_aclose(self: httpx._models.Response) -> None: - await original_aclose(self) - if context := auto_trajectory_context_var.get(None): - context.handle_httpx_response(self) - - httpx._models.Response.iter_bytes = patched_iter_bytes # ty:ignore[invalid-assignment] - httpx._models.Response.aiter_bytes = patched_aiter_bytes # ty:ignore[invalid-assignment] - httpx._models.Response.close = patched_close # ty:ignore[invalid-assignment] - httpx._models.Response.aclose = patched_aclose # ty:ignore[invalid-assignment] - - -patch_httpx() +__all__ = ["auto_trajectory", "capture_auto_trajectory"] diff --git a/src/art/preprocessing/tokenize.py b/src/art/preprocessing/tokenize.py index 9d8ca0c34..5a0daf609 100644 --- a/src/art/preprocessing/tokenize.py +++ b/src/art/preprocessing/tokenize.py @@ -525,19 +525,55 @@ def tokenize_trajectory_groups( advantage /= reward_std + 1e-6 if advantage == 0 and drop_zero_advantage_trajectories: continue - trajectory_results = tokenize_vllm_trajectory_histories( - tokenizer=tokenizer, - histories=[ - History( - messages_and_choices=trajectory.messages_and_choices, - tools=trajectory.tools, - ), - *trajectory.additional_histories, - ], - advantage=advantage, - allow_training_without_logprobs=allow_training_without_logprobs, - trajectory=trajectory, - ) + if trajectory.exchanges: + from ..trajectories._tokenize import tokenize_one + + exchange_result = tokenize_one( + trajectory, + getattr(tokenizer, "name_or_path", None), + model=None, + chat_template=None, + chat_template_kwargs=chat_template_kwargs, + tokenizer_instance=tokenizer, + ) + choice_offsets = [ + index + for index, trainable in enumerate(exchange_result.assistant_mask) + if trainable + and (index == 0 or not exchange_result.assistant_mask[index - 1]) + ] + 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 exchange_result.assistant_mask + ], + logprobs=exchange_result.logprobs, + pixel_values=None, + image_grid_thw=None, + trajectory=trajectory, + choice_offsets=choice_offsets, + extra_logprobs={}, + _tokenizer=tokenizer, + ) + ] + else: + trajectory_results = tokenize_vllm_trajectory_histories( + tokenizer=tokenizer, + histories=[ + History( + messages_and_choices=trajectory.messages_and_choices, + tools=trajectory.tools, + ), + *trajectory.additional_histories, + ], + advantage=advantage, + allow_training_without_logprobs=allow_training_without_logprobs, + trajectory=trajectory, + ) weight = 1 / ( sum(sum(result.assistant_mask) for result in trajectory_results) + 1e-6 ) diff --git a/src/art/trajectories.py b/src/art/trajectories.py deleted file mode 100644 index a929d0e71..000000000 --- a/src/art/trajectories.py +++ /dev/null @@ -1,319 +0,0 @@ -import asyncio -from contextlib import asynccontextmanager -from datetime import datetime -import time -import traceback -from typing import ( - Any, - AsyncGenerator, - Awaitable, - Iterable, - Iterator, - cast, - overload, -) - -from openai.types.chat.chat_completion import Choice -import pydantic - -from .types import Message, Messages, MessagesAndChoices, Tools - -MetadataValue = float | int | str | bool | None - - -class PydanticException(pydantic.BaseModel): - type: str - message: str - traceback: str - - -class History(pydantic.BaseModel): - messages_and_choices: MessagesAndChoices - tools: Tools | None = None - - def messages(self) -> Messages: - return get_messages(self.messages_and_choices) - - -class Trajectory(pydantic.BaseModel): - messages_and_choices: MessagesAndChoices = [] - tools: Tools | None = None - additional_histories: list[History] = [] - reward: float = 0.0 - initial_policy_version: int | None = None - final_policy_version: int | None = None - metrics: dict[str, float | int | bool] = {} - metadata: dict[str, MetadataValue] = {} - logs: list[str] = [] - start_time: datetime = pydantic.Field(default_factory=datetime.now, exclude=True) - - def log(self, message: str) -> None: - self.logs.append(message) - - def finish(self) -> "Trajectory": - duration = (datetime.now() - self.start_time).total_seconds() - self.metrics["duration"] = duration - return self - - @asynccontextmanager - async def track_duration(self, metric_name: str) -> AsyncGenerator[None, None]: - start_time = time.monotonic() - try: - yield - finally: - duration = time.monotonic() - start_time - metric_key = f"{metric_name}_duration" - self.metrics[metric_key] = self.metrics.get(metric_key, 0.0) + duration - - def __str__(self) -> str: - return f"Trajectory(reward={self.reward}, metrics={self.metrics}, metadata={self.metadata})" - - def messages(self) -> Messages: - return get_messages(self.messages_and_choices) - - # Used for logging to console - def for_logging(self) -> dict[str, Any]: - loggable_dict: dict[str, Any] = { - "reward": self.reward, - "initial_policy_version": self.initial_policy_version, - "final_policy_version": self.final_policy_version, - "metrics": self.metrics, - "metadata": self.metadata, - "messages": [], - "tools": self.tools, - "logs": self.logs, - } - for message_or_choice in self.messages_and_choices: - if isinstance(message_or_choice, Choice): - trainable = True - message: dict[str, Any] = message_or_choice.message.to_dict() - else: - trainable = False - message = cast(dict[str, Any], message_or_choice) - loggable_dict["messages"].append({**message, "trainable": trainable}) - return loggable_dict - - -def get_messages(messages_and_choices: MessagesAndChoices) -> Messages: - messages: Messages = [] - for message_or_choice in messages_and_choices: - if isinstance(message_or_choice, Choice): - content = message_or_choice.message.content or "" - tool_calls = message_or_choice.message.tool_calls or [] - assistant_message: Message = cast( - Message, - { - "role": "assistant", - "content": content, - **( - { - "tool_calls": [ - tool_call.model_dump(mode="json") - for tool_call in tool_calls - ] - } - if tool_calls - else {} - ), - }, - ) - messages.append(assistant_message) - else: - # Ensure content is always a string for tokenizer chat templates - msg = dict(message_or_choice) - if msg.get("content") is None: - msg["content"] = "" - messages.append(cast(Message, msg)) - return messages - - -class TrajectoryGroup(pydantic.BaseModel): - trajectories: list[Trajectory] - exceptions: list[PydanticException] = [] - metadata: dict[str, MetadataValue] = {} - metrics: dict[str, float | int | bool] = {} - logs: list[str] = [] - - def __init__( - self, - trajectories: ( - Iterable[Trajectory | BaseException] | Iterable[Awaitable[Trajectory]] - ), - *, - exceptions: list[BaseException] = [], - metadata: dict[str, MetadataValue] | None = None, - metrics: dict[str, float | int | bool] | None = None, - logs: list[str] | None = None, - ) -> None: - super().__init__( - trajectories=[ - trajectory - for trajectory in trajectories - if isinstance(trajectory, Trajectory) - ] - or getattr(self, "trajectories", []), - exceptions=[ - PydanticException( - type=str(type(exception)), - message=str(exception), - traceback="\n".join( - traceback.format_exception( - type(exception), exception, exception.__traceback__ - ) - ), - ) - for exception in ( - [ - exception - for exception in trajectories - if isinstance(exception, BaseException) - ] - + exceptions - ) - ], - metadata=( - metadata if metadata is not None else getattr(self, "metadata", {}) - ), - metrics=metrics if metrics is not None else getattr(self, "metrics", {}), - logs=logs if logs is not None else getattr(self, "logs", []), - ) - - def __copy__(self): - """Support for copy.copy()""" - - # Create a new instance using the constructor - # Pass shallow copies of the lists to avoid shared mutation - new_instance = self.__class__( - trajectories=self.trajectories[:], # Shallow copy of list - exceptions=[], # Will be set below - metadata=self.metadata.copy(), - metrics=self.metrics.copy(), - logs=self.logs[:], - ) - # Manually copy exceptions since they're PydanticException objects - new_instance.exceptions = self.exceptions[:] - return new_instance - - def __deepcopy__(self, memo: dict[int, Any] | None = None): - """Support for copy.deepcopy()""" - import copy - - # Initialize memo if not provided - if memo is None: - memo = {} - - # Check memo to handle circular references - if id(self) in memo: - return memo[id(self)] - - # Create a new instance with deep copies - new_instance = self.__class__( - trajectories=copy.deepcopy(self.trajectories, memo), - exceptions=[], # Will be set below - metadata=copy.deepcopy(self.metadata, memo), - metrics=copy.deepcopy(self.metrics, memo), - logs=copy.deepcopy(self.logs, memo), - ) - # Register in memo before deep copying attributes to handle circular refs - memo[id(self)] = new_instance - # Deep copy exceptions - new_instance.exceptions = copy.deepcopy(self.exceptions, memo) - return new_instance - - def log(self, message: str) -> None: - self.logs.append(message) - - def __iter__(self) -> Iterator[Trajectory]: # type: ignore - return iter(self.trajectories) - - def __len__(self) -> int: - return len(self.trajectories) - - @overload - def __new__( - cls, - trajectories: Iterable[Trajectory | BaseException], - *, - exceptions: list[BaseException] = [], - metadata: dict[str, MetadataValue] | None = None, - metrics: dict[str, float | int | bool] | None = None, - logs: list[str] | None = None, - ) -> "TrajectoryGroup": ... - - @overload - def __new__( - cls, - trajectories: Iterable[Awaitable[Trajectory]], - *, - exceptions: list[BaseException] = [], - metadata: dict[str, MetadataValue] | None = None, - metrics: dict[str, float | int | bool] | None = None, - logs: list[str] | None = None, - ) -> Awaitable["TrajectoryGroup"]: ... - - def __new__( - cls, - trajectories: ( - Iterable[Trajectory | BaseException] | Iterable[Awaitable[Trajectory]] - ), - *, - exceptions: list[BaseException] = [], - metadata: dict[str, MetadataValue] | None = None, - metrics: dict[str, float | int | bool] | None = None, - logs: list[str] | None = None, - ) -> "TrajectoryGroup | Awaitable[TrajectoryGroup]": - ts = list(trajectories) - if any(hasattr(t, "__await__") for t in ts): - - async def _( - exceptions: list[BaseException], - metadata: dict[str, MetadataValue] | None, - metrics: dict[str, float | int | bool] | None, - logs: list[str] | None, - ): - from .gather import get_gather_context, record_metrics - - context = get_gather_context() - trajectories = [] - for future in asyncio.as_completed( - cast(list[Awaitable[Trajectory]], ts) - ): - try: - trajectory = await future - trajectories.append(trajectory) - record_metrics(context, trajectory) - context.update_pbar(n=1) - except BaseException as e: - exceptions.append(e) - context.metric_sums["exceptions"] += 1 - context.update_pbar(n=0) - if context.too_many_exceptions(): - raise - return TrajectoryGroup( - trajectories=trajectories, - exceptions=exceptions, - metadata=metadata, - metrics=metrics, - logs=logs, - ) - - class CoroutineWithMetadata: - def __init__(self, coro, num_trajectories): - self.coro = coro - self._num_trajectories = num_trajectories - - def __await__(self): - return self.coro.__await__() - - coro = _(exceptions.copy(), metadata, metrics, logs) - return CoroutineWithMetadata(coro, len(ts)) # type: ignore[return-value] - else: - group = super().__new__(cls) - group.__init__( - trajectories=cast(list[Trajectory | BaseException], ts), - exceptions=exceptions, - metadata=metadata, - metrics=metrics, - logs=logs, - ) - return group diff --git a/src/art/trajectories/__init__.py b/src/art/trajectories/__init__.py new file mode 100644 index 000000000..b68c1fdb5 --- /dev/null +++ b/src/art/trajectories/__init__.py @@ -0,0 +1,413 @@ +from __future__ import annotations + +from collections.abc import AsyncGenerator, Awaitable, Coroutine, Iterable, Iterator +from contextlib import asynccontextmanager +from datetime import datetime +import time +from typing import Any, Literal, cast, overload + +from anthropic.types import Message as AnthropicMessage +from openai.types import Completion +from openai.types.chat import ChatCompletion +from openai.types.responses import Response +import pydantic +from typing_extensions import deprecated + +from ..types import Messages, MessagesAndChoices, Tools + +MetadataValue = Any + + +class ChatCompletionsRequest(pydantic.RootModel[dict[str, Any]]): + """The JSON body sent to an OpenAI-compatible Chat Completions endpoint.""" + + +class CompletionsRequest(pydantic.RootModel[dict[str, Any]]): + """The JSON body sent to an OpenAI-compatible Completions endpoint.""" + + +class ResponsesRequest(pydantic.RootModel[dict[str, Any]]): + """The JSON body sent to an OpenAI-compatible Responses endpoint.""" + + +class MessagesRequest(pydantic.RootModel[dict[str, Any]]): + """The JSON body sent to an Anthropic-compatible Messages endpoint.""" + + +class ChatCompletionsExchange(pydantic.BaseModel): + request: ChatCompletionsRequest + response: ChatCompletion + model: str | None + start_time: datetime + end_time: datetime + + +class CompletionsExchange(pydantic.BaseModel): + request: CompletionsRequest + response: Completion + model: str | None + start_time: datetime + end_time: datetime + + +class ResponsesExchange(pydantic.BaseModel): + request: ResponsesRequest + response: Response + model: str | None + start_time: datetime + end_time: datetime + + +class MessagesExchange(pydantic.BaseModel): + request: MessagesRequest + response: AnthropicMessage + model: str | None + start_time: datetime + end_time: datetime + + +class TrajectoryExchanges(pydantic.BaseModel): + chat_completions: list[ChatCompletionsExchange] = pydantic.Field( + default_factory=list + ) + completions: list[CompletionsExchange] = pydantic.Field(default_factory=list) + responses: list[ResponsesExchange] = pydantic.Field(default_factory=list) + messages: list[MessagesExchange] = pydantic.Field(default_factory=list) + + def __bool__(self) -> bool: + return any( + (self.chat_completions, self.completions, self.responses, self.messages) + ) + + +class PydanticException(pydantic.BaseModel): + type: str + message: str + traceback: str + + +class History(pydantic.BaseModel): + messages_and_choices: MessagesAndChoices + tools: Tools | None = None + + def messages(self) -> Messages: + return get_messages(self.messages_and_choices) + + +class Trajectory(pydantic.BaseModel): + exchanges: TrajectoryExchanges = pydantic.Field(default_factory=TrajectoryExchanges) + messages_and_choices: MessagesAndChoices = pydantic.Field( + default_factory=list, + ) + tools: Tools | None = None + additional_histories: list[History] = pydantic.Field( + default_factory=list, + ) + reward: float = 0.0 + initial_policy_version: int | None = None + final_policy_version: int | None = None + metrics: dict[str, float | int | bool] = pydantic.Field(default_factory=dict) + metadata: dict[str, Any] = pydantic.Field(default_factory=dict) + logs: list[str] = pydantic.Field(default_factory=list) + start_time: datetime = pydantic.Field(default_factory=datetime.now, exclude=True) + + @pydantic.model_validator(mode="after") + def validate_representation(self) -> Trajectory: + if self.exchanges and ( + self.messages_and_choices + or self.tools is not None + or self.additional_histories + ): + raise ValueError( + "A trajectory cannot contain both exchanges and legacy histories" + ) + return self + + def __enter__(self) -> Trajectory: + from ._scope import enter_trajectory + + return enter_trajectory(self) + + def __exit__(self, *exc_info: Any) -> None: + from ._scope import exit_trajectory + + exit_trajectory(self, *exc_info) + + def log(self, message: str) -> None: + self.logs.append(message) + + def finish(self) -> Trajectory: + self.metrics["duration"] = (datetime.now() - self.start_time).total_seconds() + return self + + @asynccontextmanager + async def track_duration(self, metric_name: str) -> AsyncGenerator[None, None]: + start_time = time.monotonic() + try: + yield + finally: + duration = time.monotonic() - start_time + metric_key = f"{metric_name}_duration" + self.metrics[metric_key] = self.metrics.get(metric_key, 0.0) + duration + + def __str__(self) -> str: + return f"Trajectory(reward={self.reward}, metrics={self.metrics}, metadata={self.metadata})" + + def model_dump(self, **kwargs: Any) -> dict[str, Any]: + kwargs.setdefault("exclude_defaults", True) + return super().model_dump(**kwargs) + + def model_dump_json(self, **kwargs: Any) -> str: + kwargs.setdefault("exclude_defaults", True) + return super().model_dump_json(**kwargs) + + def messages(self) -> Messages: + return get_messages(self.messages_and_choices) + + def for_logging(self) -> dict[str, Any]: + from ._compat import trajectory_for_logging + + return trajectory_for_logging(self) + + +class TrajectoryGroup(pydantic.BaseModel): + trajectories: list[Trajectory] = pydantic.Field(default_factory=list) + exceptions: list[PydanticException] = pydantic.Field(default_factory=list) + metadata: dict[str, Any] = pydantic.Field(default_factory=dict) + metrics: dict[str, float | int | bool] = pydantic.Field(default_factory=dict) + logs: list[str] = pydantic.Field(default_factory=list) + + @overload + def __new__( + cls, + trajectories: Iterable[Trajectory | BaseException] = (), + **kwargs: Any, + ) -> TrajectoryGroup: ... + + @overload + @deprecated("Use await art.trajectory_group(...) instead.") + def __new__( + cls, + trajectories: Iterable[Awaitable[Trajectory]], + **kwargs: Any, + ) -> Awaitable[TrajectoryGroup]: ... + + def __new__(cls, trajectories: Iterable[Any] = (), **kwargs: Any) -> Any: + from ._compat import new_trajectory_group + + return new_trajectory_group(cls, trajectories, kwargs) + + def __init__( + self, + trajectories: ( + Iterable[Trajectory | BaseException] | Iterable[Awaitable[Trajectory]] + ) = (), + *, + exceptions: Iterable[BaseException | PydanticException] = (), + metadata: dict[str, Any] | None = None, + metrics: dict[str, float | int | bool] | None = None, + logs: list[str] | None = None, + ) -> None: + from ._compat import init_trajectory_group + + init_trajectory_group( + self, + cast(Iterable[Trajectory | BaseException], trajectories), + exceptions=exceptions, + metadata=metadata, + metrics=metrics, + logs=logs, + ) + + def __enter__(self) -> TrajectoryGroup: + from ._scope import enter_trajectory_group + + return enter_trajectory_group(self) + + def __exit__(self, *exc_info: Any) -> None: + from ._scope import exit_trajectory_group + + exit_trajectory_group(self, *exc_info) + + def __copy__(self) -> TrajectoryGroup: + from ._compat import copy_trajectory_group + + return copy_trajectory_group(self) + + def __deepcopy__(self, memo: dict[int, Any] | None = None) -> TrajectoryGroup: + from ._compat import deepcopy_trajectory_group + + return deepcopy_trajectory_group(self, memo) + + def log(self, message: str) -> None: + self.logs.append(message) + + def __iter__(self) -> Iterator[Trajectory]: # type: ignore[override] + return iter(self.trajectories) + + def __len__(self) -> int: + return len(self.trajectories) + + def model_dump(self, **kwargs: Any) -> dict[str, Any]: + kwargs.setdefault("exclude_defaults", True) + return super().model_dump(**kwargs) + + def model_dump_json(self, **kwargs: Any) -> str: + kwargs.setdefault("exclude_defaults", True) + return super().model_dump_json(**kwargs) + + +class TokenizedTrajectory(pydantic.BaseModel): + token_ids: list[int] + logprobs: list[float] + assistant_mask: list[bool] + underlying: Trajectory + + +class TokenizedTrajectoryGroup(pydantic.BaseModel): + trajectories: list[TokenizedTrajectory] + underlying: TrajectoryGroup + + +@overload +def current_trajectory(*, required: Literal[True]) -> Trajectory: ... + + +@overload +def current_trajectory(*, required: Literal[False] = False) -> Trajectory | None: ... + + +def current_trajectory(*, required: bool = False) -> Trajectory | None: + from ._scope import get_current_trajectory + + return get_current_trajectory(required=required) + + +@overload +def current_trajectory_group(*, required: Literal[True]) -> TrajectoryGroup: ... + + +@overload +def current_trajectory_group( + *, required: Literal[False] = False +) -> TrajectoryGroup | None: ... + + +def current_trajectory_group(*, required: bool = False) -> TrajectoryGroup | None: + from ._scope import get_current_trajectory_group + + return get_current_trajectory_group(required=required) + + +async def trajectory(coroutine: Coroutine[Any, Any, Any]) -> Trajectory: + from ._scope import capture_trajectory + + return await capture_trajectory(coroutine) + + +async def trajectory_group( + trajectories: Iterable[Coroutine[Any, Any, Trajectory]], + *, + return_exceptions: bool = False, +) -> TrajectoryGroup: + from ._scope import capture_trajectory_group + + return await capture_trajectory_group( + trajectories, + return_exceptions=return_exceptions, + ) + + +def tokenize_trajectory( + trajectory: Trajectory, + base_model: str | None = None, + *, + model: str | None = None, + chat_template: str | None = None, + chat_template_kwargs: dict[str, Any] | 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, + **kwargs: Any, +) -> list[TokenizedTrajectory]: + return [tokenize_trajectory(item, base_model, **kwargs) for item in trajectories] + + +def tokenize_trajectory_group( + group: TrajectoryGroup, + base_model: str | None = None, + **kwargs: Any, +) -> TokenizedTrajectoryGroup: + return TokenizedTrajectoryGroup( + trajectories=tokenize_trajectories(group, base_model, **kwargs), + underlying=group, + ) + + +def tokenize_trajectory_groups( + groups: Iterable[TrajectoryGroup], + base_model: str | None = None, + **kwargs: Any, +) -> list[TokenizedTrajectoryGroup]: + return [tokenize_trajectory_group(group, base_model, **kwargs) for group in groups] + + +@deprecated("Use current_trajectory() instead.") +def auto_trajectory(*, required: bool = False) -> Trajectory | None: + return current_trajectory(required=required) + + +@deprecated("Use trajectory() instead.") +async def capture_auto_trajectory( + coroutine: Coroutine[Any, Any, Any], +) -> Trajectory: + return await trajectory(coroutine) + + +def get_messages(messages_and_choices: MessagesAndChoices) -> Messages: + from ._compat import messages_from_legacy_history + + return messages_from_legacy_history(messages_and_choices) + + +__all__ = [ + "ChatCompletionsRequest", + "CompletionsRequest", + "ResponsesRequest", + "MessagesRequest", + "ChatCompletionsExchange", + "CompletionsExchange", + "ResponsesExchange", + "MessagesExchange", + "TrajectoryExchanges", + "PydanticException", + "History", + "Trajectory", + "TrajectoryGroup", + "TokenizedTrajectory", + "TokenizedTrajectoryGroup", + "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/_capture/__init__.py b/src/art/trajectories/_capture/__init__.py new file mode 100644 index 000000000..01edc0b32 --- /dev/null +++ b/src/art/trajectories/_capture/__init__.py @@ -0,0 +1,23 @@ +from __future__ import annotations + +from threading import Lock + +_installed = False +_lock = Lock() + + +def install() -> None: + global _installed + if _installed: + return + with _lock: + if _installed: + return + from .aiohttp import install as install_aiohttp + from .httpx import install as install_httpx + from .requests import install as install_requests + + install_httpx() + install_aiohttp() + install_requests() + _installed = True diff --git a/src/art/trajectories/_capture/aiohttp.py b/src/art/trajectories/_capture/aiohttp.py new file mode 100644 index 000000000..660a0678b --- /dev/null +++ b/src/art/trajectories/_capture/aiohttp.py @@ -0,0 +1,81 @@ +from __future__ import annotations + +from collections.abc import AsyncIterator +from typing import Any + +import aiohttp + +from .core import CaptureState, begin, reset + + +class _CapturedStream: + def __init__(self, stream: Any, state: CaptureState) -> None: + self._stream = stream + self._state = state + + def __getattr__(self, name: str) -> Any: + return getattr(self._stream, name) + + def _record(self, value: Any) -> Any: + chunk = value[0] if isinstance(value, tuple) else value + if isinstance(chunk, bytes): + self._state.add(chunk) + if self._stream.at_eof(): + self._state.finish() + return value + + async def read(self, *args: Any, **kwargs: Any) -> bytes: + return self._record(await self._stream.read(*args, **kwargs)) + + async def readany(self) -> bytes: + return self._record(await self._stream.readany()) + + async def readline(self) -> bytes: + return self._record(await self._stream.readline()) + + async def readchunk(self) -> tuple[bytes, bool]: + return self._record(await self._stream.readchunk()) + + async def _iterate(self, iterator: Any) -> AsyncIterator[bytes]: + try: + async for chunk in iterator: + yield self._record(chunk) + finally: + self._state.finish() + + def __aiter__(self) -> AsyncIterator[bytes]: + return self._iterate(self._stream.__aiter__()) + + def iter_any(self) -> AsyncIterator[bytes]: + return self._iterate(self._stream.iter_any()) + + def iter_chunked(self, size: int) -> AsyncIterator[bytes]: + return self._iterate(self._stream.iter_chunked(size)) + + +def install() -> None: + if getattr(aiohttp.ClientSession._request, "_art_capture", False): + return + original = aiohttp.ClientSession._request + + async def request( + self: aiohttp.ClientSession, + method: str, + str_or_url: Any, + **kwargs: Any, + ) -> aiohttp.ClientResponse: + body: Any = kwargs.get("json") + if body is None: + body = kwargs.get("data") + state, token = begin(method, str(str_or_url), body) + try: + response = await original(self, method, str_or_url, **kwargs) + finally: + reset(token) + if state is not None: + state.status_code = response.status + response.content = _CapturedStream(response.content, state) # type: ignore[assignment] + return response + + request._art_capture = True # type: ignore[attr-defined] + aiohttp.ClientSession._request = request # type: ignore[method-assign] diff --git a/src/art/trajectories/_capture/core.py b/src/art/trajectories/_capture/core.py new file mode 100644 index 000000000..edeb5e5b7 --- /dev/null +++ b/src/art/trajectories/_capture/core.py @@ -0,0 +1,91 @@ +from __future__ import annotations + +import contextvars +from dataclasses import dataclass, field +from datetime import datetime +import json +import logging +from typing import Any + +from .._protocols import build_exchange, endpoint_for_url +from .._scope import get_current_trajectory + +logger = logging.getLogger(__name__) +_adapter_active: contextvars.ContextVar[bool] = contextvars.ContextVar( + "art_capture_adapter_active", default=False +) + + +@dataclass +class CaptureState: + trajectory: Any + endpoint: str + request: dict[str, Any] + start_time: datetime = field(default_factory=datetime.now) + status_code: int | None = None + body: bytearray = field(default_factory=bytearray) + captured: bool = False + + def add(self, chunk: bytes) -> None: + if not self.captured: + self.body.extend(chunk) + + def finish(self) -> None: + if self.captured: + return + self.captured = True + if self.status_code is None or not 200 <= self.status_code < 300: + return + try: + name, exchange = build_exchange( + self.endpoint, + self.request, + bytes(self.body), + start_time=self.start_time, + end_time=datetime.now(), + ) + except Exception as exc: + logger.debug("Ignoring incomplete trajectory exchange: %s", exc) + return + getattr(self.trajectory.exchanges, name).append(exchange) + + +def _json_body(value: Any) -> dict[str, Any] | None: + if isinstance(value, dict): + return value + if isinstance(value, str): + value = value.encode() + if not isinstance(value, bytes): + return None + try: + parsed = json.loads(value) + except (UnicodeDecodeError, json.JSONDecodeError): + return None + return parsed if isinstance(parsed, dict) else None + + +def begin( + method: str, + url: str, + body: Any, +) -> tuple[CaptureState | None, contextvars.Token[bool] | None]: + trajectory = get_current_trajectory(required=False) + endpoint = endpoint_for_url(url) + request = _json_body(body) + if ( + trajectory is None + or method.upper() != "POST" + or endpoint is None + or request is None + or _adapter_active.get() + ): + return None, None + return ( + CaptureState(trajectory=trajectory, endpoint=endpoint, request=request), + _adapter_active.set(True), + ) + + +def reset(token: contextvars.Token[bool] | None) -> None: + if token is not None: + _adapter_active.reset(token) diff --git a/src/art/trajectories/_capture/httpx.py b/src/art/trajectories/_capture/httpx.py new file mode 100644 index 000000000..58d6dc3a9 --- /dev/null +++ b/src/art/trajectories/_capture/httpx.py @@ -0,0 +1,106 @@ +from __future__ import annotations + +from collections.abc import AsyncIterator, Iterator +from typing import Any + +import httpx + +from .core import CaptureState, begin, reset + +_STATE = "_art_trajectory_capture" + + +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_close = httpx.Response.close + original_aclose = httpx.Response.aclose + + def send( + self: httpx.Client, request: httpx.Request, **kwargs: Any + ) -> httpx.Response: + try: + body = request.content + except httpx.RequestNotRead: + body = None + state, token = begin(request.method, str(request.url), body) + try: + response = original_send(self, request, **kwargs) + finally: + reset(token) + 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() + return response + + async def async_send( + self: httpx.AsyncClient, request: httpx.Request, **kwargs: Any + ) -> httpx.Response: + try: + body = request.content + except httpx.RequestNotRead: + body = None + state, token = begin(request.method, str(request.url), body) + try: + response = await original_async_send(self, request, **kwargs) + finally: + reset(token) + 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() + return response + + def iter_bytes( + self: httpx.Response, chunk_size: int | None = None + ) -> Iterator[bytes]: + state: CaptureState | None = getattr(self, _STATE, None) + try: + for chunk in original_iter(self, chunk_size): + if state is not None: + state.add(chunk) + yield chunk + finally: + if state is not None: + state.finish() + + async def aiter_bytes( + self: httpx.Response, chunk_size: int | None = None + ) -> AsyncIterator[bytes]: + state: CaptureState | None = getattr(self, _STATE, None) + try: + async for chunk in original_aiter(self, chunk_size): + if state is not None: + state.add(chunk) + yield chunk + finally: + if state is not None: + state.finish() + + def close(self: httpx.Response) -> None: + original_close(self) + if state := getattr(self, _STATE, None): + state.finish() + + async def aclose(self: httpx.Response) -> None: + await original_aclose(self) + if state := getattr(self, _STATE, None): + state.finish() + + send._art_capture = True # type: ignore[attr-defined] + async_send._art_capture = True # type: ignore[attr-defined] + httpx.Client.send = send # type: ignore[method-assign] + httpx.AsyncClient.send = async_send # type: ignore[method-assign] + httpx.Response.iter_bytes = iter_bytes # type: ignore[method-assign] + httpx.Response.aiter_bytes = aiter_bytes # type: ignore[method-assign] + httpx.Response.close = close # type: ignore[method-assign] + httpx.Response.aclose = aclose # type: ignore[method-assign] diff --git a/src/art/trajectories/_capture/requests.py b/src/art/trajectories/_capture/requests.py new file mode 100644 index 000000000..fb62536ca --- /dev/null +++ b/src/art/trajectories/_capture/requests.py @@ -0,0 +1,53 @@ +from __future__ import annotations + +from collections.abc import Iterator +from typing import Any + +import requests + +from .core import CaptureState, begin, reset + +_STATE = "_art_trajectory_capture" + + +def install() -> None: + if getattr(requests.Session.send, "_art_capture", False): + return + original_send = requests.Session.send + original_iter = requests.Response.iter_content + + def send( + self: requests.Session, request: requests.PreparedRequest, **kwargs: Any + ) -> requests.Response: + state, token = begin(request.method or "GET", request.url or "", request.body) + try: + response = original_send(self, request, **kwargs) + finally: + reset(token) + 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() + return response + + def iter_content( + self: requests.Response, *args: Any, **kwargs: Any + ) -> Iterator[Any]: + state: CaptureState | None = getattr(self, _STATE, None) + try: + for chunk in original_iter(self, *args, **kwargs): + if state is not None: + if isinstance(chunk, str): + chunk = chunk.encode(self.encoding or "utf-8") + if isinstance(chunk, bytes): + state.add(chunk) + yield chunk + finally: + if state is not None: + state.finish() + + send._art_capture = True # type: ignore[attr-defined] + requests.Session.send = send # type: ignore[method-assign] + requests.Response.iter_content = iter_content # type: ignore[method-assign] diff --git a/src/art/trajectories/_compat.py b/src/art/trajectories/_compat.py new file mode 100644 index 000000000..d6d08084c --- /dev/null +++ b/src/art/trajectories/_compat.py @@ -0,0 +1,198 @@ +from __future__ import annotations + +import asyncio +from collections.abc import Iterable +import copy +import traceback +from typing import Any, cast +import warnings + +from openai.types.chat.chat_completion import Choice +import pydantic + +from ..types import Message, Messages, MessagesAndChoices +from . import PydanticException, Trajectory, TrajectoryGroup + + +def exception_model( + exception: BaseException | PydanticException | dict[str, Any], +) -> PydanticException: + if isinstance(exception, PydanticException): + return exception + if isinstance(exception, dict): + return PydanticException.model_validate(exception) + return PydanticException( + type=str(type(exception)), + message=str(exception), + traceback="".join( + traceback.format_exception( + type(exception), exception, exception.__traceback__ + ) + ), + ) + + +async def _legacy_async_group( + items: list[Any], kwargs: dict[str, Any] +) -> TrajectoryGroup: + from ..gather import get_gather_context, record_metrics + + context = get_gather_context() + trajectories: list[Trajectory] = [] + exceptions = list(kwargs.pop("exceptions", ())) + for future in asyncio.as_completed(items): + try: + item = await future + trajectories.append(item) + record_metrics(context, item) + context.update_pbar(n=1) + except BaseException as exc: + exceptions.append(exc) + context.metric_sums["exceptions"] += 1 + context.update_pbar(n=0) + if context.too_many_exceptions(): + raise + return TrajectoryGroup(trajectories, exceptions=exceptions, **kwargs) + + +class _LegacyGroupCoroutine: + def __init__(self, coroutine: Any, size: int) -> None: + self.coroutine = coroutine + self._num_trajectories = size + + def __await__(self) -> Any: + return self.coroutine.__await__() + + +def new_trajectory_group( + cls: type[TrajectoryGroup], trajectories: Iterable[Any], kwargs: dict[str, Any] +) -> Any: + items = list(trajectories) + if any(hasattr(item, "__await__") for item in items): + warnings.warn( + "Awaiting TrajectoryGroup(...) is deprecated; use art.trajectory_group(...).", + DeprecationWarning, + stacklevel=2, + ) + return _LegacyGroupCoroutine( + _legacy_async_group(items, dict(kwargs)), len(items) + ) + group = object.__new__(cls) + group.__init__(items, **kwargs) + return group + + +def init_trajectory_group( + group: TrajectoryGroup, + trajectories: Iterable[Trajectory | BaseException], + *, + exceptions: Iterable[BaseException | PydanticException], + metadata: dict[str, Any] | None, + metrics: dict[str, float | int | bool] | None, + logs: list[str] | None, +) -> None: + items = list(trajectories) + normalized_trajectories = [ + item if isinstance(item, Trajectory) else Trajectory.model_validate(item) + for item in items + if not isinstance(item, BaseException) + ] + pydantic.BaseModel.__init__( + group, + trajectories=normalized_trajectories or getattr(group, "trajectories", []), + exceptions=[ + exception_model(item) + for item in [ + *(item for item in items if isinstance(item, BaseException)), + *exceptions, + ] + ] + or getattr(group, "exceptions", []), + metadata=metadata if metadata is not None else getattr(group, "metadata", {}), + metrics=metrics if metrics is not None else getattr(group, "metrics", {}), + logs=logs if logs is not None else getattr(group, "logs", []), + ) + + +def copy_trajectory_group(group: TrajectoryGroup) -> TrajectoryGroup: + copied = TrajectoryGroup( + group.trajectories[:], + metadata=group.metadata.copy(), + metrics=group.metrics.copy(), + logs=group.logs[:], + ) + copied.exceptions = group.exceptions[:] + return copied + + +def deepcopy_trajectory_group( + group: TrajectoryGroup, memo: dict[int, Any] | None +) -> TrajectoryGroup: + memo = {} if memo is None else memo + if id(group) in memo: + return memo[id(group)] + copied = TrajectoryGroup( + copy.deepcopy(group.trajectories, memo), + metadata=copy.deepcopy(group.metadata, memo), + metrics=copy.deepcopy(group.metrics, memo), + logs=copy.deepcopy(group.logs, memo), + ) + memo[id(group)] = copied + copied.exceptions = copy.deepcopy(group.exceptions, memo) + return copied + + +def messages_from_legacy_history(messages_and_choices: MessagesAndChoices) -> Messages: + messages: Messages = [] + for item in messages_and_choices: + if isinstance(item, Choice): + content = item.message.content or "" + tool_calls = item.message.tool_calls or [] + message: Message = cast( + Message, + { + "role": "assistant", + "content": content, + **( + { + "tool_calls": [ + tool_call.model_dump(mode="json") + for tool_call in tool_calls + ] + } + if tool_calls + else {} + ), + }, + ) + messages.append(message) + else: + message = dict(item) + if message.get("content") is None: + message["content"] = "" + messages.append(message) # type: ignore[arg-type] + return messages + + +def trajectory_for_logging(trajectory: Trajectory) -> dict[str, Any]: + if trajectory.exchanges: + return trajectory.model_dump(mode="json", exclude={"start_time"}) + messages = [] + for item in trajectory.messages_and_choices: + if isinstance(item, Choice): + message = item.message.to_dict() + trainable = True + else: + message = cast(dict[str, Any], item) + trainable = False + messages.append({**message, "trainable": trainable}) + return { + "reward": trajectory.reward, + "initial_policy_version": trajectory.initial_policy_version, + "final_policy_version": trajectory.final_policy_version, + "metrics": trajectory.metrics, + "metadata": trajectory.metadata, + "messages": messages, + "tools": trajectory.tools, + "logs": trajectory.logs, + } diff --git a/src/art/trajectories/_protocols.py b/src/art/trajectories/_protocols.py new file mode 100644 index 000000000..d6c02c16c --- /dev/null +++ b/src/art/trajectories/_protocols.py @@ -0,0 +1,257 @@ +from __future__ import annotations + +from collections.abc import Mapping +from datetime import datetime +import json +from typing import Any, cast +from urllib.parse import urlsplit + +from anthropic._types import NOT_GIVEN +from anthropic.lib.streaming._messages import accumulate_event +from anthropic.types import Message, RawMessageStreamEvent +from openai.types import Completion +from openai.types.chat import ChatCompletion +from openai.types.chat.chat_completion_chunk import ChatCompletionChunk +from openai.types.responses import Response +from pydantic import TypeAdapter + +from ..openai import init_chat_completion, update_chat_completion +from . import ( + ChatCompletionsExchange, + ChatCompletionsRequest, + CompletionsExchange, + CompletionsRequest, + MessagesExchange, + MessagesRequest, + ResponsesExchange, + ResponsesRequest, +) + +Endpoint = str +_ENDPOINTS = { + "/v1/chat/completions": "chat_completions", + "/v1/completions": "completions", + "/v1/responses": "responses", + "/v1/messages": "messages", +} + + +def endpoint_for_url(url: str) -> Endpoint | None: + path = urlsplit(url).path.rstrip("/") + return next( + (value for suffix, value in _ENDPOINTS.items() if path.endswith(suffix)), None + ) + + +def _sse_events(body: bytes) -> list[tuple[str | None, dict[str, Any] | str]]: + text = body.decode("utf-8").replace("\r\n", "\n") + events: list[tuple[str | None, dict[str, Any] | str]] = [] + for block in text.split("\n\n"): + event_name: str | None = None + data_lines: list[str] = [] + for line in block.splitlines(): + if line.startswith("event:"): + event_name = line[6:].strip() + elif line.startswith("data:"): + data_lines.append(line[5:].lstrip()) + if not data_lines: + continue + raw = "\n".join(data_lines) + if raw == "[DONE]": + events.append((event_name, raw)) + else: + value = json.loads(raw) + if isinstance(value, dict): + events.append((event_name, value)) + return events + + +def _chat_response(body: bytes, *, stream: bool) -> ChatCompletion: + if not stream: + return ChatCompletion.model_validate_json(body) + response: ChatCompletion | None = None + done = False + for _, payload in _sse_events(body): + if payload == "[DONE]": + done = True + continue + chunk = ChatCompletionChunk.model_validate(payload) + if response is None: + response = init_chat_completion(chunk) + update_chat_completion(response, chunk) + if response is None or not done: + raise ValueError("Incomplete Chat Completions stream") + return response + + +def _completion_response(body: bytes, *, stream: bool) -> Completion: + if not stream: + return Completion.model_validate_json(body) + chunks: list[dict[str, Any]] = [] + done = False + for _, payload in _sse_events(body): + if payload == "[DONE]": + done = True + elif isinstance(payload, dict): + chunks.append(payload) + if not chunks or not done: + raise ValueError("Incomplete Completions stream") + data = dict(chunks[0]) + choices: dict[int, dict[str, Any]] = {} + for chunk in chunks: + if isinstance(chunk.get("usage"), dict): + data["usage"] = chunk["usage"] + for raw in chunk.get("choices") or []: + if not isinstance(raw, dict): + continue + index = raw.get("index") + if not isinstance(index, int): + continue + current = choices.setdefault( + index, + {"index": index, "text": "", "finish_reason": "stop"}, + ) + current["text"] += raw.get("text") or "" + if raw.get("finish_reason") is not None: + current["finish_reason"] = raw["finish_reason"] + for key in ("token_ids", "tokens", "token_logprobs", "text_offset"): + values = raw.get(key) + if isinstance(values, list): + current.setdefault(key, []).extend(values) + logprobs = raw.get("logprobs") + if isinstance(logprobs, dict): + target = current.setdefault("logprobs", {}) + for key, values in logprobs.items(): + if isinstance(values, list): + target.setdefault(key, []).extend(values) + elif values is not None: + target[key] = values + for key, value in raw.items(): + if key not in { + "finish_reason", + "index", + "logprobs", + "text", + "text_offset", + "token_ids", + "token_logprobs", + "tokens", + }: + current[key] = value + data["object"] = "text_completion" + data["choices"] = [choices[index] for index in sorted(choices)] + return Completion.model_validate(data) + + +def _responses_response(body: bytes, *, stream: bool) -> Response: + if not stream: + return Response.model_validate_json(body) + completed: dict[str, Any] | None = None + for event_name, payload in _sse_events(body): + if isinstance(payload, dict) and ( + event_name == "response.completed" + or payload.get("type") == "response.completed" + ): + value = payload.get("response") + if isinstance(value, dict): + completed = value + if completed is None: + raise ValueError("Incomplete Responses stream") + return Response.model_validate(completed) + + +def _messages_response(body: bytes, *, stream: bool) -> Message: + if not stream: + return Message.model_validate_json(body) + adapter = TypeAdapter(RawMessageStreamEvent) + snapshot: Any = None + complete = False + token_ids: list[int] = [] + logprobs: list[float] = [] + for event_name, payload in _sse_events(body): + if not isinstance(payload, dict): + continue + if event_name and "type" not in payload: + payload = {**payload, "type": event_name} + event = adapter.validate_python(payload) + snapshot = accumulate_event( + event=event, + current_snapshot=snapshot, + output_format=NOT_GIVEN, + ) + if event.type == "message_delta": + 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 + ): + 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] + 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") + if token_ids: + data["token_ids"] = token_ids + if logprobs: + data["logprobs"] = logprobs + return Message.model_validate(data) + + +def _model(request: Mapping[str, Any], response: Any) -> str | None: + requested = request.get("model") + if isinstance(requested, str): + return requested + returned = getattr(response, "model", None) + return returned if isinstance(returned, str) else None + + +def build_exchange( + endpoint: Endpoint, + request: dict[str, Any], + body: bytes, + *, + start_time: datetime, + end_time: datetime, +) -> tuple[str, Any]: + stream = request.get("stream") is True + if endpoint == "chat_completions": + response = _chat_response(body, stream=stream) + return endpoint, ChatCompletionsExchange( + request=ChatCompletionsRequest(request), + response=response, + model=_model(request, response), + start_time=start_time, + end_time=end_time, + ) + if endpoint == "completions": + response = _completion_response(body, stream=stream) + return endpoint, CompletionsExchange( + request=CompletionsRequest(request), + response=response, + model=_model(request, response), + start_time=start_time, + end_time=end_time, + ) + if endpoint == "responses": + response = _responses_response(body, stream=stream) + return endpoint, ResponsesExchange( + request=ResponsesRequest(request), + response=response, + model=_model(request, response), + start_time=start_time, + end_time=end_time, + ) + if endpoint == "messages": + response = _messages_response(body, stream=stream) + return endpoint, MessagesExchange( + request=MessagesRequest(request), + response=response, + model=_model(request, response), + start_time=start_time, + end_time=end_time, + ) + raise ValueError(f"Unsupported trajectory endpoint: {endpoint}") diff --git a/src/art/trajectories/_scope.py b/src/art/trajectories/_scope.py new file mode 100644 index 000000000..1a56e1ac3 --- /dev/null +++ b/src/art/trajectories/_scope.py @@ -0,0 +1,100 @@ +from __future__ import annotations + +import asyncio +from collections.abc import Coroutine, Iterable +import contextvars +from typing import Any + +from . import PydanticException, Trajectory, TrajectoryGroup +from ._compat import exception_model + +_trajectories: contextvars.ContextVar[tuple[Trajectory, ...]] = contextvars.ContextVar( + "art_trajectories", default=() +) +_groups: contextvars.ContextVar[tuple[TrajectoryGroup, ...]] = contextvars.ContextVar( + "art_trajectory_groups", default=() +) + + +def get_current_trajectory(*, required: bool) -> Trajectory | None: + current = _trajectories.get() + if current: + return current[-1] + if required: + raise RuntimeError("No trajectory is active in this context") + return None + + +def get_current_trajectory_group(*, required: bool) -> TrajectoryGroup | None: + current = _groups.get() + if current: + return current[-1] + if required: + raise RuntimeError("No trajectory group is active in this context") + return None + + +def enter_trajectory(trajectory: Trajectory) -> Trajectory: + from ._capture import install + + install() + _trajectories.set((*_trajectories.get(), trajectory)) + return trajectory + + +def exit_trajectory(trajectory: Trajectory, *exc_info: Any) -> None: + current = _trajectories.get() + if not current or current[-1] is not trajectory: + raise RuntimeError("Trajectory contexts must exit in stack order") + _trajectories.set(current[:-1]) + trajectory.finish() + group = get_current_trajectory_group(required=False) + if group is not None and all(item is not trajectory for item in group.trajectories): + group.trajectories.append(trajectory) + + +def enter_trajectory_group(group: TrajectoryGroup) -> TrajectoryGroup: + _groups.set((*_groups.get(), group)) + return group + + +def exit_trajectory_group(group: TrajectoryGroup, *exc_info: Any) -> None: + current = _groups.get() + if not current or current[-1] is not group: + raise RuntimeError("TrajectoryGroup contexts must exit in stack order") + _groups.set(current[:-1]) + + +def _require_raw_coroutine(value: Any) -> None: + if isinstance(value, (asyncio.Task, asyncio.Future)) or not isinstance( + value, Coroutine + ): + raise TypeError("Expected a raw coroutine, not a Task, Future, or awaitable") + + +async def capture_trajectory(coroutine: Coroutine[Any, Any, Any]) -> Trajectory: + _require_raw_coroutine(coroutine) + with Trajectory() as captured: + await coroutine + return captured + + +async def capture_trajectory_group( + trajectories: Iterable[Coroutine[Any, Any, Trajectory]], + *, + return_exceptions: bool, +) -> TrajectoryGroup: + coroutines = list(trajectories) + for coroutine in coroutines: + _require_raw_coroutine(coroutine) + results = await asyncio.gather(*coroutines, return_exceptions=return_exceptions) + if not return_exceptions: + return TrajectoryGroup(results) # type: ignore[arg-type] + completed: list[Trajectory] = [] + exceptions: list[PydanticException] = [] + for result in results: + if isinstance(result, BaseException): + exceptions.append(exception_model(result)) + else: + completed.append(result) + return TrajectoryGroup(completed, exceptions=exceptions) diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py new file mode 100644 index 000000000..ab9d1437f --- /dev/null +++ b/src/art/trajectories/_tokenize.py @@ -0,0 +1,749 @@ +from __future__ import annotations + +from dataclasses import dataclass +import math +import re +from typing import Any, cast + +from openai.types.chat.chat_completion import Choice + +from . import ( + ChatCompletionsExchange, + CompletionsExchange, + MessagesExchange, + ResponsesExchange, + TokenizedTrajectory, + Trajectory, +) + +_TOKEN_ID = re.compile(r"token_id:(\d+)$") + + +@dataclass +class _TokenizerConfig: + base_model: str + revision: str | None = None + chat_template: str | None = None + chat_template_kwargs: dict[str, Any] | None = None + + +def _dump(value: Any) -> dict[str, Any]: + if hasattr(value, "model_dump"): + result = value.model_dump(mode="python") + return result if isinstance(result, dict) else {} + return value if isinstance(value, dict) else {} + + +def _token_id(value: Any) -> int | None: + if isinstance(value, int) and not isinstance(value, bool): + return value + if isinstance(value, str) and (match := _TOKEN_ID.fullmatch(value)): + return int(match.group(1)) + return None + + +def _pairs(values: Any) -> tuple[list[int], list[float]]: + if not isinstance(values, list): + return [], [] + token_ids: list[int] = [] + logprobs: list[float] = [] + for value in values: + data = _dump(value) + token_id = _token_id(data.get("token_id")) + if token_id is None: + token_id = _token_id(data.get("token")) + if token_id is None: + return [], [] + logprob = data.get("logprob") + token_ids.append(token_id) + logprobs.append( + float(logprob) if isinstance(logprob, (int, float)) else math.nan + ) + return token_ids, logprobs + + +def _logprob_values(values: Any) -> list[float]: + if not isinstance(values, list): + return [] + result: list[float] = [] + for value in values: + logprob = _dump(value).get("logprob") + if not isinstance(logprob, (int, float)): + return [] + result.append(float(logprob)) + return result + + +def _chat_tokens(response: Any) -> tuple[list[int] | None, list[int], list[float]]: + if len(response.choices) != 1: + raise ValueError("Trajectory tokenization requires exactly one response choice") + choice = response.choices[0] + response_data = _dump(response) + choice_data = _dump(choice) + prompt = choice_data.get("prompt_token_ids") or response_data.get( + "prompt_token_ids" + ) + prompt_ids = ( + [token for value in prompt if (token := _token_id(value)) is not None] + if isinstance(prompt, list) + else None + ) + token_ids = [ + token + for value in choice_data.get("token_ids") or [] + if (token := _token_id(value)) is not None + ] + logprob_values = getattr(getattr(choice, "logprobs", None), "content", None) + if logprob_values is None: + logprob_values = getattr(getattr(choice, "logprobs", None), "refusal", None) + values = list(logprob_values or []) + pair_ids, logprobs = _pairs(values) + if token_ids and pair_ids and token_ids != pair_ids: + raise ValueError("Response token IDs disagree with choice logprobs") + return ( + prompt_ids, + token_ids or pair_ids, + logprobs or _logprob_values(values) or [math.nan] * len(token_ids), + ) + + +def _completion_tokens( + response: Any, +) -> tuple[list[int] | None, list[int], list[float]]: + if len(response.choices) != 1: + raise ValueError("Trajectory tokenization requires exactly one response choice") + choice = response.choices[0] + response_data = _dump(response) + choice_data = _dump(choice) + prompt = choice_data.get("prompt_token_ids") or response_data.get( + "prompt_token_ids" + ) + prompt_ids = list(prompt) if isinstance(prompt, list) else None + token_ids = [ + token + for value in choice_data.get("token_ids") or [] + if (token := _token_id(value)) is not None + ] + logprobs = _dump(getattr(choice, "logprobs", None)) + tokens = logprobs.get("tokens") or [] + pair_ids = [token for value in tokens if (token := _token_id(value)) is not None] + pair_logprobs = [ + float(value) if isinstance(value, (int, float)) else math.nan + for value in logprobs.get("token_logprobs") or [] + ] + if token_ids and pair_ids and token_ids != pair_ids: + raise ValueError("Response token IDs disagree with completion logprobs") + selected = token_ids or pair_ids + if selected and len(pair_logprobs) != len(selected): + pair_logprobs = [math.nan] * len(selected) + return prompt_ids, selected, pair_logprobs + + +def _responses_tokens(response: Any) -> tuple[None, list[int], list[float]]: + data = _dump(response) + token_ids, logprobs = _pairs(data.get("raw_output_tokens")) + if token_ids: + return None, token_ids, logprobs + for output in data.get("output") or []: + for content in _dump(output).get("content") or []: + values = _dump(content).get("logprobs") + token_ids, logprobs = _pairs(values) + if token_ids: + return None, token_ids, logprobs + if logprobs := _logprob_values(values): + return None, [], logprobs + return None, [], [] + + +def _messages_tokens(response: Any) -> tuple[None, list[int], list[float]]: + data = _dump(response) + token_ids = [ + token + for value in data.get("token_ids") or [] + if (token := _token_id(value)) is not None + ] + logprobs = [ + float(value) if isinstance(value, (int, float)) else math.nan + for value in data.get("logprobs") or [] + ] + if len(logprobs) != len(token_ids): + logprobs = [math.nan] * len(token_ids) + return None, token_ids, logprobs + + +def _exchange_list(trajectory: Trajectory, model: str | None) -> list[Any]: + exchanges = [ + *trajectory.exchanges.chat_completions, + *trajectory.exchanges.completions, + *trajectory.exchanges.responses, + *trajectory.exchanges.messages, + ] + if model is not None: + exchanges = [exchange for exchange in exchanges if exchange.model == model] + if not exchanges: + raise ValueError(f"Trajectory contains no exchanges for model {model!r}") + models = {exchange.model for exchange in exchanges} + if None in models: + raise ValueError("Every tokenized exchange must identify its model") + if len(models) != 1: + raise ValueError( + "Trajectory tokenization requires exactly one model; pass model= to select one" + ) + return sorted( + exchanges, key=lambda exchange: (exchange.start_time, exchange.end_time) + ) + + +def _artifact_config(model: str) -> _TokenizerConfig: + import wandb + + artifact_path = model.removeprefix("wandb-artifact:///") + artifact = getattr(wandb, "Api")().artifact(f"{artifact_path}:latest") + 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") + renderer = metadata.get("renderer") + renderer = renderer if isinstance(renderer, dict) else {} + kwargs = renderer.get("chat_template_kwargs") + return _TokenizerConfig( + base_model=base_model, + revision=( + renderer.get("tokenizer_revision") + if isinstance(renderer.get("tokenizer_revision"), str) + else None + ), + chat_template=( + renderer.get("chat_template") + if isinstance(renderer.get("chat_template"), str) + else None + ), + chat_template_kwargs=kwargs if isinstance(kwargs, dict) else None, + ) + + +def _tokenizer_config(model: str, base_model: str | None) -> _TokenizerConfig: + if model.startswith("wandb-artifact:///"): + config = _artifact_config(model) + if base_model is not None: + if base_model != config.base_model: + config.revision = None + config.base_model = base_model + return config + if base_model is not None: + return _TokenizerConfig(base_model) + return _TokenizerConfig(model) + + +def _load_tokenizer(config: _TokenizerConfig) -> Any: + try: + from transformers import AutoTokenizer + except ImportError as exc: + raise RuntimeError( + "Tokenizer fallback requires ART's backend or tinker dependencies" + ) from exc + try: + return AutoTokenizer.from_pretrained( + config.base_model, + revision=config.revision, + ) + except Exception as exc: + raise ValueError( + f"Could not load tokenizer for {config.base_model!r}; pass base_model explicitly" + ) from exc + + +def _ids(value: Any) -> list[int]: + if hasattr(value, "input_ids"): + value = value.input_ids + if hasattr(value, "tolist"): + value = value.tolist() + if isinstance(value, dict): + value = value.get("input_ids") + if isinstance(value, list) and value and isinstance(value[0], list): + value = value[0] + if not isinstance(value, list) or any(not isinstance(item, int) for item in value): + raise TypeError("Tokenizer did not return one token ID sequence") + return value + + +def _content_text(content: Any) -> str: + if isinstance(content, str): + return content + if not isinstance(content, list): + return "" + return "".join( + block.get("text", "") + for block in content + if isinstance(block, dict) + and block.get("type") in {"input_text", "output_text", "text"} + ) + + +def _anthropic_messages(request: dict[str, Any]) -> list[dict[str, Any]]: + messages: list[dict[str, Any]] = [] + system = request.get("system") + if system: + messages.append({"role": "system", "content": _content_text(system)}) + for raw in request.get("messages") or []: + if not isinstance(raw, dict): + continue + role = raw.get("role", "user") + content = raw.get("content") + if isinstance(content, str): + messages.append({"role": role, "content": content}) + continue + text = "" + reasoning = "" + tool_calls: list[dict[str, Any]] = [] + for block in content if isinstance(content, list) else (): + if not isinstance(block, dict): + continue + kind = block.get("type") + if kind == "text": + text += str(block.get("text") or "") + elif kind == "thinking": + reasoning += str(block.get("thinking") or "") + elif kind == "tool_use": + tool_calls.append( + { + "id": block.get("id"), + "type": "function", + "function": { + "name": block.get("name"), + "arguments": __import__("json").dumps( + block.get("input") or {} + ), + }, + } + ) + elif kind == "tool_result": + if text: + messages.append({"role": role, "content": text}) + text = "" + result = block.get("content", "") + messages.append( + { + "role": "tool", + "tool_call_id": block.get("tool_use_id", block.get("id")), + "content": ( + result if isinstance(result, str) else _content_text(result) + ), + } + ) + message: dict[str, Any] = {"role": role, "content": text} + if reasoning: + message["reasoning"] = reasoning + if tool_calls: + message["tool_calls"] = tool_calls + if text or reasoning or tool_calls or role == "assistant": + messages.append(message) + return messages + + +def _responses_messages(request: dict[str, Any]) -> list[dict[str, Any]]: + messages: list[dict[str, Any]] = [] + if instructions := request.get("instructions"): + messages.append({"role": "system", "content": instructions}) + value = request.get("input") + if isinstance(value, str): + messages.append({"role": "user", "content": value}) + elif isinstance(value, list): + for item in value: + if not isinstance(item, dict): + continue + if item.get("type") == "function_call_output": + messages.append( + { + "role": "tool", + "tool_call_id": item.get("call_id"), + "content": item.get("output", ""), + } + ) + elif item.get("type") == "function_call": + messages.append( + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": item.get("call_id"), + "type": "function", + "function": { + "name": item.get("name"), + "arguments": item.get("arguments", "{}"), + }, + } + ], + } + ) + elif item.get("role"): + messages.append( + { + "role": item["role"], + "content": _content_text(item.get("content")), + } + ) + return messages + + +def _openai_tools(tools: Any, *, dialect: str) -> Any: + if not isinstance(tools, list) or dialect == "chat": + return tools + normalized = [] + for tool in tools: + if not isinstance(tool, dict) or tool.get("type", "function") != "function": + normalized.append(tool) + continue + if dialect == "messages": + function = { + "name": tool.get("name"), + "description": tool.get("description"), + "parameters": tool.get("input_schema", {}), + } + else: + function = { + "name": tool.get("name"), + "description": tool.get("description"), + "parameters": tool.get("parameters", {}), + } + normalized.append( + { + "type": "function", + "function": { + key: value for key, value in function.items() if value is not None + }, + } + ) + return normalized + + +def _request_messages( + exchange: Any, messages_override: list[dict[str, Any]] | None = None +) -> tuple[list[dict[str, Any]], Any]: + request = exchange.request.root + if isinstance(exchange, ChatCompletionsExchange): + return list(request.get("messages") or []), request.get("tools") + if isinstance(exchange, MessagesExchange): + return _anthropic_messages(request), _openai_tools( + request.get("tools"), dialect="messages" + ) + if isinstance(exchange, ResponsesExchange): + return ( + messages_override + if messages_override is not None + else _responses_messages(request), + _openai_tools(request.get("tools"), dialect="responses"), + ) + raise TypeError("Completions requests do not use chat templates") + + +def _response_message(exchange: Any) -> dict[str, Any]: + if isinstance(exchange, ChatCompletionsExchange): + return exchange.response.choices[0].message.model_dump( + mode="python", exclude_none=True + ) + if isinstance(exchange, MessagesExchange): + data = exchange.response.model_dump(mode="python") + request = {"messages": [{"role": "assistant", "content": data["content"]}]} + return _anthropic_messages(request)[0] + if isinstance(exchange, ResponsesExchange): + data = exchange.response.model_dump(mode="python") + content = [] + tool_calls = [] + for item in data.get("output") or []: + if item.get("type") == "message": + content.extend(item.get("content") or []) + elif item.get("type") == "function_call": + tool_calls.append( + { + "id": item.get("call_id"), + "type": "function", + "function": { + "name": item.get("name"), + "arguments": item.get("arguments", "{}"), + }, + } + ) + message: dict[str, Any] = { + "role": "assistant", + "content": _content_text(content), + } + if tool_calls: + message["tool_calls"] = tool_calls + return message + raise TypeError("Completions responses do not use chat templates") + + +def _template_ids( + tokenizer: Any, + exchange: Any, + *, + completed: bool, + config: _TokenizerConfig, + chat_template: str | None, + chat_template_kwargs: dict[str, Any] | None, + messages_override: list[dict[str, Any]] | None = None, +) -> list[int]: + request = exchange.request.root + if isinstance(exchange, CompletionsExchange): + prompt = request.get("prompt", "") + if isinstance(prompt, list) and all(isinstance(item, int) for item in prompt): + prompt_ids = prompt + else: + prompt_ids = _ids(tokenizer(str(prompt), add_special_tokens=False)) + if not completed: + return prompt_ids + return [ + *prompt_ids, + *_ids( + tokenizer(exchange.response.choices[0].text, add_special_tokens=False) + ), + ] + + messages, tools = _request_messages(exchange, messages_override) + if completed: + messages.append(_response_message(exchange)) + request_kwargs = request.get("chat_template_kwargs") + kwargs = { + **(config.chat_template_kwargs or {}), + **(request_kwargs if isinstance(request_kwargs, dict) else {}), + **(chat_template_kwargs or {}), + } + if isinstance(exchange, MessagesExchange) and isinstance( + thinking := 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 request.get("chat_template") or config.chat_template + result = tokenizer.apply_chat_template( + messages, + tools=tools, + tokenize=True, + add_generation_prompt=not completed, + **({"chat_template": template} if isinstance(template, str) else {}), + **kwargs, + ) + return _ids(result) + + +def _exchange_tokens(exchange: Any) -> tuple[list[int] | None, list[int], list[float]]: + if isinstance(exchange, ChatCompletionsExchange): + return _chat_tokens(exchange.response) + if isinstance(exchange, CompletionsExchange): + return _completion_tokens(exchange.response) + if isinstance(exchange, ResponsesExchange): + return _responses_tokens(exchange.response) + if isinstance(exchange, MessagesExchange): + return _messages_tokens(exchange.response) + raise TypeError(f"Unknown exchange type: {type(exchange)!r}") + + +def _visible_logprobs(exchange: Any) -> 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: + data = _dump(entry) + raw_bytes = data.get("bytes") + text = ( + bytes(raw_bytes).decode("utf-8") + if isinstance(raw_bytes, list) + else data.get("token") + ) + logprob = data.get("logprob") + if isinstance(text, str) and isinstance(logprob, (int, float)): + values.append((text, float(logprob))) + elif isinstance(exchange, CompletionsExchange): + logprobs = exchange.response.choices[0].logprobs + if logprobs is not None: + for text, logprob in zip( + logprobs.tokens or [], logprobs.token_logprobs or [], strict=False + ): + if logprob is not None: + values.append((text, float(logprob))) + elif isinstance(exchange, ResponsesExchange): + for output in _dump(exchange.response).get("output") or []: + for content in _dump(output).get("content") or []: + for entry in _dump(content).get("logprobs") or []: + data = _dump(entry) + text = data.get("token") + logprob = data.get("logprob") + if isinstance(text, str) and isinstance(logprob, (int, float)): + values.append((text, float(logprob))) + return values + + +def _align_visible_logprobs( + tokenizer: Any, completion: list[int], exchange: Any +) -> list[float] | None: + values = _visible_logprobs(exchange) + if not values or not callable(tokenizer): + return None + aligned = [math.nan] * len(completion) + cursor = 0 + for text, logprob in values: + encoded = _ids(tokenizer(text, add_special_tokens=False)) + if len(encoded) != 1: + return None + try: + index = completion.index(encoded[0], cursor) + except ValueError: + return None + aligned[index] = logprob + cursor = index + 1 + return aligned + + +def _legacy_tokenize( + trajectory: Trajectory, + base_model: str | None, + *, + chat_template: str | None, + chat_template_kwargs: dict[str, Any] | None, +) -> TokenizedTrajectory: + if trajectory.additional_histories: + raise ValueError("Tokenization requires one history") + token_ids: list[int] = [] + logprobs: list[float] = [] + assistant_mask: list[bool] = [] + for item in trajectory.messages_and_choices: + if not isinstance(item, Choice): + continue + prompt, completion, completion_logprobs = _chat_tokens( + type( + "Response", (), {"choices": [item], "model_dump": lambda self, **_: {}} + )() + ) + if prompt is None or not completion: + raise ValueError( + "Legacy fallback tokenization is unavailable without exact choice token metadata" + ) + if not token_ids: + token_ids.extend(prompt) + logprobs.extend([math.nan] * len(prompt)) + assistant_mask.extend([False] * len(prompt)) + elif prompt[: len(token_ids)] != token_ids: + raise ValueError("Legacy trajectory does not form one append-only history") + else: + suffix = prompt[len(token_ids) :] + token_ids.extend(suffix) + logprobs.extend([math.nan] * len(suffix)) + assistant_mask.extend([False] * len(suffix)) + token_ids.extend(completion) + logprobs.extend(completion_logprobs) + assistant_mask.extend([True] * len(completion)) + if not token_ids: + raise ValueError("Trajectory contains no trainable choices") + return TokenizedTrajectory( + token_ids=token_ids, + logprobs=logprobs, + assistant_mask=assistant_mask, + underlying=trajectory, + ) + + +def tokenize_one( + trajectory: Trajectory, + base_model: str | None, + *, + model: str | None, + chat_template: str | None, + chat_template_kwargs: dict[str, Any] | None, + tokenizer_instance: Any = None, +) -> TokenizedTrajectory: + if not trajectory.exchanges: + return _legacy_tokenize( + trajectory, + base_model, + chat_template=chat_template, + chat_template_kwargs=chat_template_kwargs, + ) + exchanges = _exchange_list(trajectory, model) + selected_model = cast(str, exchanges[0].model) + config = _tokenizer_config(selected_model, base_model) + tokenizer = tokenizer_instance + token_ids: list[int] = [] + logprobs: list[float] = [] + assistant_mask: list[bool] = [] + response_histories: dict[str, list[dict[str, Any]]] = {} + + for exchange in exchanges: + messages_override = None + if isinstance(exchange, ResponsesExchange): + request = exchange.request.root + messages_override = _responses_messages(request) + previous = request.get("previous_response_id") + if previous is not None: + if not isinstance(previous, str) or previous not in response_histories: + raise ValueError( + "Responses exchange refers to a previous response outside this trajectory" + ) + messages_override = [ + *response_histories[previous], + *messages_override, + ] + response_histories[exchange.response.id] = [ + *messages_override, + _response_message(exchange), + ] + prompt, completion, completion_logprobs = _exchange_tokens(exchange) + if prompt is None or not completion: + tokenizer = tokenizer or _load_tokenizer(config) + if prompt is None: + prompt = _template_ids( + tokenizer, + exchange, + completed=False, + config=config, + chat_template=chat_template, + chat_template_kwargs=chat_template_kwargs, + messages_override=messages_override, + ) + if not completion: + completed = _template_ids( + tokenizer, + exchange, + completed=True, + config=config, + chat_template=chat_template, + chat_template_kwargs=chat_template_kwargs, + messages_override=messages_override, + ) + if completed[: len(prompt)] != prompt: + raise ValueError( + "Completed response does not extend its generation prompt" + ) + completion = completed[len(prompt) :] + completion_logprobs = _align_visible_logprobs( + tokenizer, completion, exchange + ) or [math.nan] * len(completion) + if not token_ids: + token_ids.extend(prompt) + logprobs.extend([math.nan] * len(prompt)) + assistant_mask.extend([False] * 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" + ) + else: + suffix = prompt[len(token_ids) :] + token_ids.extend(suffix) + logprobs.extend([math.nan] * len(suffix)) + assistant_mask.extend([False] * len(suffix)) + if len(completion_logprobs) != len(completion): + completion_logprobs = _align_visible_logprobs( + tokenizer, completion, exchange + ) or [math.nan] * len(completion) + token_ids.extend(completion) + logprobs.extend(completion_logprobs) + assistant_mask.extend([True] * len(completion)) + + return TokenizedTrajectory( + token_ids=token_ids, + logprobs=logprobs, + assistant_mask=assistant_mask, + underlying=trajectory, + ) diff --git a/tests/unit/test_auto_trajectory.py b/tests/unit/test_auto_trajectory.py index fadd6ab98..6a5386944 100644 --- a/tests/unit/test_auto_trajectory.py +++ b/tests/unit/test_auto_trajectory.py @@ -286,38 +286,13 @@ async def say_hi() -> str | None: # Use the capture_auto_trajectory utility to capture a trajectory automatically trajectory = await art.capture_auto_trajectory(say_hi()) - assert trajectory.messages_and_choices == [ - message, - Choice(**mock_response["choices"][0]), # ty:ignore[invalid-argument-type, not-subscriptable] - message, - Choice(**mock_response["choices"][0]), # ty:ignore[invalid-argument-type, not-subscriptable] - ] - assert trajectory.tools == tools - assert trajectory.additional_histories[0].messages_and_choices == [ - message, - { - "content": "Hello! How can I assist you today?", - "role": "assistant", - }, - message, - { - "content": "Hello! How can I assist you today?", - "role": "assistant", - }, - message, - Choice(**mock_response["choices"][0]), # ty:ignore[invalid-argument-type, not-subscriptable] - ] - assert trajectory.additional_histories[0].tools is None - assert trajectory.additional_histories[1].messages_and_choices == [ - message, - Choice(**mock_response["choices"][0]), # ty:ignore[invalid-argument-type, not-subscriptable] - ] - assert trajectory.additional_histories[1].tools == tools - assert trajectory.additional_histories[2].messages_and_choices == [ - message, - mock_stream_choice, - ] - assert trajectory.additional_histories[2].tools == tools + exchanges = trajectory.exchanges.chat_completions + assert len(exchanges) == 5 + assert exchanges[0].request.root["messages"] == [message] + assert exchanges[0].request.root["tools"] == tools + assert exchanges[0].response.choices[0] == Choice(**mock_response["choices"][0]) # ty:ignore[invalid-argument-type, not-subscriptable] + assert exchanges[-1].response.choices[0] == mock_stream_choice + assert all(exchange.model == "test" for exchange in exchanges) @pytest.mark.filterwarnings("ignore::UserWarning:pydantic") @@ -385,14 +360,8 @@ async def say_hi() -> str | None: # Use the capture_auto_trajectory utility to capture a trajectory automatically trajectory = await art.capture_auto_trajectory(say_hi()) - assert trajectory.messages_and_choices == [ - message, - Choice(**mock_response["choices"][0]), # ty:ignore[invalid-argument-type, not-subscriptable] - message, - Choice(**mock_response["choices"][0]), # ty:ignore[invalid-argument-type, not-subscriptable] - ] - assert trajectory.additional_histories[0].messages_and_choices == [ - message, - mock_stream_choice, - ] - assert trajectory.additional_histories[0].tools == tools + exchanges = trajectory.exchanges.chat_completions + assert len(exchanges) == 3 + assert exchanges[0].request.root["messages"] == [message] + assert exchanges[-1].response.choices[0] == mock_stream_choice + assert all(exchange.model == "test" for exchange in exchanges) diff --git a/tests/unit/trajectories/test_capture.py b/tests/unit/trajectories/test_capture.py new file mode 100644 index 000000000..8e37b193f --- /dev/null +++ b/tests/unit/trajectories/test_capture.py @@ -0,0 +1,419 @@ +from __future__ import annotations + +import asyncio +from datetime import datetime, timedelta +import json +from typing import Any + +import aiohttp +from aiohttp import web +from anthropic import AsyncAnthropic +import httpx +from openai import AsyncOpenAI +import pytest +import pytest_asyncio +import requests + +import art +from art.trajectories._protocols import build_exchange + +CHAT: dict[str, Any] = { + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 1, + "model": "test/model", + "choices": [ + { + "index": 0, + "finish_reason": "stop", + "message": {"role": "assistant", "content": "hello"}, + "logprobs": { + "content": [ + { + "token": "token_id:2", + "logprob": -0.2, + "bytes": [104], + "top_logprobs": [], + } + ] + }, + "token_ids": [2], + "prompt_token_ids": [1], + } + ], +} + +COMPLETION: dict[str, Any] = { + "id": "cmpl-1", + "object": "text_completion", + "created": 1, + "model": "test/model", + "choices": [ + { + "index": 0, + "finish_reason": "stop", + "text": "hello", + "token_ids": [2], + "prompt_token_ids": [1], + "logprobs": { + "tokens": ["token_id:2"], + "token_logprobs": [-0.2], + "top_logprobs": [{}], + "text_offset": [0], + }, + } + ], +} + +RESPONSE: dict[str, Any] = { + "id": "resp_1", + "created_at": 1.0, + "model": "test/model", + "object": "response", + "output": [ + { + "id": "msg_1", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [ + { + "type": "output_text", + "text": "hello", + "annotations": [], + "logprobs": [], + } + ], + } + ], + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [], + "raw_output_tokens": [{"token_id": 2, "logprob": -0.2}], +} + +MESSAGE: dict[str, Any] = { + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "test/model", + "content": [{"type": "text", "text": "hello", "citations": None}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 1, "output_tokens": 1}, + "token_ids": [2], + "logprobs": [-0.2], +} + + +@pytest_asyncio.fixture +async def endpoint_server(unused_tcp_port: int): + async def handler(request: web.Request) -> web.Response: + request_body = await request.json() + if request_body.get("fail"): + return web.json_response({"error": "failed"}, status=400) + if request_body.get("incomplete"): + return web.Response( + body=_sse([(None, {"type": "incomplete"})]), + content_type="text/event-stream", + ) + bodies = { + "/v1/chat/completions": CHAT, + "/v1/completions": COMPLETION, + "/v1/responses": RESPONSE, + "/v1/messages": MESSAGE, + } + return web.json_response(bodies[request.path]) + + app = web.Application() + app.router.add_post("/v1/{tail:.*}", handler) + runner = web.AppRunner(app) + await runner.setup() + site = web.TCPSite(runner, "127.0.0.1", unused_tcp_port) + await site.start() + yield f"http://127.0.0.1:{unused_tcp_port}/v1" + await runner.cleanup() + + +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 + with art.Trajectory() as inner: + assert art.current_trajectory() is inner + assert art.current_trajectory() is outer + + async def child() -> art.Trajectory: + with art.Trajectory() as item: + await asyncio.sleep(0) + assert art.current_trajectory() is item + return item + + first, second = await asyncio.gather(child(), child()) + assert first is not second + assert art.current_trajectory() is None + with pytest.raises(RuntimeError, match="No trajectory"): + art.current_trajectory(required=True) + + +async def test_group_context_and_async_helpers() -> None: + with art.TrajectoryGroup() as group: + with art.Trajectory() as first: + pass + with art.Trajectory() as second: + pass + assert group.trajectories == [first, second] + + async def rollout() -> None: + await asyncio.sleep(0) + + captured = await art.trajectory(rollout()) + assert isinstance(captured, art.Trajectory) + task = asyncio.create_task(rollout()) + with pytest.raises(TypeError, match="raw coroutine"): + await art.trajectory(task) # type: ignore[arg-type] + await task + + async def failed() -> art.Trajectory: + raise ValueError("boom") + + successful = art.trajectory(rollout()) + result = await art.trajectory_group([successful, failed()], return_exceptions=True) + assert len(result.trajectories) == 1 + assert result.exceptions[0].message == "boom" + + +async def test_httpx_requests_and_aiohttp_capture_once(endpoint_server: str) -> None: + body = {"model": "test/model", "messages": [{"role": "user", "content": "hi"}]} + + def requests_stream() -> None: + with requests.post( + f"{endpoint_server}/chat/completions", + json=body, + stream=True, + timeout=5, + ) as response: + list(response.iter_content(chunk_size=5, decode_unicode=True)) + + with art.Trajectory() as trajectory: + async with httpx.AsyncClient() as client: + response = await client.post( + f"{endpoint_server}/chat/completions", json=body + ) + response.raise_for_status() + + await asyncio.to_thread(requests_stream) + + async with aiohttp.ClientSession() as session: + async with session.post( + f"{endpoint_server}/chat/completions", json=body + ) as response: + await response.json() + + assert len(trajectory.exchanges.chat_completions) == 3 + assert all( + exchange.response.choices[0].message.content == "hello" + for exchange in trajectory.exchanges.chat_completions + ) + + +async def test_native_openai_and_anthropic_sdks(endpoint_server: str) -> None: + openai = AsyncOpenAI(base_url=endpoint_server, api_key="test") + anthropic = AsyncAnthropic( + base_url=endpoint_server.removesuffix("/v1"), api_key="test" + ) + with art.Trajectory() as trajectory: + completion = await openai.completions.create(model="test/model", prompt="hi") + response = await openai.responses.create(model="test/model", input="hi") + message = await anthropic.messages.create( + model="test/model", + max_tokens=16, + messages=[{"role": "user", "content": "hi"}], + ) + await openai.close() + await anthropic.close() + + assert completion.choices[0].text == "hello" + assert response.output_text == "hello" + assert message.content[0].text == "hello" + assert len(trajectory.exchanges.completions) == 1 + assert len(trajectory.exchanges.responses) == 1 + assert len(trajectory.exchanges.messages) == 1 + + +async def test_failed_and_incomplete_calls_are_excluded(endpoint_server: str) -> None: + async with httpx.AsyncClient() as client: + with art.Trajectory() as trajectory: + await client.post( + f"{endpoint_server}/chat/completions", + json={"model": "test/model", "messages": [], "fail": True}, + ) + await client.post( + f"{endpoint_server}/chat/completions", + json={ + "model": "test/model", + "messages": [], + "stream": True, + "incomplete": True, + }, + ) + assert not trajectory.exchanges + + +def test_all_protocols_reconstruct_typed_responses() -> None: + now = datetime.now() + values = [ + ("chat_completions", {"model": "request-model", "messages": []}, CHAT), + ("completions", {"model": "request-model", "prompt": "hi"}, COMPLETION), + ("responses", {"input": "hi"}, RESPONSE), + ("messages", {"model": "request-model", "messages": []}, MESSAGE), + ] + for endpoint, request, response in values: + name, exchange = build_exchange( + endpoint, + request, + json.dumps(response).encode(), + start_time=now, + end_time=now + timedelta(seconds=1), + ) + assert name == endpoint + assert exchange.end_time > exchange.start_time + expected = request.get("model", "test/model") + assert exchange.model == expected + assert exchange.model_dump(mode="json")["request"] == request + + +def _sse(events: list[tuple[str | None, dict[str, Any] | str]]) -> bytes: + return "".join( + f"{f'event: {name}\n' if name else ''}data: " + f"{value if isinstance(value, str) else json.dumps(value)}\n\n" + for name, value in events + ).encode() + + +def test_all_streaming_protocols_reconstruct_final_responses() -> None: + now = datetime.now() + chat_chunk = { + "id": "chatcmpl-1", + "object": "chat.completion.chunk", + "created": 1, + "model": "test/model", + "choices": [ + { + "index": 0, + "delta": {"role": "assistant", "content": "hello"}, + "finish_reason": "stop", + "logprobs": None, + } + ], + } + completion_chunk = { + **COMPLETION, + "object": "text_completion.chunk", + "choices": [ + { + **COMPLETION["choices"][0], + "text": "hello", + "finish_reason": None, + } + ], + } + response_event = {"type": "response.completed", "response": RESPONSE} + message_events = [ + ( + "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", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 1}, + "token_ids": [2], + "logprobs": [-0.2], + }, + ), + ("message_stop", {"type": "message_stop"}), + ] + values = [ + ( + "chat_completions", + {"model": "test/model", "messages": [], "stream": True}, + _sse([(None, chat_chunk), (None, "[DONE]")]), + ), + ( + "completions", + {"model": "test/model", "prompt": "hi", "stream": True}, + _sse([(None, completion_chunk), (None, "[DONE]")]), + ), + ( + "responses", + {"model": "test/model", "input": "hi", "stream": True}, + _sse([("response.completed", response_event)]), + ), + ( + "messages", + {"model": "test/model", "messages": [], "stream": True}, + _sse(message_events), + ), + ] + for endpoint, request, body in values: + _, exchange = build_exchange( + endpoint, + request, + body, + start_time=now, + end_time=now + timedelta(seconds=1), + ) + assert exchange.model == "test/model" + if endpoint == "messages": + assert exchange.response.content[0].text == "hello" + + +def test_trajectory_rejects_mixed_representations() -> None: + _, exchange = build_exchange( + "chat_completions", + {"model": "test/model", "messages": []}, + json.dumps(CHAT).encode(), + start_time=datetime.now(), + end_time=datetime.now(), + ) + with pytest.raises(ValueError, match="both exchanges and legacy histories"): + art.Trajectory( + exchanges=art.TrajectoryExchanges(chat_completions=[exchange]), + messages_and_choices=[{"role": "user", "content": "hi"}], + ) + + +def test_metadata_accepts_json_serializable_values() -> None: + assert art.Trajectory().model_dump() == {} + trajectory = art.Trajectory(metadata={"nested": {"items": [1, "two"]}}) + assert trajectory.model_dump(mode="json")["metadata"] == { + "nested": {"items": [1, "two"]} + } diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py new file mode 100644 index 000000000..e102f56eb --- /dev/null +++ b/tests/unit/trajectories/test_tokenize.py @@ -0,0 +1,445 @@ +from __future__ import annotations + +from datetime import datetime, timedelta +import math +from types import SimpleNamespace +from typing import Any + +from anthropic.types import Message +from openai.types.chat import ChatCompletion +from openai.types.chat.chat_completion_token_logprob import ChatCompletionTokenLogprob +from openai.types.responses import Response +import pytest + +import art +from art.trajectories import ( + ChatCompletionsExchange, + ChatCompletionsRequest, + MessagesExchange, + MessagesRequest, + ResponsesExchange, + ResponsesRequest, + TrajectoryExchanges, +) + + +def _chat_exchange( + prompt: list[int], + output: list[int], + *, + model: str = "test/model", + offset: int = 0, +) -> ChatCompletionsExchange: + response = ChatCompletion.model_validate( + { + "id": f"chat-{offset}", + "object": "chat.completion", + "created": offset, + "model": model, + "choices": [ + { + "index": 0, + "finish_reason": "stop", + "message": {"role": "assistant", "content": "answer"}, + "prompt_token_ids": prompt, + "token_ids": output, + "logprobs": { + "content": [ + { + "token": f"token_id:{token}", + "logprob": -token / 10, + "bytes": [], + "top_logprobs": [], + } + for token in output + ] + }, + } + ], + } + ) + start = datetime(2026, 1, 1) + timedelta(seconds=offset) + return ChatCompletionsExchange( + request=ChatCompletionsRequest( + { + "model": model, + "messages": [{"role": "user", "content": f"turn {offset}"}], + } + ), + response=response, + model=model, + start_time=start, + end_time=start + timedelta(milliseconds=1), + ) + + +def test_exact_tokens_form_one_append_only_history_without_tokenizer() -> None: + trajectory = art.Trajectory( + exchanges=TrajectoryExchanges( + chat_completions=[ + _chat_exchange([1], [2], offset=0), + _chat_exchange([1, 2, 3], [4], offset=1), + ] + ) + ) + + tokenized = art.tokenize_trajectory(trajectory) + + assert tokenized.token_ids == [1, 2, 3, 4] + assert tokenized.assistant_mask == [False, True, False, True] + assert math.isnan(tokenized.logprobs[0]) + assert tokenized.logprobs[1] == -0.2 + assert math.isnan(tokenized.logprobs[2]) + assert tokenized.logprobs[3] == -0.4 + + +def test_branching_and_multiple_models_require_explicit_resolution() -> None: + branching = art.Trajectory( + exchanges=TrajectoryExchanges( + chat_completions=[ + _chat_exchange([1], [2], offset=0), + _chat_exchange([9], [3], offset=1), + ] + ) + ) + with pytest.raises(ValueError, match="append-only"): + art.tokenize_trajectory(branching) + + mixed = art.Trajectory( + exchanges=TrajectoryExchanges( + chat_completions=[ + _chat_exchange([1], [2], model="one", offset=0), + _chat_exchange([3], [4], model="two", offset=1), + ] + ) + ) + with pytest.raises(ValueError, match="exactly one model"): + art.tokenize_trajectory(mixed) + assert art.tokenize_trajectory(mixed, model="two").token_ids == [3, 4] + + +class _FakeTokenizer: + def __init__(self) -> None: + self.calls: list[dict[str, Any]] = [] + + def apply_chat_template(self, messages: list[dict[str, Any]], **kwargs: Any): + self.calls.append(kwargs) + return [10, 11] if messages[-1]["role"] == "assistant" else [10] + + +def test_fallback_uses_template_overrides_and_nan_logprobs(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( + { + "model": "test/model", + "messages": [{"role": "user", "content": "question"}], + "chat_template": "request-template", + "chat_template_kwargs": {"request": True}, + "thinking": {"type": "enabled", "budget_tokens": 128}, + } + ), + response=response, + model="test/model", + start_time=start, + end_time=start + timedelta(seconds=1), + ) + tokenizer = _FakeTokenizer() + monkeypatch.setattr( + "art.trajectories._tokenize._load_tokenizer", lambda _config: tokenizer + ) + + result = art.tokenize_trajectory( + art.Trajectory(exchanges=TrajectoryExchanges(messages=[exchange])), + base_model="base/model", + chat_template="explicit-template", + chat_template_kwargs={"explicit": True}, + ) + + assert result.token_ids == [10, 11] + assert result.assistant_mask == [False, True] + 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, + "add_generation_prompt": False, + "chat_template": "explicit-template", + "request": True, + "explicit": True, + "enable_thinking": True, + "thinking_budget": 128, + }, + ] + + +def test_checkpoint_fallback_uses_latest_artifact_renderer(monkeypatch) -> None: + artifact_names: list[str] = [] + + class Api: + def artifact(self, name: str): + artifact_names.append(name) + return SimpleNamespace( + metadata={ + "wandb.base_model": "base/model", + "renderer": { + "tokenizer_revision": "revision", + "chat_template": "template", + "chat_template_kwargs": {"thinking": True}, + }, + } + ) + + monkeypatch.setattr("wandb.Api", Api) + from art.trajectories._tokenize import _tokenizer_config + + config = _tokenizer_config("wandb-artifact:///entity/project/run", None) + + assert artifact_names == ["entity/project/run:latest"] + assert config.base_model == "base/model" + assert config.revision == "revision" + assert config.chat_template == "template" + assert config.chat_template_kwargs == {"thinking": True} + + +def test_anthropic_fallback_preserves_thinking_and_tool_history() -> None: + from art.trajectories._tokenize import _anthropic_messages + + messages = _anthropic_messages( + { + "system": [{"type": "text", "text": "system"}], + "messages": [ + {"role": "user", "content": "question"}, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "reason"}, + {"type": "text", "text": "calling"}, + { + "type": "tool_use", + "id": "call-1", + "name": "lookup", + "input": {"key": "value"}, + }, + ], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "call-1", + "content": [{"type": "text", "text": "result"}], + }, + {"type": "text", "text": "continue"}, + ], + }, + ], + } + ) + + assert messages == [ + {"role": "system", "content": "system"}, + {"role": "user", "content": "question"}, + { + "role": "assistant", + "content": "calling", + "reasoning": "reason", + "tool_calls": [ + { + "id": "call-1", + "type": "function", + "function": { + "name": "lookup", + "arguments": '{"key": "value"}', + }, + } + ], + }, + {"role": "tool", "tool_call_id": "call-1", "content": "result"}, + {"role": "user", "content": "continue"}, + ] + + +def test_choice_logprobs_survive_tokenizer_fallback(monkeypatch) -> None: + exchange = _chat_exchange([], []) + 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.7, + bytes=list(b"answer"), + top_logprobs=[], + ) + ] + } + ) + + class Tokenizer: + def apply_chat_template(self, messages, **kwargs): + del kwargs + return [10, 11, 12] if messages[-1]["role"] == "assistant" else [10] + + def __call__(self, text, **kwargs): + del text, kwargs + return type("Encoding", (), {"input_ids": [11]})() + + monkeypatch.setattr( + "art.trajectories._tokenize._load_tokenizer", lambda _config: Tokenizer() + ) + result = art.tokenize_trajectory( + art.Trajectory(exchanges=TrajectoryExchanges(chat_completions=[exchange])), + base_model="base/model", + ) + assert result.token_ids == [10, 11, 12] + assert result.logprobs[1] == -0.7 + assert math.isnan(result.logprobs[2]) + + +def test_json_round_trip_preserves_exchange_types() -> None: + original = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[_chat_exchange([1], [2])]) + ) + restored = art.Trajectory.model_validate_json(original.model_dump_json()) + assert restored.model_dump(mode="json") == original.model_dump(mode="json") + assert isinstance(restored.exchanges.chat_completions[0].response, ChatCompletion) + + +def _response_exchange( + response_id: str, + output_id: int, + *, + previous_response_id: str | None = None, + offset: int = 0, +) -> ResponsesExchange: + response = Response.model_validate( + { + "id": response_id, + "created_at": float(offset), + "model": "test/model", + "object": "response", + "output": [ + { + "id": f"message-{response_id}", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [ + { + "type": "output_text", + "text": "answer", + "annotations": [], + "logprobs": [], + } + ], + } + ], + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [], + "raw_output_tokens": [{"token_id": output_id, "logprob": -0.1}], + } + ) + request = {"model": "test/model", "input": f"turn {offset}"} + if previous_response_id is not None: + request["previous_response_id"] = previous_response_id + start = datetime(2026, 1, 1) + timedelta(seconds=offset) + return ResponsesExchange( + request=ResponsesRequest(request), + response=response, + model="test/model", + start_time=start, + end_time=start + timedelta(milliseconds=1), + ) + + +def test_responses_previous_response_id_resolves_local_history(monkeypatch) -> None: + class Tokenizer: + def apply_chat_template(self, messages, **kwargs): + del kwargs + return [10] if len(messages) == 1 else [10, 20, 11] + + 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) + trajectory = art.Trajectory( + exchanges=TrajectoryExchanges(responses=[first, second]) + ) + + assert art.tokenize_trajectory(trajectory, base_model="base/model").token_ids == [ + 10, + 20, + 11, + 30, + ] + + second.request.root["previous_response_id"] = "missing" + with pytest.raises(ValueError, match="outside this trajectory"): + art.tokenize_trajectory(trajectory, base_model="base/model") + + +def test_exchange_trajectories_feed_existing_training_tokenizer() -> None: + from art.preprocessing.tokenize import tokenize_trajectory_groups + + class Tokenizer: + name_or_path = "test/model" + + def decode(self, token_id: int) -> str: + return str(token_id) + + group = art.TrajectoryGroup( + [ + art.Trajectory( + exchanges=TrajectoryExchanges( + chat_completions=[_chat_exchange([1], [2])] + ), + reward=1, + ), + art.Trajectory( + exchanges=TrajectoryExchanges( + chat_completions=[_chat_exchange([1], [3])] + ), + reward=0, + ), + ] + ) + + results = list( + tokenize_trajectory_groups( + Tokenizer(), # type: ignore[arg-type] + [group], + allow_training_without_logprobs=True, + scale_rewards=False, + shuffle_group_trajectories=False, + ) + ) + + assert [result.token_ids for result in results] == [[1, 2], [1, 3]] + assert [result.assistant_mask for result in results] == [[0, 1], [0, 1]] diff --git a/uv.lock b/uv.lock index bd7c959f7..21b5bd1a4 100644 --- a/uv.lock +++ b/uv.lock @@ -311,6 +311,25 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/78/b6/6307fbef88d9b5ee7421e68d78a9f162e0da4900bc5f5793f6d3d0e34fb8/annotated_types-0.7.0-py3-none-any.whl", hash = "sha256:1f02e8b43a8fbbc3f3e0d4f0f4bfc8131bcb4eebe8849b8e5c773f3a1c582a53", size = 13643, upload-time = "2024-05-20T21:33:24.1Z" }, ] +[[package]] +name = "anthropic" +version = "0.116.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "anyio" }, + { name = "distro" }, + { name = "docstring-parser" }, + { name = "httpx" }, + { name = "jiter" }, + { name = "pydantic" }, + { name = "sniffio" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/66/a2/d31f14e28d49bae983a3634e38dfb4b31c50110b5e403596c5c6a20b23f8/anthropic-0.116.0.tar.gz", hash = "sha256:5fc248fbb9fe03ef686f8a774f81586bca31a043260aab88b387ea3660f4a396", size = 949149, upload-time = "2026-07-02T19:08:10.534Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c7/dd/2a1e81cf1b163acc340afc4ec74ed1d86f5eed1a809fabdeed3e0997b346/anthropic-0.116.0-py3-none-any.whl", hash = "sha256:6c0a7698e8d652455da3499978279bb2588c7264d0a35be3666009a4258c8256", size = 956896, upload-time = "2026-07-02T19:08:08.756Z" }, +] + [[package]] name = "antlr4-python3-runtime" version = "4.9.3" @@ -4740,10 +4759,13 @@ name = "openpipe-art" version = "0.5.18" source = { editable = "." } dependencies = [ + { name = "aiohttp" }, + { name = "anthropic" }, { name = "litellm" }, { name = "nest-asyncio" }, { name = "openai" }, { name = "polars" }, + { name = "requests" }, { name = "setproctitle" }, { name = "tblib" }, { name = "typer" }, @@ -4845,6 +4867,8 @@ dev = [ [package.metadata] requires-dist = [ { name = "accelerate", marker = "extra == 'backend'", specifier = "==1.7.0" }, + { name = "aiohttp", specifier = ">=3.10.0" }, + { name = "anthropic", specifier = ">=0.75.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" }, @@ -4888,6 +4912,7 @@ requires-dist = [ { name = "pydantic", marker = "extra == 'tinker'", specifier = ">=2.12.5" }, { name = "pytest", marker = "extra == 'backend'", specifier = ">=8.4.1" }, { name = "quack-kernels", marker = "extra == 'megatron'", specifier = "==0.3.7" }, + { name = "requests", specifier = ">=2.32.0" }, { name = "seaborn", marker = "extra == 'plotting'", specifier = ">=0.13.2" }, { name = "setproctitle", specifier = ">=1.3.6" }, { name = "setuptools", marker = "extra == 'backend'", specifier = ">=78.1.0" }, From f10e57d06f8749572d1d1e73e6606b8903d0fc89 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Tue, 14 Jul 2026 23:42:46 +0000 Subject: [PATCH 02/17] refactor: tighten trajectory typing --- src/art/preprocessing/tokenize.py | 6 +- src/art/trajectories/__init__.py | 171 ++++++++++++----- src/art/trajectories/_capture/aiohttp.py | 34 ++-- src/art/trajectories/_capture/core.py | 40 +++- src/art/trajectories/_capture/httpx.py | 34 ++-- src/art/trajectories/_capture/requests.py | 16 +- src/art/trajectories/_compat.py | 99 +++++++--- src/art/trajectories/_protocols.py | 26 +-- src/art/trajectories/_scope.py | 23 ++- src/art/trajectories/_serialization.py | 116 +++++++++++ src/art/trajectories/_tokenize.py | 223 +++++++++++++++------- tests/unit/test_auto_trajectory.py | 4 +- tests/unit/trajectories/test_capture.py | 26 ++- tests/unit/trajectories/test_tokenize.py | 37 ++-- 14 files changed, 631 insertions(+), 224 deletions(-) create mode 100644 src/art/trajectories/_serialization.py diff --git a/src/art/preprocessing/tokenize.py b/src/art/preprocessing/tokenize.py index 5a0daf609..6785e2ffc 100644 --- a/src/art/preprocessing/tokenize.py +++ b/src/art/preprocessing/tokenize.py @@ -526,15 +526,15 @@ def tokenize_trajectory_groups( if advantage == 0 and drop_zero_advantage_trajectories: continue if trajectory.exchanges: - from ..trajectories._tokenize import tokenize_one + from ..trajectories._tokenize import _as_tokenizer, tokenize_one exchange_result = tokenize_one( trajectory, - getattr(tokenizer, "name_or_path", None), + tokenizer.name_or_path, model=None, chat_template=None, chat_template_kwargs=chat_template_kwargs, - tokenizer_instance=tokenizer, + tokenizer_instance=_as_tokenizer(tokenizer), ) choice_offsets = [ index diff --git a/src/art/trajectories/__init__.py b/src/art/trajectories/__init__.py index b68c1fdb5..8fcbabcc7 100644 --- a/src/art/trajectories/__init__.py +++ b/src/art/trajectories/__init__.py @@ -1,10 +1,18 @@ from __future__ import annotations -from collections.abc import AsyncGenerator, Awaitable, Coroutine, Iterable, Iterator +from collections.abc import ( + AsyncGenerator, + Awaitable, + Coroutine, + Iterable, + Iterator, + Mapping, +) from contextlib import asynccontextmanager from datetime import datetime import time -from typing import Any, Literal, cast, overload +from types import TracebackType +from typing import Any, Literal, overload from anthropic.types import Message as AnthropicMessage from openai.types import Completion @@ -14,23 +22,25 @@ from typing_extensions import deprecated from ..types import Messages, MessagesAndChoices, Tools +from ._serialization import _CompactModel +# Deliberately open: Pydantic enforces serializability when callers dump in JSON mode. MetadataValue = Any -class ChatCompletionsRequest(pydantic.RootModel[dict[str, Any]]): +class ChatCompletionsRequest(pydantic.RootModel[dict[str, pydantic.JsonValue]]): """The JSON body sent to an OpenAI-compatible Chat Completions endpoint.""" -class CompletionsRequest(pydantic.RootModel[dict[str, Any]]): +class CompletionsRequest(pydantic.RootModel[dict[str, pydantic.JsonValue]]): """The JSON body sent to an OpenAI-compatible Completions endpoint.""" -class ResponsesRequest(pydantic.RootModel[dict[str, Any]]): +class ResponsesRequest(pydantic.RootModel[dict[str, pydantic.JsonValue]]): """The JSON body sent to an OpenAI-compatible Responses endpoint.""" -class MessagesRequest(pydantic.RootModel[dict[str, Any]]): +class MessagesRequest(pydantic.RootModel[dict[str, pydantic.JsonValue]]): """The JSON body sent to an Anthropic-compatible Messages endpoint.""" @@ -94,7 +104,7 @@ def messages(self) -> Messages: return get_messages(self.messages_and_choices) -class Trajectory(pydantic.BaseModel): +class Trajectory(_CompactModel): exchanges: TrajectoryExchanges = pydantic.Field(default_factory=TrajectoryExchanges) messages_and_choices: MessagesAndChoices = pydantic.Field( default_factory=list, @@ -107,7 +117,7 @@ class Trajectory(pydantic.BaseModel): initial_policy_version: int | None = None final_policy_version: int | None = None metrics: dict[str, float | int | bool] = pydantic.Field(default_factory=dict) - metadata: dict[str, Any] = pydantic.Field(default_factory=dict) + metadata: dict[str, MetadataValue] = pydantic.Field(default_factory=dict) logs: list[str] = pydantic.Field(default_factory=list) start_time: datetime = pydantic.Field(default_factory=datetime.now, exclude=True) @@ -128,10 +138,15 @@ def __enter__(self) -> Trajectory: return enter_trajectory(self) - def __exit__(self, *exc_info: Any) -> None: + def __exit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + ) -> None: from ._scope import exit_trajectory - exit_trajectory(self, *exc_info) + exit_trajectory(self, exc_type, exc_value, traceback) def log(self, message: str) -> None: self.logs.append(message) @@ -153,27 +168,19 @@ 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})" - def model_dump(self, **kwargs: Any) -> dict[str, Any]: - kwargs.setdefault("exclude_defaults", True) - return super().model_dump(**kwargs) - - def model_dump_json(self, **kwargs: Any) -> str: - kwargs.setdefault("exclude_defaults", True) - return super().model_dump_json(**kwargs) - def messages(self) -> Messages: return get_messages(self.messages_and_choices) - def for_logging(self) -> dict[str, Any]: + def for_logging(self) -> dict[str, object]: from ._compat import trajectory_for_logging return trajectory_for_logging(self) -class TrajectoryGroup(pydantic.BaseModel): +class TrajectoryGroup(_CompactModel): trajectories: list[Trajectory] = pydantic.Field(default_factory=list) exceptions: list[PydanticException] = pydantic.Field(default_factory=list) - metadata: dict[str, Any] = pydantic.Field(default_factory=dict) + metadata: dict[str, MetadataValue] = pydantic.Field(default_factory=dict) metrics: dict[str, float | int | bool] = pydantic.Field(default_factory=dict) logs: list[str] = pydantic.Field(default_factory=list) @@ -181,7 +188,11 @@ class TrajectoryGroup(pydantic.BaseModel): def __new__( cls, trajectories: Iterable[Trajectory | BaseException] = (), - **kwargs: Any, + *, + exceptions: Iterable[BaseException | PydanticException] = (), + metadata: dict[str, MetadataValue] | None = None, + metrics: dict[str, float | int | bool] | None = None, + logs: list[str] | None = None, ) -> TrajectoryGroup: ... @overload @@ -189,13 +200,32 @@ def __new__( def __new__( cls, trajectories: Iterable[Awaitable[Trajectory]], - **kwargs: Any, + *, + exceptions: Iterable[BaseException | PydanticException] = (), + metadata: dict[str, MetadataValue] | None = None, + metrics: dict[str, float | int | bool] | None = None, + logs: list[str] | None = None, ) -> Awaitable[TrajectoryGroup]: ... - def __new__(cls, trajectories: Iterable[Any] = (), **kwargs: Any) -> Any: + def __new__( + cls, + trajectories: Iterable[Trajectory | BaseException | Awaitable[Trajectory]] = (), + *, + exceptions: Iterable[BaseException | PydanticException] = (), + metadata: dict[str, MetadataValue] | None = None, + metrics: dict[str, float | int | bool] | None = None, + logs: list[str] | None = None, + ) -> TrajectoryGroup | Awaitable[TrajectoryGroup]: from ._compat import new_trajectory_group - return new_trajectory_group(cls, trajectories, kwargs) + return new_trajectory_group( + cls, + trajectories, + exceptions=exceptions, + metadata=metadata, + metrics=metrics, + logs=logs, + ) def __init__( self, @@ -204,15 +234,21 @@ def __init__( ) = (), *, exceptions: Iterable[BaseException | PydanticException] = (), - metadata: dict[str, Any] | None = None, + metadata: dict[str, MetadataValue] | None = None, metrics: dict[str, float | int | bool] | None = None, logs: list[str] | None = None, ) -> None: from ._compat import init_trajectory_group + items = list(trajectories) + sync_items = [ + item for item in items if isinstance(item, (Trajectory, BaseException)) + ] + if len(sync_items) != len(items): + raise TypeError("TrajectoryGroup cannot initialize from awaitables") init_trajectory_group( self, - cast(Iterable[Trajectory | BaseException], trajectories), + sync_items, exceptions=exceptions, metadata=metadata, metrics=metrics, @@ -224,17 +260,22 @@ def __enter__(self) -> TrajectoryGroup: return enter_trajectory_group(self) - def __exit__(self, *exc_info: Any) -> None: + def __exit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + ) -> None: from ._scope import exit_trajectory_group - exit_trajectory_group(self, *exc_info) + exit_trajectory_group(self, exc_type, exc_value, traceback) def __copy__(self) -> TrajectoryGroup: from ._compat import copy_trajectory_group return copy_trajectory_group(self) - def __deepcopy__(self, memo: dict[int, Any] | None = None) -> TrajectoryGroup: + def __deepcopy__(self, memo: dict[int, object] | None = None) -> TrajectoryGroup: from ._compat import deepcopy_trajectory_group return deepcopy_trajectory_group(self, memo) @@ -242,20 +283,13 @@ def __deepcopy__(self, memo: dict[int, Any] | None = None) -> TrajectoryGroup: def log(self, message: str) -> None: self.logs.append(message) + # Legacy groups iterate over trajectories rather than Pydantic field pairs. def __iter__(self) -> Iterator[Trajectory]: # type: ignore[override] return iter(self.trajectories) def __len__(self) -> int: return len(self.trajectories) - def model_dump(self, **kwargs: Any) -> dict[str, Any]: - kwargs.setdefault("exclude_defaults", True) - return super().model_dump(**kwargs) - - def model_dump_json(self, **kwargs: Any) -> str: - kwargs.setdefault("exclude_defaults", True) - return super().model_dump_json(**kwargs) - class TokenizedTrajectory(pydantic.BaseModel): token_ids: list[int] @@ -299,7 +333,7 @@ def current_trajectory_group(*, required: bool = False) -> TrajectoryGroup | Non return get_current_trajectory_group(required=required) -async def trajectory(coroutine: Coroutine[Any, Any, Any]) -> Trajectory: +async def trajectory(coroutine: Coroutine[Any, Any, object]) -> Trajectory: from ._scope import capture_trajectory return await capture_trajectory(coroutine) @@ -324,7 +358,7 @@ def tokenize_trajectory( *, model: str | None = None, chat_template: str | None = None, - chat_template_kwargs: dict[str, Any] | None = None, + chat_template_kwargs: Mapping[str, object] | None = None, ) -> TokenizedTrajectory: from ._tokenize import tokenize_one @@ -340,18 +374,39 @@ def tokenize_trajectory( def tokenize_trajectories( trajectories: Iterable[Trajectory], base_model: str | None = None, - **kwargs: Any, + *, + 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, **kwargs) for item in trajectories] + return [ + tokenize_trajectory( + item, + 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, - **kwargs: Any, + *, + 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, **kwargs), + trajectories=tokenize_trajectories( + group, + base_model, + model=model, + chat_template=chat_template, + chat_template_kwargs=chat_template_kwargs, + ), underlying=group, ) @@ -359,9 +414,31 @@ def tokenize_trajectory_group( def tokenize_trajectory_groups( groups: Iterable[TrajectoryGroup], base_model: str | None = None, - **kwargs: Any, + *, + 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, **kwargs) for group in groups] + return [ + tokenize_trajectory_group( + group, + 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: ... + + +@overload +@deprecated("Use current_trajectory() instead.") +def auto_trajectory(*, required: Literal[False] = False) -> Trajectory | None: ... @deprecated("Use current_trajectory() instead.") @@ -371,7 +448,7 @@ def auto_trajectory(*, required: bool = False) -> Trajectory | None: @deprecated("Use trajectory() instead.") async def capture_auto_trajectory( - coroutine: Coroutine[Any, Any, Any], + coroutine: Coroutine[Any, Any, object], ) -> Trajectory: return await trajectory(coroutine) diff --git a/src/art/trajectories/_capture/aiohttp.py b/src/art/trajectories/_capture/aiohttp.py index 660a0678b..87192b612 100644 --- a/src/art/trajectories/_capture/aiohttp.py +++ b/src/art/trajectories/_capture/aiohttp.py @@ -1,22 +1,29 @@ from __future__ import annotations from collections.abc import AsyncIterator -from typing import Any +from typing import Any, cast, overload import aiohttp +from yarl import URL from .core import CaptureState, begin, reset class _CapturedStream: - def __init__(self, stream: Any, state: CaptureState) -> None: + def __init__(self, stream: aiohttp.StreamReader, state: CaptureState) -> None: self._stream = stream self._state = state def __getattr__(self, name: str) -> Any: return getattr(self._stream, name) - def _record(self, value: Any) -> Any: + @overload + def _record(self, value: bytes) -> bytes: ... + + @overload + def _record(self, value: tuple[bytes, bool]) -> tuple[bytes, bool]: ... + + def _record(self, value: bytes | tuple[bytes, bool]) -> bytes | tuple[bytes, bool]: chunk = value[0] if isinstance(value, tuple) else value if isinstance(chunk, bytes): self._state.add(chunk) @@ -24,8 +31,8 @@ def _record(self, value: Any) -> Any: self._state.finish() return value - async def read(self, *args: Any, **kwargs: Any) -> bytes: - return self._record(await self._stream.read(*args, **kwargs)) + async def read(self, n: int = -1) -> bytes: + return self._record(await self._stream.read(n)) async def readany(self) -> bytes: return self._record(await self._stream.readany()) @@ -36,7 +43,7 @@ async def readline(self) -> bytes: async def readchunk(self) -> tuple[bytes, bool]: return self._record(await self._stream.readchunk()) - async def _iterate(self, iterator: Any) -> AsyncIterator[bytes]: + async def _iterate(self, iterator: AsyncIterator[bytes]) -> AsyncIterator[bytes]: try: async for chunk in iterator: yield self._record(chunk) @@ -61,10 +68,11 @@ def install() -> None: async def request( self: aiohttp.ClientSession, method: str, - str_or_url: Any, + str_or_url: str | URL, + # This private aiohttp surface is version-dependent; preserve its options. **kwargs: Any, ) -> aiohttp.ClientResponse: - body: Any = kwargs.get("json") + body: object = kwargs.get("json") if body is None: body = kwargs.get("data") state, token = begin(method, str(str_or_url), body) @@ -74,8 +82,12 @@ async def request( reset(token) if state is not None: state.status_code = response.status - response.content = _CapturedStream(response.content, state) # type: ignore[assignment] + # The proxy preserves StreamReader's runtime surface while intercepting + # reads; aiohttp exposes no protocol type for response.content. + response.content = cast( + aiohttp.StreamReader, _CapturedStream(response.content, state) + ) return response - request._art_capture = True # type: ignore[attr-defined] - aiohttp.ClientSession._request = request # type: ignore[method-assign] + setattr(request, "_art_capture", True) + setattr(aiohttp.ClientSession, "_request", request) diff --git a/src/art/trajectories/_capture/core.py b/src/art/trajectories/_capture/core.py index edeb5e5b7..840b9695e 100644 --- a/src/art/trajectories/_capture/core.py +++ b/src/art/trajectories/_capture/core.py @@ -5,9 +5,16 @@ from datetime import datetime import json import logging -from typing import Any +from typing import Any, assert_never -from .._protocols import build_exchange, endpoint_for_url +from .. import ( + ChatCompletionsExchange, + CompletionsExchange, + MessagesExchange, + ResponsesExchange, + Trajectory, +) +from .._protocols import Endpoint, Exchange, build_exchange, endpoint_for_url from .._scope import get_current_trajectory logger = logging.getLogger(__name__) @@ -18,8 +25,8 @@ @dataclass class CaptureState: - trajectory: Any - endpoint: str + trajectory: Trajectory + endpoint: Endpoint request: dict[str, Any] start_time: datetime = field(default_factory=datetime.now) status_code: int | None = None @@ -37,7 +44,7 @@ def finish(self) -> None: if self.status_code is None or not 200 <= self.status_code < 300: return try: - name, exchange = build_exchange( + _, exchange = build_exchange( self.endpoint, self.request, bytes(self.body), @@ -47,12 +54,27 @@ def finish(self) -> None: except Exception as exc: logger.debug("Ignoring incomplete trajectory exchange: %s", exc) return - getattr(self.trajectory.exchanges, name).append(exchange) + _append_exchange(self.trajectory, exchange) + + +def _append_exchange(trajectory: Trajectory, exchange: Exchange) -> None: + if isinstance(exchange, ChatCompletionsExchange): + trajectory.exchanges.chat_completions.append(exchange) + elif isinstance(exchange, CompletionsExchange): + trajectory.exchanges.completions.append(exchange) + elif isinstance(exchange, ResponsesExchange): + trajectory.exchanges.responses.append(exchange) + elif isinstance(exchange, MessagesExchange): + trajectory.exchanges.messages.append(exchange) + else: + assert_never(exchange) -def _json_body(value: Any) -> dict[str, Any] | None: +def _json_body(value: object) -> dict[str, Any] | None: if isinstance(value, dict): - return value + if not all(isinstance(key, str) for key in value): + return None + return {key: item for key, item in value.items() if isinstance(key, str)} if isinstance(value, str): value = value.encode() if not isinstance(value, bytes): @@ -67,7 +89,7 @@ def _json_body(value: Any) -> dict[str, Any] | None: def begin( method: str, url: str, - body: Any, + body: object, ) -> tuple[CaptureState | None, contextvars.Token[bool] | None]: trajectory = get_current_trajectory(required=False) endpoint = endpoint_for_url(url) diff --git a/src/art/trajectories/_capture/httpx.py b/src/art/trajectories/_capture/httpx.py index 58d6dc3a9..aac980b23 100644 --- a/src/art/trajectories/_capture/httpx.py +++ b/src/art/trajectories/_capture/httpx.py @@ -1,15 +1,23 @@ from __future__ import annotations from collections.abc import AsyncIterator, Iterator -from typing import Any import httpx +from httpx._client import UseClientDefault +from httpx._types import AuthTypes +from typing_extensions import TypedDict, Unpack from .core import CaptureState, begin, reset _STATE = "_art_trajectory_capture" +class _SendOptions(TypedDict, total=False): + stream: bool + auth: AuthTypes | UseClientDefault | None + follow_redirects: bool | UseClientDefault + + def install() -> None: if getattr(httpx.Client.send, "_art_capture", False): return @@ -21,7 +29,9 @@ def install() -> None: original_aclose = httpx.Response.aclose def send( - self: httpx.Client, request: httpx.Request, **kwargs: Any + self: httpx.Client, + request: httpx.Request, + **kwargs: Unpack[_SendOptions], ) -> httpx.Response: try: body = request.content @@ -41,7 +51,9 @@ def send( return response async def async_send( - self: httpx.AsyncClient, request: httpx.Request, **kwargs: Any + self: httpx.AsyncClient, + request: httpx.Request, + **kwargs: Unpack[_SendOptions], ) -> httpx.Response: try: body = request.content @@ -96,11 +108,11 @@ async def aclose(self: httpx.Response) -> None: if state := getattr(self, _STATE, None): state.finish() - send._art_capture = True # type: ignore[attr-defined] - async_send._art_capture = True # type: ignore[attr-defined] - httpx.Client.send = send # type: ignore[method-assign] - httpx.AsyncClient.send = async_send # type: ignore[method-assign] - httpx.Response.iter_bytes = iter_bytes # type: ignore[method-assign] - httpx.Response.aiter_bytes = aiter_bytes # type: ignore[method-assign] - httpx.Response.close = close # type: ignore[method-assign] - httpx.Response.aclose = aclose # type: ignore[method-assign] + 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, "close", close) + setattr(httpx.Response, "aclose", aclose) diff --git a/src/art/trajectories/_capture/requests.py b/src/art/trajectories/_capture/requests.py index fb62536ca..da1b0009e 100644 --- a/src/art/trajectories/_capture/requests.py +++ b/src/art/trajectories/_capture/requests.py @@ -33,11 +33,15 @@ def send( return response def iter_content( - self: requests.Response, *args: Any, **kwargs: Any - ) -> Iterator[Any]: + self: requests.Response, + chunk_size: int | None = 1, + decode_unicode: bool = False, + ) -> Iterator[str | bytes]: state: CaptureState | None = getattr(self, _STATE, None) try: - for chunk in original_iter(self, *args, **kwargs): + for chunk in original_iter( + self, chunk_size=chunk_size, decode_unicode=decode_unicode + ): if state is not None: if isinstance(chunk, str): chunk = chunk.encode(self.encoding or "utf-8") @@ -48,6 +52,6 @@ def iter_content( if state is not None: state.finish() - send._art_capture = True # type: ignore[attr-defined] - requests.Session.send = send # type: ignore[method-assign] - requests.Response.iter_content = iter_content # type: ignore[method-assign] + setattr(send, "_art_capture", True) + setattr(requests.Session, "send", send) + setattr(requests.Response, "iter_content", iter_content) diff --git a/src/art/trajectories/_compat.py b/src/art/trajectories/_compat.py index d6d08084c..98c2eb88f 100644 --- a/src/art/trajectories/_compat.py +++ b/src/art/trajectories/_compat.py @@ -1,26 +1,24 @@ from __future__ import annotations import asyncio -from collections.abc import Iterable +from collections.abc import Awaitable, Coroutine, Generator, Iterable import copy import traceback -from typing import Any, cast +from typing import Any import warnings from openai.types.chat.chat_completion import Choice import pydantic from ..types import Message, Messages, MessagesAndChoices -from . import PydanticException, Trajectory, TrajectoryGroup +from . import MetadataValue, PydanticException, Trajectory, TrajectoryGroup def exception_model( - exception: BaseException | PydanticException | dict[str, Any], + exception: BaseException | PydanticException, ) -> PydanticException: if isinstance(exception, PydanticException): return exception - if isinstance(exception, dict): - return PydanticException.model_validate(exception) return PydanticException( type=str(type(exception)), message=str(exception), @@ -33,13 +31,18 @@ def exception_model( async def _legacy_async_group( - items: list[Any], kwargs: dict[str, Any] + items: list[Awaitable[Trajectory]], + *, + exceptions: Iterable[BaseException | PydanticException], + metadata: dict[str, MetadataValue] | None, + metrics: dict[str, float | int | bool] | None, + logs: list[str] | None, ) -> TrajectoryGroup: from ..gather import get_gather_context, record_metrics context = get_gather_context() trajectories: list[Trajectory] = [] - exceptions = list(kwargs.pop("exceptions", ())) + captured_exceptions = list(exceptions) for future in asyncio.as_completed(items): try: item = await future @@ -47,38 +50,73 @@ async def _legacy_async_group( record_metrics(context, item) context.update_pbar(n=1) except BaseException as exc: - exceptions.append(exc) + captured_exceptions.append(exc) context.metric_sums["exceptions"] += 1 context.update_pbar(n=0) if context.too_many_exceptions(): raise - return TrajectoryGroup(trajectories, exceptions=exceptions, **kwargs) + return TrajectoryGroup( + trajectories, + exceptions=captured_exceptions, + metadata=metadata, + metrics=metrics, + logs=logs, + ) -class _LegacyGroupCoroutine: - def __init__(self, coroutine: Any, size: int) -> None: +class _LegacyGroupCoroutine(Awaitable[TrajectoryGroup]): + def __init__( + self, coroutine: Coroutine[Any, Any, TrajectoryGroup], size: int + ) -> None: self.coroutine = coroutine self._num_trajectories = size - def __await__(self) -> Any: + def __await__(self) -> Generator[Any, None, TrajectoryGroup]: return self.coroutine.__await__() def new_trajectory_group( - cls: type[TrajectoryGroup], trajectories: Iterable[Any], kwargs: dict[str, Any] -) -> Any: + cls: type[TrajectoryGroup], + trajectories: Iterable[Trajectory | BaseException | Awaitable[Trajectory]], + *, + exceptions: Iterable[BaseException | PydanticException], + metadata: dict[str, MetadataValue] | None, + metrics: dict[str, float | int | bool] | None, + logs: list[str] | None, +) -> TrajectoryGroup | Awaitable[TrajectoryGroup]: items = list(trajectories) - if any(hasattr(item, "__await__") for item in items): + awaitables = [item for item in items if isinstance(item, Awaitable)] + if awaitables: + if len(awaitables) != len(items): + raise TypeError("TrajectoryGroup cannot mix trajectories and awaitables") warnings.warn( "Awaiting TrajectoryGroup(...) is deprecated; use art.trajectory_group(...).", DeprecationWarning, stacklevel=2, ) return _LegacyGroupCoroutine( - _legacy_async_group(items, dict(kwargs)), len(items) + _legacy_async_group( + awaitables, + exceptions=exceptions, + metadata=metadata, + metrics=metrics, + logs=logs, + ), + len(items), ) + sync_items = [ + item for item in items if isinstance(item, (Trajectory, BaseException)) + ] + if len(sync_items) != len(items): + raise TypeError("TrajectoryGroup items must be trajectories or exceptions") group = object.__new__(cls) - group.__init__(items, **kwargs) + group.__init__( + sync_items, + exceptions=exceptions, + metadata=metadata, + metrics=metrics, + logs=logs, + ) return group @@ -87,7 +125,7 @@ def init_trajectory_group( trajectories: Iterable[Trajectory | BaseException], *, exceptions: Iterable[BaseException | PydanticException], - metadata: dict[str, Any] | None, + metadata: dict[str, MetadataValue] | None, metrics: dict[str, float | int | bool] | None, logs: list[str] | None, ) -> None: @@ -126,11 +164,13 @@ def copy_trajectory_group(group: TrajectoryGroup) -> TrajectoryGroup: def deepcopy_trajectory_group( - group: TrajectoryGroup, memo: dict[int, Any] | None + group: TrajectoryGroup, memo: dict[int, object] | None ) -> TrajectoryGroup: memo = {} if memo is None else memo - if id(group) in memo: - return memo[id(group)] + if existing := memo.get(id(group)): + if not isinstance(existing, TrajectoryGroup): + raise TypeError("TrajectoryGroup deepcopy memo contains an invalid value") + return existing copied = TrajectoryGroup( copy.deepcopy(group.trajectories, memo), metadata=copy.deepcopy(group.metadata, memo), @@ -148,8 +188,9 @@ def messages_from_legacy_history(messages_and_choices: MessagesAndChoices) -> Me if isinstance(item, Choice): content = item.message.content or "" tool_calls = item.message.tool_calls or [] - message: Message = cast( - Message, + # Response messages and request messages are parallel OpenAI models, + # but their generated Python types are unrelated. Validate at the seam. + message = pydantic.TypeAdapter(Message).validate_python( { "role": "assistant", "content": content, @@ -163,18 +204,18 @@ def messages_from_legacy_history(messages_and_choices: MessagesAndChoices) -> Me if tool_calls else {} ), - }, + } ) messages.append(message) else: - message = dict(item) + message = copy.copy(item) if message.get("content") is None: message["content"] = "" - messages.append(message) # type: ignore[arg-type] + messages.append(message) return messages -def trajectory_for_logging(trajectory: Trajectory) -> dict[str, Any]: +def trajectory_for_logging(trajectory: Trajectory) -> dict[str, object]: if trajectory.exchanges: return trajectory.model_dump(mode="json", exclude={"start_time"}) messages = [] @@ -183,7 +224,7 @@ def trajectory_for_logging(trajectory: Trajectory) -> dict[str, Any]: message = item.message.to_dict() trainable = True else: - message = cast(dict[str, Any], item) + message = dict(item) trainable = False messages.append({**message, "trainable": trainable}) return { diff --git a/src/art/trajectories/_protocols.py b/src/art/trajectories/_protocols.py index d6c02c16c..ede72e26b 100644 --- a/src/art/trajectories/_protocols.py +++ b/src/art/trajectories/_protocols.py @@ -3,12 +3,12 @@ from collections.abc import Mapping from datetime import datetime import json -from typing import Any, cast +from typing import Any, Literal from urllib.parse import urlsplit from anthropic._types import NOT_GIVEN from anthropic.lib.streaming._messages import accumulate_event -from anthropic.types import Message, RawMessageStreamEvent +from anthropic.types import Message, ParsedMessage, RawMessageStreamEvent from openai.types import Completion from openai.types.chat import ChatCompletion from openai.types.chat.chat_completion_chunk import ChatCompletionChunk @@ -27,7 +27,12 @@ ResponsesRequest, ) -Endpoint = str +Endpoint = Literal["chat_completions", "completions", "responses", "messages"] +Exchange = ( + ChatCompletionsExchange | CompletionsExchange | ResponsesExchange | MessagesExchange +) +ResponseModel = ChatCompletion | Completion | Response | Message +SSEPayload = dict[str, Any] | Literal["[DONE]"] _ENDPOINTS = { "/v1/chat/completions": "chat_completions", "/v1/completions": "completions", @@ -43,9 +48,9 @@ def endpoint_for_url(url: str) -> Endpoint | None: ) -def _sse_events(body: bytes) -> list[tuple[str | None, dict[str, Any] | str]]: +def _sse_events(body: bytes) -> list[tuple[str | None, SSEPayload]]: text = body.decode("utf-8").replace("\r\n", "\n") - events: list[tuple[str | None, dict[str, Any] | str]] = [] + events: list[tuple[str | None, SSEPayload]] = [] for block in text.split("\n\n"): event_name: str | None = None data_lines: list[str] = [] @@ -58,7 +63,7 @@ def _sse_events(body: bytes) -> list[tuple[str | None, dict[str, Any] | str]]: continue raw = "\n".join(data_lines) if raw == "[DONE]": - events.append((event_name, raw)) + events.append((event_name, "[DONE]")) else: value = json.loads(raw) if isinstance(value, dict): @@ -164,7 +169,7 @@ def _messages_response(body: bytes, *, stream: bool) -> Message: if not stream: return Message.model_validate_json(body) adapter = TypeAdapter(RawMessageStreamEvent) - snapshot: Any = None + snapshot: ParsedMessage[object] | None = None complete = False token_ids: list[int] = [] logprobs: list[float] = [] @@ -201,12 +206,11 @@ def _messages_response(body: bytes, *, stream: bool) -> Message: return Message.model_validate(data) -def _model(request: Mapping[str, Any], response: Any) -> str | None: +def _model(request: Mapping[str, object], response: ResponseModel) -> str | None: requested = request.get("model") if isinstance(requested, str): return requested - returned = getattr(response, "model", None) - return returned if isinstance(returned, str) else None + return response.model if isinstance(response.model, str) else None def build_exchange( @@ -216,7 +220,7 @@ def build_exchange( *, start_time: datetime, end_time: datetime, -) -> tuple[str, Any]: +) -> tuple[Endpoint, Exchange]: stream = request.get("stream") is True if endpoint == "chat_completions": response = _chat_response(body, stream=stream) diff --git a/src/art/trajectories/_scope.py b/src/art/trajectories/_scope.py index 1a56e1ac3..b9f5e275c 100644 --- a/src/art/trajectories/_scope.py +++ b/src/art/trajectories/_scope.py @@ -3,6 +3,7 @@ import asyncio from collections.abc import Coroutine, Iterable import contextvars +from types import TracebackType from typing import Any from . import PydanticException, Trajectory, TrajectoryGroup @@ -42,7 +43,12 @@ def enter_trajectory(trajectory: Trajectory) -> Trajectory: return trajectory -def exit_trajectory(trajectory: Trajectory, *exc_info: Any) -> None: +def exit_trajectory( + trajectory: Trajectory, + _exc_type: type[BaseException] | None, + _exc_value: BaseException | None, + _traceback: TracebackType | None, +) -> None: current = _trajectories.get() if not current or current[-1] is not trajectory: raise RuntimeError("Trajectory contexts must exit in stack order") @@ -58,21 +64,26 @@ def enter_trajectory_group(group: TrajectoryGroup) -> TrajectoryGroup: return group -def exit_trajectory_group(group: TrajectoryGroup, *exc_info: Any) -> None: +def exit_trajectory_group( + group: TrajectoryGroup, + _exc_type: type[BaseException] | None, + _exc_value: BaseException | None, + _traceback: TracebackType | None, +) -> None: current = _groups.get() if not current or current[-1] is not group: raise RuntimeError("TrajectoryGroup contexts must exit in stack order") _groups.set(current[:-1]) -def _require_raw_coroutine(value: Any) -> None: +def _require_raw_coroutine(value: object) -> None: if isinstance(value, (asyncio.Task, asyncio.Future)) or not isinstance( value, Coroutine ): raise TypeError("Expected a raw coroutine, not a Task, Future, or awaitable") -async def capture_trajectory(coroutine: Coroutine[Any, Any, Any]) -> Trajectory: +async def capture_trajectory(coroutine: Coroutine[Any, Any, object]) -> Trajectory: _require_raw_coroutine(coroutine) with Trajectory() as captured: await coroutine @@ -87,9 +98,9 @@ async def capture_trajectory_group( coroutines = list(trajectories) for coroutine in coroutines: _require_raw_coroutine(coroutine) - results = await asyncio.gather(*coroutines, return_exceptions=return_exceptions) if not return_exceptions: - return TrajectoryGroup(results) # type: ignore[arg-type] + return TrajectoryGroup(await asyncio.gather(*coroutines)) + results = await asyncio.gather(*coroutines, return_exceptions=True) completed: list[Trajectory] = [] exceptions: list[PydanticException] = [] for result in results: diff --git a/src/art/trajectories/_serialization.py b/src/art/trajectories/_serialization.py new file mode 100644 index 000000000..6da7b1846 --- /dev/null +++ b/src/art/trajectories/_serialization.py @@ -0,0 +1,116 @@ +from __future__ import annotations + +from collections.abc import Callable +from typing import Any, Literal + +from pydantic import BaseModel +from pydantic.main import IncEx + + +class _CompactModel(BaseModel): + """Pydantic model whose default dump omits fields equal to their defaults.""" + + def model_dump( + self, + *, + mode: Literal["json", "python"] | str = "python", + include: IncEx | None = None, + exclude: IncEx | None = None, + context: Any | None = None, + by_alias: bool | None = None, + exclude_unset: bool = False, + exclude_defaults: bool = True, + exclude_none: bool = False, + exclude_computed_fields: bool = False, + round_trip: bool = False, + warnings: bool | Literal["none", "warn", "error"] = True, + fallback: Callable[[Any], Any] | None = None, + serialize_as_any: bool = False, + polymorphic_serialization: bool | None = None, + ) -> dict[str, Any]: + if polymorphic_serialization is not None: + return super().model_dump( + mode=mode, + include=include, + exclude=exclude, + context=context, + by_alias=by_alias, + exclude_unset=exclude_unset, + exclude_defaults=exclude_defaults, + exclude_none=exclude_none, + exclude_computed_fields=exclude_computed_fields, + round_trip=round_trip, + warnings=warnings, + fallback=fallback, + serialize_as_any=serialize_as_any, + polymorphic_serialization=polymorphic_serialization, + ) + return super().model_dump( + mode=mode, + include=include, + exclude=exclude, + context=context, + by_alias=by_alias, + exclude_unset=exclude_unset, + exclude_defaults=exclude_defaults, + exclude_none=exclude_none, + exclude_computed_fields=exclude_computed_fields, + round_trip=round_trip, + warnings=warnings, + fallback=fallback, + serialize_as_any=serialize_as_any, + ) + + def model_dump_json( + self, + *, + indent: int | None = None, + ensure_ascii: bool = False, + include: IncEx | None = None, + exclude: IncEx | None = None, + context: Any | None = None, + by_alias: bool | None = None, + exclude_unset: bool = False, + exclude_defaults: bool = True, + exclude_none: bool = False, + exclude_computed_fields: bool = False, + round_trip: bool = False, + warnings: bool | Literal["none", "warn", "error"] = True, + fallback: Callable[[Any], Any] | None = None, + serialize_as_any: bool = False, + polymorphic_serialization: bool | None = None, + ) -> str: + if polymorphic_serialization is not None: + return super().model_dump_json( + indent=indent, + ensure_ascii=ensure_ascii, + include=include, + exclude=exclude, + context=context, + by_alias=by_alias, + exclude_unset=exclude_unset, + exclude_defaults=exclude_defaults, + exclude_none=exclude_none, + exclude_computed_fields=exclude_computed_fields, + round_trip=round_trip, + warnings=warnings, + fallback=fallback, + serialize_as_any=serialize_as_any, + polymorphic_serialization=polymorphic_serialization, + ) + return super().model_dump_json( + indent=indent, + ensure_ascii=ensure_ascii, + include=include, + exclude=exclude, + context=context, + by_alias=by_alias, + exclude_unset=exclude_unset, + exclude_defaults=exclude_defaults, + exclude_none=exclude_none, + exclude_computed_fields=exclude_computed_fields, + round_trip=round_trip, + warnings=warnings, + fallback=fallback, + serialize_as_any=serialize_as_any, + ) diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index ab9d1437f..e2ff2e99d 100644 --- a/src/art/trajectories/_tokenize.py +++ b/src/art/trajectories/_tokenize.py @@ -1,11 +1,18 @@ from __future__ import annotations +from collections.abc import Mapping from dataclasses import dataclass import math import re -from typing import Any, cast +from typing import Any, Protocol, cast +from anthropic.types import Message +from openai.types import Completion +from openai.types.chat import ChatCompletion from openai.types.chat.chat_completion import Choice +from openai.types.responses import Response +from pydantic import BaseModel +from transformers import PreTrainedTokenizerBase from . import ( ChatCompletionsExchange, @@ -15,6 +22,7 @@ TokenizedTrajectory, Trajectory, ) +from ._protocols import Exchange _TOKEN_ID = re.compile(r"token_id:(\d+)$") @@ -24,17 +32,57 @@ class _TokenizerConfig: base_model: str revision: str | None = None chat_template: str | None = None - chat_template_kwargs: dict[str, Any] | None = None + chat_template_kwargs: Mapping[str, object] | None = None -def _dump(value: Any) -> dict[str, Any]: - if hasattr(value, "model_dump"): +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: PreTrainedTokenizerBase) -> _Tokenizer: + # Transformers' annotation permits only string-valued message dictionaries, + # although its runtime API supports the structured content ART must tokenize. + return cast(_Tokenizer, 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): + return None + return {key: item for key, item in value.items() if isinstance(key, str)} + + +def _dict_list(value: object) -> list[dict[str, Any]]: + if value is None: + return [] + if not isinstance(value, list): + raise TypeError("Expected a list of JSON objects") + result: list[dict[str, Any]] = [] + for item in value: + if (mapping := _string_dict(item)) is None: + raise TypeError("Expected a list of JSON objects") + result.append(mapping) + return result + + +def _dump(value: object) -> dict[str, Any]: + if isinstance(value, BaseModel): result = value.model_dump(mode="python") return result if isinstance(result, dict) else {} - return value if isinstance(value, dict) else {} + return _string_dict(value) or {} -def _token_id(value: Any) -> int | None: +def _token_id(value: object) -> int | None: if isinstance(value, int) and not isinstance(value, bool): return value if isinstance(value, str) and (match := _TOKEN_ID.fullmatch(value)): @@ -42,7 +90,7 @@ def _token_id(value: Any) -> int | None: return None -def _pairs(values: Any) -> tuple[list[int], list[float]]: +def _pairs(values: object) -> tuple[list[int], list[float]]: if not isinstance(values, list): return [], [] token_ids: list[int] = [] @@ -62,7 +110,7 @@ def _pairs(values: Any) -> tuple[list[int], list[float]]: return token_ids, logprobs -def _logprob_values(values: Any) -> list[float]: +def _logprob_values(values: object) -> list[float]: if not isinstance(values, list): return [] result: list[float] = [] @@ -74,11 +122,9 @@ def _logprob_values(values: Any) -> list[float]: return result -def _chat_tokens(response: Any) -> tuple[list[int] | None, list[int], list[float]]: - if len(response.choices) != 1: - raise ValueError("Trajectory tokenization requires exactly one response choice") - choice = response.choices[0] - response_data = _dump(response) +def _chat_choice_tokens( + choice: Choice, response_data: dict[str, Any] +) -> tuple[list[int] | None, list[int], list[float]]: choice_data = _dump(choice) prompt = choice_data.get("prompt_token_ids") or response_data.get( "prompt_token_ids" @@ -93,9 +139,9 @@ def _chat_tokens(response: Any) -> tuple[list[int] | None, list[int], list[float for value in choice_data.get("token_ids") or [] if (token := _token_id(value)) is not None ] - logprob_values = getattr(getattr(choice, "logprobs", None), "content", None) - if logprob_values is None: - logprob_values = getattr(getattr(choice, "logprobs", None), "refusal", None) + logprob_values = None + if choice.logprobs is not None: + logprob_values = choice.logprobs.content or choice.logprobs.refusal values = list(logprob_values or []) pair_ids, logprobs = _pairs(values) if token_ids and pair_ids and token_ids != pair_ids: @@ -107,8 +153,16 @@ def _chat_tokens(response: Any) -> tuple[list[int] | None, list[int], list[float ) +def _chat_tokens( + response: ChatCompletion, +) -> tuple[list[int] | None, list[int], list[float]]: + if len(response.choices) != 1: + raise ValueError("Trajectory tokenization requires exactly one response choice") + return _chat_choice_tokens(response.choices[0], _dump(response)) + + def _completion_tokens( - response: Any, + response: Completion, ) -> tuple[list[int] | None, list[int], list[float]]: if len(response.choices) != 1: raise ValueError("Trajectory tokenization requires exactly one response choice") @@ -118,13 +172,17 @@ def _completion_tokens( prompt = choice_data.get("prompt_token_ids") or response_data.get( "prompt_token_ids" ) - prompt_ids = list(prompt) if isinstance(prompt, list) else None + prompt_ids = ( + [token for value in prompt if (token := _token_id(value)) is not None] + if isinstance(prompt, list) + else None + ) token_ids = [ token for value in choice_data.get("token_ids") or [] if (token := _token_id(value)) is not None ] - logprobs = _dump(getattr(choice, "logprobs", None)) + logprobs = _dump(choice.logprobs) tokens = logprobs.get("tokens") or [] pair_ids = [token for value in tokens if (token := _token_id(value)) is not None] pair_logprobs = [ @@ -139,7 +197,7 @@ def _completion_tokens( return prompt_ids, selected, pair_logprobs -def _responses_tokens(response: Any) -> tuple[None, list[int], list[float]]: +def _responses_tokens(response: Response) -> tuple[None, list[int], list[float]]: data = _dump(response) token_ids, logprobs = _pairs(data.get("raw_output_tokens")) if token_ids: @@ -155,7 +213,7 @@ def _responses_tokens(response: Any) -> tuple[None, list[int], list[float]]: return None, [], [] -def _messages_tokens(response: Any) -> tuple[None, list[int], list[float]]: +def _messages_tokens(response: Message) -> tuple[None, list[int], list[float]]: data = _dump(response) token_ids = [ token @@ -171,7 +229,7 @@ def _messages_tokens(response: Any) -> tuple[None, list[int], list[float]]: return None, token_ids, logprobs -def _exchange_list(trajectory: Trajectory, model: str | None) -> list[Any]: +def _exchange_list(trajectory: Trajectory, model: str | None) -> list[Exchange]: exchanges = [ *trajectory.exchanges.chat_completions, *trajectory.exchanges.completions, @@ -198,7 +256,7 @@ def _artifact_config(model: str) -> _TokenizerConfig: import wandb artifact_path = model.removeprefix("wandb-artifact:///") - artifact = getattr(wandb, "Api")().artifact(f"{artifact_path}:latest") + artifact = wandb.Api().artifact(f"{artifact_path}:latest") metadata = artifact.metadata base_model = metadata.get("base_model") or metadata.get("wandb.base_model") if not isinstance(base_model, str): @@ -235,7 +293,7 @@ def _tokenizer_config(model: str, base_model: str | None) -> _TokenizerConfig: return _TokenizerConfig(model) -def _load_tokenizer(config: _TokenizerConfig) -> Any: +def _load_tokenizer(config: _TokenizerConfig) -> _Tokenizer: try: from transformers import AutoTokenizer except ImportError as exc: @@ -243,9 +301,11 @@ def _load_tokenizer(config: _TokenizerConfig) -> Any: "Tokenizer fallback requires ART's backend or tinker dependencies" ) from exc try: - return AutoTokenizer.from_pretrained( - config.base_model, - revision=config.revision, + return _as_tokenizer( + AutoTokenizer.from_pretrained( + config.base_model, + revision=config.revision, + ) ) except Exception as exc: raise ValueError( @@ -253,31 +313,42 @@ def _load_tokenizer(config: _TokenizerConfig) -> Any: ) from exc -def _ids(value: Any) -> list[int]: - if hasattr(value, "input_ids"): - value = value.input_ids - if hasattr(value, "tolist"): - value = value.tolist() - if isinstance(value, dict): - value = value.get("input_ids") +def _ids(value: object) -> list[int]: + if (input_ids := getattr(value, "input_ids", None)) is not None: + value = input_ids + if callable(to_list := getattr(value, "tolist", None)): + value = to_list() + if mapping := _string_dict(value): + value = mapping.get("input_ids") if isinstance(value, list) and value and isinstance(value[0], list): value = value[0] - if not isinstance(value, list) or any(not isinstance(item, int) for item in value): + if not isinstance(value, list): raise TypeError("Tokenizer did not return one token ID sequence") - return value + token_ids = [ + item for item in value if isinstance(item, int) and not isinstance(item, bool) + ] + if len(token_ids) != len(value): + raise TypeError("Tokenizer did not return one token ID sequence") + return token_ids -def _content_text(content: Any) -> str: +def _content_text(content: object) -> str: if isinstance(content, str): return content if not isinstance(content, list): return "" - return "".join( - block.get("text", "") - for block in content - if isinstance(block, dict) - and block.get("type") in {"input_text", "output_text", "text"} - ) + text = "" + for block in content: + data = _string_dict(block) + if data is not None and data.get("type") in { + "input_text", + "output_text", + "text", + }: + value = data.get("text") + if isinstance(value, str): + text += value + return text def _anthropic_messages(request: dict[str, Any]) -> list[dict[str, Any]]: @@ -387,25 +458,26 @@ def _responses_messages(request: dict[str, Any]) -> list[dict[str, Any]]: return messages -def _openai_tools(tools: Any, *, dialect: str) -> Any: +def _openai_tools(tools: object, *, dialect: str) -> object: if not isinstance(tools, list) or dialect == "chat": return tools normalized = [] for tool in tools: - if not isinstance(tool, dict) or tool.get("type", "function") != "function": + data = _string_dict(tool) + if data is None or data.get("type", "function") != "function": normalized.append(tool) continue if dialect == "messages": function = { - "name": tool.get("name"), - "description": tool.get("description"), - "parameters": tool.get("input_schema", {}), + "name": data.get("name"), + "description": data.get("description"), + "parameters": data.get("input_schema", {}), } else: function = { - "name": tool.get("name"), - "description": tool.get("description"), - "parameters": tool.get("parameters", {}), + "name": data.get("name"), + "description": data.get("description"), + "parameters": data.get("parameters", {}), } normalized.append( { @@ -419,11 +491,12 @@ def _openai_tools(tools: Any, *, dialect: str) -> Any: def _request_messages( - exchange: Any, messages_override: list[dict[str, Any]] | None = None -) -> tuple[list[dict[str, Any]], Any]: + exchange: ChatCompletionsExchange | MessagesExchange | ResponsesExchange, + messages_override: list[dict[str, Any]] | None = None, +) -> tuple[list[dict[str, Any]], object]: request = exchange.request.root if isinstance(exchange, ChatCompletionsExchange): - return list(request.get("messages") or []), request.get("tools") + return _dict_list(request.get("messages")), request.get("tools") if isinstance(exchange, MessagesExchange): return _anthropic_messages(request), _openai_tools( request.get("tools"), dialect="messages" @@ -438,7 +511,9 @@ def _request_messages( raise TypeError("Completions requests do not use chat templates") -def _response_message(exchange: Any) -> dict[str, Any]: +def _response_message( + exchange: ChatCompletionsExchange | MessagesExchange | ResponsesExchange, +) -> dict[str, Any]: if isinstance(exchange, ChatCompletionsExchange): return exchange.response.choices[0].message.model_dump( mode="python", exclude_none=True @@ -476,13 +551,13 @@ def _response_message(exchange: Any) -> dict[str, Any]: def _template_ids( - tokenizer: Any, - exchange: Any, + tokenizer: _Tokenizer, + exchange: Exchange, *, completed: bool, config: _TokenizerConfig, chat_template: str | None, - chat_template_kwargs: dict[str, Any] | None, + chat_template_kwargs: Mapping[str, object] | None, messages_override: list[dict[str, Any]] | None = None, ) -> list[int]: request = exchange.request.root @@ -528,7 +603,9 @@ def _template_ids( return _ids(result) -def _exchange_tokens(exchange: Any) -> tuple[list[int] | None, list[int], list[float]]: +def _exchange_tokens( + exchange: Exchange, +) -> tuple[list[int] | None, list[int], list[float]]: if isinstance(exchange, ChatCompletionsExchange): return _chat_tokens(exchange.response) if isinstance(exchange, CompletionsExchange): @@ -540,7 +617,7 @@ def _exchange_tokens(exchange: Any) -> tuple[list[int] | None, list[int], list[f raise TypeError(f"Unknown exchange type: {type(exchange)!r}") -def _visible_logprobs(exchange: Any) -> list[tuple[str, float]]: +def _visible_logprobs(exchange: Exchange) -> list[tuple[str, float]]: values: list[tuple[str, float]] = [] if isinstance(exchange, ChatCompletionsExchange): logprobs = exchange.response.choices[0].logprobs @@ -577,10 +654,10 @@ def _visible_logprobs(exchange: Any) -> list[tuple[str, float]]: def _align_visible_logprobs( - tokenizer: Any, completion: list[int], exchange: Any + tokenizer: _Tokenizer | None, completion: list[int], exchange: Exchange ) -> list[float] | None: values = _visible_logprobs(exchange) - if not values or not callable(tokenizer): + if not values or tokenizer is None: return None aligned = [math.nan] * len(completion) cursor = 0 @@ -602,7 +679,7 @@ def _legacy_tokenize( base_model: str | None, *, chat_template: str | None, - chat_template_kwargs: dict[str, Any] | None, + chat_template_kwargs: Mapping[str, object] | None, ) -> TokenizedTrajectory: if trajectory.additional_histories: raise ValueError("Tokenization requires one history") @@ -612,11 +689,7 @@ def _legacy_tokenize( for item in trajectory.messages_and_choices: if not isinstance(item, Choice): continue - prompt, completion, completion_logprobs = _chat_tokens( - type( - "Response", (), {"choices": [item], "model_dump": lambda self, **_: {}} - )() - ) + prompt, completion, completion_logprobs = _chat_choice_tokens(item, {}) if prompt is None or not completion: raise ValueError( "Legacy fallback tokenization is unavailable without exact choice token metadata" @@ -651,8 +724,8 @@ def tokenize_one( *, model: str | None, chat_template: str | None, - chat_template_kwargs: dict[str, Any] | None, - tokenizer_instance: Any = None, + chat_template_kwargs: Mapping[str, object] | None, + tokenizer_instance: _Tokenizer | None = None, ) -> TokenizedTrajectory: if not trajectory.exchanges: return _legacy_tokenize( @@ -662,7 +735,9 @@ def tokenize_one( chat_template_kwargs=chat_template_kwargs, ) exchanges = _exchange_list(trajectory, model) - selected_model = cast(str, exchanges[0].model) + selected_model = exchanges[0].model + if selected_model is None: + raise AssertionError("_exchange_list returned an exchange without a model") config = _tokenizer_config(selected_model, base_model) tokenizer = tokenizer_instance token_ids: list[int] = [] @@ -690,9 +765,9 @@ def tokenize_one( _response_message(exchange), ] prompt, completion, completion_logprobs = _exchange_tokens(exchange) - if prompt is None or not completion: - tokenizer = tokenizer or _load_tokenizer(config) if prompt is None: + if tokenizer is None: + tokenizer = _load_tokenizer(config) prompt = _template_ids( tokenizer, exchange, @@ -703,6 +778,8 @@ def tokenize_one( messages_override=messages_override, ) if not completion: + if tokenizer is None: + tokenizer = _load_tokenizer(config) completed = _template_ids( tokenizer, exchange, diff --git a/tests/unit/test_auto_trajectory.py b/tests/unit/test_auto_trajectory.py index 6a5386944..e527abb36 100644 --- a/tests/unit/test_auto_trajectory.py +++ b/tests/unit/test_auto_trajectory.py @@ -290,7 +290,9 @@ async def say_hi() -> str | None: assert len(exchanges) == 5 assert exchanges[0].request.root["messages"] == [message] assert exchanges[0].request.root["tools"] == tools - assert exchanges[0].response.choices[0] == Choice(**mock_response["choices"][0]) # ty:ignore[invalid-argument-type, not-subscriptable] + choices = mock_response["choices"] + assert isinstance(choices, list) + assert exchanges[0].response.choices[0] == Choice.model_validate(choices[0]) assert exchanges[-1].response.choices[0] == mock_stream_choice assert all(exchange.model == "test" for exchange in exchanges) diff --git a/tests/unit/trajectories/test_capture.py b/tests/unit/trajectories/test_capture.py index 8e37b193f..6cc47145d 100644 --- a/tests/unit/trajectories/test_capture.py +++ b/tests/unit/trajectories/test_capture.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +from collections.abc import AsyncIterator from datetime import datetime, timedelta import json from typing import Any @@ -8,6 +9,7 @@ import aiohttp from aiohttp import web from anthropic import AsyncAnthropic +from anthropic.types import TextBlock import httpx from openai import AsyncOpenAI import pytest @@ -15,7 +17,8 @@ import requests import art -from art.trajectories._protocols import build_exchange +from art.trajectories import ChatCompletionsExchange, MessagesExchange +from art.trajectories._protocols import Endpoint, build_exchange CHAT: dict[str, Any] = { "id": "chatcmpl-1", @@ -107,7 +110,7 @@ @pytest_asyncio.fixture -async def endpoint_server(unused_tcp_port: int): +async def endpoint_server(unused_tcp_port: int) -> AsyncIterator[str]: async def handler(request: web.Request) -> web.Response: request_body = await request.json() if request_body.get("fail"): @@ -171,6 +174,7 @@ async def rollout() -> None: assert isinstance(captured, art.Trajectory) task = asyncio.create_task(rollout()) with pytest.raises(TypeError, match="raw coroutine"): + # Passing a Task is deliberately a static type error and a runtime error. await art.trajectory(task) # type: ignore[arg-type] await task @@ -262,7 +266,7 @@ async def test_failed_and_incomplete_calls_are_excluded(endpoint_server: str) -> def test_all_protocols_reconstruct_typed_responses() -> None: now = datetime.now() - values = [ + values: list[tuple[Endpoint, dict[str, Any], dict[str, Any]]] = [ ("chat_completions", {"model": "request-model", "messages": []}, CHAT), ("completions", {"model": "request-model", "prompt": "hi"}, COMPLETION), ("responses", {"input": "hi"}, RESPONSE), @@ -361,7 +365,7 @@ def test_all_streaming_protocols_reconstruct_final_responses() -> None: ), ("message_stop", {"type": "message_stop"}), ] - values = [ + values: list[tuple[Endpoint, dict[str, Any], bytes]] = [ ( "chat_completions", {"model": "test/model", "messages": [], "stream": True}, @@ -392,8 +396,10 @@ def test_all_streaming_protocols_reconstruct_final_responses() -> None: end_time=now + timedelta(seconds=1), ) assert exchange.model == "test/model" - if endpoint == "messages": - assert exchange.response.content[0].text == "hello" + if isinstance(exchange, MessagesExchange): + content = exchange.response.content[0] + assert isinstance(content, TextBlock) + assert content.text == "hello" def test_trajectory_rejects_mixed_representations() -> None: @@ -404,6 +410,7 @@ def test_trajectory_rejects_mixed_representations() -> None: start_time=datetime.now(), end_time=datetime.now(), ) + assert isinstance(exchange, ChatCompletionsExchange) with pytest.raises(ValueError, match="both exchanges and legacy histories"): art.Trajectory( exchanges=art.TrajectoryExchanges(chat_completions=[exchange]), @@ -413,6 +420,13 @@ def test_trajectory_rejects_mixed_representations() -> None: def test_metadata_accepts_json_serializable_values() -> None: assert art.Trajectory().model_dump() == {} + assert art.Trajectory().model_dump( + mode="json", include={"reward"}, exclude_defaults=False + ) == {"reward": 0.0} + assert ( + art.Trajectory().model_dump_json(include={"reward"}, exclude_defaults=False) + == '{"reward":0.0}' + ) trajectory = art.Trajectory(metadata={"nested": {"items": [1, "two"]}}) assert trajectory.model_dump(mode="json")["metadata"] == { "nested": {"items": [1, "two"]} diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index e102f56eb..628d3269e 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -120,14 +120,18 @@ def test_branching_and_multiple_models_require_explicit_resolution() -> None: class _FakeTokenizer: def __init__(self) -> None: - self.calls: list[dict[str, Any]] = [] + self.calls: list[dict[str, object]] = [] - def apply_chat_template(self, messages: list[dict[str, Any]], **kwargs: Any): + def apply_chat_template( + self, messages: list[dict[str, Any]], **kwargs: object + ) -> list[int]: self.calls.append(kwargs) return [10, 11] if messages[-1]["role"] == "assistant" else [10] -def test_fallback_uses_template_overrides_and_nan_logprobs(monkeypatch) -> None: +def test_fallback_uses_template_overrides_and_nan_logprobs( + monkeypatch: pytest.MonkeyPatch, +) -> None: response = Message.model_validate( { "id": "msg_1", @@ -195,11 +199,13 @@ def test_fallback_uses_template_overrides_and_nan_logprobs(monkeypatch) -> None: ] -def test_checkpoint_fallback_uses_latest_artifact_renderer(monkeypatch) -> None: +def test_checkpoint_fallback_uses_latest_artifact_renderer( + monkeypatch: pytest.MonkeyPatch, +) -> None: artifact_names: list[str] = [] class Api: - def artifact(self, name: str): + def artifact(self, name: str) -> SimpleNamespace: artifact_names.append(name) return SimpleNamespace( metadata={ @@ -283,7 +289,9 @@ def test_anthropic_fallback_preserves_thinking_and_tool_history() -> None: ] -def test_choice_logprobs_survive_tokenizer_fallback(monkeypatch) -> None: +def test_choice_logprobs_survive_tokenizer_fallback( + monkeypatch: pytest.MonkeyPatch, +) -> None: exchange = _chat_exchange([], []) logprobs = exchange.response.choices[0].logprobs assert logprobs is not None @@ -301,13 +309,15 @@ def test_choice_logprobs_survive_tokenizer_fallback(monkeypatch) -> None: ) class Tokenizer: - def apply_chat_template(self, messages, **kwargs): + 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] - def __call__(self, text, **kwargs): + def __call__(self, text: str, **kwargs: object) -> SimpleNamespace: del text, kwargs - return type("Encoding", (), {"input_ids": [11]})() + return SimpleNamespace(input_ids=[11]) monkeypatch.setattr( "art.trajectories._tokenize._load_tokenizer", lambda _config: Tokenizer() @@ -378,9 +388,13 @@ def _response_exchange( ) -def test_responses_previous_response_id_resolves_local_history(monkeypatch) -> None: +def test_responses_previous_response_id_resolves_local_history( + monkeypatch: pytest.MonkeyPatch, +) -> None: class Tokenizer: - def apply_chat_template(self, messages, **kwargs): + def apply_chat_template( + self, messages: list[dict[str, Any]], **kwargs: object + ) -> list[int]: del kwargs return [10] if len(messages) == 1 else [10, 20, 11] @@ -433,6 +447,7 @@ def decode(self, token_id: int) -> str: results = list( tokenize_trajectory_groups( + # This path only calls decode; the minimal test double is intentional. Tokenizer(), # type: ignore[arg-type] [group], allow_training_without_logprobs=True, From 770daa67d9050fc6ba2027ecbcab5bbd0be1e7e4 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 15 Jul 2026 00:42:27 +0000 Subject: [PATCH 03/17] api: make trajectory base model keyword-only --- src/art/trajectories/__init__.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/src/art/trajectories/__init__.py b/src/art/trajectories/__init__.py index 8fcbabcc7..4924eb8d1 100644 --- a/src/art/trajectories/__init__.py +++ b/src/art/trajectories/__init__.py @@ -354,8 +354,8 @@ async def trajectory_group( def tokenize_trajectory( trajectory: Trajectory, - base_model: str | None = None, *, + base_model: str | None = None, model: str | None = None, chat_template: str | None = None, chat_template_kwargs: Mapping[str, object] | None = None, @@ -373,8 +373,8 @@ def tokenize_trajectory( def tokenize_trajectories( trajectories: Iterable[Trajectory], - base_model: str | None = None, *, + base_model: str | None = None, model: str | None = None, chat_template: str | None = None, chat_template_kwargs: Mapping[str, object] | None = None, @@ -382,7 +382,7 @@ def tokenize_trajectories( return [ tokenize_trajectory( item, - base_model, + base_model=base_model, model=model, chat_template=chat_template, chat_template_kwargs=chat_template_kwargs, @@ -393,8 +393,8 @@ def tokenize_trajectories( def tokenize_trajectory_group( group: TrajectoryGroup, - base_model: str | None = None, *, + base_model: str | None = None, model: str | None = None, chat_template: str | None = None, chat_template_kwargs: Mapping[str, object] | None = None, @@ -402,7 +402,7 @@ def tokenize_trajectory_group( return TokenizedTrajectoryGroup( trajectories=tokenize_trajectories( group, - base_model, + base_model=base_model, model=model, chat_template=chat_template, chat_template_kwargs=chat_template_kwargs, @@ -413,8 +413,8 @@ def tokenize_trajectory_group( def tokenize_trajectory_groups( groups: Iterable[TrajectoryGroup], - base_model: str | None = None, *, + base_model: str | None = None, model: str | None = None, chat_template: str | None = None, chat_template_kwargs: Mapping[str, object] | None = None, @@ -422,7 +422,7 @@ def tokenize_trajectory_groups( return [ tokenize_trajectory_group( group, - base_model, + base_model=base_model, model=model, chat_template=chat_template, chat_template_kwargs=chat_template_kwargs, From 39a165b487158d10b71b787c4dfdd2dc3764f913 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 15 Jul 2026 03:30:41 +0000 Subject: [PATCH 04/17] api: type trajectory request payloads --- src/art/__init__.py | 4 +- src/art/auto_trajectory.py | 5 +- src/art/trajectories/__init__.py | 104 +++++++++++++++--- src/art/trajectories/_compat.py | 4 +- src/art/trajectories/_protocols.py | 10 +- src/art/trajectories/_tokenize.py | 16 +-- .../megatron/train_inf_mismatch/real_path.py | 6 +- tests/unit/test_auto_trajectory.py | 18 +-- tests/unit/trajectories/test_capture.py | 3 +- tests/unit/trajectories/test_tokenize.py | 28 ++--- 10 files changed, 141 insertions(+), 57 deletions(-) diff --git a/src/art/__init__.py b/src/art/__init__.py index c715dafb2..0f3b9f852 100644 --- a/src/art/__init__.py +++ b/src/art/__init__.py @@ -79,8 +79,8 @@ Trajectory, TrajectoryExchanges, TrajectoryGroup, - auto_trajectory, - capture_auto_trajectory, + auto_trajectory, # ty: ignore[deprecated] + capture_auto_trajectory, # ty: ignore[deprecated] current_trajectory, current_trajectory_group, tokenize_trajectories, diff --git a/src/art/auto_trajectory.py b/src/art/auto_trajectory.py index f4f2f64d3..9560b9af8 100644 --- a/src/art/auto_trajectory.py +++ b/src/art/auto_trajectory.py @@ -1,5 +1,8 @@ """Deprecated compatibility names for automatic trajectory capture.""" -from .trajectories import auto_trajectory, capture_auto_trajectory +from .trajectories import ( + auto_trajectory, # ty: ignore[deprecated] + capture_auto_trajectory, # ty: ignore[deprecated] +) __all__ = ["auto_trajectory", "capture_auto_trajectory"] diff --git a/src/art/trajectories/__init__.py b/src/art/trajectories/__init__.py index 4924eb8d1..3fefa13e9 100644 --- a/src/art/trajectories/__init__.py +++ b/src/art/trajectories/__init__.py @@ -12,14 +12,38 @@ from datetime import datetime import time from types import TracebackType -from typing import Any, Literal, overload +from typing import Annotated, Any, Literal, overload -from anthropic.types import Message as AnthropicMessage +from anthropic.types import ( + Message as AnthropicMessage, +) +from anthropic.types import ( + MessageParam as AnthropicMessageParam, +) +from anthropic.types import ( + TextBlockParam as AnthropicTextBlockParam, +) +from anthropic.types import ( + ThinkingConfigParam as AnthropicThinkingConfigParam, +) +from anthropic.types import ( + ToolUnionParam as AnthropicToolParam, +) from openai.types import Completion -from openai.types.chat import ChatCompletion -from openai.types.responses import Response +from openai.types.chat import ( + ChatCompletion, + ChatCompletionMessageParam, + ChatCompletionToolParam, +) +from openai.types.responses import ( + Response, + ResponseInputParam, +) +from openai.types.responses import ( + ToolParam as ResponsesToolParam, +) import pydantic -from typing_extensions import deprecated +from typing_extensions import TypedDict, deprecated from ..types import Messages, MessagesAndChoices, Tools from ._serialization import _CompactModel @@ -28,24 +52,74 @@ MetadataValue = Any -class ChatCompletionsRequest(pydantic.RootModel[dict[str, pydantic.JsonValue]]): +class ChatCompletionsRequest(TypedDict, total=False, extra_items=Any): """The JSON body sent to an OpenAI-compatible Chat Completions endpoint.""" - -class CompletionsRequest(pydantic.RootModel[dict[str, pydantic.JsonValue]]): + model: str + messages: list[ChatCompletionMessageParam] + stream: bool + tools: list[ChatCompletionToolParam] + max_completion_tokens: int + max_tokens: int + temperature: float + top_p: float + logprobs: bool + top_logprobs: int + chat_template: str + chat_template_kwargs: dict[str, Any] + + +class CompletionsRequest(TypedDict, total=False, extra_items=Any): """The JSON body sent to an OpenAI-compatible Completions endpoint.""" + model: str + prompt: str | list[str] | list[int] | list[list[int]] + stream: bool + max_tokens: int + temperature: float + top_p: float + logprobs: int + echo: bool + stop: str | list[str] + seed: int -class ResponsesRequest(pydantic.RootModel[dict[str, pydantic.JsonValue]]): + +class ResponsesRequest(TypedDict, total=False, extra_items=Any): """The JSON body sent to an OpenAI-compatible Responses endpoint.""" + model: str + input: str | ResponseInputParam + instructions: str + previous_response_id: str + stream: bool + tools: list[ResponsesToolParam] + max_output_tokens: int + temperature: float + top_p: float + chat_template: str + chat_template_kwargs: dict[str, Any] + -class MessagesRequest(pydantic.RootModel[dict[str, pydantic.JsonValue]]): +class MessagesRequest(TypedDict, total=False, extra_items=Any): """The JSON body sent to an Anthropic-compatible Messages endpoint.""" + model: str + messages: list[AnthropicMessageParam] + max_tokens: int + stream: bool + system: str | list[AnthropicTextBlockParam] + tools: list[AnthropicToolParam] + thinking: AnthropicThinkingConfigParam + temperature: float + top_p: float + top_k: int + stop_sequences: list[str] + chat_template: str + chat_template_kwargs: dict[str, Any] + class ChatCompletionsExchange(pydantic.BaseModel): - request: ChatCompletionsRequest + request: Annotated[ChatCompletionsRequest, pydantic.SkipValidation] response: ChatCompletion model: str | None start_time: datetime @@ -53,7 +127,7 @@ class ChatCompletionsExchange(pydantic.BaseModel): class CompletionsExchange(pydantic.BaseModel): - request: CompletionsRequest + request: Annotated[CompletionsRequest, pydantic.SkipValidation] response: Completion model: str | None start_time: datetime @@ -61,7 +135,7 @@ class CompletionsExchange(pydantic.BaseModel): class ResponsesExchange(pydantic.BaseModel): - request: ResponsesRequest + request: Annotated[ResponsesRequest, pydantic.SkipValidation] response: Response model: str | None start_time: datetime @@ -69,7 +143,7 @@ class ResponsesExchange(pydantic.BaseModel): class MessagesExchange(pydantic.BaseModel): - request: MessagesRequest + request: Annotated[MessagesRequest, pydantic.SkipValidation] response: AnthropicMessage model: str | None start_time: datetime @@ -284,7 +358,7 @@ def log(self, message: str) -> None: self.logs.append(message) # Legacy groups iterate over trajectories rather than Pydantic field pairs. - def __iter__(self) -> Iterator[Trajectory]: # type: ignore[override] + def __iter__(self) -> Iterator[Trajectory]: # ty: ignore[invalid-method-override] return iter(self.trajectories) def __len__(self) -> int: diff --git a/src/art/trajectories/_compat.py b/src/art/trajectories/_compat.py index 98c2eb88f..13785d2ac 100644 --- a/src/art/trajectories/_compat.py +++ b/src/art/trajectories/_compat.py @@ -85,7 +85,9 @@ def new_trajectory_group( logs: list[str] | None, ) -> TrajectoryGroup | Awaitable[TrajectoryGroup]: items = list(trajectories) - awaitables = [item for item in items if isinstance(item, Awaitable)] + awaitables = [ + item for item in items if not isinstance(item, (Trajectory, BaseException)) + ] if awaitables: if len(awaitables) != len(items): raise TypeError("TrajectoryGroup cannot mix trajectories and awaitables") diff --git a/src/art/trajectories/_protocols.py b/src/art/trajectories/_protocols.py index ede72e26b..74c76716b 100644 --- a/src/art/trajectories/_protocols.py +++ b/src/art/trajectories/_protocols.py @@ -33,7 +33,7 @@ ) ResponseModel = ChatCompletion | Completion | Response | Message SSEPayload = dict[str, Any] | Literal["[DONE]"] -_ENDPOINTS = { +_ENDPOINTS: dict[str, Endpoint] = { "/v1/chat/completions": "chat_completions", "/v1/completions": "completions", "/v1/responses": "responses", @@ -225,7 +225,7 @@ def build_exchange( if endpoint == "chat_completions": response = _chat_response(body, stream=stream) return endpoint, ChatCompletionsExchange( - request=ChatCompletionsRequest(request), + request=ChatCompletionsRequest(**request), response=response, model=_model(request, response), start_time=start_time, @@ -234,7 +234,7 @@ def build_exchange( if endpoint == "completions": response = _completion_response(body, stream=stream) return endpoint, CompletionsExchange( - request=CompletionsRequest(request), + request=CompletionsRequest(**request), response=response, model=_model(request, response), start_time=start_time, @@ -243,7 +243,7 @@ def build_exchange( if endpoint == "responses": response = _responses_response(body, stream=stream) return endpoint, ResponsesExchange( - request=ResponsesRequest(request), + request=ResponsesRequest(**request), response=response, model=_model(request, response), start_time=start_time, @@ -252,7 +252,7 @@ def build_exchange( if endpoint == "messages": response = _messages_response(body, stream=stream) return endpoint, MessagesExchange( - request=MessagesRequest(request), + request=MessagesRequest(**request), response=response, model=_model(request, response), start_time=start_time, diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index e2ff2e99d..62b81ea77 100644 --- a/src/art/trajectories/_tokenize.py +++ b/src/art/trajectories/_tokenize.py @@ -50,9 +50,11 @@ def apply_chat_template( ) -> object: ... -def _as_tokenizer(tokenizer: PreTrainedTokenizerBase) -> _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) @@ -253,10 +255,10 @@ def _exchange_list(trajectory: Trajectory, model: str | None) -> list[Exchange]: def _artifact_config(model: str) -> _TokenizerConfig: - import wandb + from wandb.apis.public import Api artifact_path = model.removeprefix("wandb-artifact:///") - artifact = wandb.Api().artifact(f"{artifact_path}:latest") + artifact = Api().artifact(f"{artifact_path}:latest") metadata = artifact.metadata base_model = metadata.get("base_model") or metadata.get("wandb.base_model") if not isinstance(base_model, str): @@ -494,7 +496,7 @@ def _request_messages( exchange: ChatCompletionsExchange | MessagesExchange | ResponsesExchange, messages_override: list[dict[str, Any]] | None = None, ) -> tuple[list[dict[str, Any]], object]: - request = exchange.request.root + request = exchange.request if isinstance(exchange, ChatCompletionsExchange): return _dict_list(request.get("messages")), request.get("tools") if isinstance(exchange, MessagesExchange): @@ -560,11 +562,11 @@ def _template_ids( chat_template_kwargs: Mapping[str, object] | None, messages_override: list[dict[str, Any]] | None = None, ) -> list[int]: - request = exchange.request.root + request = exchange.request if isinstance(exchange, CompletionsExchange): prompt = request.get("prompt", "") if isinstance(prompt, list) and all(isinstance(item, int) for item in prompt): - prompt_ids = prompt + prompt_ids = _ids(prompt) else: prompt_ids = _ids(tokenizer(str(prompt), add_special_tokens=False)) if not completed: @@ -748,7 +750,7 @@ def tokenize_one( for exchange in exchanges: messages_override = None if isinstance(exchange, ResponsesExchange): - request = exchange.request.root + request = exchange.request messages_override = _responses_messages(request) previous = request.get("previous_response_id") if previous is not None: diff --git a/tests/integration/megatron/train_inf_mismatch/real_path.py b/tests/integration/megatron/train_inf_mismatch/real_path.py index 164f28fb6..47b8e09e4 100644 --- a/tests/integration/megatron/train_inf_mismatch/real_path.py +++ b/tests/integration/megatron/train_inf_mismatch/real_path.py @@ -347,14 +347,16 @@ async def _request() -> None: top_logprobs=TOP_K, **request_kwargs, ) - if trajectory := art.auto_trajectory(): + if trajectory := art.auto_trajectory(): # ty: ignore[deprecated] logprobs = response.choices[0].logprobs trajectory.reward = reward trajectory.metrics["completion_tokens"] = ( len(logprobs.content or []) if logprobs is not None else 0 ) - return await art.capture_auto_trajectory(_request()) + return await art.capture_auto_trajectory( # ty: ignore[deprecated] + _request() + ) async def _collect_real_trajectory_groups( diff --git a/tests/unit/test_auto_trajectory.py b/tests/unit/test_auto_trajectory.py index e527abb36..9322652c9 100644 --- a/tests/unit/test_auto_trajectory.py +++ b/tests/unit/test_auto_trajectory.py @@ -280,16 +280,18 @@ async def say_hi() -> str | None: ): pass # Add ART support with a couple lines of optional code - if trajectory := art.auto_trajectory(): + if trajectory := art.auto_trajectory(): # ty: ignore[deprecated] trajectory.reward = 1.0 return chat_completion.choices[0].message.content # Use the capture_auto_trajectory utility to capture a trajectory automatically - trajectory = await art.capture_auto_trajectory(say_hi()) + trajectory = await art.capture_auto_trajectory( # ty: ignore[deprecated] + say_hi() + ) exchanges = trajectory.exchanges.chat_completions assert len(exchanges) == 5 - assert exchanges[0].request.root["messages"] == [message] - assert exchanges[0].request.root["tools"] == tools + assert exchanges[0].request["messages"] == [message] + assert exchanges[0].request["tools"] == tools choices = mock_response["choices"] assert isinstance(choices, list) assert exchanges[0].response.choices[0] == Choice.model_validate(choices[0]) @@ -356,14 +358,16 @@ async def say_hi() -> str | None: async for _ in stream: pass # Add ART support with a couple lines of optional code - if trajectory := art.auto_trajectory(): + if trajectory := art.auto_trajectory(): # ty: ignore[deprecated] trajectory.reward = 1.0 return choice.message.content # Use the capture_auto_trajectory utility to capture a trajectory automatically - trajectory = await art.capture_auto_trajectory(say_hi()) + trajectory = await art.capture_auto_trajectory( # ty: ignore[deprecated] + say_hi() + ) exchanges = trajectory.exchanges.chat_completions assert len(exchanges) == 3 - assert exchanges[0].request.root["messages"] == [message] + assert exchanges[0].request["messages"] == [message] assert exchanges[-1].response.choices[0] == mock_stream_choice assert all(exchange.model == "test" for exchange in exchanges) diff --git a/tests/unit/trajectories/test_capture.py b/tests/unit/trajectories/test_capture.py index 6cc47145d..3cc3ecbc2 100644 --- a/tests/unit/trajectories/test_capture.py +++ b/tests/unit/trajectories/test_capture.py @@ -175,7 +175,7 @@ async def rollout() -> None: task = asyncio.create_task(rollout()) with pytest.raises(TypeError, match="raw coroutine"): # Passing a Task is deliberately a static type error and a runtime error. - await art.trajectory(task) # type: ignore[arg-type] + await art.trajectory(task) # ty: ignore[invalid-argument-type] await task async def failed() -> art.Trajectory: @@ -239,6 +239,7 @@ async def test_native_openai_and_anthropic_sdks(endpoint_server: str) -> None: assert completion.choices[0].text == "hello" assert response.output_text == "hello" + assert message.content[0].type == "text" assert message.content[0].text == "hello" assert len(trajectory.exchanges.completions) == 1 assert len(trajectory.exchanges.responses) == 1 diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index 628d3269e..a83adbacc 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -61,10 +61,8 @@ def _chat_exchange( start = datetime(2026, 1, 1) + timedelta(seconds=offset) return ChatCompletionsExchange( request=ChatCompletionsRequest( - { - "model": model, - "messages": [{"role": "user", "content": f"turn {offset}"}], - } + model=model, + messages=[{"role": "user", "content": f"turn {offset}"}], ), response=response, model=model, @@ -147,13 +145,11 @@ def test_fallback_uses_template_overrides_and_nan_logprobs( start = datetime(2026, 1, 1) exchange = MessagesExchange( request=MessagesRequest( - { - "model": "test/model", - "messages": [{"role": "user", "content": "question"}], - "chat_template": "request-template", - "chat_template_kwargs": {"request": True}, - "thinking": {"type": "enabled", "budget_tokens": 128}, - } + model="test/model", + messages=[{"role": "user", "content": "question"}], + chat_template="request-template", + chat_template_kwargs={"request": True}, + thinking={"type": "enabled", "budget_tokens": 128}, ), response=response, model="test/model", @@ -218,7 +214,7 @@ def artifact(self, name: str) -> SimpleNamespace: } ) - monkeypatch.setattr("wandb.Api", Api) + monkeypatch.setattr("wandb.apis.public.Api", Api) from art.trajectories._tokenize import _tokenizer_config config = _tokenizer_config("wandb-artifact:///entity/project/run", None) @@ -375,12 +371,12 @@ def _response_exchange( "raw_output_tokens": [{"token_id": output_id, "logprob": -0.1}], } ) - request = {"model": "test/model", "input": f"turn {offset}"} + request = ResponsesRequest(model="test/model", input=f"turn {offset}") if previous_response_id is not None: request["previous_response_id"] = previous_response_id start = datetime(2026, 1, 1) + timedelta(seconds=offset) return ResponsesExchange( - request=ResponsesRequest(request), + request=request, response=response, model="test/model", start_time=start, @@ -414,7 +410,7 @@ def apply_chat_template( 30, ] - second.request.root["previous_response_id"] = "missing" + second.request["previous_response_id"] = "missing" with pytest.raises(ValueError, match="outside this trajectory"): art.tokenize_trajectory(trajectory, base_model="base/model") @@ -448,7 +444,7 @@ def decode(self, token_id: int) -> str: results = list( tokenize_trajectory_groups( # This path only calls decode; the minimal test double is intentional. - Tokenizer(), # type: ignore[arg-type] + Tokenizer(), # type: ignore[arg-type, ty:invalid-argument-type] [group], allow_training_without_logprobs=True, scale_rewards=False, From 40e88c8e142898109624e174f782794bb3f5a592 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 15 Jul 2026 03:59:10 +0000 Subject: [PATCH 05/17] fix: harden trajectory capture and tokenization --- src/art/trajectories/_capture/aiohttp.py | 5 +- src/art/trajectories/_capture/core.py | 10 +- src/art/trajectories/_capture/httpx.py | 14 +- src/art/trajectories/_capture/requests.py | 4 +- src/art/trajectories/_protocols.py | 31 ++- src/art/trajectories/_tokenize.py | 286 +++++++++++++++----- tests/unit/trajectories/test_capture.py | 166 +++++++++++- tests/unit/trajectories/test_tokenize.py | 303 ++++++++++++++++++++++ 8 files changed, 720 insertions(+), 99 deletions(-) diff --git a/src/art/trajectories/_capture/aiohttp.py b/src/art/trajectories/_capture/aiohttp.py index 87192b612..3b47a53a1 100644 --- a/src/art/trajectories/_capture/aiohttp.py +++ b/src/art/trajectories/_capture/aiohttp.py @@ -44,11 +44,14 @@ async def readchunk(self) -> tuple[bytes, bool]: return self._record(await self._stream.readchunk()) async def _iterate(self, iterator: AsyncIterator[bytes]) -> AsyncIterator[bytes]: + completed = False try: async for chunk in iterator: yield self._record(chunk) + completed = True finally: - self._state.finish() + if completed or self._state.request.get("stream") is True: + self._state.finish() def __aiter__(self) -> AsyncIterator[bytes]: return self._iterate(self._stream.__aiter__()) diff --git a/src/art/trajectories/_capture/core.py b/src/art/trajectories/_capture/core.py index 840b9695e..d3bade490 100644 --- a/src/art/trajectories/_capture/core.py +++ b/src/art/trajectories/_capture/core.py @@ -1,6 +1,7 @@ from __future__ import annotations import contextvars +from copy import deepcopy from dataclasses import dataclass, field from datetime import datetime import json @@ -44,7 +45,7 @@ def finish(self) -> None: if self.status_code is None or not 200 <= self.status_code < 300: return try: - _, exchange = build_exchange( + exchange = build_exchange( self.endpoint, self.request, bytes(self.body), @@ -74,7 +75,12 @@ def _json_body(value: object) -> dict[str, Any] | None: if isinstance(value, dict): if not all(isinstance(key, str) for key in value): return None - return {key: item for key, item in value.items() if isinstance(key, str)} + try: + return deepcopy( + {key: item for key, item in value.items() if isinstance(key, str)} + ) + except Exception: + return None if isinstance(value, str): value = value.encode() if not isinstance(value, bytes): diff --git a/src/art/trajectories/_capture/httpx.py b/src/art/trajectories/_capture/httpx.py index aac980b23..31293405c 100644 --- a/src/art/trajectories/_capture/httpx.py +++ b/src/art/trajectories/_capture/httpx.py @@ -76,36 +76,42 @@ def iter_bytes( self: httpx.Response, chunk_size: int | None = None ) -> Iterator[bytes]: state: CaptureState | None = getattr(self, _STATE, None) + completed = False try: for chunk in original_iter(self, chunk_size): if state is not None: state.add(chunk) yield chunk + completed = True finally: - if state is not None: + if state is not None and (completed or state.request.get("stream") is True): state.finish() async def aiter_bytes( self: httpx.Response, chunk_size: int | None = None ) -> AsyncIterator[bytes]: state: CaptureState | None = getattr(self, _STATE, None) + completed = False try: async for chunk in original_aiter(self, chunk_size): if state is not None: state.add(chunk) yield chunk + completed = True finally: - if state is not None: + if state is not None and (completed or state.request.get("stream") is True): state.finish() def close(self: httpx.Response) -> None: original_close(self) - if state := getattr(self, _STATE, None): + state: CaptureState | None = getattr(self, _STATE, None) + if state is not None and state.request.get("stream") is True: state.finish() async def aclose(self: httpx.Response) -> None: await original_aclose(self) - if state := getattr(self, _STATE, None): + state: CaptureState | None = getattr(self, _STATE, None) + if state is not None and state.request.get("stream") is True: state.finish() setattr(send, "_art_capture", True) diff --git a/src/art/trajectories/_capture/requests.py b/src/art/trajectories/_capture/requests.py index da1b0009e..242b2edc6 100644 --- a/src/art/trajectories/_capture/requests.py +++ b/src/art/trajectories/_capture/requests.py @@ -38,6 +38,7 @@ def iter_content( decode_unicode: bool = False, ) -> Iterator[str | bytes]: state: CaptureState | None = getattr(self, _STATE, None) + completed = False try: for chunk in original_iter( self, chunk_size=chunk_size, decode_unicode=decode_unicode @@ -48,8 +49,9 @@ def iter_content( if isinstance(chunk, bytes): state.add(chunk) yield chunk + completed = True finally: - if state is not None: + if state is not None and (completed or state.request.get("stream") is True): state.finish() setattr(send, "_art_capture", True) diff --git a/src/art/trajectories/_protocols.py b/src/art/trajectories/_protocols.py index 74c76716b..41baa1b93 100644 --- a/src/art/trajectories/_protocols.py +++ b/src/art/trajectories/_protocols.py @@ -34,10 +34,10 @@ ResponseModel = ChatCompletion | Completion | Response | Message SSEPayload = dict[str, Any] | Literal["[DONE]"] _ENDPOINTS: dict[str, Endpoint] = { - "/v1/chat/completions": "chat_completions", - "/v1/completions": "completions", - "/v1/responses": "responses", - "/v1/messages": "messages", + "/chat/completions": "chat_completions", + "/completions": "completions", + "/responses": "responses", + "/messages": "messages", } @@ -75,6 +75,7 @@ def _chat_response(body: bytes, *, stream: bool) -> ChatCompletion: if not stream: return ChatCompletion.model_validate_json(body) response: ChatCompletion | None = None + choices: dict[int, ChatCompletion] = {} done = False for _, payload in _sse_events(body): if payload == "[DONE]": @@ -82,10 +83,18 @@ def _chat_response(body: bytes, *, stream: bool) -> ChatCompletion: continue chunk = ChatCompletionChunk.model_validate(payload) if response is None: - response = init_chat_completion(chunk) - update_chat_completion(response, chunk) + response = init_chat_completion(chunk.model_copy(update={"choices": []})) + update_chat_completion(response, chunk.model_copy(update={"choices": []})) + for choice in chunk.choices: + choice_chunk = chunk.model_copy(update={"choices": [choice]}) + choice_response = choices.get(choice.index) + if choice_response is None: + choice_response = init_chat_completion(choice_chunk) + choices[choice.index] = choice_response + update_chat_completion(choice_response, choice_chunk) if response is None or not done: raise ValueError("Incomplete Chat Completions stream") + response.choices = [choices[index].choices[0] for index in sorted(choices)] return response @@ -220,11 +229,11 @@ def build_exchange( *, start_time: datetime, end_time: datetime, -) -> tuple[Endpoint, Exchange]: +) -> Exchange: stream = request.get("stream") is True if endpoint == "chat_completions": response = _chat_response(body, stream=stream) - return endpoint, ChatCompletionsExchange( + return ChatCompletionsExchange( request=ChatCompletionsRequest(**request), response=response, model=_model(request, response), @@ -233,7 +242,7 @@ def build_exchange( ) if endpoint == "completions": response = _completion_response(body, stream=stream) - return endpoint, CompletionsExchange( + return CompletionsExchange( request=CompletionsRequest(**request), response=response, model=_model(request, response), @@ -242,7 +251,7 @@ def build_exchange( ) if endpoint == "responses": response = _responses_response(body, stream=stream) - return endpoint, ResponsesExchange( + return ResponsesExchange( request=ResponsesRequest(**request), response=response, model=_model(request, response), @@ -251,7 +260,7 @@ def build_exchange( ) if endpoint == "messages": response = _messages_response(body, stream=stream) - return endpoint, MessagesExchange( + return MessagesExchange( request=MessagesRequest(**request), response=response, model=_model(request, response), diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index 62b81ea77..c61c14333 100644 --- a/src/art/trajectories/_tokenize.py +++ b/src/art/trajectories/_tokenize.py @@ -92,24 +92,61 @@ def _token_id(value: object) -> int | None: return None -def _pairs(values: object) -> tuple[list[int], list[float]]: +def _exact_token_ids(values: object, *, field: str) -> list[int] | None: + if values is None: + return None + if not isinstance(values, list): + raise ValueError(f"{field} exact token metadata must be a list") + token_ids: list[int] = [] + for value in values: + token_id = _token_id(value) + if token_id is None: + raise ValueError(f"{field} contains an invalid exact token ID") + token_ids.append(token_id) + return token_ids + + +def _pair_token_id(data: dict[str, Any], *, required: bool, field: str) -> int | None: + if "token_id" in data: + token_id = _token_id(data["token_id"]) + if token_id is None: + raise ValueError(f"{field} contains an invalid exact token ID") + return token_id + raw_token = data.get("token") + token_id = _token_id(raw_token) + if token_id is not None: + return token_id + if isinstance(raw_token, str) and raw_token.startswith("token_id:"): + raise ValueError(f"{field} contains an invalid exact token ID") + if required: + raise ValueError(f"{field} is missing an exact token ID") + return None + + +def _pairs( + values: object, *, require_token_ids: bool = False, field: str = "token pairs" +) -> tuple[list[int], list[float]]: if not isinstance(values, list): + if require_token_ids: + raise ValueError(f"{field} exact token metadata must be a list") return [], [] token_ids: list[int] = [] logprobs: list[float] = [] + complete = True for value in values: data = _dump(value) - token_id = _token_id(data.get("token_id")) - if token_id is None: - token_id = _token_id(data.get("token")) + token_id = _pair_token_id(data, required=require_token_ids, field=field) if token_id is None: - return [], [] + complete = False + continue logprob = data.get("logprob") token_ids.append(token_id) logprobs.append( - float(logprob) if isinstance(logprob, (int, float)) else math.nan + float(logprob) + if isinstance(logprob, (int, float)) and not isinstance(logprob, bool) + else math.nan ) - return token_ids, logprobs + return (token_ids, logprobs) if complete else ([], []) def _logprob_values(values: object) -> list[float]: @@ -118,7 +155,7 @@ def _logprob_values(values: object) -> list[float]: result: list[float] = [] for value in values: logprob = _dump(value).get("logprob") - if not isinstance(logprob, (int, float)): + if not isinstance(logprob, (int, float)) or isinstance(logprob, bool): return [] result.append(float(logprob)) return result @@ -128,24 +165,20 @@ def _chat_choice_tokens( choice: Choice, response_data: dict[str, Any] ) -> tuple[list[int] | None, list[int], list[float]]: choice_data = _dump(choice) - prompt = choice_data.get("prompt_token_ids") or response_data.get( - "prompt_token_ids" - ) - prompt_ids = ( - [token for value in prompt if (token := _token_id(value)) is not None] - if isinstance(prompt, list) - else None + 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") + token_ids = _exact_token_ids( + choice_data.get("token_ids"), + field="Chat Completions token_ids", ) - token_ids = [ - token - for value in choice_data.get("token_ids") or [] - if (token := _token_id(value)) is not None - ] logprob_values = None if choice.logprobs is not None: logprob_values = choice.logprobs.content or choice.logprobs.refusal values = list(logprob_values or []) - pair_ids, logprobs = _pairs(values) + pair_ids, logprobs = _pairs(values, field="Chat Completions logprobs") + token_ids = token_ids or [] if token_ids and pair_ids and token_ids != pair_ids: raise ValueError("Response token IDs disagree with choice logprobs") return ( @@ -171,22 +204,30 @@ def _completion_tokens( choice = response.choices[0] response_data = _dump(response) choice_data = _dump(choice) - prompt = choice_data.get("prompt_token_ids") or response_data.get( - "prompt_token_ids" + 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") + token_ids = _exact_token_ids( + choice_data.get("token_ids"), field="Completions token_ids" ) - prompt_ids = ( - [token for value in prompt if (token := _token_id(value)) is not None] - if isinstance(prompt, list) - else None - ) - token_ids = [ - token - for value in choice_data.get("token_ids") or [] - if (token := _token_id(value)) is not None - ] + token_ids = token_ids or [] logprobs = _dump(choice.logprobs) tokens = logprobs.get("tokens") or [] - pair_ids = [token for value in tokens if (token := _token_id(value)) is not None] + pair_ids: list[int] = [] + complete_pairs = True + for value in tokens: + token = _token_id(value) + if token is None: + if isinstance(value, str) and value.startswith("token_id:"): + raise ValueError( + "Completions logprobs contain an invalid exact token ID" + ) + complete_pairs = False + else: + pair_ids.append(token) + if not complete_pairs: + pair_ids = [] pair_logprobs = [ float(value) if isinstance(value, (int, float)) else math.nan for value in logprobs.get("token_logprobs") or [] @@ -201,27 +242,46 @@ def _completion_tokens( def _responses_tokens(response: Response) -> tuple[None, list[int], list[float]]: data = _dump(response) - token_ids, logprobs = _pairs(data.get("raw_output_tokens")) - if 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", + ) return None, token_ids, logprobs + token_ids: list[int] = [] + logprobs: list[float] = [] + saw_rendered_output = False + complete = True for output in data.get("output") or []: - for content in _dump(output).get("content") or []: - values = _dump(content).get("logprobs") - token_ids, logprobs = _pairs(values) - if token_ids: - return None, token_ids, logprobs - if logprobs := _logprob_values(values): - return None, [], logprobs + output_data = _dump(output) + if output_data.get("type") != "message": + complete = False + continue + for content in output_data.get("content") or []: + content_data = _dump(content) + text = content_data.get("text") or content_data.get("refusal") + if not isinstance(text, str) or not text: + continue + saw_rendered_output = True + pair_ids, pair_logprobs = _pairs( + content_data.get("logprobs"), field="Responses content logprobs" + ) + if not pair_ids: + complete = False + continue + token_ids.extend(pair_ids) + logprobs.extend(pair_logprobs) + if saw_rendered_output and complete: + return None, token_ids, logprobs return None, [], [] def _messages_tokens(response: Message) -> tuple[None, list[int], list[float]]: data = _dump(response) - token_ids = [ - token - for value in data.get("token_ids") or [] - if (token := _token_id(value)) is not None - ] + token_ids = ( + _exact_token_ids(data.get("token_ids"), field="Messages token_ids") or [] + ) logprobs = [ float(value) if isinstance(value, (int, float)) else math.nan for value in data.get("logprobs") or [] @@ -416,7 +476,10 @@ def _anthropic_messages(request: dict[str, Any]) -> list[dict[str, Any]]: def _responses_messages(request: dict[str, Any]) -> list[dict[str, Any]]: messages: list[dict[str, Any]] = [] - if instructions := request.get("instructions"): + instructions = request.get("instructions") + if instructions is not None and not isinstance(instructions, str): + raise ValueError("Responses instructions must be text") + if instructions: messages.append({"role": "system", "content": instructions}) value = request.get("input") if isinstance(value, str): @@ -424,16 +487,19 @@ def _responses_messages(request: dict[str, Any]) -> list[dict[str, Any]]: elif isinstance(value, list): for item in value: if not isinstance(item, dict): - continue - if item.get("type") == "function_call_output": + raise ValueError("Responses input items must be JSON objects") + kind = item.get("type") + if kind == "function_call_output": messages.append( { "role": "tool", "tool_call_id": item.get("call_id"), - "content": item.get("output", ""), + "content": _responses_input_text( + item.get("output", ""), field="function_call_output" + ), } ) - elif item.get("type") == "function_call": + elif kind == "function_call": messages.append( { "role": "assistant", @@ -450,16 +516,69 @@ def _responses_messages(request: dict[str, Any]) -> list[dict[str, Any]]: ], } ) - elif item.get("role"): + elif kind in {None, "message"} and item.get("role"): + if item.get("phase") is not None: + raise ValueError("Unsupported Responses message phase") messages.append( { "role": item["role"], - "content": _content_text(item.get("content")), + "content": _responses_input_text( + item.get("content"), field="message content" + ), } ) + else: + raise ValueError(f"Unsupported Responses input item type: {kind!r}") + elif value is not None: + raise ValueError("Responses input must be text or a list of input items") return messages +def _responses_input_text(content: object, *, field: str) -> str: + if isinstance(content, str): + return content + if not isinstance(content, list): + raise ValueError(f"Responses {field} must contain text") + text = "" + for block in content: + data = _string_dict(block) + if data is None: + raise ValueError(f"Responses {field} blocks must be JSON objects") + kind = data.get("type") + if kind not in {"input_text", "output_text", "refusal", "text"}: + raise ValueError(f"Unsupported Responses content block type: {kind!r}") + value = data.get("refusal" if kind == "refusal" else "text") + if not isinstance(value, str): + raise ValueError(f"Responses {field} blocks must contain text") + text += value + return text + + +def _responses_output_text(content: object) -> str: + if not isinstance(content, list): + raise ValueError("Responses message output content must be a list") + text = "" + for block in content: + data = _string_dict(block) + if data is None: + raise ValueError("Responses output content blocks must be JSON objects") + kind = data.get("type") + key = ( + "text" + if kind == "output_text" + else "refusal" + if kind == "refusal" + else None + ) + if key is None: + raise ValueError(f"Unsupported Responses output content type: {kind!r}") + value = data.get(key) + if not isinstance(value, str): + raise ValueError(f"Responses {kind} content must be text") + text += value + return text + + def _openai_tools(tools: object, *, dialect: str) -> object: if not isinstance(tools, list) or dialect == "chat": return tools @@ -526,12 +645,18 @@ def _response_message( return _anthropic_messages(request)[0] if isinstance(exchange, ResponsesExchange): data = exchange.response.model_dump(mode="python") - content = [] + content = "" tool_calls = [] - for item in data.get("output") or []: - if item.get("type") == "message": - content.extend(item.get("content") or []) - elif item.get("type") == "function_call": + for raw_item in data.get("output") or []: + item = _string_dict(raw_item) + if item is None: + raise ValueError("Responses output items must be JSON objects") + kind = item.get("type") + if kind == "message": + if item.get("phase") is not None: + raise ValueError("Unsupported Responses message phase") + content += _responses_output_text(item.get("content")) + elif kind == "function_call": tool_calls.append( { "id": item.get("call_id"), @@ -542,9 +667,11 @@ def _response_message( }, } ) + else: + raise ValueError(f"Unsupported Responses output item type: {kind!r}") message: dict[str, Any] = { "role": "assistant", - "content": _content_text(content), + "content": content, } if tool_calls: message["tool_calls"] = tool_calls @@ -627,11 +754,13 @@ def _visible_logprobs(exchange: Exchange) -> list[tuple[str, float]]: for entry in entries: data = _dump(entry) raw_bytes = data.get("bytes") - text = ( - bytes(raw_bytes).decode("utf-8") - if isinstance(raw_bytes, list) - else data.get("token") - ) + if isinstance(raw_bytes, list): + try: + text = bytes(raw_bytes).decode("utf-8") + except (TypeError, ValueError, UnicodeDecodeError): + return [] + else: + text = data.get("token") logprob = data.get("logprob") if isinstance(text, str) and isinstance(logprob, (int, float)): values.append((text, float(logprob))) @@ -745,9 +874,23 @@ def tokenize_one( token_ids: list[int] = [] logprobs: list[float] = [] assistant_mask: list[bool] = [] - response_histories: dict[str, list[dict[str, Any]]] = {} + response_histories: dict[str, tuple[list[dict[str, Any]], ResponsesExchange]] = {} for exchange in exchanges: + if isinstance(exchange, CompletionsExchange): + prompt = exchange.request.get("prompt") + if isinstance(prompt, list) and not all( + isinstance(item, int) and not isinstance(item, bool) for item in prompt + ): + raise ValueError( + "Trajectory tokenization does not support batched Completions prompts" + ) + if not isinstance(prompt, (str, list)): + raise ValueError("Completions prompt must be text or one token ID list") + if exchange.request.get("echo") is True: + raise ValueError( + "Trajectory tokenization does not support Completions echo=True" + ) messages_override = None if isinstance(exchange, ResponsesExchange): request = exchange.request @@ -758,14 +901,13 @@ def tokenize_one( raise ValueError( "Responses exchange refers to a previous response outside this trajectory" ) + previous_messages, previous_exchange = response_histories[previous] messages_override = [ - *response_histories[previous], + *previous_messages, + _response_message(previous_exchange), *messages_override, ] - response_histories[exchange.response.id] = [ - *messages_override, - _response_message(exchange), - ] + response_histories[exchange.response.id] = (messages_override, exchange) prompt, completion, completion_logprobs = _exchange_tokens(exchange) if prompt is None: if tokenizer is None: diff --git a/tests/unit/trajectories/test_capture.py b/tests/unit/trajectories/test_capture.py index 3cc3ecbc2..ef2cf8602 100644 --- a/tests/unit/trajectories/test_capture.py +++ b/tests/unit/trajectories/test_capture.py @@ -1,10 +1,10 @@ from __future__ import annotations import asyncio -from collections.abc import AsyncIterator +from collections.abc import AsyncGenerator, AsyncIterator, Generator from datetime import datetime, timedelta import json -from typing import Any +from typing import Any, cast import aiohttp from aiohttp import web @@ -18,7 +18,8 @@ import art from art.trajectories import ChatCompletionsExchange, MessagesExchange -from art.trajectories._protocols import Endpoint, build_exchange +from art.trajectories._capture.core import begin, reset +from art.trajectories._protocols import Endpoint, build_exchange, endpoint_for_url CHAT: dict[str, Any] = { "id": "chatcmpl-1", @@ -111,7 +112,7 @@ @pytest_asyncio.fixture async def endpoint_server(unused_tcp_port: int) -> AsyncIterator[str]: - async def handler(request: web.Request) -> web.Response: + async def handler(request: web.Request) -> web.StreamResponse: request_body = await request.json() if request_body.get("fail"): return web.json_response({"error": "failed"}, status=400) @@ -120,6 +121,18 @@ async def handler(request: web.Request) -> web.Response: body=_sse([(None, {"type": "incomplete"})]), content_type="text/event-stream", ) + if request_body.get("early_close"): + response = web.StreamResponse( + status=200, headers={"Content-Type": "application/json"} + ) + await response.prepare(request) + await response.write(json.dumps(CHAT).encode()) + await asyncio.sleep(0.05) + try: + await response.write(b" ") + except ConnectionResetError: + pass + return response bodies = { "/v1/chat/completions": CHAT, "/v1/completions": COMPLETION, @@ -265,6 +278,76 @@ async def test_failed_and_incomplete_calls_are_excluded(endpoint_server: str) -> assert not trajectory.exchanges +async def test_abandoned_transport_streams_are_excluded(endpoint_server: str) -> None: + body = { + "model": "test/model", + "messages": [], + "early_close": True, + } + + async def abandon_httpx() -> None: + async with httpx.AsyncClient() as client: + async with client.stream( + "POST", f"{endpoint_server}/chat/completions", json=body + ) as response: + iterator = cast(AsyncGenerator[bytes, None], response.aiter_bytes()) + await anext(iterator) + await iterator.aclose() + + async def abandon_aiohttp() -> None: + async with aiohttp.ClientSession() as session: + async with session.post( + f"{endpoint_server}/chat/completions", json=body + ) as response: + iterator = cast( + AsyncGenerator[bytes, None], response.content.iter_any() + ) + await anext(iterator) + await iterator.aclose() + + def abandon_requests() -> None: + with requests.post( + f"{endpoint_server}/chat/completions", + json=body, + stream=True, + timeout=5, + ) as response: + iterator = cast( + Generator[bytes, None, None], + response.iter_content(chunk_size=None), + ) + next(iterator) + iterator.close() + + with art.Trajectory() as trajectory: + await abandon_httpx() + await abandon_aiohttp() + await asyncio.to_thread(abandon_requests) + + assert not trajectory.exchanges + + +def test_capture_snapshots_nested_request_values() -> None: + request = { + "model": "test/model", + "messages": [{"role": "user", "content": "before"}], + } + with art.Trajectory() as trajectory: + state, token = begin( + "POST", "https://example.test/v1/chat/completions", request + ) + reset(token) + assert state is not None + request["messages"][0]["content"] = "after" + state.status_code = 200 + state.add(json.dumps(CHAT).encode()) + state.finish() + + assert trajectory.exchanges.chat_completions[0].request["messages"] == [ + {"role": "user", "content": "before"} + ] + + def test_all_protocols_reconstruct_typed_responses() -> None: now = datetime.now() values: list[tuple[Endpoint, dict[str, Any], dict[str, Any]]] = [ @@ -274,20 +357,42 @@ def test_all_protocols_reconstruct_typed_responses() -> None: ("messages", {"model": "request-model", "messages": []}, MESSAGE), ] for endpoint, request, response in values: - name, exchange = build_exchange( + exchange = build_exchange( endpoint, request, json.dumps(response).encode(), start_time=now, end_time=now + timedelta(seconds=1), ) - assert name == endpoint assert exchange.end_time > exchange.start_time expected = request.get("model", "test/model") assert exchange.model == expected assert exchange.model_dump(mode="json")["request"] == request +@pytest.mark.parametrize( + ("url", "endpoint"), + [ + ( + "https://azure.test/openai/deployments/model/chat/completions?api-version=1", + "chat_completions", + ), + ( + "https://generativelanguage.googleapis.com/v1beta/openai/chat/completions", + "chat_completions", + ), + ("https://gateway.test/completions", "completions"), + ("https://gateway.test/responses", "responses"), + ("https://gateway.test/messages", "messages"), + ("https://gateway.test/not-messages", None), + ], +) +def test_endpoint_detection_accepts_compatible_gateway_paths( + url: str, endpoint: Endpoint | None +) -> None: + assert endpoint_for_url(url) == endpoint + + def _sse(events: list[tuple[str | None, dict[str, Any] | str]]) -> bytes: return "".join( f"{f'event: {name}\n' if name else ''}data: " @@ -389,7 +494,7 @@ def test_all_streaming_protocols_reconstruct_final_responses() -> None: ), ] for endpoint, request, body in values: - _, exchange = build_exchange( + exchange = build_exchange( endpoint, request, body, @@ -401,10 +506,55 @@ def test_all_streaming_protocols_reconstruct_final_responses() -> None: content = exchange.response.content[0] assert isinstance(content, TextBlock) assert content.text == "hello" + assert getattr(exchange.response, "token_ids") == [2] + assert getattr(exchange.response, "logprobs") == [-0.2] + + +def test_streaming_chat_choices_are_accumulated_by_index() -> None: + now = datetime.now() + + def chunk(index: int, content: str) -> dict[str, Any]: + return { + "id": "chatcmpl-1", + "object": "chat.completion.chunk", + "created": 1, + "model": "test/model", + "choices": [ + { + "index": index, + "delta": {"role": "assistant", "content": content}, + "finish_reason": None, + "logprobs": None, + } + ], + } + + exchange = build_exchange( + "chat_completions", + {"model": "test/model", "messages": [], "stream": True, "n": 2}, + _sse( + [ + (None, chunk(1, "b")), + (None, chunk(0, "a")), + (None, chunk(1, "d")), + (None, chunk(0, "c")), + (None, "[DONE]"), + ] + ), + start_time=now, + end_time=now + timedelta(seconds=1), + ) + + assert isinstance(exchange, ChatCompletionsExchange) + assert [choice.index for choice in exchange.response.choices] == [0, 1] + assert [choice.message.content for choice in exchange.response.choices] == [ + "ac", + "bd", + ] def test_trajectory_rejects_mixed_representations() -> None: - _, exchange = build_exchange( + exchange = build_exchange( "chat_completions", {"model": "test/model", "messages": []}, json.dumps(CHAT).encode(), diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index a83adbacc..63129a08c 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -6,6 +6,7 @@ from typing import Any from anthropic.types import Message +from openai.types import Completion from openai.types.chat import ChatCompletion from openai.types.chat.chat_completion_token_logprob import ChatCompletionTokenLogprob from openai.types.responses import Response @@ -15,6 +16,8 @@ from art.trajectories import ( ChatCompletionsExchange, ChatCompletionsRequest, + CompletionsExchange, + CompletionsRequest, MessagesExchange, MessagesRequest, ResponsesExchange, @@ -71,6 +74,46 @@ def _chat_exchange( ) +def _completion_exchange( + *, + prompt: str | list[str] | list[int] | list[list[int]] = "question", + echo: bool = False, +) -> CompletionsExchange: + response = Completion.model_validate( + { + "id": "completion-1", + "object": "text_completion", + "created": 0, + "model": "test/model", + "choices": [ + { + "index": 0, + "finish_reason": "stop", + "text": "answer", + "prompt_token_ids": [1], + "token_ids": [2], + "logprobs": { + "tokens": ["token_id:2"], + "token_logprobs": [-0.2], + "top_logprobs": [{}], + "text_offset": [0], + }, + } + ], + } + ) + request = CompletionsRequest(model="test/model", prompt="question", echo=echo) + request["prompt"] = prompt + start = datetime(2026, 1, 1) + return CompletionsExchange( + request=request, + response=response, + model="test/model", + start_time=start, + end_time=start + timedelta(milliseconds=1), + ) + + def test_exact_tokens_form_one_append_only_history_without_tokenizer() -> None: trajectory = art.Trajectory( exchanges=TrajectoryExchanges( @@ -91,6 +134,76 @@ def test_exact_tokens_form_one_append_only_history_without_tokenizer() -> None: assert tokenized.logprobs[3] == -0.4 +def test_malformed_explicit_exact_token_metadata_fails_closed() -> None: + chat = _chat_exchange([1], [2]) + chat_extra = chat.response.choices[0].model_extra + assert chat_extra is not None + chat_extra["prompt_token_ids"] = [1, "invalid"] + + completion = _completion_exchange() + completion_extra = completion.response.choices[0].model_extra + assert completion_extra is not None + completion_extra["token_ids"] = [2, "invalid"] + + 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"}] + + 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( + model="test/model", + messages=[{"role": "user", "content": "question"}], + max_tokens=16, + ), + response=message_response, + model="test/model", + start_time=start, + end_time=start + timedelta(milliseconds=1), + ) + + trajectories = [ + art.Trajectory(exchanges=TrajectoryExchanges(chat_completions=[chat])), + art.Trajectory(exchanges=TrajectoryExchanges(completions=[completion])), + art.Trajectory(exchanges=TrajectoryExchanges(responses=[response])), + art.Trajectory(exchanges=TrajectoryExchanges(messages=[message])), + ] + for trajectory in trajectories: + with pytest.raises(ValueError, match="exact token"): + art.tokenize_trajectory(trajectory, base_model="base/model") + + +@pytest.mark.parametrize( + "exchange", + [ + _completion_exchange(prompt=["batched"]), + _completion_exchange(prompt=[[1, 2]]), + _completion_exchange(echo=True), + ], +) +def test_completions_reject_batch_prompts_and_echo( + exchange: CompletionsExchange, +) -> None: + with pytest.raises(ValueError, match="batched Completions|echo=True"): + art.tokenize_trajectory( + art.Trajectory(exchanges=TrajectoryExchanges(completions=[exchange])) + ) + + def test_branching_and_multiple_models_require_explicit_resolution() -> None: branching = art.Trajectory( exchanges=TrajectoryExchanges( @@ -327,6 +440,44 @@ def __call__(self, text: str, **kwargs: object) -> SimpleNamespace: assert math.isnan(result.logprobs[2]) +def test_undecodable_visible_token_bytes_fall_back_to_nan( + monkeypatch: pytest.MonkeyPatch, +) -> None: + exchange = _chat_exchange([], []) + logprobs = exchange.response.choices[0].logprobs + assert logprobs is not None + exchange.response.choices[0].logprobs = logprobs.model_copy( + update={ + "content": [ + ChatCompletionTokenLogprob( + token="ordinary-token", + logprob=-0.7, + bytes=[0xF0], + top_logprobs=[], + ) + ] + } + ) + + class Tokenizer: + def apply_chat_template( + self, messages: list[dict[str, Any]], **kwargs: object + ) -> list[int]: + del kwargs + return [10, 11] if messages[-1]["role"] == "assistant" else [10] + + monkeypatch.setattr( + "art.trajectories._tokenize._load_tokenizer", lambda _config: Tokenizer() + ) + result = art.tokenize_trajectory( + art.Trajectory(exchanges=TrajectoryExchanges(chat_completions=[exchange])), + base_model="base/model", + ) + + assert result.token_ids == [10, 11] + assert math.isnan(result.logprobs[1]) + + def test_json_round_trip_preserves_exchange_types() -> None: original = art.Trajectory( exchanges=TrajectoryExchanges(chat_completions=[_chat_exchange([1], [2])]) @@ -384,6 +535,158 @@ 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) + + def entry(token: str, token_id: int | None, logprob: float) -> dict[str, Any]: + return { + "token": token, + "logprob": logprob, + "bytes": list(("a" if token_id == 11 else "b").encode()), + "top_logprobs": [], + **({"token_id": token_id} if token_id is not None else {}), + } + + data["output"][0]["content"] = [ + { + "type": "output_text", + "text": "a", + "annotations": [], + "logprobs": [entry("token_id:11", 11, -0.1)], + }, + { + "type": "output_text", + "text": "b", + "annotations": [], + "logprobs": [ + entry( + "token_id:12" if exact_second else "b", + 12 if exact_second else None, + -0.2, + ) + ], + }, + ] + exchange.response = Response.model_validate(data) + return exchange + + +def test_responses_aggregates_complete_exact_pairs_across_content_blocks( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class Tokenizer: + def apply_chat_template( + self, messages: list[dict[str, Any]], **kwargs: object + ) -> list[int]: + del messages, kwargs + return [10] + + 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)] + ) + ), + base_model="base/model", + ) + + assert result.token_ids == [10, 11, 12] + assert result.logprobs[1:] == [-0.1, -0.2] + + +def test_responses_does_not_use_partial_exact_content_pairs( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class Tokenizer: + 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] + + def __call__(self, text: str, **kwargs: object) -> SimpleNamespace: + del kwargs + return SimpleNamespace( + input_ids=[11 if text in {"a", "token_id:11"} else 12] + ) + + 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)] + ) + ), + base_model="base/model", + ) + + assert result.token_ids == [10, 11, 12] + assert result.logprobs[1:] == [-0.1, -0.2] + + +def test_responses_rejects_only_unrenderable_prompt_history( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class Tokenizer: + def apply_chat_template( + self, messages: list[dict[str, Any]], **kwargs: object + ) -> list[int]: + del messages, kwargs + return [10] + + monkeypatch.setattr( + "art.trajectories._tokenize._load_tokenizer", lambda _config: Tokenizer() + ) + request_reasoning = _response_exchange("request-reasoning", 2) + request_reasoning.request["input"] = [ + {"id": "reasoning-1", "summary": [], "type": "reasoning"} + ] + + response_reasoning = _response_exchange("response-reasoning", 2) + data = response_reasoning.response.model_dump(mode="python") + data["output"] = [{"id": "reasoning-2", "summary": [], "type": "reasoning"}] + response_reasoning.response = Response.model_validate(data) + + with pytest.raises(ValueError, match="Unsupported Responses input"): + art.tokenize_trajectory( + art.Trajectory( + exchanges=TrajectoryExchanges(responses=[request_reasoning]) + ), + base_model="base/model", + ) + + single = art.Trajectory( + exchanges=TrajectoryExchanges(responses=[response_reasoning]) + ) + assert art.tokenize_trajectory(single, base_model="base/model").token_ids == [ + 10, + 2, + ] + + continuation = _response_exchange( + "continuation", + 3, + previous_response_id=response_reasoning.response.id, + offset=1, + ) + with pytest.raises(ValueError, match="Unsupported Responses output"): + art.tokenize_trajectory( + art.Trajectory( + exchanges=TrajectoryExchanges( + responses=[response_reasoning, continuation] + ) + ), + base_model="base/model", + ) + + def test_responses_previous_response_id_resolves_local_history( monkeypatch: pytest.MonkeyPatch, ) -> None: From 67f9ff4d44f967ea06c17aced7bd61ccc9fb4c78 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 15 Jul 2026 13:44:26 +0000 Subject: [PATCH 06/17] api: derive trajectory exchange models --- pyproject.toml | 2 +- src/art/trajectories/__init__.py | 28 ++++++++++++++++++++---- src/art/trajectories/_protocols.py | 13 ----------- tests/unit/trajectories/test_capture.py | 14 +++++++++++- tests/unit/trajectories/test_tokenize.py | 5 ----- uv.lock | 2 +- 6 files changed, 39 insertions(+), 25 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 371aeacb4..253064f16 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -6,7 +6,7 @@ readme = "README.md" requires-python = ">=3.12" dependencies = [ "aiohttp>=3.10.0", - "anthropic>=0.75.0", + "anthropic>=0.77.0", "openai>=2.14.0", "requests>=2.32.0", "typer>=0.15.2", diff --git a/src/art/trajectories/__init__.py b/src/art/trajectories/__init__.py index 3fefa13e9..98cbbdecf 100644 --- a/src/art/trajectories/__init__.py +++ b/src/art/trajectories/__init__.py @@ -121,34 +121,54 @@ class MessagesRequest(TypedDict, total=False, extra_items=Any): class ChatCompletionsExchange(pydantic.BaseModel): request: Annotated[ChatCompletionsRequest, pydantic.SkipValidation] response: ChatCompletion - model: str | None start_time: datetime end_time: datetime + @pydantic.computed_field + @property + def model(self) -> str | None: + requested = self.request.get("model") + return requested if isinstance(requested, str) else self.response.model + class CompletionsExchange(pydantic.BaseModel): request: Annotated[CompletionsRequest, pydantic.SkipValidation] response: Completion - model: str | None start_time: datetime end_time: datetime + @pydantic.computed_field + @property + def model(self) -> str | None: + requested = self.request.get("model") + return requested if isinstance(requested, str) else self.response.model + class ResponsesExchange(pydantic.BaseModel): request: Annotated[ResponsesRequest, pydantic.SkipValidation] response: Response - model: str | None start_time: datetime end_time: datetime + @pydantic.computed_field + @property + def model(self) -> str | None: + requested = self.request.get("model") + return requested if isinstance(requested, str) else self.response.model + class MessagesExchange(pydantic.BaseModel): request: Annotated[MessagesRequest, pydantic.SkipValidation] response: AnthropicMessage - model: str | None start_time: datetime end_time: datetime + @pydantic.computed_field + @property + def model(self) -> str | None: + requested = self.request.get("model") + return requested if isinstance(requested, str) else self.response.model + class TrajectoryExchanges(pydantic.BaseModel): chat_completions: list[ChatCompletionsExchange] = pydantic.Field( diff --git a/src/art/trajectories/_protocols.py b/src/art/trajectories/_protocols.py index 41baa1b93..b5b750eac 100644 --- a/src/art/trajectories/_protocols.py +++ b/src/art/trajectories/_protocols.py @@ -1,6 +1,5 @@ from __future__ import annotations -from collections.abc import Mapping from datetime import datetime import json from typing import Any, Literal @@ -31,7 +30,6 @@ Exchange = ( ChatCompletionsExchange | CompletionsExchange | ResponsesExchange | MessagesExchange ) -ResponseModel = ChatCompletion | Completion | Response | Message SSEPayload = dict[str, Any] | Literal["[DONE]"] _ENDPOINTS: dict[str, Endpoint] = { "/chat/completions": "chat_completions", @@ -215,13 +213,6 @@ def _messages_response(body: bytes, *, stream: bool) -> Message: return Message.model_validate(data) -def _model(request: Mapping[str, object], response: ResponseModel) -> str | None: - requested = request.get("model") - if isinstance(requested, str): - return requested - return response.model if isinstance(response.model, str) else None - - def build_exchange( endpoint: Endpoint, request: dict[str, Any], @@ -236,7 +227,6 @@ def build_exchange( return ChatCompletionsExchange( request=ChatCompletionsRequest(**request), response=response, - model=_model(request, response), start_time=start_time, end_time=end_time, ) @@ -245,7 +235,6 @@ def build_exchange( return CompletionsExchange( request=CompletionsRequest(**request), response=response, - model=_model(request, response), start_time=start_time, end_time=end_time, ) @@ -254,7 +243,6 @@ def build_exchange( return ResponsesExchange( request=ResponsesRequest(**request), response=response, - model=_model(request, response), start_time=start_time, end_time=end_time, ) @@ -263,7 +251,6 @@ def build_exchange( return MessagesExchange( request=MessagesRequest(**request), response=response, - model=_model(request, response), start_time=start_time, end_time=end_time, ) diff --git a/tests/unit/trajectories/test_capture.py b/tests/unit/trajectories/test_capture.py index ef2cf8602..4b49307fe 100644 --- a/tests/unit/trajectories/test_capture.py +++ b/tests/unit/trajectories/test_capture.py @@ -367,7 +367,19 @@ def test_all_protocols_reconstruct_typed_responses() -> None: assert exchange.end_time > exchange.start_time expected = request.get("model", "test/model") assert exchange.model == expected - assert exchange.model_dump(mode="json")["request"] == request + dumped = exchange.model_dump(mode="json") + assert dumped["request"] == request + assert dumped["model"] == expected + + # Older serialized exchanges may contain a stale stored model. It is ignored + # in favor of the request, with the response as fallback. + dumped["model"] = "stale/model" + restored = type(exchange).model_validate(dumped) + assert restored.model == expected + + exchange.request["model"] = "updated/model" + assert exchange.model == "updated/model" + assert exchange.model_dump(mode="json")["model"] == "updated/model" @pytest.mark.parametrize( diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index 63129a08c..809e51a6a 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -68,7 +68,6 @@ def _chat_exchange( messages=[{"role": "user", "content": f"turn {offset}"}], ), response=response, - model=model, start_time=start, end_time=start + timedelta(milliseconds=1), ) @@ -108,7 +107,6 @@ def _completion_exchange( return CompletionsExchange( request=request, response=response, - model="test/model", start_time=start, end_time=start + timedelta(milliseconds=1), ) @@ -171,7 +169,6 @@ def test_malformed_explicit_exact_token_metadata_fails_closed() -> None: max_tokens=16, ), response=message_response, - model="test/model", start_time=start, end_time=start + timedelta(milliseconds=1), ) @@ -265,7 +262,6 @@ def test_fallback_uses_template_overrides_and_nan_logprobs( thinking={"type": "enabled", "budget_tokens": 128}, ), response=response, - model="test/model", start_time=start, end_time=start + timedelta(seconds=1), ) @@ -529,7 +525,6 @@ def _response_exchange( return ResponsesExchange( request=request, response=response, - model="test/model", start_time=start, end_time=start + timedelta(milliseconds=1), ) diff --git a/uv.lock b/uv.lock index 21b5bd1a4..156052456 100644 --- a/uv.lock +++ b/uv.lock @@ -4868,7 +4868,7 @@ dev = [ requires-dist = [ { name = "accelerate", marker = "extra == 'backend'", specifier = "==1.7.0" }, { name = "aiohttp", specifier = ">=3.10.0" }, - { name = "anthropic", specifier = ">=0.75.0" }, + { 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" }, From 85a430efe04b01426ea70d7ae56cf51386be3b07 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 15 Jul 2026 14:14:23 +0000 Subject: [PATCH 07/17] fix: initialize trajectory groups once --- src/art/trajectories/__init__.py | 8 +------- src/art/trajectories/_compat.py | 17 ++++++++--------- tests/unit/trajectories/test_capture.py | 22 +++++++++++++++------- 3 files changed, 24 insertions(+), 23 deletions(-) diff --git a/src/art/trajectories/__init__.py b/src/art/trajectories/__init__.py index 98cbbdecf..21c0bd032 100644 --- a/src/art/trajectories/__init__.py +++ b/src/art/trajectories/__init__.py @@ -334,15 +334,9 @@ def __init__( ) -> None: from ._compat import init_trajectory_group - items = list(trajectories) - sync_items = [ - item for item in items if isinstance(item, (Trajectory, BaseException)) - ] - if len(sync_items) != len(items): - raise TypeError("TrajectoryGroup cannot initialize from awaitables") init_trajectory_group( self, - sync_items, + trajectories, exceptions=exceptions, metadata=metadata, metrics=metrics, diff --git a/src/art/trajectories/_compat.py b/src/art/trajectories/_compat.py index 13785d2ac..ecad617be 100644 --- a/src/art/trajectories/_compat.py +++ b/src/art/trajectories/_compat.py @@ -13,6 +13,8 @@ from ..types import Message, Messages, MessagesAndChoices from . import MetadataValue, PydanticException, Trajectory, TrajectoryGroup +_PREPARED_TRAJECTORIES = "_art_prepared_trajectories" + def exception_model( exception: BaseException | PydanticException, @@ -112,26 +114,23 @@ def new_trajectory_group( if len(sync_items) != len(items): raise TypeError("TrajectoryGroup items must be trajectories or exceptions") group = object.__new__(cls) - group.__init__( - sync_items, - exceptions=exceptions, - metadata=metadata, - metrics=metrics, - logs=logs, - ) + object.__setattr__(group, _PREPARED_TRAJECTORIES, sync_items) return group def init_trajectory_group( group: TrajectoryGroup, - trajectories: Iterable[Trajectory | BaseException], + trajectories: Iterable[Trajectory | BaseException | Awaitable[Trajectory]], *, exceptions: Iterable[BaseException | PydanticException], metadata: dict[str, MetadataValue] | None, metrics: dict[str, float | int | bool] | None, logs: list[str] | None, ) -> None: - items = list(trajectories) + prepared = group.__dict__.pop(_PREPARED_TRAJECTORIES, None) + items = prepared if isinstance(prepared, list) else list(trajectories) + if not all(isinstance(item, (Trajectory, BaseException)) for item in items): + raise TypeError("TrajectoryGroup cannot initialize from awaitables") normalized_trajectories = [ item if isinstance(item, Trajectory) else Trajectory.model_validate(item) for item in items diff --git a/tests/unit/trajectories/test_capture.py b/tests/unit/trajectories/test_capture.py index 4b49307fe..49cd8d01c 100644 --- a/tests/unit/trajectories/test_capture.py +++ b/tests/unit/trajectories/test_capture.py @@ -5,6 +5,7 @@ from datetime import datetime, timedelta import json from typing import Any, cast +from unittest.mock import Mock import aiohttp from aiohttp import web @@ -17,7 +18,7 @@ import requests import art -from art.trajectories import ChatCompletionsExchange, MessagesExchange +from art.trajectories import ChatCompletionsExchange, MessagesExchange, _compat from art.trajectories._capture.core import begin, reset from art.trajectories._protocols import Endpoint, build_exchange, endpoint_for_url @@ -200,6 +201,19 @@ async def failed() -> art.Trajectory: assert result.exceptions[0].message == "boom" +def test_sync_group_generator_initializes_once( + monkeypatch: pytest.MonkeyPatch, +) -> None: + initializer = Mock(wraps=_compat.init_trajectory_group) + monkeypatch.setattr(_compat, "init_trajectory_group", initializer) + + trajectory = art.Trajectory(reward=1) + group = art.TrajectoryGroup(item for item in [trajectory]) + + assert group.trajectories == [trajectory] + assert initializer.call_count == 1 + + async def test_httpx_requests_and_aiohttp_capture_once(endpoint_server: str) -> None: body = {"model": "test/model", "messages": [{"role": "user", "content": "hi"}]} @@ -371,12 +385,6 @@ def test_all_protocols_reconstruct_typed_responses() -> None: assert dumped["request"] == request assert dumped["model"] == expected - # Older serialized exchanges may contain a stale stored model. It is ignored - # in favor of the request, with the response as fallback. - dumped["model"] = "stale/model" - restored = type(exchange).model_validate(dumped) - assert restored.model == expected - exchange.request["model"] = "updated/model" assert exchange.model == "updated/model" assert exchange.model_dump(mode="json")["model"] == "updated/model" From 0ec473a4c257c2ce63d62d0c1619ae05fc1c5fbb Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 15 Jul 2026 14:14:23 +0000 Subject: [PATCH 08/17] fix: require exchange training logprobs --- src/art/preprocessing/tokenize.py | 11 ++++++ tests/unit/trajectories/test_tokenize.py | 44 +++++++++++++++++++++++- 2 files changed, 54 insertions(+), 1 deletion(-) diff --git a/src/art/preprocessing/tokenize.py b/src/art/preprocessing/tokenize.py index 6785e2ffc..cfa2d1f60 100644 --- a/src/art/preprocessing/tokenize.py +++ b/src/art/preprocessing/tokenize.py @@ -536,6 +536,17 @@ def tokenize_trajectory_groups( chat_template_kwargs=chat_template_kwargs, tokenizer_instance=_as_tokenizer(tokenizer), ) + if not allow_training_without_logprobs and any( + trainable and math.isnan(logprob) + for trainable, logprob in zip( + exchange_result.assistant_mask, + exchange_result.logprobs, + strict=True, + ) + ): + raise RuntimeError( + "Exchange trajectory is missing logprobs for trainable tokens" + ) choice_offsets = [ index for index, trainable in enumerate(exchange_result.assistant_mask) diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index 809e51a6a..21790037b 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -744,7 +744,7 @@ def decode(self, token_id: int) -> str: # This path only calls decode; the minimal test double is intentional. Tokenizer(), # type: ignore[arg-type, ty:invalid-argument-type] [group], - allow_training_without_logprobs=True, + allow_training_without_logprobs=False, scale_rewards=False, shuffle_group_trajectories=False, ) @@ -752,3 +752,45 @@ def decode(self, token_id: int) -> str: assert [result.token_ids for result in results] == [[1, 2], [1, 3]] assert [result.assistant_mask for result in results] == [[0, 1], [0, 1]] + + +def test_exchange_training_requires_logprobs_unless_allowed() -> None: + from art.preprocessing.tokenize import TokenizedResult, tokenize_trajectory_groups + + class Tokenizer: + name_or_path = "test/model" + + def decode(self, token_id: int) -> str: + return str(token_id) + + missing = _chat_exchange([1], [2]) + missing.response.choices[0].logprobs = None + group = art.TrajectoryGroup( + [ + art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[missing]), reward=1 + ), + art.Trajectory( + exchanges=TrajectoryExchanges( + chat_completions=[_chat_exchange([1], [3])] + ), + reward=0, + ), + ] + ) + + def tokenize(*, allow_missing: bool) -> list[TokenizedResult]: + return list( + tokenize_trajectory_groups( + # This exact-token path only calls decode. + Tokenizer(), # type: ignore[arg-type, ty:invalid-argument-type] + [group], + allow_training_without_logprobs=allow_missing, + scale_rewards=False, + shuffle_group_trajectories=False, + ) + ) + + with pytest.raises(RuntimeError, match="missing logprobs"): + tokenize(allow_missing=False) + assert len(tokenize(allow_missing=True)) == 2 From 9f14507308cabb7bf21c46278c170b4c01ba1a13 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 15 Jul 2026 14:14:23 +0000 Subject: [PATCH 09/17] feat: train tinker from exchange trajectories --- src/art/tinker_native/backend.py | 3 + src/art/tinker_native/data.py | 80 +++++++++- .../integration/test_tinker_native_backend.py | 18 ++- tests/unit/test_tinker_native_exchanges.py | 144 ++++++++++++++++++ 4 files changed, 230 insertions(+), 15 deletions(-) create mode 100644 tests/unit/test_tinker_native_exchanges.py diff --git a/src/art/tinker_native/backend.py b/src/art/tinker_native/backend.py index 0ae9093b1..145477046 100644 --- a/src/art/tinker_native/backend.py +++ b/src/art/tinker_native/backend.py @@ -331,6 +331,7 @@ async def train( state.renderer, state.tokenizer, normalize_advantages, + base_model=model.base_model, ) metrics: dict[str, float] = { @@ -591,6 +592,8 @@ async def chat_completions(body: CompletionCreateParams) -> ChatCompletion: ) ] ), + prompt_token_ids=prompt_tokens, + token_ids=list(sequence.tokens), ) ) diff --git a/src/art/tinker_native/data.py b/src/art/tinker_native/data.py index 48347660e..bf2f9a5fa 100644 --- a/src/art/tinker_native/data.py +++ b/src/art/tinker_native/data.py @@ -1,6 +1,7 @@ from __future__ import annotations import json +import math from typing import Any, Iterable, cast from openai.types.chat.chat_completion import Choice @@ -8,7 +9,14 @@ from tinker_cookbook import renderers import torch -from ..trajectories import History, Trajectory, TrajectoryGroup, get_messages +from ..trajectories import ( + History, + TokenizedTrajectory, + Trajectory, + TrajectoryGroup, + get_messages, + tokenize_trajectory, +) from ..types import MessagesAndChoices @@ -132,6 +140,8 @@ def trajectory_groups_to_datums( renderer: Any, tokenizer: Any, normalize_advantages: bool = True, + *, + base_model: str | None = None, ) -> list[tinker.Datum]: datums: list[tinker.Datum] = [] @@ -139,7 +149,7 @@ def trajectory_groups_to_datums( if not group.trajectories: continue for trajectory in group.trajectories: - if not _trajectory_has_choice(trajectory): + if not trajectory.exchanges and not _trajectory_has_choice(trajectory): raise ValueError( "Trajectory is missing a Choice object. Training requires at least one Choice " "to compute logprobs. Ensure your rollout includes an OpenAI Choice in " @@ -151,6 +161,13 @@ def trajectory_groups_to_datums( if all(advantage == 0.0 for advantage in advantages): 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 + ) + if datum is not None: + datums.append(datum) + continue for history in iter_trajectory_histories(trajectory): datum = history_to_datum(history, advantage, renderer, tokenizer) if datum is not None: @@ -255,24 +272,73 @@ def build_datum( padded_advantages = [0.0] * ob_len + [advantage] * len(completion_tokens) action_mask = [0.0] * ob_len + [1.0] * len(completion_tokens) + return _build_datum( + input_tokens, + target_tokens, + padded_logprobs, + padded_advantages, + action_mask, + ) + + +def _tokenized_trajectory_to_datum( + tokenized: TokenizedTrajectory, advantage: float +) -> tinker.Datum | None: if not ( + len(tokenized.token_ids) + == len(tokenized.logprobs) + == len(tokenized.assistant_mask) + ): + raise ValueError("Tokenized trajectory fields differ in length") + if len(tokenized.token_ids) < 2 or not any(tokenized.assistant_mask): + return None + if tokenized.assistant_mask[0]: + raise ValueError("A trainable trajectory cannot start with an assistant token") + + action_mask = tokenized.assistant_mask[1:] + if any( + trainable and math.isnan(logprob) + for trainable, logprob in zip(action_mask, tokenized.logprobs[1:], strict=True) + ): + raise ValueError("Tinker training requires logprobs for every assistant token") + return _build_datum( + tokenized.token_ids[:-1], + tokenized.token_ids[1:], + [ + logprob if trainable else 0.0 + for trainable, logprob in zip( + action_mask, tokenized.logprobs[1:], strict=True + ) + ], + [advantage if trainable else 0.0 for trainable in action_mask], + [float(trainable) for trainable in action_mask], + ) + + +def _build_datum( + input_tokens: list[int], + target_tokens: list[int], + logprobs: list[float], + advantages: list[float], + action_mask: list[float], +) -> tinker.Datum | None: + if not input_tokens or not ( len(input_tokens) == len(target_tokens) - == len(padded_logprobs) - == len(padded_advantages) + == len(logprobs) + == len(advantages) == len(action_mask) ): return None - return tinker.Datum( model_input=tinker.ModelInput.from_ints(tokens=input_tokens), loss_fn_inputs={ "target_tokens": tinker.TensorData.from_torch(torch.tensor(target_tokens)), "logprobs": tinker.TensorData.from_torch( - torch.tensor(padded_logprobs, dtype=torch.float32) + torch.tensor(logprobs, dtype=torch.float32) ), "advantages": tinker.TensorData.from_torch( - torch.tensor(padded_advantages, dtype=torch.float32) + torch.tensor(advantages, dtype=torch.float32) ), "mask": tinker.TensorData.from_torch( torch.tensor(action_mask, dtype=torch.float32) diff --git a/tests/integration/test_tinker_native_backend.py b/tests/integration/test_tinker_native_backend.py index 09ff33c47..3ad6411fd 100644 --- a/tests/integration/test_tinker_native_backend.py +++ b/tests/integration/test_tinker_native_backend.py @@ -31,13 +31,14 @@ async def simple_rollout( client: openai.AsyncOpenAI, model_name: str, prompt: str ) -> art.Trajectory: messages: art.Messages = [{"role": "user", "content": prompt}] - chat_completion = await client.chat.completions.create( - messages=messages, - model=model_name, - max_tokens=10, - timeout=60, - temperature=1, - ) + with art.Trajectory() as trajectory: + chat_completion = await client.chat.completions.create( + messages=messages, + model=model_name, + max_tokens=10, + timeout=60, + temperature=1, + ) choice = chat_completion.choices[0] content = (choice.message.content or "").lower() if "yes" in content: @@ -48,7 +49,8 @@ async def simple_rollout( reward = 0.25 else: reward = 0.0 - return art.Trajectory(messages_and_choices=[*messages, choice], reward=reward) # type: ignore[attr-defined] + trajectory.reward = reward + return trajectory @pytest.mark.skipif( diff --git a/tests/unit/test_tinker_native_exchanges.py b/tests/unit/test_tinker_native_exchanges.py new file mode 100644 index 000000000..d1bdcb8a4 --- /dev/null +++ b/tests/unit/test_tinker_native_exchanges.py @@ -0,0 +1,144 @@ +from datetime import datetime +from typing import Any + +from openai.types.chat import ChatCompletion +from openai.types.chat.chat_completion import Choice +import pytest + +import art +from art.trajectories import ChatCompletionsExchange, ChatCompletionsRequest + +pytest.importorskip("tinker") + +from art.tinker_native.data import trajectory_groups_to_datums # noqa: E402 + + +def _exchange( + prompt: list[int], output: list[int], *, logprobs: bool = True +) -> ChatCompletionsExchange: + response = ChatCompletion.model_validate( + { + "id": "chat-1", + "object": "chat.completion", + "created": 0, + "model": "test/model", + "choices": [ + { + "index": 0, + "finish_reason": "stop", + "message": {"role": "assistant", "content": "answer"}, + "prompt_token_ids": prompt, + "token_ids": output, + "logprobs": ( + { + "content": [ + { + "token": f"token_id:{token}", + "logprob": -token / 10, + "top_logprobs": [], + } + for token in output + ] + } + if logprobs + else None + ), + } + ], + } + ) + return ChatCompletionsExchange( + request=ChatCompletionsRequest(model="test/model", messages=[]), + response=response, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + +def _trajectory(reward: float, *, logprobs: bool = True) -> art.Trajectory: + return art.Trajectory( + exchanges=art.TrajectoryExchanges( + chat_completions=[ + _exchange([10], [20], logprobs=logprobs), + _exchange([10, 20, 11], [21], logprobs=logprobs), + ] + ), + reward=reward, + ) + + +def test_exchange_trajectory_builds_masked_multiturn_datum() -> None: + datums = trajectory_groups_to_datums( + [art.TrajectoryGroup([_trajectory(1), _trajectory(-1)])], + renderer=None, + tokenizer=None, + normalize_advantages=False, + ) + + assert len(datums) == 2 + datum = datums[0] + assert datum.model_input.to_ints() == [10, 20, 11] + assert datum.loss_fn_inputs["target_tokens"].to_torch().tolist() == [20, 11, 21] + assert datum.loss_fn_inputs["logprobs"].to_torch().tolist() == pytest.approx( + [-2.0, 0.0, -2.1] + ) + assert datum.loss_fn_inputs["advantages"].to_torch().tolist() == [1, 0, 1] + assert datum.loss_fn_inputs["mask"].to_torch().tolist() == [1, 0, 1] + + +def test_exchange_trajectory_requires_assistant_logprobs() -> None: + with pytest.raises(ValueError, match="requires logprobs"): + trajectory_groups_to_datums( + [art.TrajectoryGroup([_trajectory(1, logprobs=False), _trajectory(-1)])], + renderer=None, + tokenizer=None, + normalize_advantages=False, + ) + + +def test_legacy_trajectory_still_uses_history_conversion() -> None: + class Prompt: + def to_ints(self) -> list[int]: + return [10] + + class Renderer: + def build_generation_prompt(self, _messages: list[dict[str, Any]]) -> Prompt: + return Prompt() + + class Tokenizer: + def convert_tokens_to_ids(self, _token: str) -> int: + return 20 + + choice = Choice.model_validate( + { + "index": 0, + "finish_reason": "stop", + "message": {"role": "assistant", "content": "answer"}, + "logprobs": { + "content": [ + { + "token": "token_id:20", + "logprob": -2.0, + "top_logprobs": [], + } + ] + }, + } + ) + + def trajectory(reward: float) -> art.Trajectory: + return art.Trajectory( + messages_and_choices=[{"role": "user", "content": "question"}, choice], + reward=reward, + ) + + datums = trajectory_groups_to_datums( + [art.TrajectoryGroup([trajectory(1), trajectory(-1)])], + renderer=Renderer(), + tokenizer=Tokenizer(), + normalize_advantages=False, + ) + + assert len(datums) == 2 + assert datums[0].model_input.to_ints() == [10] + assert datums[0].loss_fn_inputs["target_tokens"].to_torch().tolist() == [20] From d11e9a1d6c53f70d46b1b0cda1d80eeebf599d7c Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 15 Jul 2026 15:22:02 +0000 Subject: [PATCH 10/17] fix: total trajectory completion tokens --- src/art/gather.py | 47 ++++++++--- tests/unit/trajectories/test_capture.py | 101 +++++++++++++++++++++++- 2 files changed, 135 insertions(+), 13 deletions(-) diff --git a/src/art/gather.py b/src/art/gather.py index 0dc53065a..b7608c879 100644 --- a/src/art/gather.py +++ b/src/art/gather.py @@ -176,24 +176,47 @@ async def wrap_trajectories_awaitable( def record_metrics(context: "GatherContext", trajectory: Trajectory) -> None: - logprobs = [ - message_or_choice.logprobs - for message_or_choice in trajectory.messages_and_choices - if isinstance(message_or_choice, Choice) - if message_or_choice.logprobs - ] - if logprobs: - # TODO: probably shouldn't average this - trajectory.metrics["completion_tokens"] = sum( - len(l.content or l.refusal or []) - for l in logprobs # noqa: E741 - ) / len(logprobs) + if trajectory.exchanges: + completion_tokens = _exchange_completion_tokens(trajectory) + if completion_tokens is not None: + trajectory.metrics["completion_tokens"] = completion_tokens + else: + logprobs = [ + message_or_choice.logprobs + for message_or_choice in trajectory.messages_and_choices + if isinstance(message_or_choice, Choice) + if message_or_choice.logprobs + ] + if logprobs: + trajectory.metrics["completion_tokens"] = sum( + len(logprob.content or logprob.refusal or []) for logprob in logprobs + ) context.metric_sums["reward"] += trajectory.reward context.metric_divisors["reward"] += 1 context.metric_sums.update(trajectory.metrics) context.metric_divisors.update(trajectory.metrics.keys()) +def _exchange_completion_tokens(trajectory: Trajectory) -> int | None: + counts: list[int | None] = [] + for exchange in trajectory.exchanges.chat_completions: + usage = exchange.response.usage + counts.append(usage.completion_tokens if usage is not None else None) + for exchange in trajectory.exchanges.completions: + usage = exchange.response.usage + counts.append(usage.completion_tokens if usage is not None else None) + for exchange in trajectory.exchanges.responses: + usage = exchange.response.usage + counts.append(usage.output_tokens if usage is not None else None) + counts.extend( + exchange.response.usage.output_tokens + for exchange in trajectory.exchanges.messages + ) + if not counts or any(count is None for count in counts): + return None + return sum(count for count in counts if count is not None) + + @dataclass class GatherContext: pbar: tqdm.tqdm | None = None diff --git a/tests/unit/trajectories/test_capture.py b/tests/unit/trajectories/test_capture.py index 49cd8d01c..e812e3cd9 100644 --- a/tests/unit/trajectories/test_capture.py +++ b/tests/unit/trajectories/test_capture.py @@ -18,7 +18,14 @@ import requests import art -from art.trajectories import ChatCompletionsExchange, MessagesExchange, _compat +from art.gather import GatherContext, record_metrics +from art.trajectories import ( + ChatCompletionsExchange, + CompletionsExchange, + MessagesExchange, + ResponsesExchange, + _compat, +) from art.trajectories._capture.core import begin, reset from art.trajectories._protocols import Endpoint, build_exchange, endpoint_for_url @@ -46,6 +53,7 @@ "prompt_token_ids": [1], } ], + "usage": {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3}, } COMPLETION: dict[str, Any] = { @@ -68,6 +76,7 @@ }, } ], + "usage": {"prompt_tokens": 1, "completion_tokens": 3, "total_tokens": 4}, } RESPONSE: dict[str, Any] = { @@ -94,6 +103,13 @@ "parallel_tool_calls": True, "tool_choice": "auto", "tools": [], + "usage": { + "input_tokens": 1, + "input_tokens_details": {"cached_tokens": 0}, + "output_tokens": 4, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": 5, + }, "raw_output_tokens": [{"token_id": 2, "logprob": -0.2}], } @@ -390,6 +406,89 @@ def test_all_protocols_reconstruct_typed_responses() -> None: assert exchange.model_dump(mode="json")["model"] == "updated/model" +def test_gather_metrics_sum_legacy_and_exchange_completion_tokens() -> None: + now = datetime.now() + chat = build_exchange( + "chat_completions", + {"model": "test/model", "messages": []}, + json.dumps(CHAT).encode(), + start_time=now, + end_time=now, + ) + completion = build_exchange( + "completions", + {"model": "test/model", "prompt": "hi"}, + json.dumps(COMPLETION).encode(), + start_time=now, + end_time=now, + ) + response = build_exchange( + "responses", + {"model": "test/model", "input": "hi"}, + json.dumps(RESPONSE).encode(), + start_time=now, + end_time=now, + ) + message = build_exchange( + "messages", + {"model": "test/model", "messages": []}, + json.dumps(MESSAGE).encode(), + start_time=now, + end_time=now, + ) + assert isinstance(chat, ChatCompletionsExchange) + assert isinstance(completion, CompletionsExchange) + assert isinstance(response, ResponsesExchange) + assert isinstance(message, MessagesExchange) + + trajectory = art.Trajectory( + exchanges=art.TrajectoryExchanges( + chat_completions=[chat], + completions=[completion], + responses=[response], + messages=[message], + ) + ) + record_metrics(GatherContext(), trajectory) + assert trajectory.metrics["completion_tokens"] == 10 + + first = chat.response.choices[0] + second = first.model_copy(deep=True) + assert second.logprobs is not None and second.logprobs.content is not None + second.logprobs.content *= 2 + legacy = art.Trajectory(messages_and_choices=[first, second]) + record_metrics(GatherContext(), legacy) + assert legacy.metrics["completion_tokens"] == 3 + + +def test_gather_metrics_omit_partial_exchange_usage() -> None: + now = datetime.now() + chat = build_exchange( + "chat_completions", + {"model": "test/model", "messages": []}, + json.dumps(CHAT).encode(), + start_time=now, + end_time=now, + ) + message = build_exchange( + "messages", + {"model": "test/model", "messages": []}, + json.dumps(MESSAGE).encode(), + start_time=now, + end_time=now, + ) + assert isinstance(chat, ChatCompletionsExchange) + assert isinstance(message, MessagesExchange) + chat.response.usage = None + trajectory = art.Trajectory( + exchanges=art.TrajectoryExchanges(chat_completions=[chat], messages=[message]) + ) + + record_metrics(GatherContext(), trajectory) + + assert "completion_tokens" not in trajectory.metrics + + @pytest.mark.parametrize( ("url", "endpoint"), [ From 6767051cd7a894768f0d4c82de219ed6ef97d696 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 15 Jul 2026 15:46:48 +0000 Subject: [PATCH 11/17] feat: expose protocol trajectory histories --- src/art/__init__.py | 12 + src/art/rewards/ruler.py | 16 +- src/art/trajectories/__init__.py | 117 +++++++- src/art/trajectories/_history.py | 353 ++++++++++++++++++++++++ tests/unit/trajectories/test_history.py | 321 +++++++++++++++++++++ 5 files changed, 812 insertions(+), 7 deletions(-) create mode 100644 src/art/trajectories/_history.py create mode 100644 tests/unit/trajectories/test_history.py diff --git a/src/art/__init__.py b/src/art/__init__.py index 0f3b9f852..9f4a8f836 100644 --- a/src/art/__init__.py +++ b/src/art/__init__.py @@ -70,15 +70,21 @@ from .model import Model, TrainableModel from .serverless import ServerlessBackend from .trajectories import ( + AnthropicMessagesHistory, ChatCompletionsExchange, + ChatCompletionsHistory, CompletionsExchange, + CompletionsHistory, + History, MessagesExchange, ResponsesExchange, + ResponsesHistory, TokenizedTrajectory, TokenizedTrajectoryGroup, Trajectory, TrajectoryExchanges, TrajectoryGroup, + TrajectoryHistory, auto_trajectory, # ty: ignore[deprecated] capture_auto_trajectory, # ty: ignore[deprecated] current_trajectory, @@ -135,10 +141,16 @@ "Trajectory", "TrajectoryExchanges", "TrajectoryGroup", + "History", + "TrajectoryHistory", "ChatCompletionsExchange", + "ChatCompletionsHistory", "CompletionsExchange", + "CompletionsHistory", "ResponsesExchange", + "ResponsesHistory", "MessagesExchange", + "AnthropicMessagesHistory", "TokenizedTrajectory", "TokenizedTrajectoryGroup", "trajectory", diff --git a/src/art/rewards/ruler.py b/src/art/rewards/ruler.py index 7c56dcb7d..34aa64668 100644 --- a/src/art/rewards/ruler.py +++ b/src/art/rewards/ruler.py @@ -320,6 +320,7 @@ async def ruler_score_group( for t in group.trajectories: # Create a new trajectory with the same data but fresh objects new_traj = t.__class__( + exchanges=t.exchanges.model_copy(deep=True), messages_and_choices=t.messages_and_choices.copy(), tools=t.tools.copy() if t.tools else None, additional_histories=[ @@ -333,13 +334,18 @@ async def ruler_score_group( new_trajectories.append(new_traj) # Extract message lists and preserve original rewards for comparison - message_lists: list[list[ChatCompletionMessageParam]] = [] - for traj in new_trajectories: - message_lists.append(traj.messages()) - traj.metrics["independent_reward"] = traj.reward + histories = [ + trajectory.history().as_chat_completions_history() + for trajectory in new_trajectories + ] + message_lists: list[list[ChatCompletionMessageParam]] = [ + history.messages for history in histories + ] + for trajectory in new_trajectories: + trajectory.metrics["independent_reward"] = trajectory.reward # Extract tools from first trajectory (they should all be the same) - tools = new_trajectories[0].tools if new_trajectories else None + tools = histories[0].tools if histories else None try: # Call the core ruler function to get scores diff --git a/src/art/trajectories/__init__.py b/src/art/trajectories/__init__.py index 21c0bd032..73f4f4c7a 100644 --- a/src/art/trajectories/__init__.py +++ b/src/art/trajectories/__init__.py @@ -12,7 +12,7 @@ from datetime import datetime import time from types import TracebackType -from typing import Annotated, Any, Literal, overload +from typing import Annotated, Any, Literal, TypeAlias, overload from anthropic.types import ( Message as AnthropicMessage, @@ -197,6 +197,85 @@ class History(pydantic.BaseModel): def messages(self) -> Messages: return get_messages(self.messages_and_choices) + def as_chat_completions_history(self) -> ChatCompletionsHistory: + from ._history import legacy_as_chat_completions_history + + return legacy_as_chat_completions_history(self) + + +class ChatCompletionsHistory(pydantic.BaseModel): + model: str | None + messages: Messages + tools: Tools | None = None + chat_template: str | None = None + chat_template_kwargs: dict[str, Any] | None = None + + def as_chat_completions_history(self) -> ChatCompletionsHistory: + return self + + +class AnthropicMessagesHistory(pydantic.BaseModel): + 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 + chat_template: str | None = None + chat_template_kwargs: dict[str, Any] | None = None + + def as_chat_completions_history(self) -> ChatCompletionsHistory: + from ._history import anthropic_as_chat_completions_history + + return anthropic_as_chat_completions_history(self) + + +class ResponsesHistory(pydantic.BaseModel): + model: str + input: Annotated[ + pydantic.SerializeAsAny[ResponseInputParam], pydantic.SkipValidation + ] + instructions: str | None = None + tools: Annotated[ + pydantic.SerializeAsAny[list[ResponsesToolParam] | None], + pydantic.SkipValidation, + ] = None + chat_template: str | None = None + chat_template_kwargs: dict[str, Any] | None = None + + def as_chat_completions_history(self) -> ChatCompletionsHistory: + from ._history import responses_as_chat_completions_history + + return responses_as_chat_completions_history(self) + + +class CompletionsHistory(pydantic.BaseModel): + model: str + token_ids: list[int] + sampled_spans: list[tuple[int, int]] + + def as_chat_completions_history(self) -> ChatCompletionsHistory: + raise ValueError("Raw Completions history has no chat-message structure") + + +TrajectoryHistory: TypeAlias = ( + History + | ChatCompletionsHistory + | AnthropicMessagesHistory + | ResponsesHistory + | CompletionsHistory +) + class Trajectory(_CompactModel): exchanges: TrajectoryExchanges = pydantic.Field(default_factory=TrajectoryExchanges) @@ -262,8 +341,37 @@ 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})" + def chat_completions_history( + self, *, model: str | None = None + ) -> ChatCompletionsHistory: + from ._history import chat_completions_history + + return chat_completions_history(self, model=model) + + def anthropic_messages_history( + self, *, model: str | None = None + ) -> AnthropicMessagesHistory: + from ._history import anthropic_messages_history + + return anthropic_messages_history(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 + + return completions_history(self, model=model) + + def history(self, *, model: str | None = None) -> TrajectoryHistory: + from ._history import trajectory_history + + return trajectory_history(self, model=model) + def messages(self) -> Messages: - return get_messages(self.messages_and_choices) + return self.history().as_chat_completions_history().messages def for_logging(self) -> dict[str, object]: from ._compat import trajectory_for_logging @@ -559,6 +667,11 @@ def get_messages(messages_and_choices: MessagesAndChoices) -> Messages: "TrajectoryExchanges", "PydanticException", "History", + "ChatCompletionsHistory", + "AnthropicMessagesHistory", + "ResponsesHistory", + "CompletionsHistory", + "TrajectoryHistory", "Trajectory", "TrajectoryGroup", "TokenizedTrajectory", diff --git a/src/art/trajectories/_history.py b/src/art/trajectories/_history.py new file mode 100644 index 000000000..8d1b86730 --- /dev/null +++ b/src/art/trajectories/_history.py @@ -0,0 +1,353 @@ +from __future__ import annotations + +from collections.abc import Iterable, Mapping, Sequence +import copy +from datetime import datetime +from typing import Protocol, TypeVar + +import pydantic + +from ..types import Message, Messages, Tools +from . import ( + AnthropicMessagesHistory, + ChatCompletionsExchange, + ChatCompletionsHistory, + CompletionsExchange, + CompletionsHistory, + History, + MessagesExchange, + ResponsesExchange, + ResponsesHistory, + Trajectory, + TrajectoryHistory, +) + + +class _ModelledExchange(Protocol): + start_time: datetime + end_time: datetime + + @property + def model(self) -> str | None: ... + + +_ExchangeT = TypeVar("_ExchangeT", bound=_ModelledExchange) +_ItemT = TypeVar("_ItemT") +_MESSAGES = pydantic.TypeAdapter(Messages) +_MESSAGE = pydantic.TypeAdapter(Message) +_TOOLS = pydantic.TypeAdapter(Tools | None) + + +def _select( + exchanges: Sequence[_ExchangeT], model: str | None, protocol: 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: + raise ValueError(f"Every {protocol} exchange must identify its model") + if len(models) != 1: + raise ValueError( + f"{protocol} history requires exactly one model; pass model= to select one" + ) + return sorted( + selected, key=lambda exchange: (exchange.start_time, exchange.end_time) + ) + + +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 _require_context( + requests: Sequence[Mapping[str, object]], fields: Sequence[str] +) -> None: + for field in fields: + _require_constant((request.get(field) for request in requests), field) + + +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 _only_choice(exchange: ChatCompletionsExchange | CompletionsExchange) -> None: + if len(exchange.response.choices) != 1: + raise ValueError("Multiple response choices form multiple histories") + + +def _model(exchange: _ModelledExchange) -> str: + if exchange.model is None: + raise AssertionError("_select returned an exchange without a model") + return exchange.model + + +def legacy_as_chat_completions_history(history: History) -> ChatCompletionsHistory: + return ChatCompletionsHistory( + model=None, + messages=history.messages(), + tools=copy.deepcopy(history.tools), + ) + + +def chat_completions_history( + trajectory: Trajectory, *, model: str | None +) -> ChatCompletionsHistory: + 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( + 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 + ) + ) + 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")), + ) + + +def anthropic_messages_history( + trajectory: Trajectory, *, model: str | None +) -> AnthropicMessagesHistory: + 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")), + ) + + +def _responses_input(value: object) -> list[object]: + if isinstance(value, str): + return [{"role": "user", "content": value}] + if isinstance(value, list): + return [copy.deepcopy(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: + 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" + ) + 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")), + ) + + +def completions_history( + trajectory: Trajectory, *, model: str | None +) -> CompletionsHistory: + from ._tokenize import _completion_tokens + + 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) + prompt, completion, _ = _completion_tokens(exchange.response) + if prompt is None: + raise ValueError( + "Completions history requires exact prompt and output token IDs" + ) + 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" + ) + 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, + ) + + +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} + ) + tools = _TOOLS.validate_python(_openai_tools(history.tools, dialect="messages")) + return ChatCompletionsHistory( + model=history.model, + messages=_MESSAGES.validate_python(messages), + tools=tools, + chat_template=history.chat_template, + chat_template_kwargs=copy.deepcopy(history.chat_template_kwargs), + ) + + +def responses_as_chat_completions_history( + history: ResponsesHistory, +) -> ChatCompletionsHistory: + from ._tokenize import _openai_tools, _responses_messages + + input_without_reasoning = [ + item + for item in history.input + if not isinstance(item, Mapping) or item.get("type") != "reasoning" + ] + messages = _responses_messages( + {"instructions": history.instructions, "input": input_without_reasoning} + ) + tools = _TOOLS.validate_python(_openai_tools(history.tools, dialect="responses")) + return ChatCompletionsHistory( + model=history.model, + messages=_MESSAGES.validate_python(messages), + tools=tools, + chat_template=history.chat_template, + chat_template_kwargs=copy.deepcopy(history.chat_template_kwargs), + ) + + +def trajectory_history( + trajectory: Trajectory, *, model: str | None +) -> TrajectoryHistory: + 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), + ) + + 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) + ] + 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: + 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) diff --git a/tests/unit/trajectories/test_history.py b/tests/unit/trajectories/test_history.py new file mode 100644 index 000000000..3bd822872 --- /dev/null +++ b/tests/unit/trajectories/test_history.py @@ -0,0 +1,321 @@ +from datetime import datetime, timedelta +import importlib + +from anthropic.types import Message +from openai.types import Completion +from openai.types.chat import ChatCompletion +from openai.types.responses import Response +import pytest + +import art +from art.trajectories import ( + ChatCompletionsExchange, + CompletionsExchange, + MessagesExchange, + ResponsesExchange, + TrajectoryExchanges, +) + + +def _times(offset: int = 0) -> tuple[datetime, datetime]: + start = datetime(2026, 1, 1) + timedelta(seconds=offset) + return start, start + timedelta(milliseconds=1) + + +def _chat( + messages: list[dict[str, object]], + answer: str, + *, + model: str = "test/model", + offset: int = 0, +) -> ChatCompletionsExchange: + start, end = _times(offset) + return ChatCompletionsExchange( + request={"model": model, "messages": messages}, + response=ChatCompletion.model_validate( + { + "id": f"chat-{offset}", + "object": "chat.completion", + "created": offset, + "model": model, + "choices": [ + { + "index": 0, + "finish_reason": "stop", + "message": {"role": "assistant", "content": answer}, + } + ], + } + ), + start_time=start, + end_time=end, + ) + + +def _completion( + prompt: list[int], output: list[int], *, offset: int = 0 +) -> CompletionsExchange: + start, end = _times(offset) + return CompletionsExchange( + request={"model": "test/model", "prompt": prompt}, + response=Completion.model_validate( + { + "id": f"completion-{offset}", + "object": "text_completion", + "created": offset, + "model": "test/model", + "choices": [ + { + "index": 0, + "finish_reason": "stop", + "text": "answer", + "prompt_token_ids": prompt, + "token_ids": output, + } + ], + } + ), + start_time=start, + end_time=end, + ) + + +def _message() -> MessagesExchange: + start, end = _times() + return MessagesExchange( + request={ + "model": "test/model", + "system": "Be concise", + "messages": [{"role": "user", "content": "Hello"}], + "max_tokens": 16, + }, + response=Message.model_validate( + { + "id": "message-1", + "type": "message", + "role": "assistant", + "model": "test/model", + "content": [{"type": "text", "text": "Hi"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 1, "output_tokens": 1}, + } + ), + start_time=start, + end_time=end, + ) + + +def _response( + response_id: str, + text: str, + *, + previous_response_id: str | None = None, + offset: int = 0, +) -> ResponsesExchange: + start, end = _times(offset) + request: dict[str, object] = { + "model": "test/model", + "input": f"turn {offset}", + } + if previous_response_id is not None: + request["previous_response_id"] = previous_response_id + return ResponsesExchange( + request=request, + response=Response.model_validate( + { + "id": response_id, + "created_at": float(offset), + "model": "test/model", + "object": "response", + "output": [ + { + "id": f"message-{response_id}", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [ + { + "type": "output_text", + "text": text, + "annotations": [], + "logprobs": [], + } + ], + } + ], + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [], + } + ), + start_time=start, + end_time=end, + ) + + +def test_chat_history_resolves_one_model_and_append_only_sequence() -> None: + first = _chat([{"role": "user", "content": "one"}], "first") + second = _chat( + [ + {"role": "user", "content": "one"}, + {"role": "assistant", "content": "first"}, + {"role": "user", "content": "two"}, + ], + "second", + offset=1, + ) + other = _chat( + [{"role": "user", "content": "other"}], + "other", + model="other/model", + offset=2, + ) + trajectory = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[first, second, other]) + ) + + with pytest.raises(ValueError, match="exactly one model"): + trajectory.history() + history = trajectory.history(model="test/model") + assert isinstance(history, art.ChatCompletionsHistory) + assert [message["role"] for message in history.messages] == [ + "user", + "assistant", + "user", + "assistant", + ] + 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") + second.request.pop("cache_salt") + + second.request["messages"] = [{"role": "user", "content": "branch"}] + with pytest.raises(ValueError, match="append-only"): + trajectory.chat_completions_history(model="test/model") + + +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 [message["role"] for message in messages_history.messages] == [ + "user", + "assistant", + ] + assert [message["role"] for message in message_trajectory.messages()] == [ + "system", + "user", + "assistant", + ] + + mixed = art.Trajectory( + exchanges=TrajectoryExchanges( + chat_completions=[_chat([{"role": "user", "content": "hi"}], "hi")], + messages=[_message()], + ) + ) + with pytest.raises(ValueError, match="multiple protocol histories"): + mixed.history() + assert isinstance(mixed.anthropic_messages_history(), art.AnthropicMessagesHistory) + + +def test_responses_history_expands_previous_response_chain() -> None: + trajectory = art.Trajectory( + exchanges=TrajectoryExchanges( + responses=[ + _response("response-1", "first"), + _response( + "response-2", + "second", + previous_response_id="response-1", + offset=1, + ), + ] + ) + ) + + history = trajectory.responses_history() + assert len(history.input) == 4 + assert ( + art.ResponsesHistory.model_validate_json(history.model_dump_json()) == history + ) + assert [ + message["role"] for message in history.as_chat_completions_history().messages + ] == [ + "user", + "assistant", + "user", + "assistant", + ] + + trajectory.exchanges.responses[1].request["previous_response_id"] = "missing" + with pytest.raises(ValueError, match="outside this history"): + trajectory.responses_history() + + +def test_completions_history_preserves_exact_tokens_and_sampled_spans() -> None: + trajectory = art.Trajectory( + exchanges=TrajectoryExchanges( + completions=[ + _completion([1], [2]), + _completion([1, 2, 3], [4], offset=1), + ] + ) + ) + + history = trajectory.completions_history() + assert history.token_ids == [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_legacy_messages_delegate_through_history() -> None: + trajectory = art.Trajectory( + messages_and_choices=[{"role": "user", "content": "hello"}] + ) + + assert isinstance(trajectory.history(), art.History) + assert trajectory.messages() == [{"role": "user", "content": "hello"}] + with pytest.raises(ValueError, match="do not identify a model"): + trajectory.history(model="test/model") + + +@pytest.mark.asyncio +async def test_ruler_accepts_exchange_trajectories( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from art.rewards.ruler import TrajectoryScore, ruler_score_group + + ruler_module = importlib.import_module("art.rewards.ruler") + captured: list[list[dict[str, object]]] = [] + + async def score( + message_lists: list[list[dict[str, object]]], **_: object + ) -> list[TrajectoryScore]: + captured.extend(message_lists) + return [TrajectoryScore(trajectory_id="1", explanation="good", score=0.8)] + + monkeypatch.setattr(ruler_module, "ruler", score) + trajectory = art.Trajectory( + exchanges=TrajectoryExchanges( + chat_completions=[_chat([{"role": "user", "content": "hi"}], "hello")] + ) + ) + + result = await ruler_score_group(art.TrajectoryGroup([trajectory])) + + assert result is not None + assert result.trajectories[0].exchanges == trajectory.exchanges + assert [message["role"] for message in captured[0]] == ["user", "assistant"] From 491b35b854668d668b7799d019f2206a79668ed1 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 15 Jul 2026 15:49:29 +0000 Subject: [PATCH 12/17] fix: preserve Responses reasoning in history conversion --- src/art/trajectories/_history.py | 41 ++++++++++++++++++++----- tests/unit/trajectories/test_history.py | 26 ++++++++++++---- 2 files changed, 53 insertions(+), 14 deletions(-) diff --git a/src/art/trajectories/_history.py b/src/art/trajectories/_history.py index 8d1b86730..cb9ced102 100644 --- a/src/art/trajectories/_history.py +++ b/src/art/trajectories/_history.py @@ -295,14 +295,26 @@ def responses_as_chat_completions_history( ) -> ChatCompletionsHistory: from ._tokenize import _openai_tools, _responses_messages - input_without_reasoning = [ - item - for item in history.input - if not isinstance(item, Mapping) or item.get("type") != "reasoning" - ] - messages = _responses_messages( - {"instructions": history.instructions, "input": input_without_reasoning} - ) + messages = _responses_messages({"instructions": history.instructions, "input": []}) + pending_reasoning = "" + for item in history.input: + if isinstance(item, Mapping) and item.get("type") == "reasoning": + pending_reasoning += _responses_reasoning_text(item) + continue + converted = _responses_messages({"input": [item]}) + if pending_reasoning and converted: + if converted[0].get("role") == "assistant": + converted[0]["reasoning"] = pending_reasoning + else: + messages.append( + {"role": "assistant", "content": "", "reasoning": pending_reasoning} + ) + pending_reasoning = "" + messages.extend(converted) + if pending_reasoning: + messages.append( + {"role": "assistant", "content": "", "reasoning": pending_reasoning} + ) tools = _TOOLS.validate_python(_openai_tools(history.tools, dialect="responses")) return ChatCompletionsHistory( model=history.model, @@ -313,6 +325,19 @@ def responses_as_chat_completions_history( ) +def _responses_reasoning_text(item: Mapping[str, object]) -> str: + text = "" + for field in ("content", "summary"): + blocks = item.get(field) + if not isinstance(blocks, list): + continue + for block in blocks: + value = block.get("text") if isinstance(block, Mapping) else None + if isinstance(value, str): + text += value + return text + + def trajectory_history( trajectory: Trajectory, *, model: str | None ) -> TrajectoryHistory: diff --git a/tests/unit/trajectories/test_history.py b/tests/unit/trajectories/test_history.py index 3bd822872..654b5370d 100644 --- a/tests/unit/trajectories/test_history.py +++ b/tests/unit/trajectories/test_history.py @@ -111,6 +111,7 @@ def _response( text: str, *, previous_response_id: str | None = None, + reasoning: str | None = None, offset: int = 0, ) -> ResponsesExchange: start, end = _times(offset) @@ -129,6 +130,19 @@ def _response( "model": "test/model", "object": "response", "output": [ + *( + [ + { + "id": f"reasoning-{response_id}", + "type": "reasoning", + "summary": [ + {"type": "summary_text", "text": reasoning} + ], + } + ] + if reasoning is not None + else [] + ), { "id": f"message-{response_id}", "type": "message", @@ -142,7 +156,7 @@ def _response( "logprobs": [], } ], - } + }, ], "parallel_tool_calls": True, "tool_choice": "auto", @@ -234,7 +248,7 @@ def test_responses_history_expands_previous_response_chain() -> None: trajectory = art.Trajectory( exchanges=TrajectoryExchanges( responses=[ - _response("response-1", "first"), + _response("response-1", "first", reasoning="think"), _response( "response-2", "second", @@ -246,18 +260,18 @@ def test_responses_history_expands_previous_response_chain() -> None: ) history = trajectory.responses_history() - assert len(history.input) == 4 + assert len(history.input) == 5 assert ( art.ResponsesHistory.model_validate_json(history.model_dump_json()) == history ) - assert [ - message["role"] for message in history.as_chat_completions_history().messages - ] == [ + chat_history = history.as_chat_completions_history() + assert [message["role"] for message in chat_history.messages] == [ "user", "assistant", "user", "assistant", ] + assert dict(chat_history.messages[1]).get("reasoning") == "think" trajectory.exchanges.responses[1].request["previous_response_id"] = "missing" with pytest.raises(ValueError, match="outside this history"): From 3c066996f354210cd9124e14f4b570e0bf2cf306 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 15 Jul 2026 16:40:05 +0000 Subject: [PATCH 13/17] fix: harden trajectory capture and replay --- src/art/trajectories/_capture/aiohttp.py | 24 +++ src/art/trajectories/_capture/core.py | 7 + src/art/trajectories/_capture/requests.py | 2 +- src/art/trajectories/_compat.py | 40 ++-- src/art/trajectories/_history.py | 61 +++--- src/art/trajectories/_protocols.py | 16 +- src/art/trajectories/_tokenize.py | 187 +++++++++++++---- tests/unit/trajectories/test_capture.py | 138 +++++++++++- tests/unit/trajectories/test_history.py | 26 +++ tests/unit/trajectories/test_tokenize.py | 242 ++++++++++++++++++++-- 10 files changed, 621 insertions(+), 122 deletions(-) diff --git a/src/art/trajectories/_capture/aiohttp.py b/src/art/trajectories/_capture/aiohttp.py index 3b47a53a1..1100beae5 100644 --- a/src/art/trajectories/_capture/aiohttp.py +++ b/src/art/trajectories/_capture/aiohttp.py @@ -40,6 +40,15 @@ async def readany(self) -> bytes: async def readline(self) -> bytes: return self._record(await self._stream.readline()) + async def readexactly(self, n: int) -> bytes: + return self._record(await self._stream.readexactly(n)) + + async def readuntil(self, separator: bytes = b"\n") -> bytes: + return self._record(await self._stream.readuntil(separator)) + + def read_nowait(self, n: int = -1) -> bytes: + return self._record(self._stream.read_nowait(n)) + async def readchunk(self) -> tuple[bytes, bool]: return self._record(await self._stream.readchunk()) @@ -62,6 +71,21 @@ def iter_any(self) -> AsyncIterator[bytes]: def iter_chunked(self, size: int) -> AsyncIterator[bytes]: return self._iterate(self._stream.iter_chunked(size)) + async def _iterate_chunks( + self, iterator: AsyncIterator[tuple[bytes, bool]] + ) -> AsyncIterator[tuple[bytes, bool]]: + completed = False + try: + async for chunk in iterator: + yield self._record(chunk) + completed = True + finally: + if completed or self._state.request.get("stream") is True: + self._state.finish() + + def iter_chunks(self) -> AsyncIterator[tuple[bytes, bool]]: + return self._iterate_chunks(self._stream.iter_chunks()) + def install() -> None: if getattr(aiohttp.ClientSession._request, "_art_capture", False): diff --git a/src/art/trajectories/_capture/core.py b/src/art/trajectories/_capture/core.py index d3bade490..8de4f154b 100644 --- a/src/art/trajectories/_capture/core.py +++ b/src/art/trajectories/_capture/core.py @@ -59,6 +59,13 @@ def finish(self) -> None: def _append_exchange(trajectory: Trajectory, exchange: Exchange) -> None: + if ( + trajectory.messages_and_choices + or trajectory.tools is not None + or trajectory.additional_histories + ): + logger.debug("Ignoring exchange captured into a legacy trajectory") + return if isinstance(exchange, ChatCompletionsExchange): trajectory.exchanges.chat_completions.append(exchange) elif isinstance(exchange, CompletionsExchange): diff --git a/src/art/trajectories/_capture/requests.py b/src/art/trajectories/_capture/requests.py index 242b2edc6..9af7bc94c 100644 --- a/src/art/trajectories/_capture/requests.py +++ b/src/art/trajectories/_capture/requests.py @@ -27,7 +27,7 @@ def send( if state is not None: state.status_code = response.status_code setattr(response, _STATE, state) - if not kwargs.get("stream", False): + if not kwargs.get("stream", self.stream): state.add(response.content) state.finish() return response diff --git a/src/art/trajectories/_compat.py b/src/art/trajectories/_compat.py index ecad617be..82ed0c5b7 100644 --- a/src/art/trajectories/_compat.py +++ b/src/art/trajectories/_compat.py @@ -17,10 +17,12 @@ def exception_model( - exception: BaseException | PydanticException, + exception: object, ) -> PydanticException: if isinstance(exception, PydanticException): return exception + if not isinstance(exception, BaseException): + return PydanticException.model_validate(exception) return PydanticException( type=str(type(exception)), message=str(exception), @@ -33,7 +35,7 @@ def exception_model( async def _legacy_async_group( - items: list[Awaitable[Trajectory]], + items: list[Awaitable[object]], *, exceptions: Iterable[BaseException | PydanticException], metadata: dict[str, MetadataValue] | None, @@ -48,8 +50,13 @@ async def _legacy_async_group( for future in asyncio.as_completed(items): try: item = await future - trajectories.append(item) - record_metrics(context, item) + trajectory = ( + item + if isinstance(item, Trajectory) + else Trajectory.model_validate(item) + ) + trajectories.append(trajectory) + record_metrics(context, trajectory) context.update_pbar(n=1) except BaseException as exc: captured_exceptions.append(exc) @@ -79,7 +86,7 @@ def __await__(self) -> Generator[Any, None, TrajectoryGroup]: def new_trajectory_group( cls: type[TrajectoryGroup], - trajectories: Iterable[Trajectory | BaseException | Awaitable[Trajectory]], + trajectories: Iterable[object], *, exceptions: Iterable[BaseException | PydanticException], metadata: dict[str, MetadataValue] | None, @@ -87,9 +94,7 @@ def new_trajectory_group( logs: list[str] | None, ) -> TrajectoryGroup | Awaitable[TrajectoryGroup]: items = list(trajectories) - awaitables = [ - item for item in items if not isinstance(item, (Trajectory, BaseException)) - ] + awaitables = [item for item in items if isinstance(item, Awaitable)] if awaitables: if len(awaitables) != len(items): raise TypeError("TrajectoryGroup cannot mix trajectories and awaitables") @@ -108,28 +113,23 @@ def new_trajectory_group( ), len(items), ) - sync_items = [ - item for item in items if isinstance(item, (Trajectory, BaseException)) - ] - if len(sync_items) != len(items): - raise TypeError("TrajectoryGroup items must be trajectories or exceptions") group = object.__new__(cls) - object.__setattr__(group, _PREPARED_TRAJECTORIES, sync_items) + object.__setattr__(group, _PREPARED_TRAJECTORIES, items) return group def init_trajectory_group( group: TrajectoryGroup, - trajectories: Iterable[Trajectory | BaseException | Awaitable[Trajectory]], + trajectories: Iterable[object], *, - exceptions: Iterable[BaseException | PydanticException], + exceptions: Iterable[object], metadata: dict[str, MetadataValue] | None, metrics: dict[str, float | int | bool] | None, logs: list[str] | None, ) -> None: prepared = group.__dict__.pop(_PREPARED_TRAJECTORIES, None) items = prepared if isinstance(prepared, list) else list(trajectories) - if not all(isinstance(item, (Trajectory, BaseException)) for item in items): + if any(isinstance(item, Awaitable) for item in items): raise TypeError("TrajectoryGroup cannot initialize from awaitables") normalized_trajectories = [ item if isinstance(item, Trajectory) else Trajectory.model_validate(item) @@ -154,7 +154,7 @@ def init_trajectory_group( def copy_trajectory_group(group: TrajectoryGroup) -> TrajectoryGroup: - copied = TrajectoryGroup( + copied = group.__class__( group.trajectories[:], metadata=group.metadata.copy(), metrics=group.metrics.copy(), @@ -168,11 +168,11 @@ def deepcopy_trajectory_group( group: TrajectoryGroup, memo: dict[int, object] | None ) -> TrajectoryGroup: memo = {} if memo is None else memo - if existing := memo.get(id(group)): + if (existing := memo.get(id(group))) is not None: if not isinstance(existing, TrajectoryGroup): raise TypeError("TrajectoryGroup deepcopy memo contains an invalid value") return existing - copied = TrajectoryGroup( + copied = group.__class__( copy.deepcopy(group.trajectories, memo), metadata=copy.deepcopy(group.metadata, memo), metrics=copy.deepcopy(group.metrics, memo), diff --git a/src/art/trajectories/_history.py b/src/art/trajectories/_history.py index cb9ced102..f0d915e84 100644 --- a/src/art/trajectories/_history.py +++ b/src/art/trajectories/_history.py @@ -100,6 +100,17 @@ def _model(exchange: _ModelledExchange) -> str: return exchange.model +def _require_unmixed(trajectory: Trajectory) -> None: + if trajectory.exchanges and ( + trajectory.messages_and_choices + or trajectory.tools is not None + or trajectory.additional_histories + ): + raise ValueError( + "A trajectory cannot contain both exchanges and legacy histories" + ) + + def legacy_as_chat_completions_history(history: History) -> ChatCompletionsHistory: return ChatCompletionsHistory( model=None, @@ -111,6 +122,7 @@ def legacy_as_chat_completions_history(history: History) -> ChatCompletionsHisto def chat_completions_history( trajectory: Trajectory, *, model: str | None ) -> ChatCompletionsHistory: + _require_unmixed(trajectory) if not trajectory.exchanges: if trajectory.additional_histories: raise ValueError("Trajectory contains multiple legacy histories") @@ -152,6 +164,7 @@ def chat_completions_history( 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], @@ -198,6 +211,7 @@ def _responses_input(value: object) -> list[object]: def responses_history(trajectory: Trajectory, *, model: str | None) -> ResponsesHistory: + _require_unmixed(trajectory) exchanges = _select(trajectory.exchanges.responses, model, "Responses") _require_context( [exchange.request for exchange in exchanges], @@ -241,15 +255,23 @@ def responses_history(trajectory: Trajectory, *, model: str | None) -> Responses def completions_history( trajectory: Trajectory, *, model: str | None ) -> CompletionsHistory: - from ._tokenize import _completion_tokens + 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): + prompt = _exact_token_ids( + request_prompt, field="Completions request prompt" + ) if prompt is None: raise ValueError( "Completions history requires exact prompt and output token IDs" @@ -295,26 +317,9 @@ def responses_as_chat_completions_history( ) -> ChatCompletionsHistory: from ._tokenize import _openai_tools, _responses_messages - messages = _responses_messages({"instructions": history.instructions, "input": []}) - pending_reasoning = "" - for item in history.input: - if isinstance(item, Mapping) and item.get("type") == "reasoning": - pending_reasoning += _responses_reasoning_text(item) - continue - converted = _responses_messages({"input": [item]}) - if pending_reasoning and converted: - if converted[0].get("role") == "assistant": - converted[0]["reasoning"] = pending_reasoning - else: - messages.append( - {"role": "assistant", "content": "", "reasoning": pending_reasoning} - ) - pending_reasoning = "" - messages.extend(converted) - if pending_reasoning: - messages.append( - {"role": "assistant", "content": "", "reasoning": pending_reasoning} - ) + messages = _responses_messages( + {"instructions": history.instructions, "input": history.input} + ) tools = _TOOLS.validate_python(_openai_tools(history.tools, dialect="responses")) return ChatCompletionsHistory( model=history.model, @@ -325,22 +330,10 @@ def responses_as_chat_completions_history( ) -def _responses_reasoning_text(item: Mapping[str, object]) -> str: - text = "" - for field in ("content", "summary"): - blocks = item.get(field) - if not isinstance(blocks, list): - continue - for block in blocks: - value = block.get("text") if isinstance(block, Mapping) else None - if isinstance(value, str): - text += value - return text - - def trajectory_history( trajectory: Trajectory, *, model: str | None ) -> TrajectoryHistory: + _require_unmixed(trajectory) if not trajectory.exchanges: if trajectory.additional_histories: raise ValueError("Trajectory contains multiple legacy histories") diff --git a/src/art/trajectories/_protocols.py b/src/art/trajectories/_protocols.py index b5b750eac..3f985ec27 100644 --- a/src/art/trajectories/_protocols.py +++ b/src/art/trajectories/_protocols.py @@ -12,7 +12,7 @@ from openai.types.chat import ChatCompletion from openai.types.chat.chat_completion_chunk import ChatCompletionChunk from openai.types.responses import Response -from pydantic import TypeAdapter +from pydantic import TypeAdapter, ValidationError from ..openai import init_chat_completion, update_chat_completion from . import ( @@ -63,7 +63,12 @@ def _sse_events(body: bytes) -> list[tuple[str | None, SSEPayload]]: if raw == "[DONE]": events.append((event_name, "[DONE]")) else: - value = json.loads(raw) + try: + value = json.loads(raw) + except json.JSONDecodeError: + if raw.strip().lower() in {"ping", "keepalive"}: + continue + raise if isinstance(value, dict): events.append((event_name, value)) return events @@ -79,7 +84,12 @@ def _chat_response(body: bytes, *, stream: bool) -> ChatCompletion: if payload == "[DONE]": done = True continue - chunk = ChatCompletionChunk.model_validate(payload) + try: + chunk = ChatCompletionChunk.model_validate(payload) + except ValidationError: + if not payload.get("object") and not payload.get("choices"): + continue + raise if response is None: response = init_chat_completion(chunk.model_copy(update={"choices": []})) update_chat_completion(response, chunk.model_copy(update={"choices": []})) diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index c61c14333..8f8eb550b 100644 --- a/src/art/trajectories/_tokenize.py +++ b/src/art/trajectories/_tokenize.py @@ -464,6 +464,8 @@ def _anthropic_messages(request: dict[str, Any]) -> list[dict[str, Any]]: ), } ) + else: + raise ValueError(f"Unsupported Anthropic content block type: {kind!r}") message: dict[str, Any] = {"role": role, "content": text} if reasoning: message["reasoning"] = reasoning @@ -485,50 +487,79 @@ def _responses_messages(request: dict[str, Any]) -> list[dict[str, Any]]: if isinstance(value, str): messages.append({"role": "user", "content": value}) elif isinstance(value, list): + pending_reasoning = "" + pending_tool_calls: list[dict[str, Any]] | None = None for item in value: if not isinstance(item, dict): raise ValueError("Responses input items must be JSON objects") kind = item.get("type") - if kind == "function_call_output": - messages.append( - { - "role": "tool", - "tool_call_id": item.get("call_id"), - "content": _responses_input_text( - item.get("output", ""), field="function_call_output" - ), - } - ) - elif kind == "function_call": - messages.append( - { + if kind == "reasoning": + pending_tool_calls = None + reasoning = _responses_reasoning_text(item) + if not reasoning: + raise ValueError("Responses reasoning item has no renderable text") + pending_reasoning += reasoning + continue + if kind == "function_call": + if pending_tool_calls is None: + pending_tool_calls = [] + message: dict[str, Any] = { "role": "assistant", "content": "", - "tool_calls": [ - { - "id": item.get("call_id"), - "type": "function", - "function": { - "name": item.get("name"), - "arguments": item.get("arguments", "{}"), - }, - } - ], + "tool_calls": pending_tool_calls, + } + if pending_reasoning: + message["reasoning"] = pending_reasoning + pending_reasoning = "" + messages.append(message) + pending_tool_calls.append( + { + "id": item.get("call_id"), + "type": "function", + "function": { + "name": item.get("name"), + "arguments": item.get("arguments", "{}"), + }, } ) + continue + pending_tool_calls = None + if kind == "function_call_output": + message: dict[str, Any] = { + "role": "tool", + "tool_call_id": item.get("call_id"), + "content": _responses_input_text( + item.get("output", ""), field="function_call_output" + ), + } elif kind in {None, "message"} and item.get("role"): if item.get("phase") is not None: raise ValueError("Unsupported Responses message phase") - messages.append( - { - "role": item["role"], - "content": _responses_input_text( - item.get("content"), field="message content" - ), - } - ) + message = { + "role": item["role"], + "content": _responses_input_text( + item.get("content"), field="message content" + ), + } else: raise ValueError(f"Unsupported Responses input item type: {kind!r}") + if pending_reasoning: + if message["role"] == "assistant": + message["reasoning"] = pending_reasoning + else: + messages.append( + { + "role": "assistant", + "content": "", + "reasoning": pending_reasoning, + } + ) + pending_reasoning = "" + messages.append(message) + if pending_reasoning: + messages.append( + {"role": "assistant", "content": "", "reasoning": pending_reasoning} + ) elif value is not None: raise ValueError("Responses input must be text or a list of input items") return messages @@ -579,6 +610,19 @@ def _responses_output_text(content: object) -> str: return text +def _responses_reasoning_text(item: Mapping[str, object]) -> str: + text = "" + for field in ("content", "summary"): + blocks = item.get(field) + if not isinstance(blocks, list): + continue + for block in blocks: + data = _string_dict(block) + if data is not None and isinstance(data.get("text"), str): + text += data["text"] + return text + + def _openai_tools(tools: object, *, dialect: str) -> object: if not isinstance(tools, list) or dialect == "chat": return tools @@ -646,6 +690,7 @@ def _response_message( if isinstance(exchange, ResponsesExchange): data = exchange.response.model_dump(mode="python") content = "" + reasoning = "" tool_calls = [] for raw_item in data.get("output") or []: item = _string_dict(raw_item) @@ -656,6 +701,11 @@ def _response_message( if item.get("phase") is not None: raise ValueError("Unsupported Responses message phase") content += _responses_output_text(item.get("content")) + elif kind == "reasoning": + rendered = _responses_reasoning_text(item) + if not rendered: + raise ValueError("Responses reasoning item has no renderable text") + reasoning += rendered elif kind == "function_call": tool_calls.append( { @@ -673,6 +723,8 @@ def _response_message( "role": "assistant", "content": content, } + if reasoning: + message["reasoning"] = reasoning if tool_calls: message["tool_calls"] = tool_calls return message @@ -707,7 +759,7 @@ def _template_ids( messages, tools = _request_messages(exchange, messages_override) if completed: - messages.append(_response_message(exchange)) + messages = [*messages, _response_message(exchange)] request_kwargs = request.get("chat_template_kwargs") kwargs = { **(config.chat_template_kwargs or {}), @@ -790,18 +842,42 @@ def _align_visible_logprobs( values = _visible_logprobs(exchange) if not values or tokenizer is None: return None - aligned = [math.nan] * len(completion) - cursor = 0 + token_ids: list[int] = [] + logprobs: list[float] = [] for text, logprob in values: encoded = _ids(tokenizer(text, add_special_tokens=False)) if len(encoded) != 1: return None + token_ids.append(encoded[0]) + logprobs.append(logprob) + + left: list[int] = [] + cursor = 0 + for token_id in token_ids: try: - index = completion.index(encoded[0], cursor) + index = completion.index(token_id, cursor) except ValueError: return None - aligned[index] = logprob + left.append(index) cursor = index + 1 + + right: list[int] = [] + cursor = len(completion) + for token_id in reversed(token_ids): + while cursor: + cursor -= 1 + if completion[cursor] == token_id: + right.append(cursor) + break + else: + return None + right.reverse() + if left != right: + return None + + aligned = [math.nan] * len(completion) + for index, logprob in zip(left, logprobs, strict=True): + aligned[index] = logprob return aligned @@ -837,6 +913,8 @@ def _legacy_tokenize( logprobs.extend([math.nan] * len(suffix)) assistant_mask.extend([False] * len(suffix)) token_ids.extend(completion) + if len(completion_logprobs) != len(completion): + completion_logprobs = [math.nan] * len(completion) logprobs.extend(completion_logprobs) assistant_mask.extend([True] * len(completion)) if not token_ids: @@ -858,6 +936,14 @@ def tokenize_one( chat_template_kwargs: Mapping[str, object] | None, tokenizer_instance: _Tokenizer | None = None, ) -> TokenizedTrajectory: + if trajectory.exchanges and ( + trajectory.messages_and_choices + or trajectory.tools is not None + or trajectory.additional_histories + ): + raise ValueError( + "A trajectory cannot contain both exchanges and legacy histories" + ) if not trajectory.exchanges: return _legacy_tokenize( trajectory, @@ -874,7 +960,9 @@ def tokenize_one( token_ids: list[int] = [] logprobs: list[float] = [] assistant_mask: list[bool] = [] - response_histories: dict[str, tuple[list[dict[str, Any]], ResponsesExchange]] = {} + response_histories: dict[ + str, tuple[list[dict[str, Any]] | None, ResponsesExchange] + ] = {} for exchange in exchanges: if isinstance(exchange, CompletionsExchange): @@ -891,10 +979,15 @@ def tokenize_one( raise ValueError( "Trajectory tokenization does not support Completions echo=True" ) - messages_override = None + prompt, completion, completion_logprobs = _exchange_tokens(exchange) + messages_override: list[dict[str, Any]] | None = None if isinstance(exchange, ResponsesExchange): request = exchange.request - messages_override = _responses_messages(request) + try: + messages_override = _responses_messages(request) + except ValueError: + if prompt is None: + raise previous = request.get("previous_response_id") if previous is not None: if not isinstance(previous, str) or previous not in response_histories: @@ -902,13 +995,17 @@ def tokenize_one( "Responses exchange refers to a previous response outside this trajectory" ) previous_messages, previous_exchange = response_histories[previous] - messages_override = [ - *previous_messages, - _response_message(previous_exchange), - *messages_override, - ] + 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) - prompt, completion, completion_logprobs = _exchange_tokens(exchange) if prompt is None: if tokenizer is None: tokenizer = _load_tokenizer(config) diff --git a/tests/unit/trajectories/test_capture.py b/tests/unit/trajectories/test_capture.py index e812e3cd9..7b02ca068 100644 --- a/tests/unit/trajectories/test_capture.py +++ b/tests/unit/trajectories/test_capture.py @@ -2,6 +2,7 @@ import asyncio from collections.abc import AsyncGenerator, AsyncIterator, Generator +import copy from datetime import datetime, timedelta import json from typing import Any, cast @@ -26,7 +27,7 @@ ResponsesExchange, _compat, ) -from art.trajectories._capture.core import begin, reset +from art.trajectories._capture.core import _append_exchange, begin, reset from art.trajectories._protocols import Endpoint, build_exchange, endpoint_for_url CHAT: dict[str, Any] = { @@ -230,6 +231,32 @@ def test_sync_group_generator_initializes_once( assert initializer.call_count == 1 +def test_group_pydantic_round_trip_restores_trajectories_and_exceptions() -> None: + group = art.TrajectoryGroup( + [art.Trajectory(reward=1)], + exceptions=[ValueError("boom")], + metadata={"source": "test"}, + ) + payload = group.model_dump() + + restored = art.TrajectoryGroup.model_validate_json(group.model_dump_json()) + + assert isinstance(restored, art.TrajectoryGroup) + assert restored.model_dump() == payload + assert art.TrajectoryGroup.model_validate(payload).model_dump() == payload + assert art.TrajectoryGroup(**payload).model_dump() == payload + + +def test_group_copy_preserves_subclass() -> None: + class Group(art.TrajectoryGroup): + pass + + group = Group([art.Trajectory(reward=1)]) + + assert type(copy.copy(group)) is Group + assert type(copy.deepcopy(group)) is Group + + async def test_httpx_requests_and_aiohttp_capture_once(endpoint_server: str) -> None: body = {"model": "test/model", "messages": [{"role": "user", "content": "hi"}]} @@ -264,6 +291,52 @@ def requests_stream() -> None: ) +async def test_aiohttp_capture_covers_stream_reader_consumption_methods( + endpoint_server: str, +) -> None: + body = {"model": "test/model", "messages": [{"role": "user", "content": "hi"}]} + + with art.Trajectory() as trajectory: + async with aiohttp.ClientSession() as session: + async with session.post( + f"{endpoint_server}/chat/completions", json=body + ) as response: + await response.content.readexactly(1) + await response.content.read() + async with session.post( + f"{endpoint_server}/chat/completions", json=body + ) as response: + await response.content.readuntil(b"}") + await response.content.read() + async with session.post( + f"{endpoint_server}/chat/completions", json=body + ) as response: + async for _chunk, _end_of_http_chunk in response.content.iter_chunks(): + pass + + assert len(trajectory.exchanges.chat_completions) == 3 + + +async def test_requests_session_stream_default_is_preserved( + endpoint_server: str, +) -> None: + body = {"model": "test/model", "messages": [{"role": "user", "content": "hi"}]} + + def consume(trajectory: art.Trajectory) -> None: + session = requests.Session() + session.stream = True + response = session.post( + f"{endpoint_server}/chat/completions", json=body, timeout=5 + ) + assert not trajectory.exchanges + list(response.iter_content()) + + with art.Trajectory() as trajectory: + await asyncio.to_thread(consume, trajectory) + + 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( @@ -672,6 +745,51 @@ def chunk(index: int, content: str) -> dict[str, Any]: ] +def test_streaming_chat_ignores_keepalives_and_azure_prologue() -> None: + chunk = { + "id": "chatcmpl-1", + "object": "chat.completion.chunk", + "created": 1, + "model": "test/model", + "choices": [ + { + "index": 0, + "delta": {"role": "assistant", "content": "hello"}, + "finish_reason": "stop", + "logprobs": None, + } + ], + } + body = _sse( + [ + (None, "ping"), + (None, {"id": "", "object": "", "created": 0, "model": "", "choices": []}), + (None, chunk), + (None, "[DONE]"), + ] + ) + + exchange = build_exchange( + "chat_completions", + {"model": "test/model", "messages": [], "stream": True}, + body, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + assert isinstance(exchange, ChatCompletionsExchange) + assert exchange.response.choices[0].message.content == "hello" + + with pytest.raises(json.JSONDecodeError): + build_exchange( + "chat_completions", + {"model": "test/model", "messages": [], "stream": True}, + _sse([(None, "corrupt"), (None, chunk), (None, "[DONE]")]), + start_time=datetime.now(), + end_time=datetime.now(), + ) + + def test_trajectory_rejects_mixed_representations() -> None: exchange = build_exchange( "chat_completions", @@ -688,6 +806,24 @@ def test_trajectory_rejects_mixed_representations() -> None: ) +def test_capture_does_not_mix_exchange_and_legacy_representations() -> None: + exchange = build_exchange( + "chat_completions", + {"model": "test/model", "messages": []}, + json.dumps(CHAT).encode(), + start_time=datetime.now(), + end_time=datetime.now(), + ) + assert isinstance(exchange, ChatCompletionsExchange) + trajectory = art.Trajectory( + messages_and_choices=[{"role": "user", "content": "hi"}] + ) + + _append_exchange(trajectory, exchange) + + assert not trajectory.exchanges + + def test_metadata_accepts_json_serializable_values() -> None: assert art.Trajectory().model_dump() == {} assert art.Trajectory().model_dump( diff --git a/tests/unit/trajectories/test_history.py b/tests/unit/trajectories/test_history.py index 654b5370d..e201b94c9 100644 --- a/tests/unit/trajectories/test_history.py +++ b/tests/unit/trajectories/test_history.py @@ -295,6 +295,32 @@ def test_completions_history_preserves_exact_tokens_and_sampled_spans() -> None: history.as_chat_completions_history() +def test_completions_history_uses_request_token_ids_and_rejects_echo() -> 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] + + exchange.request["echo"] = True + with pytest.raises(ValueError, match="echo=True"): + trajectory.completions_history() + + +def test_history_rejects_mutated_mixed_representation() -> None: + trajectory = art.Trajectory( + messages_and_choices=[{"role": "user", "content": "hi"}] + ) + trajectory.exchanges.chat_completions.append( + _chat([{"role": "user", "content": "hi"}], "hello") + ) + + with pytest.raises(ValueError, match="both exchanges and legacy histories"): + trajectory.history() + + def test_legacy_messages_delegate_through_history() -> None: trajectory = art.Trajectory( messages_and_choices=[{"role": "user", "content": "hello"}] diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index 21790037b..55c48ec56 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -5,7 +5,7 @@ from types import SimpleNamespace from typing import Any -from anthropic.types import Message +from anthropic.types import ImageBlockParam, Message, MessageParam from openai.types import Completion from openai.types.chat import ChatCompletion from openai.types.chat.chat_completion_token_logprob import ChatCompletionTokenLogprob @@ -436,6 +436,118 @@ def __call__(self, text: str, **kwargs: object) -> SimpleNamespace: assert math.isnan(result.logprobs[2]) +def test_ambiguous_visible_logprobs_fail_closed( + monkeypatch: pytest.MonkeyPatch, +) -> None: + exchange = _chat_exchange([], []) + 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.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 [10, 11, 12, 11] if messages[-1]["role"] == "assistant" else [10] + + def __call__(self, text: str, **kwargs: object) -> SimpleNamespace: + del text, kwargs + return SimpleNamespace(input_ids=[11]) + + monkeypatch.setattr( + "art.trajectories._tokenize._load_tokenizer", lambda _config: Tokenizer() + ) + result = art.tokenize_trajectory( + art.Trajectory(exchanges=TrajectoryExchanges(chat_completions=[exchange])), + base_model="base/model", + ) + + assert result.token_ids == [10, 11, 12, 11] + assert all(math.isnan(logprob) for logprob in result.logprobs[1:]) + + +def test_legacy_logprob_mismatch_fails_closed() -> None: + exchange = _chat_exchange([1], [2, 3]) + choice = exchange.response.choices[0] + assert choice.logprobs is not None + content = choice.logprobs.content + assert content + choice.logprobs = choice.logprobs.model_copy( + update={ + "content": [ + content[0].model_copy( + update={"token": "answer", "bytes": list(b"answer")} + ) + ] + } + ) + + 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) + + +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": { + "type": "base64", + "media_type": "image/png", + "data": "...", + }, + } + message: MessageParam = {"role": "user", "content": [image]} + exchange = MessagesExchange( + request=MessagesRequest( + model="test/model", + messages=[message], + ), + response=response, + start_time=start, + end_time=start + timedelta(seconds=1), + ) + monkeypatch.setattr( + "art.trajectories._tokenize._load_tokenizer", lambda _config: _FakeTokenizer() + ) + + with pytest.raises(ValueError, match="Unsupported Anthropic content block"): + art.tokenize_trajectory( + art.Trajectory(exchanges=TrajectoryExchanges(messages=[exchange])), + base_model="base/model", + ) + + def test_undecodable_visible_token_bytes_fall_back_to_nan( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -633,29 +745,40 @@ class Tokenizer: def apply_chat_template( self, messages: list[dict[str, Any]], **kwargs: object ) -> list[int]: - del messages, kwargs - return [10] + del kwargs + assistant_count = sum( + message["role"] == "assistant" for message in messages + ) + return [10, *range(2, 2 + assistant_count)] monkeypatch.setattr( "art.trajectories._tokenize._load_tokenizer", lambda _config: Tokenizer() ) request_reasoning = _response_exchange("request-reasoning", 2) request_reasoning.request["input"] = [ - {"id": "reasoning-1", "summary": [], "type": "reasoning"} + { + "id": "reasoning-1", + "summary": [{"type": "summary_text", "text": "request thought"}], + "type": "reasoning", + } ] response_reasoning = _response_exchange("response-reasoning", 2) data = response_reasoning.response.model_dump(mode="python") - data["output"] = [{"id": "reasoning-2", "summary": [], "type": "reasoning"}] + data["output"] = [ + { + "id": "reasoning-2", + "summary": [{"type": "summary_text", "text": "response thought"}], + "type": "reasoning", + } + ] + data.pop("raw_output_tokens", None) response_reasoning.response = Response.model_validate(data) - with pytest.raises(ValueError, match="Unsupported Responses input"): - art.tokenize_trajectory( - art.Trajectory( - exchanges=TrajectoryExchanges(responses=[request_reasoning]) - ), - base_model="base/model", - ) + art.tokenize_trajectory( + art.Trajectory(exchanges=TrajectoryExchanges(responses=[request_reasoning])), + base_model="base/model", + ) single = art.Trajectory( exchanges=TrajectoryExchanges(responses=[response_reasoning]) @@ -671,17 +794,100 @@ def apply_chat_template( previous_response_id=response_reasoning.response.id, offset=1, ) - with pytest.raises(ValueError, match="Unsupported Responses output"): + assert art.tokenize_trajectory( + art.Trajectory( + exchanges=TrajectoryExchanges(responses=[response_reasoning, continuation]) + ), + base_model="base/model", + ).token_ids == [10, 2, 3] + + +def test_responses_opaque_reasoning_requires_exact_tokens( + monkeypatch: pytest.MonkeyPatch, +) -> None: + exchange = _response_exchange("opaque-reasoning", 2) + response = exchange.response.model_dump(mode="python") + response["output"] = [ + { + "id": "reasoning-1", + "encrypted_content": "opaque", + "summary": [], + "type": "reasoning", + } + ] + exchange.response = Response.model_validate(response) + monkeypatch.setattr( + "art.trajectories._tokenize._load_tokenizer", lambda _config: _FakeTokenizer() + ) + + assert art.tokenize_trajectory( + art.Trajectory(exchanges=TrajectoryExchanges(responses=[exchange])), + base_model="base/model", + ).token_ids == [10, 2] + + response = exchange.response.model_dump(mode="python") + 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=[response_reasoning, continuation] - ) - ), + art.Trajectory(exchanges=TrajectoryExchanges(responses=[exchange])), base_model="base/model", ) +def test_responses_parallel_function_calls_form_one_assistant_turn( + monkeypatch: pytest.MonkeyPatch, +) -> None: + exchange = _response_exchange("parallel-tools", 2) + exchange.request["input"] = [ + { + "id": "reasoning-1", + "summary": [{"type": "summary_text", "text": "think"}], + "type": "reasoning", + }, + {"type": "function_call", "call_id": "one", "name": "first", "arguments": "{}"}, + { + "type": "function_call", + "call_id": "two", + "name": "second", + "arguments": "{}", + }, + ] + seen: list[list[dict[str, Any]]] = [] + + class Tokenizer(_FakeTokenizer): + def apply_chat_template( + self, messages: list[dict[str, Any]], **kwargs: object + ) -> list[int]: + seen.append(messages) + return super().apply_chat_template(messages, **kwargs) + + monkeypatch.setattr( + "art.trajectories._tokenize._load_tokenizer", lambda _config: Tokenizer() + ) + art.tokenize_trajectory( + art.Trajectory(exchanges=TrajectoryExchanges(responses=[exchange])), + base_model="base/model", + ) + + assistant = seen[0][0] + assert assistant["reasoning"] == "think" + assert [call["function"]["name"] for call in assistant["tool_calls"]] == [ + "first", + "second", + ] + + +def test_tokenization_rejects_mutated_mixed_representation() -> None: + trajectory = art.Trajectory( + messages_and_choices=[{"role": "user", "content": "hi"}] + ) + trajectory.exchanges.chat_completions.append(_chat_exchange([1], [2])) + + with pytest.raises(ValueError, match="both exchanges and legacy histories"): + art.tokenize_trajectory(trajectory) + + def test_responses_previous_response_id_resolves_local_history( monkeypatch: pytest.MonkeyPatch, ) -> None: From b16846bc1a9cc9479a1d56a130999cc5126eec1b Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 15 Jul 2026 16:41:48 +0000 Subject: [PATCH 14/17] fix: reject chat stream error events --- src/art/trajectories/_protocols.py | 17 ++++++++++++++++- tests/unit/trajectories/test_capture.py | 15 +++++++++++++++ 2 files changed, 31 insertions(+), 1 deletion(-) diff --git a/src/art/trajectories/_protocols.py b/src/art/trajectories/_protocols.py index 3f985ec27..3862d234d 100644 --- a/src/art/trajectories/_protocols.py +++ b/src/art/trajectories/_protocols.py @@ -87,7 +87,22 @@ def _chat_response(body: bytes, *, stream: bool) -> ChatCompletion: try: chunk = ChatCompletionChunk.model_validate(payload) except ValidationError: - if not payload.get("object") and not payload.get("choices"): + if ( + set(payload) + <= { + "choices", + "created", + "id", + "model", + "object", + "prompt_filter_results", + "system_fingerprint", + } + and payload.get("object") == "" + and payload.get("id") == "" + and payload.get("model") == "" + and payload.get("choices") == [] + ): continue raise if response is None: diff --git a/tests/unit/trajectories/test_capture.py b/tests/unit/trajectories/test_capture.py index 7b02ca068..2b3d6d9a1 100644 --- a/tests/unit/trajectories/test_capture.py +++ b/tests/unit/trajectories/test_capture.py @@ -789,6 +789,21 @@ def test_streaming_chat_ignores_keepalives_and_azure_prologue() -> None: end_time=datetime.now(), ) + with pytest.raises(ValueError): + build_exchange( + "chat_completions", + {"model": "test/model", "messages": [], "stream": True}, + _sse( + [ + (None, chunk), + (None, {"error": {"message": "failed"}}), + (None, "[DONE]"), + ] + ), + start_time=datetime.now(), + end_time=datetime.now(), + ) + def test_trajectory_rejects_mixed_representations() -> None: exchange = build_exchange( From 734997689fbc0ae9e7bf9fb1b84c470afc27da42 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 15 Jul 2026 21:28:21 +0000 Subject: [PATCH 15/17] fix: declare trajectory serialization dependencies --- pyproject.toml | 2 ++ uv.lock | 4 ++++ 2 files changed, 6 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 253064f16..2ee4a3783 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,7 +8,9 @@ dependencies = [ "aiohttp>=3.10.0", "anthropic>=0.77.0", "openai>=2.14.0", + "pydantic>=2.12", "requests>=2.32.0", + "typing-extensions>=4.13", "typer>=0.15.2", "litellm>=1.71.1,<=1.82.0", "weave>=0.52.24", diff --git a/uv.lock b/uv.lock index 156052456..e725c4c51 100644 --- a/uv.lock +++ b/uv.lock @@ -4765,10 +4765,12 @@ dependencies = [ { name = "nest-asyncio" }, { name = "openai" }, { name = "polars" }, + { name = "pydantic" }, { name = "requests" }, { name = "setproctitle" }, { name = "tblib" }, { name = "typer" }, + { name = "typing-extensions" }, { name = "weave" }, ] @@ -4909,6 +4911,7 @@ requires-dist = [ { name = "pyarrow", marker = "extra == 'backend'", specifier = ">=15.0.0" }, { name = "pyarrow", marker = "extra == 'tinker'", specifier = ">=15.0.0" }, { name = "pybind11", marker = "extra == 'megatron'", specifier = ">=2.13.6" }, + { name = "pydantic", specifier = ">=2.12" }, { name = "pydantic", marker = "extra == 'tinker'", specifier = ">=2.12.5" }, { name = "pytest", marker = "extra == 'backend'", specifier = ">=8.4.1" }, { name = "quack-kernels", marker = "extra == 'megatron'", specifier = "==0.3.7" }, @@ -4935,6 +4938,7 @@ requires-dist = [ { name = "transformers", marker = "extra == 'tinker'", specifier = ">=5.2.0,<=5.5.3" }, { name = "trl", marker = "extra == 'backend'", specifier = "==0.20.0" }, { name = "typer", specifier = ">=0.15.2" }, + { name = "typing-extensions", specifier = ">=4.13" }, { name = "unsloth", marker = "extra == 'backend'", specifier = "==2026.3.3" }, { name = "unsloth-zoo", marker = "extra == 'backend'", specifier = "==2026.3.1" }, { name = "uvicorn", marker = "extra == 'tinker'", specifier = ">=0.35.0" }, From 8c270646730d2785532699f7cb382de0eacc1b73 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 15 Jul 2026 21:41:45 +0000 Subject: [PATCH 16/17] fix: preserve trajectory compatibility semantics --- src/art/rewards/ruler.py | 19 ++++++------- src/art/trajectories/__init__.py | 26 ++++++++++++----- src/art/trajectories/_history.py | 13 +++++++++ src/art/trajectories/_scope.py | 9 ++++-- tests/unit/trajectories/test_capture.py | 12 ++++++++ tests/unit/trajectories/test_history.py | 36 ++++++++++++++++++++++++ tests/unit/trajectories/test_tokenize.py | 12 +++++++- 7 files changed, 105 insertions(+), 22 deletions(-) diff --git a/src/art/rewards/ruler.py b/src/art/rewards/ruler.py index 34aa64668..17233307b 100644 --- a/src/art/rewards/ruler.py +++ b/src/art/rewards/ruler.py @@ -333,21 +333,18 @@ async def ruler_score_group( ) new_trajectories.append(new_traj) - # Extract message lists and preserve original rewards for comparison - histories = [ - trajectory.history().as_chat_completions_history() - for trajectory in new_trajectories - ] - message_lists: list[list[ChatCompletionMessageParam]] = [ - history.messages for history in histories - ] for trajectory in new_trajectories: trajectory.metrics["independent_reward"] = trajectory.reward - # Extract tools from first trajectory (they should all be the same) - tools = histories[0].tools if histories else None - try: + histories = [ + trajectory.history().as_chat_completions_history() + for trajectory in new_trajectories + ] + message_lists: list[list[ChatCompletionMessageParam]] = [ + history.messages for history in histories + ] + tools = histories[0].tools if histories else None # Call the core ruler function to get scores scores = await ruler( message_lists, diff --git a/src/art/trajectories/__init__.py b/src/art/trajectories/__init__.py index 73f4f4c7a..368ded5c3 100644 --- a/src/art/trajectories/__init__.py +++ b/src/art/trajectories/__init__.py @@ -119,7 +119,9 @@ class MessagesRequest(TypedDict, total=False, extra_items=Any): class ChatCompletionsExchange(pydantic.BaseModel): - request: Annotated[ChatCompletionsRequest, pydantic.SkipValidation] + request: Annotated[ + pydantic.SerializeAsAny[ChatCompletionsRequest], pydantic.SkipValidation + ] response: ChatCompletion start_time: datetime end_time: datetime @@ -132,7 +134,9 @@ def model(self) -> str | None: class CompletionsExchange(pydantic.BaseModel): - request: Annotated[CompletionsRequest, pydantic.SkipValidation] + request: Annotated[ + pydantic.SerializeAsAny[CompletionsRequest], pydantic.SkipValidation + ] response: Completion start_time: datetime end_time: datetime @@ -145,7 +149,9 @@ def model(self) -> str | None: class ResponsesExchange(pydantic.BaseModel): - request: Annotated[ResponsesRequest, pydantic.SkipValidation] + request: Annotated[ + pydantic.SerializeAsAny[ResponsesRequest], pydantic.SkipValidation + ] response: Response start_time: datetime end_time: datetime @@ -158,7 +164,9 @@ def model(self) -> str | None: class MessagesExchange(pydantic.BaseModel): - request: Annotated[MessagesRequest, pydantic.SkipValidation] + request: Annotated[ + pydantic.SerializeAsAny[MessagesRequest], pydantic.SkipValidation + ] response: AnthropicMessage start_time: datetime end_time: datetime @@ -205,8 +213,10 @@ def as_chat_completions_history(self) -> ChatCompletionsHistory: class ChatCompletionsHistory(pydantic.BaseModel): model: str | None - messages: Messages - tools: Tools | None = None + messages: Annotated[pydantic.SerializeAsAny[Messages], pydantic.SkipValidation] + tools: Annotated[pydantic.SerializeAsAny[Tools | None], pydantic.SkipValidation] = ( + None + ) chat_template: str | None = None chat_template_kwargs: dict[str, Any] | None = None @@ -371,7 +381,9 @@ def history(self, *, model: str | None = None) -> TrajectoryHistory: return trajectory_history(self, model=model) def messages(self) -> Messages: - return self.history().as_chat_completions_history().messages + from ._history import trajectory_messages + + return trajectory_messages(self) def for_logging(self) -> dict[str, object]: from ._compat import trajectory_for_logging diff --git a/src/art/trajectories/_history.py b/src/art/trajectories/_history.py index f0d915e84..81f537c4d 100644 --- a/src/art/trajectories/_history.py +++ b/src/art/trajectories/_history.py @@ -369,3 +369,16 @@ def trajectory_history( if protocol == "responses": return responses_history(trajectory, model=model) return anthropic_messages_history(trajectory, model=model) + + +def trajectory_messages(trajectory: Trajectory) -> Messages: + if not trajectory.exchanges: + return History( + messages_and_choices=trajectory.messages_and_choices, + tools=trajectory.tools, + ).messages() + return ( + trajectory_history(trajectory, model=None) + .as_chat_completions_history() + .messages + ) diff --git a/src/art/trajectories/_scope.py b/src/art/trajectories/_scope.py index b9f5e275c..392d17ad4 100644 --- a/src/art/trajectories/_scope.py +++ b/src/art/trajectories/_scope.py @@ -46,7 +46,7 @@ def enter_trajectory(trajectory: Trajectory) -> Trajectory: def exit_trajectory( trajectory: Trajectory, _exc_type: type[BaseException] | None, - _exc_value: BaseException | None, + exc_value: BaseException | None, _traceback: TracebackType | None, ) -> None: current = _trajectories.get() @@ -55,8 +55,11 @@ def exit_trajectory( _trajectories.set(current[:-1]) trajectory.finish() group = get_current_trajectory_group(required=False) - if group is not None and all(item is not trajectory for item in group.trajectories): - group.trajectories.append(trajectory) + if group is not None: + if exc_value is not None: + group.exceptions.append(exception_model(exc_value)) + elif all(item is not trajectory for item in group.trajectories): + group.trajectories.append(trajectory) def enter_trajectory_group(group: TrajectoryGroup) -> TrajectoryGroup: diff --git a/tests/unit/trajectories/test_capture.py b/tests/unit/trajectories/test_capture.py index 2b3d6d9a1..eec66da07 100644 --- a/tests/unit/trajectories/test_capture.py +++ b/tests/unit/trajectories/test_capture.py @@ -218,6 +218,18 @@ async def failed() -> art.Trajectory: assert result.exceptions[0].message == "boom" +def test_failed_trajectory_context_records_exception_without_trajectory() -> None: + group = art.TrajectoryGroup() + + with pytest.raises(ValueError, match="boom"): + with group: + with art.Trajectory() as failed: + raise ValueError("boom") + + assert failed not in group.trajectories + assert [exception.message for exception in group.exceptions] == ["boom"] + + def test_sync_group_generator_initializes_once( monkeypatch: pytest.MonkeyPatch, ) -> None: diff --git a/tests/unit/trajectories/test_history.py b/tests/unit/trajectories/test_history.py index e201b94c9..809fa8196 100644 --- a/tests/unit/trajectories/test_history.py +++ b/tests/unit/trajectories/test_history.py @@ -272,6 +272,12 @@ 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 + ) trajectory.exchanges.responses[1].request["previous_response_id"] = "missing" with pytest.raises(ValueError, match="outside this history"): @@ -332,6 +338,19 @@ def test_legacy_messages_delegate_through_history() -> None: trajectory.history(model="test/model") +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"}]) + ], + ) + + assert trajectory.messages() == [{"role": "user", "content": "primary"}] + with pytest.raises(ValueError, match="multiple legacy histories"): + trajectory.history() + + @pytest.mark.asyncio async def test_ruler_accepts_exchange_trajectories( monkeypatch: pytest.MonkeyPatch, @@ -359,3 +378,20 @@ async def score( assert result is not None assert result.trajectories[0].exchanges == trajectory.exchanges assert [message["role"] for message in captured[0]] == ["user", "assistant"] + + +@pytest.mark.asyncio +async def test_ruler_swallow_exceptions_covers_history_projection() -> None: + from art.rewards.ruler import ruler_score_group + + group = art.TrajectoryGroup( + [ + art.Trajectory( + exchanges=TrajectoryExchanges(completions=[_completion([1], [2])]) + ) + ] + ) + + assert await ruler_score_group(group, swallow_exceptions=True) is None + with pytest.raises(ValueError, match="no chat-message structure"): + await ruler_score_group(group) diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index 55c48ec56..8adc42e39 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -587,9 +587,19 @@ def apply_chat_template( def test_json_round_trip_preserves_exchange_types() -> None: + exchange = _chat_exchange([1], [2]) + request: dict[str, Any] = { + "model": "test/model", + "messages": [ + {"role": "assistant", "content": "answer", "reasoning": "thinking"} + ], + } + exchange.request = ChatCompletionsRequest(**request) original = art.Trajectory( - exchanges=TrajectoryExchanges(chat_completions=[_chat_exchange([1], [2])]) + exchanges=TrajectoryExchanges(chat_completions=[exchange]) ) + dumped = original.model_dump(mode="json", warnings="error") + assert dumped["exchanges"]["chat_completions"][0]["request"] == request restored = art.Trajectory.model_validate_json(original.model_dump_json()) assert restored.model_dump(mode="json") == original.model_dump(mode="json") assert isinstance(restored.exchanges.chat_completions[0].response, ChatCompletion) From 41e292fb005edbd58876f913718c1b3c7eb053d8 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 15 Jul 2026 22:38:00 +0000 Subject: [PATCH 17/17] fix: pin matching FlashInfer runtime --- pyproject.toml | 3 ++- uv.lock | 5 ++++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 2ee4a3783..1a33564c0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -50,6 +50,8 @@ megatron = [ "numpy<2", "torch==2.11.0", "flash-attn-4==4.0.0b5", + "flashinfer-cubin==0.6.8.post1", + "flashinfer-python==0.6.8.post1", "ninja>=1.11.1", "quack-kernels==0.3.7", "apex", @@ -168,7 +170,6 @@ conflicts = [ ] override-dependencies = [ "click==8.2.0", - "flashinfer-python==0.6.8.post1", "megatron-core==0.17.0", "numpy<2", "nvidia-resiliency-ext<0.5", diff --git a/uv.lock b/uv.lock index e725c4c51..4c822b7d2 100644 --- a/uv.lock +++ b/uv.lock @@ -47,7 +47,6 @@ conflicts = [[ [manifest] overrides = [ { name = "click", specifier = "==8.2.0" }, - { name = "flashinfer-python", specifier = "==0.6.8.post1" }, { name = "megatron-core", specifier = "==0.17.0" }, { name = "numpy", specifier = "<2" }, { name = "nvidia-resiliency-ext", specifier = "<0.5" }, @@ -4808,6 +4807,8 @@ megatron = [ { name = "apex" }, { name = "deep-ep", marker = "sys_platform == 'linux'" }, { name = "flash-attn-4" }, + { name = "flashinfer-cubin" }, + { name = "flashinfer-python" }, { name = "megatron-bridge" }, { name = "megatron-core" }, { name = "ml-dtypes", marker = "python_full_version < '3.13'" }, @@ -4880,6 +4881,8 @@ requires-dist = [ { name = "duckdb", marker = "extra == 'backend'", specifier = ">=1.0.0" }, { name = "fastapi", marker = "extra == 'tinker'", specifier = ">=0.128.0" }, { name = "flash-attn-4", marker = "extra == 'megatron'", url = "https://files.pythonhosted.org/packages/24/f7/01ee2576ce41f9884d291ee21861ef194afc0b2b1ce3bd175fc7a6e1b133/flash_attn_4-4.0.0b5-py3-none-any.whl" }, + { name = "flashinfer-cubin", marker = "extra == 'megatron'", specifier = "==0.6.8.post1" }, + { name = "flashinfer-python", marker = "extra == 'megatron'", specifier = "==0.6.8.post1" }, { name = "gql", marker = "extra == 'backend'", specifier = ">=4.0.0" }, { name = "hf-xet", marker = "extra == 'backend'", specifier = ">=1.1.0" }, { name = "huggingface-hub", marker = "extra == 'tinker'" },