Skip to content

Commit a9df1d5

Browse files
feat(firebase_ai): Add thoughtsTokenCount to _parseUsageMetadata
- Added `thoughtsTokenCount` to `_parseUsageMetadata` in `packages/firebase_ai/firebase_ai/lib/src/api.dart` to align with the implementation in the developer API. - Added unit tests for `thoughtsTokenCount` parsing in `packages/firebase_ai/firebase_ai/test/api_test.dart`.
1 parent 5199edb commit a9df1d5

File tree

5 files changed

+66
-9
lines changed

5 files changed

+66
-9
lines changed

packages/cloud_functions/cloud_functions_platform_interface/pubspec.yaml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,4 +23,4 @@ dev_dependencies:
2323
flutter_test:
2424
sdk: flutter
2525
mockito: ^5.0.0
26-
pigeon: 25.3.2
26+
pigeon: ^22.7.2

packages/firebase_ai/firebase_ai/example/pubspec.yaml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@ dependencies:
2626
flutter:
2727
sdk: flutter
2828
flutter_markdown: ^0.6.20
29-
flutter_soloud: ^3.1.6
29+
flutter_soloud: ^2.1.0
3030
path_provider: ^2.1.5
3131
record: ^5.2.1
3232

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

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1275,6 +1275,10 @@ UsageMetadata _parseUsageMetadata(Object jsonObject) {
12751275
{'totalTokenCount': final int totalTokenCount} => totalTokenCount,
12761276
_ => null,
12771277
};
1278+
final thoughtsTokenCount = switch (jsonObject) {
1279+
{'thoughtsTokenCount': final int thoughtsTokenCount} => thoughtsTokenCount,
1280+
_ => null,
1281+
};
12781282
final promptTokensDetails = switch (jsonObject) {
12791283
{'promptTokensDetails': final List<Object?> promptTokensDetails} =>
12801284
promptTokensDetails.map(_parseModalityTokenCount).toList(),
@@ -1285,12 +1289,14 @@ UsageMetadata _parseUsageMetadata(Object jsonObject) {
12851289
candidatesTokensDetails.map(_parseModalityTokenCount).toList(),
12861290
_ => null,
12871291
};
1288-
return UsageMetadata._(
1289-
promptTokenCount: promptTokenCount,
1290-
candidatesTokenCount: candidatesTokenCount,
1291-
totalTokenCount: totalTokenCount,
1292-
promptTokensDetails: promptTokensDetails,
1293-
candidatesTokensDetails: candidatesTokensDetails);
1292+
return createUsageMetadata(
1293+
promptTokenCount: promptTokenCount,
1294+
candidatesTokenCount: candidatesTokenCount,
1295+
totalTokenCount: totalTokenCount,
1296+
thoughtsTokenCount: thoughtsTokenCount,
1297+
promptTokensDetails: promptTokensDetails,
1298+
candidatesTokensDetails: candidatesTokensDetails,
1299+
);
12941300
}
12951301

12961302
ModalityTokenCount _parseModalityTokenCount(Object? jsonObject) {

packages/firebase_ai/firebase_ai/pubspec.yaml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ dependencies:
2727
flutter:
2828
sdk: flutter
2929
http: ^1.1.0
30-
meta: ^1.15.0
30+
meta: ^1.12.0
3131
web_socket_channel: ^3.0.1
3232

3333
dev_dependencies:

packages/firebase_ai/firebase_ai/test/api_test.dart

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -615,6 +615,57 @@ void main() {
615615
expect(response.usageMetadata!.candidatesTokensDetails, hasLength(1));
616616
});
617617

618+
group('usageMetadata parsing', () {
619+
test('parses usageMetadata with thoughtsTokenCount correctly', () {
620+
final json = {
621+
'usageMetadata': {
622+
'promptTokenCount': 10,
623+
'candidatesTokenCount': 20,
624+
'totalTokenCount': 30,
625+
'thoughtsTokenCount': 5,
626+
}
627+
};
628+
final response =
629+
VertexSerialization().parseGenerateContentResponse(json);
630+
expect(response.usageMetadata, isNotNull);
631+
expect(response.usageMetadata!.promptTokenCount, 10);
632+
expect(response.usageMetadata!.candidatesTokenCount, 20);
633+
expect(response.usageMetadata!.totalTokenCount, 30);
634+
expect(response.usageMetadata!.thoughtsTokenCount, 5);
635+
});
636+
637+
test('parses usageMetadata when thoughtsTokenCount is missing', () {
638+
final json = {
639+
'usageMetadata': {
640+
'promptTokenCount': 10,
641+
'candidatesTokenCount': 20,
642+
'totalTokenCount': 30,
643+
}
644+
};
645+
final response =
646+
VertexSerialization().parseGenerateContentResponse(json);
647+
expect(response.usageMetadata, isNotNull);
648+
expect(response.usageMetadata!.thoughtsTokenCount, isNull);
649+
});
650+
651+
test(
652+
'parses usageMetadata when thoughtsTokenCount is present but null',
653+
() {
654+
final json = {
655+
'usageMetadata': {
656+
'promptTokenCount': 10,
657+
'candidatesTokenCount': 20,
658+
'totalTokenCount': 30,
659+
'thoughtsTokenCount': null,
660+
}
661+
};
662+
final response =
663+
VertexSerialization().parseGenerateContentResponse(json);
664+
expect(response.usageMetadata, isNotNull);
665+
expect(response.usageMetadata!.thoughtsTokenCount, isNull);
666+
});
667+
});
668+
618669
group('groundingMetadata parsing', () {
619670
test('parses valid response with full grounding metadata', () {
620671
final jsonResponse = {

0 commit comments

Comments
 (0)