mirror of
https://github.com/RooVetGit/Roo-Code.git
synced 2026-09-06 08:18:39 +00:00
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:
parent
5bc3af1bc3
commit
e03c03ee14
6 changed files with 83 additions and 22 deletions
|
|
@ -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
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}),
|
||||
})
|
||||
|
||||
|
|
|
|||
|
|
@ -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[] = []
|
||||
|
|
|
|||
|
|
@ -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[] = []
|
||||
|
|
|
|||
|
|
@ -14,7 +14,7 @@ export interface CodeIndexConfig {
|
|||
openAiCompatibleOptions?: { baseUrl: string; apiKey: string; modelDimension?: number }
|
||||
qdrantUrl?: string
|
||||
qdrantApiKey?: string
|
||||
searchMinScore?: number
|
||||
searchMinScore: number
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue