Skip to content

Commit d7151b5

Browse files
goyaljaiclaude
goyaljai
and
claude
committed
Address review comments: simplify grounding to boolean flag, remove unsupported API features
- Replace GoogleGroundingConfig sealed class with simple groundingEnabled: Boolean in GoogleParams - Remove LLMCapability.Grounding (grounding is provider-specific, not a model trait) - Remove timeRangeFilter/Interval and searchTypes — tested on generativelanguage.googleapis.com, neither changed model behavior (timeRangeFilter proved no-op via Iran/US war 2026 date test) - Replace JsonObject grounding fields with typed GoogleSearch/GoogleTool internal classes - Add GoogleGroundingLiveTest integration test - Use shouldNotBeNull() instead of !! in tests Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
1 parent 86cf3b1 commit d7151b5

7 files changed

Lines changed: 71 additions & 168 deletions

File tree

‎integration-tests/src/jvmTest/kotlin/ai/koog/integration/tests/GoogleGroundingLiveTest.kt‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
package ai.koog.integration.tests
2+
3+
import ai.koog.integration.tests.utils.TestCredentials.readTestGoogleAIKeyFromEnv
4+
import ai.koog.prompt.dsl.prompt
5+
import ai.koog.prompt.executor.clients.google.GoogleLLMClient
6+
import ai.koog.prompt.executor.clients.google.GoogleModels
7+
import ai.koog.prompt.executor.clients.google.GoogleParams
8+
import ai.koog.prompt.executor.llms.SingleLLMPromptExecutor
9+
import ai.koog.prompt.message.Message
10+
import io.kotest.matchers.collections.shouldNotBeEmpty
11+
import io.kotest.matchers.types.shouldBeInstanceOf
12+
import kotlinx.coroutines.test.runTest
13+
import org.junit.jupiter.api.Test
14+
import kotlin.time.Duration.Companion.seconds
15+
16+
class GoogleGroundingLiveTest {
17+
18+
private val client = GoogleLLMClient(readTestGoogleAIKeyFromEnv())
19+
private val executor = SingleLLMPromptExecutor(client)
20+
21+
@Test
22+
fun `grounding returns non-empty response for Gemini 2_5 Flash`() = runTest(timeout = 60.seconds) {
23+
val p = prompt("grounding-test", params = GoogleParams(groundingEnabled = true)) {
24+
user("Iran vs US war 2026, what is happening?")
25+
}
26+
val response = executor.execute(p, GoogleModels.Gemini2_5Flash)
27+
response.shouldNotBeEmpty()
28+
response.first().shouldBeInstanceOf<Message.Assistant>()
29+
check((response.first() as Message.Assistant).content.isNotBlank())
30+
}
31+
32+
@Test
33+
fun `grounding with disabled flag returns response without search`() = runTest(timeout = 60.seconds) {
34+
val p = prompt("grounding-off-test", params = GoogleParams(groundingEnabled = false)) {
35+
user("What is the capital of France?")
36+
}
37+
val response = executor.execute(p, GoogleModels.Gemini2_5Flash)
38+
response.shouldNotBeEmpty()
39+
response.first().shouldBeInstanceOf<Message.Assistant>()
40+
}
41+
}

‎prompt/prompt-executor/prompt-executor-clients/prompt-executor-google-client/src/commonMain/kotlin/ai/koog/prompt/executor/clients/google/GoogleLLMClient.kt‎

Lines changed: 5 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ import ai.koog.prompt.executor.clients.google.models.GoogleModelsResponse
2525
import ai.koog.prompt.executor.clients.google.models.GooglePart
2626
import ai.koog.prompt.executor.clients.google.models.GoogleRequest
2727
import ai.koog.prompt.executor.clients.google.models.GoogleResponse
28+
import ai.koog.prompt.executor.clients.google.models.GoogleSearch
2829
import ai.koog.prompt.executor.clients.google.models.GoogleTool
2930
import ai.koog.prompt.executor.clients.google.models.GoogleToolConfig
3031
import ai.koog.prompt.executor.clients.google.structure.GoogleBasicJsonSchemaGenerator
@@ -506,28 +507,10 @@ public open class GoogleLLMClient @JvmOverloads constructor(
506507
null -> null
507508
}
508509

