token counts

Author: nicohrubec

dev-packages/node-integration-tests/suites/tracing/google-genai/scenario-embeddings.mjs

@@ -6,7 +6,7 @@ function startMockGoogleGenAIServer() {
   const app = express();
   app.use(express.json());
 
-  app.post('/v1beta/models/:model\\:embedContent', (req, res) => {
+  app.post('/v1beta/models/:model\\:batchEmbedContents', (req, res) => {
     const model = req.params.model;
 
     if (model === 'error-model') {
@@ -18,14 +18,11 @@ function startMockGoogleGenAIServer() {
       embeddings: [
         {
           values: [0.1, 0.2, 0.3, 0.4, 0.5],
-          statistics: {
-            tokenCount: 8,
-            truncated: false,
-          },
         },
       ],
-      metadata: {
-        billableCharacterCount: 30,
+      usageMetadata: {
+        promptTokenCount: 8,
+        totalTokenCount: 8,
       },
     });
   });

packages/core/src/tracing/google-genai/embeddings.ts

@@ -38,15 +38,27 @@ export function addEmbeddingsRequestAttributes(span: Span, params: Record<string
 
 /**
  * Add response attributes from the Google GenAI embedContent response.
- * The EmbedContentResponse has no usageMetadata/candidates/modelVersion.
- * Token counts come from embeddings[].statistics.tokenCount.
+ * Token counts come from usageMetadata (Gemini API) or embeddings[].statistics.tokenCount (Vertex AI).
  * @see https://ai.google.dev/api/embeddings#EmbedContentResponse
  */
 export function addEmbedContentResponseAttributes(span: Span, response: unknown): void {
   if (!response || typeof response !== 'object') return;
 
   const embedResponse = response as GoogleGenAIEmbedContentResponse;
 
+  // Try usageMetadata first (same shape as GenerateContentResponse)
+  if (embedResponse.usageMetadata && typeof embedResponse.usageMetadata === 'object') {
+    const usage = embedResponse.usageMetadata;
+    if (typeof usage.promptTokenCount === 'number') {
+      span.setAttribute(GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE, usage.promptTokenCount);
+    }
+    if (typeof usage.totalTokenCount === 'number') {
+      span.setAttribute(GEN_AI_USAGE_TOTAL_TOKENS_ATTRIBUTE, usage.totalTokenCount);
+    }
+    return;
+  }
+
+  // Fallback: sum token counts from individual embedding statistics (Vertex AI)
   if (Array.isArray(embedResponse.embeddings)) {
     let totalTokenCount = 0;
     for (const embedding of embedResponse.embeddings) {

packages/core/src/tracing/google-genai/index.ts

@@ -31,7 +31,13 @@ import { buildMethodPath, extractSystemInstructions, getFinalOperationName, getS
 import { CHAT_PATH, CHATS_CREATE_METHOD, GOOGLE_GENAI_SYSTEM_NAME } from './constants';
 import { addEmbedContentResponseAttributes, addEmbeddingsRequestAttributes } from './embeddings';
 import { instrumentStream } from './streaming';
-import type { Candidate, ContentPart, GoogleGenAIIstrumentedMethod, GoogleGenAIOptions, GoogleGenAIResponse } from './types';
+import type {
+  Candidate,
+  ContentPart,
+  GoogleGenAIIstrumentedMethod,
+  GoogleGenAIOptions,
+  GoogleGenAIResponse,
+} from './types';
 import type { ContentListUnion, ContentUnion, Message, PartListUnion } from './utils';
 import { contentUnionToMessages, isEmbeddingsMethod, isStreamingMethod, shouldInstrument } from './utils';
 

packages/core/src/tracing/google-genai/types.ts

@@ -210,6 +210,13 @@ export type GoogleGenAIEmbedContentResponse = {
     /** Billable character count. */
     billableCharacterCount?: number;
   };
+  /** Usage metadata (same shape as GenerateContentResponse). */
+  usageMetadata?: {
+    /** Number of tokens in the request. */
+    promptTokenCount?: number;
+    /** Total token count. */
+    totalTokenCount?: number;
+  };
 };
 
 // Export the response type for use in instrumentation