diff --git a/packages/types/src/codebase-index.ts b/packages/types/src/codebase-index.ts index be7778f538..d08a65e854 100644 --- a/packages/types/src/codebase-index.ts +++ b/packages/types/src/codebase-index.ts @@ -36,6 +36,8 @@ export const codebaseIndexConfigSchema = z.object({ // OpenAI Compatible specific fields codebaseIndexOpenAiCompatibleBaseUrl: z.string().optional(), codebaseIndexOpenAiCompatibleModelDimension: z.number().optional(), + // Custom query instruction for instruction-aware embedding models + codebaseIndexCustomQueryInstruction: z.string().optional(), }) export type CodebaseIndexConfig = z.infer diff --git a/src/services/code-index/__tests__/config-manager.spec.ts b/src/services/code-index/__tests__/config-manager.spec.ts index 9fc096ba74..4ee57caf9d 100644 --- a/src/services/code-index/__tests__/config-manager.spec.ts +++ b/src/services/code-index/__tests__/config-manager.spec.ts @@ -141,6 +141,42 @@ describe("CodeIndexConfigManager", () => { }) }) + it("should load custom query instruction from globalState", async () => { + const mockGlobalState = { + codebaseIndexEnabled: true, + codebaseIndexQdrantUrl: "http://qdrant.local", + codebaseIndexEmbedderProvider: "openai", + codebaseIndexEmbedderModelId: "text-embedding-3-large", + codebaseIndexCustomQueryInstruction: "Represent this code for searching:", + } + mockContextProxy.getGlobalState.mockReturnValue(mockGlobalState) + + // Mock both sync and async secret access + setupSecretMocks({ + codeIndexOpenAiKey: "test-openai-key", + codeIndexQdrantApiKey: "test-qdrant-key", + }) + + const result = await configManager.loadConfiguration() + + expect(result.currentConfig).toEqual({ + isConfigured: true, + embedderProvider: "openai", + modelId: "text-embedding-3-large", + modelDimension: undefined, + openAiOptions: { openAiNativeApiKey: "test-openai-key" }, + ollamaOptions: { ollamaBaseUrl: undefined }, + geminiOptions: undefined, + mistralOptions: undefined, + openAiCompatibleOptions: undefined, + vercelAiGatewayOptions: undefined, + qdrantUrl: "http://qdrant.local", + qdrantApiKey: "test-qdrant-key", + searchMinScore: 0.4, + customQueryInstruction: "Represent this code for searching:", + }) + }) + it("should load OpenAI Compatible configuration from globalState and secrets", async () => { const mockGlobalState = { codebaseIndexEnabled: true, @@ -1808,5 +1844,104 @@ describe("CodeIndexConfigManager", () => { expect(configManager.currentModelDimension).toBe(undefined) }) }) + + describe("currentCustomQueryInstruction", () => { + it("should return custom query instruction when set", async () => { + mockContextProxy.getGlobalState.mockReturnValue({ + codebaseIndexEnabled: true, + codebaseIndexQdrantUrl: "http://localhost:6333", + codebaseIndexEmbedderProvider: "openai", + codebaseIndexEmbedderModelId: "text-embedding-3-small", + codebaseIndexCustomQueryInstruction: "Represent this code for retrieval:", + }) + mockContextProxy.getSecret.mockImplementation((key: string) => { + if (key === "codeIndexOpenAiKey") return "test-key" + return undefined + }) + + configManager = new CodeIndexConfigManager(mockContextProxy) + await configManager.loadConfiguration() + + expect(configManager.currentCustomQueryInstruction).toBe("Represent this code for retrieval:") + }) + + it("should return undefined when custom query instruction is not set", async () => { + mockContextProxy.getGlobalState.mockReturnValue({ + codebaseIndexEnabled: true, + codebaseIndexQdrantUrl: "http://localhost:6333", + codebaseIndexEmbedderProvider: "openai", + codebaseIndexEmbedderModelId: "text-embedding-3-small", + // No custom query instruction + }) + mockContextProxy.getSecret.mockImplementation((key: string) => { + if (key === "codeIndexOpenAiKey") return "test-key" + return undefined + }) + + configManager = new CodeIndexConfigManager(mockContextProxy) + await configManager.loadConfiguration() + + expect(configManager.currentCustomQueryInstruction).toBe(undefined) + }) + + it("should return empty string when custom query instruction is explicitly empty", async () => { + mockContextProxy.getGlobalState.mockReturnValue({ + codebaseIndexEnabled: true, + codebaseIndexQdrantUrl: "http://localhost:6333", + codebaseIndexEmbedderProvider: "openai", + codebaseIndexEmbedderModelId: "text-embedding-3-small", + codebaseIndexCustomQueryInstruction: "", + }) + mockContextProxy.getSecret.mockImplementation((key: string) => { + if (key === "codeIndexOpenAiKey") return "test-key" + return undefined + }) + + configManager = new CodeIndexConfigManager(mockContextProxy) + await configManager.loadConfiguration() + + expect(configManager.currentCustomQueryInstruction).toBe("") + }) + + it("should handle custom query instruction with different providers", async () => { + // Test with Ollama + mockContextProxy.getGlobalState.mockReturnValue({ + codebaseIndexEnabled: true, + codebaseIndexQdrantUrl: "http://localhost:6333", + codebaseIndexEmbedderProvider: "ollama", + codebaseIndexEmbedderBaseUrl: "http://localhost:11434", + codebaseIndexEmbedderModelId: "nomic-embed-text", + codebaseIndexCustomQueryInstruction: "search_query: ", + }) + + configManager = new CodeIndexConfigManager(mockContextProxy) + await configManager.loadConfiguration() + + expect(configManager.currentCustomQueryInstruction).toBe("search_query: ") + + // Test with OpenAI Compatible + mockContextProxy.getGlobalState.mockImplementation((key: string) => { + if (key === "codebaseIndexConfig") { + return { + codebaseIndexEnabled: true, + codebaseIndexQdrantUrl: "http://localhost:6333", + codebaseIndexEmbedderProvider: "openai-compatible", + codebaseIndexEmbedderModelId: "custom-model", + codebaseIndexOpenAiCompatibleBaseUrl: "https://api.example.com/v1", + codebaseIndexCustomQueryInstruction: "query: ", + } + } + return undefined + }) + setupSecretMocks({ + codebaseIndexOpenAiCompatibleApiKey: "test-key", + }) + + const newManager = new CodeIndexConfigManager(mockContextProxy) + await newManager.loadConfiguration() + + expect(newManager.currentCustomQueryInstruction).toBe("query: ") + }) + }) }) }) diff --git a/src/services/code-index/__tests__/service-factory.spec.ts b/src/services/code-index/__tests__/service-factory.spec.ts index 1d8f7ba478..6c59dfd852 100644 --- a/src/services/code-index/__tests__/service-factory.spec.ts +++ b/src/services/code-index/__tests__/service-factory.spec.ts @@ -194,6 +194,8 @@ describe("CodeIndexServiceFactory", () => { "https://api.example.com/v1", "test-api-key", testModelId, + undefined, // maxItemTokens + undefined, // customQueryInstruction ) }) @@ -216,7 +218,9 @@ describe("CodeIndexServiceFactory", () => { expect(MockedOpenAICompatibleEmbedder).toHaveBeenCalledWith( "https://api.example.com/v1", "test-api-key", - undefined, + undefined, // modelId + undefined, // maxItemTokens + undefined, // customQueryInstruction ) }) @@ -279,7 +283,7 @@ describe("CodeIndexServiceFactory", () => { factory.createEmbedder() // Assert - expect(MockedGeminiEmbedder).toHaveBeenCalledWith("test-gemini-api-key", undefined) + expect(MockedGeminiEmbedder).toHaveBeenCalledWith("test-gemini-api-key", undefined, undefined) }) it("should create GeminiEmbedder with specified modelId", () => { @@ -297,7 +301,7 @@ describe("CodeIndexServiceFactory", () => { factory.createEmbedder() // Assert - expect(MockedGeminiEmbedder).toHaveBeenCalledWith("test-gemini-api-key", "text-embedding-004") + expect(MockedGeminiEmbedder).toHaveBeenCalledWith("test-gemini-api-key", "text-embedding-004", undefined) }) it("should throw error when Gemini API key is missing", () => { diff --git a/src/services/code-index/config-manager.ts b/src/services/code-index/config-manager.ts index 2c0e8bb5c9..7b2179074b 100644 --- a/src/services/code-index/config-manager.ts +++ b/src/services/code-index/config-manager.ts @@ -24,6 +24,7 @@ export class CodeIndexConfigManager { private qdrantApiKey?: string private searchMinScore?: number private searchMaxResults?: number + private customQueryInstruction?: string constructor(private readonly contextProxy: ContextProxy) { // Initialize with current configuration to avoid false restart triggers @@ -61,6 +62,7 @@ export class CodeIndexConfigManager { codebaseIndexEmbedderModelId, codebaseIndexSearchMinScore, codebaseIndexSearchMaxResults, + codebaseIndexCustomQueryInstruction, } = codebaseIndexConfig const openAiKey = this.contextProxy?.getSecret("codeIndexOpenAiKey") ?? "" @@ -78,6 +80,7 @@ export class CodeIndexConfigManager { this.qdrantApiKey = qdrantApiKey ?? "" this.searchMinScore = codebaseIndexSearchMinScore this.searchMaxResults = codebaseIndexSearchMaxResults + this.customQueryInstruction = codebaseIndexCustomQueryInstruction // Validate and set model dimension const rawDimension = codebaseIndexConfig.codebaseIndexEmbedderModelDimension @@ -150,6 +153,7 @@ export class CodeIndexConfigManager { qdrantUrl?: string qdrantApiKey?: string searchMinScore?: number + customQueryInstruction?: string } requiresRestart: boolean }> { @@ -195,6 +199,7 @@ export class CodeIndexConfigManager { qdrantUrl: this.qdrantUrl, qdrantApiKey: this.qdrantApiKey, searchMinScore: this.currentSearchMinScore, + customQueryInstruction: this.customQueryInstruction, }, requiresRestart, } @@ -399,6 +404,7 @@ export class CodeIndexConfigManager { qdrantApiKey: this.qdrantApiKey, searchMinScore: this.currentSearchMinScore, searchMaxResults: this.currentSearchMaxResults, + customQueryInstruction: this.customQueryInstruction, } } @@ -480,4 +486,12 @@ export class CodeIndexConfigManager { public get currentSearchMaxResults(): number { return this.searchMaxResults ?? DEFAULT_MAX_SEARCH_RESULTS } + + /** + * Gets the configured custom query instruction for instruction-aware embedding models. + * Returns the user-configured instruction or undefined if not set. + */ + public get currentCustomQueryInstruction(): string | undefined { + return this.customQueryInstruction + } } diff --git a/src/services/code-index/embedders/__tests__/custom-query-instruction.spec.ts b/src/services/code-index/embedders/__tests__/custom-query-instruction.spec.ts new file mode 100644 index 0000000000..25dc70f03a --- /dev/null +++ b/src/services/code-index/embedders/__tests__/custom-query-instruction.spec.ts @@ -0,0 +1,386 @@ +// Test suite for custom query instruction functionality across all embedders +import { OpenAiEmbedder } from "../openai" +import { CodeIndexOllamaEmbedder } from "../ollama" +import { OpenAICompatibleEmbedder } from "../openai-compatible" +import { GeminiEmbedder } from "../gemini" +import { MistralEmbedder } from "../mistral" +import { VercelAiGatewayEmbedder } from "../vercel-ai-gateway" +import { getModelQueryPrefix } from "../../../../shared/embeddingModels" + +// Mock the embeddingModels module +vi.mock("../../../../shared/embeddingModels", () => ({ + getModelQueryPrefix: vi.fn(), + getModelDimension: vi.fn(), + getDefaultModelId: vi.fn(), +})) + +const mockedGetModelQueryPrefix = vi.mocked(getModelQueryPrefix) + +describe("Custom Query Instruction Support", () => { + beforeEach(() => { + vi.clearAllMocks() + // Setup default mock behavior + mockedGetModelQueryPrefix.mockReturnValue(undefined) + }) + + describe("OpenAI Embedder", () => { + it("should use custom query instruction when provided", async () => { + const customInstruction = "Represent this code for searching:" + const embedder = new OpenAiEmbedder({ + openAiNativeApiKey: "test-api-key", + openAiEmbeddingModelId: "text-embedding-3-small", + customQueryInstruction: customInstruction, + }) + + // Mock the OpenAI API call + const mockFetch = vi.fn().mockResolvedValue({ + ok: true, + json: async () => ({ + data: [{ embedding: new Array(1536).fill(0.1) }], + }), + }) + global.fetch = mockFetch + + await embedder.createEmbeddings(["test query"]) + + // Check that the custom instruction was used + const fetchCall = mockFetch.mock.calls[0] + const requestBody = JSON.parse(fetchCall[1].body) + expect(requestBody.input).toBe(customInstruction + "test query") + expect(mockedGetModelQueryPrefix).not.toHaveBeenCalled() + }) + + it("should fall back to model-specific prefix when no custom instruction", async () => { + const modelPrefix = "search_query: " + mockedGetModelQueryPrefix.mockReturnValue(modelPrefix) + + const embedder = new OpenAiEmbedder({ + openAiNativeApiKey: "test-api-key", + openAiEmbeddingModelId: "text-embedding-3-small", + // No custom instruction + }) + + // Mock the OpenAI API call + const mockFetch = vi.fn().mockResolvedValue({ + ok: true, + json: async () => ({ + data: [{ embedding: new Array(1536).fill(0.1) }], + }), + }) + global.fetch = mockFetch + + await embedder.createEmbeddings(["test query"]) + + // Check that the model prefix was used + const fetchCall = mockFetch.mock.calls[0] + const requestBody = JSON.parse(fetchCall[1].body) + expect(requestBody.input).toBe(modelPrefix + "test query") + expect(mockedGetModelQueryPrefix).toHaveBeenCalledWith("openai", "text-embedding-3-small") + }) + + it("should use empty string custom instruction when explicitly set", async () => { + const embedder = new OpenAiEmbedder({ + openAiNativeApiKey: "test-api-key", + openAiEmbeddingModelId: "text-embedding-3-small", + customQueryInstruction: "", // Explicitly empty custom instruction + }) + + // Mock the OpenAI API call + const mockFetch = vi.fn().mockResolvedValue({ + ok: true, + json: async () => ({ + data: [{ embedding: new Array(1536).fill(0.1) }], + }), + }) + global.fetch = mockFetch + + await embedder.createEmbeddings(["test query"]) + + // Check that no prefix was added + const fetchCall = mockFetch.mock.calls[0] + const requestBody = JSON.parse(fetchCall[1].body) + expect(requestBody.input).toBe("test query") + expect(mockedGetModelQueryPrefix).not.toHaveBeenCalled() + }) + }) + + describe("Ollama Embedder", () => { + it("should use custom query instruction when provided", async () => { + const customInstruction = "query: " + const embedder = new CodeIndexOllamaEmbedder({ + ollamaBaseUrl: "http://localhost:11434", + ollamaModelId: "nomic-embed-text", + customQueryInstruction: customInstruction, + }) + + // Mock the Ollama API call + const mockFetch = vi.fn().mockResolvedValue({ + ok: true, + json: async () => ({ + embeddings: [new Array(768).fill(0.1)], + }), + }) + global.fetch = mockFetch + + await embedder.createEmbeddings(["test query"]) + + // Check that the custom instruction was used + const fetchCall = mockFetch.mock.calls[0] + const requestBody = JSON.parse(fetchCall[1].body) + expect(requestBody.input).toBe(customInstruction + "test query") + expect(mockedGetModelQueryPrefix).not.toHaveBeenCalled() + }) + + it("should fall back to model-specific prefix when no custom instruction", async () => { + const modelPrefix = "search_query: " + mockedGetModelQueryPrefix.mockReturnValue(modelPrefix) + + const embedder = new CodeIndexOllamaEmbedder({ + ollamaBaseUrl: "http://localhost:11434", + ollamaModelId: "nomic-embed-code", + // No custom instruction + }) + + // Mock the Ollama API call + const mockFetch = vi.fn().mockResolvedValue({ + ok: true, + json: async () => ({ + embeddings: [new Array(768).fill(0.1)], + }), + }) + global.fetch = mockFetch + + await embedder.createEmbeddings(["test query"]) + + // Check that the model prefix was used + const fetchCall = mockFetch.mock.calls[0] + const requestBody = JSON.parse(fetchCall[1].body) + expect(requestBody.input).toBe(modelPrefix + "test query") + expect(mockedGetModelQueryPrefix).toHaveBeenCalledWith("ollama", "nomic-embed-code") + }) + }) + + describe("OpenAI Compatible Embedder", () => { + it("should use custom query instruction when provided", async () => { + const customInstruction = "Represent for retrieval: " + const embedder = new OpenAICompatibleEmbedder( + "https://api.example.com", + "custom-model", + "test-api-key", + 1024, // dimension + customInstruction, + ) + + // Mock the API call + const mockFetch = vi.fn().mockResolvedValue({ + ok: true, + json: async () => ({ + data: [{ embedding: new Array(1024).fill(0.1) }], + }), + }) + global.fetch = mockFetch + + await embedder.createEmbeddings(["test query"]) + + // Check that the custom instruction was used + const fetchCall = mockFetch.mock.calls[0] + const requestBody = JSON.parse(fetchCall[1].body) + expect(requestBody.input).toBe(customInstruction + "test query") + expect(mockedGetModelQueryPrefix).not.toHaveBeenCalled() + }) + + it("should fall back to model-specific prefix when no custom instruction", async () => { + const modelPrefix = "search_document: " + mockedGetModelQueryPrefix.mockReturnValue(modelPrefix) + + const embedder = new OpenAICompatibleEmbedder( + "https://api.example.com", + "custom-model", + "test-api-key", + 1024, // dimension + undefined, // No custom instruction + ) + + // Mock the API call + const mockFetch = vi.fn().mockResolvedValue({ + ok: true, + json: async () => ({ + data: [{ embedding: new Array(1024).fill(0.1) }], + }), + }) + global.fetch = mockFetch + + await embedder.createEmbeddings(["test query"]) + + // Check that the model prefix was used + const fetchCall = mockFetch.mock.calls[0] + const requestBody = JSON.parse(fetchCall[1].body) + expect(requestBody.input).toBe(modelPrefix + "test query") + expect(mockedGetModelQueryPrefix).toHaveBeenCalledWith("openai-compatible", "custom-model") + }) + }) + + describe("Gemini Embedder", () => { + it("should use custom query instruction when provided", async () => { + const customInstruction = "Code search query: " + const embedder = new GeminiEmbedder("test-api-key", "text-embedding-004", customInstruction) + + // Mock the Gemini API call + const mockFetch = vi.fn().mockResolvedValue({ + ok: true, + json: async () => ({ + embedding: { + values: new Array(768).fill(0.1), + }, + }), + }) + global.fetch = mockFetch + + await embedder.createEmbeddings(["test query"]) + + // Check that the custom instruction was used + const fetchCall = mockFetch.mock.calls[0] + const requestBody = JSON.parse(fetchCall[1].body) + expect(requestBody.content.parts[0].text).toBe(customInstruction + "test query") + expect(mockedGetModelQueryPrefix).not.toHaveBeenCalled() + }) + + it("should fall back to model-specific prefix when no custom instruction", async () => { + const modelPrefix = "retrieval: " + mockedGetModelQueryPrefix.mockReturnValue(modelPrefix) + + const embedder = new GeminiEmbedder( + "test-api-key", + "text-embedding-004", + undefined, // No custom instruction + ) + + // Mock the Gemini API call + const mockFetch = vi.fn().mockResolvedValue({ + ok: true, + json: async () => ({ + embedding: { + values: new Array(768).fill(0.1), + }, + }), + }) + global.fetch = mockFetch + + await embedder.createEmbeddings(["test query"]) + + // Check that the model prefix was used + const fetchCall = mockFetch.mock.calls[0] + const requestBody = JSON.parse(fetchCall[1].body) + expect(requestBody.content.parts[0].text).toBe(modelPrefix + "test query") + expect(mockedGetModelQueryPrefix).toHaveBeenCalledWith("gemini", "text-embedding-004") + }) + }) + + describe("Mistral Embedder", () => { + it("should use custom query instruction when provided", async () => { + const customInstruction = "Embedding query: " + const embedder = new MistralEmbedder("test-api-key", "mistral-embed", customInstruction) + + // Mock the Mistral API call + const mockFetch = vi.fn().mockResolvedValue({ + ok: true, + json: async () => ({ + data: [{ embedding: new Array(1024).fill(0.1) }], + }), + }) + global.fetch = mockFetch + + await embedder.createEmbeddings(["test query"]) + + // Check that the custom instruction was used + const fetchCall = mockFetch.mock.calls[0] + const requestBody = JSON.parse(fetchCall[1].body) + expect(requestBody.input[0]).toBe(customInstruction + "test query") + expect(mockedGetModelQueryPrefix).not.toHaveBeenCalled() + }) + + it("should fall back to model-specific prefix when no custom instruction", async () => { + const modelPrefix = "search_query: " + mockedGetModelQueryPrefix.mockReturnValue(modelPrefix) + + const embedder = new MistralEmbedder( + "test-api-key", + "mistral-embed", + undefined, // No custom instruction + ) + + // Mock the Mistral API call + const mockFetch = vi.fn().mockResolvedValue({ + ok: true, + json: async () => ({ + data: [{ embedding: new Array(1024).fill(0.1) }], + }), + }) + global.fetch = mockFetch + + await embedder.createEmbeddings(["test query"]) + + // Check that the model prefix was used + const fetchCall = mockFetch.mock.calls[0] + const requestBody = JSON.parse(fetchCall[1].body) + expect(requestBody.input[0]).toBe(modelPrefix + "test query") + expect(mockedGetModelQueryPrefix).toHaveBeenCalledWith("mistral", "mistral-embed") + }) + }) + + describe("Vercel AI Gateway Embedder", () => { + it("should use custom query instruction when provided", async () => { + const customInstruction = "Search: " + const embedder = new VercelAiGatewayEmbedder( + "test-api-key", + "openai:text-embedding-3-small", + customInstruction, + ) + + // Mock the Vercel AI Gateway API call + const mockFetch = vi.fn().mockResolvedValue({ + ok: true, + json: async () => ({ + data: [{ embedding: new Array(1536).fill(0.1) }], + }), + }) + global.fetch = mockFetch + + await embedder.createEmbeddings(["test query"]) + + // Check that the custom instruction was used + const fetchCall = mockFetch.mock.calls[0] + const requestBody = JSON.parse(fetchCall[1].body) + expect(requestBody.input).toBe(customInstruction + "test query") + expect(mockedGetModelQueryPrefix).not.toHaveBeenCalled() + }) + + it("should fall back to model-specific prefix when no custom instruction", async () => { + const modelPrefix = "query: " + mockedGetModelQueryPrefix.mockReturnValue(modelPrefix) + + const embedder = new VercelAiGatewayEmbedder( + "test-api-key", + "openai:text-embedding-3-small", + undefined, // No custom instruction + ) + + // Mock the Vercel AI Gateway API call + const mockFetch = vi.fn().mockResolvedValue({ + ok: true, + json: async () => ({ + data: [{ embedding: new Array(1536).fill(0.1) }], + }), + }) + global.fetch = mockFetch + + await embedder.createEmbeddings(["test query"]) + + // Check that the model prefix was used + const fetchCall = mockFetch.mock.calls[0] + const requestBody = JSON.parse(fetchCall[1].body) + expect(requestBody.input).toBe(modelPrefix + "test query") + expect(mockedGetModelQueryPrefix).toHaveBeenCalledWith("vercel-ai-gateway", "openai:text-embedding-3-small") + }) + }) +}) diff --git a/src/services/code-index/embedders/gemini.ts b/src/services/code-index/embedders/gemini.ts index 7e795875c9..55ca1cfe6e 100644 --- a/src/services/code-index/embedders/gemini.ts +++ b/src/services/code-index/embedders/gemini.ts @@ -23,8 +23,9 @@ export class GeminiEmbedder implements IEmbedder { * Creates a new Gemini embedder * @param apiKey The Gemini API key for authentication * @param modelId The model ID to use (defaults to gemini-embedding-001) + * @param customQueryInstruction Optional custom query instruction for instruction-aware models */ - constructor(apiKey: string, modelId?: string) { + constructor(apiKey: string, modelId?: string, customQueryInstruction?: string) { if (!apiKey) { throw new Error(t("embeddings:validation.apiKeyRequired")) } @@ -38,6 +39,7 @@ export class GeminiEmbedder implements IEmbedder { apiKey, this.modelId, GEMINI_MAX_ITEM_TOKENS, + customQueryInstruction, ) } diff --git a/src/services/code-index/embedders/mistral.ts b/src/services/code-index/embedders/mistral.ts index c23bcbba1d..30d0bb8a44 100644 --- a/src/services/code-index/embedders/mistral.ts +++ b/src/services/code-index/embedders/mistral.ts @@ -22,8 +22,9 @@ export class MistralEmbedder implements IEmbedder { * Creates a new Mistral embedder * @param apiKey The Mistral API key for authentication * @param modelId The model ID to use (defaults to codestral-embed-2505) + * @param customQueryInstruction Optional custom query instruction for instruction-aware models */ - constructor(apiKey: string, modelId?: string) { + constructor(apiKey: string, modelId?: string, customQueryInstruction?: string) { if (!apiKey) { throw new Error(t("embeddings:validation.apiKeyRequired")) } @@ -37,6 +38,7 @@ export class MistralEmbedder implements IEmbedder { apiKey, this.modelId, MAX_ITEM_TOKENS, // This is the max token limit (8191), not the embedding dimension + customQueryInstruction, ) } diff --git a/src/services/code-index/embedders/ollama.ts b/src/services/code-index/embedders/ollama.ts index 9688a15ff0..b00b45619f 100644 --- a/src/services/code-index/embedders/ollama.ts +++ b/src/services/code-index/embedders/ollama.ts @@ -17,8 +17,9 @@ const OLLAMA_VALIDATION_TIMEOUT_MS = 30000 // 30 seconds for validation requests export class CodeIndexOllamaEmbedder implements IEmbedder { private readonly baseUrl: string private readonly defaultModelId: string + private readonly customQueryInstruction?: string - constructor(options: ApiHandlerOptions) { + constructor(options: ApiHandlerOptions & { customQueryInstruction?: string }) { // Ensure ollamaBaseUrl and ollamaModelId exist on ApiHandlerOptions or add defaults let baseUrl = options.ollamaBaseUrl || "http://localhost:11434" @@ -27,6 +28,7 @@ export class CodeIndexOllamaEmbedder implements IEmbedder { this.baseUrl = baseUrl this.defaultModelId = options.ollamaModelId || "nomic-embed-text:latest" + this.customQueryInstruction = options.customQueryInstruction } /** @@ -39,8 +41,8 @@ export class CodeIndexOllamaEmbedder implements IEmbedder { const modelToUse = model || this.defaultModelId const url = `${this.baseUrl}/api/embed` // Endpoint as specified - // Apply model-specific query prefix if required - const queryPrefix = getModelQueryPrefix("ollama", modelToUse) + // Use custom query instruction if provided, otherwise fall back to model-specific prefix + const queryPrefix = this.customQueryInstruction || getModelQueryPrefix("ollama", modelToUse) const processedTexts = queryPrefix ? texts.map((text, index) => { // Prevent double-prefixing diff --git a/src/services/code-index/embedders/openai-compatible.ts b/src/services/code-index/embedders/openai-compatible.ts index 6eaf2b6c2c..f2be2dbc97 100644 --- a/src/services/code-index/embedders/openai-compatible.ts +++ b/src/services/code-index/embedders/openai-compatible.ts @@ -39,6 +39,7 @@ export class OpenAICompatibleEmbedder implements IEmbedder { private readonly apiKey: string private readonly isFullUrl: boolean private readonly maxItemTokens: number + private readonly customQueryInstruction?: string // Global rate limiting state shared across all instances private static globalRateLimitState = { @@ -56,8 +57,15 @@ export class OpenAICompatibleEmbedder implements IEmbedder { * @param apiKey The API key for authentication * @param modelId Optional model identifier (defaults to "text-embedding-3-small") * @param maxItemTokens Optional maximum tokens per item (defaults to MAX_ITEM_TOKENS) + * @param customQueryInstruction Optional custom query instruction for instruction-aware models */ - constructor(baseUrl: string, apiKey: string, modelId?: string, maxItemTokens?: number) { + constructor( + baseUrl: string, + apiKey: string, + modelId?: string, + maxItemTokens?: number, + customQueryInstruction?: string, + ) { if (!baseUrl) { throw new Error(t("embeddings:validation.baseUrlRequired")) } @@ -83,6 +91,7 @@ export class OpenAICompatibleEmbedder implements IEmbedder { // Cache the URL type check for performance this.isFullUrl = this.isFullEndpointUrl(baseUrl) this.maxItemTokens = maxItemTokens || MAX_ITEM_TOKENS + this.customQueryInstruction = customQueryInstruction } /** @@ -94,8 +103,8 @@ export class OpenAICompatibleEmbedder implements IEmbedder { async createEmbeddings(texts: string[], model?: string): Promise { const modelToUse = model || this.defaultModelId - // Apply model-specific query prefix if required - const queryPrefix = getModelQueryPrefix("openai-compatible", modelToUse) + // Use custom query instruction if provided, otherwise fall back to model-specific prefix + const queryPrefix = this.customQueryInstruction || getModelQueryPrefix("openai-compatible", modelToUse) const processedTexts = queryPrefix ? texts.map((text, index) => { // Prevent double-prefixing diff --git a/src/services/code-index/embedders/openai.ts b/src/services/code-index/embedders/openai.ts index b993e280d9..bb58d00829 100644 --- a/src/services/code-index/embedders/openai.ts +++ b/src/services/code-index/embedders/openai.ts @@ -21,12 +21,13 @@ import { handleOpenAIError } from "../../../api/providers/utils/openai-error-han export class OpenAiEmbedder extends OpenAiNativeHandler implements IEmbedder { private embeddingsClient: OpenAI private readonly defaultModelId: string + private readonly customQueryInstruction?: string /** * Creates a new OpenAI embedder * @param options API handler options */ - constructor(options: ApiHandlerOptions & { openAiEmbeddingModelId?: string }) { + constructor(options: ApiHandlerOptions & { openAiEmbeddingModelId?: string; customQueryInstruction?: string }) { super(options) const apiKey = this.options.openAiNativeApiKey ?? "not-provided" @@ -39,6 +40,7 @@ export class OpenAiEmbedder extends OpenAiNativeHandler implements IEmbedder { } this.defaultModelId = options.openAiEmbeddingModelId || "text-embedding-3-small" + this.customQueryInstruction = options.customQueryInstruction } /** @@ -50,8 +52,8 @@ export class OpenAiEmbedder extends OpenAiNativeHandler implements IEmbedder { async createEmbeddings(texts: string[], model?: string): Promise { const modelToUse = model || this.defaultModelId - // Apply model-specific query prefix if required - const queryPrefix = getModelQueryPrefix("openai", modelToUse) + // Use custom query instruction if provided, otherwise fall back to model-specific prefix + const queryPrefix = this.customQueryInstruction || getModelQueryPrefix("openai", modelToUse) const processedTexts = queryPrefix ? texts.map((text, index) => { // Prevent double-prefixing diff --git a/src/services/code-index/embedders/vercel-ai-gateway.ts b/src/services/code-index/embedders/vercel-ai-gateway.ts index ec0888cd8b..d25f4ac707 100644 --- a/src/services/code-index/embedders/vercel-ai-gateway.ts +++ b/src/services/code-index/embedders/vercel-ai-gateway.ts @@ -31,8 +31,9 @@ export class VercelAiGatewayEmbedder implements IEmbedder { * Creates a new Vercel AI Gateway embedder * @param apiKey The Vercel AI Gateway API key for authentication * @param modelId The model ID to use (defaults to mistral/codestral-embed) + * @param customQueryInstruction Optional custom query instruction for instruction-aware models */ - constructor(apiKey: string, modelId?: string) { + constructor(apiKey: string, modelId?: string, customQueryInstruction?: string) { if (!apiKey) { throw new Error(t("embeddings:validation.apiKeyRequired")) } @@ -46,6 +47,7 @@ export class VercelAiGatewayEmbedder implements IEmbedder { apiKey, this.modelId, MAX_ITEM_TOKENS, + customQueryInstruction, ) } diff --git a/src/services/code-index/interfaces/config.ts b/src/services/code-index/interfaces/config.ts index f168e26869..a097d36e29 100644 --- a/src/services/code-index/interfaces/config.ts +++ b/src/services/code-index/interfaces/config.ts @@ -19,6 +19,7 @@ export interface CodeIndexConfig { qdrantApiKey?: string searchMinScore?: number searchMaxResults?: number + customQueryInstruction?: string } /** diff --git a/src/services/code-index/service-factory.ts b/src/services/code-index/service-factory.ts index 6d69e1f0b6..3e284b62d2 100644 --- a/src/services/code-index/service-factory.ts +++ b/src/services/code-index/service-factory.ts @@ -46,6 +46,7 @@ export class CodeIndexServiceFactory { return new OpenAiEmbedder({ ...config.openAiOptions, openAiEmbeddingModelId: config.modelId, + customQueryInstruction: config.customQueryInstruction, }) } else if (provider === "ollama") { if (!config.ollamaOptions?.ollamaBaseUrl) { @@ -54,6 +55,7 @@ export class CodeIndexServiceFactory { return new CodeIndexOllamaEmbedder({ ...config.ollamaOptions, ollamaModelId: config.modelId, + customQueryInstruction: config.customQueryInstruction, }) } else if (provider === "openai-compatible") { if (!config.openAiCompatibleOptions?.baseUrl || !config.openAiCompatibleOptions?.apiKey) { @@ -63,22 +65,28 @@ export class CodeIndexServiceFactory { config.openAiCompatibleOptions.baseUrl, config.openAiCompatibleOptions.apiKey, config.modelId, + undefined, // maxItemTokens + config.customQueryInstruction, ) } else if (provider === "gemini") { if (!config.geminiOptions?.apiKey) { throw new Error(t("embeddings:serviceFactory.geminiConfigMissing")) } - return new GeminiEmbedder(config.geminiOptions.apiKey, config.modelId) + return new GeminiEmbedder(config.geminiOptions.apiKey, config.modelId, config.customQueryInstruction) } else if (provider === "mistral") { if (!config.mistralOptions?.apiKey) { throw new Error(t("embeddings:serviceFactory.mistralConfigMissing")) } - return new MistralEmbedder(config.mistralOptions.apiKey, config.modelId) + return new MistralEmbedder(config.mistralOptions.apiKey, config.modelId, config.customQueryInstruction) } else if (provider === "vercel-ai-gateway") { if (!config.vercelAiGatewayOptions?.apiKey) { throw new Error(t("embeddings:serviceFactory.vercelAiGatewayConfigMissing")) } - return new VercelAiGatewayEmbedder(config.vercelAiGatewayOptions.apiKey, config.modelId) + return new VercelAiGatewayEmbedder( + config.vercelAiGatewayOptions.apiKey, + config.modelId, + config.customQueryInstruction, + ) } throw new Error( diff --git a/webview-ui/src/components/chat/CodeIndexPopover.tsx b/webview-ui/src/components/chat/CodeIndexPopover.tsx index 45bf4224a1..4a40871ae2 100644 --- a/webview-ui/src/components/chat/CodeIndexPopover.tsx +++ b/webview-ui/src/components/chat/CodeIndexPopover.tsx @@ -64,6 +64,7 @@ interface LocalCodeIndexSettings { codebaseIndexEmbedderModelDimension?: number // Generic dimension for all providers codebaseIndexSearchMaxResults?: number codebaseIndexSearchMinScore?: number + codebaseIndexCustomQueryInstruction?: string // Secret settings (start empty, will be loaded separately) codeIndexOpenAiKey?: string @@ -187,6 +188,7 @@ export const CodeIndexPopover: React.FC = ({ codebaseIndexEmbedderModelDimension: undefined, codebaseIndexSearchMaxResults: CODEBASE_INDEX_DEFAULTS.DEFAULT_SEARCH_RESULTS, codebaseIndexSearchMinScore: CODEBASE_INDEX_DEFAULTS.DEFAULT_SEARCH_MIN_SCORE, + codebaseIndexCustomQueryInstruction: "", codeIndexOpenAiKey: "", codeIndexQdrantApiKey: "", codebaseIndexOpenAiCompatibleBaseUrl: "", @@ -222,6 +224,7 @@ export const CodeIndexPopover: React.FC = ({ codebaseIndexConfig.codebaseIndexSearchMaxResults ?? CODEBASE_INDEX_DEFAULTS.DEFAULT_SEARCH_RESULTS, codebaseIndexSearchMinScore: codebaseIndexConfig.codebaseIndexSearchMinScore ?? CODEBASE_INDEX_DEFAULTS.DEFAULT_SEARCH_MIN_SCORE, + codebaseIndexCustomQueryInstruction: codebaseIndexConfig.codebaseIndexCustomQueryInstruction || "", codeIndexOpenAiKey: "", codeIndexQdrantApiKey: "", codebaseIndexOpenAiCompatibleBaseUrl: codebaseIndexConfig.codebaseIndexOpenAiCompatibleBaseUrl || "", @@ -1287,6 +1290,27 @@ export const CodeIndexPopover: React.FC = ({ + + {/* Custom Query Instruction */} +
+
+ + + + +
+ + updateSetting("codebaseIndexCustomQueryInstruction", e.target.value) + } + placeholder={t("settings:codeIndex.customQueryInstructionPlaceholder")} + className="w-full" + /> +
)} diff --git a/webview-ui/src/i18n/locales/en/settings.json b/webview-ui/src/i18n/locales/en/settings.json index dfccc49cc4..8d6214b4f7 100644 --- a/webview-ui/src/i18n/locales/en/settings.json +++ b/webview-ui/src/i18n/locales/en/settings.json @@ -89,6 +89,9 @@ "qdrantKeyLabel": "Qdrant Key:", "qdrantApiKeyLabel": "Qdrant API Key", "qdrantApiKeyPlaceholder": "Enter your Qdrant API key (optional)", + "customQueryInstructionLabel": "Custom Query Instruction", + "customQueryInstructionPlaceholder": "Enter custom query instruction (e.g., 'Represent this code for retrieving similar code:')", + "customQueryInstructionHint": "Custom instruction to prepend to queries when embedding. Leave empty to use model defaults.", "setupConfigLabel": "Setup", "advancedConfigLabel": "Advanced Configuration", "searchMinScoreLabel": "Search Score Threshold",