Skip to content

Commit 761b2ea

Browse files
Jopo-JPJopo-JP
committed
feat: Enhance Gemini embedder with configurable dimensions and validation for embedding models
1 parent 8a3dcfb commit 761b2ea

File tree

3 files changed

+121
-13
lines changed

3 files changed

+121
-13
lines changed

‎packages/types/src/codebase-index.ts‎

Lines changed: 17 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -42,11 +42,24 @@ export type CodebaseIndexConfig = z.infer<typeof codebaseIndexConfigSchema>
4242
* CodebaseIndexModels
4343
*/
4444

45+
const modelProfileSchema = z.object({
46+
/** The fixed dimension for the model, or a fallback for models with variable dimensions. */
47+
dimension: z.number(),
48+
scoreThreshold: z.number().optional(),
49+
queryPrefix: z.string().optional(),
50+
/** The minimum dimension supported by a variable-dimension model. */
51+
minDimension: z.number().optional(),
52+
/** The maximum dimension supported by a variable-dimension model. */
53+
maxDimension: z.number().optional(),
54+
/** The default dimension for a variable-dimension model, used for UI presentation. */
55+
defaultDimension: z.number().optional(),
56+
})
57+
4558
export const codebaseIndexModelsSchema = z.object({
46-
openai: z.record(z.string(), z.object({ dimension: z.number() })).optional(),
47-
ollama: z.record(z.string(), z.object({ dimension: z.number() })).optional(),
48-
"openai-compatible": z.record(z.string(), z.object({ dimension: z.number() })).optional(),
49-
gemini: z.record(z.string(), z.object({ dimension: z.number() })).optional(),
59+
openai: z.record(z.string(), modelProfileSchema).optional(),
60+
ollama: z.record(z.string(), modelProfileSchema).optional(),
61+
"openai-compatible": z.record(z.string(), modelProfileSchema).optional(),
62+
gemini: z.record(z.string(), modelProfileSchema).optional(),
5063
})
5164

5265
export type CodebaseIndexModels = z.infer<typeof codebaseIndexModelsSchema>

‎src/shared/embeddingModels.ts‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,9 +5,13 @@
55
export type EmbedderProvider = "openai" | "ollama" | "openai-compatible" | "gemini" // Add other providers as needed
66

77
export interface EmbeddingModelProfile {
8+
/** The fixed dimension for the model, or a fallback for models with variable dimensions. */
89
dimension: number
910
scoreThreshold?: number // Model-specific minimum score threshold for semantic search
1011
queryPrefix?: string // Optional prefix required by the model for queries
12+
minDimension?: number // The minimum dimension supported by a variable-dimension model.
13+
maxDimension?: number // The maximum dimension supported by a variable-dimension model.
14+
defaultDimension?: number // The default dimension for a variable-dimension model, used for UI presentation.
1115
// Add other model-specific properties if needed, e.g., context window size
1216
}
1317

‎webview-ui/src/components/chat/CodeIndexPopover.tsx‎

Lines changed: 100 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -71,7 +71,7 @@ interface LocalCodeIndexSettings {
7171
}
7272

