fix(embeddings): preserve Gemini native token usage (#18244)

changelog.d/fixes/18244-gemini-embedding-usage-metadata.md

@@ -0,0 +1 @@
+- **fix(embeddings):** Preserve Gemini native `usageMetadata.promptTokenCount` in OpenAI-compatible `prompt_tokens` and `total_tokens` for single and batch embeddings, without double-counting batch vectors. Missing or invalid counts safely remain zero; embedding values and indexes are unchanged (#18244).

open-sse/handlers/embeddingStructuredInput.ts

@@ -211,13 +211,34 @@ function embeddingValues(entry: unknown): unknown[] {
   return Array.isArray(values) ? values : [];
 }
 
+/**
+ * Gemini embedContent and batchEmbedContents report one usageMetadata object
+ * at the response root. Read its prompt count without estimating from vector
+ * length or adding per-vector values (which would overcount batch requests).
+ */
+function geminiEmbeddingUsage(data: Record<string, unknown>): {
+  prompt_tokens: number;
+  total_tokens: number;
+} {
+  const metadata = data.usageMetadata;
+  const reported =
+    metadata && typeof metadata === "object" && !Array.isArray(metadata)
+      ? (metadata as Record<string, unknown>).promptTokenCount
+      : undefined;
+  const promptTokens =
+    typeof reported === "number" && Number.isSafeInteger(reported) && reported >= 0
+      ? reported
+      : 0;
+  return { prompt_tokens: promptTokens, total_tokens: promptTokens };
+}
+
 function normalizeGeminiEmbedContentResponse(
   data: Record<string, unknown>
 ): Record<string, unknown> {
   return {
     object: "list",
     data: [{ object: "embedding", embedding: embeddingValues(data.embedding), index: 0 }],
-    usage: { prompt_tokens: 0, total_tokens: 0 },
+    usage: geminiEmbeddingUsage(data),
   };
 }
 
@@ -230,7 +251,7 @@ function normalizeGeminiBatchResponse(data: Record<string, unknown>): Record<str
       embedding: embeddingValues(entry),
       index,
     })),
-    usage: { prompt_tokens: 0, total_tokens: 0 },
+    usage: geminiEmbeddingUsage(data),
   };
 }
 

tests/unit/embeddings-multimodal-7956.test.ts

@@ -391,3 +391,87 @@ test("handleEmbedding clearly rejects modalities not advertised by the resolved
     globalThis.fetch = originalFetch;
   }
 });
+
+test("#18244: normalize native Gemini embedding usage metadata for single and batch requests", async () => {
+  const { prepareStructuredEmbeddingRequest } =
+    await import("../../open-sse/handlers/embeddingStructuredInput.ts");
+  const provider = {
+    id: "gemini",
+    baseUrl: "https://generativelanguage.googleapis.com/v1beta/openai/embeddings",
+    authType: "apikey",
+    authHeader: "bearer",
+    structuredInputProtocol: "gemini-embed-content" as const,
+    models: [],
+  };
+  const options = { fetchMedia: async () => { throw new Error("unexpected fetch"); } };
+  const single = await prepareStructuredEmbeddingRequest(
+    provider, "gemini-embedding-2", { input: "hello" }, "key", options
+  );
+  const batch = await prepareStructuredEmbeddingRequest(
+    provider, "gemini-embedding-2", { input: ["hello", "world"] }, "key", options
+  );
+  const originalSingle = {
+    embedding: { values: [0.25, -0.75] },
+    usageMetadata: { promptTokenCount: 8, promptTokenDetails: [{ modality: "TEXT", tokenCount: 8 }] },
+  };
+  const originalBatch = {
+    embeddings: [{ values: [0.25] }, { values: [-0.75] }],
+    usageMetadata: { promptTokenCount: 21, promptTokenDetails: [{ modality: "TEXT", tokenCount: 21 }] },
+  };
+  assert.deepEqual(single.normalizeResponse?.(originalSingle), {
+    object: "list",
+    data: [{ object: "embedding", embedding: [0.25, -0.75], index: 0 }],
+    usage: { prompt_tokens: 8, total_tokens: 8 },
+  });
+  assert.deepEqual(batch.normalizeResponse?.(originalBatch), {
+    object: "list",
+    data: [
+      { object: "embedding", embedding: [0.25], index: 0 },
+      { object: "embedding", embedding: [-0.75], index: 1 },
+    ],
+    usage: { prompt_tokens: 21, total_tokens: 21 },
+  });
+  // Metadata is response-wide. Preserve vectors and positions; never multiply by batch size.
+  assert.deepEqual(originalBatch.embeddings, [{ values: [0.25] }, { values: [-0.75] }]);
+});
+
+test("#18244: missing or invalid Gemini usage metadata never invents token counts", async () => {
+  const { prepareStructuredEmbeddingRequest } =
+    await import("../../open-sse/handlers/embeddingStructuredInput.ts");
+  const provider = {
+    id: "gemini",
+    baseUrl: "https://generativelanguage.googleapis.com/v1beta/openai/embeddings",
+    authType: "apikey",
+    authHeader: "bearer",
+    structuredInputProtocol: "gemini-embed-content" as const,
+    models: [],
+  };
+  const options = { fetchMedia: async () => { throw new Error("unexpected fetch"); } };
+  const single = await prepareStructuredEmbeddingRequest(
+    provider, "gemini-embedding-2", { input: "hello" }, "key", options
+  );
+  const batch = await prepareStructuredEmbeddingRequest(
+    provider, "gemini-embedding-2", { input: ["hello", "world"] }, "key", options
+  );
+  const missingOrInvalid = [
+    undefined, null, [], "tokens", {}, { promptTokenCount: "8" },
+    { promptTokenCount: -1 }, { promptTokenCount: 1.5 },
+    { promptTokenCount: Number.MAX_SAFE_INTEGER + 1 },
+    { promptTokenCount: Number.POSITIVE_INFINITY },
+  ];
+  for (const metadata of missingOrInvalid) {
+    const expected = { prompt_tokens: 0, total_tokens: 0 };
+    assert.deepEqual(
+      single.normalizeResponse?.({ embedding: { values: [0.25] }, usageMetadata: metadata })?.usage,
+      expected
+    );
+    assert.deepEqual(
+      batch.normalizeResponse?.({ embeddings: [{ values: [0.25] }], usageMetadata: metadata })?.usage,
+      expected
+    );
+  }
+  assert.deepEqual(
+    single.normalizeResponse?.({ embedding: { values: [0.25] }, usageMetadata: { promptTokenCount: 0 } })?.usage,
+    { prompt_tokens: 0, total_tokens: 0 }
+  );
+});

