Skip to content

Commit f8ebbaf

Browse files
authored
fix(firebase_ai): Add GroundingMetadata parsing for Developer API (#17657)
* fix(firebase_ai): Add `GroundingMetadata` parsing for Developer API Adds parsing of `GroundingMetadata` to the `developer/api.dart` response parsing, along with tests that mirror the existing ones in `api_test.dart`. The `GoogleSearch` tool is already being serialized for requests, but I added tests to `google_ai_generative_model_test.dart` to confirm this. * Fix import CI issues
1 parent 8c0802d commit f8ebbaf

File tree

4 files changed

+401
-2
lines changed

4 files changed

+401
-2
lines changed

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

Lines changed: 122 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,14 +24,20 @@ import '../api.dart'
2424
FinishReason,
2525
GenerateContentResponse,
2626
GenerationConfig,
27+
GroundingChunk,
28+
GroundingMetadata,
29+
GroundingSupport,
2730
HarmBlockThreshold,
2831
HarmCategory,
2932
HarmProbability,
3033
PromptFeedback,
3134
SafetyRating,
3235
SafetySetting,
36+
SearchEntryPoint,
37+
Segment,
3338
SerializationStrategy,
3439
UsageMetadata,
40+
WebGroundingChunk,
3541
createUsageMetadata;
3642
import '../content.dart'
3743
show Content, FunctionCall, InlineDataPart, Part, TextPart;
@@ -206,6 +212,11 @@ Candidate _parseCandidate(Object? jsonObject) {
206212
{'finishMessage': final String finishMessage} => finishMessage,
207213
_ => null
208214
},
215+
groundingMetadata: switch (jsonObject) {
216+
{'groundingMetadata': final Object groundingMetadata} =>
217+
_parseGroundingMetadata(groundingMetadata),
218+
_ => null
219+
},
209220
);
210221
}
211222

@@ -299,6 +310,117 @@ Citation _parseCitationSource(Object? jsonObject) {
299310
);
300311
}
301312