509-
val groundingTool: GoogleTool? = googleParams.groundingConfig?.let { config ->
510-
require(model.supports(LLMCapability.Grounding)) {
511-
"Model ${model.id} does not support grounding"
512-
}
513-
when (config) {
514-
is GoogleGroundingConfig.GoogleSearch -> GoogleTool(googleSearch = buildJsonObject {})
515-
is GoogleGroundingConfig.GoogleSearchRetrieval -> GoogleTool(
516-
googleSearchRetrieval = buildJsonObject {
517-
if (config.dynamicThreshold != null) {
518-
put(
519-
"dynamicRetrievalConfig",
520-
buildJsonObject {
521-
// MODE_DYNAMIC is required for the threshold to take effect;
522-
// without it the API uses MODE_UNSPECIFIED and ignores the threshold.
523-
put("mode", "MODE_DYNAMIC")
524-
put("dynamicThreshold", config.dynamicThreshold)
525-
}
526-
)
527-
}
528-
}
529-
)
530-
}
510+
val groundingTool: GoogleTool? = if (googleParams.groundingEnabled) {
511+
GoogleTool(googleSearch = GoogleSearch())
512+
} else {
513+
null
531514
}
532515

533516
val allTools = when {

‎prompt/prompt-executor/prompt-executor-clients/prompt-executor-google-client/src/commonMain/kotlin/ai/koog/prompt/executor/clients/google/GoogleModels.kt‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -64,11 +64,10 @@ public object GoogleModels : LLModelDefinitions {
6464
)
6565

6666
/**
67-
* Full capabilities including standard, multimodal, tools, native structured output, and grounding
67+
* Full capabilities including standard, multimodal, tools, native structured output
6868
*/
6969
private val fullCapabilities: List<LLMCapability> =
70-
standardCapabilities + multimodalCapabilities + toolCapabilities + structuredOutputCapabilities +
71-
listOf(LLMCapability.Grounding)
70+
standardCapabilities + multimodalCapabilities + toolCapabilities + structuredOutputCapabilities
7271

7372
/**
7473
* Gemini 2.0 Flash is a fast, efficient model for a wide range of tasks.

‎prompt/prompt-executor/prompt-executor-clients/prompt-executor-google-client/src/commonMain/kotlin/ai/koog/prompt/executor/clients/google/GoogleParams.kt‎

Lines changed: 9 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -4,46 +4,6 @@ import ai.koog.prompt.executor.clients.google.models.GoogleThinkingConfig
44
import ai.koog.prompt.params.LLMParams
55
import kotlinx.serialization.json.JsonElement
66

7-
/**
8-
* Controls grounding with Google Search for Gemini models.
9-
*
10-
* Grounding augments model responses with real-time information from Google Search.
11-
* Use [GoogleSearch] for Gemini 2.0+ models, [GoogleSearchRetrieval] for Gemini 1.5 models.
12-
*
13-
* ### Example:
14-
* ```kotlin
15-
* val prompt = prompt("search") { user("Who won Euro 2024?") }
16-
* val response = executor.execute(prompt, GoogleModels.Gemini2_5Flash,
17-
* params = GoogleParams(groundingConfig = GoogleGroundingConfig.GoogleSearch))
18-
* ```
19-
*/
20-
public sealed class GoogleGroundingConfig {
21-
22-
/**
23-
* Enables grounding via the native `google_search` tool.
24-
* Supported by Gemini 2.0+ models.
25-
*/
26-
public data object GoogleSearch : GoogleGroundingConfig()
27-
28-
/**
29-
* Enables grounding via `googleSearchRetrieval` with optional dynamic retrieval.
30-
* Supported by Gemini 1.5 models.
31-
*
32-
* @property dynamicThreshold Confidence threshold in [0.0, 1.0] below which
33-
* the model will trigger a retrieval request. Lower values retrieve more often.
34-
* Defaults to the API default when null.
35-
*/
36-
public data class GoogleSearchRetrieval(
37-
val dynamicThreshold: Double? = null,
38-
) : GoogleGroundingConfig() {
39-
init {
40-
require(dynamicThreshold == null || dynamicThreshold in 0.0..1.0) {
41-
"dynamicThreshold must be in [0.0, 1.0], but was $dynamicThreshold"
42-
}
43-
}
44-
}
45-
}
46-
477
internal fun LLMParams.toGoogleParams(): GoogleParams {
488
if (this is GoogleParams) return this
499
return GoogleParams(
@@ -75,8 +35,8 @@ internal fun LLMParams.toGoogleParams(): GoogleParams {
7535
* @property topK The maximum number of tokens to consider when sampling.
7636
* @property thinkingConfig Controls whether the model should expose its chain-of-thought
7737
* and how many tokens it may spend on it (see [GoogleThinkingConfig]).
78-
* @property groundingConfig Enables grounding with Google Search to augment responses with
79-
* real-time information (see [GoogleGroundingConfig]). Requires [ai.koog.prompt.llm.LLMCapability.Grounding].
38+
* @property groundingEnabled Enables grounding with Google Search to augment responses with
39+
* real-time information. Supported by Gemini 2.0+ models.
8040
*/
8141
@Suppress("LongParameterList")
8242
public class GoogleParams(
@@ -91,7 +51,7 @@ public class GoogleParams(
9151
public val topP: Double? = null,
9252
public val topK: Int? = null,
9353
public val thinkingConfig: GoogleThinkingConfig? = null,
94-
public val groundingConfig: GoogleGroundingConfig? = null,
54+
public val groundingEnabled: Boolean = false,
9555
) : LLMParams(
9656
temperature,
9757
maxTokens,
@@ -135,7 +95,7 @@ public class GoogleParams(
13595
topP = topP,
13696
topK = topK,
13797
thinkingConfig = thinkingConfig,
138-
groundingConfig = groundingConfig,
98+
groundingEnabled = groundingEnabled,
13999
)
140100

141101
/**
@@ -153,7 +113,7 @@ public class GoogleParams(
153113
topP: Double? = this.topP,
154114
topK: Int? = this.topK,
155115
thinkingConfig: GoogleThinkingConfig? = this.thinkingConfig,
156-
groundingConfig: GoogleGroundingConfig? = this.groundingConfig,
116+
groundingEnabled: Boolean = this.groundingEnabled,
157117
): GoogleParams = GoogleParams(
158118
temperature = temperature,
159119
maxTokens = maxTokens,
@@ -166,7 +126,7 @@ public class GoogleParams(
166126
topP = topP,
167127
topK = topK,
168128
thinkingConfig = thinkingConfig,
169-
groundingConfig = groundingConfig,
129+
groundingEnabled = groundingEnabled,
170130
)
171131

172132
override fun equals(other: Any?): Boolean = when {
@@ -184,13 +144,13 @@ public class GoogleParams(
184144
topP == other.topP &&
185145
topK == other.topK &&
186146
thinkingConfig == other.thinkingConfig &&
187-
groundingConfig == other.groundingConfig
147+
groundingEnabled == other.groundingEnabled
188148
}
189149

190150
override fun hashCode(): Int = listOf(
191151
temperature, maxTokens, numberOfChoices,
192152
speculation, schema, toolChoice, user,
193-
additionalProperties, topP, topK, thinkingConfig, groundingConfig
153+
additionalProperties, topP, topK, thinkingConfig, groundingEnabled
194154
).fold(0) { acc, element ->
195155
31 * acc + (element?.hashCode() ?: 0)
196156
}
@@ -208,7 +168,7 @@ public class GoogleParams(
208168
append(", topP=$topP")
209169
append(", topK=$topK")
210170
append(", thinkingConfig=$thinkingConfig")
211-
append(", groundingConfig=$groundingConfig")
171+
append(", groundingEnabled=$groundingEnabled")
212172
append(")")
213173
}
214174
}

‎prompt/prompt-executor/prompt-executor-clients/prompt-executor-google-client/src/commonMain/kotlin/ai/koog/prompt/executor/clients/google/models/GoogleGenerateContent.kt‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -249,12 +249,13 @@ internal sealed interface GoogleData {
249249
* The next conversation turn may contain a `FunctionResponse` with the `Content.role` "function" generation context for
250250
* the next model turn.
251251
*/
252+
@Serializable
253+
internal class GoogleSearch
254+
252255
@Serializable
253256
internal class GoogleTool(
254257
val functionDeclarations: List<GoogleFunctionDeclaration>? = null,
255-
@SerialName("google_search")
256-
val googleSearch: JsonObject? = null,
257-
val googleSearchRetrieval: JsonObject? = null,
258+
val googleSearch: GoogleSearch? = null,
258259
)
259260

260261
/**

‎prompt/prompt-executor/prompt-executor-clients/prompt-executor-google-client/src/jvmTest/kotlin/ai/koog/prompt/executor/clients/google/GoogleLLMClientTest.kt‎

Lines changed: 10 additions & 79 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ import ai.koog.prompt.params.LLMParams
2222
import ai.koog.prompt.streaming.StreamFrame
2323
import io.kotest.matchers.collections.shouldContain
2424
import io.kotest.matchers.collections.shouldHaveSize
25+
import io.kotest.matchers.nulls.shouldNotBeNull
2526
import io.kotest.matchers.shouldBe
2627
import io.kotest.matchers.shouldNotBe
2728
import io.kotest.matchers.types.shouldBeInstanceOf
@@ -666,53 +667,26 @@ class GoogleLLMClientTest {
666667
}
667668

668669
@Test
669-
fun `createGoogleRequest injects google_search tool when GoogleSearch grounding is set`() {
670+
fun `createGoogleRequest injects googleSearch tool when groundingEnabled is true`() {
670671
val client = GoogleLLMClient(apiKey = "apiKey")
671672
val model = GoogleModels.Gemini2_5Flash
672673

673674
val request = client.createGoogleRequest(
674675
prompt = Prompt(
675676
messages = emptyList(),
676677
id = "id",
677-
params = GoogleParams(groundingConfig = GoogleGroundingConfig.GoogleSearch)
678+
params = GoogleParams(groundingEnabled = true)
678679
),
679680
model = model,
680681
tools = emptyList()
681682
)
682683

683-
val tools = request.tools
684-
tools shouldNotBe null
685-
tools!!.shouldHaveSize(1)
684+
val tools = request.tools.shouldNotBeNull()
685+
tools.shouldHaveSize(1)
686686
tools.first().googleSearch shouldNotBe null
687687
tools.first().functionDeclarations shouldBe null
688688
}
689689

690-
@Test
691-
fun `createGoogleRequest injects googleSearchRetrieval tool with threshold when set`() {
692-
val client = GoogleLLMClient(apiKey = "apiKey")
693-
val model = GoogleModels.Gemini2_5Flash
694-
695-
val request = client.createGoogleRequest(
696-
prompt = Prompt(
697-
messages = emptyList(),
698-
id = "id",
699-
params = GoogleParams(groundingConfig = GoogleGroundingConfig.GoogleSearchRetrieval(dynamicThreshold = 0.3))
700-
),
701-
model = model,
702-
tools = emptyList()
703-
)
704-
705-
val tools = request.tools
706-
tools shouldNotBe null
707-
tools!!.shouldHaveSize(1)
708-
val retrieval = tools.first().googleSearchRetrieval
709-
retrieval shouldNotBe null
710-
val retrievalConfig = retrieval!!["dynamicRetrievalConfig"]?.jsonObject
711-
retrievalConfig shouldNotBe null
712-
retrievalConfig!!["mode"]?.jsonPrimitive?.content shouldBe "MODE_DYNAMIC"
713-
retrievalConfig["dynamicThreshold"]?.jsonPrimitive?.content shouldBe "0.3"
714-
}
715-
716690
@Test
717691
fun `createGoogleRequest merges grounding tool with function tools`() {
718692
val client = GoogleLLMClient(apiKey = "apiKey")
@@ -724,58 +698,15 @@ class GoogleLLMClientTest {
724698
prompt = Prompt(
725699
messages = emptyList(),
726700
id = "id",
727-
params = GoogleParams(groundingConfig = GoogleGroundingConfig.GoogleSearch)
701+
params = GoogleParams(groundingEnabled = true)
728702
),
729703
model = model,
730704
tools = listOf(tool)
731705
)
732706

733-
val tools = request.tools
734-
tools shouldNotBe null
735-
tools!!.shouldHaveSize(2)
736-
val hasGrounding = tools.any { it.googleSearch != null }
737-
val hasFunctions = tools.any { it.functionDeclarations != null }
738-
hasGrounding shouldBe true
739-
hasFunctions shouldBe true
740-
}
741-
742-
@Test
743-
fun `createGoogleRequest throws when grounding set on model that does not support it`() {
744-
val client = GoogleLLMClient(apiKey = "apiKey")
745-
val modelWithoutGrounding = GoogleModels.Embeddings.GeminiEmbedding001
746-
747-
val result = runCatching {
748-
client.createGoogleRequest(
749-
prompt = Prompt(
750-
messages = emptyList(),
751-
id = "id",
752-
params = GoogleParams(groundingConfig = GoogleGroundingConfig.GoogleSearch)
753-
),
754-
model = modelWithoutGrounding,
755-
tools = emptyList()
756-
)
757-
}
758-
result.isFailure shouldBe true
759-
}
760-
761-
@Test
762-
fun `GoogleSearchRetrieval rejects dynamicThreshold above 1_0`() {
763-
val result = runCatching { GoogleGroundingConfig.GoogleSearchRetrieval(dynamicThreshold = 1.5) }
764-
result.isFailure shouldBe true
765-
result.exceptionOrNull()!!.message shouldBe "dynamicThreshold must be in [0.0, 1.0], but was 1.5"
766-
}
767-
768-
@Test
769-
fun `GoogleSearchRetrieval rejects negative dynamicThreshold`() {
770-
val result = runCatching { GoogleGroundingConfig.GoogleSearchRetrieval(dynamicThreshold = -0.1) }
771-
result.isFailure shouldBe true
772-
}
773-
774-
@Test
775-
fun `GoogleSearchRetrieval accepts null and boundary dynamicThreshold values`() {
776-
GoogleGroundingConfig.GoogleSearchRetrieval(dynamicThreshold = null)
777-
GoogleGroundingConfig.GoogleSearchRetrieval(dynamicThreshold = 0.0)
778-
GoogleGroundingConfig.GoogleSearchRetrieval(dynamicThreshold = 1.0)
779-
GoogleGroundingConfig.GoogleSearchRetrieval(dynamicThreshold = 0.5)
707+
val tools = request.tools.shouldNotBeNull()
708+
tools.shouldHaveSize(2)
709+
tools.any { it.googleSearch != null } shouldBe true
710+
tools.any { it.functionDeclarations != null } shouldBe true
780711
}
781712
}

‎prompt/prompt-llm/src/commonMain/kotlin/ai/koog/prompt/llm/LLMCapability.kt‎

Lines changed: 0 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -200,18 +200,6 @@ public sealed class LLMCapability(public val id: String) {
200200
}
201201
}
202202

203-
/**
204-
* Represents the ability to ground model responses using live web search results.
205-
*
206-
* Models with this capability can augment their responses with real-time information
207-
* retrieved via Google Search. Enabled via [ai.koog.prompt.executor.clients.google.GoogleGroundingConfig]
208-
* in [ai.koog.prompt.executor.clients.google.GoogleParams].
209-
*
210-
* Supported by Gemini 2.0+ models.
211-
*/
212-
@Serializable
213-
public data object Grounding : LLMCapability("grounding")
214-
215203
/**
216204
* Represents an OpenAI-related API endpoint for Large Language Model (LLM) operations.
217205
*

0 commit comments

Comments
 (0)