Skip to content

Commit a348a2a

Browse files
feat: add user-configurable search score threshold slider for semantic search (RooCodeInc#5027) (RooCodeInc#5041)
Co-authored-by: Daniel Riccio <[email protected]>
1 parent 933f28f commit a348a2a

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

46 files changed

+630
-41
lines changed

packages/types/src/codebase-index.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ export const codebaseIndexConfigSchema = z.object({
1010
codebaseIndexEmbedderProvider: z.enum(["openai", "ollama", "openai-compatible"]).optional(),
1111
codebaseIndexEmbedderBaseUrl: z.string().optional(),
1212
codebaseIndexEmbedderModelId: z.string().optional(),
13+
codebaseIndexSearchMinScore: z.number().min(0).max(1).optional(),
1314
})
1415

1516
export type CodebaseIndexConfig = z.infer<typeof codebaseIndexConfigSchema>

src/services/code-index/__tests__/config-manager.spec.ts

Lines changed: 147 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -709,6 +709,153 @@ describe("CodeIndexConfigManager", () => {
709709
const result = await configManager.loadConfiguration()
710710
expect(result.requiresRestart).toBe(false)
711711
})
712+
713+
describe("currentSearchMinScore priority system", () => {
714+
it("should return user-configured score when set", async () => {
715+
mockContextProxy.getGlobalState.mockReturnValue({
716+
codebaseIndexEnabled: true,
717+
codebaseIndexQdrantUrl: "http://qdrant.local",
718+
codebaseIndexEmbedderProvider: "openai",
719+
codebaseIndexEmbedderModelId: "text-embedding-3-small",
720+
codebaseIndexSearchMinScore: 0.8, // User setting
721+
})
722+
mockContextProxy.getSecret.mockImplementation((key: string) => {
723+
if (key === "codeIndexOpenAiKey") return "test-key"
724+
return undefined
725+
})
726+
727+
await configManager.loadConfiguration()
728+
expect(configManager.currentSearchMinScore).toBe(0.8)
729+
})
730+
731+
it("should fall back to model-specific threshold when user setting is undefined", async () => {
732+
mockContextProxy.getGlobalState.mockReturnValue({
733+
codebaseIndexEnabled: true,
734+
codebaseIndexQdrantUrl: "http://qdrant.local",
735+
codebaseIndexEmbedderProvider: "ollama",
736+
codebaseIndexEmbedderModelId: "nomic-embed-code",
737+
// No codebaseIndexSearchMinScore - user hasn't configured it
738+
})
739+
740+
await configManager.loadConfiguration()
741+
// nomic-embed-code has a specific threshold of 0.15
742+
expect(configManager.currentSearchMinScore).toBe(0.15)
743+
})
744+
745+
it("should fall back to default SEARCH_MIN_SCORE when neither user setting nor model threshold exists", async () => {
746+
mockContextProxy.getGlobalState.mockReturnValue({
747+
codebaseIndexEnabled: true,
748+
codebaseIndexQdrantUrl: "http://qdrant.local",
749+
codebaseIndexEmbedderProvider: "openai",
750+
codebaseIndexEmbedderModelId: "unknown-model", // Model not in profiles
751+
// No codebaseIndexSearchMinScore
752+
})
753+
mockContextProxy.getSecret.mockImplementation((key: string) => {
754+
if (key === "codeIndexOpenAiKey") return "test-key"
755+
return undefined
756+
})
757+
758+
await configManager.loadConfiguration()
759+
// Should fall back to default SEARCH_MIN_SCORE (0.4)
760+
expect(configManager.currentSearchMinScore).toBe(0.4)
761+
})
762+
763+
it("should respect user setting of 0 (edge case)", async () => {
764+
mockContextProxy.getGlobalState.mockReturnValue({
765+
codebaseIndexEnabled: true,
766+
codebaseIndexQdrantUrl: "http://qdrant.local",
767+
codebaseIndexEmbedderProvider: "ollama",
768+
codebaseIndexEmbedderModelId: "nomic-embed-code",
769+
codebaseIndexSearchMinScore: 0, // User explicitly sets 0
770+
})
771+
772+
await configManager.loadConfiguration()
773+
// Should return 0, not fall back to model threshold (0.15)
774+
expect(configManager.currentSearchMinScore).toBe(0)
775+
})
776+
777+
it("should use model-specific threshold with openai-compatible provider", async () => {
778+
mockContextProxy.getGlobalState.mockImplementation((key: string) => {
779+
if (key === "codebaseIndexConfig") {
780+
return {
781+
codebaseIndexEnabled: true,
782+
codebaseIndexQdrantUrl: "http://qdrant.local",
783+
codebaseIndexEmbedderProvider: "openai-compatible",
784+
codebaseIndexEmbedderModelId: "nomic-embed-code",
785+
// No codebaseIndexSearchMinScore
786+
}
787+
}
788+
if (key === "codebaseIndexOpenAiCompatibleBaseUrl") return "https://api.example.com/v1"
789+
return undefined
790+
})
791+
mockContextProxy.getSecret.mockImplementation((key: string) => {
792+
if (key === "codebaseIndexOpenAiCompatibleApiKey") return "test-api-key"
793+
return undefined
794+
})
795+
796+
await configManager.loadConfiguration()
797+
// openai-compatible provider also has nomic-embed-code with 0.15 threshold
798+
expect(configManager.currentSearchMinScore).toBe(0.15)
799+
})
800+
801+
it("should use default model ID when modelId is not specified", async () => {
802+
mockContextProxy.getGlobalState.mockReturnValue({
803+
codebaseIndexEnabled: true,
804+
codebaseIndexQdrantUrl: "http://qdrant.local",
805+
codebaseIndexEmbedderProvider: "openai",
806+
// No modelId specified
807+
// No codebaseIndexSearchMinScore
808+
})
809+
mockContextProxy.getSecret.mockImplementation((key: string) => {
810+
if (key === "codeIndexOpenAiKey") return "test-key"
811+
return undefined
812+
})
813+
814+
await configManager.loadConfiguration()
815+
// Should use default model (text-embedding-3-small) threshold (0.4)
816+
expect(configManager.currentSearchMinScore).toBe(0.4)
817+
})
818+
819+
it("should handle priority correctly: user > model > default", async () => {
820+
// Test 1: User setting takes precedence
821+
mockContextProxy.getGlobalState.mockReturnValue({
822+
codebaseIndexEnabled: true,
823+
codebaseIndexQdrantUrl: "http://qdrant.local",
824+
codebaseIndexEmbedderProvider: "ollama",
825+
codebaseIndexEmbedderModelId: "nomic-embed-code", // Has 0.15 threshold
826+
codebaseIndexSearchMinScore: 0.9, // User overrides
827+
})
828+
829+
await configManager.loadConfiguration()
830+
expect(configManager.currentSearchMinScore).toBe(0.9) // User setting wins
831+
832+
// Test 2: Model threshold when no user setting
833+
mockContextProxy.getGlobalState.mockReturnValue({
834+
codebaseIndexEnabled: true,
835+
codebaseIndexQdrantUrl: "http://qdrant.local",
836+
codebaseIndexEmbedderProvider: "ollama",
837+
codebaseIndexEmbedderModelId: "nomic-embed-code",
838+
// No user setting
839+
})
840+
841+
const newManager = new CodeIndexConfigManager(mockContextProxy)
842+
await newManager.loadConfiguration()
843+
expect(newManager.currentSearchMinScore).toBe(0.15) // Model threshold
844+
845+
// Test 3: Default when neither exists
846+
mockContextProxy.getGlobalState.mockReturnValue({
847+
codebaseIndexEnabled: true,
848+
codebaseIndexQdrantUrl: "http://qdrant.local",
849+
codebaseIndexEmbedderProvider: "openai",
850+
codebaseIndexEmbedderModelId: "custom-unknown-model",
851+
// No user setting, unknown model
852+
})
853+
854+
const anotherManager = new CodeIndexConfigManager(mockContextProxy)
855+
await anotherManager.loadConfiguration()
856+
expect(anotherManager.currentSearchMinScore).toBe(0.4) // Default
857+
})
858+
})
712859
})
713860

