Add protocol-aware trajectory capture and tokenization

.gitignore

@@ -14,7 +14,7 @@ wandb/
 docs/node_modules/
 dist/
 replays/
-trajectories/
+/trajectories/
 .DS_Store
 .local/
 .claude/settings.local.json

pyproject.toml

@@ -5,7 +5,12 @@ description = "The OpenPipe Agent Reinforcement Training (ART) library"
 readme = "README.md"
 requires-python = ">=3.12"
 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",
@@ -45,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",
@@ -163,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",

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,33 @@
 )
 from .model import Model, TrainableModel
 from .serverless import ServerlessBackend
-from .trajectories import Trajectory, TrajectoryGroup
+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,
+    current_trajectory_group,
+    tokenize_trajectories,
+    tokenize_trajectory,
+    tokenize_trajectory_group,
+    tokenize_trajectory_groups,
+    trajectory,
+    trajectory_group,
+)
 from .types import (
     LocalTrainResult,
     MegatronRuntimeConfig,
@@ -90,6 +115,8 @@
     "dev",
     "auto_trajectory",
     "capture_auto_trajectory",
+    "current_trajectory",
+    "current_trajectory_group",
     "gather_trajectories",
     "gather_trajectory_groups",
     "trajectory_group_batches",
@@ -112,7 +139,26 @@
     "TrainConfig",
     "TrainResult",
     "Trajectory",
+    "TrajectoryExchanges",
     "TrajectoryGroup",
+    "History",
+    "TrajectoryHistory",
+    "ChatCompletionsExchange",
+    "ChatCompletionsHistory",
+    "CompletionsExchange",
+    "CompletionsHistory",
+    "ResponsesExchange",
+    "ResponsesHistory",
+    "MessagesExchange",
+    "AnthropicMessagesHistory",
+    "TokenizedTrajectory",
+    "TokenizedTrajectoryGroup",
+    "trajectory",
+    "trajectory_group",
+    "tokenize_trajectory",
+    "tokenize_trajectories",
+    "tokenize_trajectory_group",
+    "tokenize_trajectory_groups",
     "capture_yielded_trajectory",
     "yield_trajectory",
 ]

src/art/auto_trajectory.py

@@ -1,207 +1,8 @@
-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 .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")
+from .trajectories import (
+    auto_trajectory,  # ty: ignore[deprecated]
+    capture_auto_trajectory,  # ty: ignore[deprecated]
 )
 
-
-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"]

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