Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion packages/kilo-indexing/src/indexing/manager.ts
Original file line number Diff line number Diff line change
Expand Up @@ -549,7 +549,7 @@ export class CodeIndexManager {
const ignoreInstance = await loadIgnore(this.workspacePath)
const config = this._configManager!.getConfig()
const baseline = prepared ?? (await this.createBaseline(factory))
const { embedder, vectorStore, scanner, fileWatcher } = factory.createServices(this._cacheManager!, ignoreInstance)
const { embedder, vectorStore, scanner, fileWatcher } = await factory.createServices(this._cacheManager!, ignoreInstance)
fileWatcher.setOverlay?.(baseline?.overlay)
log.info("created indexing services", {
workspacePath: this.workspacePath,
Expand Down
52 changes: 39 additions & 13 deletions packages/kilo-indexing/src/indexing/service-factory.ts
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import type { Ignore } from "ignore"
import path from "path"

import { getDefaultModelId } from "./model-registry"
import { getDefaultModelId, getModelDimension } from "./model-registry"
import { resolveEmbeddingProfile } from "./embedding-profile"

import { OpenAiEmbedder } from "./embedders/openai"
Expand Down Expand Up @@ -61,9 +61,10 @@ export class CodeIndexServiceFactory {
}
}

public createEmbedder(): IEmbedder {
public createEmbedder(dimensionOverride?: number): IEmbedder {
const config = this.configManager.getConfig()
const provider = config.embedderProvider
const dimension = dimensionOverride ?? config.modelDimension

if (provider === "kilo") {
if (!config.kiloOptions?.apiKey) throw new Error("Kilo API key is required for embedding.")
Expand All @@ -73,7 +74,7 @@ export class CodeIndexServiceFactory {
baseUrl: config.kiloOptions.baseUrl,
organizationId: config.kiloOptions.organizationId,
modelId: config.modelId,
dimensions: config.modelDimension,
dimensions: dimension,
})
}
if (provider === "openai") {
Expand All @@ -82,7 +83,7 @@ export class CodeIndexServiceFactory {
}
if (provider === "ollama") {
if (!config.ollamaOptions?.baseUrl) throw new Error("Ollama base URL is required for embedding.")
return new CodeIndexOllamaEmbedder(config.ollamaOptions.baseUrl, config.modelId, config.modelDimension)
return new CodeIndexOllamaEmbedder(config.ollamaOptions.baseUrl, config.modelId, dimension)
}
if (provider === "openai-compatible") {
if (!config.openAiCompatibleOptions?.baseUrl) throw new Error("OpenAI-compatible base URL is required.")
Expand Down Expand Up @@ -116,7 +117,7 @@ export class CodeIndexServiceFactory {
config.modelId,
undefined,
config.openRouterOptions.specificProvider,
config.modelDimension,
dimension,
)
}
if (provider === "voyage") {
Expand Down Expand Up @@ -169,9 +170,9 @@ export class CodeIndexServiceFactory {
}
}

public createVectorStore(workspacePath = this.workspacePath): IVectorStore {
public createVectorStore(workspacePath = this.workspacePath, dimensionOverride?: number): IVectorStore {
const config = this.configManager.getConfig()
const profile = resolveEmbeddingProfile(config.embedderProvider, config.modelId, config.modelDimension)
const profile = resolveEmbeddingProfile(config.embedderProvider, config.modelId, dimensionOverride ?? config.modelDimension)

if (!profile || profile.dimension <= 0) {
throw new Error(
Expand Down Expand Up @@ -246,38 +247,63 @@ export class CodeIndexServiceFactory {
)
}

public createServices(
public async detectEmbeddingDimension(embedder: IEmbedder): Promise<number> {
log.info("detecting embedding dimension", { provider: embedder.embedderInfo.name })
const response = await embedder.createEmbeddings(["test"])
const dimension = response.embeddings[0]?.length
if (!dimension || dimension <= 0) {
throw new Error(`Failed to detect embedding dimension: invalid response from ${embedder.embedderInfo.name}`)
}
log.info("detected embedding dimension", { provider: embedder.embedderInfo.name, dimension })
return dimension
}

public async createServices(
cacheManager: CacheManager,
ignoreInstance: Ignore,
): {
): Promise<{
embedder: IEmbedder
vectorStore: IVectorStore
parser: ICodeParser
scanner: DirectoryScanner
fileWatcher: IFileWatcher
} {
}> {
if (!this.configManager.isFeatureConfigured) {
throw new Error("Code indexing is not configured. Save your settings to start indexing.")
}

const config = this.configManager.getConfig()
const modelId = config.modelId ?? getDefaultModelId(config.embedderProvider)
const registryDimension = getModelDimension(config.embedderProvider, modelId)

log.info("creating indexing services", {
workspacePath: this.workspacePath,
provider: config.embedderProvider,
vectorStore: config.vectorStoreProvider,
model: config.modelId ?? getDefaultModelId(config.embedderProvider),
model: modelId,
configured: config.isConfigured,
})

const embedder = this.createEmbedder()
const vectorStore = this.createVectorStore()
let effectiveDimension: number | undefined
if (config.modelDimension) {
effectiveDimension = config.modelDimension
} else if (registryDimension) {
effectiveDimension = registryDimension
} else {
const detectionEmbedder = this.createEmbedder()
effectiveDimension = await this.detectEmbeddingDimension(detectionEmbedder)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

WARNING: Runtime detection still does not cover baseline-backed worktrees

This fallback runs after CodeIndexManager._recreateServices() has already called createBaseline(factory), and createBaseline() still builds its vector store with no detected override. For an unknown custom model in a worktree session, initialization will still fail with Cannot determine vector dimension... before execution ever reaches this branch.


Reply with @kilocode-bot fix it to have Kilo Code address this issue.

}

const embedder = this.createEmbedder(effectiveDimension)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

WARNING: Detected dimensions are not applied to openai-compatible embeddings

effectiveDimension is passed into createEmbedder() here, but the openai-compatible branch above still constructs OpenAICompatibleEmbedder without its dimensions option. That means the vector store is created at the detected size while later embedding requests still use the endpoint's default size, which will trip a vector-length mismatch as soon as indexing starts for unknown custom models.


Reply with @kilocode-bot fix it to have Kilo Code address this issue.

const vectorStore = this.createVectorStore(this.workspacePath, effectiveDimension)
const parser = codeParser
const scanner = this.createDirectoryScanner(embedder, vectorStore, parser, ignoreInstance)
const fileWatcher = this.createFileWatcher(embedder, vectorStore, cacheManager, ignoreInstance)

log.info("indexing services created", {
workspacePath: this.workspacePath,
provider: embedder.embedderInfo.name,
dimension: effectiveDimension,
})

return { embedder, vectorStore, parser, scanner, fileWatcher }
Expand Down
Original file line number Diff line number Diff line change
@@ -1,11 +1,13 @@
import { describe, expect, test, mock, beforeEach } from "bun:test"
import path from "path"
import type { Ignore } from "ignore"
import { mockEmbeddingsCreate, openAIMockFactory, setOpenAIConstructorHook } from "./embedders/__helpers__/openai-mock"

mock.module("openai", openAIMockFactory)
import { CodeIndexServiceFactory } from "../../../src/indexing/service-factory"
import { CodeIndexConfigManager } from "../../../src/indexing/config-manager"
import { CacheManager } from "../../../src/indexing/cache-manager"
import type { IEmbedder } from "../../../src/indexing/interfaces"

const workspacePath = "/tmp/ws"
const cacheDirectory = "/tmp/cache"
Expand Down Expand Up @@ -221,4 +223,70 @@ describe("CodeIndexServiceFactory", () => {
dimensions: 1024,
})
})

test("uses registry default for known model without explicit dimension", async () => {
const factory = createFactory({
embedderProvider: "openai",
modelId: "text-embedding-3-small",
modelDimension: undefined,
})

let embedderCreatedWithDimension: number | undefined
const originalCreateEmbedder = factory.createEmbedder.bind(factory)
factory.createEmbedder = function (override?: number) {
embedderCreatedWithDimension = override
return originalCreateEmbedder(override)
}

const services = await factory.createServices({} as CacheManager, {} as Ignore)

expect(embedderCreatedWithDimension).toBe(1536)
expect((services.vectorStore as unknown as { vectorSize: number }).vectorSize).toBe(1536)
})

test("uses explicit dimension over registry default", async () => {
const factory = createFactory({
embedderProvider: "openai",
modelId: "text-embedding-3-small",
modelDimension: 12346789,
})

let embedderCreatedWithDimension: number | undefined
const originalCreateEmbedder = factory.createEmbedder.bind(factory)
factory.createEmbedder = function (override?: number) {
embedderCreatedWithDimension = override
return originalCreateEmbedder(override)
}

const services = await factory.createServices({} as CacheManager, {} as Ignore)

expect(embedderCreatedWithDimension).toBe(12346789)
expect((services.vectorStore as unknown as { vectorSize: number }).vectorSize).toBe(12346789)
})

test("detects dimension for unknown model without explicit dimension", async () => {
const factory = createFactory({
embedderProvider: "openai-compatible",
modelId: "custom/unknown-model",
modelDimension: undefined,
openAiCompatibleBaseUrl: "http://localhost:1234/v1",
})

let detectionCalled = false
const originalDetect = factory.detectEmbeddingDimension.bind(factory)
;(factory as any).detectEmbeddingDimension = async function (embedder: IEmbedder) {
detectionCalled = true
return originalDetect(embedder)
}

mockEmbeddingsCreate.mockResolvedValue({
data: [{ embedding: [0.1, 0.2, 0.3, 0.4, 0.5] }],
usage: { prompt_tokens: 1, total_tokens: 1 },
})

const services = await factory.createServices({} as CacheManager, {} as Ignore)

expect(detectionCalled).toBe(true)
expect((services.vectorStore as unknown as { vectorSize: number }).vectorSize).toBe(5)
})
})