313+
GroundingMetadata _parseGroundingMetadata(Object? jsonObject) {
314+
if (jsonObject is! Map) {
315+
throw unhandledFormat('GroundingMetadata', jsonObject);
316+
}
317+
318+
final searchEntryPoint = switch (jsonObject) {
319+
{'searchEntryPoint': final Object? searchEntryPoint} =>
320+
_parseSearchEntryPoint(searchEntryPoint),
321+
_ => null,
322+
};
323+
final groundingChunks = switch (jsonObject) {
324+
{'groundingChunks': final List<Object?> groundingChunks} =>
325+
groundingChunks.map(_parseGroundingChunk).toList(),
326+
_ => null,
327+
} ??
328+
[];
329+
// Filters out null elements, which are returned from _parseGroundingSupport when
330+
// segment is null.
331+
final groundingSupport = switch (jsonObject) {
332+
{'groundingSupport': final List<Object?> groundingSupport} =>
333+
groundingSupport
334+
.map(_parseGroundingSupport)
335+
.whereType<GroundingSupport>()
336+
.toList(),
337+
_ => null,
338+
} ??
339+
[];
340+
final webSearchQueries = switch (jsonObject) {
341+
{'webSearchQueries': final List<String>? webSearchQueries} =>
342+
webSearchQueries,
343+
_ => null,
344+
} ??
345+
[];
346+
347+
return GroundingMetadata(
348+
searchEntryPoint: searchEntryPoint,
349+
groundingChunks: groundingChunks,
350+
groundingSupport: groundingSupport,
351+
webSearchQueries: webSearchQueries);
352+
}
353+
354+
Segment _parseSegment(Object? jsonObject) {
355+
if (jsonObject is! Map) {
356+
throw unhandledFormat('Segment', jsonObject);
357+
}
358+
359+
return Segment(
360+
partIndex: (jsonObject['partIndex'] as int?) ?? 0,
361+
startIndex: (jsonObject['startIndex'] as int?) ?? 0,
362+
endIndex: (jsonObject['endIndex'] as int?) ?? 0,
363+
text: (jsonObject['text'] as String?) ?? '');
364+
}
365+
366+
WebGroundingChunk _parseWebGroundingChunk(Object? jsonObject) {
367+
if (jsonObject is! Map) {
368+
throw unhandledFormat('WebGroundingChunk', jsonObject);
369+
}
370+
371+
return WebGroundingChunk(
372+
uri: jsonObject['uri'] as String?,
373+
title: jsonObject['title'] as String?,
374+
domain: jsonObject['domain'] as String?,
375+
);
376+
}
377+
378+
GroundingChunk _parseGroundingChunk(Object? jsonObject) {
379+
if (jsonObject is! Map) {
380+
throw unhandledFormat('GroundingChunk', jsonObject);
381+
}
382+
383+
return GroundingChunk(
384+
web: jsonObject['web'] != null
385+
? _parseWebGroundingChunk(jsonObject['web'])
386+
: null,
387+
);
388+
}
389+
390+
GroundingSupport? _parseGroundingSupport(Object? jsonObject) {
391+
if (jsonObject is! Map) {
392+
throw unhandledFormat('GroundingSupport', jsonObject);
393+
}
394+
395+
final segment = switch (jsonObject) {
396+
{'segment': final Object? segment} => _parseSegment(segment),
397+
_ => null,
398+
};
399+
if (segment == null) {
400+
return null;
401+
}
402+
403+
return GroundingSupport(
404+
segment: segment,
405+
groundingChunkIndices:
406+
(jsonObject['groundingChunkIndices'] as List<int>?) ?? []);
407+
}
408+
409+
SearchEntryPoint _parseSearchEntryPoint(Object? jsonObject) {
410+
if (jsonObject is! Map) {
411+
throw unhandledFormat('SearchEntryPoint', jsonObject);
412+
}
413+
414+
final renderedContent = jsonObject['renderedContent'] as String?;
415+
if (renderedContent == null) {
416+
throw unhandledFormat('SearchEntryPoint', jsonObject);
417+
}
418+
419+
return SearchEntryPoint(
420+
renderedContent: renderedContent,
421+
);
422+
}
423+
302424
Content _parseGoogleAIContent(Object jsonObject) {
303425
return switch (jsonObject) {
304426
{'parts': final List<Object?> parts} => Content(

packages/firebase_ai/firebase_ai/test/api_test.dart

Lines changed: 51 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,9 +14,10 @@
1414

1515
import 'dart:convert';
1616

17-
import 'package:firebase_ai/firebase_ai.dart';
1817
import 'package:firebase_ai/src/api.dart';
19-
18+
import 'package:firebase_ai/src/content.dart';
19+
import 'package:firebase_ai/src/error.dart';
20+
import 'package:firebase_ai/src/schema.dart';
2021
import 'package:flutter_test/flutter_test.dart';
2122

2223
// --- Mock/Helper Implementations ---
@@ -947,6 +948,54 @@ void main() {
947948
throwsA(isA<FirebaseAISdkException>().having(
948949
(e) => e.message, 'message', contains('WebGroundingChunk'))));
949950
});
951+
952+
test(
953+
'parses groundingSupport and filters out entries without a segment',
954+
() {
955+
final jsonResponse = {
956+
'candidates': [
957+
{
958+
'content': {
959+
'parts': [
960+
{'text': 'Test'}
961+
]
962+
},
963+
'finishReason': 'STOP',
964+
'groundingMetadata': {
965+
'groundingSupport': [
966+
// Valid entry
967+
{
968+
'segment': {
969+
'startIndex': 0,
970+
'endIndex': 4,
971+
'text': 'Test'
972+
},
973+
'groundingChunkIndices': [0]
974+
},
975+
// Invalid entry - missing segment
976+
{
977+
'groundingChunkIndices': [1]
978+
},
979+
// Invalid entry - empty object
980+
{}
981+
]
982+
}
983+
}
984+
]
985+
};
986+
987+
final response =
988+
VertexSerialization().parseGenerateContentResponse(jsonResponse);
989+
final groundingMetadata = response.candidates.first.groundingMetadata;
990+
991+
expect(groundingMetadata, isNotNull);
992+
// The invalid entries should be filtered out.
993+
expect(groundingMetadata!.groundingSupport, hasLength(1));
994+
995+
final validSupport = groundingMetadata.groundingSupport.first;
996+
expect(validSupport.segment.text, 'Test');
997+
expect(validSupport.groundingChunkIndices, [0]);
998+
});
950999
});
9511000

9521001
test('parses JSON with no candidates (empty list)', () {

0 commit comments

Comments
 (0)