search grounding WIP

Author: dlarocque

packages/firebase_vertexai/firebase_vertexai/lib/firebase_vertexai.dart

@@ -53,11 +53,12 @@ export 'src/error.dart'
         ServerException,
         UnsupportedUserLocation;
 export 'src/firebase_vertexai.dart' show FirebaseVertexAI;
-export 'src/function_calling.dart'
+export 'src/tool.dart'
     show
         FunctionCallingConfig,
         FunctionCallingMode,
         FunctionDeclaration,
+        GoogleSearch,
         Tool,
         ToolConfig;
 export 'src/imagen_api.dart'

packages/firebase_vertexai/firebase_vertexai/lib/src/api.dart

@@ -169,7 +169,8 @@ final class Candidate {
   // TODO: token count?
   // ignore: public_member_api_docs
   Candidate(this.content, this.safetyRatings, this.citationMetadata,
-      this.finishReason, this.finishMessage);
+      this.finishReason, this.finishMessage,
+      {this.groundingMetadata});
 
   /// Generated content returned from the model.
   final Content content;
@@ -194,6 +195,8 @@ final class Candidate {
   /// Message for finish reason.
   final String? finishMessage;
 
+  final GroundingMetadata? groundingMetadata;
+
   /// The concatenation of the text parts of [content], if any.
   ///
   /// If this candidate was finished for a reason of [FinishReason.recitation]
@@ -225,6 +228,59 @@ final class Candidate {
   }
 }
 
+final class Segment {
+  Segment(
+      {required this.partIndex,
+      required this.startIndex,
+      required this.endIndex,
+      required this.text});
+
+  final int partIndex;
+  final int startIndex;
+  final int endIndex;
+  final String text;
+}
+
+final class WebGroundingChunk {
+  WebGroundingChunk({this.uri, this.title, this.domain});
+
+  final String? uri;
+  final String? title;
+  final String? domain;
+}
+
+final class GroundingChunk {
+  GroundingChunk({this.web});
+
+  final WebGroundingChunk? web;
+}
+
+final class GroundingSupport {
+  GroundingSupport({this.segment, required this.groundingChunkIndices});
+
+  final Segment? segment;
+  final List<int> groundingChunkIndices;
+}
+
+final class SearchEntrypoint {
+  SearchEntrypoint({this.renderedContent});
+
+  final String? renderedContent;
+}
+
+final class GroundingMetadata {
+  GroundingMetadata(
+      {this.searchEntrypoint,
+      required this.groundingChunks,
+      required this.groundingSupport,
+      required this.webSearchQueries});
+
+  final SearchEntrypoint? searchEntrypoint;
+  final List<GroundingChunk> groundingChunks;
+  final List<GroundingSupport> groundingSupport;
+  final List<String> webSearchQueries;
+}
+
 /// Safety rating for a piece of content.
 ///
 /// The safety rating contains the category of harm and the harm probability
@@ -938,29 +994,33 @@ Candidate _parseCandidate(Object? jsonObject) {
   }
 
   return Candidate(
-    jsonObject.containsKey('content')
-        ? parseContent(jsonObject['content'] as Object)
-        : Content(null, []),
-    switch (jsonObject) {
-      {'safetyRatings': final List<Object?> safetyRatings} =>
-        safetyRatings.map(_parseSafetyRating).toList(),
-      _ => null
-    },
-    switch (jsonObject) {
-      {'citationMetadata': final Object citationMetadata} =>
-        _parseCitationMetadata(citationMetadata),
-      _ => null
-    },
-    switch (jsonObject) {
-      {'finishReason': final Object finishReason} =>
-        FinishReason._parseValue(finishReason),
-      _ => null
-    },
-    switch (jsonObject) {
-      {'finishMessage': final String finishMessage} => finishMessage,
-      _ => null
-    },
-  );
+      jsonObject.containsKey('content')
+          ? parseContent(jsonObject['content'] as Object)
+          : Content(null, []),
+      switch (jsonObject) {
+        {'safetyRatings': final List<Object?> safetyRatings} =>
+          safetyRatings.map(_parseSafetyRating).toList(),
+        _ => null
+      },
+      switch (jsonObject) {
+        {'citationMetadata': final Object citationMetadata} =>
+          _parseCitationMetadata(citationMetadata),
+        _ => null
+      },
+      switch (jsonObject) {
+        {'finishReason': final Object finishReason} =>
+          FinishReason._parseValue(finishReason),
+        _ => null
+      },
+      switch (jsonObject) {
+        {'finishMessage': final String finishMessage} => finishMessage,
+        _ => null
+      },
+      groundingMetadata: switch (jsonObject) {
+        {'groundingMetadata': final Object groundingMetadata} =>
+          _parseGroundingMetadata(groundingMetadata),
+        _ => null
+      });
 }
 
 PromptFeedback _parsePromptFeedback(Object jsonObject) {
@@ -1074,3 +1134,97 @@ Citation _parseCitationSource(Object? jsonObject) {
     jsonObject['license'] as String?,
   );
 }
