Skip to content

Commit 5583e0f

Browse files
author
Roo
committed
feat: add custom base URL support for Gemini embedding provider
- Add codebaseIndexGeminiBaseUrl field to type definitions - Update ConfigManager to handle Gemini base URL configuration - Add UI field for Gemini base URL in CodeIndexPopover - Add validation for optional Gemini base URL - Update service factory to pass base URL to GeminiEmbedder - Add comprehensive test coverage for new functionality Fixes #5744
1 parent 29b7d06 commit 5583e0f

File tree

9 files changed

+432
-9
lines changed

9 files changed

+432
-9
lines changed

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,8 @@ export const codebaseIndexConfigSchema = z.object({
3434
// OpenAI Compatible specific fields
3535
codebaseIndexOpenAiCompatibleBaseUrl: z.string().optional(),
3636
codebaseIndexOpenAiCompatibleModelDimension: z.number().optional(),
37+
// Gemini specific fields
38+
codebaseIndexGeminiBaseUrl: z.string().optional(),
3739
})
3840

3941
export type CodebaseIndexConfig = z.infer<typeof codebaseIndexConfigSchema>
@@ -62,6 +64,7 @@ export const codebaseIndexProviderSchema = z.object({
6264
codebaseIndexOpenAiCompatibleApiKey: z.string().optional(),
6365
codebaseIndexOpenAiCompatibleModelDimension: z.number().optional(),
6466
codebaseIndexGeminiApiKey: z.string().optional(),
67+
codebaseIndexGeminiBaseUrl: z.string().optional(),
6568
})
6669

6770
export type CodebaseIndexProvider = z.infer<typeof codebaseIndexProviderSchema>

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

Lines changed: 145 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1239,6 +1239,151 @@ describe("CodeIndexConfigManager", () => {
12391239
expect(configManager.isFeatureConfigured).toBe(true)
12401240
})
12411241

