Add AzureOpenAIClientBuilderCustomizer interface

Introduces a new customization point for Azure OpenAI client configuration
through the AzureOpenAIClientBuilderCustomizer interface. This allows
applications to customize the OpenAIClientBuilder while preserving the default
auto-configuration behavior.

- Add AzureOpenAIClientBuilderCustomizer interface (since 1.0.0-M6)
- Modify all OpenAIClientBuilder bean creation methods to apply customizers
- Add integration test verifying ordered customizer application

Signed-off-by: Manuel Andreo Garcia <manuel@magware.dev>
Signed-off-by: Soby Chacko <soby.chacko@broadcom.com>
This commit is contained in:
Manuel Andreo Garcia
2025-01-21 15:23:39 +01:00
committed by Soby Chacko
parent e92616b10e
commit bdd0aa118a
3 changed files with 61 additions and 5 deletions

View File

@@ -0,0 +1,21 @@
package org.springframework.ai.autoconfigure.azure.openai;
import com.azure.ai.openai.OpenAIClientBuilder;
/**
* Callback interface that can be implemented by beans wishing to customize the
* {@link OpenAIClientBuilder} whilst retaining the default auto-configuration.
*
* @author Manuel Andreo Garcia
* @since 1.0.0-M6
*/
@FunctionalInterface
public interface AzureOpenAIClientBuilderCustomizer {
/**
* Customize the {@link OpenAIClientBuilder}.
* @param clientBuilder the {@link OpenAIClientBuilder} to customize
*/
void customize(OpenAIClientBuilder clientBuilder);
}

View File

@@ -54,6 +54,7 @@ import org.springframework.util.StringUtils;
*
* @author Piotr Olaszewski
* @author Soby Chacko
* @author Manuel Andreo Garcia
*/
@AutoConfiguration
@ConditionalOnClass({ OpenAIClientBuilder.class, AzureOpenAiChatModel.class })
@@ -66,7 +67,9 @@ public class AzureOpenAiAutoConfiguration {
@Bean
@ConditionalOnMissingBean // ({ OpenAIClient.class, TokenCredential.class })
public OpenAIClientBuilder openAIClientBuilder(AzureOpenAiConnectionProperties connectionProperties) {
public OpenAIClientBuilder openAIClientBuilder(AzureOpenAiConnectionProperties connectionProperties,
ObjectProvider<AzureOpenAIClientBuilderCustomizer> customizers) {
if (StringUtils.hasText(connectionProperties.getApiKey())) {
Assert.hasText(connectionProperties.getEndpoint(), "Endpoint must not be empty");
@@ -77,17 +80,21 @@ public class AzureOpenAiAutoConfiguration {
.map(entry -> new Header(entry.getKey(), entry.getValue()))
.collect(Collectors.toList());
ClientOptions clientOptions = new ClientOptions().setApplicationId(APPLICATION_ID).setHeaders(headers);
return new OpenAIClientBuilder().endpoint(connectionProperties.getEndpoint())
OpenAIClientBuilder clientBuilder = new OpenAIClientBuilder().endpoint(connectionProperties.getEndpoint())
.credential(new AzureKeyCredential(connectionProperties.getApiKey()))
.clientOptions(clientOptions);
applyOpenAIClientBuilderCustomizers(clientBuilder, customizers);
return clientBuilder;
}
// Connect to OpenAI (e.g. not the Azure OpenAI). The deploymentName property is
// used as OpenAI model name.
if (StringUtils.hasText(connectionProperties.getOpenAiApiKey())) {
return new OpenAIClientBuilder().endpoint("https://api.openai.com/v1")
OpenAIClientBuilder clientBuilder = new OpenAIClientBuilder().endpoint("https://api.openai.com/v1")
.credential(new KeyCredential(connectionProperties.getOpenAiApiKey()))
.clientOptions(new ClientOptions().setApplicationId(APPLICATION_ID));
applyOpenAIClientBuilderCustomizers(clientBuilder, customizers);
return clientBuilder;
}
throw new IllegalArgumentException("Either API key or OpenAI API key must not be empty");
@@ -97,14 +104,16 @@ public class AzureOpenAiAutoConfiguration {
@ConditionalOnMissingBean
@ConditionalOnBean(TokenCredential.class)
public OpenAIClientBuilder openAIClientWithTokenCredential(AzureOpenAiConnectionProperties connectionProperties,
TokenCredential tokenCredential) {
TokenCredential tokenCredential, ObjectProvider<AzureOpenAIClientBuilderCustomizer> customizers) {
Assert.notNull(tokenCredential, "TokenCredential must not be null");
Assert.hasText(connectionProperties.getEndpoint(), "Endpoint must not be empty");
return new OpenAIClientBuilder().endpoint(connectionProperties.getEndpoint())
OpenAIClientBuilder clientBuilder = new OpenAIClientBuilder().endpoint(connectionProperties.getEndpoint())
.credential(tokenCredential)
.clientOptions(new ClientOptions().setApplicationId(APPLICATION_ID));
applyOpenAIClientBuilderCustomizers(clientBuilder, customizers);
return clientBuilder;
}
@Bean
@@ -169,4 +178,9 @@ public class AzureOpenAiAutoConfiguration {
return new AzureOpenAiAudioTranscriptionModel(openAIClient.buildClient(), audioProperties.getOptions());
}
private void applyOpenAIClientBuilderCustomizers(OpenAIClientBuilder clientBuilder,
ObjectProvider<AzureOpenAIClientBuilderCustomizer> customizers) {
customizers.orderedStream().forEach(customizer -> customizer.customize(clientBuilder));
}
}

View File

@@ -20,6 +20,7 @@ import java.lang.reflect.Field;
import java.net.URI;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.stream.Collectors;
import com.azure.ai.openai.OpenAIClient;
@@ -33,6 +34,9 @@ import com.azure.core.http.HttpRequest;
import com.azure.core.http.HttpResponse;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
import org.springframework.ai.autoconfigure.azure.openai.AzureOpenAIClientBuilderCustomizer;
import reactor.core.publisher.Flux;
import org.springframework.ai.autoconfigure.azure.openai.AzureOpenAiAutoConfiguration;
@@ -59,6 +63,7 @@ import static org.assertj.core.api.Assertions.assertThat;
* @author Christian Tzolov
* @author Piotr Olaszewski
* @author Soby Chacko
* @author Manuel Andreo Garcia
* @since 0.8.0
*/
@EnabledIfEnvironmentVariable(named = "AZURE_OPENAI_API_KEY", matches = ".+")
@@ -228,4 +233,20 @@ class AzureOpenAiAutoConfigurationIT {
.run(context -> assertThat(context.getBeansOfType(AzureOpenAiAudioTranscriptionModel.class)).isNotEmpty());
}
@Test
void openAIClientBuilderCustomizer() {
AtomicBoolean firstCustomizationApplied = new AtomicBoolean(false);
AtomicBoolean secondCustomizationApplied = new AtomicBoolean(false);
this.contextRunner
.withBean("first", AzureOpenAIClientBuilderCustomizer.class,
() -> clientBuilder -> firstCustomizationApplied.set(true))
.withBean("second", AzureOpenAIClientBuilderCustomizer.class,
() -> clientBuilder -> secondCustomizationApplied.set(true))
.run(context -> {
context.getBean(OpenAIClientBuilder.class);
assertThat(firstCustomizationApplied.get()).isTrue();
assertThat(secondCustomizationApplied.get()).isTrue();
});
}
}