braintrust 3.24.0 → 3.26.0

package/dist/index.mjs

@@ -91,6 +91,7 @@ import {
   initExperiment,
   initLogger,
   injectTraceContext,
+  instrumentCloudflareAIChatAgent,
   isEmpty,
   isObject,
   isTemplateFormat,
@@ -139,13 +140,15 @@ import {
   wrapMastraAgent,
   wrapTraced,
   zodToJsonSchema
-} from "./chunk-XE5FS7QY.mjs";
+} from "./chunk-XIZOM2HO.mjs";
 import {
+  INSTRUMENTATION_NAMES,
   __export,
   aiSDKChannels,
   anthropicChannels,
   bedrockRuntimeChannels,
   claudeAgentSDKChannels,
+  cloudflareAgentsChannels,
   cohereChannels,
   cursorSDKChannels,
   genkitChannels,
@@ -155,6 +158,8 @@ import {
   groqChannels,
   harnessAgentChannels,
   huggingFaceChannels,
+  huggingFaceTransformersChannels,
+  isSupportedHuggingFaceTransformersTask,
   isomorph_default,
   langSmithChannels,
   mistralChannels,
@@ -163,8 +168,9 @@ import {
   openRouterAgentChannels,
   openRouterChannels,
   piCodingAgentChannels,
-  strandsAgentSDKChannels
-} from "./chunk-7F6GCRHH.mjs";
+  strandsAgentSDKChannels,
+  withSpanInstrumentationName
+} from "./chunk-PN7EWK66.mjs";
 
 // src/exports.ts
 var exports_exports = {};
@@ -315,6 +321,8 @@ __export(exports_exports, {
   wrapAnthropic: () => wrapAnthropic,
   wrapBedrockRuntime: () => wrapBedrockRuntime,
   wrapClaudeAgentSDK: () => wrapClaudeAgentSDK,
+  wrapCloudflareAIChat: () => wrapCloudflareAIChat,
+  wrapCloudflareAgent: () => wrapCloudflareAgent,
   wrapCohere: () => wrapCohere,
   wrapCopilotClient: () => wrapCopilotClient,
   wrapCursorSDK: () => wrapCursorSDK,
@@ -323,6 +331,7 @@ __export(exports_exports, {
   wrapGoogleGenAI: () => wrapGoogleGenAI,
   wrapGroq: () => wrapGroq,
   wrapHuggingFace: () => wrapHuggingFace,
+  wrapHuggingFaceTransformers: () => wrapHuggingFaceTransformers,
   wrapLangSmithClient: () => wrapLangSmithClient,
   wrapLangSmithRunTrees: () => wrapLangSmithRunTrees,
   wrapLangSmithTraceable: () => wrapLangSmithTraceable,
@@ -941,6 +950,12 @@ var wrapAgentClass = (AgentClass, options = {}) => {
           const original = Reflect.get(instanceTarget, prop, instanceTarget);
           if (harnessAgent && typeof original === "function") {
             switch (prop) {
+              case "createSession":
+                return wrapHarnessAgentCreateSession(
+                  // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
+                  original,
+                  harnessAgent
+                );
               case "generate":
                 return wrapHarnessAgentGenerate(
                   original,
@@ -990,6 +1005,19 @@ var wrapAgentClass = (AgentClass, options = {}) => {
     }
   });
 };
+var wrapHarnessAgentCreateSession = (createSession, instance) => {
+  const wrapper = function(params) {
+    return harnessAgentChannels.createSession.tracePromise(
+      () => params === void 0 ? createSession.call(instance) : createSession.call(instance, params),
+      createAISDKChannelContext(params ?? {}, { self: instance })
+    );
+  };
+  Object.defineProperty(wrapper, "name", {
+    value: "HarnessAgent.createSession",
+    writable: false
+  });
+  return wrapper;
+};
 var wrapHarnessAgentGenerate = (generate, instance, channel, name, options) => makeGenerateTextWrapper(
   channel,
   name,
@@ -1259,12 +1287,17 @@ var BraintrustLanguageModelWrapper = class {
   // For the first cut, do not support custom span_info arguments. We can
   // propagate those via async local storage
   async doGenerate(options) {
-    const span = startSpan({
-      name: "Chat Completion",
-      spanAttributes: {
-        type: "llm"
-      }
-    });
+    const span = startSpan(
+      withSpanInstrumentationName(
+        {
+          name: "Chat Completion",
+          spanAttributes: {
+            type: "llm"
+          }
+        },
+        INSTRUMENTATION_NAMES.AI_SDK
+      )
+    );
     const { prompt, mode, ...rest } = options;
     const startTime = getCurrentUnixTimestamp();
     try {
@@ -1295,12 +1328,17 @@ var BraintrustLanguageModelWrapper = class {
   async doStream(options) {
     const { prompt, mode, ...rest } = options;
     const startTime = getCurrentUnixTimestamp();
-    const span = startSpan({
-      name: "Chat Completion",
-      spanAttributes: {
-        type: "llm"
-      }
-    });
+    const span = startSpan(
+      withSpanInstrumentationName(
+        {
+          name: "Chat Completion",
+          spanAttributes: {
+            type: "llm"
+          }
+        },
+        INSTRUMENTATION_NAMES.AI_SDK
+      )
+    );
     span.log({
       input: postProcessPrompt(prompt),
       metadata: {
@@ -1655,7 +1693,9 @@ function BraintrustMiddleware(config = {}) {
           }
         }
       };
-      const span = startSpan(spanArgs);
+      const span = startSpan(
+        withSpanInstrumentationName(spanArgs, INSTRUMENTATION_NAMES.AI_SDK)
+      );
       try {
         const result = await doGenerate();
         const metadata = {};
@@ -1717,7 +1757,9 @@ function BraintrustMiddleware(config = {}) {
           }
         }
       };
-      const span = startSpan(spanArgs);
+      const span = startSpan(
+        withSpanInstrumentationName(spanArgs, INSTRUMENTATION_NAMES.AI_SDK)
+      );
       try {
         const { stream, ...rest } = await doStream();
         const textChunks = [];
@@ -1924,7 +1966,12 @@ var EveBridge = class {
     };
     const span = withCurrent(
       NOOP_SPAN,
-      () => _internalStartSpanWithInitialMerge({ ...args, startTime })
+      () => _internalStartSpanWithInitialMerge(
+        withSpanInstrumentationName(
+          { ...args, startTime },
+          INSTRUMENTATION_NAMES.EVE
+        )
+      )
     );
     if (typeof rowId !== "string") {
       return span;
@@ -3996,6 +4043,67 @@ function patchAgentSessionClass(AgentSession) {
   });
 }
 
+// src/wrappers/cloudflare-agent.ts
+var WRAPPED_RUN_AGENT_TOOL = /* @__PURE__ */ Symbol.for(
+  "braintrust.cloudflare-agents.wrapped-run-agent-tool"
+);
+function wrapCloudflareAgent(Agent) {
+  if (typeof Agent !== "function") {
+    warnUnsupportedAgent();
+    return Agent;
+  }
+  const prototypeDescriptor = Object.getOwnPropertyDescriptor(
+    Agent,
+    "prototype"
+  );
+  const prototype = prototypeDescriptor && "value" in prototypeDescriptor ? prototypeDescriptor.value : void 0;
+  if (!isObjectLike(prototype)) {
+    warnUnsupportedAgent();
+    return Agent;
+  }
+  if (ownValue(prototype, WRAPPED_RUN_AGENT_TOOL) === true) {
+    return Agent;
+  }
+  const descriptor = Object.getOwnPropertyDescriptor(prototype, "runAgentTool");
+  if (!descriptor || typeof descriptor.value !== "function") {
+    warnUnsupportedAgent();
+    return Agent;
+  }
+  const originalRunAgentTool = descriptor.value;
+  Object.defineProperty(prototype, "runAgentTool", {
+    ...descriptor,
+    value: function wrappedRunAgentTool(...args) {
+      const options = args[1];
+      if (ownValue(options, "detached")) {
+        return Reflect.apply(originalRunAgentTool, this, args);
+      }
+      return cloudflareAgentsChannels.runAgentTool.tracePromise(
+        () => Reflect.apply(originalRunAgentTool, this, args),
+        { arguments: args, self: this }
+      );
+    }
+  });
+  Object.defineProperty(prototype, WRAPPED_RUN_AGENT_TOOL, {
+    configurable: false,
+    enumerable: false,
+    value: true
+  });
+  return Agent;
+}
+function ownValue(value, key) {
+  if (!isObjectLike(value)) {
+    return void 0;
+  }
+  const descriptor = Object.getOwnPropertyDescriptor(value, key);
+  return descriptor && "value" in descriptor ? descriptor.value : void 0;
+}
+function isObjectLike(value) {
+  return typeof value === "object" && value !== null || typeof value === "function";
+}
+function warnUnsupportedAgent() {
+  debugLogger.warn("Unsupported Cloudflare Agents Agent class. Not wrapping.");
+}
+
 // src/wrappers/strands-agent-sdk.ts
 var WRAPPED_CLASS = /* @__PURE__ */ Symbol.for("braintrust.strands-agent-sdk.wrapped-class");
 var WRAPPED_INSTANCE = /* @__PURE__ */ Symbol.for(
@@ -4164,6 +4272,81 @@ async function consumeAsyncGenerator(generator) {
   return result.value;
 }
 
+// src/wrappers/cloudflare-ai-chat.ts
+var wrappedClasses = /* @__PURE__ */ new WeakMap();
+function wrapCloudflareAIChat(module) {
+  if (!module || typeof module !== "object") {
+    return module;
+  }
+  const candidate = module;
+  let AIChatAgent;
+  try {
+    AIChatAgent = Reflect.get(candidate, "AIChatAgent");
+  } catch (error) {
+    debugLogger.debug("Failed to inspect @cloudflare/ai-chat module:", error);
+    return module;
+  }
+  if (typeof AIChatAgent !== "function") {
+    debugLogger.warn(
+      "Unsupported @cloudflare/ai-chat module. AIChatAgent was not found; not wrapping."
+    );
+    return module;
+  }
+  const target = isModuleNamespace5(module) ? Object.setPrototypeOf({}, module) : candidate;
+  return new Proxy(target, {
+    get(target2, property, receiver) {
+      const value = Reflect.get(target2, property, receiver);
+      if (property === "AIChatAgent" && typeof value === "function") {
+        return wrapAIChatAgentClass(
+          value
+        );
+      }
+      return value;
+    }
+  });
+}
+function wrapAIChatAgentClass(AgentClass) {
+  const cached = wrappedClasses.get(AgentClass);
+  if (cached) {
+    return cached;
+  }
+  const wrapped = new Proxy(AgentClass, {
+    get(target, property) {
+      const value = Reflect.get(target, property, target);
+      return typeof value === "function" ? value.bind(target) : value;
+    },
+    construct(target, args, newTarget) {
+      const instance = Reflect.construct(
+        target,
+        args,
+        newTarget
+      );
+      return instrumentCloudflareAIChatAgent(instance);
+    }
+  });
+  wrappedClasses.set(AgentClass, wrapped);
+  wrappedClasses.set(wrapped, wrapped);
+  return wrapped;
+}
+function isModuleNamespace5(value) {
+  if (!value || typeof value !== "object") {
+    return false;
+  }
+  try {
+    if (Reflect.get(value, Symbol.toStringTag) === "Module") {
+      return true;
+    }
+    const keys = Object.keys(value);
+    if (keys.length === 0) {
+      return false;
+    }
+    const descriptor = Object.getOwnPropertyDescriptor(value, keys[0]);
+    return descriptor ? !descriptor.configurable && !descriptor.writable : false;
+  } catch {
+    return false;
+  }
+}
+
 // src/wrappers/google-genai.ts
 function wrapGoogleGenAI(googleGenAI) {
   if (!googleGenAI || typeof googleGenAI !== "object") {
@@ -4816,6 +4999,120 @@ function wrapFeatureExtraction(original, endpointUrl) {
   };
 }
 
+// src/wrappers/huggingface-transformers.ts
+var PIPELINE_CONSTRUCTOR_KEYS = /* @__PURE__ */ new Set([
+  "TextGenerationPipeline",
+  "Text2TextGenerationPipeline",
+  "SummarizationPipeline",
+  "FeatureExtractionPipeline",
+  "QuestionAnsweringPipeline"
+]);
+var wrappedValues = /* @__PURE__ */ new WeakMap();
+function wrapHuggingFaceTransformers(transformers) {
+  const pipelineTask = transformers?.task;
+  if (typeof transformers === "function" && isSupportedHuggingFaceTransformersTask(pipelineTask)) {
+    return wrapPipeline(transformers);
+  }
+  if (!isSupportedModule(transformers)) {
+    return transformers;
+  }
+  const existing = wrappedValues.get(transformers);
+  if (existing) {
+    return existing;
+  }
+  const proxy = new Proxy(Object.create(transformers), {
+    get(_target, property, receiver) {
+      const value = Reflect.get(transformers, property, receiver);
+      if (property === "pipeline" && typeof value === "function") {
+        return wrapPipelineFactory(value);
+      }
+      if (typeof property === "string" && PIPELINE_CONSTRUCTOR_KEYS.has(property) && typeof value === "function") {
+        return wrapPipelineConstructor(
+          value
+        );
+      }
+      return value;
+    }
+  });
+  wrappedValues.set(transformers, proxy);
+  wrappedValues.set(proxy, proxy);
+  return proxy;
+}
+function isSupportedModule(value) {
+  if (value === null || typeof value !== "object" && typeof value !== "function") {
+    return false;
+  }
+  if (typeof Reflect.get(value, "pipeline") === "function") {
+    return true;
+  }
+  for (const key of PIPELINE_CONSTRUCTOR_KEYS) {
+    if (typeof Reflect.get(value, key) === "function") {
+      return true;
+    }
+  }
+  return false;
+}
+function wrapPipelineFactory(factory) {
+  const existing = wrappedValues.get(factory);
+  if (existing) {
+    return existing;
+  }
+  const wrapped = function(...args) {
+    const [task] = args;
+    const context = {
+      arguments: args
+    };
+    return huggingFaceTransformersChannels.pipeline.tracePromise(() => Reflect.apply(factory, this, args), context).then((pipeline) => {
+      if (isSupportedHuggingFaceTransformersTask(pipeline.task ?? task)) {
+        return wrapPipeline(pipeline);
+      }
+      return pipeline;
+    });
+  };
+  wrappedValues.set(factory, wrapped);
+  wrappedValues.set(wrapped, wrapped);
+  return wrapped;
+}
+function wrapPipelineConstructor(constructor) {
+  const existing = wrappedValues.get(constructor);
+  if (existing) {
+    return existing;
+  }
+  const proxy = new Proxy(constructor, {
+    construct(target, args, newTarget) {
+      const pipeline = Reflect.construct(target, args, newTarget);
+      if (isSupportedHuggingFaceTransformersTask(pipeline.task)) {
+        return wrapPipeline(pipeline);
+      }
+      return pipeline;
+    }
+  });
+  wrappedValues.set(constructor, proxy);
+  wrappedValues.set(proxy, proxy);
+  return proxy;
+}
+function wrapPipeline(pipeline) {
+  const existing = wrappedValues.get(pipeline);
+  if (existing) {
+    return existing;
+  }
+  const proxy = new Proxy(pipeline, {
+    apply(target, thisArg, args) {
+      const context = {
+        arguments: args,
+        self: target
+      };
+      return huggingFaceTransformersChannels.pipelineCall.tracePromise(
+        () => Reflect.apply(target, thisArg, args),
+        context
+      );
+    }
+  });
+  wrappedValues.set(pipeline, proxy);
+  wrappedValues.set(proxy, proxy);
+  return proxy;
+}
+
 // src/wrappers/openrouter-agent.ts
 function wrapOpenRouterAgent(agent) {
   const candidate = agent;
@@ -5532,7 +5829,7 @@ function wrapNamespaceExport(namespace, exportName, marker, wrap2) {
     );
     return namespace;
   }
-  const target = isModuleNamespace5(namespace) ? Object.setPrototypeOf({}, namespace) : candidate;
+  const target = isModuleNamespace6(namespace) ? Object.setPrototypeOf({}, namespace) : candidate;
   const moduleNamespace = target !== candidate;
   let wrappedExport;
   return new Proxy(target, {
@@ -5566,7 +5863,7 @@ function wrapNamespaceExport(namespace, exportName, marker, wrap2) {
     }
   });
 }
-function isModuleNamespace5(value) {
+function isModuleNamespace6(value) {
   if (!value || typeof value !== "object") {
     return false;
   }
@@ -8836,26 +9133,30 @@ async function runEvaluatorInternal(experiment, evaluator, progressReporter, fil
         }
       }
     })();
+    let timeoutId;
+    let abortHandler;
+    const cleanupCancellation = () => {
+      if (timeoutId !== void 0) {
+        clearTimeout(timeoutId);
+        timeoutId = void 0;
+      }
+      if (abortHandler && evaluator.signal) {
+        evaluator.signal.removeEventListener("abort", abortHandler);
+        abortHandler = void 0;
+      }
+    };
     const cancel = async () => {
       await new Promise((_, reject2) => {
         if (cancelled) {
           reject2(new InternalAbortError("Evaluator already cancelled"));
           return;
         }
-        let timeoutId;
-        let abortHandler;
         const rejectOnce = (error) => {
           if (cancelled) {
             return;
           }
           cancelled = true;
-          if (timeoutId) {
-            clearTimeout(timeoutId);
-            timeoutId = void 0;
-          }
-          if (abortHandler && evaluator.signal) {
-            evaluator.signal.removeEventListener("abort", abortHandler);
-          }
+          cleanupCancellation();
           reject2(error);
         };
         if (evaluator.timeout) {
@@ -8895,6 +9196,7 @@ async function runEvaluatorInternal(experiment, evaluator, progressReporter, fil
       }
       throw e;
     } finally {
+      cleanupCancellation();
       if (!collectResults) {
         collectedResults.length = 0;
       }
@@ -10158,6 +10460,8 @@ export {
   wrapAnthropic,
   wrapBedrockRuntime,
   wrapClaudeAgentSDK,
+  wrapCloudflareAIChat,
+  wrapCloudflareAgent,
   wrapCohere,
   wrapCopilotClient,
   wrapCursorSDK,
@@ -10166,6 +10470,7 @@ export {
   wrapGoogleGenAI,
   wrapGroq,
   wrapHuggingFace,
+  wrapHuggingFaceTransformers,
   wrapLangSmithClient,
   wrapLangSmithRunTrees,
   wrapLangSmithTraceable,