💄 style: Add cached token count to usage of GoogleAI and VertexAI (#8545)

samurai00 · web-flow · commit be8df5a7ff8d · 2025-07-24T10:56:03.000+08:00
diff --git a/src/libs/model-runtime/utils/streams/google-ai.test.ts b/src/libs/model-runtime/utils/streams/google-ai.test.ts
@@ -186,6 +186,52 @@ describe('GoogleGenerativeAIStream', () => {
     ]);
   });
 
+  it('should handle token count with cached token count', async () => {
+    vi.spyOn(uuidModule, 'nanoid').mockReturnValueOnce('1');
+
+    const data = {
+      candidates: [{ content: { role: 'model' }, finishReason: 'STOP', index: 0 }],
+      usageMetadata: {
+        promptTokenCount: 15725,
+        candidatesTokenCount: 1053,
+        totalTokenCount: 16778,
+        cachedContentTokenCount: 14286,
+        promptTokensDetails: [{ modality: 'TEXT', tokenCount: 15725 }],
+        cacheTokensDetails: [{ modality: 'TEXT', tokenCount: 14286 }],
+      },
+      modelVersion: 'gemini-2.0-flash-exp',
+    };
+
+    const mockGoogleStream = new ReadableStream({
+      start(controller) {
+        controller.enqueue(data);
+
+        controller.close();
+      },
+    });
+
+    const protocolStream = GoogleGenerativeAIStream(mockGoogleStream);
+
+    const decoder = new TextDecoder();
+    const chunks = [];
+
+    // @ts-ignore
+    for await (const chunk of protocolStream) {
+      chunks.push(decoder.decode(chunk, { stream: true }));
+    }
+
+    expect(chunks).toEqual([
+      // stop
+      'id: chat_1\n',
+      'event: stop\n',
+      `data: "STOP"\n\n`,
+      // usage
+      'id: chat_1\n',
+      'event: usage\n',
+      `data: {"inputCachedTokens":14286,"inputTextTokens":15725,"outputTextTokens":1053,"totalInputTokens":15725,"totalOutputTokens":1053,"totalTokens":16778}\n\n`,
+    ]);
+  });
+
   it('should handle stop with content', async () => {
     vi.spyOn(uuidModule, 'nanoid').mockReturnValueOnce('1');
 
diff --git a/src/libs/model-runtime/utils/streams/google-ai.ts b/src/libs/model-runtime/utils/streams/google-ai.ts
@@ -33,10 +33,10 @@ const transformGoogleGenerativeAIStream = (
       { data: candidate.finishReason, id: context?.id, type: 'stop' },
       {
         data: {
-          // TODO: Google SDK 0.24.0 don't have promptTokensDetails types
-          inputImageTokens: usage.promptTokensDetails?.find((i: any) => i.modality === 'IMAGE')
+          inputCachedTokens: usage.cachedContentTokenCount,
+          inputImageTokens: usage.promptTokensDetails?.find((i) => i.modality === 'IMAGE')
             ?.tokenCount,
-          inputTextTokens: usage.promptTokensDetails?.find((i: any) => i.modality === 'TEXT')
+          inputTextTokens: usage.promptTokensDetails?.find((i) => i.modality === 'TEXT')
             ?.tokenCount,
           outputReasoningTokens: reasoningTokens,
           outputTextTokens,
@@ -79,7 +79,7 @@ const transformGoogleGenerativeAIStream = (
     // 首先检查是否为 reasoning 内容 (thought: true)
     if (Array.isArray(candidate.content?.parts) && candidate.content.parts.length > 0) {
       for (const part of candidate.content.parts) {
-        if (part && part.text && (part as any).thought === true) {
+        if (part && part.text && part.thought === true) {
           return { data: part.text, id: context.id, type: 'reasoning' };
         }
       }
diff --git a/src/libs/model-runtime/utils/streams/vertex-ai.ts b/src/libs/model-runtime/utils/streams/vertex-ai.ts
@@ -31,13 +31,11 @@ const transformVertexAIStream = (
       { data: candidate.finishReason, id: context?.id, type: 'stop' },
       {
         data: {
-          // TODO: Google SDK 0.24.0 don't have promptTokensDetails types
-          inputImageTokens: (usage as any).promptTokensDetails?.find(
-            (i: any) => i.modality === 'IMAGE',
-          )?.tokenCount,
-          inputTextTokens: (usage as any).promptTokensDetails?.find(
-            (i: any) => i.modality === 'TEXT',
-          )?.tokenCount,
+          inputCachedTokens: usage.cachedContentTokenCount,
+          inputImageTokens: usage.promptTokensDetails?.find((i) => i.modality === 'IMAGE')
+            ?.tokenCount,
+          inputTextTokens: usage.promptTokensDetails?.find((i) => i.modality === 'TEXT')
+            ?.tokenCount,
           outputReasoningTokens,
           outputTextTokens,
           totalInputTokens: usage.promptTokenCount,
@@ -56,7 +54,7 @@ const transformVertexAIStream = (
     candidate.content.parts.length > 0
   ) {
     for (const part of candidate.content.parts) {
-      if (part && part.text && (part as any).thought === true) {
+      if (part && part.text && part.thought === true) {
         return { data: part.text, id: context.id, type: 'reasoning' };
       }
     }
