Dataset pipelines (#384)

Author: ankrgyl

py/src/braintrust/__init__.py

@@ -62,6 +62,7 @@ def is_equal(expected, output):
 
 from .audit import *
 from .auto import auto_instrument as auto_instrument
+from .dataset_pipeline import *
 from .framework import *
 from .framework2 import *
 from .functions.invoke import *

py/src/braintrust/dataset_pipeline.py

@@ -0,0 +1,109 @@
+from collections.abc import Awaitable, Sequence
+from dataclasses import dataclass
+from typing import Any, Generic, Literal, Protocol, TypeAlias, TypeVar
+
+from typing_extensions import NotRequired, TypedDict
+
+from .generated_types import ObjectReference
+from .logger import Metadata
+from .trace import Trace
+
+
+__all__ = [
+    "DatasetPipeline",
+    "DatasetPipelineDefinition",
+    "DatasetPipelineRow",
+    "DatasetPipelineScope",
+    "DatasetPipelineSource",
+    "DatasetPipelineTarget",
+    "DatasetPipelineTransform",
+    "DatasetPipelineTransformArgs",
+    "DatasetPipelineTransformResult",
+    "get_registered_dataset_pipelines",
+]
+
+
+DatasetPipelineScope: TypeAlias = Literal["span", "trace"]
+
+
+class DatasetPipelineSource(TypedDict, total=False):
+    project_id: str
+    project_name: str
+    org_name: str
+    filter: str
+    scope: DatasetPipelineScope
+
+
+class DatasetPipelineTarget(TypedDict):
+    dataset_name: str
+    project_id: NotRequired[str]
+    project_name: NotRequired[str]
+    org_name: NotRequired[str]
+    description: NotRequired[str]
+    metadata: NotRequired[Metadata]
+
+
+class DatasetPipelineRow(TypedDict, total=False):
+    id: str
+    input: Any | None
+    expected: Any | None
+    tags: Sequence[str] | None
+    metadata: Metadata | None
+    origin: ObjectReference
+
+
+Row = TypeVar("Row", bound=DatasetPipelineRow, covariant=True)
+
+
+class DatasetPipelineTransformArgs(TypedDict, total=False):
+    input: Any | None
+    output: Any | None
+    metadata: Metadata | None
+    expected: Any | None
+    trace: Trace
+
+
+DatasetPipelineTransformResult: TypeAlias = Row | Sequence[Row] | None
+
+
+class DatasetPipelineTransform(Protocol[Row]):
+    def __call__(
+        self,
+        input: Any | None = None,
+        output: Any | None = None,
+        metadata: Metadata | None = None,
+        expected: Any | None = None,
+        trace: Trace | None = None,
+    ) -> DatasetPipelineTransformResult[Row] | Awaitable[DatasetPipelineTransformResult[Row]]: ...
+
+
+@dataclass(frozen=True)
+class DatasetPipelineDefinition(Generic[Row]):
+    source: DatasetPipelineSource
+    transform: DatasetPipelineTransform[Row]
+    target: DatasetPipelineTarget
+    name: str | None = None
+
+
+_DATASET_PIPELINES: list[DatasetPipelineDefinition[Any]] = []
+
+
+def get_registered_dataset_pipelines() -> list[DatasetPipelineDefinition[Any]]:
+    return list(_DATASET_PIPELINES)
+
+
+def DatasetPipeline(
+    name: str | None = None,
+    *,
+    source: DatasetPipelineSource,
+    transform: DatasetPipelineTransform[DatasetPipelineRow],
+    target: DatasetPipelineTarget,
+) -> DatasetPipelineDefinition[DatasetPipelineRow]:
+    definition = DatasetPipelineDefinition(
+        name=name,
+        source=source.copy(),
+        transform=transform,
+        target=target.copy(),
+    )
+    _DATASET_PIPELINES.append(definition)
+    return definition

py/src/braintrust/test_trace.py

@@ -41,6 +41,35 @@ async def fetch_fn(span_type):
         assert len(result) == 3
         assert {s.span_id for s in result} == {"span-1", "span-2", "span-3"}
 
+    @pytest.mark.asyncio
+    async def test_fetch_preserves_span_result_fields(self):
+        """Test that fetched spans preserve fields needed for full trace attachments."""
+        mock_spans = [
+            make_span(
+                "span-1",
+                "tool",
+                expected={"answer": "ok"},
+                error={"message": "boom"},
+                metrics={"start": 1, "end": 2},
+                scores={"quality": 0},
+                tags=["debug"],
+            )
+        ]
+
+        async def fetch_fn(span_type):
+            del span_type
+            return mock_spans
+
+        fetcher = CachedSpanFetcher(fetch_fn=fetch_fn)
+        result = await fetcher.get_spans()
+
+        assert result[0].expected == {"answer": "ok"}
+        assert result[0].error == {"message": "boom"}
+        assert result[0].metrics == {"start": 1, "end": 2}
+        assert result[0].scores == {"quality": 0}
+        assert result[0].tags == ["debug"]
+        assert result[0].to_dict()["error"] == {"message": "boom"}
+
     @pytest.mark.asyncio
     async def test_fetch_specific_span_types(self):
         """Test fetching specific span types when filter specified."""