From 41107e3695108fa3880c705662bc611a66d77e24 Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Fri, 25 Sep 2026 14:26:04 +0000 Subject: [PATCH 1/3] fix: honor dataset pipeline trace scope in the Python runner `scope: "trace"` was silently TypeScript-only. The TS runner threads the source scope into `transformArgsForCandidate`, but the Python runner never read it: `transform_args_for_candidate` always took the span path, so a trace-scoped Python pipeline received `input`/`output`/`expected`/`metadata` as `None` (and, via the SDK's positional fallback in `call_user_fn`, a single-parameter transform received `None` instead of the trace). The Python runner now resolves the merged source scope once per transform batch and passes only `trace` for trace scope, matching the TS runner byte-for-byte. The TS runner's `hydrateDiscoveryRefs` takes the already-merged source so the two stay diffable. Adds one runner smoke test covering both scopes against a stub braintrust package, reusing the spawn/assert helpers with the existing TypeScript runner test, and documents the scope contract in the README. An unknown scope value is left to the `PipelineScope` enum, which already rejects it while parsing the inspect response. Co-Authored-By: Claude Opus 5 --- README.md | 5 + scripts/dataset-pipeline-runner.py | 18 +- scripts/dataset-pipeline-runner.ts | 7 +- src/datasets/pipeline.rs | 265 ++++++++++++++++++++++++++--- 4 files changed, 262 insertions(+), 33 deletions(-) diff --git a/README.md b/README.md index 7f0894c8..5d3d78be 100644 --- a/README.md +++ b/README.md @@ -300,6 +300,11 @@ bt datasets pipeline push ./pipeline.ts bt datasets pipeline run ./pipeline.py --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 ` controls how many source refs to discover. diff --git a/scripts/dataset-pipeline-runner.py b/scripts/dataset-pipeline-runner.py index 24d394f9..44039abb 100644 --- a/scripts/dataset-pipeline-runner.py +++ b/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, diff --git a/scripts/dataset-pipeline-runner.ts b/scripts/dataset-pipeline-runner.ts index 585643b4..45212fa8 100644 --- a/scripts/dataset-pipeline-runner.ts +++ b/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 { requireBraintrustRuntime(braintrust); - const source = requirePipelineSource(pipeline, sourceOverride); const state = await stateForOrg(braintrust, source.orgName); const tracesByRootSpanId = new Map(); 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, ); diff --git a/src/datasets/pipeline.rs b/src/datasets/pipeline.rs index b9fde2a1..ecc85680 100644 --- a/src/datasets/pipeline.rs +++ b/src/datasets/pipeline.rs @@ -2525,32 +2525,15 @@ export default DatasetPipeline({ "sourceProjectId": "source-project-id", "attachmentDir": attachment_dir, }); - let mut child = Command::new("node") + let mut command = Command::new("node"); + command .arg("--experimental-strip-types") .arg(&runner_path) .arg(&pipeline_path) - .current_dir(root.path()) - .env("BT_DATASET_PIPELINE_STAGE", "transform") - .stdin(Stdio::piped()) - .stdout(Stdio::piped()) - .stderr(Stdio::piped()) - .spawn() - .expect("spawn node runner"); - child - .stdin - .as_mut() - .expect("runner stdin") - .write_all(request.to_string().as_bytes()) - .expect("write runner request"); - let output = child.wait_with_output().expect("runner output"); - assert!( - output.status.success(), - "runner failed\nstdout:\n{}\nstderr:\n{}", - String::from_utf8_lossy(&output.stdout), - String::from_utf8_lossy(&output.stderr) - ); + .current_dir(root.path()); + let output = run_pipeline_transform(command, &request); + let response = pipeline_transform_response(&output); - let response: Value = serde_json::from_slice(&output.stdout).expect("runner JSON response"); assert_eq!(response["rowCount"], json!(1)); assert_eq!(response["rows"][0]["id"], json!("source-row")); assert_eq!( @@ -2761,4 +2744,242 @@ export default DatasetPipeline({ pull_artifact.spec_dir.join("transformed.jsonl") ); } + + fn write_fake_python_braintrust_package(root: &Path) -> PathBuf { + let package_root = root.join("fake_site_packages"); + let package_dir = package_root.join("braintrust"); + fs::create_dir_all(&package_dir).expect("create fake braintrust package"); + fs::write( + package_dir.join("__init__.py"), + r#" +from .dataset_pipeline import DatasetPipeline +"#, + ) + .expect("write fake braintrust __init__"); + fs::write( + package_dir.join("dataset_pipeline.py"), + r#" +_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 +"#, + ) + .expect("write fake dataset_pipeline module"); + fs::write( + package_dir.join("framework.py"), + r#" +import inspect + + +# The real SDK dispatches by signature; the pipelines below 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 +"#, + ) + .expect("write fake framework module"); + fs::write( + package_dir.join("logger.py"), + r#" +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 +"#, + ) + .expect("write fake logger module"); + fs::write( + package_dir.join("trace.py"), + r#" +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"}, + } + ] +"#, + ) + .expect("write fake trace module"); + package_root + } + + /// Runs an already-configured runner command through the transform stage. + fn run_pipeline_transform(mut command: Command, request: &Value) -> std::process::Output { + let mut child = command + .env("BT_DATASET_PIPELINE_STAGE", "transform") + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn() + .expect("spawn pipeline runner"); + child + .stdin + .as_mut() + .expect("runner stdin") + .write_all(request.to_string().as_bytes()) + .expect("write runner request"); + child.wait_with_output().expect("runner output") + } + + fn pipeline_transform_response(output: &std::process::Output) -> Value { + assert!( + output.status.success(), + "runner failed\nstdout:\n{}\nstderr:\n{}", + String::from_utf8_lossy(&output.stdout), + String::from_utf8_lossy(&output.stderr) + ); + serde_json::from_slice(&output.stdout).expect("runner JSON response") + } + + /// Writes the Python runner plus a stub `braintrust` package into `root` and + /// returns the command that runs `pipeline_source` against them. + fn python_pipeline_transform_command(root: &Path, pipeline_source: &str) -> Option { + let python = python_runner::resolve_python_interpreter(None, &[])?; + let package_root = write_fake_python_braintrust_package(root); + let runner_path = root.join(PY_RUNNER_FILE); + fs::write(&runner_path, PY_RUNNER_SOURCE).expect("write python runner source"); + let pipeline_path = root.join("pipeline.py"); + fs::write(&pipeline_path, pipeline_source).expect("write pipeline"); + + let mut command = Command::new(python); + command + .arg(&runner_path) + .arg(&pipeline_path) + .current_dir(root) + .env("PYTHONPATH", &package_root); + Some(command) + } + + /// One transform that reports the exact arg set it was handed, so both scopes + /// can assert the contract with the same pipeline. + const SCOPE_PROBE_PIPELINE: &str = r#" +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, +) +"#; + + #[test] + fn python_runner_passes_scoped_transform_args() { + // Trace-scoped discovery emits refs without a span row id; span-scoped refs + // carry one. The span case also overrides the scope the pipeline declares. + let cases = [ + ( + "trace", + json!({ + "refs": [{ "root_span_id": "root-span" }], + "sourceProjectId": "source-project-id", + }), + json!(["trace"]), + Value::Null, + json!("root-span"), + ), + ( + "trace", + json!({ + "refs": [{ "root_span_id": "root-span", "id": "source-row" }], + "sourceProjectId": "source-project-id", + "source": { "projectName": "test-project", "scope": "span" }, + }), + json!(["expected", "id", "input", "metadata", "output", "trace"]), + json!({ "prompt": "hello" }), + json!("source-row"), + ), + ]; + + for (declared_scope, request, expected_args, expected_span_input, expected_id) in cases { + let root = tempfile::tempdir().expect("tempdir"); + let Some(command) = python_pipeline_transform_command( + root.path(), + &SCOPE_PROBE_PIPELINE.replace("__SCOPE__", declared_scope), + ) else { + eprintln!( + "Skipping python_runner_passes_scoped_transform_args (python not installed)." + ); + return; + }; + + let output = run_pipeline_transform(command, &request); + let response = pipeline_transform_response(&output); + + assert_eq!(response["candidates"], json!(1)); + assert_eq!(response["rowCount"], json!(1)); + let row = &response["rows"][0]; + assert_eq!(row["input"]["args"], expected_args); + assert_eq!(row["input"]["span_input"], expected_span_input); + assert_eq!(row["input"]["root_span_id"], json!("root-span")); + assert_eq!(row["id"], expected_id); + } + } + + #[test] + fn pipeline_inspect_rejects_unknown_source_scope() { + let error = serde_json::from_value::(json!({ + "source": { "projectName": "test-project", "scope": "traces" }, + "target": { "projectName": "test-target-project", "datasetName": "test-dataset" }, + })) + .expect_err("unknown scope should not deserialize"); + + assert!( + error.to_string().contains("`span`") && error.to_string().contains("`trace`"), + "error should name the valid scopes: {error}" + ); + } } From 06b69b21277de78a997a1da3d9af918b438104eb Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Sat, 26 Sep 2026 09:11:10 +0000 Subject: [PATCH 2/3] test: move dataset pipeline Python fixtures into real files Addresses review feedback: the stub braintrust package and the scope-probe pipeline are plain .py files under src/datasets/pipeline-test-fixtures/, loaded with include_str! and written out as before, so they can be edited with syntax highlighting instead of as embedded raw strings. Co-Authored-By: Claude Opus 5 --- .../braintrust/__init__.py | 1 + .../braintrust/dataset_pipeline.py | 12 ++ .../braintrust/framework.py | 10 ++ .../braintrust/logger.py | 15 ++ .../braintrust/trace.py | 25 +++ .../scope_probe_pipeline.py | 25 +++ src/datasets/pipeline.rs | 147 ++++-------------- 7 files changed, 119 insertions(+), 116 deletions(-) create mode 100644 src/datasets/pipeline-test-fixtures/braintrust/__init__.py create mode 100644 src/datasets/pipeline-test-fixtures/braintrust/dataset_pipeline.py create mode 100644 src/datasets/pipeline-test-fixtures/braintrust/framework.py create mode 100644 src/datasets/pipeline-test-fixtures/braintrust/logger.py create mode 100644 src/datasets/pipeline-test-fixtures/braintrust/trace.py create mode 100644 src/datasets/pipeline-test-fixtures/scope_probe_pipeline.py diff --git a/src/datasets/pipeline-test-fixtures/braintrust/__init__.py b/src/datasets/pipeline-test-fixtures/braintrust/__init__.py new file mode 100644 index 00000000..649918ad --- /dev/null +++ b/src/datasets/pipeline-test-fixtures/braintrust/__init__.py @@ -0,0 +1 @@ +from .dataset_pipeline import DatasetPipeline diff --git a/src/datasets/pipeline-test-fixtures/braintrust/dataset_pipeline.py b/src/datasets/pipeline-test-fixtures/braintrust/dataset_pipeline.py new file mode 100644 index 00000000..5c938bbf --- /dev/null +++ b/src/datasets/pipeline-test-fixtures/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 diff --git a/src/datasets/pipeline-test-fixtures/braintrust/framework.py b/src/datasets/pipeline-test-fixtures/braintrust/framework.py new file mode 100644 index 00000000..aa0f7089 --- /dev/null +++ b/src/datasets/pipeline-test-fixtures/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 diff --git a/src/datasets/pipeline-test-fixtures/braintrust/logger.py b/src/datasets/pipeline-test-fixtures/braintrust/logger.py new file mode 100644 index 00000000..09c4a3c3 --- /dev/null +++ b/src/datasets/pipeline-test-fixtures/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 diff --git a/src/datasets/pipeline-test-fixtures/braintrust/trace.py b/src/datasets/pipeline-test-fixtures/braintrust/trace.py new file mode 100644 index 00000000..37f0b8c3 --- /dev/null +++ b/src/datasets/pipeline-test-fixtures/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"}, + } + ] diff --git a/src/datasets/pipeline-test-fixtures/scope_probe_pipeline.py b/src/datasets/pipeline-test-fixtures/scope_probe_pipeline.py new file mode 100644 index 00000000..c0b39420 --- /dev/null +++ b/src/datasets/pipeline-test-fixtures/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, +) diff --git a/src/datasets/pipeline.rs b/src/datasets/pipeline.rs index ecc85680..63918105 100644 --- a/src/datasets/pipeline.rs +++ b/src/datasets/pipeline.rs @@ -2745,103 +2745,39 @@ export default DatasetPipeline({ ); } + /// Stub `braintrust` package the Python runner imports, in + /// `pipeline-test-fixtures/braintrust/`. + const STUB_BRAINTRUST_MODULES: [(&str, &str); 5] = [ + ( + "__init__.py", + include_str!("pipeline-test-fixtures/braintrust/__init__.py"), + ), + ( + "dataset_pipeline.py", + include_str!("pipeline-test-fixtures/braintrust/dataset_pipeline.py"), + ), + ( + "framework.py", + include_str!("pipeline-test-fixtures/braintrust/framework.py"), + ), + ( + "logger.py", + include_str!("pipeline-test-fixtures/braintrust/logger.py"), + ), + ( + "trace.py", + include_str!("pipeline-test-fixtures/braintrust/trace.py"), + ), + ]; + fn write_fake_python_braintrust_package(root: &Path) -> PathBuf { let package_root = root.join("fake_site_packages"); let package_dir = package_root.join("braintrust"); fs::create_dir_all(&package_dir).expect("create fake braintrust package"); - fs::write( - package_dir.join("__init__.py"), - r#" -from .dataset_pipeline import DatasetPipeline -"#, - ) - .expect("write fake braintrust __init__"); - fs::write( - package_dir.join("dataset_pipeline.py"), - r#" -_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 -"#, - ) - .expect("write fake dataset_pipeline module"); - fs::write( - package_dir.join("framework.py"), - r#" -import inspect - - -# The real SDK dispatches by signature; the pipelines below 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 -"#, - ) - .expect("write fake framework module"); - fs::write( - package_dir.join("logger.py"), - r#" -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 -"#, - ) - .expect("write fake logger module"); - fs::write( - package_dir.join("trace.py"), - r#" -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"}, - } - ] -"#, - ) - .expect("write fake trace module"); + for (name, source) in STUB_BRAINTRUST_MODULES { + fs::write(package_dir.join(name), source) + .unwrap_or_else(|err| panic!("write fake braintrust {name}: {err}")); + } package_root } @@ -2892,29 +2828,8 @@ class LocalTrace: Some(command) } - /// One transform that reports the exact arg set it was handed, so both scopes - /// can assert the contract with the same pipeline. - const SCOPE_PROBE_PIPELINE: &str = r#" -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, -) -"#; + const SCOPE_PROBE_PIPELINE: &str = + include_str!("pipeline-test-fixtures/scope_probe_pipeline.py"); #[test] fn python_runner_passes_scoped_transform_args() { From 32f9a669ff26902e1fd5a6e9f59ae5c0fec8a175 Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Sat, 26 Sep 2026 09:14:31 +0000 Subject: [PATCH 3/3] test: move dataset pipeline TypeScript fixtures into real files Same treatment as the Python fixtures: the stub braintrust npm package and the JSON-attachment pipeline move out of embedded raw strings into plain files, loaded with include_str!. Fixtures are now grouped by runtime under src/datasets/pipeline-test-fixtures/{python,node}/. Co-Authored-By: Claude Opus 5 --- .../node/braintrust/index.cjs | 57 ++++++ .../node/braintrust/index.mjs | 20 ++ .../node/braintrust/package.json | 10 + .../node/json_attachment_pipeline.ts | 24 +++ .../{ => python}/braintrust/__init__.py | 0 .../braintrust/dataset_pipeline.py | 0 .../{ => python}/braintrust/framework.py | 0 .../{ => python}/braintrust/logger.py | 0 .../{ => python}/braintrust/trace.py | 0 .../{ => python}/scope_probe_pipeline.py | 0 src/datasets/pipeline.rs | 172 ++++-------------- 11 files changed, 150 insertions(+), 133 deletions(-) create mode 100644 src/datasets/pipeline-test-fixtures/node/braintrust/index.cjs create mode 100644 src/datasets/pipeline-test-fixtures/node/braintrust/index.mjs create mode 100644 src/datasets/pipeline-test-fixtures/node/braintrust/package.json create mode 100644 src/datasets/pipeline-test-fixtures/node/json_attachment_pipeline.ts rename src/datasets/pipeline-test-fixtures/{ => python}/braintrust/__init__.py (100%) rename src/datasets/pipeline-test-fixtures/{ => python}/braintrust/dataset_pipeline.py (100%) rename src/datasets/pipeline-test-fixtures/{ => python}/braintrust/framework.py (100%) rename src/datasets/pipeline-test-fixtures/{ => python}/braintrust/logger.py (100%) rename src/datasets/pipeline-test-fixtures/{ => python}/braintrust/trace.py (100%) rename src/datasets/pipeline-test-fixtures/{ => python}/scope_probe_pipeline.py (100%) diff --git a/src/datasets/pipeline-test-fixtures/node/braintrust/index.cjs b/src/datasets/pipeline-test-fixtures/node/braintrust/index.cjs new file mode 100644 index 00000000..a0ad25db --- /dev/null +++ b/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, +}; diff --git a/src/datasets/pipeline-test-fixtures/node/braintrust/index.mjs b/src/datasets/pipeline-test-fixtures/node/braintrust/index.mjs new file mode 100644 index 00000000..f7637920 --- /dev/null +++ b/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"); + } +} diff --git a/src/datasets/pipeline-test-fixtures/node/braintrust/package.json b/src/datasets/pipeline-test-fixtures/node/braintrust/package.json new file mode 100644 index 00000000..f4771614 --- /dev/null +++ b/src/datasets/pipeline-test-fixtures/node/braintrust/package.json @@ -0,0 +1,10 @@ +{ + "name": "braintrust", + "type": "module", + "exports": { + ".": { + "import": "./index.mjs", + "require": "./index.cjs" + } + } +} diff --git a/src/datasets/pipeline-test-fixtures/node/json_attachment_pipeline.ts b/src/datasets/pipeline-test-fixtures/node/json_attachment_pipeline.ts new file mode 100644 index 00000000..2fc7fd56 --- /dev/null +++ b/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 }, + ), + }, + }; + }, +}); diff --git a/src/datasets/pipeline-test-fixtures/braintrust/__init__.py b/src/datasets/pipeline-test-fixtures/python/braintrust/__init__.py similarity index 100% rename from src/datasets/pipeline-test-fixtures/braintrust/__init__.py rename to src/datasets/pipeline-test-fixtures/python/braintrust/__init__.py diff --git a/src/datasets/pipeline-test-fixtures/braintrust/dataset_pipeline.py b/src/datasets/pipeline-test-fixtures/python/braintrust/dataset_pipeline.py similarity index 100% rename from src/datasets/pipeline-test-fixtures/braintrust/dataset_pipeline.py rename to src/datasets/pipeline-test-fixtures/python/braintrust/dataset_pipeline.py diff --git a/src/datasets/pipeline-test-fixtures/braintrust/framework.py b/src/datasets/pipeline-test-fixtures/python/braintrust/framework.py similarity index 100% rename from src/datasets/pipeline-test-fixtures/braintrust/framework.py rename to src/datasets/pipeline-test-fixtures/python/braintrust/framework.py diff --git a/src/datasets/pipeline-test-fixtures/braintrust/logger.py b/src/datasets/pipeline-test-fixtures/python/braintrust/logger.py similarity index 100% rename from src/datasets/pipeline-test-fixtures/braintrust/logger.py rename to src/datasets/pipeline-test-fixtures/python/braintrust/logger.py diff --git a/src/datasets/pipeline-test-fixtures/braintrust/trace.py b/src/datasets/pipeline-test-fixtures/python/braintrust/trace.py similarity index 100% rename from src/datasets/pipeline-test-fixtures/braintrust/trace.py rename to src/datasets/pipeline-test-fixtures/python/braintrust/trace.py diff --git a/src/datasets/pipeline-test-fixtures/scope_probe_pipeline.py b/src/datasets/pipeline-test-fixtures/python/scope_probe_pipeline.py similarity index 100% rename from src/datasets/pipeline-test-fixtures/scope_probe_pipeline.py rename to src/datasets/pipeline-test-fixtures/python/scope_probe_pipeline.py diff --git a/src/datasets/pipeline.rs b/src/datasets/pipeline.rs index 63918105..9a04e014 100644 --- a/src/datasets/pipeline.rs +++ b/src/datasets/pipeline.rs @@ -2380,134 +2380,11 @@ mod tests { } let root = tempfile::tempdir().expect("tempdir"); - let node_modules = root.path().join("node_modules").join("braintrust"); - fs::create_dir_all(&node_modules).expect("create fake braintrust package"); - fs::write( - node_modules.join("package.json"), - r#"{"name":"braintrust","type":"module","exports":{".":{"import":"./index.mjs","require":"./index.cjs"}}}"#, - ) - .expect("write fake package.json"); - fs::write( - node_modules.join("index.cjs"), - r#" -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, -}; -"#, - ) - .expect("write fake braintrust cjs module"); - fs::write( - node_modules.join("index.mjs"), - r#" -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"); - } -} -"#, - ) - .expect("write fake braintrust esm module"); - - let runner_path = root.path().join("dataset-pipeline-runner.ts"); + write_fake_node_braintrust_package(root.path()); + let runner_path = root.path().join(RUNNER_FILE); fs::write(&runner_path, RUNNER_SOURCE).expect("write runner source"); let pipeline_path = root.path().join("pipeline.ts"); - fs::write( - &pipeline_path, - r#" -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 }, - ), - }, - }; - }, -}); -"#, - ) - .expect("write pipeline"); + fs::write(&pipeline_path, JSON_ATTACHMENT_PIPELINE).expect("write pipeline"); let attachment_dir = root.path().join("attachments"); let request = json!({ @@ -2746,30 +2623,59 @@ export default DatasetPipeline({ } /// Stub `braintrust` package the Python runner imports, in - /// `pipeline-test-fixtures/braintrust/`. + /// `pipeline-test-fixtures/python/braintrust/`. const STUB_BRAINTRUST_MODULES: [(&str, &str); 5] = [ ( "__init__.py", - include_str!("pipeline-test-fixtures/braintrust/__init__.py"), + include_str!("pipeline-test-fixtures/python/braintrust/__init__.py"), ), ( "dataset_pipeline.py", - include_str!("pipeline-test-fixtures/braintrust/dataset_pipeline.py"), + include_str!("pipeline-test-fixtures/python/braintrust/dataset_pipeline.py"), ), ( "framework.py", - include_str!("pipeline-test-fixtures/braintrust/framework.py"), + include_str!("pipeline-test-fixtures/python/braintrust/framework.py"), ), ( "logger.py", - include_str!("pipeline-test-fixtures/braintrust/logger.py"), + include_str!("pipeline-test-fixtures/python/braintrust/logger.py"), ), ( "trace.py", - include_str!("pipeline-test-fixtures/braintrust/trace.py"), + include_str!("pipeline-test-fixtures/python/braintrust/trace.py"), ), ]; + /// Stub `braintrust` package the TypeScript runner imports, in + /// `pipeline-test-fixtures/node/braintrust/`. + const STUB_NODE_BRAINTRUST_FILES: [(&str, &str); 3] = [ + ( + "package.json", + include_str!("pipeline-test-fixtures/node/braintrust/package.json"), + ), + ( + "index.cjs", + include_str!("pipeline-test-fixtures/node/braintrust/index.cjs"), + ), + ( + "index.mjs", + include_str!("pipeline-test-fixtures/node/braintrust/index.mjs"), + ), + ]; + + const JSON_ATTACHMENT_PIPELINE: &str = + include_str!("pipeline-test-fixtures/node/json_attachment_pipeline.ts"); + + fn write_fake_node_braintrust_package(root: &Path) { + let package_dir = root.join("node_modules").join("braintrust"); + fs::create_dir_all(&package_dir).expect("create fake braintrust package"); + for (name, source) in STUB_NODE_BRAINTRUST_FILES { + fs::write(package_dir.join(name), source) + .unwrap_or_else(|err| panic!("write fake braintrust {name}: {err}")); + } + } + fn write_fake_python_braintrust_package(root: &Path) -> PathBuf { let package_root = root.join("fake_site_packages"); let package_dir = package_root.join("braintrust"); @@ -2829,7 +2735,7 @@ export default DatasetPipeline({ } const SCOPE_PROBE_PIPELINE: &str = - include_str!("pipeline-test-fixtures/scope_probe_pipeline.py"); + include_str!("pipeline-test-fixtures/python/scope_probe_pipeline.py"); #[test] fn python_runner_passes_scoped_transform_args() {