diff --git a/packages/firebase_ai/firebase_ai/lib/src/api.dart b/packages/firebase_ai/firebase_ai/lib/src/api.dart index 7a482c087f1d..83dbb82407ed 100644 --- 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 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 index a21b17a0ee56..ed5596868d62 100644 --- 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 = {