[FirebaseAI] Add support for Grounding with Google Search

firebase-ai/src/main/kotlin/com/google/firebase/ai/type/Candidate.kt

@@ -22,7 +22,6 @@ import com.google.firebase.ai.common.util.FirstOrdinalSerializer
 import java.util.Calendar
 import kotlinx.serialization.ExperimentalSerializationApi
 import kotlinx.serialization.KSerializer
-import kotlinx.serialization.SerialName
 import kotlinx.serialization.Serializable
 import kotlinx.serialization.json.JsonNames
 
@@ -39,7 +38,8 @@ internal constructor(
   public val content: Content,
   public val safetyRatings: List<SafetyRating>,
   public val citationMetadata: CitationMetadata?,
-  public val finishReason: FinishReason?
+  public val finishReason: FinishReason?,
+  public val groundingMetadata: GroundingMetadata?
 ) {
 
   @Serializable
@@ -48,48 +48,22 @@ internal constructor(
     val finishReason: FinishReason.Internal? = null,
     val safetyRatings: List<SafetyRating.Internal>? = null,
     val citationMetadata: CitationMetadata.Internal? = null,
-    val groundingMetadata: GroundingMetadata? = null,
+    val groundingMetadata: GroundingMetadata.Internal? = null
   ) {
     internal fun toPublic(): Candidate {
       val safetyRatings = safetyRatings?.mapNotNull { it.toPublic() }.orEmpty()
       val citations = citationMetadata?.toPublic()
       val finishReason = finishReason?.toPublic()
+      val groundingMetadata = groundingMetadata?.toPublic()
 
       return Candidate(
         this.content?.toPublic() ?: content("model") {},
         safetyRatings,
         citations,
-        finishReason
+        finishReason,
+        groundingMetadata
       )
     }
-
-    @Serializable
-    internal data class GroundingMetadata(
-      @SerialName("web_search_queries") val webSearchQueries: List<String>?,
-      @SerialName("search_entry_point") val searchEntryPoint: SearchEntryPoint?,
-      @SerialName("retrieval_queries") val retrievalQueries: List<String>?,
-      @SerialName("grounding_attribution") val groundingAttribution: List<GroundingAttribution>?,
-    ) {
-
-      @Serializable
-      internal data class SearchEntryPoint(
-        @SerialName("rendered_content") val renderedContent: String?,
-        @SerialName("sdk_blob") val sdkBlob: String?,
-      )
-
-      @Serializable
-      internal data class GroundingAttribution(
-        val segment: Segment,
-        @SerialName("confidence_score") val confidenceScore: Float?,
-      ) {
-
-        @Serializable
-        internal data class Segment(
-          @SerialName("start_index") val startIndex: Int,
-          @SerialName("end_index") val endIndex: Int,
-        )
-      }
-    }
   }
 }
 
