axiom 0.48.1 → 0.49.1

package/dist/bin.js

@@ -2,7 +2,7 @@
 import {
   extractOverrides,
   loadEvalCommand
-} from "./chunk-3KVDJV5T.js";
+} from "./chunk-4S2HLGZW.js";
 import {
   OAuth,
   fetchOrganizations,
@@ -15,11 +15,11 @@ import {
   verifyToken,
   waitForCallback
 } from "./chunk-3VKWOZAQ.js";
-import "./chunk-P4YKSKIZ.js";
+import "./chunk-3MMFYPFN.js";
 import {
   AxiomCLIError
 } from "./chunk-ISSDOC43.js";
-import "./chunk-WOJCPKV5.js";
+import "./chunk-WO6AVN5J.js";
 import "./chunk-MM5FFQJT.js";
 import "./chunk-4TKUTT24.js";
 import "./chunk-KEXKKQVW.js";
@@ -389,7 +389,7 @@ import { Command } from "commander";
 var loadVersionCommand = (program2) => {
   return program2.addCommand(
     new Command("version").description("cli version").action(() => {
-      console.log("0.48.1");
+      console.log("0.49.1");
     })
   );
 };
@@ -399,7 +399,7 @@ var { loadEnvConfig } = pkg;
 loadEnvConfig(process.cwd());
 var { cleanedArgv, overrides } = extractOverrides(process.argv.slice(2));
 var program = new Command2();
-program.name("axiom").description("Axiom's CLI to manage your objects and run evals").version("0.48.1");
+program.name("axiom").description("Axiom's CLI to manage your objects and run evals").version("0.49.1");
 program.hook("preAction", async (_, actionCommand) => {
   const commandName = actionCommand.name();
   const parentCommand = actionCommand.parent;

package/dist/{chunk-P4YKSKIZ.js → chunk-3MMFYPFN.js}

@@ -5,7 +5,7 @@ import {
   getRedactionPolicy,
   handleMaybeRedactedAttribute,
   package_default
-} from "./chunk-WOJCPKV5.js";
+} from "./chunk-WO6AVN5J.js";
 import {
   isValidName
 } from "./chunk-MM5FFQJT.js";
@@ -260,10 +260,16 @@ function recordSpanError(span, err) {
     span.setAttribute(Attr.HTTP.Response.StatusCode, err.status);
   }
 }
+var spanContextStore = /* @__PURE__ */ new WeakMap();
+var spanContextFactories = {
+  v1: () => ({ version: "v1", originalPrompt: [], rawCall: void 0 }),
+  v2: () => ({ version: "v2", originalPrompt: [], originalV2Prompt: void 0 }),
+  v3: () => ({ version: "v3", originalPrompt: [], originalV3Prompt: void 0 })
+};
 async function withSpanHandling(modelId, operation, options) {
   const bag = propagation.getActiveBaggage();
   const isWithinWithSpan = bag?.getEntry(WITHSPAN_BAGGAGE_KEY)?.value === "true";
-  const spanContext = options?.version === "v3" ? { version: "v3", originalPrompt: [], originalV3Prompt: void 0 } : options?.version === "v2" ? { version: "v2", originalPrompt: [], originalV2Prompt: void 0 } : { version: "v1", originalPrompt: [], rawCall: void 0 };
+  const createSpanContext = spanContextFactories[options?.version ?? "v1"];
   const name = createGenAISpanName(Attr.GenAI.Operation.Name_Values.Chat, modelId);
   if (isWithinWithSpan) {
     const activeSpan = trace.getActiveSpan();
@@ -271,6 +277,11 @@ async function withSpanHandling(modelId, operation, options) {
       throw new Error("Expected active span when within withSpan");
     }
     activeSpan.updateName(name);
+    let spanContext = spanContextStore.get(activeSpan);
+    if (!spanContext) {
+      spanContext = createSpanContext();
+      spanContextStore.set(activeSpan, spanContext);
+    }
     const lease = {
       owned: false,
       end: () => {
@@ -294,7 +305,7 @@ async function withSpanHandling(modelId, operation, options) {
           owned: true,
           end: () => span.end()
         };
-        return await operation(span, spanContext, lease);
+        return await operation(span, createSpanContext(), lease);
       },
       {
         manualEnd: options?.streaming ?? false,
@@ -1307,6 +1318,19 @@ var appendPromptMetadataToSpan = (span, messages) => {
     if (axiomMeta.version) span.setAttribute(Attr.GenAI.PromptMetadata.Version, axiomMeta.version);
   }
 };
+function setStreamStepAttributes(span, values) {
+  const inputTokens = ensureNumber(values.inputTokens);
+  if (inputTokens !== void 0) {
+    span.setAttribute(Attr.GenAI.Usage.InputTokens, inputTokens);
+  }
+  const outputTokens = ensureNumber(values.outputTokens);
+  if (outputTokens !== void 0) {
+    span.setAttribute(Attr.GenAI.Usage.OutputTokens, outputTokens);
+  }
+  if (values.finishReason) {
+    span.setAttribute(Attr.GenAI.Response.FinishReasons, JSON.stringify([values.finishReason]));
+  }
+}
 function axiomAIMiddlewareV1() {
   return {
     wrapGenerate: async ({ doGenerate, params, model }) => {
@@ -1357,17 +1381,23 @@ function axiomAIMiddlewareV1() {
                 },
                 async flush(controller) {
                   try {
+                    const statsResult = stats.result;
                     await setPostCallAttributesV1(
                       span,
                       {
                         ...head,
-                        ...stats.result,
+                        ...statsResult,
                         toolCalls: toolAggregator.result.length > 0 ? toolAggregator.result : void 0,
                         text: textAggregator.text
                       },
                       context3,
                       model
                     );
+                    setStreamStepAttributes(childSpan, {
+                      inputTokens: statsResult.usage?.promptTokens,
+                      outputTokens: statsResult.usage?.completionTokens,
+                      finishReason: statsResult.finishReason
+                    });
                     childSpan.end();
                     if (lease.owned) lease.end();
                     controller.terminate();
@@ -1452,14 +1482,20 @@ function axiomAIMiddlewareV2() {
                 },
                 async flush(controller) {
                   try {
+                    const statsResult = stats.result;
                     const streamResult = {
-                      ...stats.result,
+                      ...statsResult,
                       content: [
                         ...textAggregator.text ? [{ type: "text", text: textAggregator.text }] : [],
                         ...toolAggregator.result
                       ]
                     };
                     await setPostCallAttributesV2(span, streamResult, context3, model);
+                    setStreamStepAttributes(childSpan, {
+                      inputTokens: statsResult.usage?.inputTokens,
+                      outputTokens: statsResult.usage?.outputTokens,
+                      finishReason: statsResult.finishReason
+                    });
                     childSpan.end();
                     if (lease.owned) lease.end();
                     controller.terminate();
@@ -1557,11 +1593,15 @@ async function setPostCallAttributesV1(span, result, context3, _model) {
   }
   const inputTokens = ensureNumber(result.usage?.promptTokens);
   if (inputTokens !== void 0) {
-    span.setAttribute(Attr.GenAI.Usage.InputTokens, inputTokens);
+    const curr = context3.accumulatedInputTokens ?? 0;
+    context3.accumulatedInputTokens = curr + inputTokens;
+    span.setAttribute(Attr.GenAI.Usage.InputTokens, context3.accumulatedInputTokens);
   }
   const outputTokens = ensureNumber(result.usage?.completionTokens);
   if (outputTokens !== void 0) {
-    span.setAttribute(Attr.GenAI.Usage.OutputTokens, outputTokens);
+    const curr = context3.accumulatedOutputTokens ?? 0;
+    context3.accumulatedOutputTokens = curr + outputTokens;
+    span.setAttribute(Attr.GenAI.Usage.OutputTokens, context3.accumulatedOutputTokens);
   }
   if (result.finishReason) {
     span.setAttribute(Attr.GenAI.Response.FinishReasons, JSON.stringify([result.finishReason]));
@@ -1599,22 +1639,23 @@ async function setPostCallAttributesV2(span, result, context3, _model) {
   const toolCalls = result.content?.filter(
     (c) => c.type === "tool-call"
   );
-  const alreadySet = span.attributes?.[Attr.GenAI.Response.FinishReasons] !== void 0;
-  if (!alreadySet) {
-    if (result.response?.id) {
-      span.setAttribute(Attr.GenAI.Response.ID, result.response.id);
-    }
-    if (result.response?.modelId) {
-      span.setAttribute(Attr.GenAI.Response.Model, result.response.modelId);
-    }
-    const inputTokens = ensureNumber(result.usage?.inputTokens);
-    if (inputTokens !== void 0) {
-      span.setAttribute(Attr.GenAI.Usage.InputTokens, inputTokens);
-    }
-    const outputTokens = ensureNumber(result.usage?.outputTokens);
-    if (outputTokens !== void 0) {
-      span.setAttribute(Attr.GenAI.Usage.OutputTokens, outputTokens);
-    }
+  if (result.response?.id) {
+    span.setAttribute(Attr.GenAI.Response.ID, result.response.id);
+  }
+  if (result.response?.modelId) {
+    span.setAttribute(Attr.GenAI.Response.Model, result.response.modelId);
+  }
+  const inputTokens = ensureNumber(result.usage?.inputTokens);
+  if (inputTokens !== void 0) {
+    const curr = context3.accumulatedInputTokens ?? 0;
+    context3.accumulatedInputTokens = curr + inputTokens;
+    span.setAttribute(Attr.GenAI.Usage.InputTokens, context3.accumulatedInputTokens);
+  }
+  const outputTokens = ensureNumber(result.usage?.outputTokens);
+  if (outputTokens !== void 0) {
+    const curr = context3.accumulatedOutputTokens ?? 0;
+    context3.accumulatedOutputTokens = curr + outputTokens;
+    span.setAttribute(Attr.GenAI.Usage.OutputTokens, context3.accumulatedOutputTokens);
   }
   if (toolCalls && toolCalls.length > 0) {
     const originalPrompt = context3.originalPrompt || [];
@@ -1648,7 +1689,7 @@ async function setPostCallAttributesV2(span, result, context3, _model) {
       redactionPolicy.captureMessageContent
     );
   }
-  if (result.finishReason && !alreadySet) {
+  if (result.finishReason) {
     span.setAttribute(Attr.GenAI.Response.FinishReasons, JSON.stringify([result.finishReason]));
   }
 }
@@ -1737,6 +1778,11 @@ function axiomAIMiddlewareV3() {
                       context3,
                       model
                     );
+                    setStreamStepAttributes(childSpan, {
+                      inputTokens: statsResult.usage?.inputTokens?.total,
+                      outputTokens: statsResult.usage?.outputTokens?.total,
+                      finishReason: statsResult.finishReason?.unified
+                    });
                     childSpan.end();
                     if (lease.owned) lease.end();
                     controller.terminate();
@@ -1788,22 +1834,23 @@ async function setPostCallAttributesV3(span, result, context3, _model) {
   const toolCalls = result.content?.filter(
     (c) => c.type === "tool-call"
   );
-  const alreadySet = span.attributes?.[Attr.GenAI.Response.FinishReasons] !== void 0;
-  if (!alreadySet) {
-    if (result.response?.id) {
-      span.setAttribute(Attr.GenAI.Response.ID, result.response.id);
-    }
-    if (result.response?.modelId) {
-      span.setAttribute(Attr.GenAI.Response.Model, result.response.modelId);
-    }
-    const inputTokens = ensureNumber(result.usage?.inputTokens?.total);
-    if (inputTokens !== void 0) {
-      span.setAttribute(Attr.GenAI.Usage.InputTokens, inputTokens);
-    }
-    const outputTokens = ensureNumber(result.usage?.outputTokens?.total);
-    if (outputTokens !== void 0) {
-      span.setAttribute(Attr.GenAI.Usage.OutputTokens, outputTokens);
-    }
+  if (result.response?.id) {
+    span.setAttribute(Attr.GenAI.Response.ID, result.response.id);
+  }
+  if (result.response?.modelId) {
+    span.setAttribute(Attr.GenAI.Response.Model, result.response.modelId);
+  }
+  const inputTokens = ensureNumber(result.usage?.inputTokens?.total);
+  if (inputTokens !== void 0) {
+    const curr = context3.accumulatedInputTokens ?? 0;
+    context3.accumulatedInputTokens = curr + inputTokens;
+    span.setAttribute(Attr.GenAI.Usage.InputTokens, context3.accumulatedInputTokens);
+  }
+  const outputTokens = ensureNumber(result.usage?.outputTokens?.total);
+  if (outputTokens !== void 0) {
+    const curr = context3.accumulatedOutputTokens ?? 0;
+    context3.accumulatedOutputTokens = curr + outputTokens;
+    span.setAttribute(Attr.GenAI.Usage.OutputTokens, context3.accumulatedOutputTokens);
   }
   if (toolCalls && toolCalls.length > 0) {
     const originalPrompt = context3.originalPrompt || [];
@@ -1837,7 +1884,7 @@ async function setPostCallAttributesV3(span, result, context3, _model) {
       redactionPolicy.captureMessageContent
     );
   }
-  if (result.finishReason && !alreadySet) {
+  if (result.finishReason) {
     span.setAttribute(
       Attr.GenAI.Response.FinishReasons,
       JSON.stringify([result.finishReason.unified])
@@ -3099,4 +3146,4 @@ export {
   withEvalContext,
   getConfigScope
 };
-//# sourceMappingURL=chunk-P4YKSKIZ.js.map
+//# sourceMappingURL=chunk-3MMFYPFN.js.map