Skip to content

Commit e03c03e

Browse files
hannesrudolphhannesrudolph
committed
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
1 parent 5bc3af1 commit e03c03e

File tree

6 files changed

+83
-22
lines changed

6 files changed

+83
-22
lines changed

‎src/services/code-index/config-manager.ts‎

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ import { ContextProxy } from "../../core/config/ContextProxy"
33
import { EmbedderProvider } from "./interfaces/manager"
44
import { CodeIndexConfig, PreviousConfigSnapshot } from "./interfaces/config"
55
import { SEARCH_MIN_SCORE } from "./constants"
6-
import { getDefaultModelId, getModelDimension } from "../../shared/embeddingModels"
6+
import { getDefaultModelId, getModelDimension, getModelScoreThreshold } from "../../shared/embeddingModels"
77

88
/**
99
* Manages configuration state and validation for the code indexing feature.
@@ -18,7 +18,6 @@ export class CodeIndexConfigManager {
1818
private openAiCompatibleOptions?: { baseUrl: string; apiKey: string; modelDimension?: number }
1919
private qdrantUrl?: string = "http://localhost:6333"
2020
private qdrantApiKey?: string
21-
private searchMinScore?: number
2221

2322
constructor(private readonly contextProxy: ContextProxy) {
2423
// Initialize with current configuration to avoid false restart triggers
@@ -61,7 +60,6 @@ export class CodeIndexConfigManager {
6160
this.qdrantUrl = codebaseIndexQdrantUrl
6261
this.qdrantApiKey = qdrantApiKey ?? ""
6362
this.openAiOptions = { openAiNativeApiKey: openAiKey }
64-
this.searchMinScore = SEARCH_MIN_SCORE
6563

6664
// Set embedder provider with support for openai-compatible
6765
if (codebaseIndexEmbedderProvider === "ollama") {
@@ -139,7 +137,7 @@ export class CodeIndexConfigManager {
139137
openAiCompatibleOptions: this.openAiCompatibleOptions,
140138
qdrantUrl: this.qdrantUrl,
141139
qdrantApiKey: this.qdrantApiKey,
142-
searchMinScore: this.searchMinScore,
140+
searchMinScore: this.currentSearchMinScore,
143141
},
144142
requiresRestart,
145143
}
@@ -294,7 +292,7 @@ export class CodeIndexConfigManager {
294292
openAiCompatibleOptions: this.openAiCompatibleOptions,
295293
qdrantUrl: this.qdrantUrl,
296294
qdrantApiKey: this.qdrantApiKey,
297-
searchMinScore: this.searchMinScore,
295+
searchMinScore: this.currentSearchMinScore,
298296
}
299297
}
300298

@@ -337,9 +335,12 @@ export class CodeIndexConfigManager {
337335
}
338336

339337
/**
340-
* Gets the configured minimum search score.
338+
* Gets the configured minimum search score based on the current model.
339+
* Falls back to the constant SEARCH_MIN_SCORE if no model-specific threshold is found.
341340
*/
342-
public get currentSearchMinScore(): number | undefined {
343-
return this.searchMinScore
341+
public get currentSearchMinScore(): number {
342+
const currentModelId = this.modelId ?? getDefaultModelId(this.embedderProvider)
343+
const modelSpecificThreshold = getModelScoreThreshold(this.embedderProvider, currentModelId)
344+
return modelSpecificThreshold ?? SEARCH_MIN_SCORE
344345
}
345346
}

‎src/services/code-index/embedders/ollama.ts‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import { ApiHandlerOptions } from "../../../shared/api"
22
import { EmbedderInfo, EmbeddingResponse, IEmbedder } from "../interfaces"
3+
import { getModelQueryPrefix } from "../../../shared/embeddingModels"
34
import { t } from "../../../i18n"
45

