fix(indexing): support OpenRouter Gemini embeddings

This commit is contained in:
marius-kilocode
2026-05-21 13:06:36 +02:00
parent d31b25eeab
commit 205e22ee46
10 changed files with 123 additions and 31 deletions
+5
View File
@@ -0,0 +1,5 @@
---
"@kilocode/kilo-indexing": patch
---
Support OpenRouter Gemini embedding preview indexing and honor configured embedding dimensions when sizing vector stores.
@@ -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 {
@@ -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",
@@ -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,
@@ -48,6 +48,7 @@ const profiles: Record<string, Record<string, ModelProfile>> = {
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": {
@@ -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" }))
@@ -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,
},
{
@@ -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,
})
})
})
+3
View File
@@ -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]
}
+9 -2
View File
@@ -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<string | null> {
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<boolean> {
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
}