Skip to content

Commit f0858ad

Browse files
filiphrfiliphr
committed
Set ApiKey as late as possible
Signed-off-by: Filip Hrisafov <[email protected]>
1 parent f05c376 commit f0858ad

File tree

2 files changed

+137
-18
lines changed

2 files changed

+137
-18
lines changed

‎models/spring-ai-anthropic/src/main/java/org/springframework/ai/anthropic/api/AnthropicApi.java‎

Lines changed: 21 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -98,6 +98,8 @@ public static Builder builder() {
9898

9999
private final WebClient webClient;
100100

101+
private final ApiKey apiKey;
102+
101103
/**
102104
* Create a new client api.
103105
* @param baseUrl api base URL.
@@ -120,17 +122,12 @@ private AnthropicApi(String baseUrl, String completionsPath, ApiKey anthropicApi
120122
};
121123

122124
this.completionsPath = completionsPath;
125+
this.apiKey = anthropicApiKey;
123126

124127
this.restClient = restClientBuilder.clone()
125128
.baseUrl(baseUrl)
126129
.defaultHeaders(jsonContentHeaders)
127130
.defaultStatusHandler(responseErrorHandler)
128-
.defaultRequest(requestHeadersSpec -> {
129-
String value = anthropicApiKey.getValue();
130-
if (StringUtils.hasText(value)) {
131-
requestHeadersSpec.header(HEADER_X_API_KEY, value);
132-
}
133-
})
134131
.build();
135132

136133
this.webClient = webClientBuilder.clone()
@@ -140,12 +137,6 @@ private AnthropicApi(String baseUrl, String completionsPath, ApiKey anthropicApi
140137
resp -> resp.bodyToMono(String.class)
141138
.flatMap(it -> Mono.error(new RuntimeException(
142139
"Response exception, Status: [" + resp.statusCode() + "], Body:[" + it + "]"))))
143-
.defaultRequest(requestHeadersSpec -> {
144-
String value = anthropicApiKey.getValue();
145-
if (StringUtils.hasText(value)) {
146-
requestHeadersSpec.header(HEADER_X_API_KEY, value);
147-
}
148-
})
149140
.build();
150141
}
151142

@@ -175,7 +166,15 @@ public ResponseEntity<ChatCompletionResponse> chatCompletionEntity(ChatCompletio
175166

176167
return this.restClient.post()
177168
.uri(this.completionsPath)
178-
.headers(headers -> headers.addAll(additionalHttpHeader))
169+
.headers(headers -> {
170+
headers.addAll(additionalHttpHeader);
171+
if (!headers.containsKey(HEADER_X_API_KEY)) {
172+
String apiKeyValue = this.apiKey.getValue();
173+
if (StringUtils.hasText(apiKeyValue)) {
174+
headers.add(HEADER_X_API_KEY, apiKeyValue);
175+
}
176+
}
177+
})
179178
.body(chatRequest)
180179
.retrieve()
181180
.toEntity(ChatCompletionResponse.class);
@@ -211,7 +210,15 @@ public Flux<ChatCompletionResponse> chatCompletionStream(ChatCompletionRequest c
211210

212211
return this.webClient.post()
213212
.uri(this.completionsPath)
214-
.headers(headers -> headers.addAll(additionalHttpHeader))
213+
.headers(headers -> {
214+
headers.addAll(additionalHttpHeader);
215+
if (!headers.containsKey(HEADER_X_API_KEY)) {
216+
String apiKeyValue = this.apiKey.getValue();
217+
if (StringUtils.hasText(apiKeyValue)) {
218+
headers.add(HEADER_X_API_KEY, apiKeyValue);
219+
}
220+
}
221+
})
215222
.body(Mono.just(chatRequest), ChatCompletionRequest.class)
216223
.retrieve()
217224
.bodyToFlux(String.class)

‎models/spring-ai-anthropic/src/test/java/org/springframework/ai/anthropic/api/AnthropicApiBuilderTests.java‎

Lines changed: 116 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -36,13 +36,16 @@
3636
import org.springframework.http.HttpStatus;
3737
import org.springframework.http.MediaType;
3838
import org.springframework.http.ResponseEntity;
39+
import org.springframework.util.LinkedMultiValueMap;
40+
import org.springframework.util.MultiValueMap;
3941
import org.springframework.web.client.ResponseErrorHandler;
4042
import org.springframework.web.client.RestClient;
4143
import org.springframework.web.reactive.function.client.WebClient;
4244

4345
import okhttp3.mockwebserver.MockResponse;
4446
import okhttp3.mockwebserver.MockWebServer;
4547
import okhttp3.mockwebserver.RecordedRequest;
48+
import org.opentest4j.AssertionFailedError;
4649

4750
public class AnthropicApiBuilderTests {
4851

@@ -191,6 +194,50 @@ void dynamicApiKeyRestClient() throws InterruptedException {
191194
assertThat(recordedRequest.getHeader("x-api-key")).isEqualTo("key2");
192195
}
193196

197+
@Test
198+
void dynamicApiKeyRestClientWithAdditionalApiKeyHeader() throws InterruptedException {
199+
AnthropicApi api = AnthropicApi.builder()
200+
.apiKey(() -> {
201+
throw new AssertionFailedError("Should not be called, API key is provided in headers");
202+
})
203+
.baseUrl(mockWebServer.url("/").toString())
204+
.build();
205+
206+
MockResponse mockResponse = new MockResponse().setResponseCode(200)
207+
.addHeader(HttpHeaders.CONTENT_TYPE, MediaType.APPLICATION_JSON_VALUE)
208+
.setBody("""
209+
{
210+
"id": "msg_1nZdL29xx5MUA1yADyHTEsnR8uuvGzszyY",
211+
"type": "message",
212+
"role": "assistant",
213+
"content": [],
214+
"model": "claude-opus-3-latest",
215+
"stop_reason": null,
216+
"stop_sequence": null,
217+
"usage": {
218+
"input_tokens": 25,
219+
"output_tokens": 1
220+
}
221+
}
222+
""");
223+
mockWebServer.enqueue(mockResponse);
224+
225+
AnthropicApi.AnthropicMessage chatCompletionMessage = new AnthropicApi.AnthropicMessage(
226+
List.of(new AnthropicApi.ContentBlock("Hello world")), AnthropicApi.Role.USER);
227+
AnthropicApi.ChatCompletionRequest request = AnthropicApi.ChatCompletionRequest.builder()
228+
.model(AnthropicApi.ChatModel.CLAUDE_3_OPUS)
229+
.temperature(0.8)
230+
.messages(List.of(chatCompletionMessage))
231+
.build();
232+
MultiValueMap<String, String> additionalHeaders = new LinkedMultiValueMap<>();
233+
additionalHeaders.add("x-api-key", "additional-key");
234+
ResponseEntity<AnthropicApi.ChatCompletionResponse> response = api.chatCompletionEntity(request, additionalHeaders);
235+
assertThat(response.getStatusCode()).isEqualTo(HttpStatus.OK);
236+
RecordedRequest recordedRequest = mockWebServer.takeRequest();
237+
assertThat(recordedRequest.getHeader(HttpHeaders.AUTHORIZATION)).isNull();
238+
assertThat(recordedRequest.getHeader("x-api-key")).isEqualTo("additional-key");
239+
}
240+
194241
@Test
195242
void dynamicApiKeyWebClient() throws InterruptedException {
196243
Queue<ApiKey> apiKeys = new LinkedList<>(List.of(new SimpleApiKey("key1"), new SimpleApiKey("key2")));
@@ -203,8 +250,23 @@ void dynamicApiKeyWebClient() throws InterruptedException {
203250
.addHeader(HttpHeaders.CONTENT_TYPE, MediaType.TEXT_EVENT_STREAM_VALUE)
204251
.setBody(
205252
"""
206-
{"type": "message_start", "message": {"id": "msg_1nZdL29xx5MUA1yADyHTEsnR8uuvGzszyY", "type": "message", "role": "assistant", "content": [], "model": "claude-opus-4-20250514", "stop_reason": null, "stop_sequence": null, "usage": {"input_tokens": 25, "output_tokens": 1}}}
207-
""");
253+
{
254+
"type": "message_start",
255+
"message": {
256+
"id": "msg_1nZdL29xx5MUA1yADyHTEsnR8uuvGzszyY",
257+
"type": "message",
258+
"role": "assistant",
259+
"content": [],
260+
"model": "claude-opus-4-20250514",
261+
"stop_reason": null,
262+
"stop_sequence": null,
263+
"usage": {
264+
"input_tokens": 25,
265+
"output_tokens": 1
266+
}
267+
}
268+
}
269+
""".replace("\n", ""));
208270
mockWebServer.enqueue(mockResponse);
209271
mockWebServer.enqueue(mockResponse);
210272

@@ -216,20 +278,70 @@ void dynamicApiKeyWebClient() throws InterruptedException {
216278
.messages(List.of(chatCompletionMessage))
217279
.stream(true)
218280
.build();
219-
List<AnthropicApi.ChatCompletionResponse> response = api.chatCompletionStream(request)
281+
api.chatCompletionStream(request)
220282
.collectList()
221283
.block();
222284
RecordedRequest recordedRequest = mockWebServer.takeRequest();
223285
assertThat(recordedRequest.getHeader(HttpHeaders.AUTHORIZATION)).isNull();
224286
assertThat(recordedRequest.getHeader("x-api-key")).isEqualTo("key1");
225287

226-
response = api.chatCompletionStream(request).collectList().block();
288+
api.chatCompletionStream(request).collectList().block();
227289

228290
recordedRequest = mockWebServer.takeRequest();
229291
assertThat(recordedRequest.getHeader(HttpHeaders.AUTHORIZATION)).isNull();
230292
assertThat(recordedRequest.getHeader("x-api-key")).isEqualTo("key2");
231293
}
232294

295+
@Test
296+
void dynamicApiKeyWebClientWithAdditionalApiKey() throws InterruptedException {
297+
Queue<ApiKey> apiKeys = new LinkedList<>(List.of(new SimpleApiKey("key1"), new SimpleApiKey("key2")));
298+
AnthropicApi api = AnthropicApi.builder()
299+
.apiKey(() -> Objects.requireNonNull(apiKeys.poll()).getValue())
300+
.baseUrl(mockWebServer.url("/").toString())
301+
.build();
302+
303+
MockResponse mockResponse = new MockResponse().setResponseCode(200)
304+
.addHeader(HttpHeaders.CONTENT_TYPE, MediaType.TEXT_EVENT_STREAM_VALUE)
305+
.setBody(
306+
"""
307+
{
308+
"type": "message_start",
309+
"message": {
310+
"id": "msg_1nZdL29xx5MUA1yADyHTEsnR8uuvGzszyY",
311+
"type": "message",
312+
"role": "assistant",
313+
"content": [],
314+
"model": "claude-opus-4-20250514",
315+
"stop_reason": null,
316+
"stop_sequence": null,
317+
"usage": {
318+
"input_tokens": 25,
319+
"output_tokens": 1
320+
}
321+
}
322+
}
323+
""".replace("\n", ""));
324+
mockWebServer.enqueue(mockResponse);
325+
326+
AnthropicApi.AnthropicMessage chatCompletionMessage = new AnthropicApi.AnthropicMessage(
327+
List.of(new AnthropicApi.ContentBlock("Hello world")), AnthropicApi.Role.USER);
328+
AnthropicApi.ChatCompletionRequest request = AnthropicApi.ChatCompletionRequest.builder()
329+
.model(AnthropicApi.ChatModel.CLAUDE_3_OPUS)
330+
.temperature(0.8)
331+
.messages(List.of(chatCompletionMessage))
332+
.stream(true)
333+
.build();
334+
MultiValueMap<String, String> additionalHeaders = new LinkedMultiValueMap<>();
335+
additionalHeaders.add("x-api-key", "additional-key");
336+
337+
api.chatCompletionStream(request, additionalHeaders)
338+
.collectList()
339+
.block();
340+
RecordedRequest recordedRequest = mockWebServer.takeRequest();
341+
assertThat(recordedRequest.getHeader(HttpHeaders.AUTHORIZATION)).isNull();
342+
assertThat(recordedRequest.getHeader("x-api-key")).isEqualTo("additional-key");
343+
}
344+
233345
}
234346

235347
}

0 commit comments

Comments
 (0)