ADFA-6054 | Route embeddings to selected backend and support batch indexing (#104)

* feat(ai): route embeddings to the backend the user selected

Implement EmbeddingBackend in ai-agent-openai (/v1/embeddings, reordered by
index) and ai-agent-gemini (batchEmbedContents, split at the per-call cap);
ai-core resolves it as a chat turn does. Vector-Search drops the lexical
fallback, batches, and stamps provenance so a change reindexes (schema v2).

Refs: ADFA-6054

* fix(ai): address review on embedding indexing

Scope the vector index per project and restore it from the database, bound the query embedding to the search budget, and report embedding failures as failures.

* fix(ai-core): persist the shown backend once it is configured

Author: jatezzz

plugins/AI-Agent-Gemini/README.md

@@ -59,7 +59,8 @@ Every source file sits in a package named for its layer; nothing is loose at the
 root of `com/itsaky/androidide/plugins/aiagentgemini/`.
 
 - `plugin/GeminiPlugin.kt` — plugin entry point; registers the backend with ai-core
-- `backend/GeminiBackend.kt` — the REST transport, streaming (SSE) and model catalog
+- `backend/GeminiBackend.kt` — the REST transport, streaming (SSE), embeddings and model catalog
+- `backend/GeminiEmbeddingProtocol.kt` — the `batchEmbedContents` body, per-call cap and positional reply (pure)
 - `errors/GeminiErrorFormatter.kt` — turns an API failure into one translated sentence
 - `security/SecureApiKeyStore.kt` — this plugin's Keystore alias, over the IDE's `KeystoreSecretStore`
 - `preferences/GeminiPreferences.kt` — this plugin's settings store, plus the

plugins/AI-Agent-Gemini/src/main/AndroidManifest.xml

@@ -44,7 +44,7 @@
              26.35 would name a release that cannot run this plugin. Do not lower it. -->
         <meta-data
             android:name="plugin.min_ide_version"
-            android:value="26.36" />
+            android:value="26.39" />
 
         <meta-data
             android:name="plugin.max_ide_version"

plugins/AI-Agent-Gemini/src/main/kotlin/com/itsaky/androidide/plugins/aiagentgemini/backend/GeminiBackend.kt

@@ -0,0 +1,86 @@
+package com.itsaky.androidide.plugins.aiagentgemini.backend
+
+import java.io.IOException
+import org.json.JSONArray
+import org.json.JSONObject
+
+/**
+ * The `models/{model}:batchEmbedContents` wire format, as pure functions.
+ *
+ * Kept apart from [GeminiBackend] so the two rules that are easy to get silently wrong are
+ * unit-testable without a key: the API answers **positionally**, with no `index` field of the kind
+ * OpenAI sends, and it caps how many requests one call may carry.
+ */
+internal object GeminiEmbeddingProtocol {
+
+    /**
+     * Requests per call, as the API's documented cap.
+     *
+     * A call over the cap is rejected whole, so a batch the caller asked for is split here rather
+     * than handed to the server and hoped for.
+     */
+    const val MAX_REQUESTS_PER_CALL = 100
+
+    /**
+     * Splits [texts] into calls of at most [MAX_REQUESTS_PER_CALL], preserving order.
+     *
+     * @param texts the batch as the caller asked for it
+     * @return the calls to issue, in order; empty when [texts] is empty
+     */
+    fun batches(texts: List<String>): List<List<String>> = texts.chunked(MAX_REQUESTS_PER_CALL)
+
+    /**
+     * Builds the request body for one call.
+     *
+     * Each entry repeats the model because the API requires it per request, even though every
+     * request in a call must name the same one.
+     *
+     * @param model the embedding model id, without the `models/` prefix
+     * @param texts the call's texts, in order
+     * @return the request JSON
+     */
+    fun body(model: String, texts: List<String>): JSONObject {
+        val requests = JSONArray()
+        for (text in texts) {
+            requests.put(
+                JSONObject()
+                    .put("model", "models/$model")
+                    .put(
+                        "content",
+                        JSONObject().put("parts", JSONArray().put(JSONObject().put("text", text)))
+                    )
+            )
+        }
+        return JSONObject().put("requests", requests)
+    }
+
+    /**
+     * Reads a response back into the caller's order.
+     *
+     * The answer carries no index, so position *is* the contract: a response of the wrong length
+     * cannot be realigned and has to fail here, rather than shift every later vector onto the
+     * wrong chunk.
+     *
+     * @param response the parsed response body
+     * @param expected how many texts the call asked for
+     * @return the vectors, in request order
+     * @throws IOException when the answer is the wrong length or carries an unusable vector
+     */
+    fun vectors(response: JSONObject, expected: Int): List<FloatArray> {
+        val embeddings = response.optJSONArray("embeddings")
+            ?: throw IOException("Embeddings response carried no embeddings array")
+        if (embeddings.length() != expected) {
+            throw IOException(
+                "Embeddings response carried ${embeddings.length()} vectors for $expected inputs"
+            )
+        }
+
+        return (0 until expected).map { position ->
+            val values = embeddings.optJSONObject(position)?.optJSONArray("values")
+            if (values == null || values.length() == 0) {
+                throw IOException("Embeddings response element $position carried no vector")
+            }
+            FloatArray(values.length()) { index -> values.getDouble(index).toFloat() }
+        }
+    }
+}

plugins/AI-Agent-Gemini/src/main/kotlin/com/itsaky/androidide/plugins/aiagentgemini/backend/GeminiEmbeddingProtocol.kt

@@ -0,0 +1,21 @@
+package com.itsaky.androidide.plugins.aiagentgemini.backend
+
+/**
+ * One key's catalog, split into the models each picker may offer.
+ *
+ * Carried as one value because both halves come from one paginated `ListModels` walk, filtered on
+ * the capability each model declares: fetching them separately would pay for the walk twice and
+ * let the two pickers describe different snapshots.
+ *
+ * @param chat models that advertise `generateContent`
+ * @param embedding models that advertise `embedContent`
+ */
+internal data class ModelCatalog(
+    val chat: List<String>,
+    val embedding: List<String>,
+) {
+    companion object {
+        /** What a key that lists nothing offers. */
+        val EMPTY = ModelCatalog(emptyList(), emptyList())
+    }
+}

plugins/AI-Agent-Gemini/src/main/kotlin/com/itsaky/androidide/plugins/aiagentgemini/backend/ModelCatalog.kt

@@ -14,6 +14,9 @@ internal object NetworkTags {
 
     /** The ListModels catalog — `"GECT"`. */
     const val CATALOG = 0x47454354
+
+    /** Batch embedding — `"GEEM"`. */
+    const val EMBEDDING = 0x4745454D
 }
 
 /**

plugins/AI-Agent-Gemini/src/main/kotlin/com/itsaky/androidide/plugins/aiagentgemini/backend/NetworkTags.kt

@@ -54,6 +54,7 @@ class GeminiPlugin : IPlugin, DocumentationExtension {
         // Tags for the controls on this backend's settings pane (see GeminiSettingsFragment).
         const val TOOLTIP_TAG_SETTINGS_GEMINI_KEY = "ai_gemini_key"
         const val TOOLTIP_TAG_SETTINGS_GEMINI_MODEL = "ai_gemini_model"
+        const val TOOLTIP_TAG_SETTINGS_GEMINI_EMBEDDING_MODEL = "ai_gemini_embedding_model"
         const val TOOLTIP_TAG_SETTINGS_GET_KEY = "ai_gemini_get_free_key"
 
         @Volatile
@@ -268,6 +269,22 @@ class GeminiPlugin : IPlugin, DocumentationExtension {
                 current models is shown instead.</p>
             """.trimIndent(),
         ),
+        PluginTooltipEntry(
+            tag = TOOLTIP_TAG_SETTINGS_GEMINI_EMBEDDING_MODEL,
+            summary = "Which model turns your code into vectors for semantic search. Never used for chat.",
+            detail = """
+                <p>Semantic search compares meaning rather than words, which it
+                does by embedding every chunk of the project with this model. It
+                is a separate setting because no Gemini model does both: this list
+                holds the models that advertise <code>embedContent</code>, and the
+                <b>Gemini Model</b> list above holds those that advertise
+                <code>generateContent</code>.</p>
+                <p>Changing it changes the vector space, so the project is indexed
+                again from scratch. Vectors from two different models are not
+                comparable, and mixing them would quietly return worse results
+                rather than fail.</p>
+            """.trimIndent(),
+        ),
         PluginTooltipEntry(
             tag = TOOLTIP_TAG_SETTINGS_GET_KEY,
             summary = "Opens Google AI Studio in your browser, where API keys are free to create.",

plugins/AI-Agent-Gemini/src/main/kotlin/com/itsaky/androidide/plugins/aiagentgemini/plugin/GeminiPlugin.kt

@@ -23,6 +23,14 @@ internal object GeminiPreferences {
     const val KEY_API_KEY_VERIFIED = "gemini_api_key_verified"
     const val KEY_MODEL = "gemini_model"
 
+    /**
+     * Embedding model id, e.g. `gemini-embedding-001`.
+     *
+     * Separate from [KEY_MODEL] because no Gemini model advertises both `generateContent` and
+     * `embedContent`, so one setting could never serve both.
+     */
+    const val KEY_EMBEDDING_MODEL = "gemini_embedding_model"
+
     /**
      * Why the last request was refused for credential reasons, or absent. Diagnostics rather than a
      * setting, so deliberately outside [OWNED_KEYS]: there is nothing here worth carrying over from
@@ -42,7 +50,7 @@ internal object GeminiPreferences {
 
     /** Everything this backend owns; anything else in the old shared file is not ours to take. */
     private val OWNED_KEYS = listOf(
-        KEY_API_KEY, KEY_API_KEY_TIMESTAMP, KEY_API_KEY_VERIFIED, KEY_MODEL,
+        KEY_API_KEY, KEY_API_KEY_TIMESTAMP, KEY_API_KEY_VERIFIED, KEY_MODEL, KEY_EMBEDDING_MODEL,
     )
 
     /**

plugins/AI-Agent-Gemini/src/main/kotlin/com/itsaky/androidide/plugins/aiagentgemini/preferences/GeminiPreferences.kt

@@ -8,8 +8,20 @@ package com.itsaky.androidide.plugins.aiagentgemini.settings
  */
 sealed interface CatalogResult {
 
-    /** The backend answered. [models] may be empty, which is itself suspicious for a valid key. */
-    data class Success(val models: List<String>) : CatalogResult
+    /**
+     * The backend answered. Either list may be empty, which for [models] is itself suspicious for
+     * a valid key.
+     *
+     * Both halves come from one paginated walk, split by the method each model declares, so they
+     * describe the same snapshot of the same key.
+     *
+     * @param models the chat-capable models the key can reach
+     * @param embeddingModels the embedding-capable models the key can reach
+     */
+    data class Success(
+        val models: List<String>,
+        val embeddingModels: List<String>,
+    ) : CatalogResult
 
     /** No "gemini" backend was resolvable — ai-core or ai-agent-gemini is missing, disabled,
      * or not yet active. */

plugins/AI-Agent-Gemini/src/main/kotlin/com/itsaky/androidide/plugins/aiagentgemini/settings/CatalogResult.kt

@@ -2,6 +2,7 @@ package com.itsaky.androidide.plugins.aiagentgemini.settings
 
 import com.itsaky.androidide.plugins.PluginLogger
 import com.itsaky.androidide.plugins.aiagentgemini.backend.GeminiBackend
+import com.itsaky.androidide.plugins.aiagentgemini.backend.ModelCatalog
 import com.itsaky.androidide.plugins.aiagentgemini.logging.LOG_PREFIX
 import com.itsaky.androidide.plugins.aiagentgemini.plugin.GeminiPlugin
 import java.util.concurrent.CancellationException
@@ -63,10 +64,10 @@ class BackendGeminiCatalogGateway(
         get() = GeminiPlugin.getContext()?.logger
 
     override fun listModelsForSavedKey(): CatalogResult =
-        await { it.listModels() }
+        await { it.listCatalog() }
 
     override fun listModels(apiKey: String): CatalogResult =
-        await { it.listModels(apiKey) }
+        await { it.listCatalog(apiKey) }
 
     /**
      * Runs [request] against the backend and awaits its future.
@@ -76,7 +77,7 @@ class BackendGeminiCatalogGateway(
      * @param request the catalog call to make; picks which credential is used
      */
     private fun await(
-        request: (GeminiBackend) -> CompletableFuture<List<String>>
+        request: (GeminiBackend) -> CompletableFuture<ModelCatalog>
     ): CatalogResult {
         val backend = try {
             backendProvider()
@@ -93,7 +94,9 @@ class BackendGeminiCatalogGateway(
         }
 
         return try {
-            CatalogResult.Success(future.get(LIST_MODELS_TIMEOUT_SECONDS, TimeUnit.SECONDS).orEmpty())
+            val catalog = future.get(LIST_MODELS_TIMEOUT_SECONDS, TimeUnit.SECONDS)
+                ?: ModelCatalog.EMPTY
+            CatalogResult.Success(catalog.chat, catalog.embedding)
         } catch (e: ExecutionException) {
             // The API failure the backend reported; its message carries the HTTP status.
             CatalogResult.Failed(e.cause ?: e)