Simplify trajectory capture scopes (#779)

* simplify trajectory group capture scopes

* isolate trajectory group input iteration

Author: bradhilton

src/art/__init__.py

@@ -108,7 +108,7 @@
     auto_trajectory,  # ty: ignore[deprecated]
     capture_auto_trajectory,  # ty: ignore[deprecated]
     current_trajectory,
-    current_trajectory_group,
+    no_capture,
     trajectory,
     trajectory_group,
 )
@@ -132,7 +132,7 @@
     "auto_trajectory",
     "capture_auto_trajectory",
     "current_trajectory",
-    "current_trajectory_group",
+    "no_capture",
     "gather_trajectories",
     "gather_trajectory_groups",
     "trajectory_group_batches",

src/art/trajectories/__init__.py

@@ -8,7 +8,7 @@
     Iterator,
     Mapping,
 )
-from contextlib import asynccontextmanager
+from contextlib import AbstractContextManager, asynccontextmanager
 from dataclasses import dataclass
 from datetime import datetime
 from enum import IntFlag
@@ -721,21 +721,6 @@ def __init__(
             logs=logs,
         )
 
-    def __enter__(self) -> TrajectoryGroup:
-        from ._scope import enter_trajectory_group
-
-        return enter_trajectory_group(self)
-
-    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_type, exc_value, traceback)
-
     def __copy__(self) -> TrajectoryGroup:
         from ._compat import copy_trajectory_group
 
@@ -859,20 +844,12 @@ def current_trajectory(*, require: bool = False) -> Trajectory | None:
     return get_current_trajectory(required=require)
 
 
-@overload
-def current_trajectory_group(*, require: Literal[True]) -> TrajectoryGroup: ...
-
-
-@overload
-def current_trajectory_group(
-    *, require: Literal[False] = False
-) -> TrajectoryGroup | None: ...
-
+def no_capture() -> AbstractContextManager[None]:
+    """Hide enclosing trajectory capture while allowing new nested scopes."""
 
-def current_trajectory_group(*, require: bool = False) -> TrajectoryGroup | None:
-    from ._scope import get_current_trajectory_group
+    from ._scope import no_capture as capture_barrier
 
-    return get_current_trajectory_group(required=require)
+    return capture_barrier()
 
 
 async def trajectory(coroutine: Coroutine[Any, Any, object]) -> Trajectory:
@@ -957,7 +934,7 @@ def get_messages(messages_and_choices: MessagesAndChoices) -> Messages:
     "TokenFlag",
     "MetadataValue",
     "current_trajectory",
-    "current_trajectory_group",
+    "no_capture",
     "trajectory",
     "trajectory_group",
     "auto_trajectory",

src/art/trajectories/_scope.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
 import asyncio
-from collections.abc import Coroutine, Iterable
+from collections.abc import Coroutine, Iterable, Iterator
+from contextlib import contextmanager
 import contextvars
 from types import TracebackType
 from typing import Any
@@ -12,9 +13,6 @@
 _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:
@@ -26,15 +24,6 @@ def get_current_trajectory(*, required: bool) -> Trajectory | None:
     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
 
@@ -46,37 +35,25 @@ 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()
     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:
-        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:
-    _groups.set((*_groups.get(), group))
-    return group
 
+@contextmanager
+def no_capture() -> Iterator[None]:
+    """Hide enclosing trajectory capture while allowing new nested scopes."""
 
-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])
+    token = _trajectories.set(())
+    try:
+        yield
+    finally:
+        _trajectories.reset(token)
 
 
 def _require_raw_coroutine(value: object) -> None:
@@ -98,12 +75,16 @@ async def capture_trajectory_group(
     *,
     return_exceptions: bool,
 ) -> TrajectoryGroup:
-    coroutines = list(trajectories)
-    for coroutine in coroutines:
-        _require_raw_coroutine(coroutine)
+    with no_capture():
+        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(await asyncio.gather(*coroutines))
-    results = await asyncio.gather(*coroutines, return_exceptions=True)
+        return TrajectoryGroup(results)
     completed: list[Trajectory] = []
     exceptions: list[PydanticException] = []
     for result in results:

tests/unit/trajectories/test_capture.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import asyncio
-from collections.abc import AsyncGenerator, AsyncIterator, Generator
+from collections.abc import AsyncGenerator, AsyncIterator, Coroutine, Generator
 import copy
 from datetime import datetime, timedelta
 import gzip
@@ -226,44 +226,144 @@ async def child() -> art.Trajectory:
         art.current_trajectory(require=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]
