braintrust 0.0.140 → 0.0.141-dev1

package/dist/index.mjs

@@ -316,6 +316,125 @@ var LazyValue = class {
 // src/logger.ts
 import Mustache from "mustache";
 import { z } from "zod";
+
+// src/stream.ts
+import { callEventSchema } from "@braintrust/core/typespecs";
+import {
+  createParser
+} from "eventsource-parser";
+var BraintrustStream = class _BraintrustStream {
+  stream;
+  constructor(baseStream) {
+    this.stream = baseStream.pipeThrough(btStreamParser());
+  }
+  copy() {
+    const [newStream, copyStream] = this.stream.tee();
+    this.stream = copyStream;
+    return new _BraintrustStream(newStream);
+  }
+  toReadableStream() {
+    return this.stream;
+  }
+};
+function btStreamParser() {
+  const decoder = new TextDecoder();
+  let parser;
+  return new TransformStream({
+    async start(controller) {
+      parser = createParser((event) => {
+        if (event.type === "reconnect-interval") {
+          return;
+        }
+        const parsed = callEventSchema.safeParse(event);
+        if (!parsed.success) {
+          throw new Error(`Failed to parse event: ${parsed.error}`);
+        }
+        switch (event.event) {
+          case "text_delta":
+            controller.enqueue({
+              type: "text_delta",
+              data: JSON.parse(event.data)
+            });
+            break;
+          case "json_delta":
+            controller.enqueue({
+              type: "json_delta",
+              data: event.data
+            });
+            break;
+          case "done":
+            break;
+        }
+      });
+    },
+    async transform(chunk, controller) {
+      if (chunk instanceof Uint8Array) {
+        parser.feed(decoder.decode(chunk));
+      } else if (typeof chunk === "string") {
+        parser.feed(chunk);
+      } else {
+        controller.enqueue(chunk);
+      }
+    },
+    async flush(controller) {
+      controller.terminate();
+    }
+  });
+}
+function createFinalValuePassThroughStream(onFinal) {
+  const decoder = new TextDecoder();
+  const textChunks = [];
+  const jsonChunks = [];
+  const transformStream = new TransformStream({
+    transform(chunk, controller) {
+      if (typeof chunk === "string") {
+        textChunks.push(chunk);
+      } else if (chunk instanceof Uint8Array) {
+        textChunks.push(decoder.decode(chunk));
+      } else {
+        const chunkType = chunk.type;
+        switch (chunkType) {
+          case "text_delta":
+            textChunks.push(chunk.data);
+            break;
+          case "json_delta":
+            jsonChunks.push(chunk.data);
+            break;
+          default:
+            const _type = chunkType;
+            throw new Error(`Unknown chunk type ${_type}`);
+        }
+        controller.enqueue(chunk);
+      }
+    },
+    flush(controller) {
+      if (jsonChunks.length > 0) {
+        onFinal(JSON.parse(jsonChunks.join("")));
+      } else if (textChunks.length > 0) {
+        onFinal(textChunks.join(""));
+      } else {
+        onFinal(void 0);
+      }
+      controller.terminate();
+    }
+  });
+  return transformStream;
+}
+function devNullWritableStream() {
+  return new WritableStream({
+    write(chunk) {
+    },
+    close() {
+    },
+    abort(reason) {
+    },
+    start(controller) {
+    }
+  });
+}
+
+// src/logger.ts
+import { waitUntil } from "@vercel/functions";
 var NoopSpan = class {
   id;
   kind = "span";
@@ -352,6 +471,7 @@ var loginSchema = z.strictObject({
   appPublicUrl: z.string(),
   orgName: z.string(),
   logUrl: z.string(),
+  proxyUrl: z.string(),
   loginToken: z.string(),
   orgId: z.string().nullish(),
   gitMetadataSettings: gitMetadataSettingsSchema.nullish()
@@ -390,11 +510,13 @@ var BraintrustState = class _BraintrustState {
   orgId = null;
   orgName = null;
   logUrl = null;
+  proxyUrl = null;
   loggedIn = false;
   gitMetadataSettings;
   fetch = globalThis.fetch;
   _apiConn = null;
   _logConn = null;
+  _proxyConn = null;
   resetLoginInfo() {
     this.appUrl = null;
     this.appPublicUrl = null;
@@ -402,10 +524,12 @@ var BraintrustState = class _BraintrustState {
     this.orgId = null;
     this.orgName = null;
     this.logUrl = null;
+    this.proxyUrl = null;
     this.loggedIn = false;
     this.gitMetadataSettings = void 0;
     this._apiConn = null;
     this._logConn = null;
+    this._proxyConn = null;
   }
   copyLoginInfo(other) {
     this.appUrl = other.appUrl;
@@ -414,10 +538,12 @@ var BraintrustState = class _BraintrustState {
     this.orgId = other.orgId;
     this.orgName = other.orgName;
     this.logUrl = other.logUrl;
+    this.proxyUrl = other.proxyUrl;
     this.loggedIn = other.loggedIn;
     this.gitMetadataSettings = other.gitMetadataSettings;
     this._apiConn = other._apiConn;
     this._logConn = other._logConn;
+    this._proxyConn = other._proxyConn;
   }
   serialize() {
     if (!this.loggedIn) {
@@ -425,7 +551,7 @@ var BraintrustState = class _BraintrustState {
         "Cannot serialize BraintrustState without being logged in"
       );
     }
-    if (!this.appUrl || !this.appPublicUrl || !this.logUrl || !this.orgName || !this.loginToken || !this.loggedIn) {
+    if (!this.appUrl || !this.appPublicUrl || !this.logUrl || !this.proxyUrl || !this.orgName || !this.loginToken || !this.loggedIn) {
       throw new Error(
         "Cannot serialize BraintrustState without all login attributes"
       );
@@ -437,6 +563,7 @@ var BraintrustState = class _BraintrustState {
       orgId: this.orgId,
       orgName: this.orgName,
       logUrl: this.logUrl,
+      proxyUrl: this.proxyUrl,
       gitMetadataSettings: this.gitMetadataSettings
     };
   }
@@ -459,6 +586,7 @@ var BraintrustState = class _BraintrustState {
     state.logConn().set_token(state.loginToken);
     state.logConn().make_long_lived();
     state.apiConn().set_token(state.loginToken);
+    state.proxyConn().set_token(state.loginToken);
     state.loggedIn = true;
     state.loginReplaceLogConn(state.logConn());
     return state;
@@ -497,6 +625,15 @@ var BraintrustState = class _BraintrustState {
     }
     return this._logConn;
   }
+  proxyConn() {
+    if (!this._proxyConn) {
+      if (!this.proxyUrl) {
+        throw new Error("Must initialize proxyUrl before requesting proxyConn");
+      }
+      this._proxyConn = new HTTPConnection(this.proxyUrl, this.fetch);
+    }
+    return this._proxyConn;
+  }
   bgLogger() {
     return this._bgLogger;
   }
@@ -809,6 +946,11 @@ var Logger = class {
   parentObjectType() {
     return SpanObjectTypeV2.PROJECT_LOGS;
   }
+  triggerWaitUntilFlush() {
+    if (!this.state.bgLogger().syncFlush) {
+      return waitUntil(this.state.bgLogger().flush());
+    }
+  }
   /**
    * Log a single event. The event will be batched and uploaded behind the scenes if `logOptions.asyncFlush` is true.
    *
@@ -834,6 +976,7 @@ var Logger = class {
     this.lastStartTime = span.end();
     const ret = span.id;
     if (this.asyncFlush === true) {
+      this.triggerWaitUntilFlush();
       return ret;
     } else {
       return (async () => {
@@ -861,6 +1004,7 @@ var Logger = class {
       () => span.end()
     );
     if (this.asyncFlush) {
+      this.triggerWaitUntilFlush();
       return ret;
     } else {
       return (async () => {
@@ -1292,6 +1436,8 @@ function init(projectOrOptions, optionalOptions) {
     appUrl,
     apiKey,
     orgName,
+    forceLogin,
+    fetch,
     metadata,
     gitMetadataSettings,
     projectId,
@@ -1309,7 +1455,7 @@ function init(projectOrOptions, optionalOptions) {
     }
     const lazyMetadata2 = new LazyValue(
       async () => {
-        await state.login({ apiKey, appUrl, orgName });
+        await state.login({ apiKey, appUrl, orgName, fetch, forceLogin });
         const args = {
           project_name: project,
           project_id: projectId,
@@ -1480,6 +1626,8 @@ function initDataset(projectOrOptions, optionalOptions) {
     appUrl,
     apiKey,
     orgName,
+    fetch,
+    forceLogin,
     projectId,
     useOutput: legacy,
     state: stateArg
@@ -1490,7 +1638,9 @@ function initDataset(projectOrOptions, optionalOptions) {
       await state.login({
         orgName,
         apiKey,
-        appUrl
+        appUrl,
+        fetch,
+        forceLogin
       });
       const args = {
         org_id: state.orgId,
@@ -1569,6 +1719,7 @@ function initLogger(options = {}) {
     apiKey,
     orgName,
     forceLogin,
+    fetch,
     state: stateArg
   } = options || {};
   const computeMetadataArgs = {
@@ -1582,7 +1733,8 @@ function initLogger(options = {}) {
         orgName,
         apiKey,
         appUrl,
-        forceLogin
+        forceLogin,
+        fetch
       });
       return computeLoggerMetadata(state, computeMetadataArgs);
     }
@@ -1606,6 +1758,8 @@ async function loadPrompt({
   appUrl,
   apiKey,
   orgName,
+  fetch,
+  forceLogin,
   state: stateArg
 }) {
   if (isEmpty(projectName) && isEmpty(projectId)) {
@@ -1618,7 +1772,9 @@ async function loadPrompt({
   state.login({
     orgName,
     apiKey,
-    appUrl
+    appUrl,
+    fetch,
+    forceLogin
   });
   const args = {
     project_name: projectName,
@@ -1661,6 +1817,7 @@ async function login(options = {}) {
   }
   await _globalState.login(options);
   globalThis.__inherited_braintrust_state = _globalState;
+  return _globalState;
 }
 async function loginToState(options = {}) {
   const {
@@ -1699,6 +1856,7 @@ async function loginToState(options = {}) {
   }
   conn.make_long_lived();
   state.apiConn().set_token(apiKey);
+  state.proxyConn().set_token(apiKey);
   state.loginToken = conn.token;
   state.loggedIn = true;
   state.loginReplaceLogConn(conn);
@@ -1870,6 +2028,12 @@ function _check_org_info(state, org_info, org_name) {
       state.orgId = org.id;
       state.orgName = org.name;
       state.logUrl = isomorph_default.getEnv("BRAINTRUST_API_URL") ?? org.api_url;
+      state.proxyUrl = isomorph_default.getEnv("BRAINTRUST_PROXY_URL") ?? org.proxy_url;
+      if (state.proxyUrl) {
+        const url = new URL(state.proxyUrl);
+        url.pathname = "";
+        state.proxyUrl = url.toString();
+      }
       state.gitMetadataSettings = org.git_metadata || void 0;
       break;
     }
@@ -1881,7 +2045,9 @@ function _check_org_info(state, org_info, org_name) {
   }
 }
 function _urljoin(...parts) {
-  return parts.map((x) => x.replace(/^\//, "")).join("/");
+  return parts.map(
+    (x, i) => x.replace(/^\//, "").replace(i < parts.length - 1 ? /\/$/ : "", "")
+  ).join("/");
 }
 function validateTags(tags) {
   const seen = /* @__PURE__ */ new Set();
@@ -2380,12 +2546,32 @@ var SpanImpl = class _SpanImpl {
     const sanitized = validateAndSanitizeExperimentLogPartialArgs(event ?? {});
     let sanitizedAndInternalData = { ...internalData };
     mergeDicts(sanitizedAndInternalData, sanitized);
+    const serializableInternalData = {};
+    const lazyInternalData = {};
+    for (const [key, value] of Object.entries(sanitizedAndInternalData)) {
+      if (value instanceof BraintrustStream) {
+        const streamCopy = value.copy();
+        lazyInternalData[key] = new LazyValue(async () => {
+          return await new Promise((resolve) => {
+            streamCopy.toReadableStream().pipeThrough(createFinalValuePassThroughStream(resolve)).pipeTo(devNullWritableStream());
+          });
+        });
+      } else if (value instanceof ReadableStream) {
+        lazyInternalData[key] = new LazyValue(async () => {
+          return await new Promise((resolve) => {
+            value.pipeThrough(createFinalValuePassThroughStream(resolve)).pipeTo(devNullWritableStream());
+          });
+        });
+      } else {
+        serializableInternalData[key] = value;
+      }
+    }
     let partialRecord = {
       id: this.id,
       span_id: this.spanId,
       root_span_id: this.rootSpanId,
       span_parents: this.spanParents,
-      ...sanitizedAndInternalData,
+      ...serializableInternalData,
       [IS_MERGE_FIELD]: this.isMerge
     };
     const serializedPartialRecord = JSON.stringify(partialRecord, (k, v) => {
@@ -2409,6 +2595,14 @@ var SpanImpl = class _SpanImpl {
     }
     const computeRecord = async () => ({
       ...partialRecord,
+      ...Object.fromEntries(
+        await Promise.all(
+          Object.entries(lazyInternalData).map(async ([key, value]) => [
+            key,
+            await value.get()
+          ])
+        )
+      ),
       ...new SpanComponentsV2({
         objectType: this.parentObjectType,
         objectId: await this.parentObjectId.get()
@@ -2773,6 +2967,55 @@ function configureNode() {
   _internalSetInitialState();
 }
 
+// src/functions/invoke.ts
+import {
+  INVOKE_API_VERSION
+} from "@braintrust/core/typespecs";
+async function invoke(args) {
+  const {
+    orgName,
+    apiKey,
+    appUrl,
+    forceLogin,
+    fetch,
+    arg,
+    parent: parentArg,
+    state: stateArg,
+    stream,
+    schema,
+    ...functionId
+  } = args;
+  const state = stateArg ?? _internalGetGlobalState();
+  await state.login({
+    orgName,
+    apiKey,
+    appUrl,
+    forceLogin
+  });
+  const parent = parentArg ? typeof parentArg === "string" ? parentArg : await parentArg.export() : await currentSpan().export();
+  const request = {
+    ...functionId,
+    arg,
+    parent,
+    stream,
+    api_version: INVOKE_API_VERSION
+  };
+  const resp = await state.proxyConn().post(`function/invoke`, request, {
+    headers: {
+      Accept: stream ? "text/event-stream" : "application/json"
+    }
+  });
+  if (stream) {
+    if (!resp.body) {
+      throw new Error("Received empty stream body");
+    }
+    return new BraintrustStream(resp.body);
+  } else {
+    const data = await resp.json();
+    return schema ? schema.parse(data) : data;
+  }
+}
+
 // src/framework.ts
 import chalk from "chalk";
 import { SpanTypeAttribute as SpanTypeAttribute2, mergeDicts as mergeDicts2 } from "@braintrust/core";
@@ -2817,7 +3060,7 @@ var BarProgressReporter = class {
 // src/framework.ts
 import pluralize from "pluralize";
 
-// ../../node_modules/.pnpm/async@3.2.5/node_modules/async/dist/async.mjs
+// ../node_modules/.pnpm/async@3.2.5/node_modules/async/dist/async.mjs
 function initialParams(fn) {
   return function(...args) {
     var callback = args.pop();
@@ -3792,93 +4035,6 @@ function queue(worker, concurrency) {
     _worker(items[0], cb);
   }, concurrency, 1);
 }
-var Heap = class {
-  constructor() {
-    this.heap = [];
-    this.pushCount = Number.MIN_SAFE_INTEGER;
-  }
-  get length() {
-    return this.heap.length;
-  }
-  empty() {
-    this.heap = [];
-    return this;
-  }
-  percUp(index) {
-    let p;
-    while (index > 0 && smaller(this.heap[index], this.heap[p = parent(index)])) {
-      let t = this.heap[index];
-      this.heap[index] = this.heap[p];
-      this.heap[p] = t;
-      index = p;
-    }
-  }
-  percDown(index) {
-    let l;
-    while ((l = leftChi(index)) < this.heap.length) {
-      if (l + 1 < this.heap.length && smaller(this.heap[l + 1], this.heap[l])) {
-        l = l + 1;
-      }
-      if (smaller(this.heap[index], this.heap[l])) {
-        break;
-      }
-      let t = this.heap[index];
-      this.heap[index] = this.heap[l];
-      this.heap[l] = t;
-      index = l;
-    }
-  }
-  push(node) {
-    node.pushCount = ++this.pushCount;
-    this.heap.push(node);
-    this.percUp(this.heap.length - 1);
-  }
-  unshift(node) {
-    return this.heap.push(node);
-  }
-  shift() {
-    let [top] = this.heap;
-    this.heap[0] = this.heap[this.heap.length - 1];
-    this.heap.pop();
-    this.percDown(0);
-    return top;
-  }
-  toArray() {
-    return [...this];
-  }
-  *[Symbol.iterator]() {
-    for (let i = 0; i < this.heap.length; i++) {
-      yield this.heap[i].data;
-    }
-  }
-  remove(testFn) {
-    let j = 0;
-    for (let i = 0; i < this.heap.length; i++) {
-      if (!testFn(this.heap[i])) {
-        this.heap[j] = this.heap[i];
-        j++;
-      }
-    }
-    this.heap.splice(j);
-    for (let i = parent(this.heap.length - 1); i >= 0; i--) {
-      this.percDown(i);
-    }
-    return this;
-  }
-};
-function leftChi(i) {
-  return (i << 1) + 1;
-}
-function parent(i) {
-  return (i + 1 >> 1) - 1;
-}
-function smaller(x, y) {
-  if (x.priority !== y.priority) {
-    return x.priority < y.priority;
-  } else {
-    return x.pushCount < y.pushCount;
-  }
-}
 function race(tasks, callback) {
   callback = once(callback);
   if (!Array.isArray(tasks))
@@ -4999,44 +5155,19 @@ function postProcessOutput(text, tool_calls, finish_reason) {
 }
 
 // src/index.ts
-import {
-  ParentExperimentIds,
-  ParentProjectLogIds,
-  IdField as IdField2,
-  InputField,
-  InputsField,
-  OtherExperimentLogFields,
-  ExperimentLogPartialArgs as ExperimentLogPartialArgs2,
-  ExperimentLogFullArgs as ExperimentLogFullArgs2,
-  LogFeedbackFullArgs as LogFeedbackFullArgs2,
-  LogCommentFullArgs,
-  CommentEvent,
-  DatasetRecord as DatasetRecord2
-} from "@braintrust/core";
 configureNode();
 export {
   BaseExperiment,
   BraintrustState,
-  CommentEvent,
+  BraintrustStream,
   Dataset,
-  DatasetRecord2 as DatasetRecord,
   Eval,
   Experiment,
-  ExperimentLogFullArgs2 as ExperimentLogFullArgs,
-  ExperimentLogPartialArgs2 as ExperimentLogPartialArgs,
-  IdField2 as IdField,
-  InputField,
-  InputsField,
   LEGACY_CACHED_HEADER,
   LazyValue,
-  LogCommentFullArgs,
-  LogFeedbackFullArgs2 as LogFeedbackFullArgs,
   Logger,
   NOOP_SPAN,
   NoopSpan,
-  OtherExperimentLogFields,
-  ParentExperimentIds,
-  ParentProjectLogIds,
   Prompt,
   ReadonlyExperiment,
   Reporter,
@@ -5045,15 +5176,18 @@ export {
   _internalGetGlobalState,
   _internalSetInitialState,
   buildLocalSummary,
+  createFinalValuePassThroughStream,
   currentExperiment,
   currentLogger,
   currentSpan,
+  devNullWritableStream,
   flush,
   getSpanParentObject,
   init,
   initDataset,
   initExperiment,
   initLogger,
+  invoke,
   loadPrompt,
   log,
   login,