Skip to content

Commit d089fb2

Browse files
roomoteroomote
committed
feat: add custom headers support for OpenAI compatible embeddings
- Add codebaseIndexOpenAiCompatibleHeaders field to types - Update config manager to handle custom headers - Modify OpenAICompatibleEmbedder to accept and use custom headers - Update service factory to pass headers to embedder - Add comprehensive test coverage for custom headers Fixes #8909
1 parent ff0c65a commit d089fb2

File tree

6 files changed

+142
-13
lines changed

6 files changed

+142
-13
lines changed

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ export const codebaseIndexConfigSchema = z.object({
3636
// OpenAI Compatible specific fields
3737
codebaseIndexOpenAiCompatibleBaseUrl: z.string().optional(),
3838
codebaseIndexOpenAiCompatibleModelDimension: z.number().optional(),
39+
codebaseIndexOpenAiCompatibleHeaders: z.record(z.string()).optional(),
3940
})
4041

4142
export type CodebaseIndexConfig = z.infer<typeof codebaseIndexConfigSchema>
@@ -65,6 +66,7 @@ export const codebaseIndexProviderSchema = z.object({
6566
codebaseIndexOpenAiCompatibleBaseUrl: z.string().optional(),
6667
codebaseIndexOpenAiCompatibleApiKey: z.string().optional(),
6768
codebaseIndexOpenAiCompatibleModelDimension: z.number().optional(),
69+
codebaseIndexOpenAiCompatibleHeaders: z.record(z.string()).optional(),
6870
codebaseIndexGeminiApiKey: z.string().optional(),
6971
codebaseIndexMistralApiKey: z.string().optional(),
7072
codebaseIndexVercelAiGatewayApiKey: z.string().optional(),

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

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ export class CodeIndexConfigManager {
1616
private modelDimension?: number
1717
private openAiOptions?: ApiHandlerOptions
1818
private ollamaOptions?: ApiHandlerOptions
19-
private openAiCompatibleOptions?: { baseUrl: string; apiKey: string }
19+
private openAiCompatibleOptions?: { baseUrl: string; apiKey: string; headers?: Record<string, string> }
2020
private geminiOptions?: { apiKey: string }
2121
private mistralOptions?: { apiKey: string }
2222
private vercelAiGatewayOptions?: { apiKey: string }
@@ -68,6 +68,7 @@ export class CodeIndexConfigManager {
6868
// Fix: Read OpenAI Compatible settings from the correct location within codebaseIndexConfig
6969
const openAiCompatibleBaseUrl = codebaseIndexConfig.codebaseIndexOpenAiCompatibleBaseUrl ?? ""
7070
const openAiCompatibleApiKey = this.contextProxy?.getSecret("codebaseIndexOpenAiCompatibleApiKey") ?? ""
71+
const openAiCompatibleHeaders = codebaseIndexConfig.codebaseIndexOpenAiCompatibleHeaders ?? undefined
7172
const geminiApiKey = this.contextProxy?.getSecret("codebaseIndexGeminiApiKey") ?? ""
7273
const mistralApiKey = this.contextProxy?.getSecret("codebaseIndexMistralApiKey") ?? ""
7374
const vercelAiGatewayApiKey = this.contextProxy?.getSecret("codebaseIndexVercelAiGatewayApiKey") ?? ""
@@ -123,6 +124,7 @@ export class CodeIndexConfigManager {
123124
? {
124125
baseUrl: openAiCompatibleBaseUrl,
125126
apiKey: openAiCompatibleApiKey,
127+
headers: openAiCompatibleHeaders,
126128
}
127129
: undefined
128130

@@ -143,7 +145,7 @@ export class CodeIndexConfigManager {
143145
modelDimension?: number
144146
openAiOptions?: ApiHandlerOptions
145147
ollamaOptions?: ApiHandlerOptions
146-
openAiCompatibleOptions?: { baseUrl: string; apiKey: string }
148+
openAiCompatibleOptions?: { baseUrl: string; apiKey: string; headers?: Record<string, string> }
147149
geminiOptions?: { apiKey: string }
148150
mistralOptions?: { apiKey: string }
149151
vercelAiGatewayOptions?: { apiKey: string }

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

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -112,6 +112,31 @@ describe("OpenAICompatibleEmbedder", () => {
112112
expect(embedder).toBeDefined()
113113
})
114114

115+
it("should create embedder with custom headers", () => {
116+
const customHeaders = {
117+
"X-Custom-Header": "custom-value",
118+
"X-Another-Header": "another-value",
119+
}
120+
embedder = new OpenAICompatibleEmbedder(testBaseUrl, testApiKey, testModelId, undefined, customHeaders)
121+
122+
expect(MockedOpenAI).toHaveBeenCalledWith({
123+
baseURL: testBaseUrl,
124+
apiKey: testApiKey,
125+
defaultHeaders: customHeaders,
126+
})
127+
expect(embedder).toBeDefined()
128+
})
129+
130+
it("should create embedder without custom headers when not provided", () => {
131+
embedder = new OpenAICompatibleEmbedder(testBaseUrl, testApiKey, testModelId, undefined, undefined)
132+
133+
expect(MockedOpenAI).toHaveBeenCalledWith({
134+
baseURL: testBaseUrl,
135+
apiKey: testApiKey,
136+
})
137+
expect(embedder).toBeDefined()
138+
})
139+
115140
it("should throw error when baseUrl is missing", () => {
116141
expect(() => new OpenAICompatibleEmbedder("", testApiKey, testModelId)).toThrow(
117142
"embeddings:validation.baseUrlRequired",
@@ -813,6 +838,81 @@ describe("OpenAICompatibleEmbedder", () => {
813838
expect(baseResult.embeddings[0]).toEqual([0.4, 0.5, 0.6])
814839
})
815840

841+
it("should include custom headers in direct fetch requests", async () => {
842+
const testTexts = ["Test text"]
843+
const customHeaders = {
844+
"X-Custom-Header": "custom-value",
845+
"X-API-Version": "v2",
846+
}
847+
const base64String = createBase64Embedding([0.1, 0.2, 0.3])
848+
849+
// Test Azure URL with custom headers (direct fetch)
850+
const azureEmbedder = new OpenAICompatibleEmbedder(
851+
azureUrl,
852+
testApiKey,
853+
testModelId,
854+
undefined,
855+
customHeaders,
856+
)
857+
const mockFetchResponse = createMockResponse({
858+
data: [{ embedding: base64String }],
859+
usage: { prompt_tokens: 10, total_tokens: 15 },
860+
})
861+
;(global.fetch as MockedFunction<typeof fetch>).mockResolvedValue(mockFetchResponse as any)
862+
863+
const azureResult = await azureEmbedder.createEmbeddings(testTexts)
864+
expect(global.fetch).toHaveBeenCalledWith(
865+
azureUrl,
866+
expect.objectContaining({
867+
method: "POST",
868+
headers: expect.objectContaining({
869+
"Content-Type": "application/json",
870+
"api-key": testApiKey,
871+
Authorization: `Bearer ${testApiKey}`,
872+
"X-Custom-Header": "custom-value",
873+
"X-API-Version": "v2",
874+
}),
875+
}),
876+
)
877+
expect(mockEmbeddingsCreate).not.toHaveBeenCalled()
878+
expectEmbeddingValues(azureResult.embeddings[0], [0.1, 0.2, 0.3])
879+
})
880+
881+
it("should handle custom headers that override default headers", async () => {
882+
const testTexts = ["Test text"]
883+
const customHeaders = {
884+
"api-key": "override-key", // Override the default api-key
885+
"X-Custom-Header": "custom-value",
886+
}
887+
const base64String = createBase64Embedding([0.1, 0.2, 0.3])
888+
889+
const azureEmbedder = new OpenAICompatibleEmbedder(
890+
azureUrl,
891+
testApiKey,
892+
testModelId,
893+
undefined,
894+
customHeaders,
895+
)
896+
const mockFetchResponse = createMockResponse({
897+
data: [{ embedding: base64String }],
898+
usage: { prompt_tokens: 10, total_tokens: 15 },
899+
})
900+
;(global.fetch as MockedFunction<typeof fetch>).mockResolvedValue(mockFetchResponse as any)
901+
902+
const azureResult = await azureEmbedder.createEmbeddings(testTexts)
903+
expect(global.fetch).toHaveBeenCalledWith(
904+
azureUrl,
905+
expect.objectContaining({
906+
method: "POST",
907+
headers: expect.objectContaining({
908+
"api-key": "override-key", // Custom header overrides default
909+
"X-Custom-Header": "custom-value",
910+
}),
911+
}),
912+
)
913+
expectEmbeddingValues(azureResult.embeddings[0], [0.1, 0.2, 0.3])
914+
})
915+
816916
it.each([
817917
[401, "Authentication failed. Please check your API key."],
818918
[500, "Failed to create embeddings after 3 attempts"],

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

Lines changed: 33 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ export class OpenAICompatibleEmbedder implements IEmbedder {
3737
private readonly defaultModelId: string
3838
private readonly baseUrl: string
3939
private readonly apiKey: string
40+
private readonly customHeaders?: Record<string, string>
4041
private readonly isFullUrl: boolean
4142
private readonly maxItemTokens: number
4243

@@ -56,8 +57,15 @@ export class OpenAICompatibleEmbedder implements IEmbedder {
5657
* @param apiKey The API key for authentication
5758
* @param modelId Optional model identifier (defaults to "text-embedding-3-small")
5859
* @param maxItemTokens Optional maximum tokens per item (defaults to MAX_ITEM_TOKENS)
60+
* @param customHeaders Optional custom headers to include in requests
5961
*/
60-
constructor(baseUrl: string, apiKey: string, modelId?: string, maxItemTokens?: number) {
62+
constructor(
63+
baseUrl: string,
64+
apiKey: string,
65+
modelId?: string,
66+
maxItemTokens?: number,
67+
customHeaders?: Record<string, string>,
68+
) {
6169
if (!baseUrl) {
6270
throw new Error(t("embeddings:validation.baseUrlRequired"))
6371
}
@@ -67,13 +75,21 @@ export class OpenAICompatibleEmbedder implements IEmbedder {
6775

6876
this.baseUrl = baseUrl
6977
this.apiKey = apiKey
78+
this.customHeaders = customHeaders
7079

7180
// Wrap OpenAI client creation to handle invalid API key characters
7281
try {
73-
this.embeddingsClient = new OpenAI({
82+
// If custom headers are provided, we need to use defaultHeaders in OpenAI config
83+
const openAIConfig: any = {
7484
baseURL: baseUrl,
7585
apiKey: apiKey,
76-
})
86+
}
87+
88+
if (customHeaders) {
89+
openAIConfig.defaultHeaders = customHeaders
90+
}
91+
92+
this.embeddingsClient = new OpenAI(openAIConfig)
7793
} catch (error) {
7894
// Use the error handler to transform ByteString conversion errors
7995
throw handleOpenAIError(error, "OpenAI Compatible")
@@ -204,15 +220,22 @@ export class OpenAICompatibleEmbedder implements IEmbedder {
204220
batchTexts: string[],
205221
model: string,
206222
): Promise<OpenAIEmbeddingResponse> {
223+
const headers: Record<string, string> = {
224+
"Content-Type": "application/json",
225+
// Azure OpenAI uses 'api-key' header, while OpenAI uses 'Authorization'
226+
// We'll try 'api-key' first for Azure compatibility
227+
"api-key": this.apiKey,
228+
Authorization: `Bearer ${this.apiKey}`,
229+
}
230+
231+
// Add custom headers if provided
232+
if (this.customHeaders) {
233+
Object.assign(headers, this.customHeaders)
234+
}
235+
207236
const response = await fetch(url, {
208237
method: "POST",
209-
headers: {
210-
"Content-Type": "application/json",
211-
// Azure OpenAI uses 'api-key' header, while OpenAI uses 'Authorization'
212-
// We'll try 'api-key' first for Azure compatibility
213-
"api-key": this.apiKey,
214-
Authorization: `Bearer ${this.apiKey}`,
215-
},
238+
headers,
216239
body: JSON.stringify({
217240
input: batchTexts,
218241
model: model,

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ export interface CodeIndexConfig {
1111
modelDimension?: number // Generic dimension property for all providers
1212
openAiOptions?: ApiHandlerOptions
1313
ollamaOptions?: ApiHandlerOptions
14-
openAiCompatibleOptions?: { baseUrl: string; apiKey: string }
14+
openAiCompatibleOptions?: { baseUrl: string; apiKey: string; headers?: Record<string, string> }
1515
geminiOptions?: { apiKey: string }
1616
mistralOptions?: { apiKey: string }
1717
vercelAiGatewayOptions?: { apiKey: string }

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,8 @@ export class CodeIndexServiceFactory {
6363
config.openAiCompatibleOptions.baseUrl,
6464
config.openAiCompatibleOptions.apiKey,
6565
config.modelId,
66+
undefined, // maxItemTokens (use default)
67+
config.openAiCompatibleOptions.headers,
6668
)
6769
} else if (provider === "gemini") {
6870
if (!config.geminiOptions?.apiKey) {

0 commit comments

Comments
 (0)