Skip to content

Commit 1793936

Browse files
cynthiajoancynthiajoan
committed
add test and function calling
1 parent 3e6963d commit 1793936

File tree

4 files changed

+344
-7
lines changed

4 files changed

+344
-7
lines changed

‎packages/firebase_ai/firebase_ai/example/lib/pages/server_template_page.dart‎

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -206,6 +206,9 @@ class _ServerTemplatePageState extends State<ServerTemplatePage> {
206206
Content.text(message),
207207
inputs: {
208208
'customerName': message,
209+
'orientation': 'PORTRAIT',
210+
'useFlash': true,
211+
'zoom': 2,
209212
},
210213
);
211214

@@ -215,14 +218,16 @@ class _ServerTemplatePageState extends State<ServerTemplatePage> {
215218
if (functionCalls!.isNotEmpty) {
216219
final functionCall = functionCalls.first;
217220
if (functionCall.name == 'takePicture') {
221+
ByteData catBytes = await rootBundle.load('assets/images/cat.jpg');
222+
var imageBytes = catBytes.buffer.asUint8List();
218223
final functionResult = {
219-
'orientation': 'LANDSCAPE',
220-
'useFlash': true,
221-
'zoom': 2,
224+
'aspectRatio': '16:9',
225+
'mimeType': 'image/jpeg',
226+
'data': base64Encode(imageBytes),
222227
};
223228
var functionResponse = await _chatFunctionSession?.sendMessage(
224229
Content.functionResponse(functionCall.name, functionResult),
225-
inputs: functionResult,
230+
inputs: {},
226231
);
227232
_messages
228233
.add(MessageData(text: functionResponse?.text, fromUser: false));
@@ -255,7 +260,7 @@ class _ServerTemplatePageState extends State<ServerTemplatePage> {
255260
try {
256261
_messages.add(MessageData(text: message, fromUser: true));
257262
var response = await _templateImagenModel?.generateImages(
258-
'new-imagen',
263+
'portrait-googleai',
259264
inputs: {
260265
'animal': message,
261266
},

‎packages/firebase_ai/firebase_ai/lib/src/server_template/template_generative_model.dart‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,32 @@ part of '../base_model.dart';
1818
/// A generative model that connects to a remote server template.
1919
@experimental
2020
final class TemplateGenerativeModel extends BaseTemplateApiClientModel {
21+
@internal
22+
TemplateGenerativeModel.internal({
23+
required String location,
24+
required FirebaseApp app,
25+
required bool useVertexBackend,
26+
bool? useLimitedUseAppCheckTokens,
27+
FirebaseAppCheck? appCheck,
28+
FirebaseAuth? auth,
29+
http.Client? httpClient,
30+
}) : super(
31+
serializationStrategy: useVertexBackend
32+
? VertexSerialization()
33+
: DeveloperSerialization(),
34+
modelUri: useVertexBackend
35+
? _VertexUri(app: app, model: '', location: location)
36+
: _GoogleAIUri(app: app, model: ''),
37+
client: HttpApiClient(
38+
apiKey: app.options.apiKey,
39+
httpClient: httpClient,
40+
requestHeaders: BaseModel.firebaseTokens(
41+
appCheck, auth, app, useLimitedUseAppCheckTokens)),
42+
templateUri: useVertexBackend
43+
? _TemplateVertexUri(app: app, location: location)
44+
: _TemplateGoogleAIUri(app: app),
45+
);
46+
2147
TemplateGenerativeModel._({
2248
required String location,
2349
required FirebaseApp app,

‎packages/firebase_ai/firebase_ai/lib/src/server_template/template_imagen_model.dart‎

Lines changed: 25 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,15 +17,38 @@ part of '../base_model.dart';
1717
/// An image model that connects to a remote server template.
1818
@experimental
1919
final class TemplateImagenModel extends BaseTemplateApiClientModel {
20+
@internal
21+
TemplateImagenModel.internal(
22+
{required FirebaseApp app,
23+
required String location,
24+
required bool useVertexBackend,
25+
bool? useLimitedUseAppCheckTokens,
26+
FirebaseAppCheck? appCheck,
27+
FirebaseAuth? auth,
28+
http.Client? httpClient})
29+
: super(
30+
serializationStrategy: VertexSerialization(),
31+
modelUri: useVertexBackend
32+
? _VertexUri(app: app, model: '', location: location)
33+
: _GoogleAIUri(app: app, model: ''),
34+
client: HttpApiClient(
35+
apiKey: app.options.apiKey,
36+
httpClient: httpClient,
37+
requestHeaders: BaseModel.firebaseTokens(
38+
appCheck, auth, app, useLimitedUseAppCheckTokens)),
39+
templateUri: useVertexBackend
40+
? _TemplateVertexUri(app: app, location: location)
41+
: _TemplateGoogleAIUri(app: app),
42+
);
43+
2044
TemplateImagenModel._(
2145
{required FirebaseApp app,
2246
required String location,
2347
required bool useVertexBackend,
2448
bool? useLimitedUseAppCheckTokens,
2549
FirebaseAppCheck? appCheck,
2650
FirebaseAuth? auth})
27-
:
28-
super(
51+
: super(
2952
serializationStrategy: VertexSerialization(),
3053
modelUri: useVertexBackend
3154
? _VertexUri(app: app, model: '', location: location)

‎packages/firebase_ai/firebase_ai/test/server_template_test.dart‎

Lines changed: 283 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,283 @@
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+
15+
import 'dart:convert';
16+
17+
import 'package:firebase_ai/firebase_ai.dart';
18+
import 'package:firebase_core/firebase_core.dart';
19+
import 'package:flutter_test/flutter_test.dart';
20+
import 'package:http/http.dart' as http;
21+
import 'package:http/testing.dart';
22+
23+
import 'mock.dart';
24+
25+
// A response for generateContent and generateContentStream.
26+
final _arbitraryGenerateContentResponse = {
27+
'candidates': [
28+
{
29+
'content': {
30+
'role': 'model',
31+
'parts': [
32+
{'text': 'Some response'},
33+
],
34+
},
35+
},
36+
],
37+
};
38+
39+
// A response for Imagen's generateImages.
40+
final _arbitraryImagenResponse = {
41+
'predictions': [
42+
{
43+
'mimeType': 'image/png',
44+
'bytesBase64Encoded':
45+
'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNkYAAAAAYAAjCB0C8AAAAASUVORK5CYII='
46+
}
47+
]
48+
};
49+
50+
void main() {
51+
setupFirebaseVertexAIMocks();
52+
late FirebaseApp app;
53+
setUpAll(() async {
54+
app = await Firebase.initializeApp();
55+
});
56+
57+
group('TemplateGenerativeModel', () {
58+
const templateId = 'my-template';
59+
const location = 'us-central1';
60+
61+
TemplateGenerativeModel createModel(http.Client client,
62+
{bool useVertexBackend = true}) {
63+
// ignore: invalid_use_of_internal_member
64+
return TemplateGenerativeModel.internal(
65+
app: app,
66+
location: location,
67+
useVertexBackend: useVertexBackend,
68+
httpClient: client,
69+
);
70+
}
71+
72+
test('generateContent can make successful request', () async {
73+
final mockHttp = MockClient((request) async {
74+
final body = jsonDecode(request.body) as Map<String, Object?>;
75+
expect(request.url.path,
76+
endsWith('/templates/$templateId:templateGenerateContent'));
77+
expect(body['inputs'], {'prompt': 'Some prompt'});
78+
return http.Response(jsonEncode(_arbitraryGenerateContentResponse), 200,
79+
headers: {'content-type': 'application/json'});
80+
});
81+
82+
final model = createModel(mockHttp);
83+
final response = await model
84+
.generateContent(templateId, inputs: {'prompt': 'Some prompt'});
85+
expect(response.text, 'Some response');
86+
});
87+
88+
test('generateContentStream can make successful request', () async {
89+
final mockHttp = MockClient((request) async {
90+
final body = jsonDecode(request.body) as Map<String, Object?>;
91+
expect(request.url.path,
92+
endsWith('/templates/$templateId:templateStreamGenerateContent'));
93+
expect(body['inputs'], {'prompt': 'Some prompt'});
94+
final responsePayload = jsonEncode(_arbitraryGenerateContentResponse);
95+
final stream = Stream.value(utf8.encode('data: $responsePayload'));
96+
final streamedResponse = http.StreamedResponse(stream, 200,
97+
headers: {'content-type': 'application/json'});
98+
return http.Response.fromStream(streamedResponse);
99+
});
100+
101+
final model = createModel(mockHttp);
102+
final responseStream = model
103+
.generateContentStream(templateId, inputs: {'prompt': 'Some prompt'});
104+
final response = await responseStream.first;
105+
expect(response.text, 'Some response');
106+
});
107+
108+
test('templateGenerateContentWithHistory includes history', () async {
109+
final history = [
110+
Content.text('Hi!'),
111+
Content.model([const TextPart('Hello there.')]),
112+
];
113+
final mockHttp = MockClient((request) async {
114+
final body = jsonDecode(request.body) as Map<String, Object?>;
115+
final contents = body['history'] as List;
116+
expect(contents, hasLength(2));
117+
expect(contents[0]['parts'][0]['text'], 'Hi!');
118+
expect(contents[1]['role'], 'model');
119+
return http.Response(jsonEncode(_arbitraryGenerateContentResponse), 200,
120+
headers: {'content-type': 'application/json'});
121+
});
122+
final model = createModel(mockHttp);
123+
await model.templateGenerateContentWithHistory(history, templateId);
124+
});
125+
126+
test('templateGenerateContentWithHistoryStream includes history', () async {
127+
final history = [
128+
Content.text('Hi!'),
129+
Content.model([const TextPart('Hello there.')]),
130+
];
131+
final mockHttp = MockClient((request) async {
132+
final body = jsonDecode(request.body) as Map<String, Object?>;
133+
final contents = body['history'] as List;
134+
expect(contents, hasLength(2));
135+
expect(contents[0]['parts'][0]['text'], 'Hi!');
136+
expect(contents[1]['role'], 'model');
137+
final responsePayload = jsonEncode(_arbitraryGenerateContentResponse);
138+
final stream = Stream.value(utf8.encode('data: $responsePayload'));
139+
final streamedResponse = http.StreamedResponse(stream, 200,
140+
headers: {'content-type': 'application/json'});
141+
return http.Response.fromStream(streamedResponse);
142+
});
143+
final model = createModel(mockHttp);
144+
final responseStream =
145+
model.templateGenerateContentWithHistoryStream(history, templateId);
146+
await responseStream.drain();
147+
});
148+
});
149+
150+
group('TemplateImagenModel', () {
151+
const templateId = 'my-imagen-template';
152+
const location = 'us-central1';
153+
154+
TemplateImagenModel createModel(http.Client client,
155+
{bool useVertexBackend = true}) {
156+
// ignore: invalid_use_of_internal_member
157+
return TemplateImagenModel.internal(
158+
app: app,
159+
location: location,
160+
useVertexBackend: useVertexBackend,
161+
httpClient: client,
162+
);
163+
}
164+
165+
test('generateImages can make successful request', () async {
166+
final mockHttp = MockClient((request) async {
167+
final body = jsonDecode(request.body) as Map<String, Object?>;
168+
expect(request.url.path, endsWith('/templates/$templateId:predict'));
169+
expect(body['inputs'], {'prompt': 'A cat'});
170+
return http.Response(jsonEncode(_arbitraryImagenResponse), 200,
171+
headers: {'content-type': 'application/json'});
172+
});
173+
final model = createModel(mockHttp);
174+
final response =
175+
await model.generateImages(templateId, inputs: {'prompt': 'A cat'});
176+
expect(response.images, hasLength(1));
177+
expect(response.images.first, isA<ImagenInlineImage>());
178+
});
179+
});
180+
181+
group('TemplateChatSession', () {
182+
const templateId = 'my-chat-template';
183+
late TemplateGenerativeModel model;
184+
185+
test('sendMessage adds to history', () async {
186+
final mockHttp = MockClient((request) async {
187+
return http.Response(jsonEncode(_arbitraryGenerateContentResponse), 200,
188+
headers: {'content-type': 'application/json'});
189+
});
190+
// ignore: invalid_use_of_internal_member
191+
model = TemplateGenerativeModel.internal(
192+
app: app,
193+
location: 'us-central1',
194+
useVertexBackend: true,
195+
httpClient: mockHttp,
196+
);
197+
198+
final chat = model.startChat(templateId);
199+
expect(chat.history, isEmpty);
200+
final response = await chat.sendMessage(Content.text('Hi'));
201+
expect(chat.history, hasLength(2));
202+
expect(chat.history.first.parts.first, isA<TextPart>());
203+
expect((chat.history.first.parts.first as TextPart).text, 'Hi');
204+
expect(chat.history.last.role, 'model');
205+
expect(chat.history.last.parts.first, isA<TextPart>());
206+
expect((chat.history.last.parts.first as TextPart).text, response.text);
207+
});
208+
209+
test('sendMessageStream adds to history', () async {
210+
final mockHttp = MockClient((request) async {
211+
final stream = Stream.fromIterable([
212+
'data: ${jsonEncode({
213+
'candidates': [
214+
{
215+
'content': {
216+
'role': 'model',
217+
'parts': [
218+
{'text': 'Some '},
219+
],
220+
},
221+
},
222+
],
223+
})}',
224+
'data: ${jsonEncode({
225+
'candidates': [
226+
{
227+
'content': {
228+
'role': 'model',
229+
'parts': [
230+
{'text': 'response'},
231+
],
232+
},
233+
},
234+
],
235+
})}'
236+
].map(utf8.encode));
237+
final streamedResponse = http.StreamedResponse(stream, 200,
238+
headers: {'content-type': 'application/json'});
239+
return http.Response.fromStream(streamedResponse);
240+
});
241+
// ignore: invalid_use_of_internal_member
242+
model = TemplateGenerativeModel.internal(
243+
app: app,
244+
location: 'us-central1',
245+
useVertexBackend: true,
246+
httpClient: mockHttp,
247+
);
248+
final chat = model.startChat(templateId);
249+
expect(chat.history, isEmpty);
250+
final responseStream = chat.sendMessageStream(Content.text('Hi'));
251+
final responses = await responseStream.toList();
252+
expect(responses, hasLength(2));
253+
expect(chat.history, hasLength(2));
254+
expect(chat.history.first.parts.first, isA<TextPart>());
255+
expect((chat.history.first.parts.first as TextPart).text, 'Hi');
256+
expect(chat.history.last.role, 'model');
257+
expect(chat.history.last.parts.first, isA<TextPart>());
258+
expect((chat.history.last.parts.first as TextPart).text, 'Some response');
259+
});
260+
261+
test('sendMessage with initial history', () async {
262+
final mockHttp = MockClient((request) async {
263+
return http.Response(jsonEncode(_arbitraryGenerateContentResponse), 200,
264+
headers: {'content-type': 'application/json'});
265+
});
266+
// ignore: invalid_use_of_internal_member
267+
model = TemplateGenerativeModel.internal(
268+
app: app,
269+
location: 'us-central1',
270+
useVertexBackend: true,
271+
httpClient: mockHttp,
272+
);
273+
final history = [
274+
Content.text('Hi!'),
275+
Content.model([const TextPart('Hello there.')]),
276+
];
277+
final chat = model.startChat(templateId, history: history);
278+
expect(chat.history, hasLength(2));
279+
await chat.sendMessage(Content.text('How are you?'));
280+
expect(chat.history, hasLength(4));
281+
});
282+
});
283+
}

0 commit comments

Comments
 (0)