Add Google Search and Thinking support to firebase-ai

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

@@ -91,6 +91,7 @@ private constructor(
   internal val responseMimeType: String?,
   internal val responseSchema: Schema?,
   internal val responseModalities: List<ResponseModality>?,
+  internal val thinkingConfig: ThinkingConfig?,
 ) {
 
   /**
@@ -121,6 +122,8 @@ private constructor(
    *
    * @property responseModalities See [GenerationConfig.responseModalities].
    *
+   * @property thinkingConfig See [GenerationConfig.thinkingConfig].
+   *
    * @see [generationConfig]
    */
   public class Builder {
@@ -135,6 +138,7 @@ private constructor(
     @JvmField public var responseMimeType: String? = null
     @JvmField public var responseSchema: Schema? = null
     @JvmField public var responseModalities: List<ResponseModality>? = null
+    @JvmField public var thinkingConfig: ThinkingConfig? = null
 
     public fun setTemperature(temperature: Float?): Builder = apply {
       this.temperature = temperature
@@ -165,6 +169,9 @@ private constructor(
     public fun setResponseModalities(responseModalities: List<ResponseModality>?): Builder = apply {
       this.responseModalities = responseModalities
     }
+    public fun setThinkingConfig(thinkingConfig: ThinkingConfig?): Builder = apply {
+      this.thinkingConfig = thinkingConfig
+    }
 
     /** Create a new [GenerationConfig] with the attached arguments. */
     public fun build(): GenerationConfig =
@@ -179,7 +186,8 @@ private constructor(
         frequencyPenalty = frequencyPenalty,
         responseMimeType = responseMimeType,
         responseSchema = responseSchema,
-        responseModalities = responseModalities
+        responseModalities = responseModalities,
+        thinkingConfig = thinkingConfig
       )
   }
 
@@ -195,7 +203,8 @@ private constructor(
       presencePenalty = presencePenalty,
       responseMimeType = responseMimeType,
       responseSchema = responseSchema?.toInternal(),
-      responseModalities = responseModalities?.map { it.toInternal() }
+      responseModalities = responseModalities?.map { it.toInternal() },
+      thinkingConfig = thinkingConfig?.toInternal()
     )
 
   @Serializable
@@ -210,7 +219,8 @@ private constructor(
     @SerialName("presence_penalty") val presencePenalty: Float? = null,
     @SerialName("frequency_penalty") val frequencyPenalty: Float? = null,
     @SerialName("response_schema") val responseSchema: Schema.Internal? = null,
-    @SerialName("response_modalities") val responseModalities: List<String>? = null
+    @SerialName("response_modalities") val responseModalities: List<String>? = null,
+    @SerialName("thinking_config") val thinkingConfig: ThinkingConfig.Internal? = null
   )
 
   public companion object {

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

@@ -0,0 +1,37 @@
+/*
+ * Copyright 2024 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.firebase.ai.type
+
+import kotlinx.serialization.SerialName
+import kotlinx.serialization.Serializable
+
+/**
+ * Configuration for the "Thinking" feature, which allows the model to process information more
+ * deeply before responding.
+ *
+ * @property includeThoughts Whether to include the model's thoughts in the response.
+ */
+public class ThinkingConfig(
+  public val includeThoughts: Boolean = false,
+) {
+  internal fun toInternal() = Internal(includeThoughts = includeThoughts)
+
+  @Serializable
+  internal data class Internal(
+    @SerialName("include_thoughts") val includeThoughts: Boolean,
+  )
+}

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

@@ -16,6 +16,7 @@
 
 package com.google.firebase.ai.type
 
+import kotlinx.serialization.SerialName
 import kotlinx.serialization.Serializable
 import kotlinx.serialization.json.JsonObject
 
@@ -26,13 +27,22 @@ 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>? = null,
+  internal val googleSearchRetrieval: Boolean = false
+) {
+  internal fun toInternal() =
+    Internal(
+      functionDeclarations = functionDeclarations?.map { it.toInternal() },
+      googleSearchRetrieval = if (googleSearchRetrieval) JsonObject(emptyMap()) else null
+    )
+
   @Serializable
   internal data class Internal(
-    val functionDeclarations: List<FunctionDeclaration.Internal>? = null,
+    @SerialName("function_declarations") val functionDeclarations: List<FunctionDeclaration.Internal>? = null,
     // This is a json object because it is not possible to make a data class with no parameters.
-    val codeExecution: JsonObject? = null,
+    @SerialName("code_execution") val codeExecution: JsonObject? = null,
+    @SerialName("google_search_retrieval") val googleSearchRetrieval: JsonObject? = null,
   )
   public companion object {
 
@@ -43,7 +53,13 @@ internal constructor(internal val functionDeclarations: List<FunctionDeclaration
      */
     @JvmStatic
     public fun functionDeclarations(functionDeclarations: List<FunctionDeclaration>): Tool {
-      return Tool(functionDeclarations)
+      return Tool(functionDeclarations = functionDeclarations)
+    }
+
+    /** Creates a [Tool] instance that provides the model with access to Google Search. */
+    @JvmStatic
+    public fun googleSearchRetrieval(): Tool {
+      return Tool(googleSearchRetrieval = true)
     }
   }
 }