Fix AI SDK usage typings

Author: lforst

js/src/instrumentation/plugins/ai-sdk-plugin.ts

@@ -3964,9 +3964,14 @@ export function extractTokenMetrics(
     return metrics;
   }
 
+  const inputTokenDetails =
+    typeof usage.inputTokens === "object" ? usage.inputTokens : undefined;
+  const outputTokenDetails =
+    typeof usage.outputTokens === "object" ? usage.outputTokens : undefined;
+
   // Extract token counts
   const promptTokens = firstNumber(
-    usage.inputTokens?.total,
+    inputTokenDetails?.total,
     usage.inputTokens,
     usage.promptTokens,
     usage.prompt_tokens,
@@ -3976,7 +3981,7 @@ export function extractTokenMetrics(
   }
 
   const completionTokens = firstNumber(
-    usage.outputTokens?.total,
+    outputTokenDetails?.total,
     usage.outputTokens,
     usage.completionTokens,
     usage.completion_tokens,
@@ -3997,7 +4002,7 @@ export function extractTokenMetrics(
   }
 
   const promptCachedTokens = firstNumber(
-    usage.inputTokens?.cacheRead,
+    inputTokenDetails?.cacheRead,
     usage.inputTokenDetails?.cacheReadTokens,
     usage.cachedInputTokens,
     usage.promptCachedTokens,
@@ -4008,7 +4013,7 @@ export function extractTokenMetrics(
   }
 
   const promptCacheCreationTokens = firstNumber(
-    usage.inputTokens?.cacheWrite,
+    inputTokenDetails?.cacheWrite,
     usage.inputTokenDetails?.cacheWriteTokens,
     usage.promptCacheCreationTokens,
     usage.prompt_cache_creation_tokens,
@@ -4035,7 +4040,7 @@ export function extractTokenMetrics(
   }
 
   const reasoningTokenCount = firstNumber(
-    usage.outputTokens?.reasoning,
+    outputTokenDetails?.reasoning,
     usage.reasoningTokens,
     usage.completionReasoningTokens,
     usage.completion_reasoning_tokens,

js/src/vendor-sdk-types/ai-sdk-common.ts

@@ -12,7 +12,7 @@ interface AISDKTokenBucket {
   [key: string]: unknown;
 }
 
-type AISDKTokenCount = number & AISDKTokenBucket;
+type AISDKTokenCount = number | AISDKTokenBucket;
 
 export interface AISDKUsage {
   inputTokens?: AISDKTokenCount;

js/src/wrappers/ai-sdk-evaluate.test.ts

@@ -19,41 +19,45 @@ import { wrapAISDK } from "./ai-sdk/ai-sdk";
 
 configureNode();
 
+const questions = {
+  category: {
+    type: "choice",
+    instructions: "Which team should handle this request?",
+    criteria: { billing: null, technical: null },
+  },
+  urgency: {
+    type: "score",
+    instructions: "How urgent is this request?",
+    criteria: ["Low", "High"],
+  },
+  duplicate: {
+    type: "boolean",
+    instructions: "Was there a duplicate charge?",
+  },
+};
+
+const answers = {
+  category: {
+    type: "choice",
+    choice: "billing",
+    probabilities: { billing: 0.9, technical: 0.1 },
+  },
+  urgency: {
+    type: "score",
+    score: 0.75,
+    probabilities: { 0: 0.25, 1: 0.75 },
+  },
+  duplicate: { type: "boolean", probability: 0.95 },
+};
+
 const params: AISDKEvaluateParams = {
   model: "typesafe-ai/jev",
   state: { message: "Charged twice" },
-  questions: {
-    category: {
-      type: "choice",
-      instructions: "Which team should handle this request?",
-      criteria: { billing: null, technical: null },
-    },
-    urgency: {
-      type: "score",
-      instructions: "How urgent is this request?",
-      criteria: ["Low", "High"],
-    },
-    duplicate: {
-      type: "boolean",
-      instructions: "Was there a duplicate charge?",
-    },
-  },
+  questions,
 };
 
 const result: AISDKEvaluationResult = {
-  answers: {
-    category: {
-      type: "choice",
-      choice: "billing",
-      probabilities: { billing: 0.9, technical: 0.1 },
-    },
-    urgency: {
-      type: "score",
-      score: 0.75,
-      probabilities: { 0: 0.25, 1: 0.75 },
-    },
-    duplicate: { type: "boolean", probability: 0.95 },
-  },
+  answers,
   usage: { inputTokens: 20, outputTokens: 5, totalTokens: 25 },
   response: { modelId: "typesafe-ai/jev-1.13.0" },
   providerMetadata: {
@@ -120,24 +124,24 @@ describe("AI SDK evaluate instrumentation", () => {
       input: {
         state: params.state,
         questions: [
-          { ...params.questions.category, id: "category" },
-          { ...params.questions.urgency, id: "urgency" },
-          { ...params.questions.duplicate, id: "duplicate" },
+          { ...questions.category, id: "category" },
+          { ...questions.urgency, id: "urgency" },
+          { ...questions.duplicate, id: "duplicate" },
         ],
       },
       output: {
         answers: [
           {
-            ...result.answers.category,
+            ...answers.category,
             confidence: 0.9,
             id: "category",
           },
           {
-            ...result.answers.urgency,
+            ...answers.urgency,
             confidence: 0.8,
             id: "urgency",
           },
-          { ...result.answers.duplicate, id: "duplicate" },
+          { ...answers.duplicate, id: "duplicate" },
         ],
       },
       metadata: {
@@ -187,7 +191,7 @@ describe("AI SDK evaluate instrumentation", () => {
         }),
       ]),
     );
-    expect(result.answers.category).toHaveProperty("probabilities", {
+    expect(answers.category).toHaveProperty("probabilities", {
       billing: 0.9,
       technical: 0.1,
     });