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:
committed by
Soby Chacko
parent
e92616b10e
commit
bdd0aa118a
@@ -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);
|
||||
|
||||
}
|
||||
@@ -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));
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
@@ -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();
|
||||
});
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user