+
+GroundingMetadata _parseGroundingMetadata(Object? jsonObject) {
+  if (jsonObject is! Map) {
+    throw unhandledFormat('GroundingMetadata', jsonObject);
+  }
+
+  final searchEntrypoint = switch (jsonObject) {
+    {'searchEntrypoint': final Object? searchEntrypoint} =>
+      _parseSearchEntrypoint(searchEntrypoint),
+    _ => null,
+  };
+  final groundingChunks = switch (jsonObject) {
+        {'groundingChunks': final List<Object?> groundingChunks} =>
+          groundingChunks.map(_parseGroundingChunk).toList(),
+        _ => null,
+      } ??
+      [];
+  final groundingSupport = switch (jsonObject) {
+        {'groundingSupport': final List<Object?> groundingSupport} =>
+          groundingSupport.map(_parseGroundingSupport).toList(),
+        _ => null,
+      } ??
+      [];
+  final webSearchQueries = switch (jsonObject) {
+        {'webSearchQueries': final List<String>? webSearchQueries} =>
+          webSearchQueries,
+        _ => null,
+      } ??
+      [];
+
+  return GroundingMetadata(
+      searchEntrypoint: searchEntrypoint,
+      groundingChunks: groundingChunks,
+      groundingSupport: groundingSupport,
+      webSearchQueries: webSearchQueries);
+}
+
+Segment _parseSegment(Object? jsonObject) {
+  if (jsonObject is! Map) {
+    throw unhandledFormat('Segment', jsonObject);
+  }
+
+  return Segment(
+      partIndex: (jsonObject['partIndex'] as int?) ?? 0,
+      startIndex: (jsonObject['startIndex'] as int?) ?? 0,
+      endIndex: (jsonObject['endIndex'] as int?) ?? 0,
+      text: (jsonObject['text'] as String?) ?? '');
+}
+
+WebGroundingChunk _parseWebGroundingChunk(Object? jsonObject) {
+  if (jsonObject is! Map) {
+    throw unhandledFormat('WebGroundingChunk', jsonObject);
+  }
+
+  return WebGroundingChunk(
+    uri: jsonObject['uri'] as String?,
+    title: jsonObject['title'] as String?,
+    domain: jsonObject['domain'] as String?,
+  );
+}
+
+GroundingChunk _parseGroundingChunk(Object? jsonObject) {
+  if (jsonObject is! Map) {
+    throw unhandledFormat('GroundingChunk', jsonObject);
+  }
+
+  return GroundingChunk(
+    web: jsonObject['web'] != null
+        ? _parseWebGroundingChunk(jsonObject['web'])
+        : null,
+  );
+}
+
+GroundingSupport _parseGroundingSupport(Object? jsonObject) {
+  if (jsonObject is! Map) {
+    throw unhandledFormat('GroundingSupport', jsonObject);
+  }
+
+  return GroundingSupport(
+      segment: jsonObject['segment'] != null
+          ? _parseSegment(jsonObject['segment'])
+          : null,
+      groundingChunkIndices:
+          (jsonObject['groundingChunkIndices'] as List<int>?) ?? []);
+}
+
+SearchEntrypoint _parseSearchEntrypoint(Object? jsonObject) {
+  if (jsonObject is! Map) {
+    throw unhandledFormat('SearchEntrypoint', jsonObject);
+  }
+
+  return SearchEntrypoint(
+      renderedContent: (jsonObject['renderedContent'] as String?) ?? '');
+}

packages/firebase_vertexai/firebase_vertexai/lib/src/base_model.dart

@@ -26,7 +26,7 @@ import 'package:web_socket_channel/io.dart';
 import 'api.dart';
 import 'client.dart';
 import 'content.dart';
-import 'function_calling.dart';
+import 'tool.dart';
 import 'imagen_api.dart';
 import 'imagen_content.dart';
 import 'live_api.dart';

packages/firebase_vertexai/firebase_vertexai/lib/src/function_calling.dart renamed to packages/firebase_vertexai/firebase_vertexai/lib/src/tool.dart

@@ -21,12 +21,16 @@ import 'schema.dart';
 /// knowledge and scope of the model.
 final class Tool {
   // ignore: public_member_api_docs
-  Tool._(this._functionDeclarations);
+  Tool._(this._functionDeclarations, this._googleSearch);
 
   /// Returns a [Tool] instance with list of [FunctionDeclaration].
   static Tool functionDeclarations(
       List<FunctionDeclaration> functionDeclarations) {
-    return Tool._(functionDeclarations);
+    return Tool._(functionDeclarations, null);
+  }
+
+  static Tool googleSearch({GoogleSearch googleSearch = const GoogleSearch()}) {
+    return Tool._(null, googleSearch);
   }
 
   /// A list of `FunctionDeclarations` available to the model that can be used
@@ -38,15 +42,25 @@ final class Tool {
   /// [FunctionResponse]
   /// with the role "function" generation context for the next model turn.
   final List<FunctionDeclaration>? _functionDeclarations;
+  final GoogleSearch? _googleSearch;
 
   /// Convert to json object.
   Map<String, Object> toJson() => {
         if (_functionDeclarations case final _functionDeclarations?)
           'functionDeclarations':
               _functionDeclarations.map((f) => f.toJson()).toList(),
+        if (_googleSearch case final _googleSearch?)
+          'googleSearch': _googleSearch.toJson()
       };
 }
 
+final class GoogleSearch {
+  const GoogleSearch();
+
+  // Convert to json object.
+  Map<String, Object> toJson() => {};
+}
+
 /// Structured representation of a function declaration as defined by the
 /// [OpenAPI 3.03 specification](https://spec.openapis.org/oas/v3.0.3).
 ///