diff --git a/firebase-ai/src/main/kotlin/com/google/firebase/ai/type/GenerationConfig.kt b/firebase-ai/src/main/kotlin/com/google/firebase/ai/type/GenerationConfig.kt index 1c2d2680bb1..9acdc534cfb 100644 --- a/firebase-ai/src/main/kotlin/com/google/firebase/ai/type/GenerationConfig.kt +++ b/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?, + 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? = 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?): 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? = null + @SerialName("response_modalities") val responseModalities: List? = null, + @SerialName("thinking_config") val thinkingConfig: ThinkingConfig.Internal? = null ) public companion object { diff --git a/firebase-ai/src/main/kotlin/com/google/firebase/ai/type/ThinkingConfig.kt b/firebase-ai/src/main/kotlin/com/google/firebase/ai/type/ThinkingConfig.kt new file mode 100644 index 00000000000..9ccfde3797d --- /dev/null +++ b/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, + ) +} diff --git a/firebase-ai/src/main/kotlin/com/google/firebase/ai/type/Tool.kt b/firebase-ai/src/main/kotlin/com/google/firebase/ai/type/Tool.kt index 83391166bd4..7b69b3f42e7 100644 --- a/firebase-ai/src/main/kotlin/com/google/firebase/ai/type/Tool.kt +++ b/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?) { - internal fun toInternal() = Internal(functionDeclarations?.map { it.toInternal() } ?: emptyList()) +internal constructor( + internal val functionDeclarations: List? = 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? = null, + @SerialName("function_declarations") val functionDeclarations: List? = 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): 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) } } }