+def test_no_capture_hides_enclosing_trajectory_but_allows_nested_capture() -> None:
+    def capture_exchange() -> None:
+        state, token = begin(
+            "POST",
+            "https://example.test/v1/chat/completions",
+            {"model": "test/model", "messages": []},
+        )
+        reset(token)
+        if state is not None:
+            state.status_code = 200
+            state.add(json.dumps(CHAT).encode())
+            state.finish()
+
+    with art.Trajectory() as outer:
+        capture_exchange()
+        with art.no_capture():
+            assert art.current_trajectory() is None
+            with pytest.raises(RuntimeError, match="No trajectory"):
+                art.current_trajectory(require=True)
+            capture_exchange()
+            with art.Trajectory() as inner:
+                capture_exchange()
+                with art.no_capture():
+                    assert art.current_trajectory() is None
+                    capture_exchange()
+                    with art.Trajectory() as nested:
+                        capture_exchange()
+                    assert art.current_trajectory() is None
+                assert art.current_trajectory() is inner
+                capture_exchange()
+            assert art.current_trajectory() is None
+        assert art.current_trajectory() is outer
+        capture_exchange()
+
+        with pytest.raises(ValueError, match="restore"):
+            with art.no_capture():
+                raise ValueError("restore")
+        assert art.current_trajectory() is outer
+
+    assert len(outer.exchanges.chat_completions) == 2
+    assert len(inner.exchanges.chat_completions) == 2
+    assert len(nested.exchanges.chat_completions) == 1
+
+
+def test_no_capture_uses_the_scope_active_when_a_request_begins() -> None:
+    with art.Trajectory() as trajectory:
+        state, token = begin(
+            "POST",
+            "https://example.test/v1/chat/completions",
+            {"model": "test/model", "messages": []},
+        )
+        reset(token)
+        assert state is not None
+
+        with art.no_capture():
+            state.status_code = 200
+            state.add(json.dumps(CHAT).encode())
+            state.finish()
+            ignored, ignored_token = begin(
+                "POST",
+                "https://example.test/v1/chat/completions",
+                {"model": "test/model", "messages": []},
+            )
+            reset(ignored_token)
+            assert ignored is None
+
+    assert len(trajectory.exchanges.chat_completions) == 1
+
+
+async def test_no_capture_context_is_copied_when_tasks_are_created() -> None:
+    async def current_after_scheduling() -> art.Trajectory | None:
+        await asyncio.sleep(0)
+        return art.current_trajectory()
+
+    with art.Trajectory() as outer:
+        inherited = asyncio.create_task(current_after_scheduling())
+        with art.no_capture():
+            detached = asyncio.create_task(current_after_scheduling())
+            assert await current_after_scheduling() is None
+        assert await inherited is outer
+        assert await detached is None
+
+
+async def test_async_helpers_and_group_aggregation_are_isolated() -> None:
+    def capture_exchange() -> None:
+        state, token = begin(
+            "POST",
+            "https://example.test/v1/chat/completions",
+            {"model": "test/model", "messages": []},
+        )
+        reset(token)
+        if state is not None:
+            state.status_code = 200
+            state.add(json.dumps(CHAT).encode())
+            state.finish()
 
     async def rollout() -> None:
         await asyncio.sleep(0)
+        capture_exchange()
 
     captured = await art.trajectory(rollout())
     assert isinstance(captured, art.Trajectory)
+    assert len(captured.exchanges.chat_completions) == 1
     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)  # ty: ignore[invalid-argument-type]
     await task
 
+    async def unscoped() -> art.Trajectory:
+        assert art.current_trajectory() is None
+        capture_exchange()
+        return art.Trajectory()
+
     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"
-
-
-def test_failed_trajectory_context_records_exception_without_trajectory() -> None:
-    group = art.TrajectoryGroup()
+    def generated() -> Generator[Coroutine[Any, Any, art.Trajectory], None, None]:
+        capture_exchange()
+        yield art.trajectory(rollout())
 
-    with pytest.raises(ValueError, match="boom"):
-        with group:
-            with art.Trajectory() as failed:
-                raise ValueError("boom")
+    with art.Trajectory() as outer:
+        successful = art.trajectory(rollout())
+        result = await art.trajectory_group(
+            [successful, unscoped(), failed()],
+            return_exceptions=True,
+        )
+        generated_result = await art.trajectory_group(generated())
+        with pytest.raises(ValueError, match="boom"):
+            await art.trajectory_group([failed()])
+        assert art.current_trajectory() is outer
 
-    assert failed not in group.trajectories
-    assert [exception.message for exception in group.exceptions] == ["boom"]
+    assert len(result.trajectories) == 2
+    assert len(result.trajectories[0].exchanges.chat_completions) == 1
+    assert not result.trajectories[1].exchanges
+    assert result.exceptions[0].message == "boom"
+    assert len(generated_result.trajectories[0].exchanges.chat_completions) == 1
+    assert not outer.exchanges
 
 
 def test_sync_group_generator_initializes_once(