I have added parsing for promptTokensDetails and candidatesTokensDetails in the _parseUsageMetadata function.

Author: google-labs-jules[bot]

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

@@ -20,13 +20,15 @@ import '../api.dart'
         Candidate,
         Citation,
         CitationMetadata,
+        ContentModality,
         CountTokensResponse,
         FinishReason,
         GenerateContentResponse,
         GenerationConfig,
         HarmBlockThreshold,
         HarmCategory,
         HarmProbability,
+        ModalityTokenCount,
         PromptFeedback,
         SafetyRating,
         SafetySetting,
@@ -38,6 +40,31 @@ import '../content.dart'
 import '../error.dart';
 import '../tool.dart' show Tool, ToolConfig;
 
+ContentModality _parseContentModality(Object jsonObject) {
+  return switch (jsonObject) {
+    'MODALITY_UNSPECIFIED' => ContentModality.unspecified,
+    'TEXT' => ContentModality.text,
+    'IMAGE' => ContentModality.image,
+    'VIDEO' => ContentModality.video,
+    'AUDIO' => ContentModality.audio,
+    'DOCUMENT' => ContentModality.document,
+    _ => throw unhandledFormat('ContentModality', jsonObject),
+  };
+}
+
+ModalityTokenCount _parseModalityTokenCount(Object? jsonObject) {
+  if (jsonObject is! Map) {
+    throw unhandledFormat('ModalityTokenCount', jsonObject);
+  }
+  final modality = _parseContentModality(jsonObject['modality']);
+
+  if (jsonObject.containsKey('tokenCount')) {
+    return ModalityTokenCount(modality, jsonObject['tokenCount'] as int);
+  } else {
+    return ModalityTokenCount(modality, 0);
+  }
+}
+
 HarmProbability _parseHarmProbability(Object jsonObject) =>
     switch (jsonObject) {
       'UNSPECIFIED' => HarmProbability.unknown,
@@ -251,13 +278,23 @@ UsageMetadata _parseUsageMetadata(Object jsonObject) {
     {'thoughtsTokenCount': final int thoughtsTokenCount} => thoughtsTokenCount,
     _ => null,
   };
+  final promptTokensDetails = switch (jsonObject) {
+    {'promptTokensDetails': final List<Object?> promptTokensDetails} =>
+      promptTokensDetails.map(_parseModalityTokenCount).toList(),
+    _ => null,
+  };
+  final candidatesTokensDetails = switch (jsonObject) {
+    {'candidatesTokensDetails': final List<Object?> candidatesTokensDetails} =>
+      candidatesTokensDetails.map(_parseModalityTokenCount).toList(),
+    _ => null,
+  };
   return createUsageMetadata(
     promptTokenCount: promptTokenCount,
     candidatesTokenCount: candidatesTokenCount,
     totalTokenCount: totalTokenCount,
     thoughtsTokenCount: thoughtsTokenCount,
-    promptTokensDetails: null,
-    candidatesTokensDetails: null,
+    promptTokensDetails: promptTokensDetails,
+    candidatesTokensDetails: candidatesTokensDetails,
   );
 }
 

packages/firebase_ai/firebase_ai/test/developer_api_test.dart

@@ -153,6 +153,40 @@ void main() {
         expect((part as InlineDataPart).mimeType, 'application/octet-stream');
         expect(part.bytes, inlineData);
       });
+
+      test(
+          'parses usageMetadata with prompt and candidate token details correctly',
+          () {
+        final jsonResponse = {
+          'usageMetadata': {
+            'promptTokenCount': 10,
+            'candidatesTokenCount': 25,
+            'totalTokenCount': 35,
+            'promptTokensDetails': [
+              {'modality': 'TEXT', 'tokenCount': 10}
+            ],
+            'candidatesTokensDetails': [
+              {'modality': 'TEXT', 'tokenCount': 25}
+            ],
+          }
+        };
+
+        final response =
+            DeveloperSerialization().parseGenerateContentResponse(jsonResponse);
+
+        expect(response.usageMetadata, isNotNull);
+        expect(response.usageMetadata!.promptTokenCount, 10);
+        expect(response.usageMetadata!.candidatesTokenCount, 25);
+        expect(response.usageMetadata!.totalTokenCount, 35);
+        expect(response.usageMetadata!.promptTokensDetails, isNotNull);
+        expect(response.usageMetadata!.promptTokensDetails, hasLength(1));
+        expect(response.usageMetadata!.promptTokensDetails!.first.tokenCount,
+            10);
+        expect(response.usageMetadata!.candidatesTokensDetails, isNotNull);
+        expect(response.usageMetadata!.candidatesTokensDetails, hasLength(1));
+        expect(response.usageMetadata!.candidatesTokensDetails!.first.tokenCount,
+            25);
+      });
     });
   });
 }