braintrust 0.0.140 → 0.0.141-dev

package/dist/browser.mjs

@@ -97,6 +97,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";
@@ -133,6 +252,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()
@@ -171,11 +291,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;
@@ -183,10 +305,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;
@@ -195,10 +319,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) {
@@ -206,7 +332,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"
       );
@@ -218,6 +344,7 @@ var BraintrustState = class _BraintrustState {
       orgId: this.orgId,
       orgName: this.orgName,
       logUrl: this.logUrl,
+      proxyUrl: this.proxyUrl,
       gitMetadataSettings: this.gitMetadataSettings
     };
   }
@@ -240,6 +367,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;
@@ -278,6 +406,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;
   }
@@ -590,6 +727,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.
    *
@@ -615,6 +757,7 @@ var Logger = class {
     this.lastStartTime = span.end();
     const ret = span.id;
     if (this.asyncFlush === true) {
+      this.triggerWaitUntilFlush();
       return ret;
     } else {
       return (async () => {
@@ -642,6 +785,7 @@ var Logger = class {
       () => span.end()
     );
     if (this.asyncFlush) {
+      this.triggerWaitUntilFlush();
       return ret;
     } else {
       return (async () => {
@@ -1073,6 +1217,8 @@ function init(projectOrOptions, optionalOptions) {
     appUrl,
     apiKey,
     orgName,
+    forceLogin,
+    fetch,
     metadata,
     gitMetadataSettings,
     projectId,
@@ -1090,7 +1236,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,
@@ -1261,6 +1407,8 @@ function initDataset(projectOrOptions, optionalOptions) {
     appUrl,
     apiKey,
     orgName,
+    fetch,
+    forceLogin,
     projectId,
     useOutput: legacy,
     state: stateArg
@@ -1271,7 +1419,9 @@ function initDataset(projectOrOptions, optionalOptions) {
       await state.login({
         orgName,
         apiKey,
-        appUrl
+        appUrl,
+        fetch,
+        forceLogin
       });
       const args = {
         org_id: state.orgId,
@@ -1350,6 +1500,7 @@ function initLogger(options = {}) {
     apiKey,
     orgName,
     forceLogin,
+    fetch,
     state: stateArg
   } = options || {};
   const computeMetadataArgs = {
@@ -1363,7 +1514,8 @@ function initLogger(options = {}) {
         orgName,
         apiKey,
         appUrl,
-        forceLogin
+        forceLogin,
+        fetch
       });
       return computeLoggerMetadata(state, computeMetadataArgs);
     }
@@ -1387,6 +1539,8 @@ async function loadPrompt({
   appUrl,
   apiKey,
   orgName,
+  fetch,
+  forceLogin,
   state: stateArg
 }) {
   if (isEmpty(projectName) && isEmpty(projectId)) {
@@ -1399,7 +1553,9 @@ async function loadPrompt({
   state.login({
     orgName,
     apiKey,
-    appUrl
+    appUrl,
+    fetch,
+    forceLogin
   });
   const args = {
     project_name: projectName,
@@ -1442,6 +1598,7 @@ async function login(options = {}) {
   }
   await _globalState.login(options);
   globalThis.__inherited_braintrust_state = _globalState;
+  return _globalState;
 }
 async function loginToState(options = {}) {
   const {
@@ -1480,6 +1637,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);
@@ -1651,6 +1809,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;
     }
@@ -1662,7 +1826,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();
@@ -2161,12 +2327,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) => {
@@ -2190,6 +2376,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()
@@ -2559,6 +2753,55 @@ function configureBrowser() {
   browserConfigured = true;
 }
 
+// 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/wrappers/oai.ts
 import { SpanTypeAttribute as SpanTypeAttribute2 } from "@braintrust/core";
 import { mergeDicts as mergeDicts2 } from "@braintrust/core";
@@ -2915,6 +3158,7 @@ var WrapperStream = class {
 configureBrowser();
 export {
   BraintrustState,
+  BraintrustStream,
   Dataset,
   Experiment,
   LEGACY_CACHED_HEADER,
@@ -2927,15 +3171,18 @@ export {
   X_CACHED_HEADER,
   _internalGetGlobalState,
   _internalSetInitialState,
+  createFinalValuePassThroughStream,
   currentExperiment,
   currentLogger,
   currentSpan,
+  devNullWritableStream,
   flush,
   getSpanParentObject,
   init,
   initDataset,
   initExperiment,
   initLogger,
+  invoke,
   loadPrompt,
   log,
   login,