KG-314 Add GeminiEmbedding to GoogleModels.kt (#1235)

Related to [KG-314](https://youtrack.jetbrains.com/issue/KG-314)

## Motivation and Context
This PR enables text embedding support for Google's Gemini models within
the Koog framework, resolving issue #713.

The `GoogleLLMClient` has been updated to implement the
`LLMEmbeddingProvider` interface, allowing users to generate vector
embeddings alongside existing chat functionality. This includes the
addition of the `gemini-embedding-001` model definition and the
necessary API integration for the `embedContent` endpoint.

**Testing:**
Functionality has been verified by adding the Gemini embedding model to
the standard embedding integration test suite (`integration_testEmbed`).

## Breaking Changes
None.

---

#### Type of the changes
- [x] New feature (non-breaking change which adds functionality)
- [ ] Bug fix (non-breaking change which fixes an issue)
- [ ] Breaking change (fix or feature that would cause existing
functionality to change)
- [ ] Documentation update
- [x] Tests improvement
- [ ] Refactoring

#### Checklist
- [x] The pull request has a description of the proposed change
- [x] I read the [Contributing
Guidelines](https://github.com/JetBrains/koog/blob/main/CONTRIBUTING.md)
before opening the pull request
- [x] The pull request uses **`develop`** as the base branch
- [x] Tests for the changes have been added
- [x] All new and existing tests passed

##### Additional steps for pull requests adding a new feature
- [x] An issue describing the proposed change exists
- [x] The pull request includes a link to the issue
- [x] The change was discussed and approved in the issue
- [ ] Docs have been added / updated

---------

Co-authored-by: Malte Heuser <[email protected]>

Author: mltheuser

embeddings/embeddings-llm/build.gradle.kts

@@ -26,6 +26,7 @@ kotlin {
             dependencies {
                 implementation(kotlin("test"))
                 implementation(project(":prompt:prompt-executor:prompt-executor-clients:prompt-executor-openai-client"))
+                implementation(project(":prompt:prompt-executor:prompt-executor-clients:prompt-executor-google-client"))
                 implementation(libs.kotlinx.coroutines.core)
                 implementation(libs.kotlinx.coroutines.test)
             }

embeddings/embeddings-llm/src/commonTest/kotlin/ai/koog/embeddings/local/LLMEmbedderTest.kt

@@ -2,6 +2,7 @@ package ai.koog.embeddings.local
 
 import ai.koog.embeddings.base.Vector
 import ai.koog.prompt.executor.clients.LLMEmbeddingProvider
+import ai.koog.prompt.executor.clients.google.GoogleModels
 import ai.koog.prompt.executor.clients.openai.OpenAIModels
 import ai.koog.prompt.llm.LLModel
 import kotlinx.coroutines.test.runTest
@@ -14,6 +15,7 @@ class LLMEmbedderTest {
     val modelsList = listOf(
         OpenAIModels.Embeddings.TextEmbedding3Small,
         OllamaEmbeddingModels.NOMIC_EMBED_TEXT,
+        GoogleModels.Embeddings.GeminiEmbedding001,
     )
 
     @Test

integration-tests/src/jvmTest/kotlin/ai/koog/integration/tests/utils/Models.kt

@@ -64,6 +64,7 @@ object Models {
             BedrockModels.Embeddings.AmazonTitanEmbedText,
             OpenAIModels.Embeddings.TextEmbedding3Large,
             MistralAIModels.Embeddings.MistralEmbed,
+            GoogleModels.Embeddings.GeminiEmbedding001,
         )
     }
 

koog-ktor/src/commonMain/kotlin/ai/koog/ktor/utils/LLMModelParser.kt

@@ -260,6 +260,7 @@ private val GOOGLE_MODELS_MAP = mapOf(
     "gemini2_5pro" to GoogleModels.Gemini2_5Pro,
     "gemini2_5flash" to GoogleModels.Gemini2_5Flash,
     "gemini2_5flashlite" to GoogleModels.Gemini2_5FlashLite,
+    "gemini_embedding001" to GoogleModels.Embeddings.GeminiEmbedding001,
 )
 
 private val MISTRAL_MODELS_MAP = mapOf(

koog-ktor/src/commonTest/kotlin/ai/koog/ktor/ModelIdentifierParsingTest.kt

@@ -249,6 +249,11 @@ class ModelIdentifierParsingTest {
         assertNotNull(gemini25FlashLite)
         assertEquals(LLMProvider.Google, gemini25FlashLite.provider)
         assertEquals(GoogleModels.Gemini2_5FlashLite, gemini25FlashLite)
+
+        val geminiEmbedding001 = getModelFromIdentifier("google.gemini_embedding001")
+        assertNotNull(geminiEmbedding001)
+        assertEquals(LLMProvider.Google, geminiEmbedding001.provider)
+        assertEquals(GoogleModels.Embeddings.GeminiEmbedding001, geminiEmbedding001)
     }
 
     // MistralAI model identifier tests

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

@@ -10,9 +10,12 @@ import ai.koog.prompt.dsl.Prompt
 import ai.koog.prompt.executor.clients.ConnectionTimeoutConfig
 import ai.koog.prompt.executor.clients.LLMClient
 import ai.koog.prompt.executor.clients.LLMClientException
+import ai.koog.prompt.executor.clients.LLMEmbeddingProvider
 import ai.koog.prompt.executor.clients.google.models.GoogleCandidate
 import ai.koog.prompt.executor.clients.google.models.GoogleContent
 import ai.koog.prompt.executor.clients.google.models.GoogleData
+import ai.koog.prompt.executor.clients.google.models.GoogleEmbeddingRequest
+import ai.koog.prompt.executor.clients.google.models.GoogleEmbeddingResponse
 import ai.koog.prompt.executor.clients.google.models.GoogleFunctionCallingConfig
 import ai.koog.prompt.executor.clients.google.models.GoogleFunctionCallingMode
 import ai.koog.prompt.executor.clients.google.models.GoogleFunctionDeclaration
@@ -79,6 +82,7 @@ public class GoogleClientSettings(
     public val defaultPath: String = "v1beta/models",
     public val generateContentMethod: String = "generateContent",
     public val streamGenerateContentMethod: String = "streamGenerateContent",
+    public val embedContentMethod: String = "embedContent"
 )
 
 /**
@@ -97,7 +101,7 @@ public open class GoogleLLMClient(
     private val settings: GoogleClientSettings = GoogleClientSettings(),
     baseClient: HttpClient = HttpClient(),
     private val clock: Clock = Clock.System
-) : LLMClient {
+) : LLMClient, LLMEmbeddingProvider {
 
     @OptIn(InternalStructuredOutputApi::class)
     private companion object {
@@ -757,4 +761,38 @@ public open class GoogleLLMClient(
     override fun close() {
         httpClient.close()
     }
+
+    override suspend fun embed(text: String, model: LLModel): List<Double> {
+        require(model.capabilities.contains(LLMCapability.Embed)) {
+            "Model ${model.id} does not support embedding."
+        }
+
+        logger.debug { "Embedding text with model: ${model.id}" }
+
+        val request = GoogleEmbeddingRequest(
+            model = "models/${model.id}",
+            content = GoogleContent(
+                parts = listOf(GooglePart.Text(text))
+            )
+        )
+
+        try {
+            val response = httpClient.post(
+                path = "${settings.defaultPath}/${model.id}:${settings.embedContentMethod}",
+                request = request,
+                requestBodyType = GoogleEmbeddingRequest::class,
+                responseType = GoogleEmbeddingResponse::class,
+            )
+
+            return response.embedding.values
+        } catch (e: CancellationException) {
+            throw e
+        } catch (e: Exception) {
+            throw LLMClientException(
+                clientName = clientName,
+                message = e.message,
+                cause = e
+            )
+        }
+    }
 }

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

@@ -150,4 +150,23 @@ public object GoogleModels : LLModelDefinitions {
         contextLength = 1_048_576,
         maxOutputTokens = 65_536,
     )
+
+    /**
+     * Models for generating text embeddings.
+     */
+    public object Embeddings {
+        /**
+         * Gemini embedding model for generating embeddings for words, phrases, and sentences.
+         *
+         * Input token limit: 2048
+         *
+         * @see <a href="https://ai.google.dev/gemini-api/docs/embeddings#model-versions">
+         */
+        public val GeminiEmbedding001: LLModel = LLModel(
+            provider = LLMProvider.Google,
+            id = "gemini-embedding-001",
+            capabilities = listOf(LLMCapability.Embed),
+            contextLength = 2048,
+        )
+    }
 }

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

@@ -0,0 +1,19 @@
+package ai.koog.prompt.executor.clients.google.models
+
+import kotlinx.serialization.Serializable
+
+@Serializable
+internal data class GoogleEmbeddingRequest(
+    val model: String,
+    val content: GoogleContent
+)
+
+@Serializable
+internal data class GoogleEmbeddingResponse(
+    val embedding: GoogleEmbeddingData
+)
+
+@Serializable
+internal data class GoogleEmbeddingData(
+    val values: List<Double>
+)