feat: add ATIF rollout ingestion (#495)

Author: eric-tramel

docs/assets/recipes/trace_ingestion/agent_rollout_distillation.py

@@ -411,8 +411,9 @@ def build_arg_parser() -> ArgumentParser:
         type=Path,
         default=None,
         help=(
-            "Optional directory containing rollout JSONL files. When omitted, `claude_code` defaults to "
-            "~/.claude/projects and `codex` defaults to ~/.codex/sessions."
+            "Optional directory containing rollout trace files. `atif` expects standalone JSON trajectory files "
+            "and requires `--trace-dir`. When omitted, `claude_code` defaults to ~/.claude/projects and `codex` "
+            "defaults to ~/.codex/sessions."
         ),
     )
     parser.add_argument("--model-alias", type=str, default="nvidia-super")
@@ -459,6 +460,8 @@ def build_seed_source(
     trace_dir: Path | None,
     rollout_format: dd.AgentRolloutFormat,
 ) -> dd.AgentRolloutSeedSource:
+    if rollout_format == dd.AgentRolloutFormat.ATIF and trace_dir is None:
+        raise ValueError("--trace-dir is required when --format atif.")
     seed_source_kwargs: dict[str, str | dd.AgentRolloutFormat] = {"format": rollout_format}
     if trace_dir is not None:
         seed_source_kwargs["path"] = str(trace_dir)

docs/concepts/seed-datasets.md

@@ -171,7 +171,7 @@ Path: {{ relative_path }}
 
 ### 🤖 AgentRolloutSeedSource
 
-Parse agent rollout trace files (e.g. from Claude Code or Codex) into a structured seed dataset. Each trace becomes one seed row with normalized metadata and the full message history, ready for distillation or analysis pipelines.
+Parse agent rollout trace files (e.g. from ATIF, Claude Code, or Codex) into a structured seed dataset. Each trace becomes one seed row with normalized metadata and the full message history, ready for distillation or analysis pipelines.
 
 ```python
 seed_source = dd.AgentRolloutSeedSource(
@@ -181,11 +181,13 @@ seed_source = dd.AgentRolloutSeedSource(
 config_builder.with_seed_dataset(seed_source)
 ```
 
-When `path` is omitted, built-in defaults are used:
+When `path` is omitted, built-in defaults are used for the vendor-native formats:
 
 - **Claude Code** → `~/.claude/projects`
 - **Codex** → `~/.codex/sessions`
 
+ATIF rollouts use standalone `.json` trajectory files and require an explicit `path`.
+
 You can override both the path and file pattern:
 
 ```python
@@ -196,10 +198,19 @@ seed_source = dd.AgentRolloutSeedSource(
 )
 ```
 
+For ATIF trajectories:
+
+```python
+seed_source = dd.AgentRolloutSeedSource(
+    format=dd.AgentRolloutFormat.ATIF,
+    path="my_atif_traces/",
+)
+```
+
 `AgentRolloutSeedSource` exposes a rich set of seeded columns:
 
 - `trace_id` — unique identifier for the trace
-- `source_kind` — the rollout format (e.g. `"claude_code"`, `"codex"`)
+- `source_kind` — the rollout format (e.g. `"atif"`, `"claude_code"`, `"codex"`)
 - `source_path` — full path to the source file
 - `root_session_id` — top-level session identifier
 - `agent_id` — agent identifier (if present)

docs/recipes/cards.md

@@ -107,7 +107,7 @@ Each recipe is a self-contained example that can be run independently.
 
     **Demonstrates:**
 
-    - `AgentRolloutSeedSource` across Claude Code and Codex rollout formats
+    - `AgentRolloutSeedSource` across ATIF, Claude Code, and Codex rollout formats
     - Using normalized trace columns in generation prompts
     - Distilling agent traces into reusable structured records
 

docs/recipes/trace_ingestion/agent_rollout_distillation.md

@@ -1,8 +1,9 @@
 [Download Code :octicons-download-24:](../../assets/recipes/trace_ingestion/agent_rollout_distillation.py){ .md-button download="agent_rollout_distillation.py" }
 
 This recipe ingests built-in agent rollout traces with `AgentRolloutSeedSource(...)`, selecting the format with
-`--format` and optionally overriding the input directory with `--trace-dir`. It works with `claude_code` and `codex`;
-both can use their default locations when `--trace-dir` is omitted. The pipeline turns each imported trace into a
+`--format` and optionally overriding the input directory with `--trace-dir`. It works with `atif`, `claude_code`,
+and `codex`; `atif` expects standalone `.json` trajectory files and requires `--trace-dir`, while `claude_code` and
+`codex` can use their default locations when `--trace-dir` is omitted. The pipeline turns each imported trace into a
 compact task digest, a standalone instruction-response pair for coding-assistant SFT, and a judge-scored quality
 signal you can use for downstream filtering. It supports both full dataset creation and in-memory preview mode via
 `--preview`.

