feat: Add Dataset pipelines (experimental)

.changeset/dataset-pipeline.md

@@ -0,0 +1,5 @@
+---
+"braintrust": patch
+---
+
+feat: Add Dataset pipelines (experimental)

js/src/dataset-pipeline.ts

@@ -0,0 +1,148 @@
+import type { Trace } from "./trace";
+
+type DatasetPipelineRow = {
+  input?: unknown;
+  output?: unknown;
+  expected?: unknown;
+  tags?: string[];
+  metadata?: Record<string, unknown>;
+  id?: string;
+};
+
+type DatasetPipelineTransformResult =
+  | DatasetPipelineRow
+  | DatasetPipelineRow[]
+  | null
+  | undefined;
+
+type DatasetPipelineDefinition<Scope extends "span" | "trace"> = {
+  /** The name of the dataset pipeline as it will show up in Braintrust */
+  name?: string;
+  /** Information about what spans/traces should be passed into the dataset pipeline. */
+  source: {
+    /** What project to take spans/traces from. Has precedence over `projectName`. */
+    projectId?: string;
+    /** What project to take spans/traces from. */
+    projectName?: string;
+    /** What organization to take spans/traces from. */
+    orgName?: string;
+    /** An optional BTQL filter to filter spans by. Not providing this filter means all spans/traces are eligible for the pipeline. */
+    filter?: string;
+    /**
+     * Whether to pass exclusively spans or entire traces to the pipeline. Affects the transform input arguments.
+     *
+     * Defaults to: `"span"`
+     */
+    scope?: Scope;
+  };
+  /**
+   * A transformation function that either receives a span or a trace, depending on what scope was defined in the `source.scope` option.
+   *
+   * Can return one or more new rows for the target dataset, or `null`/`undefined` if no new row should be inserted.
+   */
+  transform: (
+    transformInput: Scope extends "span"
+      ? {
+          input: unknown;
+          output: unknown;
+          expected: unknown;
+          metadata?: Record<string, unknown>;
+          trace: Trace;
+        }
+      : { trace: Trace },
+  ) => DatasetPipelineTransformResult | Promise<DatasetPipelineTransformResult>;
+  /** Information about the target dataset */
+  target: {
+    /** Id of the project where the dataset currently lives or should be created or updated. */
+    projectId?: string;
+    /** Name of the project where the dataset currently lives or should be created or updated. */
+    projectName?: string;
+    /** Organization name of the project where the dataset currently lives or should be created or updated. */
+    orgName?: string;
+    /** Name of the dataset. Either the current name or new name if the dataset is created or updated. */
+    datasetName: string;
+    /** Description of the dataset when the dataset is created or updated. */
+    description?: string;
+    /** Metadata of the dataset when the dataset is created or updated. */
+    metadata?: Record<string, unknown>;
+  };
+};
+
+/**
+ * This is the interface for pipelines that is exposed to `bt`
+ */
+type DatasetPipelineBtDefinition = {
+  name?: string;
+  source: {
+    projectId?: string;
+    projectName?: string;
+    orgName?: string;
+    filter?: string;
+    scope: "span" | "trace";
+  };
+  transform: (
+    transformInput:
+      | {
+          input: unknown;
+          output: unknown;
+          expected: unknown;
+          metadata?: Record<string, unknown>;
+          trace: Trace;
+        }
+      | { trace: Trace },
+  ) => DatasetPipelineTransformResult | Promise<DatasetPipelineTransformResult>;
+  target: {
+    projectId?: string;
+    projectName?: string;
+    orgName?: string;
+    datasetName: string;
+    description?: string;
+    metadata?: Record<string, unknown>;
+  };
+};
+
+declare global {
+  // DO NOT CHANGE THE NAME OR INTERFACE OF THIS GLOBAL IN A NON-BACKWARDS COMPATIBLE WAY: `bt` CLI depends on it
+  var __braintrust_dataset_pipelines: DatasetPipelineBtDefinition[] | undefined;
+}
+
+/**
+ * Creates a runnable dataset pipeline.
+ *
+ * Dataset pipelines can be used to take trace data stored in Braintrust, filter and transform it, and directly feed it back into a Braintrust dataset.
+ *
+ * You can run a dataset pipeline with the `bt` CLI using `bt datasets pipeline run some-file-path.ts --limit 100`.
+ * The limit option controls how many spans/traces (depending on the `definition.source.scope` option) are discovered for the pipeline.
+ *
+ * @experimental - The API for this function is not yet stabilized and may change or be removed across non-major versions. Functionality is not guaranteed.
+ */
+export function DatasetPipeline<Scope extends "span" | "trace">(
+  definition: DatasetPipelineDefinition<Scope>,
+): void {
+  if (!globalThis.__braintrust_dataset_pipelines) {
+    globalThis.__braintrust_dataset_pipelines = [];
+  }
+
+  const storedDefinition: DatasetPipelineBtDefinition = {
+    name: definition.name,
+    source: {
+      projectId: definition.source.projectId,
+      projectName: definition.source.projectName,
+      orgName: definition.source.orgName,
+      filter: definition.source.filter,
+      scope: definition.source.scope ?? "span",
+    },
+    // eslint-disable-next-line @typescript-eslint/consistent-type-assertions, @typescript-eslint/no-explicit-any
+    transform: definition.transform as any,
+    target: {
+      projectId: definition.target.projectId,
+      projectName: definition.target.projectName,
+      orgName: definition.target.orgName,
+      datasetName: definition.target.datasetName,
+      description: definition.target.description,
+      metadata: definition.target.metadata,
+    },
+  };
+
+  globalThis.__braintrust_dataset_pipelines.push(storedDefinition);
+}

