Fix formatting and add API key for embeddings

Signed-off-by: Filip Hrisafov <[email protected]>

Author: filiphr

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

@@ -175,12 +175,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.");
 
-		return this.restClient.post().uri(this.completionsPath).headers(headers -> {
-			headers.addAll(additionalHttpHeader);
-			if (!headers.containsKey(HttpHeaders.AUTHORIZATION) && !(this.apiKey instanceof NoopApiKey)) {
-				headers.setBearerAuth(this.apiKey.getValue());
-			}
-		}).body(chatRequest).retrieve().toEntity(ChatCompletion.class);
+		// @formatter:off
+		return this.restClient.post()
+			.uri(this.completionsPath)
+			.headers(headers -> {
+				headers.addAll(additionalHttpHeader);
+				addDefaultHeadersIfMissing(headers);
+			})
+			.body(chatRequest)
+			.retrieve()
+			.toEntity(ChatCompletion.class);
+		// @formatter:on
 	}
 
 	/**
@@ -209,12 +214,13 @@ public Flux<ChatCompletionChunk> chatCompletionStream(ChatCompletionRequest chat
 
 		AtomicBoolean isInsideTool = new AtomicBoolean(false);
 
-		return this.webClient.post().uri(this.completionsPath).headers(headers -> {
-			headers.addAll(additionalHttpHeader);
-			if (!headers.containsKey(HttpHeaders.AUTHORIZATION) && !(this.apiKey instanceof NoopApiKey)) {
-				headers.setBearerAuth(this.apiKey.getValue());
-			}
-		})
+		// @formatter:off
+		return this.webClient.post()
+			.uri(this.completionsPath)
+			.headers(headers -> {
+				headers.addAll(additionalHttpHeader);
+				addDefaultHeadersIfMissing(headers);
+			}) // @formatter:on
 			.body(Mono.just(chatRequest), ChatCompletionRequest.class)
 			.retrieve()
 			.bodyToFlux(String.class)
@@ -288,13 +294,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/test/java/org/springframework/ai/openai/api/OpenAiApiBuilderTests.java

@@ -372,6 +372,57 @@ void dynamicApiKeyWebClientWithAdditionalAuthorizationHeader() throws Interrupte
 			assertThat(recordedRequest.getHeader(HttpHeaders.AUTHORIZATION)).isEqualTo("Bearer additional-key");
 		}
 
+		@Test
+		void dynamicApiKeyRestClientEmbeddings() throws InterruptedException {
+			Queue<ApiKey> apiKeys = new LinkedList<>(List.of(new SimpleApiKey("key1"), new SimpleApiKey("key2")));
+			OpenAiApi api = OpenAiApi.builder()
+				.apiKey(() -> Objects.requireNonNull(apiKeys.poll()).getValue())
+				.baseUrl(mockWebServer.url("/").toString())
+				.build();
+
+			MockResponse mockResponse = new MockResponse().setResponseCode(200)
+				.addHeader(HttpHeaders.CONTENT_TYPE, MediaType.APPLICATION_JSON_VALUE)
+				.setBody("""
+						{
+							"object": "list",
+							"data": [
+						    	{
+						    		"object": "embedding",
+								"index": 0,
+						       "embedding": [
+						         -0.005540426,
+						         0.0047363234,
+						         -0.015009919,
+						         -0.027093535,
+						         -0.015173893,
+						         0.015173893,
+						         -0.017608276
+						         ]
+						     }
+						   ],
+						   "model": "text-embedding-ada-002-v2",
+						   "usage": {
+						     "prompt_tokens": 2,
+						     "total_tokens": 2
+						   }
+						}
+						""");
+			mockWebServer.enqueue(mockResponse);
+			mockWebServer.enqueue(mockResponse);
+
+			OpenAiApi.EmbeddingRequest<String> request = new OpenAiApi.EmbeddingRequest<>("Hello world");
+			ResponseEntity<OpenAiApi.EmbeddingList<OpenAiApi.Embedding>> response = api.embeddings(request);
+			assertThat(response.getStatusCode()).isEqualTo(HttpStatus.OK);
+			RecordedRequest recordedRequest = mockWebServer.takeRequest();
+			assertThat(recordedRequest.getHeader(HttpHeaders.AUTHORIZATION)).isEqualTo("Bearer key1");
+
+			response = api.embeddings(request);
+			assertThat(response.getStatusCode()).isEqualTo(HttpStatus.OK);
+
+			recordedRequest = mockWebServer.takeRequest();
+			assertThat(recordedRequest.getHeader(HttpHeaders.AUTHORIZATION)).isEqualTo("Bearer key2");
+		}
+
 	}
 
 }