fix: honor dataset pipeline trace scope in the Python runner

README.md

@@ -300,6 +300,11 @@ bt datasets pipeline push ./pipeline.ts
 bt datasets pipeline run ./pipeline.py --project "<source project>" --limit 100
 ```
 
+The pipeline source `scope` selects what each transform call receives:
+
+- `"span"` (the default) discovers spans and passes `id`, `input`, `output`, `expected`, `metadata`, and `trace`.
+- `"trace"` discovers whole traces and passes only `trace`.
+
 Useful flags:
 
 - `--limit <n>` controls how many source refs to discover.

scripts/dataset-pipeline-runner.py

@@ -332,12 +332,10 @@ def ref_span_row_id(ref: Any) -> str | None:
 
 
 def hydrate_discovery_refs(
-    pipeline: Any,
-    source_override: Any,
+    source: dict[str, Any],
     source_project_id: str,
     refs: list[Any],
 ) -> list[dict[str, Any]]:
-    source = merged_source(pipeline, source_override)
     state = state_for_org(source.get("orgName"))
     candidates: list[dict[str, Any]] = []
     traces_by_root_span_id: dict[str, LocalTrace] = {}
@@ -385,7 +383,13 @@ async def source_row_for_candidate(candidate: dict[str, Any]) -> Any | None:
     raise RuntimeError(f"Source span row {row_id!r} was not found in hydrated trace.")
 
 
-async def transform_args_for_candidate(candidate: dict[str, Any]) -> dict[str, Any]:
+async def transform_args_for_candidate(
+    candidate: dict[str, Any],
+    scope: str,
+) -> dict[str, Any]:
+    if scope == "trace":
+        return {"trace": candidate["trace"]}
+
     row = await source_row_for_candidate(candidate)
     args = {
         "input": span_attr(row, "input"),
@@ -447,13 +451,15 @@ async def transform_refs(
     if max_concurrency <= 0:
         raise RuntimeError("maxConcurrency must be a positive integer.")
     transform = pipeline_transform(pipeline)
-    candidates = hydrate_discovery_refs(pipeline, source_override, source_project_id, refs)
+    source = merged_source(pipeline, source_override)
+    scope = source.get("scope") or "span"
+    candidates = hydrate_discovery_refs(source, source_project_id, refs)
     transformed_rows: list[list[dict[str, Any]]] = [[] for _ in candidates]
     semaphore = asyncio.Semaphore(max_concurrency)
 
     async def run_one(index: int, candidate: dict[str, Any]) -> None:
         async with semaphore:
-            transform_args = await transform_args_for_candidate(candidate)
+            transform_args = await transform_args_for_candidate(candidate, scope)
             result = await call_user_fn(
                 asyncio.get_running_loop(),
                 transform,

scripts/dataset-pipeline-runner.ts

@@ -530,13 +530,11 @@ function refSpanRowId(ref: DiscoveryRef): string | undefined {
 
 async function hydrateDiscoveryRefs(
   braintrust: BraintrustModule,
-  pipeline: DatasetPipelineDefinition,
-  sourceOverride: PipelineSource | undefined,
+  source: PipelineSource,
   sourceProjectId: string,
   refs: unknown[],
 ): Promise<HydratedCandidate[]> {
   requireBraintrustRuntime(braintrust);
-  const source = requirePipelineSource(pipeline, sourceOverride);
   const state = await stateForOrg(braintrust, source.orgName);
   const tracesByRootSpanId = new Map<string, unknown>();
   return refs.map((ref) => {
@@ -680,8 +678,7 @@ async function transformRefs(
   const scope = source.scope ?? "span";
   const candidates = await hydrateDiscoveryRefs(
     braintrust,
-    pipeline,
-    sourceOverride,
+    source,
     sourceProjectId,
     refs,
   );

src/datasets/pipeline-test-fixtures/node/braintrust/index.cjs

@@ -0,0 +1,57 @@
+class OriginalJSONAttachment {
+  constructor() {
+    throw new Error("original JSONAttachment should be shimmed");
+  }
+}
+
+module.exports = {
+  DatasetPipeline(definition) {
+    globalThis.__braintrust_dataset_pipelines ??= [];
+    globalThis.__braintrust_dataset_pipelines.push({
+      ...definition,
+      source: {
+        ...definition.source,
+        scope: definition.source.scope ?? "span",
+      },
+    });
+  },
+  LocalTrace: class {
+    constructor(options) {
+      this.options = options;
+    }
+    getConfiguration() {
+      return { root_span_id: this.options.rootSpanId };
+    }
+    async getSpans() {
+      return [
+        {
+          id: "source-row",
+          span_id: "source-span",
+          input: { prompt: "hello" },
+          output: { answer: "world" },
+          expected: "ok",
+          metadata: { topic: "smoke" },
+        },
+      ];
+    }
+  },
+  _internalGetGlobalState() {
+    return {
+      loggedIn: true,
+      orgName: "source-org",
+      login: async function () {
+        return this;
+      },
+    };
+  },
+  loginToState: async function ({ orgName }) {
+    return {
+      loggedIn: true,
+      orgName,
+      login: async function () {
+        return this;
+      },
+    };
+  },
+  JSONAttachment: OriginalJSONAttachment,
+};

src/datasets/pipeline-test-fixtures/node/braintrust/index.mjs

@@ -0,0 +1,20 @@
+export function DatasetPipeline(definition) {
+  globalThis.__braintrust_dataset_pipelines ??= [];
+  globalThis.__braintrust_dataset_pipelines.push({
+    ...definition,
+    source: {
+      ...definition.source,
+      scope: definition.source.scope ?? "span",
+    },
+  });
+}
+
+export class JSONAttachment {
+  constructor(data, options) {
+    const hook = globalThis.__BT_DATASET_PIPELINE_DEFER_JSON_ATTACHMENT__;
+    if (hook) {
+      return hook(data, options);
+    }
+    throw new Error("dataset pipeline deferred JSON hook was not installed");
+  }
+}

src/datasets/pipeline-test-fixtures/node/braintrust/package.json

@@ -0,0 +1,10 @@
+{
+  "name": "braintrust",
+  "type": "module",
+  "exports": {
+    ".": {
+      "import": "./index.mjs",
+      "require": "./index.cjs"
+    }
+  }
+}

src/datasets/pipeline-test-fixtures/node/json_attachment_pipeline.ts

@@ -0,0 +1,24 @@
+import { DatasetPipeline, JSONAttachment } from "braintrust";
+
+export default DatasetPipeline({
+  name: "ts-json-attachment-smoke",
+  source: { projectName: "source-project", scope: "span" },
+  target: { projectName: "target-project", datasetName: "traces" },
+  transform: (args) => {
+    if (args.id !== "source-row") {
+      throw new Error(`expected source row id, got ${args.id}`);
+    }
+    return {
+      id: undefined,
+      origin: { object_type: "dataset", object_id: "wrong", id: "wrong" },
+      input: {
+        source_id: args.id,
+        source_input: args.input,
+        full_trace: new JSONAttachment(
+          { ok: true, root: args.trace.getConfiguration().root_span_id },
+          { filename: "trace.json", pretty: true },
+        ),
+      },
+    };
+  },
+});

src/datasets/pipeline-test-fixtures/python/braintrust/__init__.py

@@ -0,0 +1 @@
+from .dataset_pipeline import DatasetPipeline

src/datasets/pipeline-test-fixtures/python/braintrust/dataset_pipeline.py

@@ -0,0 +1,12 @@
+_DATASET_PIPELINES = []
+
+
+def DatasetPipeline(name=None, source=None, target=None, transform=None):
+    pipeline = {
+        "name": name,
+        "source": dict(source or {}),
+        "target": dict(target or {}),
+        "transform": transform,
+    }
+    _DATASET_PIPELINES.append(pipeline)
+    return pipeline

src/datasets/pipeline-test-fixtures/python/braintrust/framework.py

@@ -0,0 +1,10 @@
+import inspect
+
+
+# The real SDK dispatches by signature; the fixture pipelines take **kwargs,
+# where that dispatch is equivalent to forwarding every argument.
+async def call_user_fn(loop, fn, **kwargs):
+    result = fn(**kwargs)
+    if inspect.isawaitable(result):
+        return await result
+    return result

src/datasets/pipeline-test-fixtures/python/braintrust/logger.py

@@ -0,0 +1,15 @@
+class _FakeState:
+    def login(self, **kwargs):
+        return self
+
+
+_STATE = _FakeState()
+
+
+def _internal_get_global_state():
+    return _STATE
+
+
+# Imported by the runner at module load; only reached for a cross-org source.
+def login_to_state(org_name=None):
+    return _STATE

src/datasets/pipeline-test-fixtures/python/braintrust/trace.py

@@ -0,0 +1,25 @@
+class LocalTrace:
+    def __init__(
+        self,
+        object_type=None,
+        object_id=None,
+        root_span_id=None,
+        ensure_spans_flushed=None,
+        state=None,
+    ):
+        self.root_span_id = root_span_id
+
+    def get_configuration(self):
+        return {"root_span_id": self.root_span_id}
+
+    async def get_spans(self, include_scorers=False):
+        return [
+            {
+                "id": "source-row",
+                "span_id": "source-span",
+                "input": {"prompt": "hello"},
+                "output": {"answer": "world"},
+                "expected": "ok",
+                "metadata": {"topic": "smoke"},
+            }
+        ]

src/datasets/pipeline-test-fixtures/python/scope_probe_pipeline.py

@@ -0,0 +1,25 @@
+"""Reports the exact transform arg set the runner hands it.
+
+`__SCOPE__` is replaced with the scope under test, so both scopes assert the
+transform-args contract with the same pipeline.
+"""
+
+from braintrust import DatasetPipeline
+
+
+def transform(**kwargs):
+    return {
+        "input": {
+            "args": sorted(kwargs),
+            "span_input": kwargs.get("input"),
+            "root_span_id": kwargs["trace"].get_configuration()["root_span_id"],
+        }
+    }
+
+
+DatasetPipeline(
+    name="py-scope-smoke",
+    source={"project_name": "test-project", "scope": "__SCOPE__"},
+    target={"project_name": "test-target-project", "dataset_name": "test-dataset"},
+    transform=transform,
+)