braintrust 0.0.109 → 0.0.110

package/dist/browser.js

@@ -95,28 +95,77 @@ function mergeDicts(mergeInto, mergeFrom) {
 function constructJsonArray(items) {
   return `[${items.join(",")}]`;
 }
-var batchRecords = function* (allItems, batchSize, maxRequestSize) {
-  while (true) {
-    const items = [];
-    let itemsLen = 0;
-    while (items.length < batchSize && itemsLen < maxRequestSize / 2) {
-      let item = null;
-      if (allItems.length > 0) {
-        item = allItems.pop();
-      } else {
-        break;
+function mapAt(m, k) {
+  const ret = m.get(k);
+  if (ret === void 0) {
+    throw new Error(`Map does not contain key ${k}`);
+  }
+  return ret;
+}
+function depthFirstSearch(args) {
+  var _a;
+  const { graph, firstVisitF, lastVisitF } = args;
+  for (const vs of graph.values()) {
+    for (const v of vs.values()) {
+      if (!graph.has(v)) {
+        throw new Error(`Outgoing vertex ${v} must be a key in the graph`);
       }
-      const itemS = JSON.stringify(item);
-      items.push(itemS);
-      itemsLen += itemS.length;
     }
-    if (items.length === 0) {
-      break;
-    }
-    yield items;
   }
-};
-function generateMergedRowKey(row) {
+  const firstVisitedVertices = /* @__PURE__ */ new Set();
+  const visitationOrder = (_a = args.visitationOrder) != null ? _a : [...graph.keys()];
+  const events = visitationOrder.map((vertex) => ({ eventType: "first", vertex, extras: {} })).reverse();
+  while (events.length) {
+    const { eventType, vertex, extras } = events.pop();
+    if (eventType === "last") {
+      lastVisitF == null ? void 0 : lastVisitF(vertex);
+      continue;
+    }
+    if (firstVisitedVertices.has(vertex)) {
+      continue;
+    }
+    firstVisitedVertices.add(vertex);
+    firstVisitF == null ? void 0 : firstVisitF(vertex, { parentVertex: extras.parentVertex });
+    events.push({ eventType: "last", vertex, extras: {} });
+    mapAt(graph, vertex).forEach((child) => {
+      events.push({
+        eventType: "first",
+        vertex: child,
+        extras: { parentVertex: vertex }
+      });
+    });
+  }
+}
+function undirectedConnectedComponents(graph) {
+  const directedGraph = new Map(
+    [...graph.vertices].map((v) => [v, /* @__PURE__ */ new Set()])
+  );
+  for (const [i, j] of graph.edges) {
+    mapAt(directedGraph, i).add(j);
+    mapAt(directedGraph, j).add(i);
+  }
+  let labelCounter = 0;
+  const vertexLabels = /* @__PURE__ */ new Map();
+  const firstVisitF = (vertex, args) => {
+    const label = (args == null ? void 0 : args.parentVertex) !== void 0 ? mapAt(vertexLabels, args == null ? void 0 : args.parentVertex) : labelCounter++;
+    vertexLabels.set(vertex, label);
+  };
+  depthFirstSearch({ graph: directedGraph, firstVisitF });
+  const output = Array.from({ length: labelCounter }).map(() => []);
+  for (const [vertex, label] of vertexLabels.entries()) {
+    output[label].push(vertex);
+  }
+  return output;
+}
+function topologicalSort(graph, visitationOrder) {
+  const reverseOrdering = [];
+  const lastVisitF = (vertex) => {
+    reverseOrdering.push(vertex);
+  };
+  depthFirstSearch({ graph, lastVisitF, visitationOrder });
+  return reverseOrdering.reverse();
+}
+function generateMergedRowKey(row, useParentIdForId) {
   return JSON.stringify(
     [
       "org_id",
@@ -125,24 +174,22 @@ function generateMergedRowKey(row) {
       "dataset_id",
       "prompt_session_id",
       "log_id",
-      "id"
+      (useParentIdForId != null ? useParentIdForId : false) ? PARENT_ID_FIELD : "id"
     ].map((k) => row[k])
   );
 }
 function mergeRowBatch(rows) {
-  const out = [];
-  const remainingRows = [];
   for (const row of rows) {
     if (row.id === void 0) {
-      out.push(row);
-    } else {
-      remainingRows.push(row);
+      throw new Error(
+        "Logged row is missing an id. This is an internal braintrust error. Please contact us at info@braintrustdata.com for help"
+      );
     }
   }
-  const rowGroups = {};
-  for (const row of remainingRows) {
+  const rowGroups = /* @__PURE__ */ new Map();
+  for (const row of rows) {
     const key = generateMergedRowKey(row);
-    const existingRow = rowGroups[key];
+    const existingRow = rowGroups.get(key);
     if (existingRow !== void 0 && row[IS_MERGE_FIELD]) {
       const preserveNoMerge = !existingRow[IS_MERGE_FIELD];
       mergeDicts(existingRow, row);
@@ -150,11 +197,102 @@ function mergeRowBatch(rows) {
         delete existingRow[IS_MERGE_FIELD];
       }
     } else {
-      rowGroups[key] = row;
+      rowGroups.set(key, row);
+    }
+  }
+  const merged = [...rowGroups.values()];
+  const rowToLabel = new Map(
+    merged.map((r, i) => [generateMergedRowKey(r), i])
+  );
+  const graph = new Map(
+    Array.from({ length: merged.length }).map((_, i) => [i, /* @__PURE__ */ new Set()])
+  );
+  merged.forEach((r, i) => {
+    const parentId = r[PARENT_ID_FIELD];
+    if (!parentId) {
+      return;
+    }
+    const parentRowKey = generateMergedRowKey(
+      r,
+      true
+      /* useParentIdForId */
+    );
+    const parentLabel = rowToLabel.get(parentRowKey);
+    if (parentLabel !== void 0) {
+      mapAt(graph, parentLabel).add(i);
+    }
+  });
+  const connectedComponents = undirectedConnectedComponents({
+    vertices: new Set(graph.keys()),
+    edges: new Set(
+      [...graph.entries()].flatMap(
+        ([k, vs]) => [...vs].map((v) => {
+          const ret = [k, v];
+          return ret;
+        })
+      )
+    )
+  });
+  const buckets = connectedComponents.map(
+    (cc) => topologicalSort(
+      graph,
+      cc
+      /* visitationOrder */
+    )
+  );
+  return buckets.map((bucket) => bucket.map((i) => merged[i]));
+}
+function batchItems(args) {
+  var _a, _b;
+  let { items } = args;
+  const batchMaxNumItems = (_a = args.batchMaxNumItems) != null ? _a : Number.POSITIVE_INFINITY;
+  const batchMaxNumBytes = (_b = args.batchMaxNumBytes) != null ? _b : Number.POSITIVE_INFINITY;
+  const output = [];
+  let nextItems = [];
+  let batchSet = [];
+  let batch = [];
+  let batchLen = 0;
+  function addToBatch(item) {
+    batch.push(item);
+    batchLen += item.length;
+  }
+  function flushBatch() {
+    batchSet.push(batch);
+    batch = [];
+    batchLen = 0;
+  }
+  while (items.length) {
+    for (const bucket of items) {
+      let i = 0;
+      for (const item of bucket) {
+        if (batch.length === 0 || item.length + batchLen < batchMaxNumBytes && batch.length < batchMaxNumItems) {
+          addToBatch(item);
+        } else if (i === 0) {
+          flushBatch();
+          addToBatch(item);
+        } else {
+          break;
+        }
+        ++i;
+      }
+      if (i < bucket.length) {
+        nextItems.push(bucket.slice(i));
+      }
+      if (batchLen >= batchMaxNumBytes || batch.length > batchMaxNumItems) {
+        flushBatch();
+      }
+    }
+    if (batch.length) {
+      flushBatch();
     }
+    if (batchSet.length) {
+      output.push(batchSet);
+      batchSet = [];
+    }
+    items = nextItems;
+    nextItems = [];
   }
-  out.push(...Object.values(rowGroups));
-  return out;
+  return output;
 }
 var DEFAULT_IS_LEGACY_DATASET = true;
 function ensureDatasetRecord(r, legacy) {
@@ -277,13 +415,9 @@ var LazyValue = class {
 // src/logger.ts
 var NoopSpan = class {
   id;
-  span_id;
-  root_span_id;
   kind = "span";
   constructor() {
     this.id = "";
-    this.span_id = "";
-    this.root_span_id = "";
   }
   log(_) {
   }
@@ -761,14 +895,17 @@ var BackgroundLogger = class {
     if (allItems.length === 0) {
       return;
     }
-    const postPromises = [];
-    for (const batch of batchRecords(
-      allItems,
-      batchSize,
-      this.maxRequestSize
-    )) {
-      postPromises.push(
-        (async () => {
+    const allItemsStr = allItems.map(
+      (bucket) => bucket.map((item) => JSON.stringify(item))
+    );
+    const batchSets = batchItems({
+      items: allItemsStr,
+      batchMaxNumItems: batchSize,
+      batchMaxNumBytes: this.maxRequestSize / 2
+    });
+    for (const batchSet of batchSets) {
+      const postPromises = batchSet.map(
+        (batch) => (async () => {
           try {
             await this.submitLogsRequest(batch);
             return { type: "success" };
@@ -777,14 +914,14 @@ var BackgroundLogger = class {
           }
         })()
       );
-    }
-    const results = await Promise.all(postPromises);
-    const failingResultErrors = results.map((r) => r.type === "success" ? void 0 : r.value).filter((r) => r !== void 0);
-    if (failingResultErrors.length) {
-      throw new AggregateError(
-        failingResultErrors,
-        `Encountered the following errors while logging:`
-      );
+      const results = await Promise.all(postPromises);
+      const failingResultErrors = results.map((r) => r.type === "success" ? void 0 : r.value).filter((r) => r !== void 0);
+      if (failingResultErrors.length) {
+        throw new AggregateError(
+          failingResultErrors,
+          `Encountered the following errors while logging:`
+        );
+      }
     }
     if (this.items.length > 0) {
       await this.flushOnce(args);
@@ -1817,16 +1954,24 @@ var SpanImpl = class _SpanImpl {
     this.parentObject = args.parentObject;
     this.parentIds = args.parentIds;
     const id = args.event?.id ?? v4_default();
-    const span_id = v4_default();
-    this.rowIds = {
-      id,
-      span_id,
-      root_span_id: "parentSpanInfo" in args && args.parentSpanInfo?.root_span_id ? args.parentSpanInfo.root_span_id : span_id
-    };
-    if ("parentSpanInfo" in args && args.parentSpanInfo?.span_id) {
-      this.internalData.span_parents = [args.parentSpanInfo.span_id];
-    } else if ("parentId" in args && !isEmpty(args.parentId)) {
-      this.rowIds[PARENT_ID_FIELD] = args.parentId;
+    this.rowIds = { id };
+    const parentSpanInfo = "parentSpanInfo" in args ? args.parentSpanInfo : void 0;
+    const parentId = "parentId" in args ? args.parentId : void 0;
+    if (parentSpanInfo && parentId) {
+      throw new Error(
+        "Only one of parentSpanInfo and parentId may be specified"
+      );
+    }
+    if (parentId) {
+      this.rowIds[PARENT_ID_FIELD] = parentId;
+    } else {
+      this.rowIds.span_id = v4_default();
+      if (parentSpanInfo) {
+        this.rowIds.root_span_id = parentSpanInfo.root_span_id;
+        this.internalData.span_parents = [parentSpanInfo.span_id];
+      } else {
+        this.rowIds.root_span_id = this.rowIds.span_id;
+      }
     }
     this.isMerge = false;
     const { id: _id, ...eventRest } = args.event ?? {};
@@ -1836,37 +1981,29 @@ var SpanImpl = class _SpanImpl {
   get id() {
     return this.rowIds.id;
   }
-  get span_id() {
-    return this.rowIds.span_id;
-  }
-  get root_span_id() {
-    return this.rowIds.root_span_id;
-  }
   log(event) {
     const sanitized = validateAndSanitizeExperimentLogPartialArgs(event);
     let sanitizedAndInternalData = { ...this.internalData };
     mergeDicts(sanitizedAndInternalData, sanitized);
     this.internalData = {};
-    if (sanitizedAndInternalData.metrics?.end) {
-      this.loggedEndTime = sanitizedAndInternalData.metrics?.end;
-    }
-    const parentIds = new LazyValue(async () => {
-      const { kind, ...ids } = await this.parentIds.get();
-      return ids;
-    });
     if (sanitizedAndInternalData.tags && sanitizedAndInternalData.tags.length > 0 && this.rowIds.span_id !== this.rowIds.root_span_id) {
       throw new Error("Tags can only be logged to the root span");
     }
-    const serializedSanitizedAndInternalData = JSON.stringify(
-      sanitizedAndInternalData
-    );
-    sanitizedAndInternalData = JSON.parse(serializedSanitizedAndInternalData);
+    let partialRecord = {
+      ...sanitizedAndInternalData,
+      ...this.rowIds,
+      [IS_MERGE_FIELD]: this.isMerge
+    };
+    if (partialRecord.metrics?.end) {
+      this.loggedEndTime = partialRecord.metrics?.end;
+    }
+    const serializedPartialRecord = JSON.stringify(partialRecord);
+    partialRecord = JSON.parse(serializedPartialRecord);
     const record = new LazyValue(async () => {
+      const { kind, ...parentIds } = await this.parentIds.get();
       return {
-        ...sanitizedAndInternalData,
-        ...this.rowIds,
-        ...await parentIds.get(),
-        [IS_MERGE_FIELD]: this.isMerge
+        ...partialRecord,
+        ...parentIds
       };
     });
     this.bgLogger.log([record]);
@@ -1892,14 +2029,24 @@ var SpanImpl = class _SpanImpl {
     );
   }
   startSpan(args) {
+    const parentId = args?.parentId ?? (this.rowIds[PARENT_ID_FIELD] ? this.id : void 0);
+    const parentSpanInfo = (() => {
+      if (parentId)
+        return void 0;
+      if (!(this.rowIds.span_id && this.rowIds.root_span_id)) {
+        throw new Error("Impossible");
+      }
+      return {
+        span_id: this.rowIds.span_id,
+        root_span_id: this.rowIds.root_span_id
+      };
+    })();
     return new _SpanImpl({
       parentObject: this.parentObject,
       parentIds: this.parentIds,
       bgLogger: this.bgLogger,
-      parentSpanInfo: {
-        span_id: this.rowIds.span_id,
-        root_span_id: this.rowIds.root_span_id
-      },
+      parentSpanInfo,
+      parentId,
       ...args
     });
   }
@@ -2226,7 +2373,9 @@ function wrapChatCompletion(completion) {
     if (params.stream) {
       const startTime = getCurrentUnixTimestamp();
       const ret = await completion(params, options);
-      return new WrapperStream(span, startTime, ret);
+      const wrapperStream = new WrapperStream(span, startTime, ret.iterator());
+      ret.iterator = () => wrapperStream[Symbol.asyncIterator]();
+      return ret;
     } else {
       try {
         const ret = await completion(