diff --git a/src/services/code-index/__tests__/config-manager.spec.ts b/src/services/code-index/__tests__/config-manager.spec.ts index 6d0e59e827a..2d6e704d76c 100644 --- a/src/services/code-index/__tests__/config-manager.spec.ts +++ b/src/services/code-index/__tests__/config-manager.spec.ts @@ -8,6 +8,17 @@ import { PreviousConfigSnapshot } from "../interfaces/config" // Mock ContextProxy vi.mock("../../../core/config/ContextProxy") +// Mock embeddingModels module +vi.mock("../../../shared/embeddingModels") + +// Import mocked functions +import { getDefaultModelId, getModelDimension, getModelScoreThreshold } from "../../../shared/embeddingModels" + +// Type the mocked functions +const mockedGetDefaultModelId = vi.mocked(getDefaultModelId) +const mockedGetModelDimension = vi.mocked(getModelDimension) +const mockedGetModelScoreThreshold = vi.mocked(getModelScoreThreshold) + describe("CodeIndexConfigManager", () => { let mockContextProxy: any let configManager: CodeIndexConfigManager @@ -339,6 +350,14 @@ describe("CodeIndexConfigManager", () => { }) it("should NOT require restart when models have same dimensions", async () => { + // Mock both models to have same dimension + mockedGetModelDimension.mockImplementation((provider, modelId) => { + if (modelId === "text-embedding-3-small" || modelId === "text-embedding-ada-002") { + return 1536 + } + return undefined + }) + // Initial state with text-embedding-3-small (1536D) mockContextProxy.getGlobalState.mockReturnValue({ codebaseIndexEnabled: true, @@ -794,6 +813,14 @@ describe("CodeIndexConfigManager", () => { }) it("should fall back to model-specific threshold when user setting is undefined", async () => { + // Mock the model score threshold + mockedGetModelScoreThreshold.mockImplementation((provider, modelId) => { + if (provider === "ollama" && modelId === "nomic-embed-code") { + return 0.15 + } + return undefined + }) + mockContextProxy.getGlobalState.mockReturnValue({ codebaseIndexEnabled: true, codebaseIndexQdrantUrl: "http://qdrant.local", @@ -840,6 +867,14 @@ describe("CodeIndexConfigManager", () => { }) it("should use model-specific threshold with openai-compatible provider", async () => { + // Mock the model score threshold + mockedGetModelScoreThreshold.mockImplementation((provider, modelId) => { + if (provider === "openai-compatible" && modelId === "nomic-embed-code") { + return 0.15 + } + return undefined + }) + mockContextProxy.getGlobalState.mockImplementation((key: string) => { if (key === "codebaseIndexConfig") { return { @@ -882,6 +917,14 @@ describe("CodeIndexConfigManager", () => { }) it("should handle priority correctly: user > model > default", async () => { + // Mock the model score threshold + mockedGetModelScoreThreshold.mockImplementation((provider, modelId) => { + if (provider === "ollama" && modelId === "nomic-embed-code") { + return 0.15 + } + return undefined + }) + // Test 1: User setting takes precedence mockContextProxy.getGlobalState.mockReturnValue({ codebaseIndexEnabled: true, @@ -1501,6 +1544,13 @@ describe("CodeIndexConfigManager", () => { }) describe("loadConfiguration", () => { + beforeEach(() => { + // Set default mock behaviors + mockedGetDefaultModelId.mockReturnValue("text-embedding-3-small") + mockedGetModelDimension.mockReturnValue(undefined) + mockedGetModelScoreThreshold.mockReturnValue(undefined) + }) + it("should load configuration and return proper structure", async () => { const mockConfigValues = { codebaseIndexEnabled: true, @@ -1634,5 +1684,131 @@ describe("CodeIndexConfigManager", () => { configManager = new CodeIndexConfigManager(mockContextProxy) expect(configManager.isConfigured()).toBe(false) }) + + describe("currentModelDimension", () => { + beforeEach(() => { + vi.clearAllMocks() + }) + + it("should return model's built-in dimension when available", async () => { + // Mock getModelDimension to return a built-in dimension + mockedGetModelDimension.mockReturnValue(1536) + + mockContextProxy.getGlobalState.mockReturnValue({ + codebaseIndexEnabled: true, + codebaseIndexEmbedderProvider: "openai", + codebaseIndexEmbedderModelId: "text-embedding-3-small", + codebaseIndexEmbedderModelDimension: 2048, // Custom dimension should be ignored + codebaseIndexQdrantUrl: "http://localhost:6333", + }) + mockContextProxy.getSecret.mockImplementation((key: string) => { + if (key === "codeIndexOpenAiKey") return "test-key" + return undefined + }) + + configManager = new CodeIndexConfigManager(mockContextProxy) + await configManager.loadConfiguration() + + // Should return model's built-in dimension, not custom + expect(configManager.currentModelDimension).toBe(1536) + expect(mockedGetModelDimension).toHaveBeenCalledWith("openai", "text-embedding-3-small") + }) + + it("should use custom dimension only when model has no built-in dimension", async () => { + // Mock getModelDimension to return undefined (no built-in dimension) + mockedGetModelDimension.mockReturnValue(undefined) + + mockContextProxy.getGlobalState.mockReturnValue({ + codebaseIndexEnabled: true, + codebaseIndexEmbedderProvider: "openai-compatible", + codebaseIndexEmbedderModelId: "custom-model", + codebaseIndexEmbedderModelDimension: 2048, // Custom dimension should be used + codebaseIndexQdrantUrl: "http://localhost:6333", + }) + mockContextProxy.getSecret.mockImplementation((key: string) => { + if (key === "codebaseIndexOpenAiCompatibleApiKey") return "test-key" + return undefined + }) + + configManager = new CodeIndexConfigManager(mockContextProxy) + await configManager.loadConfiguration() + + // Should use custom dimension as fallback + expect(configManager.currentModelDimension).toBe(2048) + expect(mockedGetModelDimension).toHaveBeenCalledWith("openai-compatible", "custom-model") + }) + + it("should return undefined when neither model dimension nor custom dimension is available", async () => { + // Mock getModelDimension to return undefined + mockedGetModelDimension.mockReturnValue(undefined) + + mockContextProxy.getGlobalState.mockReturnValue({ + codebaseIndexEnabled: true, + codebaseIndexEmbedderProvider: "openai-compatible", + codebaseIndexEmbedderModelId: "unknown-model", + // No custom dimension set + codebaseIndexQdrantUrl: "http://localhost:6333", + }) + mockContextProxy.getSecret.mockImplementation((key: string) => { + if (key === "codebaseIndexOpenAiCompatibleApiKey") return "test-key" + return undefined + }) + + configManager = new CodeIndexConfigManager(mockContextProxy) + await configManager.loadConfiguration() + + // Should return undefined + expect(configManager.currentModelDimension).toBe(undefined) + expect(mockedGetModelDimension).toHaveBeenCalledWith("openai-compatible", "unknown-model") + }) + + it("should use default model ID when modelId is not specified", async () => { + // Mock getDefaultModelId and getModelDimension + mockedGetDefaultModelId.mockReturnValue("text-embedding-3-small") + mockedGetModelDimension.mockReturnValue(1536) + + mockContextProxy.getGlobalState.mockReturnValue({ + codebaseIndexEnabled: true, + codebaseIndexEmbedderProvider: "openai", + // No modelId specified + codebaseIndexQdrantUrl: "http://localhost:6333", + }) + mockContextProxy.getSecret.mockImplementation((key: string) => { + if (key === "codeIndexOpenAiKey") return "test-key" + return undefined + }) + + configManager = new CodeIndexConfigManager(mockContextProxy) + await configManager.loadConfiguration() + + // Should use default model ID + expect(configManager.currentModelDimension).toBe(1536) + expect(mockedGetDefaultModelId).toHaveBeenCalledWith("openai") + expect(mockedGetModelDimension).toHaveBeenCalledWith("openai", "text-embedding-3-small") + }) + + it("should ignore invalid custom dimension (0 or negative)", async () => { + // Mock getModelDimension to return undefined + mockedGetModelDimension.mockReturnValue(undefined) + + mockContextProxy.getGlobalState.mockReturnValue({ + codebaseIndexEnabled: true, + codebaseIndexEmbedderProvider: "openai-compatible", + codebaseIndexEmbedderModelId: "custom-model", + codebaseIndexEmbedderModelDimension: 0, // Invalid dimension + codebaseIndexQdrantUrl: "http://localhost:6333", + }) + mockContextProxy.getSecret.mockImplementation((key: string) => { + if (key === "codebaseIndexOpenAiCompatibleApiKey") return "test-key" + return undefined + }) + + configManager = new CodeIndexConfigManager(mockContextProxy) + await configManager.loadConfiguration() + + // Should return undefined since custom dimension is invalid + expect(configManager.currentModelDimension).toBe(undefined) + }) + }) }) }) diff --git a/src/services/code-index/__tests__/service-factory.spec.ts b/src/services/code-index/__tests__/service-factory.spec.ts index 373b0e3e827..1d8f7ba4786 100644 --- a/src/services/code-index/__tests__/service-factory.spec.ts +++ b/src/services/code-index/__tests__/service-factory.spec.ts @@ -420,10 +420,42 @@ describe("CodeIndexServiceFactory", () => { ) }) - it("should prioritize manual modelDimension over getModelDimension for OpenAI Compatible provider", () => { + it("should prioritize getModelDimension over manual modelDimension for OpenAI Compatible provider", () => { // Arrange const testModelId = "custom-model" const manualDimension = 1024 + const modelDimension = 768 + const testConfig = { + embedderProvider: "openai-compatible", + modelId: testModelId, + modelDimension: manualDimension, // This should be ignored when model has built-in dimension + openAiCompatibleOptions: { + baseUrl: "https://api.example.com/v1", + apiKey: "test-api-key", + }, + qdrantUrl: "http://localhost:6333", + qdrantApiKey: "test-key", + } + mockConfigManager.getConfig.mockReturnValue(testConfig as any) + mockGetModelDimension.mockReturnValue(modelDimension) // This should be used + + // Act + factory.createVectorStore() + + // Assert + expect(mockGetModelDimension).toHaveBeenCalledWith("openai-compatible", testModelId) + expect(MockedQdrantVectorStore).toHaveBeenCalledWith( + "/test/workspace", + "http://localhost:6333", + modelDimension, // Should use model's built-in dimension, not manual + "test-key", + ) + }) + + it("should use manual modelDimension only when model has no built-in dimension", () => { + // Arrange + const testModelId = "unknown-model" + const manualDimension = 1024 const testConfig = { embedderProvider: "openai-compatible", modelId: testModelId, @@ -436,17 +468,17 @@ describe("CodeIndexServiceFactory", () => { qdrantApiKey: "test-key", } mockConfigManager.getConfig.mockReturnValue(testConfig as any) - mockGetModelDimension.mockReturnValue(768) // This should be ignored + mockGetModelDimension.mockReturnValue(undefined) // Model has no built-in dimension // Act factory.createVectorStore() // Assert - expect(mockGetModelDimension).not.toHaveBeenCalled() + expect(mockGetModelDimension).toHaveBeenCalledWith("openai-compatible", testModelId) expect(MockedQdrantVectorStore).toHaveBeenCalledWith( "/test/workspace", "http://localhost:6333", - manualDimension, + manualDimension, // Should use manual dimension as fallback "test-key", ) }) diff --git a/src/services/code-index/config-manager.ts b/src/services/code-index/config-manager.ts index f022aec7806..9958f456c3e 100644 --- a/src/services/code-index/config-manager.ts +++ b/src/services/code-index/config-manager.ts @@ -398,10 +398,19 @@ export class CodeIndexConfigManager { /** * Gets the current model dimension being used for embeddings. - * Returns the explicitly configured dimension or undefined if not set. + * Returns the model's built-in dimension if available, otherwise falls back to custom dimension. */ public get currentModelDimension(): number | undefined { - return this.modelDimension + // First try to get the model-specific dimension + const modelId = this.modelId ?? getDefaultModelId(this.embedderProvider) + const modelDimension = getModelDimension(this.embedderProvider, modelId) + + // Only use custom dimension if model doesn't have a built-in dimension + if (!modelDimension && this.modelDimension && this.modelDimension > 0) { + return this.modelDimension + } + + return modelDimension } /** diff --git a/src/services/code-index/service-factory.ts b/src/services/code-index/service-factory.ts index ec8b1e7ade7..b7951db7acf 100644 --- a/src/services/code-index/service-factory.ts +++ b/src/services/code-index/service-factory.ts @@ -108,12 +108,12 @@ export class CodeIndexServiceFactory { let vectorSize: number | undefined - // First check if a manual dimension is provided (works for all providers) - if (config.modelDimension && config.modelDimension > 0) { + // First try to get the model-specific dimension from profiles + vectorSize = getModelDimension(provider, modelId) + + // Only use manual dimension if model doesn't have a built-in dimension + if (!vectorSize && config.modelDimension && config.modelDimension > 0) { vectorSize = config.modelDimension - } else { - // Fall back to model-specific dimension from profiles - vectorSize = getModelDimension(provider, modelId) } if (vectorSize === undefined || vectorSize <= 0) {