Resolve OpenAI ApiKey for every request

 - Set ApiKey as late as possible Signed-off-by: Filip Hrisafov <filip.hrisafov@gmail.com>

Author: filiphr

models/spring-ai-openai/src/main/java/org/springframework/ai/openai/api/OpenAiApi.java

@@ -62,6 +62,7 @@
  * @author Thomas Vitale
  * @author David Frizelle
  * @author Alexandros Pappas
+ * @author Filip Hrisafov
  */
 public class OpenAiApi {
 
@@ -128,10 +129,6 @@ public OpenAiApi(String baseUrl, ApiKey apiKey, MultiValueMap<String, String> he
 
 // @formatter:off
 Consumer<HttpHeaders> finalHeaders = h -> {
-if (!(apiKey instanceof NoopApiKey)) {
-h.setBearerAuth(apiKey.getValue());
-}
-
 h.setContentType(MediaType.APPLICATION_JSON);
 h.addAll(headers);
 };
@@ -179,12 +176,17 @@ public ResponseEntity<ChatCompletion> chatCompletionEntity(ChatCompletionRequest
 Assert.isTrue(!chatRequest.stream(), "Request must set the stream property to false.");
 Assert.notNull(additionalHttpHeader, "The additional HTTP headers can not be null.");
 
+// @formatter:off
 return this.restClient.post()
 .uri(this.completionsPath)
-.headers(headers -> headers.addAll(additionalHttpHeader))
+.headers(headers -> {
+headers.addAll(additionalHttpHeader);
+addDefaultHeadersIfMissing(headers);
+})
 .body(chatRequest)
 .retrieve()
 .toEntity(ChatCompletion.class);
+// @formatter:on
 }
 
 /**
@@ -213,9 +215,13 @@ public Flux<ChatCompletionChunk> chatCompletionStream(ChatCompletionRequest chat
 
 AtomicBoolean isInsideTool = new AtomicBoolean(false);
 
+// @formatter:off
 return this.webClient.post()
 .uri(this.completionsPath)
-.headers(headers -> headers.addAll(additionalHttpHeader))
+.headers(headers -> {
+headers.addAll(additionalHttpHeader);
+addDefaultHeadersIfMissing(headers);
+}) // @formatter:on
 .body(Mono.just(chatRequest), ChatCompletionRequest.class)
 .retrieve()
 .bodyToFlux(String.class)
@@ -289,13 +295,20 @@ public <T> ResponseEntity<EmbeddingList<Embedding>> embeddings(EmbeddingRequest<
 
 return this.restClient.post()
 .uri(this.embeddingsPath)
+.headers(this::addDefaultHeadersIfMissing)
 .body(embeddingRequest)
 .retrieve()
 .toEntity(new ParameterizedTypeReference<>() {
 
 });
 }
 
+private void addDefaultHeadersIfMissing(HttpHeaders headers) {
+if (!headers.containsKey(HttpHeaders.AUTHORIZATION) && !(this.apiKey instanceof NoopApiKey)) {
+headers.setBearerAuth(this.apiKey.getValue());
+}
+}
+
 // Package-private getters for mutate/copy
 String getBaseUrl() {
 return this.baseUrl;

models/spring-ai-openai/src/main/java/org/springframework/ai/openai/api/OpenAiAudioApi.java

@@ -49,6 +49,7 @@
  * @author Christian Tzolov
  * @author Ilayaperumal Gopinathan
  * @author Jonghoon Park
+ * @author Filip Hrisafov
  * @since 0.8.1
  */
 public class OpenAiAudioApi {
@@ -71,20 +72,30 @@ public OpenAiAudioApi(String baseUrl, ApiKey apiKey, MultiValueMap<String, Strin
 ResponseErrorHandler responseErrorHandler) {
 
 Consumer<HttpHeaders> authHeaders = h -> {
-if (!(apiKey instanceof NoopApiKey)) {
-h.setBearerAuth(apiKey.getValue());
-}
 h.addAll(headers);
-// h.setContentType(MediaType.APPLICATION_JSON);
 };
 
+// @formatter:off
 this.restClient = restClientBuilder.clone()
 .baseUrl(baseUrl)
 .defaultHeaders(authHeaders)
 .defaultStatusHandler(responseErrorHandler)
+.defaultRequest(requestHeadersSpec -> {
+if (!(apiKey instanceof NoopApiKey)) {
+requestHeadersSpec.header(HttpHeaders.AUTHORIZATION, "Bearer " + apiKey.getValue());
+}
+})
 .build();
 
-this.webClient = webClientBuilder.clone().baseUrl(baseUrl).defaultHeaders(authHeaders).build();
+this.webClient = webClientBuilder.clone()
+.baseUrl(baseUrl)
+.defaultHeaders(authHeaders)
+.defaultRequest(requestHeadersSpec -> {
+if (!(apiKey instanceof NoopApiKey)) {
+requestHeadersSpec.header(HttpHeaders.AUTHORIZATION, "Bearer " + apiKey.getValue());
+}
+})
+.build(); // @formatter:on
 }
 
 public static Builder builder() {

models/spring-ai-openai/src/main/java/org/springframework/ai/openai/api/OpenAiImageApi.java

@@ -27,6 +27,7 @@
 import org.springframework.ai.model.SimpleApiKey;
 import org.springframework.ai.openai.api.common.OpenAiApiConstants;
 import org.springframework.ai.retry.RetryUtils;
+import org.springframework.http.HttpHeaders;
 import org.springframework.http.MediaType;
 import org.springframework.http.ResponseEntity;
 import org.springframework.util.Assert;
@@ -40,6 +41,7 @@
  *
  * @see <a href= "https://platform.openai.com/docs/api-reference/images">Images</a>
  * @author lambochen
+ * @author Filip Hrisafov
  */
 public class OpenAiImageApi {
 
@@ -62,15 +64,18 @@ public OpenAiImageApi(String baseUrl, ApiKey apiKey, MultiValueMap<String, Strin
 RestClient.Builder restClientBuilder, ResponseErrorHandler responseErrorHandler) {
 
 // @formatter:off
-this.restClient = restClientBuilder.baseUrl(baseUrl)
+this.restClient = restClientBuilder.clone()
+.baseUrl(baseUrl)
 .defaultHeaders(h -> {
-if (!(apiKey instanceof NoopApiKey)) {
-h.setBearerAuth(apiKey.getValue());
-}
 h.setContentType(MediaType.APPLICATION_JSON);
 h.addAll(headers);
 })
 .defaultStatusHandler(responseErrorHandler)
+.defaultRequest(requestHeadersSpec -> {
+if (!(apiKey instanceof NoopApiKey)) {
+requestHeadersSpec.header(HttpHeaders.AUTHORIZATION, "Bearer " + apiKey.getValue());
+}
+})
 .build();
 // @formatter:on
 

models/spring-ai-openai/src/main/java/org/springframework/ai/openai/api/OpenAiModerationApi.java

@@ -27,6 +27,7 @@
 import org.springframework.ai.model.SimpleApiKey;
 import org.springframework.ai.openai.api.common.OpenAiApiConstants;
 import org.springframework.ai.retry.RetryUtils;
+import org.springframework.http.HttpHeaders;
 import org.springframework.http.MediaType;
 import org.springframework.http.ResponseEntity;
 import org.springframework.util.Assert;
@@ -40,6 +41,7 @@
  *
  * @author Ahmed Yousri
  * @author Ilayaperumal Gopinathan
+ * @author Filip Hrisafov
  * @see <a href=
  * "https://platform.openai.com/docs/api-reference/moderations">https://platform.openai.com/docs/api-reference/moderations</a>
  */
@@ -64,13 +66,20 @@ public OpenAiModerationApi(String baseUrl, ApiKey apiKey, MultiValueMap<String,
 
 this.objectMapper = new ObjectMapper().configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false);
 
-this.restClient = restClientBuilder.baseUrl(baseUrl).defaultHeaders(h -> {
-if (!(apiKey instanceof NoopApiKey)) {
-h.setBearerAuth(apiKey.getValue());
-}
-h.setContentType(MediaType.APPLICATION_JSON);
-h.addAll(headers);
-}).defaultStatusHandler(responseErrorHandler).build();
+// @formatter:off
+this.restClient = restClientBuilder.clone()
+.baseUrl(baseUrl)
+.defaultHeaders(h -> {
+h.setContentType(MediaType.APPLICATION_JSON);
+h.addAll(headers);
+})
+.defaultStatusHandler(responseErrorHandler)
+.defaultRequest(requestHeadersSpec -> {
+if (!(apiKey instanceof NoopApiKey)) {
+requestHeadersSpec.header(HttpHeaders.AUTHORIZATION, "Bearer " + apiKey.getValue());
+}
+})
+.build(); // @formatter:on
 }
 
 public ResponseEntity<OpenAiModerationResponse> createModeration(OpenAiModerationRequest openAiModerationRequest) {