diff --git a/src/services/code-index/config-manager.ts b/src/services/code-index/config-manager.ts index 678cec36a1..d73e6b8695 100644 --- a/src/services/code-index/config-manager.ts +++ b/src/services/code-index/config-manager.ts @@ -3,7 +3,7 @@ import { ContextProxy } from "../../core/config/ContextProxy" import { EmbedderProvider } from "./interfaces/manager" import { CodeIndexConfig, PreviousConfigSnapshot } from "./interfaces/config" import { SEARCH_MIN_SCORE } from "./constants" -import { getDefaultModelId, getModelDimension } from "../../shared/embeddingModels" +import { getDefaultModelId, getModelDimension, getModelScoreThreshold } from "../../shared/embeddingModels" /** * Manages configuration state and validation for the code indexing feature. @@ -18,7 +18,6 @@ export class CodeIndexConfigManager { private openAiCompatibleOptions?: { baseUrl: string; apiKey: string; modelDimension?: number } private qdrantUrl?: string = "http://localhost:6333" private qdrantApiKey?: string - private searchMinScore?: number constructor(private readonly contextProxy: ContextProxy) { // Initialize with current configuration to avoid false restart triggers @@ -61,7 +60,6 @@ export class CodeIndexConfigManager { this.qdrantUrl = codebaseIndexQdrantUrl this.qdrantApiKey = qdrantApiKey ?? "" this.openAiOptions = { openAiNativeApiKey: openAiKey } - this.searchMinScore = SEARCH_MIN_SCORE // Set embedder provider with support for openai-compatible if (codebaseIndexEmbedderProvider === "ollama") { @@ -139,7 +137,7 @@ export class CodeIndexConfigManager { openAiCompatibleOptions: this.openAiCompatibleOptions, qdrantUrl: this.qdrantUrl, qdrantApiKey: this.qdrantApiKey, - searchMinScore: this.searchMinScore, + searchMinScore: this.currentSearchMinScore, }, requiresRestart, } @@ -294,7 +292,7 @@ export class CodeIndexConfigManager { openAiCompatibleOptions: this.openAiCompatibleOptions, qdrantUrl: this.qdrantUrl, qdrantApiKey: this.qdrantApiKey, - searchMinScore: this.searchMinScore, + searchMinScore: this.currentSearchMinScore, } } @@ -337,9 +335,12 @@ export class CodeIndexConfigManager { } /** - * Gets the configured minimum search score. + * Gets the configured minimum search score based on the current model. + * Falls back to the constant SEARCH_MIN_SCORE if no model-specific threshold is found. */ - public get currentSearchMinScore(): number | undefined { - return this.searchMinScore + public get currentSearchMinScore(): number { + const currentModelId = this.modelId ?? getDefaultModelId(this.embedderProvider) + const modelSpecificThreshold = getModelScoreThreshold(this.embedderProvider, currentModelId) + return modelSpecificThreshold ?? SEARCH_MIN_SCORE } } diff --git a/src/services/code-index/embedders/ollama.ts b/src/services/code-index/embedders/ollama.ts index 748ed188a4..7f3e08b99e 100644 --- a/src/services/code-index/embedders/ollama.ts +++ b/src/services/code-index/embedders/ollama.ts @@ -1,5 +1,6 @@ import { ApiHandlerOptions } from "../../../shared/api" import { EmbedderInfo, EmbeddingResponse, IEmbedder } from "../interfaces" +import { getModelQueryPrefix } from "../../../shared/embeddingModels" import { t } from "../../../i18n" /** @@ -25,6 +26,10 @@ 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) + const processedTexts = queryPrefix ? texts.map((text) => `${queryPrefix}${text}`) : texts + try { // Note: Standard Ollama API uses 'prompt' for single text, not 'input' for array. // Implementing based on user's specific request structure. @@ -35,7 +40,7 @@ export class CodeIndexOllamaEmbedder implements IEmbedder { }, body: JSON.stringify({ model: modelToUse, - input: texts, // Using 'input' as requested + input: processedTexts, // Using 'input' as requested }), }) diff --git a/src/services/code-index/embedders/openai-compatible.ts b/src/services/code-index/embedders/openai-compatible.ts index 0983cc297f..98f324c380 100644 --- a/src/services/code-index/embedders/openai-compatible.ts +++ b/src/services/code-index/embedders/openai-compatible.ts @@ -6,7 +6,7 @@ import { MAX_BATCH_RETRIES as MAX_RETRIES, INITIAL_RETRY_DELAY_MS as INITIAL_DELAY_MS, } from "../constants" -import { getDefaultModelId } from "../../../shared/embeddingModels" +import { getDefaultModelId, getModelQueryPrefix } from "../../../shared/embeddingModels" import { t } from "../../../i18n" interface EmbeddingItem { @@ -59,9 +59,14 @@ 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) + const processedTexts = queryPrefix ? texts.map((text) => `${queryPrefix}${text}`) : texts + const allEmbeddings: number[][] = [] const usage = { promptTokens: 0, totalTokens: 0 } - const remainingTexts = [...texts] + const remainingTexts = [...processedTexts] while (remainingTexts.length > 0) { const currentBatch: string[] = [] diff --git a/src/services/code-index/embedders/openai.ts b/src/services/code-index/embedders/openai.ts index d0dc132df7..9ff9f8f4aa 100644 --- a/src/services/code-index/embedders/openai.ts +++ b/src/services/code-index/embedders/openai.ts @@ -8,6 +8,7 @@ import { MAX_BATCH_RETRIES as MAX_RETRIES, INITIAL_RETRY_DELAY_MS as INITIAL_DELAY_MS, } from "../constants" +import { getModelQueryPrefix } from "../../../shared/embeddingModels" import { t } from "../../../i18n" /** @@ -36,9 +37,14 @@ 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) + const processedTexts = queryPrefix ? texts.map((text) => `${queryPrefix}${text}`) : texts + const allEmbeddings: number[][] = [] const usage = { promptTokens: 0, totalTokens: 0 } - const remainingTexts = [...texts] + const remainingTexts = [...processedTexts] while (remainingTexts.length > 0) { const currentBatch: string[] = [] diff --git a/src/services/code-index/interfaces/config.ts b/src/services/code-index/interfaces/config.ts index 0843120fd9..584bac3c57 100644 --- a/src/services/code-index/interfaces/config.ts +++ b/src/services/code-index/interfaces/config.ts @@ -14,7 +14,7 @@ export interface CodeIndexConfig { openAiCompatibleOptions?: { baseUrl: string; apiKey: string; modelDimension?: number } qdrantUrl?: string qdrantApiKey?: string - searchMinScore?: number + searchMinScore: number } /** diff --git a/src/shared/embeddingModels.ts b/src/shared/embeddingModels.ts index cd7c1d4e6b..c78dc6c487 100644 --- a/src/shared/embeddingModels.ts +++ b/src/shared/embeddingModels.ts @@ -6,6 +6,8 @@ export type EmbedderProvider = "openai" | "ollama" | "openai-compatible" // Add export interface EmbeddingModelProfile { dimension: number + scoreThreshold?: number // Model-specific minimum score threshold for semantic search + queryPrefix?: string // Optional prefix required by the model for queries // Add other model-specific properties if needed, e.g., context window size } @@ -18,21 +20,31 @@ export type EmbeddingModelProfiles = { // Example profiles - expand this list as needed export const EMBEDDING_MODEL_PROFILES: EmbeddingModelProfiles = { openai: { - "text-embedding-3-small": { dimension: 1536 }, - "text-embedding-3-large": { dimension: 3072 }, - "text-embedding-ada-002": { dimension: 1536 }, + "text-embedding-3-small": { dimension: 1536, scoreThreshold: 0.4 }, + "text-embedding-3-large": { dimension: 3072, scoreThreshold: 0.4 }, + "text-embedding-ada-002": { dimension: 1536, scoreThreshold: 0.4 }, }, ollama: { - "nomic-embed-text": { dimension: 768 }, - "mxbai-embed-large": { dimension: 1024 }, - "all-minilm": { dimension: 384 }, + "nomic-embed-text": { dimension: 768, scoreThreshold: 0.4 }, + "nomic-embed-code": { + dimension: 3584, + scoreThreshold: 0.15, + queryPrefix: "Represent this query for searching relevant code: ", + }, + "mxbai-embed-large": { dimension: 1024, scoreThreshold: 0.4 }, + "all-minilm": { dimension: 384, scoreThreshold: 0.4 }, // Add default Ollama model if applicable, e.g.: // 'default': { dimension: 768 } // Assuming a default dimension }, "openai-compatible": { - "text-embedding-3-small": { dimension: 1536 }, - "text-embedding-3-large": { dimension: 3072 }, - "text-embedding-ada-002": { dimension: 1536 }, + "text-embedding-3-small": { dimension: 1536, scoreThreshold: 0.4 }, + "text-embedding-3-large": { dimension: 3072, scoreThreshold: 0.4 }, + "text-embedding-ada-002": { dimension: 1536, scoreThreshold: 0.4 }, + "nomic-embed-code": { + dimension: 3584, + scoreThreshold: 0.15, + queryPrefix: "Represent this query for searching relevant code: ", + }, }, } @@ -59,6 +71,38 @@ export function getModelDimension(provider: EmbedderProvider, modelId: string): return modelProfile.dimension } +/** + * Retrieves the score threshold for a given provider and model ID. + * @param provider The embedder provider (e.g., "openai"). + * @param modelId The specific model ID (e.g., "text-embedding-3-small"). + * @returns The score threshold or undefined if the model is not found. + */ +export function getModelScoreThreshold(provider: EmbedderProvider, modelId: string): number | undefined { + const providerProfiles = EMBEDDING_MODEL_PROFILES[provider] + if (!providerProfiles) { + return undefined + } + + const modelProfile = providerProfiles[modelId] + return modelProfile?.scoreThreshold +} + +/** + * Retrieves the query prefix for a given provider and model ID. + * @param provider The embedder provider (e.g., "openai"). + * @param modelId The specific model ID (e.g., "nomic-embed-code"). + * @returns The query prefix or undefined if the model doesn't require one. + */ +export function getModelQueryPrefix(provider: EmbedderProvider, modelId: string): string | undefined { + const providerProfiles = EMBEDDING_MODEL_PROFILES[provider] + if (!providerProfiles) { + return undefined + } + + const modelProfile = providerProfiles[modelId] + return modelProfile?.queryPrefix +} + /** * Gets the default *specific* embedding model ID based on the provider. * Does not include the provider prefix.