feat(ai): add support for setting a thinking budget

Author: thatfiredev

firebase-ai/CHANGELOG.md

@@ -1,5 +1,6 @@
 # Unreleased
 
+* [feat] Added support for setting thinking budget for the Gemini models that support it. (#6990)
 * [fixed] Fixed `FirebaseAI.getInstance` StackOverflowException (#6971)
 * [fixed] Fixed an issue that was causing the SDK to send empty `FunctionDeclaration` descriptions to the API. 
 

firebase-ai/api.txt

@@ -362,6 +362,7 @@ package com.google.firebase.ai.type {
     method public com.google.firebase.ai.type.GenerationConfig.Builder setResponseSchema(com.google.firebase.ai.type.Schema? responseSchema);
     method public com.google.firebase.ai.type.GenerationConfig.Builder setStopSequences(java.util.List<java.lang.String>? stopSequences);
     method public com.google.firebase.ai.type.GenerationConfig.Builder setTemperature(Float? temperature);
+    method public com.google.firebase.ai.type.GenerationConfig.Builder setThinkingConfig(com.google.firebase.ai.type.ThinkingConfig? thinkingConfig);
     method public com.google.firebase.ai.type.GenerationConfig.Builder setTopK(Integer? topK);
     method public com.google.firebase.ai.type.GenerationConfig.Builder setTopP(Float? topP);
     field public Integer? candidateCount;
@@ -373,6 +374,7 @@ package com.google.firebase.ai.type {
     field public com.google.firebase.ai.type.Schema? responseSchema;
     field public java.util.List<java.lang.String>? stopSequences;
     field public Float? temperature;
+    field public com.google.firebase.ai.type.ThinkingConfig? thinkingConfig;
     field public Integer? topK;
     field public Float? topP;
   }
@@ -881,6 +883,23 @@ package com.google.firebase.ai.type {
     property public final String text;
   }
 
+  public final class ThinkingConfig {
+    ctor public ThinkingConfig(Integer? thinkingBudget);
+    method public Integer? getThinkingBudget();
+    property public final Integer? thinkingBudget;
+  }
+
+  public static final class ThinkingConfig.Builder {
+    ctor public ThinkingConfig.Builder();
+    method public com.google.firebase.ai.type.ThinkingConfig build();
+    method public com.google.firebase.ai.type.ThinkingConfig.Builder setThinkingBudget(Integer? thinkingBudget);
+    field public Integer? thinkingBudget;
+  }
+
+  public final class ThinkingConfigKt {
+    method public static com.google.firebase.ai.type.ThinkingConfig thinkingConfig(kotlin.jvm.functions.Function1<? super com.google.firebase.ai.type.ThinkingConfig.Builder,kotlin.Unit> init);
+  }
+
   public final class Tool {
     method public static com.google.firebase.ai.type.Tool functionDeclarations(java.util.List<com.google.firebase.ai.type.FunctionDeclaration> functionDeclarations);
     field public static final com.google.firebase.ai.type.Tool.Companion Companion;

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?,
 ) {
 
   /**
@@ -135,6 +136,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 +167,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 +184,8 @@ private constructor(
         frequencyPenalty = frequencyPenalty,
         responseMimeType = responseMimeType,
         responseSchema = responseSchema,
-        responseModalities = responseModalities
+        responseModalities = responseModalities,
+        thinkingConfig = thinkingConfig
       )
   }
 
@@ -195,7 +201,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 +217,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,56 @@
+/*
+ * Copyright 2025 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 parameters for thinking features. */
+public class ThinkingConfig(public val thinkingBudget: Int?) {
+
+  public class Builder() {
+    @JvmField public var thinkingBudget: Int? = null
+
+    /** The number of thoughts tokens that the model should generate. */
+    public fun setThinkingBudget(thinkingBudget: Int?): Builder = apply {
+      this.thinkingBudget = thinkingBudget
+    }
+
+    public fun build(): ThinkingConfig = ThinkingConfig(thinkingBudget = thinkingBudget)
+  }
+
+  internal fun toInternal() = Internal(thinkingBudget)
+
+  @Serializable
+  internal data class Internal(@SerialName("thinking_budget") val thinkingBudget: Int?)
+}
+
+/**
+ * Helper method to construct a [ThinkingConfig] in a DSL-like manner.
+ *
+ * Example Usage:
+ * ```
+ * thinkingConfig {
+ *   thinkingBudget = 0 // disable thinking
+ * }
+ * ```
+ */
+public fun thinkingConfig(init: ThinkingConfig.Builder.() -> Unit): ThinkingConfig {
+  val builder = ThinkingConfig.Builder()
+  builder.init()
+  return builder.build()
+}

firebase-ai/src/test/java/com/google/firebase/ai/type/ThinkingConfigTest.kt

@@ -0,0 +1,50 @@
+/*
+ * Copyright 2025 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 io.kotest.assertions.json.shouldEqualJson
+import io.kotest.matchers.equals.shouldBeEqual
+import kotlinx.serialization.encodeToString
+import kotlinx.serialization.json.Json
+import org.junit.Test
+
+internal class ThinkingConfigTest {
+
+  @Test
+  fun `Basic ThinkingConfig`() {
+    val thinkingConfig = ThinkingConfig.Builder().setThinkingBudget(1024).build()
+
+    val expectedJson =
+      """
+      {
+          "thinking_budget": 1024
+      }
+      """
+        .trimIndent()
+
+    Json.encodeToString(thinkingConfig.toInternal()).shouldEqualJson(expectedJson)
+  }
+
+  @Test
+  fun `thinkingConfig DSL correctly delegates to ThinkingConfig#Builder`() {
+    val thinkingConfig = ThinkingConfig.Builder().setThinkingBudget(1024).build()
+
+    val thinkingConfigDsl = thinkingConfig { thinkingBudget = 1024 }
+
+    thinkingConfig.thinkingBudget?.shouldBeEqual(thinkingConfigDsl.thinkingBudget as Int)
+  }
+}