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

This should resolve #6990

---------

Co-authored-by: Rodrigo Lazo Paz <[email protected]>

Author: thatfiredev

firebase-ai/CHANGELOG.md

@@ -1,8 +1,9 @@
 # Unreleased
+* [feature] Added support for configuring the "thinking" budget when using Gemini
+  2.5 series models. (#6990)
 * [feature] **Breaking Change**: Add support for Grounding with Google Search (#7042).
     * **Action Required:** Update all references of `groundingAttributions`, `webSearchQueries`, `retrievalQueries` in `GroundingMetadata` to be non-optional.
 
-
 # 16.2.0
 * [changed] Deprecate the `totalBillableCharacters` field (only usable with pre-2.0 models). (#7042)
 * [feature] Added support for extra schema properties like `title`, `minItems`, `maxItems`, `minimum`
@@ -33,4 +34,3 @@
 
  Note: This feature is in Public Preview, which means that it is not subject to any SLA or
  deprecation policy and could change in backwards-incompatible ways.
-

firebase-ai/api.txt

@@ -364,6 +364,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;
@@ -375,6 +376,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;
   }
@@ -997,6 +999,19 @@ package com.google.firebase.ai.type {
     property public final String text;
   }
 
+  public final class ThinkingConfig {
+  }
+
+  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(int 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);
     method public static com.google.firebase.ai.type.Tool googleSearch(com.google.firebase.ai.type.GoogleSearch googleSearch = com.google.firebase.ai.type.GoogleSearch());
@@ -1019,16 +1034,18 @@ package com.google.firebase.ai.type {
   }
 
   public final class UsageMetadata {
-    ctor public UsageMetadata(int promptTokenCount, Integer? candidatesTokenCount, int totalTokenCount, java.util.List<com.google.firebase.ai.type.ModalityTokenCount> promptTokensDetails, java.util.List<com.google.firebase.ai.type.ModalityTokenCount> candidatesTokensDetails);
+    ctor public UsageMetadata(int promptTokenCount, Integer? candidatesTokenCount, int totalTokenCount, java.util.List<com.google.firebase.ai.type.ModalityTokenCount> promptTokensDetails, java.util.List<com.google.firebase.ai.type.ModalityTokenCount> candidatesTokensDetails, int thoughtsTokenCount);
     method public Integer? getCandidatesTokenCount();
     method public java.util.List<com.google.firebase.ai.type.ModalityTokenCount> getCandidatesTokensDetails();
     method public int getPromptTokenCount();
     method public java.util.List<com.google.firebase.ai.type.ModalityTokenCount> getPromptTokensDetails();
+    method public int getThoughtsTokenCount();
     method public int getTotalTokenCount();
     property public final Integer? candidatesTokenCount;
     property public final java.util.List<com.google.firebase.ai.type.ModalityTokenCount> candidatesTokensDetails;
     property public final int promptTokenCount;
     property public final java.util.List<com.google.firebase.ai.type.ModalityTokenCount> promptTokensDetails;
+    property public final int thoughtsTokenCount;
     property public final int totalTokenCount;
   }
 

firebase-ai/gradle.properties

@@ -12,5 +12,5 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-version=16.3.0
+version=17.0.0
 latestReleasedVersion=16.2.0

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,64 @@
+/*
+ * 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
+private constructor(
+  internal val thinkingBudget: Int? = null,
+) {
+
+  public class Builder() {
+    @JvmField
+    @set:JvmSynthetic // hide void setter from Java
+    public var thinkingBudget: Int? = null
+
+    /**
+     * Indicates the thinking budget in tokens. 0 is DISABLED. -1 is AUTOMATIC. The default values
+     * and allowed ranges are model dependent.
+     */
+    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/main/kotlin/com/google/firebase/ai/type/UsageMetadata.kt

@@ -28,13 +28,15 @@ import kotlinx.serialization.Serializable
  * prompt.
  * @param candidatesTokensDetails The breakdown, by modality, of how many tokens are consumed by the
  * candidates.
+ * @param thoughtsTokenCount The number of tokens used by the model's internal "thinking" process.
  */
 public class UsageMetadata(
   public val promptTokenCount: Int,
   public val candidatesTokenCount: Int?,
   public val totalTokenCount: Int,
   public val promptTokensDetails: List<ModalityTokenCount>,
   public val candidatesTokensDetails: List<ModalityTokenCount>,
+  public val thoughtsTokenCount: Int,
 ) {
 
   @Serializable
@@ -44,6 +46,7 @@ public class UsageMetadata(
     val totalTokenCount: Int? = null,
     val promptTokensDetails: List<ModalityTokenCount.Internal>? = null,
     val candidatesTokensDetails: List<ModalityTokenCount.Internal>? = null,
+    val thoughtsTokenCount: Int? = null,
   ) {
 
     internal fun toPublic(): UsageMetadata =
@@ -52,7 +55,8 @@ public class UsageMetadata(
         candidatesTokenCount ?: 0,
         totalTokenCount ?: 0,
         promptTokensDetails = promptTokensDetails?.map { it.toPublic() } ?: emptyList(),
-        candidatesTokensDetails = candidatesTokensDetails?.map { it.toPublic() } ?: emptyList()
+        candidatesTokensDetails = candidatesTokensDetails?.map { it.toPublic() } ?: emptyList(),
+        thoughtsTokenCount ?: 0
       )
   }
 }

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)
+  }
+}