@tanstack/ai 0.32.0 → 0.34.0

package/dist/esm/middlewares/otel.js

@@ -1,5 +1,16 @@
-import { SpanStatusCode, SpanKind, trace, context } from "@opentelemetry/api";
+import { SpanKind, SpanStatusCode, trace, context } from "@opentelemetry/api";
 import { MAX_TOKENS_KEYS, NESTED_MAX_TOKENS_KEY } from "../utilities/sampling-keys.js";
+import { firstNumber } from "../utilities/numbers.js";
+import { errorMessage, errorTypeName } from "../utilities/errors.js";
+import { usageAttributes } from "./usage-attributes.js";
+const OPERATION_NAME = {
+  chat: "chat",
+  image: "image_generation",
+  video: "video_generation",
+  audio: "audio_generation",
+  tts: "text_to_speech",
+  transcription: "transcription"
+};
 const stateByCtx = /* @__PURE__ */ new WeakMap();
 const DEFAULT_MAX_CONTENT_LENGTH = 1e5;
 const REDACTION_FAILED_SENTINEL = "[redaction_failed]";
@@ -51,68 +62,6 @@ function messageEventName(role) {
       return `gen_ai.${role}.message`;
   }
 }
-function firstNumber(...candidates) {
-  for (const candidate of candidates) {
-    if (typeof candidate === "number" && Number.isFinite(candidate)) {
-      return candidate;
-    }
-  }
-  return void 0;
-}
-function usageAttributes(usage) {
-  const attrs = {
-    "gen_ai.usage.input_tokens": usage.promptTokens,
-    "gen_ai.usage.output_tokens": usage.completionTokens
-  };
-  const optional = [
-    ["gen_ai.usage.total_tokens", usage.totalTokens],
-    ["gen_ai.usage.cost", usage.cost],
-    [
-      "gen_ai.usage.cache_read.input_tokens",
-      usage.promptTokensDetails?.cachedTokens
-    ],
-    [
-      "gen_ai.usage.cache_creation.input_tokens",
-      usage.promptTokensDetails?.cacheWriteTokens
-    ],
-    [
-      "gen_ai.usage.reasoning.output_tokens",
-      usage.completionTokensDetails?.reasoningTokens
-    ],
-    ["tanstack.ai.usage.duration_seconds", usage.durationSeconds],
-    ["tanstack.ai.usage.upstream_cost", usage.costDetails?.upstreamCost],
-    [
-      "tanstack.ai.usage.upstream_input_cost",
-      usage.costDetails?.upstreamInputCost
-    ],
-    [
-      "tanstack.ai.usage.upstream_output_cost",
-      usage.costDetails?.upstreamOutputCost
-    ]
-  ];
-  for (const [key, value] of optional) {
-    const num = firstNumber(value);
-    if (num !== void 0) attrs[key] = num;
-  }
-  return attrs;
-}
-function errorMessage(err) {
-  if (err instanceof Error) return err.message;
-  if (typeof err === "string") return err;
-  if (err && typeof err === "object" && "message" in err) {
-    const m = err.message;
-    if (typeof m === "string") return m;
-  }
-  return void 0;
-}
-function errorTypeName(err) {
-  if (err instanceof Error) return err.name || "Error";
-  if (err && typeof err === "object" && "name" in err) {
-    const n = err.name;
-    if (typeof n === "string") return n;
-  }
-  return "Error";
-}
 function safeCall(label, fn) {
   try {
     return fn();
@@ -180,17 +129,69 @@ function otelMiddleware(options) {
     span.end();
     state.currentIterationSpan = null;
   };
+  const mediaSpans = /* @__PURE__ */ new WeakMap();
+  const recordMediaDuration = (ctx, durationMs, errorType) => {
+    if (!durationHistogram) return;
+    durationHistogram.record(durationMs / 1e3, {
+      "gen_ai.system": ctx.provider,
+      "gen_ai.operation.name": OPERATION_NAME[ctx.activity],
+      "gen_ai.request.model": ctx.model,
+      ...errorType ? { "error.type": errorType } : {}
+    });
+  };
+  const startMediaSpan = (ctx) => {
+    safeCall("otel.onStart", () => {
+      const operationName = OPERATION_NAME[ctx.activity];
+      const info = { kind: "generation", ctx };
+      const name = safeCall("otel.spanNameFormatter", () => spanNameFormatter?.(info)) ?? `${operationName} ${ctx.model}`;
+      const baseOptions = {
+        kind: SpanKind.CLIENT,
+        attributes: {
+          "gen_ai.system": ctx.provider,
+          "gen_ai.operation.name": operationName,
+          "gen_ai.request.model": ctx.model
+        }
+      };
+      const spanOptions = safeCall(
+        "otel.onBeforeSpanStart",
+        () => onBeforeSpanStart?.(info, baseOptions)
+      ) ?? baseOptions;
+      const span = tracer.startSpan(name, spanOptions);
+      const enriched = safeCall(
+        "otel.attributeEnricher",
+        () => attributeEnricher?.(info)
+      );
+      if (enriched) span.setAttributes(enriched);
+      mediaSpans.set(ctx, span);
+    });
+  };
+  const endMediaSpan = (ctx, finalize) => {
+    const span = mediaSpans.get(ctx);
+    mediaSpans.delete(ctx);
+    if (!span) return;
+    finalize(span);
+    safeCall(
+      "otel.onSpanEnd",
+      () => onSpanEnd?.({ kind: "generation", ctx }, span)
+    );
+    span.end();
+  };
   return {
     name: "otel",
     onStart(ctx) {
+      if (ctx.activity !== "chat") {
+        startMediaSpan(ctx);
+        return;
+      }
+      const chatCtx = ctx;
       safeCall("otel.onStart", () => {
-        const info = { kind: "chat", ctx };
-        const name = safeCall("otel.spanNameFormatter", () => spanNameFormatter?.(info)) ?? `chat ${ctx.model}`;
+        const info = { kind: "chat", ctx: chatCtx };
+        const name = safeCall("otel.spanNameFormatter", () => spanNameFormatter?.(info)) ?? `chat ${chatCtx.model}`;
         const baseOptions = {
           kind: SpanKind.INTERNAL,
           attributes: {
-            "gen_ai.system": ctx.provider,
-            "gen_ai.request.model": ctx.model
+            "gen_ai.system": chatCtx.provider,
+            "gen_ai.request.model": chatCtx.model
             // NOTE: `gen_ai.operation.name` is deliberately NOT set on the
             // root span. The root represents a `chat()` invocation that may
             // span multiple model calls; only iteration spans correspond to
@@ -209,14 +210,15 @@ function otelMiddleware(options) {
           () => attributeEnricher?.(info)
         );
         if (enriched) rootSpan.setAttributes(enriched);
-        stateByCtx.set(ctx, {
+        stateByCtx.set(chatCtx, {
           rootSpan,
           currentIterationSpan: null,
           toolSpans: /* @__PURE__ */ new Map(),
           iterationCount: 0,
           assistantTextBuffer: "",
           assistantTextBufferTruncated: false,
-          startTime: Date.now()
+          startTime: Date.now(),
+          lastFinishReason: null
         });
       });
     },
@@ -348,6 +350,7 @@ function otelMiddleware(options) {
           appendAssistantText(state, chunk.delta);
         }
         if (chunk.type !== "RUN_FINISHED") return;
+        if (chunk.finishReason) state.lastFinishReason = chunk.finishReason;
         const span = state.currentIterationSpan;
         if (!span) return;
         if (chunk.finishReason) {
@@ -376,14 +379,22 @@ function otelMiddleware(options) {
       return void 0;
     },
     onUsage(ctx, usage) {
+      if (ctx.activity !== "chat") {
+        safeCall("otel.onUsage", () => {
+          const span = mediaSpans.get(ctx);
+          if (span) span.setAttributes(usageAttributes(usage));
+        });
+        return;
+      }
+      const chatCtx = ctx;
       safeCall("otel.onUsage", () => {
-        const state = stateByCtx.get(ctx);
+        const state = stateByCtx.get(chatCtx);
         if (!state) return;
         if (tokenHistogram) {
           const metricAttrs = {
-            "gen_ai.system": ctx.provider,
+            "gen_ai.system": chatCtx.provider,
             "gen_ai.operation.name": "chat",
-            "gen_ai.request.model": ctx.model
+            "gen_ai.request.model": chatCtx.model
           };
           tokenHistogram.record(usage.promptTokens, {
             ...metricAttrs,
@@ -506,8 +517,23 @@ function otelMiddleware(options) {
       });
     },
     onError(ctx, info) {
+      if (ctx.activity !== "chat") {
+        safeCall("otel.onError", () => {
+          const message = errorMessage(info.error);
+          endMediaSpan(ctx, (span) => {
+            span.recordException(info.error);
+            span.setStatus({
+              code: SpanStatusCode.ERROR,
+              ...message !== void 0 ? { message } : {}
+            });
+          });
+          recordMediaDuration(ctx, info.duration, errorTypeName(info.error));
+        });
+        return;
+      }
+      const chatCtx = ctx;
       safeCall("otel.onError", () => {
-        const state = stateByCtx.get(ctx);
+        const state = stateByCtx.get(chatCtx);
         if (!state) return;
         const errType = errorTypeName(info.error);
         const message = errorMessage(info.error);
@@ -525,7 +551,7 @@ function otelMiddleware(options) {
             () => onSpanEnd?.(
               {
                 kind: "iteration",
-                ctx,
+                ctx: chatCtx,
                 iteration: state.iterationCount - 1
               },
               iterationSpan
@@ -543,7 +569,7 @@ function otelMiddleware(options) {
             () => onSpanEnd?.(
               {
                 kind: "tool",
-                ctx,
+                ctx: chatCtx,
                 toolCallId: id,
                 toolName,
                 iteration: state.iterationCount - 1
@@ -561,23 +587,37 @@ function otelMiddleware(options) {
         });
         if (durationHistogram) {
           durationHistogram.record(info.duration / 1e3, {
-            "gen_ai.system": ctx.provider,
+            "gen_ai.system": chatCtx.provider,
             "gen_ai.operation.name": "chat",
-            "gen_ai.request.model": ctx.model,
+            "gen_ai.request.model": chatCtx.model,
             "error.type": errType
           });
         }
         safeCall(
           "otel.onSpanEnd",
-          () => onSpanEnd?.({ kind: "chat", ctx }, state.rootSpan)
+          () => onSpanEnd?.({ kind: "chat", ctx: chatCtx }, state.rootSpan)
         );
         state.rootSpan.end();
-        stateByCtx.delete(ctx);
+        stateByCtx.delete(chatCtx);
       });
     },
     onAbort(ctx, info) {
+      if (ctx.activity !== "chat") {
+        safeCall("otel.onAbort", () => {
+          endMediaSpan(ctx, (span) => {
+            span.setAttribute("tanstack.ai.completion.reason", "cancelled");
+            span.setStatus({
+              code: SpanStatusCode.ERROR,
+              message: info.reason ?? "cancelled"
+            });
+          });
+          recordMediaDuration(ctx, info.duration, "cancelled");
+        });
+        return;
+      }
+      const chatCtx = ctx;
       safeCall("otel.onAbort", () => {
-        const state = stateByCtx.get(ctx);
+        const state = stateByCtx.get(chatCtx);
         if (!state) return;
         const closeCancelled = (span) => {
           span.setAttribute("tanstack.ai.completion.reason", "cancelled");
@@ -591,7 +631,7 @@ function otelMiddleware(options) {
             () => onSpanEnd?.(
               {
                 kind: "iteration",
-                ctx,
+                ctx: chatCtx,
                 iteration: state.iterationCount - 1
               },
               iterationSpan
@@ -608,7 +648,7 @@ function otelMiddleware(options) {
             () => onSpanEnd?.(
               {
                 kind: "tool",
-                ctx,
+                ctx: chatCtx,
                 toolCallId: id,
                 toolName,
                 iteration: state.iterationCount - 1
@@ -622,23 +662,33 @@ function otelMiddleware(options) {
         closeCancelled(state.rootSpan);
         if (durationHistogram) {
           durationHistogram.record(info.duration / 1e3, {
-            "gen_ai.system": ctx.provider,
+            "gen_ai.system": chatCtx.provider,
             "gen_ai.operation.name": "chat",
-            "gen_ai.request.model": ctx.model,
+            "gen_ai.request.model": chatCtx.model,
             "error.type": "cancelled"
           });
         }
         safeCall(
           "otel.onSpanEnd",
-          () => onSpanEnd?.({ kind: "chat", ctx }, state.rootSpan)
+          () => onSpanEnd?.({ kind: "chat", ctx: chatCtx }, state.rootSpan)
         );
         state.rootSpan.end();
-        stateByCtx.delete(ctx);
+        stateByCtx.delete(chatCtx);
       });
     },
     onFinish(ctx, info) {
+      if (ctx.activity !== "chat") {
+        safeCall("otel.onFinish", () => {
+          endMediaSpan(ctx, (span) => {
+            if (info.usage) span.setAttributes(usageAttributes(info.usage));
+          });
+          recordMediaDuration(ctx, info.duration);
+        });
+        return;
+      }
+      const chatCtx = ctx;
       safeCall("otel.onFinish", () => {
-        const state = stateByCtx.get(ctx);
+        const state = stateByCtx.get(chatCtx);
         if (!state) return;
         for (const [id, entry] of state.toolSpans) {
           const { span, toolName } = entry;
@@ -648,7 +698,7 @@ function otelMiddleware(options) {
             () => onSpanEnd?.(
               {
                 kind: "tool",
-                ctx,
+                ctx: chatCtx,
                 toolCallId: id,
                 toolName,
                 iteration: state.iterationCount - 1
@@ -659,20 +709,20 @@ function otelMiddleware(options) {
           span.end();
           state.toolSpans.delete(id);
         }
-        closeIterationSpan(state, ctx);
+        closeIterationSpan(state, chatCtx);
         if (durationHistogram) {
           durationHistogram.record(info.duration / 1e3, {
-            "gen_ai.system": ctx.provider,
+            "gen_ai.system": chatCtx.provider,
             "gen_ai.operation.name": "chat",
-            "gen_ai.request.model": ctx.model
+            "gen_ai.request.model": chatCtx.model
           });
         }
         if (info.usage) {
           state.rootSpan.setAttributes(usageAttributes(info.usage));
         }
-        if (info.finishReason) {
+        if (state.lastFinishReason) {
           state.rootSpan.setAttribute("gen_ai.response.finish_reasons", [
-            info.finishReason
+            state.lastFinishReason
           ]);
         }
         state.rootSpan.setAttribute(
@@ -681,10 +731,10 @@ function otelMiddleware(options) {
         );
         safeCall(
           "otel.onSpanEnd",
-          () => onSpanEnd?.({ kind: "chat", ctx }, state.rootSpan)
+          () => onSpanEnd?.({ kind: "chat", ctx: chatCtx }, state.rootSpan)
         );
         state.rootSpan.end();
-        stateByCtx.delete(ctx);
+        stateByCtx.delete(chatCtx);
       });
     }
   };