fix(firebaseai): Fix usageMetadata.thoughtsTokenCount (#17608)

andrewheard · google-labs-jules[bot] · web-flow · commit fe9ddd331d0e · 2025-08-08T14:13:32.000-04:00
Co-authored-by: google-labs-jules[bot] &lt;161369871+google-labs-jules[bot]@users.noreply.github.com&gt;
diff --git a/packages/firebase_ai/firebase_ai/lib/src/api.dart b/packages/firebase_ai/firebase_ai/lib/src/api.dart
@@ -1275,6 +1275,10 @@ UsageMetadata _parseUsageMetadata(Object jsonObject) {
     {'totalTokenCount': final int totalTokenCount} => totalTokenCount,
     _ => null,
   };
+  final thoughtsTokenCount = switch (jsonObject) {
+    {'thoughtsTokenCount': final int thoughtsTokenCount} => thoughtsTokenCount,
+    _ => null,
+  };
   final promptTokensDetails = switch (jsonObject) {
     {'promptTokensDetails': final List<Object?> promptTokensDetails} =>
       promptTokensDetails.map(_parseModalityTokenCount).toList(),
@@ -1285,12 +1289,14 @@ UsageMetadata _parseUsageMetadata(Object jsonObject) {
       candidatesTokensDetails.map(_parseModalityTokenCount).toList(),
     _ => null,
   };
-  return UsageMetadata._(
-      promptTokenCount: promptTokenCount,
-      candidatesTokenCount: candidatesTokenCount,
-      totalTokenCount: totalTokenCount,
-      promptTokensDetails: promptTokensDetails,
-      candidatesTokensDetails: candidatesTokensDetails);
+  return createUsageMetadata(
+    promptTokenCount: promptTokenCount,
+    candidatesTokenCount: candidatesTokenCount,
+    totalTokenCount: totalTokenCount,
+    thoughtsTokenCount: thoughtsTokenCount,
+    promptTokensDetails: promptTokensDetails,
+    candidatesTokensDetails: candidatesTokensDetails,
+  );
 }
 
 ModalityTokenCount _parseModalityTokenCount(Object? jsonObject) {
diff --git a/packages/firebase_ai/firebase_ai/test/api_test.dart b/packages/firebase_ai/firebase_ai/test/api_test.dart
@@ -615,6 +615,40 @@ void main() {
         expect(response.usageMetadata!.candidatesTokensDetails, hasLength(1));
       });
 
+      group('usageMetadata parsing', () {
+        test('parses usageMetadata when thoughtsTokenCount is set', () {
+          final json = {
+            'usageMetadata': {
+              'promptTokenCount': 10,
+              'candidatesTokenCount': 20,
+              'totalTokenCount': 30,
+              'thoughtsTokenCount': 5,
+            }
+          };
+          final response =
+              VertexSerialization().parseGenerateContentResponse(json);
+          expect(response.usageMetadata, isNotNull);
+          expect(response.usageMetadata!.promptTokenCount, 10);
+          expect(response.usageMetadata!.candidatesTokenCount, 20);
+          expect(response.usageMetadata!.totalTokenCount, 30);
+          expect(response.usageMetadata!.thoughtsTokenCount, 5);
+        });
+
+        test('parses usageMetadata when thoughtsTokenCount is missing', () {
+          final json = {
+            'usageMetadata': {
+              'promptTokenCount': 10,
+              'candidatesTokenCount': 20,
+              'totalTokenCount': 30,
+            }
+          };
+          final response =
+              VertexSerialization().parseGenerateContentResponse(json);
+          expect(response.usageMetadata, isNotNull);
+          expect(response.usageMetadata!.thoughtsTokenCount, isNull);
+        });
+      });
+
       group('groundingMetadata parsing', () {
         test('parses valid response with full grounding metadata', () {
           final jsonResponse = {
