Skip to content

Commit f3102f8

Browse files
committed
feat: add API key support for Ollama Embedder Provider
- Add ollamaApiKey parameter to CodeIndexOllamaEmbedder constructor - Include Authorization header in all API requests when API key is provided - Update CodeIndexConfigManager to read and store Ollama API key from secrets - Add UI field for Ollama API key input in CodeIndexPopover component - Update type definitions to include codebaseIndexOllamaApiKey secret - Add comprehensive tests for API key authentication - Maintain backward compatibility for local Ollama instances without auth Closes #8737
1 parent a8f87d2 commit f3102f8

File tree

9 files changed

+346
-11
lines changed

9 files changed

+346
-11
lines changed

packages/types/src/codebase-index.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,7 @@ export const codebaseIndexProviderSchema = z.object({
6868
codebaseIndexGeminiApiKey: z.string().optional(),
6969
codebaseIndexMistralApiKey: z.string().optional(),
7070
codebaseIndexVercelAiGatewayApiKey: z.string().optional(),
71+
codebaseIndexOllamaApiKey: z.string().optional(),
7172
})
7273

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

packages/types/src/global-settings.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -199,6 +199,7 @@ export const SECRET_STATE_KEYS = [
199199
"codebaseIndexGeminiApiKey",
200200
"codebaseIndexMistralApiKey",
201201
"codebaseIndexVercelAiGatewayApiKey",
202+
"codebaseIndexOllamaApiKey",
202203
"huggingFaceApiKey",
203204
"sambaNovaApiKey",
204205
"zaiApiKey",

src/core/webview/webviewMessageHandler.ts

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2493,6 +2493,12 @@ export const webviewMessageHandler = async (
24932493
settings.codebaseIndexVercelAiGatewayApiKey,
24942494
)
24952495
}
2496+
if (settings.codebaseIndexOllamaApiKey !== undefined) {
2497+
await provider.contextProxy.storeSecret(
2498+
"codebaseIndexOllamaApiKey",
2499+
settings.codebaseIndexOllamaApiKey,
2500+
)
2501+
}
24962502