packages/data-designer-config/src/data_designer/config/seed_source.py

@@ -8,7 +8,7 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Literal
 
-from pydantic import BaseModel, Field, PrivateAttr, field_validator
+from pydantic import BaseModel, Field, PrivateAttr, field_validator, model_validator
 from typing_extensions import Self
 
 from data_designer.config.errors import InvalidFilePathError
@@ -192,11 +192,14 @@ def _validate_filesystem_seed_source_file_pattern(value: str | None) -> str | No
 
 
 class AgentRolloutFormat(StrEnum):
+    ATIF = "atif"
     CLAUDE_CODE = "claude_code"
     CODEX = "codex"
 
 
-def get_agent_rollout_format_defaults(fmt: AgentRolloutFormat) -> tuple[str, str]:
+def get_agent_rollout_format_defaults(fmt: AgentRolloutFormat) -> tuple[str | None, str]:
+    if fmt == AgentRolloutFormat.ATIF:
+        return (None, "*.json")
     if fmt == AgentRolloutFormat.CLAUDE_CODE:
         return (get_claude_code_default_path(), "*.jsonl")
     if fmt == AgentRolloutFormat.CODEX:
@@ -215,8 +218,8 @@ class AgentRolloutSeedSource(FileSystemSeedSource):
     path: str | None = Field(
         None,
         description=(
-            "Directory containing agent rollout artifacts. When omitted, built-in defaults are used: "
-            "Claude Code defaults to ~/.claude/projects and Codex defaults to ~/.codex/sessions. "
+            "Directory containing agent rollout artifacts. This field is required for ATIF trajectories. "
+            "When omitted, built-in defaults are used for formats that define one. "
             "Relative paths are resolved from the current working directory when the config is loaded, "
             "not from the config file location."
         ),
@@ -225,16 +228,26 @@ class AgentRolloutSeedSource(FileSystemSeedSource):
     file_pattern: str | None = Field(
         None,
         description=(
-            "Case-sensitive filename pattern used to match agent rollout files. When omitted, defaults to '*.jsonl'."
+            "Case-sensitive filename pattern used to match agent rollout files. When omitted, "
+            "ATIF defaults to '*.json' while Claude Code and Codex default to '*.jsonl'."
         ),
     )
 
+    @model_validator(mode="after")
+    def validate_runtime_path_source(self) -> Self:
+        default_path, _ = get_agent_rollout_format_defaults(self.format)
+        if self.path is None and default_path is None:
+            raise ValueError(f"🛑 AgentRolloutSeedSource.path is required for format {self.format.value!r}.")
+        return self
+
     @property
     def runtime_path(self) -> str:
         if self._runtime_path is not None:
             return self._runtime_path
         default_path, _ = get_agent_rollout_format_defaults(self.format)
         resolved_path = self.path if self.path is not None else default_path
+        if resolved_path is None:
+            raise ValueError(f"🛑 AgentRolloutSeedSource.path is required for format {self.format.value!r}.")
         self._runtime_path = _resolve_filesystem_runtime_path(resolved_path)
         return self._runtime_path
 

packages/data-designer-config/tests/config/test_seed_source.py

@@ -226,6 +226,12 @@ def test_filesystem_seed_sources_reject_path_like_file_patterns(
 @pytest.mark.parametrize(
     ("rollout_format", "file_pattern", "error_message"),
     [
+        pytest.param(
+            AgentRolloutFormat.ATIF,
+            "nested/trace.json",
+            "match file names, not relative paths",
+            id="atif-posix",
+        ),
         pytest.param(
             AgentRolloutFormat.CLAUDE_CODE,
             "",
@@ -250,6 +256,23 @@ def test_agent_rollout_seed_source_rejects_invalid_file_patterns(
         AgentRolloutSeedSource(path=str(tmp_path), file_pattern=file_pattern, format=rollout_format)
 
 
+def test_agent_rollout_seed_source_requires_explicit_atif_path() -> None:
+    with pytest.raises(ValueError, match="path is required for format 'atif'"):
+        AgentRolloutSeedSource(format=AgentRolloutFormat.ATIF)
+
+
+def test_agent_rollout_seed_source_uses_default_atif_file_pattern(tmp_path: Path) -> None:
+    trace_dir = tmp_path / "atif"
+    trace_dir.mkdir()
+
+    source = AgentRolloutSeedSource(path=str(trace_dir), format=AgentRolloutFormat.ATIF)
+
+    assert source.seed_type == "agent_rollout"
+    assert source.resolved_file_pattern == "*.json"
+    assert source.recursive is True
+    assert source.format == AgentRolloutFormat.ATIF
+
+
 def test_agent_rollout_seed_source_uses_default_claude_path(
     monkeypatch: pytest.MonkeyPatch,
     tmp_path: Path,