7373
// Validation schema for codebase index settings
74-
const createValidationSchema = (provider: EmbedderProvider, t: any) => {
74+
const createValidationSchema = (provider: EmbedderProvider, t: any, models: any) => {
7575
const baseSchema = z.object({
7676
codebaseIndexEnabled: z.boolean(),
7777
codebaseIndexQdrantUrl: z
@@ -115,12 +115,32 @@ const createValidationSchema = (provider: EmbedderProvider, t: any) => {
115115
})
116116

117117
case "gemini":
118-
return baseSchema.extend({
119-
codebaseIndexGeminiApiKey: z.string().min(1, t("settings:codeIndex.validation.geminiApiKeyRequired")),
120-
codebaseIndexEmbedderModelId: z
121-
.string()
122-
.min(1, t("settings:codeIndex.validation.modelSelectionRequired")),
123-
})
118+
return baseSchema
119+
.extend({
120+
codebaseIndexGeminiApiKey: z
121+
.string()
122+
.min(1, t("settings:codeIndex.validation.geminiApiKeyRequired")),
123+
codebaseIndexEmbedderModelId: z
124+
.string()
125+
.min(1, t("settings:codeIndex.validation.modelSelectionRequired")),
126+
codebaseIndexEmbedderModelDimension: z.number().optional(),
127+
})
128+
.refine(
129+
(data) => {
130+
const model = models?.gemini?.[data.codebaseIndexEmbedderModelId || ""]
131+
if (model?.minDimension && model?.maxDimension && data.codebaseIndexEmbedderModelDimension) {
132+
return (
133+
data.codebaseIndexEmbedderModelDimension >= model.minDimension &&
134+
data.codebaseIndexEmbedderModelDimension <= model.maxDimension
135+
)
136+
}
137+
return true
138+
},
139+
{
140+
message: t("settings:codeIndex.validation.invalidDimension"),
141+
path: ["codebaseIndexEmbedderModelDimension"],
142+
},
143+
)
124144

125145
default:
126146
return baseSchema
@@ -188,7 +208,13 @@ export const CodeIndexPopover: React.FC<CodeIndexPopoverProps> = ({
188208
codebaseIndexEmbedderBaseUrl: codebaseIndexConfig.codebaseIndexEmbedderBaseUrl || "",
189209
codebaseIndexEmbedderModelId: codebaseIndexConfig.codebaseIndexEmbedderModelId || "",
190210
codebaseIndexEmbedderModelDimension:
191-
codebaseIndexConfig.codebaseIndexEmbedderModelDimension || undefined,
211+
// This order is critical to prevent a UI race condition. After saving,
212+
// the component's local `currentSettings` is updated immediately, while
213+
// the global `codebaseIndexConfig` might still be stale. Prioritizing
214+
// `currentSettings` ensures the UI reflects the saved value instantly.
215+
currentSettings.codebaseIndexEmbedderModelDimension ||
216+
codebaseIndexConfig.codebaseIndexEmbedderModelDimension ||
217+
undefined,
192218
codebaseIndexSearchMaxResults:
193219
codebaseIndexConfig.codebaseIndexSearchMaxResults ?? CODEBASE_INDEX_DEFAULTS.DEFAULT_SEARCH_RESULTS,
194220
codebaseIndexSearchMinScore:
@@ -349,7 +375,7 @@ export const CodeIndexPopover: React.FC<CodeIndexPopoverProps> = ({
349375

350376
// Validation function
351377
const validateSettings = (): boolean => {
352-
const schema = createValidationSchema(currentSettings.codebaseIndexEmbedderProvider, t)
378+
const schema = createValidationSchema(currentSettings.codebaseIndexEmbedderProvider, t, codebaseIndexModels)
353379

354380
// Prepare data for validation
355381
const dataToValidate: any = {}
@@ -916,6 +942,71 @@ export const CodeIndexPopover: React.FC<CodeIndexPopoverProps> = ({
916942
</p>
917943
)}
918944
</div>
945+
{(() => {
946+
const selectedModelProfile =
947+
codebaseIndexModels?.gemini?.[
948+
currentSettings.codebaseIndexEmbedderModelId
949+
]
950+
951+
// Conditionally render the dimension slider only for Gemini models
952+
// that explicitly define a min and max dimension in their profile.
953+
if (
954+
selectedModelProfile?.minDimension &&
955+
selectedModelProfile?.maxDimension
956+
) {
957+
return (
958+
<div className="space-y-2">
959+
<div className="flex items-center gap-2">
960+
<label className="text-sm font-medium">
961+
{t("settings:codeIndex.modelDimensionLabel")}
962+
</label>
963+
</div>
964+
<div className="flex items-center gap-2">
965+
<Slider
966+
min={selectedModelProfile.minDimension}
967+
max={selectedModelProfile.maxDimension}
968+
step={1}
969+
value={[
970+
currentSettings.codebaseIndexEmbedderModelDimension ??
971+
selectedModelProfile.defaultDimension ??
972+
selectedModelProfile.minDimension,
973+
]}
974+
onValueChange={(values) =>
975+
updateSetting(
976+
"codebaseIndexEmbedderModelDimension",
977+
values[0],
978+
)
979+
}
980+
className="flex-1"
981+
data-testid="model-dimension-slider"
982+
/>
983+
<span className="w-12 text-center">
984+
{currentSettings.codebaseIndexEmbedderModelDimension ??
985+
selectedModelProfile.defaultDimension ??
986+
selectedModelProfile.minDimension}
987+
</span>
988+
<VSCodeButton
989+
appearance="icon"
990+
title={t("settings:codeIndex.resetToDefault")}
991+
onClick={() =>
992+
updateSetting(
993+
"codebaseIndexEmbedderModelDimension",
994+
selectedModelProfile.defaultDimension,
995+
)
996+
}>
997+
<span className="codicon codicon-discard" />
998+
</VSCodeButton>
999+
</div>
1000+
{formErrors.codebaseIndexEmbedderModelDimension && (
1001+
<p className="text-xs text-vscode-errorForeground mt-1 mb-0">
1002+
{formErrors.codebaseIndexEmbedderModelDimension}
1003+
</p>
1004+
)}
1005+
</div>
1006+
)
1007+
}
1008+
return null
1009+
})()}
9191010
</>
9201011
)}
9211012

0 commit comments

Comments
 (0)