@@ -252,7 +226,7 @@ public class FinishReason private constructor(public val name: String, public va
   @Serializable(Internal.Serializer::class)
   internal enum class Internal {
     UNKNOWN,
-    @SerialName("FINISH_REASON_UNSPECIFIED") UNSPECIFIED,
+    UNSPECIFIED,
     STOP,
     MAX_TOKENS,
     SAFETY,
@@ -317,3 +291,203 @@ public class FinishReason private constructor(public val name: String, public va
     public val MALFORMED_FUNCTION_CALL: FinishReason = FinishReason("MALFORMED_FUNCTION_CALL", 9)
   }
 }
+
+/**
+ * Metadata returned to the client when grounding is enabled.
+ *
+ * If using Grounding with Google Search, you are required to comply with the
+ * [Service Specific Terms](https://cloud.google.com/terms/service-terms) for "Grounding with Google
+ * Search".
+ *
+ * @property webSearchQueries A list of web search queries that the model performed to gather the
+ * grounding information. These can be used to allow users to explore the search results themselves.
+ * @property searchEntryPoint Google search entry point for web searches. This contains an HTML/CSS
+ * snippet that **must** be embedded in an app to display a Google Search Entry point for follow-up
+ * web searches related to the model's "Grounded Response". To ensure proper rendering, it's
+ * recommended to display this content within a `WebView`.
+ * @property groundingChunks A list of [GroundingChunk] objects. Each chunk represents a piece of
+ * retrieved content (e.g. from a web page) that the model used to ground its response.
+ * @property groundingSupports A list of [GroundingSupport] objects. Each object details how
+ * specific segments of the model's response are supported by the `groundingChunks`.
+ */
+public class GroundingMetadata(
+  public val webSearchQueries: List<String>?,
+  public val searchEntryPoint: SearchEntryPoint?,
+  public val retrievalQueries: List<String>?,
+  @Deprecated("Use groundingChunks instead")
+  public val groundingAttribution: List<GroundingAttribution>?,
+  public val groundingChunks: List<GroundingChunk>,
+  public val groundingSupports: List<GroundingSupport>,
+) {
+  @Serializable
+  internal data class Internal(
+    val webSearchQueries: List<String>?,
+    val searchEntryPoint: SearchEntryPoint.Internal?,
+    val retrievalQueries: List<String>?,
+    @Deprecated("Use groundingChunks instead")
+    val groundingAttribution: List<GroundingAttribution.Internal>?,
+    val groundingChunks: List<GroundingChunk.Internal>?,
+    val groundingSupports: List<GroundingSupport.Internal>?,
+  ) {
+    internal fun toPublic() =
+      GroundingMetadata(
+        webSearchQueries = webSearchQueries,
+        searchEntryPoint = searchEntryPoint?.toPublic(),
+        retrievalQueries = retrievalQueries,
+        groundingAttribution = groundingAttribution?.map { it.toPublic() },
+        groundingChunks = groundingChunks?.map { it.toPublic() }.orEmpty(),
+        groundingSupports = groundingSupports?.map { it.toPublic() }.orEmpty()
+      )
+  }
+}
+
+/**
+ * A class representing the Google Search entry point.
+ *
+ * When using this feature, you are required to comply with the
+ * [Service Specific Terms](https://cloud.google.com/terms/service-terms) for "Grounding with Google
+ * Search".
+ *
+ * @property renderedContent An HTML/CSS snippet that can be embedded in your app. The snippet is
+ * designed to avoid undesired interaction with the rest of the page's CSS. To ensure proper
+ * rendering, it's recommended to display this content within a `WebView`.
+ * @property sdkBlob A blob of data for the client SDK to render the search entry point.
+ */
+public class SearchEntryPoint(
+  public val renderedContent: String?,
+  public val sdkBlob: String?,
+) {
+  @Serializable
+  internal data class Internal(
+    val renderedContent: String?,
+    val sdkBlob: String?,
+  ) {
+    internal fun toPublic() = SearchEntryPoint(renderedContent = renderedContent, sdkBlob = sdkBlob)
+  }
+}
+
+/**
+ * Represents a chunk of retrieved data that supports a claim in the model's response. This is part
+ * of the grounding information provided when grounding is enabled.
+ *
+ * When using this feature, you are required to comply with the
+ * [Service Specific Terms](https://cloud.google.com/terms/service-terms) for "Grounding with Google
+ * Search".
+ *
+ * @property web Contains details if the grounding chunk is from a web source.
+ */
+public class GroundingChunk(
+  public val web: WebGroundingChunk?,
+) {
+  @Serializable
+  internal data class Internal(
+    val web: WebGroundingChunk.Internal?,
+  ) {
+    internal fun toPublic() = GroundingChunk(web = web?.toPublic())
+  }
+}
+
+/**
+ * A grounding chunk from the web.
+ *
+ * When using this feature, you are required to comply with the
+ * [Service Specific Terms](https://cloud.google.com/terms/service-terms) for "Grounding with Google
+ * Search".
+ *
+ * @property uri The URI of the retrieved web page.
+ * @property title The title of the retrieved web page.
+ * @property domain The domain of the original URI from which the content was retrieved (e.g.,
+ * `example.com`). This is only populated when using the Vertex AI Gemini API.
+ */
+public class WebGroundingChunk(
+  public val uri: String?,
+  public val title: String?,
+  public val domain: String?
+) {
+  @Serializable
+  internal data class Internal(val uri: String?, val title: String?, val domain: String?) {
+    internal fun toPublic() = WebGroundingChunk(uri = uri, title = title, domain = domain)
+  }
+}
+
+/**
+ * Provides information about how a specific segment of the model's response is supported by the
+ * retrieved grounding chunks.
+ *
+ * When using this feature, you are required to comply with the
+ * [Service Specific Terms](https://cloud.google.com/terms/service-terms) for "Grounding with Google
+ * Search".
+ *
+ * @property segment Specifies the segment of the model's response content that this grounding
+ * support pertains to.
+ * @property groundingChunkIndices A list of indices that refer to specific [GroundingChunk] objects
+ * within the [GroundingMetadata.groundingChunks] array. These referenced chunks are the sources
+ * that support the claim made in the associated `segment` of the response.
+ */
+public class GroundingSupport(
+  public val segment: Segment?,
+  public val groundingChunkIndices: List<Int>,
+) {
+  @Serializable
+  internal data class Internal(
+    val segment: Segment.Internal?,
+    val groundingChunkIndices: List<Int>?,
+  ) {
+    internal fun toPublic() =
+      GroundingSupport(
+        segment = segment?.toPublic(),
+        groundingChunkIndices = groundingChunkIndices.orEmpty(),
+      )
+  }
+}
+
+@Deprecated("Use GroundingChunk instead")
+public class GroundingAttribution(
+  public val segment: Segment,
+  public val confidenceScore: Float?,
+) {
+  @Deprecated("Use GroundingChunk instead")
+  @Serializable
+  internal data class Internal(
+    val segment: Segment.Internal,
+    val confidenceScore: Float?,
+  ) {
+    internal fun toPublic() =
+      GroundingAttribution(segment = segment.toPublic(), confidenceScore = confidenceScore)
+  }
+}
+
+/**
+ * Represents a specific segment within a [Content] object, often used to pinpoint the exact
+ * location of text or data that grounding information refers to.
+ *
+ * @property startIndex The zero-based start index of the segment within the specified [Part],
+ * measured in UTF-8 bytes. This offset is inclusive.
+ * @property endIndex The zero-based end index of the segment within the specified [Part], measured
+ * in UTF-8 bytes. This offset is exclusive.
+ * @property partIndex The zero-based index of the [Part] object within the `parts` array of its
+ * parent [Content] object. This identifies which part of the content the segment belongs to.
+ * @property text The text content of the segment.
+ */
+public class Segment(
+  public val startIndex: Int,
+  public val endIndex: Int,
+  public val partIndex: Int,
+  public val text: String,
+) {
+  @Serializable
+  internal data class Internal(
+    val startIndex: Int?,
+    val endIndex: Int?,
+    val partIndex: Int?,
+    val text: String?,
+  ) {
+    internal fun toPublic() =
+      Segment(
+        startIndex = startIndex ?: 0,
+        endIndex = endIndex ?: 0,
+        partIndex = partIndex ?: 0,
+        text = text ?: ""
+      )
+  }
+}

firebase-ai/src/main/kotlin/com/google/firebase/ai/type/GoogleSearch.kt

@@ -0,0 +1,13 @@
+package com.google.firebase.ai.type
+
+import kotlinx.serialization.Serializable
+
+/**
+ * A tool that allows the generative model to connect to Google Search to access and incorporate
+ * up-to-date information from the web into its responses.
+ *
+ * When this tool is used, the model's responses may include "Grounded Results" which are subject to
+ * the Grounding with Google Search terms outlined in the
+ * [Service Specific Terms](https://cloud.google.com/terms/service-terms).
+ */
+@Serializable public class GoogleSearch {}

firebase-ai/src/main/kotlin/com/google/firebase/ai/type/Tool.kt

@@ -26,11 +26,19 @@ import kotlinx.serialization.json.JsonObject
  * @param functionDeclarations The set of functions that this tool allows the model access to
  */
 public class Tool
-internal constructor(internal val functionDeclarations: List<FunctionDeclaration>?) {
-  internal fun toInternal() = Internal(functionDeclarations?.map { it.toInternal() } ?: emptyList())
+internal constructor(
+  internal val functionDeclarations: List<FunctionDeclaration>?,
+  internal val googleSearch: GoogleSearch?
+) {
+  internal fun toInternal() =
+    Internal(
+      functionDeclarations?.map { it.toInternal() } ?: emptyList(),
+      googleSearch = this.googleSearch
+    )
   @Serializable
   internal data class Internal(
     val functionDeclarations: List<FunctionDeclaration.Internal>? = null,
+    val googleSearch: GoogleSearch? = null,
     // This is a json object because it is not possible to make a data class with no parameters.
     val codeExecution: JsonObject? = null,
   )
@@ -43,7 +51,26 @@ internal constructor(internal val functionDeclarations: List<FunctionDeclaration
      */
     @JvmStatic
     public fun functionDeclarations(functionDeclarations: List<FunctionDeclaration>): Tool {
-      return Tool(functionDeclarations)
+      return Tool(functionDeclarations, null)
+    }
+
+    /**
+     * Creates a [Tool] instance that enables the model to use Grounding with Google Search.
+     *
+     * This allows the model to connect to Google Search to access and incorporate up-to-date
+     * information from the web into its responses.
+     *
+     * When this tool is used, the model's responses may include "Grounded Results" which are
+     * subject to the Grounding with Google Search terms outlined in the
+     * [Service Specific Terms](https://cloud.google.com/terms/service-terms).
+     *
+     * @param googleSearch An empty [GoogleSearch] object. The presence of this object in the list
+     * of tools enables the model to use Google Search.
+     * @return a [Tool] configured for Google Search.
+     */
+    @JvmStatic
+    public fun googleSearch(googleSearch: GoogleSearch = GoogleSearch()): Tool {
+      return Tool(null, googleSearch)
     }
   }
 }

firebase-ai/src/test/java/com/google/firebase/ai/DevAPIUnarySnapshotTests.kt

@@ -22,6 +22,9 @@ import com.google.firebase.ai.type.ResponseStoppedException
 import com.google.firebase.ai.type.ServerException
 import com.google.firebase.ai.util.goldenDevAPIUnaryFile
 import io.kotest.assertions.throwables.shouldThrow
+import io.kotest.matchers.collections.shouldNotBeEmpty
+import io.kotest.matchers.nulls.shouldBeNull
+import io.kotest.matchers.nulls.shouldNotBeNull
 import io.kotest.matchers.should
 import io.kotest.matchers.shouldBe
 import io.kotest.matchers.shouldNotBe
@@ -95,4 +98,22 @@ internal class DevAPIUnarySnapshotTests {
     goldenDevAPIUnaryFile("unary-failure-unknown-model.json", HttpStatusCode.NotFound) {
       withTimeout(testTimeout) { shouldThrow<ServerException> { model.generateContent("prompt") } }
     }
+
+  // This test case can be removed once https://b.corp.google.com/issues/422779395 (internal) is
+  // fixed.
+  @Test
+  fun `google search grounding empty grounding chunks`() =
+    goldenDevAPIUnaryFile("unary-success-google-search-grounding-empty-grounding-chunks.json") {
+      withTimeout(testTimeout) {
+        val response = model.generateContent("prompt")
+
+        response.candidates.shouldNotBeEmpty()
+        val candidate = response.candidates.first()
+        val groundingMetadata = candidate.groundingMetadata
+        groundingMetadata.shouldNotBeNull()
+
+        groundingMetadata.groundingChunks.shouldNotBeEmpty()
+        groundingMetadata.groundingChunks.forEach { it.web.shouldBeNull() }
+      }
+    }
 }