Skip to content

Commit fe9ddd3

Browse files
andrewheardgoogle-labs-jules[bot]
andrewheard
and
google-labs-jules[bot]
authored
fix(firebaseai): Fix usageMetadata.thoughtsTokenCount (#17608)
Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com>
1 parent 467eaa1 commit fe9ddd3

2 files changed

Lines changed: 46 additions & 6 deletions

File tree

‎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/test/api_test.dart‎

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

618+
group('usageMetadata parsing', () {
619+
test('parses usageMetadata when thoughtsTokenCount is set', () {
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+
618652
group('groundingMetadata parsing', () {
619653
test('parses valid response with full grounding metadata', () {
620654
final jsonResponse = {

0 commit comments

Comments
 (0)