From 205e22ee4672305d3cb2e0c34b607a4950f8f4e8 Mon Sep 17 00:00:00 2001 From: marius-kilocode Date: Thu, 21 May 2026 13:06:36 +0200 Subject: [PATCH] fix(indexing): support OpenRouter Gemini embeddings --- .changeset/gentle-indexing-dimensions.md | 5 +++ .../src/indexing/config-manager.ts | 5 +-- .../src/indexing/embedders/openrouter.ts | 38 +++++++++++++---- .../src/indexing/embedding-profile.ts | 2 +- .../src/indexing/model-registry.ts | 1 + .../kilocode/indexing/config-manager.test.ts | 14 +++++++ .../indexing/embedders/openrouter.test.ts | 41 +++++++++++-------- .../kilocode/indexing/service-factory.test.ts | 34 ++++++++++++++- packages/kilo-vscode/script/launch.ts | 3 ++ packages/kilo-vscode/script/local-bin.ts | 11 ++++- 10 files changed, 123 insertions(+), 31 deletions(-) create mode 100644 .changeset/gentle-indexing-dimensions.md diff --git a/.changeset/gentle-indexing-dimensions.md b/.changeset/gentle-indexing-dimensions.md new file mode 100644 index 0000000000..1541b934e5 --- /dev/null +++ b/.changeset/gentle-indexing-dimensions.md @@ -0,0 +1,5 @@ +--- +"@kilocode/kilo-indexing": patch +--- + +Support OpenRouter Gemini embedding preview indexing and honor configured embedding dimensions when sizing vector stores. diff --git a/packages/kilo-indexing/src/indexing/config-manager.ts b/packages/kilo-indexing/src/indexing/config-manager.ts index 1c5124de2d..fd8c63210d 100644 --- a/packages/kilo-indexing/src/indexing/config-manager.ts +++ b/packages/kilo-indexing/src/indexing/config-manager.ts @@ -302,10 +302,9 @@ export class CodeIndexConfigManager { } public get currentModelDimension(): number | undefined { + if (this.modelDimension && this.modelDimension > 0) return this.modelDimension const id = this.modelId ?? getDefaultModelId(this.embedderProvider) - const dim = getModelDimension(this.embedderProvider, id) - if (!dim && this.modelDimension && this.modelDimension > 0) return this.modelDimension - return dim + return getModelDimension(this.embedderProvider, id) } public get currentSearchMinScore(): number { diff --git a/packages/kilo-indexing/src/indexing/embedders/openrouter.ts b/packages/kilo-indexing/src/indexing/embedders/openrouter.ts index c055964a5c..e774018fb9 100644 --- a/packages/kilo-indexing/src/indexing/embedders/openrouter.ts +++ b/packages/kilo-indexing/src/indexing/embedders/openrouter.ts @@ -25,7 +25,8 @@ interface EmbeddingItem { } interface OpenRouterEmbeddingResponse { - data: EmbeddingItem[] + data?: EmbeddingItem[] + error?: string | { code?: string | number; message?: string } usage?: { prompt_tokens?: number total_tokens?: number @@ -193,10 +194,7 @@ export class OpenRouterEmbedder implements IEmbedder { const requestParams: any = { input: batchTexts, model: model, - // OpenAI package (as of v4.78.1) has a parsing issue that truncates embedding dimensions to 256 - // when processing numeric arrays, which breaks compatibility with models using larger dimensions. - // By requesting base64 encoding, we bypass the package's parser and handle decoding ourselves. - encoding_format: "base64", + encoding_format: "float", } if (this.dimensions !== undefined) { @@ -213,8 +211,22 @@ export class OpenRouterEmbedder implements IEmbedder { } const response = (await this.embeddingsClient.embeddings.create(requestParams)) as OpenRouterEmbeddingResponse + const err = response.error + const msg = typeof err === "string" ? err : err?.message + const code = typeof err === "object" && err ? err.code : undefined + if (!response.data || response.data.length === 0) { + log.warn("OpenRouter embedder batch returned invalid response", { + location: "OpenRouterEmbedder:_embedBatchWithRetries", + model, + dimensions: this.dimensions, + provider: this.specificProvider, + code, + err: msg, + }) + throw new Error(msg ?? "Invalid response from OpenRouter embedding endpoint") + } - // Convert base64 embeddings to float32 arrays + // Normalize base64 embeddings if OpenRouter returns them despite the float request. const processedEmbeddings = response.data.map((item: EmbeddingItem) => { if (typeof item.embedding === "string") { const buffer = Buffer.from(item.embedding, "base64") @@ -292,7 +304,7 @@ export class OpenRouterEmbedder implements IEmbedder { const requestParams: any = { input: testTexts, model: modelToUse, - encoding_format: "base64", + encoding_format: "float", } if (this.dimensions !== undefined) { @@ -315,6 +327,18 @@ export class OpenRouterEmbedder implements IEmbedder { // Check if we got a valid response if (!response?.data || response.data.length === 0) { + const err = response?.error + const msg = typeof err === "string" ? err : err?.message + const code = typeof err === "object" && err ? err.code : undefined + log.warn("OpenRouter embedder validation returned invalid response", { + location: "OpenRouterEmbedder:validateConfiguration", + model: modelToUse, + dimensions: this.dimensions, + provider: this.specificProvider, + dataCount: response?.data?.length ?? 0, + code, + err: msg, + }) return { valid: false, error: "Invalid response from OpenRouter embedding endpoint", diff --git a/packages/kilo-indexing/src/indexing/embedding-profile.ts b/packages/kilo-indexing/src/indexing/embedding-profile.ts index 843288e80a..9d37d8eacb 100644 --- a/packages/kilo-indexing/src/indexing/embedding-profile.ts +++ b/packages/kilo-indexing/src/indexing/embedding-profile.ts @@ -20,7 +20,7 @@ export function resolveEmbeddingProfile( modelDimension?: number, ): EmbeddingProfile | undefined { const id = modelId ?? getDefaultModelId(provider) - const dim = getModelDimension(provider, id) ?? parseDimension(modelDimension) + const dim = parseDimension(modelDimension) ?? getModelDimension(provider, id) if (!dim) return undefined return { provider, diff --git a/packages/kilo-indexing/src/indexing/model-registry.ts b/packages/kilo-indexing/src/indexing/model-registry.ts index 12a427d210..d02b488cbe 100644 --- a/packages/kilo-indexing/src/indexing/model-registry.ts +++ b/packages/kilo-indexing/src/indexing/model-registry.ts @@ -48,6 +48,7 @@ const profiles: Record> = { openrouter: { "openai/text-embedding-3-small": { dimension: 1536, scoreThreshold: 0.4 }, "openai/text-embedding-3-large": { dimension: 3072, scoreThreshold: 0.4 }, + "google/gemini-embedding-2-preview": { dimension: 3072, scoreThreshold: 0.35 }, }, "openai-compatible": {}, "vercel-ai-gateway": { diff --git a/packages/kilo-indexing/test/kilocode/indexing/config-manager.test.ts b/packages/kilo-indexing/test/kilocode/indexing/config-manager.test.ts index cc56020ab7..4ba3338df4 100644 --- a/packages/kilo-indexing/test/kilocode/indexing/config-manager.test.ts +++ b/packages/kilo-indexing/test/kilocode/indexing/config-manager.test.ts @@ -83,6 +83,20 @@ describe("CodeIndexConfigManager", () => { expect(cfg.currentModelDimension).toBe(2048) }) + test("uses configured dimension before static model metadata", () => { + const cfg = new CodeIndexConfigManager( + createInput({ + embedderProvider: "openrouter", + openAiKey: undefined, + openRouterApiKey: "or-test", + modelId: "google/gemini-embedding-2-preview", + modelDimension: 1536, + }), + ) + + expect(cfg.currentModelDimension).toBe(1536) + }) + describe("loadConfiguration restart checks", () => { test("requires restart when model changes with same dimension", () => { const cfg = new CodeIndexConfigManager(createInput({ modelId: "text-embedding-3-small" })) diff --git a/packages/kilo-indexing/test/kilocode/indexing/embedders/openrouter.test.ts b/packages/kilo-indexing/test/kilocode/indexing/embedders/openrouter.test.ts index 729e5c7787..27f39c7180 100644 --- a/packages/kilo-indexing/test/kilocode/indexing/embedders/openrouter.test.ts +++ b/packages/kilo-indexing/test/kilocode/indexing/embedders/openrouter.test.ts @@ -68,14 +68,10 @@ describe("OpenRouterEmbedder", () => { }) test("should create embeddings successfully", async () => { - // Create base64 encoded embedding with values that can be exactly represented in Float32 - const testEmbedding = new Float32Array([0.25, 0.5, 0.75]) - const base64String = Buffer.from(testEmbedding.buffer).toString("base64") - const mockResponse = { data: [ { - embedding: base64String, + embedding: [0.25, 0.5, 0.75], }, ], usage: { @@ -91,7 +87,7 @@ describe("OpenRouterEmbedder", () => { expect(mockEmbeddingsCreate).toHaveBeenCalledWith({ input: ["test text"], model: defaultModel, - encoding_format: "base64", + encoding_format: "float", }) expect(result.embeddings).toHaveLength(1) expect(result.embeddings[0]).toEqual([0.25, 0.5, 0.75]) @@ -156,7 +152,7 @@ describe("OpenRouterEmbedder", () => { expect(mockEmbeddingsCreate).toHaveBeenCalledWith({ input: ["test"], model: customModel, - encoding_format: "base64", + encoding_format: "float", }) }) @@ -187,7 +183,7 @@ describe("OpenRouterEmbedder", () => { expect(mockEmbeddingsCreate).toHaveBeenCalledWith({ input: ["test"], model: defaultModel, - encoding_format: "base64", + encoding_format: "float", provider: { order: [specificProvider], only: [specificProvider], @@ -221,7 +217,7 @@ describe("OpenRouterEmbedder", () => { expect(mockEmbeddingsCreate).toHaveBeenCalledWith({ input: ["test"], model: defaultModel, - encoding_format: "base64", + encoding_format: "float", dimensions: 1024, }) }) @@ -257,7 +253,7 @@ describe("OpenRouterEmbedder", () => { expect(mockEmbeddingsCreate).toHaveBeenCalledWith({ input: ["test"], model: defaultModel, - encoding_format: "base64", + encoding_format: "float", }) }) }) @@ -271,13 +267,10 @@ describe("OpenRouterEmbedder", () => { }) test("should validate configuration successfully", async () => { - const testEmbedding = new Float32Array([0.25, 0.5]) - const base64String = Buffer.from(testEmbedding.buffer).toString("base64") - const mockResponse = { data: [ { - embedding: base64String, + embedding: [0.25, 0.5], }, ], usage: { @@ -296,7 +289,7 @@ describe("OpenRouterEmbedder", () => { { input: ["test"], model: defaultModel, - encoding_format: "base64", + encoding_format: "float", }, { timeout: REMOTE_EMBEDDER_VALIDATION_TIMEOUT_MS, @@ -305,6 +298,20 @@ describe("OpenRouterEmbedder", () => { ) }) + test("should reject responses without embedding data", async () => { + mockEmbeddingsCreate.mockResolvedValue({ + error: { + code: 404, + message: "No successful provider responses.", + }, + }) + + const result = await embedder.validateConfiguration() + + expect(result.valid).toBe(false) + expect(result.error).toBe("Invalid response from OpenRouter embedding endpoint") + }) + test("should handle validation failure", async () => { const authError = new Error("Invalid API key") ;(authError as any).status = 401 @@ -346,7 +353,7 @@ describe("OpenRouterEmbedder", () => { { input: ["test"], model: defaultModel, - encoding_format: "base64", + encoding_format: "float", provider: { order: [specificProvider], only: [specificProvider], @@ -388,7 +395,7 @@ describe("OpenRouterEmbedder", () => { { input: ["test"], model: defaultModel, - encoding_format: "base64", + encoding_format: "float", dimensions: 1024, }, { diff --git a/packages/kilo-indexing/test/kilocode/indexing/service-factory.test.ts b/packages/kilo-indexing/test/kilocode/indexing/service-factory.test.ts index 67aff168c9..86d50ccc4a 100644 --- a/packages/kilo-indexing/test/kilocode/indexing/service-factory.test.ts +++ b/packages/kilo-indexing/test/kilocode/indexing/service-factory.test.ts @@ -149,11 +149,42 @@ describe("CodeIndexServiceFactory", () => { expect(mockEmbeddingsCreate).toHaveBeenCalledWith({ input: ["hello"], model: "openai/text-embedding-3-small", - encoding_format: "base64", + encoding_format: "float", dimensions: 1024, }) }) + test("creates vector store for OpenRouter Gemini embedding preview", () => { + const factory = createFactory({ + embedderProvider: "openrouter", + openAiKey: undefined, + openRouterApiKey: "or-test", + modelId: "google/gemini-embedding-2-preview", + vectorStoreProvider: "lancedb", + }) + + const store = factory.createVectorStore() as unknown as { vectorSize: number } + + expect(store).toBeDefined() + expect(store.vectorSize).toBe(3072) + }) + + test("uses configured dimension before static model metadata for vector stores", () => { + const factory = createFactory({ + embedderProvider: "openrouter", + openAiKey: undefined, + openRouterApiKey: "or-test", + modelId: "openai/text-embedding-3-small", + modelDimension: 1024, + vectorStoreProvider: "lancedb", + }) + + const store = factory.createVectorStore() as unknown as { vectorSize: number } + + expect(store).toBeDefined() + expect(store.vectorSize).toBe(1024) + }) + test("creates Kilo embedder with Cloud-provided model", async () => { const factory = createFactory({ embedderProvider: "kilo", @@ -177,6 +208,7 @@ describe("CodeIndexServiceFactory", () => { input: ["hello"], model: "mistralai/mistral-embed-2312", encoding_format: "base64", + dimensions: 1024, }) }) }) diff --git a/packages/kilo-vscode/script/launch.ts b/packages/kilo-vscode/script/launch.ts index caa0a8b493..9973e7a394 100644 --- a/packages/kilo-vscode/script/launch.ts +++ b/packages/kilo-vscode/script/launch.ts @@ -330,6 +330,9 @@ async function launch() { // Strip Electron/VS Code env vars so the spawned instance doesn't attach // to the current Electron process (e.g. when launched from a VS Code task). const env = cleanEnv(process.env) + if (mode === "dev") { + env.KILO_INDEXING_LOG = "1" + } for (const key of Object.keys(env)) { if (key.startsWith("ELECTRON_") || key.startsWith("VSCODE_")) delete env[key] } diff --git a/packages/kilo-vscode/script/local-bin.ts b/packages/kilo-vscode/script/local-bin.ts index a8d2f669bd..bf85ec1c9a 100644 --- a/packages/kilo-vscode/script/local-bin.ts +++ b/packages/kilo-vscode/script/local-bin.ts @@ -22,6 +22,7 @@ const kiloVscodeDir = join(import.meta.dir, "..") const packagesDir = join(kiloVscodeDir, "..") const opencodeDir = join(packagesDir, "opencode") const coreDir = join(packagesDir, "core") +const indexingDir = join(packagesDir, "kilo-indexing") const targetBinDir = join(kiloVscodeDir, "bin") const binName = process.platform === "win32" ? "kilo.exe" : "kilo" @@ -36,7 +37,8 @@ async function cliSourceHash(): Promise { try { const opencodeResult = await $`git log -1 --format=%H -- .`.cwd(opencodeDir).quiet() const coreResult = await $`git log -1 --format=%H -- .`.cwd(coreDir).quiet() - return `${opencodeResult.text().trim()}-${coreResult.text().trim()}` || null + const indexingResult = await $`git log -1 --format=%H -- .`.cwd(indexingDir).quiet() + return `${opencodeResult.text().trim()}-${coreResult.text().trim()}-${indexingResult.text().trim()}` || null } catch { return null } @@ -46,7 +48,12 @@ async function isDirty(): Promise { try { const opencodeResult = await $`git status --porcelain -- .`.cwd(opencodeDir).quiet() const coreResult = await $`git status --porcelain -- .`.cwd(coreDir).quiet() - return opencodeResult.text().trim().length > 0 || coreResult.text().trim().length > 0 + const indexingResult = await $`git status --porcelain -- .`.cwd(indexingDir).quiet() + return ( + opencodeResult.text().trim().length > 0 || + coreResult.text().trim().length > 0 || + indexingResult.text().trim().length > 0 + ) } catch { return false }