Dataset pipelines

js/src/dataset-pipeline.ts

@@ -0,0 +1,101 @@
+import type { ObjectReferenceType as ObjectReference } from "./generated_types";
+import type {
+  BaseMetadata,
+  Dataset,
+  DefaultMetadataType,
+  FullInitDatasetOptions,
+} from "./logger";
+import type { Trace } from "./trace";
+
+export type DatasetPipelineScope = "span" | "trace";
+
+export type DatasetPipelineSource = {
+  projectId?: string;
+  projectName?: string;
+  orgName?: string;
+  filter?: string;
+  scope?: DatasetPipelineScope;
+};
+
+type DatasetPipelineInitDatasetOptions = FullInitDatasetOptions<false>;
+
+export type DatasetPipelineOrigin = ObjectReference;
+
+export type DatasetPipelineTarget = {
+  projectId?: DatasetPipelineInitDatasetOptions["projectId"];
+  projectName?: DatasetPipelineInitDatasetOptions["project"];
+  orgName?: DatasetPipelineInitDatasetOptions["orgName"];
+  datasetName: NonNullable<DatasetPipelineInitDatasetOptions["dataset"]>;
+  description?: DatasetPipelineInitDatasetOptions["description"];
+  metadata?: DatasetPipelineInitDatasetOptions["metadata"];
+};
+
+export type DatasetPipelineRow = Parameters<Dataset["insert"]>[0];
+
+export type DatasetPipelineTransformArgs<
+  Input = unknown,
+  Output = unknown,
+  Expected = unknown,
+  Metadata extends BaseMetadata = DefaultMetadataType,
+> = {
+  input?: Input;
+  output?: Output;
+  expected?: Expected;
+  metadata?: Metadata extends void ? Record<string, unknown> : Metadata;
+  trace?: Trace;
+};
+
+export type DatasetPipelineTransformResult =
+  | DatasetPipelineRow
+  | DatasetPipelineRow[]
+  | null
+  | undefined;
+
+export type DatasetPipelineTransform = (
+  args: DatasetPipelineTransformArgs,
+) => DatasetPipelineTransformResult | Promise<DatasetPipelineTransformResult>;
+
+export type DatasetPipelineDefinition = {
+  name?: string;
+  source: DatasetPipelineSource;
+  transform: DatasetPipelineTransform;
+  target: DatasetPipelineTarget;
+};
+
+const DATASET_PIPELINE_MARKER = "__braintrustDatasetPipeline";
+
+declare global {
+  // eslint-disable-next-line no-var
+  var __braintrust_dataset_pipelines: DatasetPipelineDefinition[] | undefined;
+}
+
+function registry(): DatasetPipelineDefinition[] {
+  if (!globalThis.__braintrust_dataset_pipelines) {
+    globalThis.__braintrust_dataset_pipelines = [];
+  }
+  return globalThis.__braintrust_dataset_pipelines;
+}
+
+export function getRegisteredDatasetPipelines(): DatasetPipelineDefinition[] {
+  return [...registry()];
+}
+
+export function isDatasetPipelineDefinition(
+  value: unknown,
+): value is DatasetPipelineDefinition {
+  return (
+    typeof value === "object" &&
+    value !== null &&
+    DATASET_PIPELINE_MARKER in value
+  );
+}
+
+export function DatasetPipeline(
+  definition: DatasetPipelineDefinition,
+): DatasetPipelineDefinition {
+  const registered = Object.assign(definition, {
+    [DATASET_PIPELINE_MARKER]: true,
+  });
+  registry().push(registered);
+  return registered;
+}

js/src/exports.ts

@@ -215,6 +215,24 @@ export {
   defaultErrorScoreHandler,
 } from "./framework";
 
+export type {
+  DatasetPipelineDefinition,
+  DatasetPipelineOrigin,
+  DatasetPipelineRow,
+  DatasetPipelineScope,
+  DatasetPipelineSource,
+  DatasetPipelineTarget,
+  DatasetPipelineTransform,
+  DatasetPipelineTransformArgs,
+  DatasetPipelineTransformResult,
+} from "./dataset-pipeline";
+
+export {
+  DatasetPipeline,
+  getRegisteredDatasetPipelines,
+  isDatasetPipelineDefinition,
+} from "./dataset-pipeline";
+
 export type {
   CodeOpts,
   CreateProjectOpts,
@@ -246,7 +264,7 @@ export {
 } from "./prompt-schemas";
 
 export type { Trace, SpanData, GetThreadOptions } from "./trace";
-export { SpanFetcher, CachedSpanFetcher } from "./trace";
+export { LocalTrace, SpanFetcher, CachedSpanFetcher } from "./trace";
 
 export type {
   ParentExperimentIds,

js/src/logger.ts

@@ -79,6 +79,7 @@ import {
   type RepoInfoType as RepoInfo,
   type PromptBlockDataType as PromptBlockData,
   type ResponseFormatJsonSchemaType as ResponseFormatJsonSchema,
+  type ObjectReferenceType as ObjectReference,
 } from "./generated_types";
 
 const BRAINTRUST_ATTACHMENT =
@@ -4263,8 +4264,8 @@ export function initDataset<
     legacy,
     _internal_btql,
     resolvedVersion instanceof LazyValue ||
-      normalizedEnvironment !== undefined ||
-      normalizedSnapshotName !== undefined
+    normalizedEnvironment !== undefined ||
+    normalizedSnapshotName !== undefined
       ? {
           ...(resolvedVersion instanceof LazyValue
             ? {
@@ -6041,9 +6042,9 @@ export type WithTransactionId<R> = R & {
 export const DEFAULT_FETCH_BATCH_SIZE = 1000;
 export const MAX_BTQL_ITERATIONS = 10000;
 
-export class ObjectFetcher<RecordType> implements AsyncIterable<
-  WithTransactionId<RecordType>
-> {
+export class ObjectFetcher<RecordType>
+  implements AsyncIterable<WithTransactionId<RecordType>>
+{
   private _fetchedData: WithTransactionId<RecordType>[] | undefined = undefined;
 
   constructor(
@@ -6900,6 +6901,7 @@ export class SpanImpl implements Span {
       const cachedSpan: CachedSpan = {
         input: partialRecord.input,
         output: partialRecord.output,
+        expected: partialRecord.expected,
         metadata: partialRecord.metadata,
         span_id: this._spanId,
         span_parents: this._spanParents,
@@ -7367,6 +7369,7 @@ export class Dataset<
     metadata,
     tags,
     output,
+    origin,
     isMerge,
   }: {
     id: string;
@@ -7375,6 +7378,7 @@ export class Dataset<
     metadata?: Record<string, unknown>;
     tags?: string[];
     output?: unknown;
+    origin?: ObjectReference;
     isMerge?: boolean;
   }): LazyValue<BackgroundLogEvent> {
     return new LazyValue(async () => {
@@ -7389,6 +7393,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,
@@ -7412,6 +7417,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.
@@ -7423,13 +7429,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?: ObjectReference;
   }): string {
     this.validateEvent({ metadata, expected, output, tags });
 
@@ -7442,6 +7450,7 @@ export class Dataset<
         metadata,
         tags,
         output,
+        origin,
         isMerge: false,
       }),
     );

js/src/span-cache.ts

@@ -36,6 +36,7 @@ function canUseSpanCache(): boolean {
 export interface CachedSpan {
   input?: unknown;
   output?: unknown;
+  expected?: unknown;
   metadata?: Record<string, unknown>;
   span_id: string;
   span_parents?: string[];