braintrust 0.2.1 → 0.2.2

package/dist/index.mjs

@@ -6413,6 +6413,110 @@ function traced(callback, args) {
     })();
   }
 }
+function isGeneratorFunction(fn) {
+  return Object.prototype.toString.call(fn) === "[object GeneratorFunction]";
+}
+function isAsyncGeneratorFunction(fn) {
+  return Object.prototype.toString.call(fn) === "[object AsyncGeneratorFunction]";
+}
+function wrapTracedSyncGenerator(fn, spanArgs, noTraceIO) {
+  const wrapper = function* (...fnArgs) {
+    const span = startSpan(spanArgs);
+    try {
+      if (!noTraceIO) {
+        span.log({ input: fnArgs });
+      }
+      const envValue = isomorph_default.getEnv("BRAINTRUST_MAX_GENERATOR_ITEMS");
+      const maxItems = envValue !== void 0 ? Number(envValue) : 1e3;
+      if (!noTraceIO && maxItems !== 0) {
+        let collected = [];
+        let truncated = false;
+        const gen = generatorWithCurrent(span, fn.apply(this, fnArgs));
+        try {
+          for (const value of gen) {
+            if (maxItems === -1 || !truncated && collected.length < maxItems) {
+              collected.push(value);
+            } else {
+              truncated = true;
+              collected = [];
+              console.warn(
+                `Generator output exceeded limit of ${maxItems} items, output not logged. Increase BRAINTRUST_MAX_GENERATOR_ITEMS or set to -1 to disable limit.`
+              );
+            }
+            yield value;
+          }
+          if (!truncated) {
+            span.log({ output: collected });
+          }
+        } catch (error2) {
+          logError(span, error2);
+          if (!truncated && collected.length > 0) {
+            span.log({ output: collected });
+          }
+          throw error2;
+        }
+      } else {
+        const gen = generatorWithCurrent(span, fn.apply(this, fnArgs));
+        for (const value of gen) {
+          yield value;
+        }
+      }
+    } finally {
+      span.end();
+    }
+  };
+  Object.defineProperty(wrapper, "name", { value: fn.name });
+  return wrapper;
+}
+function wrapTracedAsyncGenerator(fn, spanArgs, noTraceIO) {
+  const wrapper = async function* (...fnArgs) {
+    const span = startSpan(spanArgs);
+    try {
+      if (!noTraceIO) {
+        span.log({ input: fnArgs });
+      }
+      const envValue = isomorph_default.getEnv("BRAINTRUST_MAX_GENERATOR_ITEMS");
+      const maxItems = envValue !== void 0 ? Number(envValue) : 1e3;
+      if (!noTraceIO && maxItems !== 0) {
+        let collected = [];
+        let truncated = false;
+        const gen = asyncGeneratorWithCurrent(span, fn.apply(this, fnArgs));
+        try {
+          for await (const value of gen) {
+            if (maxItems === -1 || !truncated && collected.length < maxItems) {
+              collected.push(value);
+            } else {
+              truncated = true;
+              collected = [];
+              console.warn(
+                `Generator output exceeded limit of ${maxItems} items, output not logged. Increase BRAINTRUST_MAX_GENERATOR_ITEMS or set to -1 to disable limit.`
+              );
+            }
+            yield value;
+          }
+          if (!truncated) {
+            span.log({ output: collected });
+          }
+        } catch (error2) {
+          logError(span, error2);
+          if (!truncated && collected.length > 0) {
+            span.log({ output: collected });
+          }
+          throw error2;
+        }
+      } else {
+        const gen = asyncGeneratorWithCurrent(span, fn.apply(this, fnArgs));
+        for await (const value of gen) {
+          yield value;
+        }
+      }
+    } finally {
+      span.end();
+    }
+  };
+  Object.defineProperty(wrapper, "name", { value: fn.name });
+  return wrapper;
+}
 function wrapTraced(fn, args) {
   const spanArgs = {
     name: fn.name,
@@ -6421,6 +6525,13 @@ function wrapTraced(fn, args) {
   };
   const hasExplicitInput = args && args.event && "input" in args.event && args.event.input !== void 0;
   const hasExplicitOutput = args && args.event && args.event.output !== void 0;
+  const noTraceIO = args?.noTraceIO || hasExplicitInput || hasExplicitOutput;
+  if (isGeneratorFunction(fn)) {
+    return wrapTracedSyncGenerator(fn, spanArgs, !!noTraceIO);
+  }
+  if (isAsyncGeneratorFunction(fn)) {
+    return wrapTracedAsyncGenerator(fn, spanArgs, !!noTraceIO);
+  }
   if (args?.asyncFlush) {
     return (...fnArgs) => traced((span) => {
       if (!hasExplicitInput) {
@@ -6506,6 +6617,52 @@ function startSpanAndIsLogger(args) {
 function withCurrent(span, callback, state = void 0) {
   return (state ?? _globalState).currentSpan.run(span, () => callback(span));
 }
+function* generatorWithCurrent(span, gen, state = void 0) {
+  let nextValue;
+  while (true) {
+    const result = withCurrent(
+      span,
+      () => {
+        try {
+          return gen.next(nextValue);
+        } catch (e) {
+          return { value: void 0, done: true, error: e };
+        }
+      },
+      state
+    );
+    if ("error" in result) {
+      throw result.error;
+    }
+    if (result.done) {
+      return result.value;
+    }
+    nextValue = yield result.value;
+  }
+}
+async function* asyncGeneratorWithCurrent(span, gen, state = void 0) {
+  let nextValue;
+  while (true) {
+    const result = await withCurrent(
+      span,
+      async () => {
+        try {
+          return await gen.next(nextValue);
+        } catch (e) {
+          return { value: void 0, done: true, error: e };
+        }
+      },
+      state
+    );
+    if ("error" in result) {
+      throw result.error;
+    }
+    if (result.done) {
+      return result.value;
+    }
+    nextValue = yield result.value;
+  }
+}
 function withParent(parent, callback, state = void 0) {
   return (state ?? _globalState).currentParent.run(parent, () => callback());
 }
@@ -8002,7 +8159,9 @@ var _exportsForTestingOnly = {
   clearTestBackgroundLogger,
   simulateLoginForTests,
   simulateLogoutForTests,
-  setInitialTestState
+  setInitialTestState,
+  isGeneratorFunction,
+  isAsyncGeneratorFunction
 };
 
 // src/node.ts
@@ -10611,6 +10770,8 @@ function responsesProxy(openai) {
         return responsesCreateProxy(target.create.bind(target));
       } else if (name === "stream") {
         return responsesStreamProxy(target.stream.bind(target));
+      } else if (name === "parse") {
+        return responsesParseProxy(target.parse.bind(target));
       }
       return Reflect.get(target, name, receiver);
     }
@@ -10655,6 +10816,36 @@ function parseEventFromResponseCreateResult(result) {
     metrics: parseMetricsFromUsage(result?.usage)
   };
 }
+function parseSpanFromResponseParseParams(params) {
+  const input = [{ role: "user", content: params.input }];
+  if (params.instructions) {
+    input.push({ role: "system", content: params.instructions });
+  }
+  const spanArgs = {
+    name: "openai.responses.parse",
+    spanAttributes: {
+      type: "llm"
+    },
+    event: {
+      input,
+      metadata: {
+        ...filterFrom(params, ["input", "instructions"]),
+        provider: "openai"
+      }
+    },
+    startTime: getCurrentUnixTimestamp()
+  };
+  return {
+    span: startSpan(spanArgs),
+    start: spanArgs.startTime
+  };
+}
+function parseEventFromResponseParseResult(result) {
+  return {
+    output: result?.output_parsed || result?.output_text || "",
+    metrics: parseMetricsFromUsage(result?.usage)
+  };
+}
 function traceResponseCreateStream(stream, timedSpan) {
   const span = timedSpan.span;
   let ttft = -1;
@@ -10729,6 +10920,16 @@ function responsesStreamProxy(target) {
     }
   });
 }
+function responsesParseProxy(target) {
+  const hooks = {
+    name: "openai.responses.parse",
+    toSpanFunc: parseSpanFromResponseParseParams,
+    resultToEventFunc: parseEventFromResponseParseResult,
+    traceStreamFunc: traceResponseCreateStream
+    // Reuse the same stream tracing
+  };
+  return proxyCreate(target, hooks);
+}
 var TOKEN_NAME_MAP = {
   input_tokens: "prompt_tokens",
   output_tokens: "completion_tokens",
@@ -11698,14 +11899,26 @@ function postProcessPrompt(prompt) {
   });
 }
 function postProcessOutput(text, toolCalls, finishReason) {
-  return {
-    index: 0,
-    message: {
-      role: "assistant",
-      content: text ?? (toolCalls ? toolCalls.length === 1 && toolCalls[0].toolName === "json" ? toolCalls[0].args : JSON.stringify(toolCalls) : "")
-    },
-    finish_reason: finishReason
-  };
+  return [
+    {
+      index: 0,
+      message: {
+        role: "assistant",
+        content: text ?? "",
+        ...toolCalls && toolCalls.length > 0 ? {
+          tool_calls: toolCalls.map((toolCall) => ({
+            id: toolCall.toolCallId,
+            function: {
+              name: toolCall.toolName,
+              arguments: toolCall.args
+            },
+            type: "function"
+          }))
+        } : {}
+      },
+      finish_reason: finishReason
+    }
+  ];
 }
 
 // src/wrappers/anthropic.ts

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "braintrust",
-  "version": "0.2.1",
+  "version": "0.2.2",
   "description": "SDK for integrating Braintrust",
   "repository": {
     "type": "git",
@@ -92,7 +92,7 @@
   },
   "dependencies": {
     "@ai-sdk/provider": "^1.1.3",
-    "@braintrust/core": "0.0.91",
+    "@braintrust/core": "0.0.92",
     "@next/env": "^14.2.3",
     "@vercel/functions": "^1.0.2",
     "argparse": "^2.0.1",