Skip to content

Commit 18f5614

Browse files
cynthiajoancynthiajoan
authored
feat(firebaseai): add think feature (#17409)
* add think feature * revert unrelated change * add thinking related test * Update to the thinking feature * Add more unit test, and move thinkingConfig to generationConfig from base class * no need for empty constructor * remove unnecessary documentation * remove unnecessary import * revert content.dart since it's not related * fix analyzer * add final to the thinking budget
1 parent 9fda0bb commit 18f5614

File tree

4 files changed

+193
-4
lines changed

4 files changed

+193
-4
lines changed

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

Lines changed: 31 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -149,6 +149,7 @@ final class UsageMetadata {
149149
{this.promptTokenCount,
150150
this.candidatesTokenCount,
151151
this.totalTokenCount,
152+
this.thoughtsTokenCount,
152153
this.promptTokensDetails,
153154
this.candidatesTokensDetails});
154155

@@ -161,6 +162,9 @@ final class UsageMetadata {
161162
/// Total token count for the generation request (prompt + candidates).
162163
final int? totalTokenCount;
163164

165+
/// Number of tokens present in thoughts output.
166+
final int? thoughtsTokenCount;
167+
164168
/// List of modalities that were processed in the request input.
165169
final List<ModalityTokenCount>? promptTokensDetails;
166170

@@ -175,13 +179,15 @@ UsageMetadata createUsageMetadata({
175179
required int? promptTokenCount,
176180
required int? candidatesTokenCount,
177181
required int? totalTokenCount,
182+
required int? thoughtsTokenCount,
178183
required List<ModalityTokenCount>? promptTokensDetails,
179184
required List<ModalityTokenCount>? candidatesTokensDetails,
180185
}) =>
181186
UsageMetadata._(
182187
promptTokenCount: promptTokenCount,
183188
candidatesTokenCount: candidatesTokenCount,
184189
totalTokenCount: totalTokenCount,
190+
thoughtsTokenCount: thoughtsTokenCount,
185191
promptTokensDetails: promptTokensDetails,
186192
candidatesTokensDetails: candidatesTokensDetails);
187193

@@ -700,10 +706,25 @@ enum ResponseModalities {
700706
const ResponseModalities(this._jsonString);
701707
final String _jsonString;
702708

703-
/// Convert to json format
709+
// ignore: public_member_api_docs
704710
String toJson() => _jsonString;
705711
}
706712

713+
/// Config for thinking features.
714+
class ThinkingConfig {
715+
// ignore: public_member_api_docs
716+
ThinkingConfig({this.thinkingBudget});
717+
718+
/// The number of thoughts tokens that the model should generate.
719+
final int? thinkingBudget;
720+
721+
// ignore: public_member_api_docs
722+
Map<String, Object?> toJson() => {
723+
if (thinkingBudget case final thinkingBudget?)
724+
'thinkingBudget': thinkingBudget,
725+
};
726+
}
727+
707728
/// Configuration options for model generation and outputs.
708729
abstract class BaseGenerationConfig {
709730
// ignore: public_member_api_docs
@@ -829,6 +850,7 @@ final class GenerationConfig extends BaseGenerationConfig {
829850
super.responseModalities,
830851
this.responseMimeType,
831852
this.responseSchema,
853+
this.thinkingConfig,
832854
});
833855

834856
/// The set of character sequences (up to 5) that will stop output generation.
@@ -850,6 +872,12 @@ final class GenerationConfig extends BaseGenerationConfig {
850872
/// a schema; currently this is limited to `application/json`.
851873
final Schema? responseSchema;
852874

875+
/// Config for thinking features.
876+
///
877+
/// An error will be returned if this field is set for models that don't
878+
/// support thinking.
879+
final ThinkingConfig? thinkingConfig;
880+
853881
@override
854882
Map<String, Object?> toJson() => {
855883
...super.toJson(),
@@ -860,6 +888,8 @@ final class GenerationConfig extends BaseGenerationConfig {
860888
'responseMimeType': responseMimeType,
861889
if (responseSchema case final responseSchema?)
862890
'responseSchema': responseSchema.toJson(),
891+
if (thinkingConfig case final thinkingConfig?)
892+
'thinkingConfig': thinkingConfig.toJson(),
863893
};
864894
}
865895

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -244,10 +244,15 @@ UsageMetadata _parseUsageMetadata(Object jsonObject) {
244244
{'totalTokenCount': final int totalTokenCount} => totalTokenCount,
245245
_ => null,
246246
};
247+
final thoughtsTokenCount = switch (jsonObject) {
248+
{'thoughtsTokenCount': final int thoughtsTokenCount} => thoughtsTokenCount,
249+
_ => null,
250+
};
247251
return createUsageMetadata(
248252
promptTokenCount: promptTokenCount,
249253
candidatesTokenCount: candidatesTokenCount,
250254
totalTokenCount: totalTokenCount,
255+
thoughtsTokenCount: thoughtsTokenCount,
251256
promptTokensDetails: null,
252257
candidatesTokensDetails: null,
253258
);

‎packages/firebase_ai/firebase_ai/test/api_test.dart

Lines changed: 31 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,6 @@
1111
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
1212
// See the License for the specific language governing permissions and
1313
// limitations under the License.
14-
1514
import 'package:firebase_ai/firebase_ai.dart';
1615
import 'package:firebase_ai/src/api.dart';
1716

@@ -393,6 +392,7 @@ void main() {
393392
group('GenerationConfig & BaseGenerationConfig', () {
394393
test('GenerationConfig toJson with all fields', () {
395394
final schema = Schema.object(properties: {});
395+
final thinkingConfig = ThinkingConfig(thinkingBudget: 100);
396396
final config = GenerationConfig(
397397
candidateCount: 1,
398398
stopSequences: ['\n', 'stop'],
@@ -404,6 +404,7 @@ void main() {
404404
frequencyPenalty: 0.4,
405405
responseMimeType: 'application/json',
406406
responseSchema: schema,
407+
thinkingConfig: thinkingConfig,
407408
);
408409
expect(config.toJson(), {
409410
'candidateCount': 1,
@@ -415,8 +416,8 @@ void main() {
415416
'frequencyPenalty': 0.4,
416417
'stopSequences': ['\n', 'stop'],
417418
'responseMimeType': 'application/json',
418-
'responseSchema': schema
419-
.toJson(), // Schema itself not schema.toJson() in the provided code
419+
'responseSchema': schema.toJson(),
420+
'thinkingConfig': {'thinkingBudget': 100},
420421
});
421422
});
422423

@@ -435,6 +436,33 @@ void main() {
435436
'responseMimeType': 'text/plain',
436437
});
437438
});
439+
440+
test('GenerationConfig toJson without thinkingConfig', () {
441+
final config = GenerationConfig(temperature: 0.5);
442+
expect(config.toJson(), {'temperature': 0.5});
443+
});
444+
});
445+
446+
group('ThinkingConfig', () {
447+
test('toJson with thinkingBudget set', () {
448+
final config = ThinkingConfig(thinkingBudget: 123);
449+
expect(config.toJson(), {'thinkingBudget': 123});
450+
});
451+
452+
test('toJson with thinkingBudget null', () {
453+
final config = ThinkingConfig();
454+
// Expecting the key to be absent or the value to be explicitly null,
455+
// depending on implementation. Current implementation omits the key.
456+
expect(config.toJson(), {});
457+
});
458+
459+
test('constructor initializes thinkingBudget', () {
460+
final config = ThinkingConfig(thinkingBudget: 456);
461+
expect(config.thinkingBudget, 456);
462+
463+
final configNull = ThinkingConfig();
464+
expect(configNull.thinkingBudget, isNull);
465+
});
438466
});
439467

440468
group('Parsing Functions', () {

‎packages/firebase_ai/firebase_ai/test/developer_api_test.dart

Lines changed: 126 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,126 @@
1+
// Copyright 2025 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
import 'package:firebase_ai/src/developer/api.dart';
15+
import 'package:flutter_test/flutter_test.dart';
16+
17+
void main() {
18+
group('DeveloperSerialization', () {
19+
group('parseGenerateContentResponse', () {
20+
test('parses usageMetadata with thoughtsTokenCount correctly', () {
21+
final jsonResponse = {
22+
'candidates': [
23+
{
24+
'content': {
25+
'role': 'model',
26+
'parts': [
27+
{'text': 'Some generated text.'}
28+
]
29+
},
30+
'finishReason': 'STOP',
31+
}
32+
],
33+
'usageMetadata': {
34+
'promptTokenCount': 10,
35+
'candidatesTokenCount': 5,
36+
'totalTokenCount': 15,
37+
'thoughtsTokenCount': 3,
38+
}
39+
};
40+
final response =
41+
DeveloperSerialization().parseGenerateContentResponse(jsonResponse);
42+
expect(response.usageMetadata, isNotNull);
43+
expect(response.usageMetadata!.promptTokenCount, 10);
44+
expect(response.usageMetadata!.candidatesTokenCount, 5);
45+
expect(response.usageMetadata!.totalTokenCount, 15);
46+
expect(response.usageMetadata!.thoughtsTokenCount, 3);
47+
});
48+
49+
test('parses usageMetadata when thoughtsTokenCount is missing', () {
50+
final jsonResponse = {
51+
'candidates': [
52+
{
53+
'content': {
54+
'role': 'model',
55+
'parts': [
56+
{'text': 'Some generated text.'}
57+
]
58+
},
59+
'finishReason': 'STOP',
60+
}
61+
],
62+
'usageMetadata': {
63+
'promptTokenCount': 10,
64+
'candidatesTokenCount': 5,
65+
'totalTokenCount': 15,
66+
// thoughtsTokenCount is missing
67+
}
68+
};
69+
final response =
70+
DeveloperSerialization().parseGenerateContentResponse(jsonResponse);
71+
expect(response.usageMetadata, isNotNull);
72+
expect(response.usageMetadata!.promptTokenCount, 10);
73+
expect(response.usageMetadata!.candidatesTokenCount, 5);
74+
expect(response.usageMetadata!.totalTokenCount, 15);
75+
expect(response.usageMetadata!.thoughtsTokenCount, isNull);
76+
});
77+
78+
test('parses usageMetadata when thoughtsTokenCount is present but null',
79+
() {
80+
final jsonResponse = {
81+
'candidates': [
82+
{
83+
'content': {
84+
'role': 'model',
85+
'parts': [
86+
{'text': 'Some generated text.'}
87+
]
88+
},
89+
'finishReason': 'STOP',
90+
}
91+
],
92+
'usageMetadata': {
93+
'promptTokenCount': 10,
94+
'candidatesTokenCount': 5,
95+
'totalTokenCount': 15,
96+
'thoughtsTokenCount': null,
97+
}
98+
};
99+
final response =
100+
DeveloperSerialization().parseGenerateContentResponse(jsonResponse);
101+
expect(response.usageMetadata, isNotNull);
102+
expect(response.usageMetadata!.thoughtsTokenCount, isNull);
103+
});
104+
105+
test('parses response when usageMetadata is missing', () {
106+
final jsonResponse = {
107+
'candidates': [
108+
{
109+
'content': {
110+
'role': 'model',
111+
'parts': [
112+
{'text': 'Some generated text.'}
113+
]
114+
},
115+
'finishReason': 'STOP',
116+
}
117+
],
118+
// usageMetadata is missing
119+
};
120+
final response =
121+
DeveloperSerialization().parseGenerateContentResponse(jsonResponse);
122+
expect(response.usageMetadata, isNull);
123+
});
124+
});
125+
});
126+
}

0 commit comments

Comments
 (0)