🔐 feat: Implement Entra ID authentication for Azure OpenAI integration

- Added support for Entra ID authentication in OpenAIClient and related services.
- Updated header management to conditionally use Entra ID access tokens or API keys based on environment configuration.
- Introduced utility functions for Entra ID token retrieval and credential management.
- Enhanced tests to verify Entra ID authentication flow and its integration with Azure configurations.
This commit is contained in:
victorbjorkgren 2025-09-12 17:29:43 +02:00 committed by victorbjor
parent a1471c2f37
commit 9288e84454
9 changed files with 212 additions and 18 deletions

View file

@ -6,8 +6,10 @@ const {
Tokenizer,
createFetch,
resolveHeaders,
shouldUseEntraId,
constructAzureURL,
getModelMaxTokens,
getEntraIdAccessToken,
genAzureChatCompletion,
getModelMaxOutputTokens,
createStreamEventHandlers,
@ -614,7 +616,7 @@ class OpenAIClient extends BaseClient {
return (reply ?? '').trim();
}
initializeLLM({
async initializeLLM({
model = openAISettings.model.default,
modelName,
temperature = 0.2,
@ -753,7 +755,14 @@ class OpenAIClient extends BaseClient {
this.options.defaultQuery = azureOptions.azureOpenAIApiVersion
? { 'api-version': azureOptions.azureOpenAIApiVersion }
: undefined;
this.options.headers['api-key'] = this.apiKey;
if (shouldUseEntraId()) {
this.options.headers = {
...this.options.headers,
Authorization: `Bearer ${await getEntraIdAccessToken()}`,
};
} else {
this.options.headers['api-key'] = this.apiKey;
}
}
}
@ -812,7 +821,7 @@ ${convo}
try {
this.abortController = new AbortController();
const llm = this.initializeLLM({
const llm = await this.initializeLLM({
...modelOptions,
conversationId,
context: 'title',
@ -961,7 +970,7 @@ ${convo}
const initialPromptTokens = this.maxContextTokens - remainingContextTokens;
logger.debug('[OpenAIClient] initialPromptTokens', initialPromptTokens);
const llm = this.initializeLLM({
const llm = await this.initializeLLM({
model,
temperature: 0.2,
context: 'summary',
@ -1187,7 +1196,14 @@ ${convo}
this.options.defaultQuery = azureOptions.azureOpenAIApiVersion
? { 'api-version': azureOptions.azureOpenAIApiVersion }
: undefined;
this.options.headers['api-key'] = this.apiKey;
if (shouldUseEntraId()) {
this.options.headers = {
...this.options.headers,
Authorization: `Bearer ${await getEntraIdAccessToken()}`,
};
} else {
this.options.headers['api-key'] = this.apiKey;
}
}
}
@ -1208,7 +1224,14 @@ ${convo}
: this.azureEndpoint.split(/(?<!\/)\/(chat|completion)\//)[0];
opts.defaultQuery = { 'api-version': this.azure.azureOpenAIApiVersion };
opts.defaultHeaders = { ...opts.defaultHeaders, 'api-key': this.apiKey };
if (shouldUseEntraId()) {
opts.defaultHeaders = {
...opts.defaultHeaders,
Authorization: `Bearer ${await getEntraIdAccessToken()}`,
};
} else {
opts.defaultHeaders = { ...opts.defaultHeaders, 'api-key': this.apiKey };
}
}
if (this.isOmni === true && modelOptions.max_tokens != null) {