tests/unit/gemini-embedding-2-multimodal.test.ts

@@ -308,3 +308,65 @@ test("handleEmbedding keeps gemini-embedding-001 text batches on the OpenAI shim
     globalThis.fetch = originalFetch;
   }
 });
+
+test("#18244: Gemini native embedContent usage reaches the API response", async () => {
+  const originalFetch = globalThis.fetch;
+  globalThis.fetch = async () => new Response(JSON.stringify({
+    embedding: { values: [0.125, -0.25] },
+    usageMetadata: {
+      promptTokenCount: 8,
+      promptTokenDetails: [{ modality: "TEXT", tokenCount: 8 }],
+    },
+  }), { status: 200, headers: { "content-type": "application/json" } });
+
+  try {
+    const result = await handleEmbedding({
+      body: { model: "gemini/gemini-embedding-2", input: "account for these tokens" },
+      credentials: { apiKey: "test-gemini-token" },
+      log: null,
+    });
+    assert.equal(result.success, true, result.error);
+    const data = result.data as {
+      data: Array<{ object: string; embedding: number[]; index: number }>;
+      usage: { prompt_tokens: number; total_tokens: number };
+    };
+    assert.deepEqual(data.usage, { prompt_tokens: 8, total_tokens: 8 });
+    assert.deepEqual(data.data, [
+      { object: "embedding", embedding: [0.125, -0.25], index: 0 },
+    ]);
+  } finally {
+    globalThis.fetch = originalFetch;
+  }
+});
+
+test("#18244: Gemini native batch token usage is accounted once, not once per embedding", async () => {
+  const originalFetch = globalThis.fetch;
+  globalThis.fetch = async () => new Response(JSON.stringify({
+    embeddings: [{ values: [0.1, 0.2] }, { values: [-0.3, 0.4] }, { values: [0.5, -0.6] }],
+    usageMetadata: {
+      promptTokenCount: 19,
+      promptTokenDetails: [{ modality: "TEXT", tokenCount: 19 }],
+    },
+  }), { status: 200, headers: { "content-type": "application/json" } });
+
+  try {
+    const result = await handleEmbedding({
+      body: { model: "gemini/gemini-embedding-2", input: ["first", "second", "third"] },
+      credentials: { apiKey: "test-gemini-token" },
+      log: null,
+    });
+    assert.equal(result.success, true, result.error);
+    const data = result.data as {
+      data: Array<{ object: string; embedding: number[]; index: number }>;
+      usage: { prompt_tokens: number; total_tokens: number };
+    };
+    assert.deepEqual(data.usage, { prompt_tokens: 19, total_tokens: 19 });
+    assert.deepEqual(data.data, [
+      { object: "embedding", embedding: [0.1, 0.2], index: 0 },
+      { object: "embedding", embedding: [-0.3, 0.4], index: 1 },
+      { object: "embedding", embedding: [0.5, -0.6], index: 2 },
+    ]);
+  } finally {
+    globalThis.fetch = originalFetch;
+  }
+});