714861
describe("empty/missing API key handling", () => {

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

Lines changed: 18 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.
@@ -34,10 +34,10 @@ export class CodeIndexConfigManager {
3434
const codebaseIndexConfig = this.contextProxy?.getGlobalState("codebaseIndexConfig") ?? {
3535
codebaseIndexEnabled: false,
3636
codebaseIndexQdrantUrl: "http://localhost:6333",
37-
codebaseIndexSearchMinScore: 0.4,
3837
codebaseIndexEmbedderProvider: "openai",
3938
codebaseIndexEmbedderBaseUrl: "",
4039
codebaseIndexEmbedderModelId: "",
40+
codebaseIndexSearchMinScore: undefined,
4141
}
4242

4343
const {
@@ -46,6 +46,7 @@ export class CodeIndexConfigManager {
4646
codebaseIndexEmbedderProvider,
4747
codebaseIndexEmbedderBaseUrl,
4848
codebaseIndexEmbedderModelId,
49+
codebaseIndexSearchMinScore,
4950
} = codebaseIndexConfig
5051

5152
const openAiKey = this.contextProxy?.getSecret("codeIndexOpenAiKey") ?? ""
@@ -60,8 +61,8 @@ export class CodeIndexConfigManager {
6061
this.isEnabled = codebaseIndexEnabled || false
6162
this.qdrantUrl = codebaseIndexQdrantUrl
6263
this.qdrantApiKey = qdrantApiKey ?? ""
64+
this.searchMinScore = codebaseIndexSearchMinScore
6365
this.openAiOptions = { openAiNativeApiKey: openAiKey }
64-
this.searchMinScore = SEARCH_MIN_SCORE
6566

6667
// Set embedder provider with support for openai-compatible
6768
if (codebaseIndexEmbedderProvider === "ollama") {
@@ -139,7 +140,7 @@ export class CodeIndexConfigManager {
139140
openAiCompatibleOptions: this.openAiCompatibleOptions,
140141
qdrantUrl: this.qdrantUrl,
141142
qdrantApiKey: this.qdrantApiKey,
142-
searchMinScore: this.searchMinScore,
143+
searchMinScore: this.currentSearchMinScore,
143144
},
144145
requiresRestart,
145146
}
@@ -294,7 +295,7 @@ export class CodeIndexConfigManager {
294295
openAiCompatibleOptions: this.openAiCompatibleOptions,
295296
qdrantUrl: this.qdrantUrl,
296297
qdrantApiKey: this.qdrantApiKey,
297-
searchMinScore: this.searchMinScore,
298+
searchMinScore: this.currentSearchMinScore,
298299
}
299300
}
300301

@@ -337,9 +338,18 @@ export class CodeIndexConfigManager {
337338
}
338339

339340
/**
340-
* Gets the configured minimum search score.
341+
* Gets the configured minimum search score based on user setting, model-specific threshold, or fallback.
342+
* Priority: 1) User setting, 2) Model-specific threshold, 3) Default SEARCH_MIN_SCORE constant.
341343
*/
342-
public get currentSearchMinScore(): number | undefined {
343-
return this.searchMinScore
344+
public get currentSearchMinScore(): number {
345+
// First check if user has configured a custom score threshold
346+
if (this.searchMinScore !== undefined) {
347+
return this.searchMinScore
348+
}
349+
350+
// Fall back to model-specific threshold
351+
const currentModelId = this.modelId ?? getDefaultModelId(this.embedderProvider)
352+
const modelSpecificThreshold = getModelScoreThreshold(this.embedderProvider, currentModelId)
353+
return modelSpecificThreshold ?? SEARCH_MIN_SCORE
344354
}
345355
}

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

Lines changed: 28 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
import { ApiHandlerOptions } from "../../../shared/api"
22
import { EmbedderInfo, EmbeddingResponse, IEmbedder } from "../interfaces"
3+
import { getModelQueryPrefix } from "../../../shared/embeddingModels"
4+
import { MAX_ITEM_TOKENS } from "../constants"
35
import { t } from "../../../i18n"
46

57
/**
@@ -25,6 +27,31 @@ export class CodeIndexOllamaEmbedder implements IEmbedder {
2527
const modelToUse = model || this.defaultModelId
2628
const url = `${this.baseUrl}/api/embed` // Endpoint as specified
2729

30+
// Apply model-specific query prefix if required
31+
const queryPrefix = getModelQueryPrefix("ollama", modelToUse)
32+
const processedTexts = queryPrefix
33+
? texts.map((text, index) => {
34+
// Prevent double-prefixing
35+
if (text.startsWith(queryPrefix)) {
36+
return text
37+
}
38+
const prefixedText = `${queryPrefix}${text}`
39+
const estimatedTokens = Math.ceil(prefixedText.length / 4)
40+
if (estimatedTokens > MAX_ITEM_TOKENS) {
41+
console.warn(
42+
t("embeddings:textWithPrefixExceedsTokenLimit", {
43+
index,
44+
estimatedTokens,
45+
maxTokens: MAX_ITEM_TOKENS,
46+
}),
47+
)
48+
// Return original text if adding prefix would exceed limit
49+
return text
50+
}
51+
return prefixedText
52+
})
53+
: texts
54+
2855
try {
2956
// Note: Standard Ollama API uses 'prompt' for single text, not 'input' for array.
3057
// Implementing based on user's specific request structure.
@@ -35,7 +62,7 @@ export class CodeIndexOllamaEmbedder implements IEmbedder {
3562
},
3663
body: JSON.stringify({
3764
model: modelToUse,
38-
input: texts, // Using 'input' as requested
65+
input: processedTexts, // Using 'input' as requested
3966
}),
4067
})
4168

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

Lines changed: 28 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,35 @@ 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
66+
? texts.map((text, index) => {
67+
// Prevent double-prefixing
68+
if (text.startsWith(queryPrefix)) {
69+
return text
70+
}
71+
const prefixedText = `${queryPrefix}${text}`
72+
const estimatedTokens = Math.ceil(prefixedText.length / 4)
73+
if (estimatedTokens > MAX_ITEM_TOKENS) {
74+
console.warn(
75+
t("embeddings:textWithPrefixExceedsTokenLimit", {
76+
index,
77+
estimatedTokens,
78+
maxTokens: MAX_ITEM_TOKENS,
79+
}),
80+
)
81+
// Return original text if adding prefix would exceed limit
82+
return text
83+
}
84+
return prefixedText
85+
})
86+
: texts
87+
6288
const allEmbeddings: number[][] = []
6389
const usage = { promptTokens: 0, totalTokens: 0 }
64-
const remainingTexts = [...texts]
90+
const remainingTexts = [...processedTexts]
6591

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

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

Lines changed: 28 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,35 @@ 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
44+
? texts.map((text, index) => {
45+
// Prevent double-prefixing
46+
if (text.startsWith(queryPrefix)) {
47+
return text
48+
}
49+
const prefixedText = `${queryPrefix}${text}`
50+
const estimatedTokens = Math.ceil(prefixedText.length / 4)
51+
if (estimatedTokens > MAX_ITEM_TOKENS) {
52+
console.warn(
53+
t("embeddings:textWithPrefixExceedsTokenLimit", {
54+
index,
55+
estimatedTokens,
56+
maxTokens: MAX_ITEM_TOKENS,
57+
}),
58+
)
59+
// Return original text if adding prefix would exceed limit
60+
return text
61+
}
62+
return prefixedText
63+
})
64+
: texts
65+
3966
const allEmbeddings: number[][] = []
4067
const usage = { promptTokens: 0, totalTokens: 0 }
41-
const remainingTexts = [...texts]
68+
const remainingTexts = [...processedTexts]
4269

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

0 commit comments

Comments
 (0)