Skip to content

Commit c67442d

Browse files
sobychackomarkpollack
authored andcommitted
Add custom header support for Azure OpenAI
- Adds configuration properties to allow custom header specification - Implements mechanism to apply custom headers to Azure OpenAI requests - Enhances flexibility for users to customize API interactions These changes allow users to add necessary headers for authentication, tracking, or other purposes when interacting with Azure OpenAI services. Resolves #1284
1 parent e644bf7 commit c67442d

File tree

3 files changed

+28
-12
lines changed

3 files changed

+28
-12
lines changed

spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/azure/openai/AzureOpenAiAutoConfiguration.java

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,8 @@
1616
package org.springframework.ai.autoconfigure.azure.openai;
1717

1818
import java.util.List;
19+
import java.util.Map;
20+
import java.util.stream.Collectors;
1921

2022
import org.springframework.ai.azure.openai.AzureOpenAiAudioTranscriptionModel;
2123
import org.springframework.ai.azure.openai.AzureOpenAiChatModel;
@@ -40,6 +42,7 @@
4042
import com.azure.core.credential.KeyCredential;
4143
import com.azure.core.credential.TokenCredential;
4244
import com.azure.core.util.ClientOptions;
45+
import com.azure.core.util.Header;
4346

4447
/**
4548
* @author Piotr Olaszewski
@@ -57,14 +60,19 @@ public class AzureOpenAiAutoConfiguration {
5760
@Bean
5861
@ConditionalOnMissingBean({ OpenAIClient.class, TokenCredential.class })
5962
public OpenAIClient openAIClient(AzureOpenAiConnectionProperties connectionProperties) {
60-
6163
if (StringUtils.hasText(connectionProperties.getApiKey())) {
6264

6365
Assert.hasText(connectionProperties.getEndpoint(), "Endpoint must not be empty");
6466

67+
Map<String, String> customHeaders = connectionProperties.getCustomHeaders();
68+
List<Header> headers = customHeaders.entrySet()
69+
.stream()
70+
.map(entry -> new Header(entry.getKey(), entry.getValue()))
71+
.collect(Collectors.toList());
72+
ClientOptions clientOptions = new ClientOptions().setApplicationId(APPLICATION_ID).setHeaders(headers);
6573
return new OpenAIClientBuilder().endpoint(connectionProperties.getEndpoint())
6674
.credential(new AzureKeyCredential(connectionProperties.getApiKey()))
67-
.clientOptions(new ClientOptions().setApplicationId(APPLICATION_ID))
75+
.clientOptions(clientOptions)
6876
.buildClient();
6977
}
7078

spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/azure/openai/AzureOpenAiConnectionProperties.java

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2023 - 2024 the original author or authors.
2+
* Copyright 2023-2024 the original author or authors.
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* you may not use this file except in compliance with the License.
@@ -13,8 +13,12 @@
1313
* See the License for the specific language governing permissions and
1414
* limitations under the License.
1515
*/
16+
1617
package org.springframework.ai.autoconfigure.azure.openai;
1718

19+
import java.util.HashMap;
20+
import java.util.Map;
21+
1822
import org.springframework.boot.context.properties.ConfigurationProperties;
1923

2024
@ConfigurationProperties(AzureOpenAiConnectionProperties.CONFIG_PREFIX)
@@ -40,6 +44,8 @@ public class AzureOpenAiConnectionProperties {
4044
*/
4145
private String endpoint;
4246

47+
private Map<String, String> customHeaders = new HashMap<>();
48+
4349
public String getEndpoint() {
4450
return this.endpoint;
4551
}
@@ -64,4 +70,12 @@ public void setOpenAiApiKey(String openAiApiKey) {
6470
this.openAiApiKey = openAiApiKey;
6571
}
6672

73+
public Map<String, String> getCustomHeaders() {
74+
return customHeaders;
75+
}
76+
77+
public void setCustomHeaders(Map<String, String> customHeaders) {
78+
this.customHeaders = customHeaders;
79+
}
80+
6781
}

spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/azure/AzureOpenAiAutoConfigurationIT.java

Lines changed: 3 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,9 @@
1515
*/
1616
package org.springframework.ai.autoconfigure.azure;
1717

18+
import com.azure.ai.openai.OpenAIClient;
19+
import com.azure.ai.openai.implementation.OpenAIClientImpl;
20+
import com.azure.core.http.*;
1821
import org.junit.jupiter.api.Test;
1922
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
2023
import org.springframework.ai.autoconfigure.azure.openai.AzureOpenAiAutoConfiguration;
@@ -34,15 +37,6 @@
3437
import org.springframework.core.io.ClassPathResource;
3538
import org.springframework.core.io.Resource;
3639
import org.springframework.util.ReflectionUtils;
37-
38-
import com.azure.ai.openai.OpenAIClient;
39-
import com.azure.ai.openai.implementation.OpenAIClientImpl;
40-
import com.azure.core.http.HttpHeader;
41-
import com.azure.core.http.HttpHeaderName;
42-
import com.azure.core.http.HttpMethod;
43-
import com.azure.core.http.HttpPipeline;
44-
import com.azure.core.http.HttpRequest;
45-
import com.azure.core.http.HttpResponse;
4640
import reactor.core.publisher.Flux;
4741

4842
import java.lang.reflect.Field;

0 commit comments

Comments
 (0)