1242+
it("should load Gemini configuration with custom base URL from globalState", async () => {
1243+
const mockGlobalState = {
1244+
codebaseIndexEnabled: true,
1245+
codebaseIndexQdrantUrl: "http://qdrant.local",
1246+
codebaseIndexEmbedderProvider: "gemini",
1247+
codebaseIndexEmbedderModelId: "text-embedding-004",
1248+
codebaseIndexGeminiBaseUrl: "https://custom-gemini-proxy.example.com/v1beta/openai/",
1249+
}
1250+
mockContextProxy.getGlobalState.mockImplementation((key: string) => {
1251+
if (key === "codebaseIndexConfig") return mockGlobalState
1252+
return undefined
1253+
})
1254+
1255+
setupSecretMocks({
1256+
codebaseIndexGeminiApiKey: "test-gemini-key",
1257+
codeIndexQdrantApiKey: "test-qdrant-key",
1258+
})
1259+
1260+
const result = await configManager.loadConfiguration()
1261+
1262+
expect(result.currentConfig).toEqual({
1263+
isConfigured: true,
1264+
embedderProvider: "gemini",
1265+
modelId: "text-embedding-004",
1266+
modelDimension: undefined,
1267+
openAiOptions: { openAiNativeApiKey: "" },
1268+
ollamaOptions: { ollamaBaseUrl: undefined },
1269+
openAiCompatibleOptions: undefined,
1270+
geminiOptions: {
1271+
apiKey: "test-gemini-key",
1272+
baseUrl: "https://custom-gemini-proxy.example.com/v1beta/openai/",
1273+
},
1274+
qdrantUrl: "http://qdrant.local",
1275+
qdrantApiKey: "test-qdrant-key",
1276+
searchMinScore: 0.4,
1277+
})
1278+
})
1279+
1280+
it("should load Gemini configuration without custom base URL (default)", async () => {
1281+
const mockGlobalState = {
1282+
codebaseIndexEnabled: true,
1283+
codebaseIndexQdrantUrl: "http://qdrant.local",
1284+
codebaseIndexEmbedderProvider: "gemini",
1285+
codebaseIndexEmbedderModelId: "text-embedding-004",
1286+
// No custom base URL specified
1287+
}
1288+
mockContextProxy.getGlobalState.mockImplementation((key: string) => {
1289+
if (key === "codebaseIndexConfig") return mockGlobalState
1290+
return undefined
1291+
})
1292+
1293+
setupSecretMocks({
1294+
codebaseIndexGeminiApiKey: "test-gemini-key",
1295+
codeIndexQdrantApiKey: "test-qdrant-key",
1296+
})
1297+
1298+
const result = await configManager.loadConfiguration()
1299+
1300+
expect(result.currentConfig).toEqual({
1301+
isConfigured: true,
1302+
embedderProvider: "gemini",
1303+
modelId: "text-embedding-004",
1304+
modelDimension: undefined,
1305+
openAiOptions: { openAiNativeApiKey: "" },
1306+
ollamaOptions: { ollamaBaseUrl: undefined },
1307+
openAiCompatibleOptions: undefined,
1308+
geminiOptions: {
1309+
apiKey: "test-gemini-key",
1310+
baseUrl: "", // Should be empty string when not specified
1311+
},
1312+
qdrantUrl: "http://qdrant.local",
1313+
qdrantApiKey: "test-qdrant-key",
1314+
searchMinScore: 0.4,
1315+
})
1316+
})
1317+
1318+
it("should detect restart requirement when Gemini base URL changes", async () => {
1319+
// Initial state with no custom base URL
1320+
mockContextProxy.getGlobalState.mockImplementation((key: string) => {
1321+
if (key === "codebaseIndexConfig") {
1322+
return {
1323+
codebaseIndexEnabled: true,
1324+
codebaseIndexQdrantUrl: "http://qdrant.local",
1325+
codebaseIndexEmbedderProvider: "gemini",
1326+
codebaseIndexEmbedderModelId: "text-embedding-004",
1327+
}
1328+
}
1329+
return undefined
1330+
})
1331+
setupSecretMocks({
1332+
codebaseIndexGeminiApiKey: "test-gemini-key",
1333+
codeIndexQdrantApiKey: "test-qdrant-key",
1334+
})
1335+
1336+
await configManager.loadConfiguration()
1337+
1338+
// Change to custom base URL
1339+
mockContextProxy.getGlobalState.mockImplementation((key: string) => {
1340+
if (key === "codebaseIndexConfig") {
1341+
return {
1342+
codebaseIndexEnabled: true,
1343+
codebaseIndexQdrantUrl: "http://qdrant.local",
1344+
codebaseIndexEmbedderProvider: "gemini",
1345+
codebaseIndexEmbedderModelId: "text-embedding-004",
1346+
codebaseIndexGeminiBaseUrl: "https://custom-gemini-proxy.example.com/v1beta/openai/",
1347+
}
1348+
}
1349+
return undefined
1350+
})
1351+
1352+
const result = await configManager.loadConfiguration()
1353+
expect(result.requiresRestart).toBe(true)
1354+
})
1355+
1356+
it("should detect restart requirement when Gemini API key changes", async () => {
1357+
// Initial state
1358+
mockContextProxy.getGlobalState.mockImplementation((key: string) => {
1359+
if (key === "codebaseIndexConfig") {
1360+
return {
1361+
codebaseIndexEnabled: true,
1362+
codebaseIndexQdrantUrl: "http://qdrant.local",
1363+
codebaseIndexEmbedderProvider: "gemini",
1364+
codebaseIndexEmbedderModelId: "text-embedding-004",
1365+
codebaseIndexGeminiBaseUrl: "https://custom-gemini-proxy.example.com/v1beta/openai/",
1366+
}
1367+
}
1368+
return undefined
1369+
})
1370+
setupSecretMocks({
1371+
codebaseIndexGeminiApiKey: "old-gemini-key",
1372+
codeIndexQdrantApiKey: "test-qdrant-key",
1373+
})
1374+
1375+
await configManager.loadConfiguration()
1376+
1377+
// Change Gemini API key
1378+
setupSecretMocks({
1379+
codebaseIndexGeminiApiKey: "new-gemini-key",
1380+
codeIndexQdrantApiKey: "test-qdrant-key",
1381+
})
1382+
1383+
const result = await configManager.loadConfiguration()
1384+
expect(result.requiresRestart).toBe(true)
1385+
})
1386+
12421387
it("should return false when Gemini API key is missing", async () => {
12431388
mockContextProxy.getGlobalState.mockImplementation((key: string) => {
12441389
if (key === "codebaseIndexConfig") {

‎src/services/code-index/__tests__/service-factory.spec.ts‎

Lines changed: 44 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -279,7 +279,7 @@ describe("CodeIndexServiceFactory", () => {
279279
factory.createEmbedder()
280280

281281
// Assert
282-
expect(MockedGeminiEmbedder).toHaveBeenCalledWith("test-gemini-api-key", undefined)
282+
expect(MockedGeminiEmbedder).toHaveBeenCalledWith("test-gemini-api-key", undefined, undefined)
283283
})
284284

285285
it("should create GeminiEmbedder with specified modelId", () => {
@@ -297,7 +297,49 @@ describe("CodeIndexServiceFactory", () => {
297297
factory.createEmbedder()
298298

299299
// Assert
300-
expect(MockedGeminiEmbedder).toHaveBeenCalledWith("test-gemini-api-key", "text-embedding-004")
300+
expect(MockedGeminiEmbedder).toHaveBeenCalledWith("test-gemini-api-key", "text-embedding-004", undefined)
301+
})
302+
303+
it("should create GeminiEmbedder with custom base URL", () => {
304+
// Arrange
305+
const testConfig = {
306+
embedderProvider: "gemini",
307+
modelId: "text-embedding-004",
308+
geminiOptions: {
309+
apiKey: "test-gemini-api-key",
310+
baseUrl: "https://custom-gemini-proxy.example.com/v1beta/openai/",
311+
},
312+
}
313+
mockConfigManager.getConfig.mockReturnValue(testConfig as any)
314+
315+
// Act
316+
factory.createEmbedder()
317+
318+
// Assert
319+
expect(MockedGeminiEmbedder).toHaveBeenCalledWith(
320+
"test-gemini-api-key",
321+
"text-embedding-004",
322+
"https://custom-gemini-proxy.example.com/v1beta/openai/",
323+
)
324+
})
325+
326+
it("should create GeminiEmbedder without base URL when not specified", () => {
327+
// Arrange
328+
const testConfig = {
329+
embedderProvider: "gemini",
330+
modelId: "text-embedding-004",
331+
geminiOptions: {
332+
apiKey: "test-gemini-api-key",
333+
// baseUrl not specified
334+
},
335+
}
336+
mockConfigManager.getConfig.mockReturnValue(testConfig as any)
337+
338+
// Act
339+
factory.createEmbedder()
340+
341+
// Assert
342+
expect(MockedGeminiEmbedder).toHaveBeenCalledWith("test-gemini-api-key", "text-embedding-004", undefined)
301343
})
302344

303345
it("should throw error when Gemini API key is missing", () => {

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

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ export class CodeIndexConfigManager {
1717
private openAiOptions?: ApiHandlerOptions
1818
private ollamaOptions?: ApiHandlerOptions
1919
private openAiCompatibleOptions?: { baseUrl: string; apiKey: string }
20-
private geminiOptions?: { apiKey: string }
20+
private geminiOptions?: { apiKey: string; baseUrl?: string }
2121
private qdrantUrl?: string = "http://localhost:6333"
2222
private qdrantApiKey?: string
2323
private searchMinScore?: number
@@ -49,6 +49,9 @@ export class CodeIndexConfigManager {
4949
codebaseIndexEmbedderModelId: "",
5050
codebaseIndexSearchMinScore: undefined,
5151
codebaseIndexSearchMaxResults: undefined,
52+
codebaseIndexOpenAiCompatibleBaseUrl: "",
53+
codebaseIndexEmbedderModelDimension: undefined,
54+
codebaseIndexGeminiBaseUrl: "",
5255
}
5356

5457
const {
@@ -67,6 +70,7 @@ export class CodeIndexConfigManager {
6770
const openAiCompatibleBaseUrl = codebaseIndexConfig.codebaseIndexOpenAiCompatibleBaseUrl ?? ""
6871
const openAiCompatibleApiKey = this.contextProxy?.getSecret("codebaseIndexOpenAiCompatibleApiKey") ?? ""
6972
const geminiApiKey = this.contextProxy?.getSecret("codebaseIndexGeminiApiKey") ?? ""
73+
const geminiBaseUrl = codebaseIndexConfig.codebaseIndexGeminiBaseUrl ?? ""
7074

7175
// Update instance variables with configuration
7276
this.codebaseIndexEnabled = codebaseIndexEnabled ?? true
@@ -118,7 +122,7 @@ export class CodeIndexConfigManager {
118122
}
119123
: undefined
120124

121-
this.geminiOptions = geminiApiKey ? { apiKey: geminiApiKey } : undefined
125+
this.geminiOptions = geminiApiKey ? { apiKey: geminiApiKey, baseUrl: geminiBaseUrl } : undefined
122126
}
123127

124128
/**
@@ -134,7 +138,7 @@ export class CodeIndexConfigManager {
134138
openAiOptions?: ApiHandlerOptions
135139
ollamaOptions?: ApiHandlerOptions
136140
openAiCompatibleOptions?: { baseUrl: string; apiKey: string }
137-
geminiOptions?: { apiKey: string }
141+
geminiOptions?: { apiKey: string; baseUrl?: string }
138142
qdrantUrl?: string
139143
qdrantApiKey?: string
140144
searchMinScore?: number
@@ -153,6 +157,7 @@ export class CodeIndexConfigManager {
153157
openAiCompatibleBaseUrl: this.openAiCompatibleOptions?.baseUrl ?? "",
154158
openAiCompatibleApiKey: this.openAiCompatibleOptions?.apiKey ?? "",
155159
geminiApiKey: this.geminiOptions?.apiKey ?? "",
160+
geminiBaseUrl: this.geminiOptions?.baseUrl ?? "",
156161
qdrantUrl: this.qdrantUrl ?? "",
157162
qdrantApiKey: this.qdrantApiKey ?? "",
158163
}
@@ -241,6 +246,7 @@ export class CodeIndexConfigManager {
241246
const prevOpenAiCompatibleApiKey = prev?.openAiCompatibleApiKey ?? ""
242247
const prevModelDimension = prev?.modelDimension
243248
const prevGeminiApiKey = prev?.geminiApiKey ?? ""
249+
const prevGeminiBaseUrl = prev?.geminiBaseUrl ?? ""
244250
const prevQdrantUrl = prev?.qdrantUrl ?? ""
245251
const prevQdrantApiKey = prev?.qdrantApiKey ?? ""
246252

@@ -277,6 +283,7 @@ export class CodeIndexConfigManager {
277283
const currentOpenAiCompatibleApiKey = this.openAiCompatibleOptions?.apiKey ?? ""
278284
const currentModelDimension = this.modelDimension
279285
const currentGeminiApiKey = this.geminiOptions?.apiKey ?? ""
286+
const currentGeminiBaseUrl = this.geminiOptions?.baseUrl ?? ""
280287
const currentQdrantUrl = this.qdrantUrl ?? ""
281288
const currentQdrantApiKey = this.qdrantApiKey ?? ""
282289

@@ -295,6 +302,10 @@ export class CodeIndexConfigManager {
295302
return true
296303
}
297304

305+
if (prevGeminiApiKey !== currentGeminiApiKey || prevGeminiBaseUrl !== currentGeminiBaseUrl) {
306+
return true
307+
}
308+
298309
// Check for model dimension changes (generic for all providers)
299310
if (prevModelDimension !== currentModelDimension) {
300311
return true

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

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -23,18 +23,22 @@ export class GeminiEmbedder implements IEmbedder {
2323
* Creates a new Gemini embedder
2424
* @param apiKey The Gemini API key for authentication
2525
* @param modelId The model ID to use (defaults to gemini-embedding-001)
26+
* @param baseUrl The base URL for the Gemini API (defaults to Google's official endpoint)
2627
*/
27-
constructor(apiKey: string, modelId?: string) {
28+
constructor(apiKey: string, modelId?: string, baseUrl?: string) {
2829
if (!apiKey) {
2930
throw new Error(t("embeddings:validation.apiKeyRequired"))
3031
}
3132

3233
// Use provided model or default
3334
this.modelId = modelId || GeminiEmbedder.DEFAULT_MODEL
3435

36+
// Use provided base URL or default
37+
const geminiBaseUrl = baseUrl || GeminiEmbedder.GEMINI_BASE_URL
38+
3539
// Create an OpenAI Compatible embedder with Gemini's configuration
3640
this.openAICompatibleEmbedder = new OpenAICompatibleEmbedder(
37-
GeminiEmbedder.GEMINI_BASE_URL,
41+
geminiBaseUrl,
3842
apiKey,
3943
this.modelId,
4044
GEMINI_MAX_ITEM_TOKENS,

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ export interface CodeIndexConfig {
1212
openAiOptions?: ApiHandlerOptions
1313
ollamaOptions?: ApiHandlerOptions
1414
openAiCompatibleOptions?: { baseUrl: string; apiKey: string }
15-
geminiOptions?: { apiKey: string }
15+
geminiOptions?: { apiKey: string; baseUrl?: string }
1616
qdrantUrl?: string
1717
qdrantApiKey?: string
1818
searchMinScore?: number
@@ -33,6 +33,7 @@ export type PreviousConfigSnapshot = {
3333
openAiCompatibleBaseUrl?: string
3434
openAiCompatibleApiKey?: string
3535
geminiApiKey?: string
36+
geminiBaseUrl?: string
3637
qdrantUrl?: string
3738
qdrantApiKey?: string
3839
}

‎src/services/code-index/service-factory.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -63,7 +63,7 @@ export class CodeIndexServiceFactory {
6363
if (!config.geminiOptions?.apiKey) {
6464
throw new Error(t("embeddings:serviceFactory.geminiConfigMissing"))
6565
}
66-
return new GeminiEmbedder(config.geminiOptions.apiKey, config.modelId)
66+
return new GeminiEmbedder(config.geminiOptions.apiKey, config.modelId, config.geminiOptions.baseUrl)
6767
}
6868

6969
throw new Error(

0 commit comments

Comments
 (0)