fix(indexing): support OpenRouter Gemini embeddings
This commit is contained in:
@@ -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,
|
||||
})
|
||||
})
|
||||
})
|
||||
|
||||
@@ -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]
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user