Spring/spring-ai
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Add Azure Search user-agent, test OpenAI client header

Browse Source 
* Adding test to verify the user-agent header in the Azure OpenAi chat client
* Adding the user-agent header to the search client in Azure vector store
This commit is contained in:
Soby Chacko
2024-09-19 12:07:50 -04:00
committed by Mark Pollack
parent 035036c7f2
commit e644bf77da
2 changed files with 45 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/azure/AzureVectorStoreAutoConfiguration.java
View File

@@ -17,6 +17,7 @@
package org.springframework.ai.autoconfigure.vectorstore.azure;
import com.azure.core.credential.AzureKeyCredential;
import com.azure.core.util.ClientOptions;
import com.azure.search.documents.indexes.SearchIndexClient;
import com.azure.search.documents.indexes.SearchIndexClientBuilder;
@@ -47,11 +48,16 @@ import org.springframework.context.annotation.Bean;
@ConditionalOnProperty(prefix = "spring.ai.vectorstore.azure", value = { "url", "api-key", "index-name" })
public class AzureVectorStoreAutoConfiguration {
private final static String APPLICATION_ID = "spring-ai";
@Bean
@ConditionalOnMissingBean
public SearchIndexClient searchIndexClient(AzureVectorStoreProperties properties) {
ClientOptions clientOptions = new ClientOptions();
clientOptions.setApplicationId(APPLICATION_ID);
return new SearchIndexClientBuilder().endpoint(properties.getUrl())
.credential(new AzureKeyCredential(properties.getApiKey()))
.clientOptions(clientOptions)
.buildClient();
}

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

@@ -1,5 +1,5 @@
/*
* Copyright 2023 - 2024 the original author or authors.
* Copyright 2023-2024 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
@@ -33,8 +33,20 @@ import org.springframework.boot.autoconfigure.AutoConfigurations;
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
import org.springframework.core.io.ClassPathResource;
import org.springframework.core.io.Resource;
import org.springframework.util.ReflectionUtils;
import com.azure.ai.openai.OpenAIClient;
import com.azure.ai.openai.implementation.OpenAIClientImpl;
import com.azure.core.http.HttpHeader;
import com.azure.core.http.HttpHeaderName;
import com.azure.core.http.HttpMethod;
import com.azure.core.http.HttpPipeline;
import com.azure.core.http.HttpRequest;
import com.azure.core.http.HttpResponse;
import reactor.core.publisher.Flux;
import java.lang.reflect.Field;
import java.net.URI;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
@@ -44,11 +56,12 @@ import static org.assertj.core.api.Assertions.assertThat;
/**
* @author Christian Tzolov
* @author Piotr Olaszewski
* @author Soby Chacko
* @since 0.8.0
*/
@EnabledIfEnvironmentVariable(named = "AZURE_OPENAI_API_KEY", matches = ".+")
@EnabledIfEnvironmentVariable(named = "AZURE_OPENAI_ENDPOINT", matches = ".+")
public class AzureOpenAiAutoConfigurationIT {
class AzureOpenAiAutoConfigurationIT {
private static String CHAT_MODEL_NAME = "gpt-4o";
@@ -79,7 +92,7 @@ public class AzureOpenAiAutoConfigurationIT {
"Tell me about 3 famous pirates from the Golden Age of Piracy and why they did.");
@Test
public void chatCompletion() {
void chatCompletion() {
contextRunner.run(context -> {
AzureOpenAiChatModel chatModel = context.getBean(AzureOpenAiChatModel.class);
ChatResponse response = chatModel.call(new Prompt(List.of(userMessage, systemMessage)));
@@ -88,7 +101,26 @@ public class AzureOpenAiAutoConfigurationIT {
}
@Test
public void chatCompletionStreaming() {
void httpRequestContainsUserAgentHeader() {
contextRunner.run(context -> {
OpenAIClient openAIClient = context.getBean(OpenAIClient.class);
Field serviceClientField = ReflectionUtils.findField(OpenAIClient.class, "serviceClient");
assertThat(serviceClientField).isNotNull();
ReflectionUtils.makeAccessible(serviceClientField);
OpenAIClientImpl oaci = (OpenAIClientImpl) ReflectionUtils.getField(serviceClientField, openAIClient);
assertThat(oaci).isNotNull();
HttpPipeline httpPipeline = oaci.getHttpPipeline();
HttpResponse httpResponse = httpPipeline
.send(new HttpRequest(HttpMethod.POST, new URI(System.getenv("AZURE_OPENAI_ENDPOINT")).toURL()))
.block();
assertThat(httpResponse).isNotNull();
HttpHeader httpHeader = httpResponse.getRequest().getHeaders().get(HttpHeaderName.USER_AGENT);
assertThat(httpHeader.getValue().startsWith("spring-ai azsdk-java-azure-ai-openai/")).isTrue();
});
}
@Test
void chatCompletionStreaming() {
contextRunner.run(context -> {
AzureOpenAiChatModel chatModel = context.getBean(AzureOpenAiChatModel.class);
@@ -140,7 +172,7 @@ public class AzureOpenAiAutoConfigurationIT {
}
@Test
public void chatActivation() {
void chatActivation() {
// Disable the chat auto-configuration.
contextRunner.withPropertyValues("spring.ai.azure.openai.chat.enabled=false").run(context -> {
@@ -159,7 +191,7 @@ public class AzureOpenAiAutoConfigurationIT {
}
@Test
public void embeddingActivation() {
void embeddingActivation() {
// Disable the embedding auto-configuration.
contextRunner.withPropertyValues("spring.ai.azure.openai.embedding.enabled=false").run(context -> {
@@ -178,7 +210,7 @@ public class AzureOpenAiAutoConfigurationIT {
}
@Test
public void audioTranscriptionActivation() {
void audioTranscriptionActivation() {
// Disable the transcription auto-configuration.
contextRunner.withPropertyValues("spring.ai.azure.openai.audio.transcription.enabled=false").run(context -> {