fix: add nomic-embed-code support with model-specific score thresholds and query prefixes (#5027)

- Add scoreThreshold and queryPrefix properties to embedding model profiles
- Implement nomic-embed-code model with 0.15 threshold and required query prefix
- Update config manager to use model-specific score thresholds dynamically
- Modify all embedders to apply query prefixes when required
- Maintain backward compatibility for existing models
- Fix search functionality for nomic-embed-code embeddings
This commit is contained in:
hannesrudolph 2025-06-23 06:19:36 -05:00
parent 5bc3af1bc3
commit e03c03ee14
6 changed files with 83 additions and 22 deletions

View file

@ -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
}
}

View file

@ -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
}),
})

View file

@ -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<EmbeddingResponse> {
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[] = []

View file

@ -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<EmbeddingResponse> {
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[] = []

View file

@ -14,7 +14,7 @@ export interface CodeIndexConfig {
openAiCompatibleOptions?: { baseUrl: string; apiKey: string; modelDimension?: number }
qdrantUrl?: string
qdrantApiKey?: string
searchMinScore?: number
searchMinScore: number
}
/**

View file

@ -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.