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

Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com>

Author: 2 people

packages/firebase_ai/firebase_ai/lib/src/api.dart

@@ -1299,6 +1299,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(),
@@ -1309,12 +1313,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) {

packages/firebase_ai/firebase_ai/test/api_test.dart

@@ -650,6 +650,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 = {