56
/**
@@ -25,6 +26,10 @@ export class CodeIndexOllamaEmbedder implements IEmbedder {
2526
const modelToUse = model || this.defaultModelId
2627
const url = `${this.baseUrl}/api/embed` // Endpoint as specified
2728

29+
// Apply model-specific query prefix if required
30+
const queryPrefix = getModelQueryPrefix("ollama", modelToUse)
31+
const processedTexts = queryPrefix ? texts.map((text) => `${queryPrefix}${text}`) : texts
32+
2833
try {
2934
// Note: Standard Ollama API uses 'prompt' for single text, not 'input' for array.
3035
// Implementing based on user's specific request structure.
@@ -35,7 +40,7 @@ export class CodeIndexOllamaEmbedder implements IEmbedder {
3540
},
3641
body: JSON.stringify({
3742
model: modelToUse,
38-
input: texts, // Using 'input' as requested
43+
input: processedTexts, // Using 'input' as requested
3944
}),
4045
})
4146

‎src/services/code-index/embedders/openai-compatible.ts‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ import {
66
MAX_BATCH_RETRIES as MAX_RETRIES,
77
INITIAL_RETRY_DELAY_MS as INITIAL_DELAY_MS,
88
} from "../constants"
9-
import { getDefaultModelId } from "../../../shared/embeddingModels"
9+
import { getDefaultModelId, getModelQueryPrefix } from "../../../shared/embeddingModels"
1010
import { t } from "../../../i18n"
1111

1212
interface EmbeddingItem {
@@ -59,9 +59,14 @@ export class OpenAICompatibleEmbedder implements IEmbedder {
5959
*/
6060
async createEmbeddings(texts: string[], model?: string): Promise<EmbeddingResponse> {
6161
const modelToUse = model || this.defaultModelId
62+
63+
// Apply model-specific query prefix if required
64+
const queryPrefix = getModelQueryPrefix("openai-compatible", modelToUse)
65+
const processedTexts = queryPrefix ? texts.map((text) => `${queryPrefix}${text}`) : texts
66+
6267
const allEmbeddings: number[][] = []
6368
const usage = { promptTokens: 0, totalTokens: 0 }
64-
const remainingTexts = [...texts]
69+
const remainingTexts = [...processedTexts]
6570

6671
while (remainingTexts.length > 0) {
6772
const currentBatch: string[] = []

‎src/services/code-index/embedders/openai.ts‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import {
88
MAX_BATCH_RETRIES as MAX_RETRIES,
99
INITIAL_RETRY_DELAY_MS as INITIAL_DELAY_MS,
1010
} from "../constants"
11+
import { getModelQueryPrefix } from "../../../shared/embeddingModels"
1112
import { t } from "../../../i18n"
1213

1314
/**
@@ -36,9 +37,14 @@ export class OpenAiEmbedder extends OpenAiNativeHandler implements IEmbedder {
3637
*/
3738
async createEmbeddings(texts: string[], model?: string): Promise<EmbeddingResponse> {
3839
const modelToUse = model || this.defaultModelId
40+
41+
// Apply model-specific query prefix if required
42+
const queryPrefix = getModelQueryPrefix("openai", modelToUse)
43+
const processedTexts = queryPrefix ? texts.map((text) => `${queryPrefix}${text}`) : texts
44+
3945
const allEmbeddings: number[][] = []
4046
const usage = { promptTokens: 0, totalTokens: 0 }
41-
const remainingTexts = [...texts]
47+
const remainingTexts = [...processedTexts]
4248

4349
while (remainingTexts.length > 0) {
4450
const currentBatch: string[] = []

‎src/services/code-index/interfaces/config.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ export interface CodeIndexConfig {
1414
openAiCompatibleOptions?: { baseUrl: string; apiKey: string; modelDimension?: number }
1515
qdrantUrl?: string
1616
qdrantApiKey?: string
17-
searchMinScore?: number
17+
searchMinScore: number
1818
}
1919

2020
/**

‎src/shared/embeddingModels.ts‎

Lines changed: 53 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@ export type EmbedderProvider = "openai" | "ollama" | "openai-compatible" // Add
66

77
export interface EmbeddingModelProfile {
88
dimension: number
9+
scoreThreshold?: number // Model-specific minimum score threshold for semantic search
10+
queryPrefix?: string // Optional prefix required by the model for queries
911
// Add other model-specific properties if needed, e.g., context window size
1012
}
1113

@@ -18,21 +20,31 @@ export type EmbeddingModelProfiles = {
1820
// Example profiles - expand this list as needed
1921
export const EMBEDDING_MODEL_PROFILES: EmbeddingModelProfiles = {
2022
openai: {
21-
"text-embedding-3-small": { dimension: 1536 },
22-
"text-embedding-3-large": { dimension: 3072 },
23-
"text-embedding-ada-002": { dimension: 1536 },
23+
"text-embedding-3-small": { dimension: 1536, scoreThreshold: 0.4 },
24+
"text-embedding-3-large": { dimension: 3072, scoreThreshold: 0.4 },
25+
"text-embedding-ada-002": { dimension: 1536, scoreThreshold: 0.4 },
2426
},
2527
ollama: {
26-
"nomic-embed-text": { dimension: 768 },
27-
"mxbai-embed-large": { dimension: 1024 },
28-
"all-minilm": { dimension: 384 },
28+
"nomic-embed-text": { dimension: 768, scoreThreshold: 0.4 },
29+
"nomic-embed-code": {
30+
dimension: 3584,
31+
scoreThreshold: 0.15,
32+
queryPrefix: "Represent this query for searching relevant code: ",
33+
},
34+
"mxbai-embed-large": { dimension: 1024, scoreThreshold: 0.4 },
35+
"all-minilm": { dimension: 384, scoreThreshold: 0.4 },
2936
// Add default Ollama model if applicable, e.g.:
3037
// 'default': { dimension: 768 } // Assuming a default dimension
3138
},
3239
"openai-compatible": {
33-
"text-embedding-3-small": { dimension: 1536 },
34-
"text-embedding-3-large": { dimension: 3072 },
35-
"text-embedding-ada-002": { dimension: 1536 },
40+
"text-embedding-3-small": { dimension: 1536, scoreThreshold: 0.4 },
41+
"text-embedding-3-large": { dimension: 3072, scoreThreshold: 0.4 },
42+
"text-embedding-ada-002": { dimension: 1536, scoreThreshold: 0.4 },
43+
"nomic-embed-code": {
44+
dimension: 3584,
45+
scoreThreshold: 0.15,
46+
queryPrefix: "Represent this query for searching relevant code: ",
47+
},
3648
},
3749
}
3850

@@ -59,6 +71,38 @@ export function getModelDimension(provider: EmbedderProvider, modelId: string):
5971
return modelProfile.dimension
6072
}
6173

74+
/**
75+
* Retrieves the score threshold for a given provider and model ID.
76+
* @param provider The embedder provider (e.g., "openai").
77+
* @param modelId The specific model ID (e.g., "text-embedding-3-small").
78+
* @returns The score threshold or undefined if the model is not found.
79+
*/
80+
export function getModelScoreThreshold(provider: EmbedderProvider, modelId: string): number | undefined {
81+
const providerProfiles = EMBEDDING_MODEL_PROFILES[provider]
82+
if (!providerProfiles) {
83+
return undefined
84+
}
85+
86+
const modelProfile = providerProfiles[modelId]
87+
return modelProfile?.scoreThreshold
88+
}
89+
90+
/**
91+
* Retrieves the query prefix for a given provider and model ID.
92+
* @param provider The embedder provider (e.g., "openai").
93+
* @param modelId The specific model ID (e.g., "nomic-embed-code").
94+
* @returns The query prefix or undefined if the model doesn't require one.
95+
*/
96+
export function getModelQueryPrefix(provider: EmbedderProvider, modelId: string): string | undefined {
97+
const providerProfiles = EMBEDDING_MODEL_PROFILES[provider]
98+
if (!providerProfiles) {
99+
return undefined
100+
}
101+
102+
const modelProfile = providerProfiles[modelId]
103+
return modelProfile?.queryPrefix
104+
}
105+
62106
/**
63107
* Gets the default *specific* embedding model ID based on the provider.
64108
* Does not include the provider prefix.

0 commit comments

Comments
 (0)