js/src/exports.ts

@@ -225,6 +225,8 @@ export {
   defaultErrorScoreHandler,
 } from "./framework";
 
+export { DatasetPipeline } from "./dataset-pipeline";
+
 export type {
   CodeOpts,
   CreateProjectOpts,

js/src/logger-json-attachment.test.ts

@@ -40,6 +40,49 @@ describe("JSONAttachment", () => {
     expect(attachment.reference.filename).toBe("custom.json");
   });
 
+  it("should defer to the dataset pipeline hook when installed", () => {
+    const globalWithHook = globalThis as typeof globalThis & {
+      __BT_DATASET_PIPELINE_DEFER_JSON_ATTACHMENT__?: (
+        data: unknown,
+        options?: { filename?: string; pretty?: boolean },
+      ) => object;
+    };
+    const previous =
+      globalWithHook.__BT_DATASET_PIPELINE_DEFER_JSON_ATTACHMENT__;
+    try {
+      globalWithHook.__BT_DATASET_PIPELINE_DEFER_JSON_ATTACHMENT__ = (
+        data,
+        options,
+      ) => {
+        const reference = {
+          type: "braintrust_deferred_attachment",
+          kind: "json",
+          filename: options?.filename,
+          content_type: "application/json",
+          pretty: options?.pretty,
+          data,
+        };
+        return { reference };
+      };
+
+      const attachment = new JSONAttachment(
+        { test: "data" },
+        { filename: "trace.json", pretty: true },
+      );
+
+      expect(attachment.reference).toEqual({
+        type: "braintrust_deferred_attachment",
+        kind: "json",
+        filename: "trace.json",
+        content_type: "application/json",
+        pretty: true,
+        data: { test: "data" },
+      });
+    } finally {
+      globalWithHook.__BT_DATASET_PIPELINE_DEFER_JSON_ATTACHMENT__ = previous;
+    }
+  });
+
   it("should pretty print when requested", async () => {
     const testData = { a: 1, b: 2 };
     const attachment = new JSONAttachment(testData, { pretty: true });

js/src/logger.ts

@@ -79,6 +79,7 @@ import {
   type RepoInfoType as RepoInfo,
   type PromptBlockDataType as PromptBlockData,
   type ResponseFormatJsonSchemaType as ResponseFormatJsonSchema,
+  type ObjectReferenceType,
 } from "./generated_types";
 
 const BRAINTRUST_ATTACHMENT =
@@ -1392,6 +1393,20 @@ export abstract class BaseAttachment {
   abstract debugInfo(): Record<string, unknown>;
 }
 
+type DatasetPipelineDeferredJSONAttachmentHook = (
+  data: unknown,
+  options?: { filename?: string; pretty?: boolean },
+) => object;
+
+declare global {
+  // Set by the bt dataset pipeline runner so JSONAttachment can be represented
+  // as a destination-uploaded marker during transform.
+  // eslint-disable-next-line no-var
+  var __BT_DATASET_PIPELINE_DEFER_JSON_ATTACHMENT__:
+    | DatasetPipelineDeferredJSONAttachmentHook
+    | undefined;
+}
+
 /**
  * Represents an attachment to be uploaded and the associated metadata.
  * `Attachment` objects can be inserted anywhere in an event, allowing you to
@@ -1857,6 +1872,20 @@ export class JSONAttachment extends Attachment {
     },
   ) {
     const { filename = "data.json", pretty = false, state } = options ?? {};
+    const deferredJsonAttachment =
+      globalThis.__BT_DATASET_PIPELINE_DEFER_JSON_ATTACHMENT__;
+    if (deferredJsonAttachment) {
+      super({
+        data: new Blob([]),
+        filename,
+        contentType: "application/json",
+        state,
+      });
+      return deferredJsonAttachment(data, {
+        filename,
+        pretty,
+      }) as unknown as JSONAttachment;
+    }
 
     // Serialize the JSON data
     const jsonString = pretty
@@ -6895,9 +6924,15 @@ export class SpanImpl implements Span {
       const cachedSpan: CachedSpan = {
         input: partialRecord.input,
         output: partialRecord.output,
+        expected: partialRecord.expected,
+        error: partialRecord.error,
+        scores: partialRecord.scores,
+        metrics: partialRecord.metrics,
         metadata: partialRecord.metadata,
+        tags: partialRecord.tags,
         span_id: this._spanId,
         span_parents: this._spanParents,
+        is_root: this._spanId === this._rootSpanId,
         span_attributes: partialRecord.span_attributes,
       };
       this._state.spanCache.queueWrite(
@@ -7362,6 +7397,7 @@ export class Dataset<
     metadata,
     tags,
     output,
+    origin,
     isMerge,
   }: {
     id: string;
@@ -7370,6 +7406,7 @@ export class Dataset<
     metadata?: Record<string, unknown>;
     tags?: string[];
     output?: unknown;
+    origin?: ObjectReferenceType;
     isMerge?: boolean;
   }): LazyValue<BackgroundLogEvent> {
     return new LazyValue(async () => {
@@ -7384,6 +7421,7 @@ export class Dataset<
         dataset_id,
         created: !isMerge ? new Date().toISOString() : undefined, //if we're merging/updating an event we will not add this ts
         metadata,
+        origin,
         ...(!!isMerge
           ? {
               [IS_MERGE_FIELD]: true,
@@ -7407,6 +7445,7 @@ export class Dataset<
    * about anything else that's relevant, that you can use to help find and analyze examples later. For example, you could log the
    * `prompt`, example's `id`, or anything else that would be useful to slice/dice later. The values in `metadata` can be any
    * JSON-serializable type, but its keys must be strings.
+   * @param event.origin (Optional) a reference to the source object this dataset record was derived from.
    * @param event.id (Optional) a unique identifier for the event. If you don't provide one, Braintrust will generate one for you.
    * @param event.output: (Deprecated) The output of your application. Use `expected` instead.
    * @returns The `id` of the logged record.
@@ -7418,13 +7457,15 @@ export class Dataset<
     tags,
     id,
     output,
+    origin,
   }: {
     readonly input?: unknown;
     readonly expected?: unknown;
     readonly tags?: string[];
     readonly metadata?: Record<string, unknown>;
     readonly id?: string;
     readonly output?: unknown;
+    readonly origin?: ObjectReferenceType;
   }): string {
     this.validateEvent({ metadata, expected, output, tags });
 
@@ -7437,6 +7478,7 @@ export class Dataset<
         metadata,
         tags,
         output,
+        origin,
         isMerge: false,
       }),
     );

js/src/span-cache.ts

@@ -36,9 +36,15 @@ function canUseSpanCache(): boolean {
 export interface CachedSpan {
   input?: unknown;
   output?: unknown;
+  expected?: unknown;
+  error?: unknown;
+  scores?: Record<string, unknown>;
+  metrics?: Record<string, unknown>;
   metadata?: Record<string, unknown>;
+  tags?: string[];
   span_id: string;
   span_parents?: string[];
+  is_root?: boolean | null;
   span_attributes?: {
     name?: string;
     type?: string;