24972503
// Send success response first - settings are saved regardless of validation
24982504
await provider.postMessageToWebview({
@@ -2630,6 +2636,7 @@ export const webviewMessageHandler = async (
26302636
const hasVercelAiGatewayApiKey = !!(await provider.context.secrets.get(
26312637
"codebaseIndexVercelAiGatewayApiKey",
26322638
))
2639+
const hasOllamaApiKey = !!(await provider.context.secrets.get("codebaseIndexOllamaApiKey"))
26332640

26342641
provider.postMessageToWebview({
26352642
type: "codeIndexSecretStatus",
@@ -2640,6 +2647,7 @@ export const webviewMessageHandler = async (
26402647
hasGeminiApiKey,
26412648
hasMistralApiKey,
26422649
hasVercelAiGatewayApiKey,
2650+
hasOllamaApiKey,
26432651
},
26442652
})
26452653
break

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

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,7 @@ export class CodeIndexConfigManager {
7171
const geminiApiKey = this.contextProxy?.getSecret("codebaseIndexGeminiApiKey") ?? ""
7272
const mistralApiKey = this.contextProxy?.getSecret("codebaseIndexMistralApiKey") ?? ""
7373
const vercelAiGatewayApiKey = this.contextProxy?.getSecret("codebaseIndexVercelAiGatewayApiKey") ?? ""
74+
const ollamaApiKey = this.contextProxy?.getSecret("codebaseIndexOllamaApiKey") ?? ""
7475

7576
// Update instance variables with configuration
7677
this.codebaseIndexEnabled = codebaseIndexEnabled ?? true
@@ -116,6 +117,7 @@ export class CodeIndexConfigManager {
116117

117118
this.ollamaOptions = {
118119
ollamaBaseUrl: codebaseIndexEmbedderBaseUrl,
120+
ollamaApiKey: ollamaApiKey || undefined,
119121
}
120122

121123
this.openAiCompatibleOptions =
@@ -162,6 +164,7 @@ export class CodeIndexConfigManager {
162164
modelDimension: this.modelDimension,
163165
openAiKey: this.openAiOptions?.openAiNativeApiKey ?? "",
164166
ollamaBaseUrl: this.ollamaOptions?.ollamaBaseUrl ?? "",
167+
ollamaApiKey: this.ollamaOptions?.ollamaApiKey ?? "",
165168
openAiCompatibleBaseUrl: this.openAiCompatibleOptions?.baseUrl ?? "",
166169
openAiCompatibleApiKey: this.openAiCompatibleOptions?.apiKey ?? "",
167170
geminiApiKey: this.geminiOptions?.apiKey ?? "",
@@ -263,6 +266,7 @@ export class CodeIndexConfigManager {
263266
const prevProvider = prev?.embedderProvider ?? "openai"
264267
const prevOpenAiKey = prev?.openAiKey ?? ""
265268
const prevOllamaBaseUrl = prev?.ollamaBaseUrl ?? ""
269+
const prevOllamaApiKey = prev?.ollamaApiKey ?? ""
266270
const prevOpenAiCompatibleBaseUrl = prev?.openAiCompatibleBaseUrl ?? ""
267271
const prevOpenAiCompatibleApiKey = prev?.openAiCompatibleApiKey ?? ""
268272
const prevModelDimension = prev?.modelDimension
@@ -301,6 +305,7 @@ export class CodeIndexConfigManager {
301305
// Authentication changes (API keys)
302306
const currentOpenAiKey = this.openAiOptions?.openAiNativeApiKey ?? ""
303307
const currentOllamaBaseUrl = this.ollamaOptions?.ollamaBaseUrl ?? ""
308+
const currentOllamaApiKey = this.ollamaOptions?.ollamaApiKey ?? ""
304309
const currentOpenAiCompatibleBaseUrl = this.openAiCompatibleOptions?.baseUrl ?? ""
305310
const currentOpenAiCompatibleApiKey = this.openAiCompatibleOptions?.apiKey ?? ""
306311
const currentModelDimension = this.modelDimension
@@ -314,7 +319,7 @@ export class CodeIndexConfigManager {
314319
return true
315320
}
316321

317-
if (prevOllamaBaseUrl !== currentOllamaBaseUrl) {
322+
if (prevOllamaBaseUrl !== currentOllamaBaseUrl || prevOllamaApiKey !== currentOllamaApiKey) {
318323
return true
319324
}
320325

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

Lines changed: 268 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -81,6 +81,15 @@ describe("CodeIndexOllamaEmbedder", () => {
8181
expect(embedderWithDefaults.embedderInfo.name).toBe("ollama")
8282
})
8383

84+
it("should initialize with API key when provided", () => {
85+
const embedderWithApiKey = new CodeIndexOllamaEmbedder({
86+
ollamaModelId: "nomic-embed-text",
87+
ollamaBaseUrl: "http://localhost:11434",
88+
ollamaApiKey: "test-api-key-123",
89+
})
90+
expect(embedderWithApiKey.embedderInfo.name).toBe("ollama")
91+
})
92+
8493
it("should normalize URLs with trailing slashes", async () => {
8594
// Create embedder with URL that has a trailing slash
8695
const embedderWithTrailingSlash = new CodeIndexOllamaEmbedder({
@@ -166,6 +175,128 @@ describe("CodeIndexOllamaEmbedder", () => {
166175
})
167176
})
168177

178+
describe("API Key Authentication", () => {
179+
it("should include Authorization header when API key is provided", async () => {
180+
const embedderWithApiKey = new CodeIndexOllamaEmbedder({
181+
ollamaModelId: "nomic-embed-text",
182+
ollamaBaseUrl: "http://localhost:11434",
183+
ollamaApiKey: "test-api-key-123",
184+
})
185+
186+
// Mock successful /api/tags call
187+
mockFetch.mockImplementationOnce(() =>
188+
Promise.resolve({
189+
ok: true,
190+
status: 200,
191+
json: () =>
192+
Promise.resolve({
193+
models: [{ name: "nomic-embed-text" }],
194+
}),
195+
} as Response),
196+
)
197+
198+
// Mock successful /api/embed test call
199+
mockFetch.mockImplementationOnce(() =>
200+
Promise.resolve({
201+
ok: true,
202+
status: 200,
203+
json: () =>
204+
Promise.resolve({
205+
embeddings: [[0.1, 0.2, 0.3]],
206+
}),
207+
} as Response),
208+
)
209+
210+
await embedderWithApiKey.validateConfiguration()
211+
212+
// Check that Authorization header is included in both calls
213+
expect(mockFetch).toHaveBeenCalledTimes(2)
214+
215+
// First call to /api/tags
216+
expect(mockFetch.mock.calls[0][1]?.headers).toEqual({
217+
"Content-Type": "application/json",
218+
Authorization: "Bearer test-api-key-123",
219+
})
220+
221+
// Second call to /api/embed
222+
expect(mockFetch.mock.calls[1][1]?.headers).toEqual({
223+
"Content-Type": "application/json",
224+
Authorization: "Bearer test-api-key-123",
225+
})
226+
})
227+
228+
it("should not include Authorization header when API key is not provided", async () => {
229+
// Mock successful /api/tags call
230+
mockFetch.mockImplementationOnce(() =>
231+
Promise.resolve({
232+
ok: true,
233+
status: 200,
234+
json: () =>
235+
Promise.resolve({
236+
models: [{ name: "nomic-embed-text" }],
237+
}),
238+
} as Response),
239+
)
240+
241+
// Mock successful /api/embed test call
242+
mockFetch.mockImplementationOnce(() =>
243+
Promise.resolve({
244+
ok: true,
245+
status: 200,
246+
json: () =>
247+
Promise.resolve({
248+
embeddings: [[0.1, 0.2, 0.3]],
249+
}),
250+
} as Response),
251+
)
252+
253+
await embedder.validateConfiguration()
254+
255+
// Check that Authorization header is NOT included
256+
expect(mockFetch).toHaveBeenCalledTimes(2)
257+
258+
// First call to /api/tags
259+
expect(mockFetch.mock.calls[0][1]?.headers).toEqual({
260+
"Content-Type": "application/json",
261+
})
262+
263+
// Second call to /api/embed
264+
expect(mockFetch.mock.calls[1][1]?.headers).toEqual({
265+
"Content-Type": "application/json",
266+
})
267+
})
268+
269+
it("should handle authentication errors with API key", async () => {
270+
const embedderWithApiKey = new CodeIndexOllamaEmbedder({
271+
ollamaModelId: "nomic-embed-text",
272+
ollamaBaseUrl: "http://localhost:11434",
273+
ollamaApiKey: "invalid-api-key",
274+
})
275+
276+
// Mock 401 Unauthorized response
277+
mockFetch.mockImplementationOnce(() =>
278+
Promise.resolve({
279+
ok: false,
280+
status: 401,
281+
} as Response),
282+
)
283+
284+
const result = await embedderWithApiKey.validateConfiguration()
285+
286+
expect(result.valid).toBe(false)
287+
expect(result.error).toBe("embeddings:ollama.serviceUnavailable")
288+
expect(mockFetch).toHaveBeenCalledWith(
289+
"http://localhost:11434/api/tags",
290+
expect.objectContaining({
291+
headers: {
292+
"Content-Type": "application/json",
293+
Authorization: "Bearer invalid-api-key",
294+
},
295+
}),
296+
)
297+
})
298+
})
299+
169300
describe("validateConfiguration", () => {
170301
it("should validate successfully when service is available and model exists", async () => {
171302
// Mock successful /api/tags call
@@ -323,5 +454,142 @@ describe("CodeIndexOllamaEmbedder", () => {
323454
expect(result.valid).toBe(false)
324455
expect(result.error).toBe("Network timeout")
325456
})
457+
458+
describe("createEmbeddings", () => {
459+
it("should create embeddings successfully without API key", async () => {
460+
const texts = ["Hello world", "Test embedding"]
461+
462+
// Mock successful /api/embed call
463+
mockFetch.mockImplementationOnce(() =>
464+
Promise.resolve({
465+
ok: true,
466+
status: 200,
467+
json: () =>
468+
Promise.resolve({
469+
embeddings: [
470+
[0.1, 0.2, 0.3],
471+
[0.4, 0.5, 0.6],
472+
],
473+
}),
474+
} as Response),
475+
)
476+
477+
const result = await embedder.createEmbeddings(texts)
478+
479+
expect(result).toEqual({
480+
embeddings: [
481+
[0.1, 0.2, 0.3],
482+
[0.4, 0.5, 0.6],
483+
],
484+
})
485+
486+
expect(mockFetch).toHaveBeenCalledWith(
487+
"http://localhost:11434/api/embed",
488+
expect.objectContaining({
489+
method: "POST",
490+
headers: {
491+
"Content-Type": "application/json",
492+
},
493+
body: JSON.stringify({
494+
model: "nomic-embed-text",
495+
input: texts,
496+
}),
497+
}),
498+
)
499+
})
500+
501+
it("should create embeddings with API key in Authorization header", async () => {
502+
const embedderWithApiKey = new CodeIndexOllamaEmbedder({
503+
ollamaModelId: "nomic-embed-text",
504+
ollamaBaseUrl: "http://localhost:11434",
505+
ollamaApiKey: "test-api-key-123",
506+
})
507+
508+
const texts = ["Hello world", "Test embedding"]
509+
510+
// Mock successful /api/embed call
511+
mockFetch.mockImplementationOnce(() =>
512+
Promise.resolve({
513+
ok: true,
514+
status: 200,
515+
json: () =>
516+
Promise.resolve({
517+
embeddings: [
518+
[0.1, 0.2, 0.3],
519+
[0.4, 0.5, 0.6],
520+
],
521+
}),
522+
} as Response),
523+
)
524+
525+
const result = await embedderWithApiKey.createEmbeddings(texts)
526+
527+
expect(result).toEqual({
528+
embeddings: [
529+
[0.1, 0.2, 0.3],
530+
[0.4, 0.5, 0.6],
531+
],
532+
})
533+
534+
// Verify Authorization header is included
535+
expect(mockFetch).toHaveBeenCalledWith(
536+
"http://localhost:11434/api/embed",
537+
expect.objectContaining({
538+
method: "POST",
539+
headers: {
540+
"Content-Type": "application/json",
541+
Authorization: "Bearer test-api-key-123",
542+
},
543+
body: JSON.stringify({
544+
model: "nomic-embed-text",
545+
input: texts,
546+
}),
547+
}),
548+
)
549+
})
550+
551+
it("should handle authentication error when creating embeddings", async () => {
552+
const embedderWithApiKey = new CodeIndexOllamaEmbedder({
553+
ollamaModelId: "nomic-embed-text",
554+
ollamaBaseUrl: "http://localhost:11434",
555+
ollamaApiKey: "invalid-api-key",
556+
})
557+
558+
const texts = ["Hello world"]
559+
560+
// Mock 401 Unauthorized response
561+
mockFetch.mockImplementationOnce(() =>
562+
Promise.resolve({
563+
ok: false,
564+
status: 401,
565+
statusText: "Unauthorized",
566+
} as Response),
567+
)
568+
569+
await expect(embedderWithApiKey.createEmbeddings(texts)).rejects.toThrow(
570+
"embeddings:ollama.embeddingFailed",
571+
)
572+
573+
// Verify request included the API key
574+
expect(mockFetch).toHaveBeenCalledWith(
575+
"http://localhost:11434/api/embed",
576+
expect.objectContaining({
577+
headers: {
578+
"Content-Type": "application/json",
579+
Authorization: "Bearer invalid-api-key",
580+
},
581+
}),
582+
)
583+
})
584+
585+
it("should handle network errors when creating embeddings", async () => {
586+
const texts = ["Hello world"]
587+
588+
// Mock network error
589+
mockFetch.mockRejectedValueOnce(new Error("Network error"))
590+
591+
await expect(embedder.createEmbeddings(texts)).rejects.toThrow("embeddings:ollama.embeddingFailed")
592+
})
593+
})
326594
})
327595
})

0 commit comments

Comments
 (0)