Skip to content

Commit 00d8008

Browse files
filiphrfiliphr
committed
Add support for ApiKey for Anthropic and use it dynamically for every request
1 parent 925f176 commit 00d8008

File tree

2 files changed

+258
-4
lines changed

2 files changed

+258
-4
lines changed

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

Lines changed: 23 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -34,8 +34,10 @@
3434
import reactor.core.publisher.Mono;
3535

3636
import org.springframework.ai.anthropic.api.StreamHelper.ChatCompletionResponseBuilder;
37+
import org.springframework.ai.model.ApiKey;
3738
import org.springframework.ai.model.ChatModelDescription;
3839
import org.springframework.ai.model.ModelOptionsUtils;
40+
import org.springframework.ai.model.SimpleApiKey;
3941
import org.springframework.ai.observation.conventions.AiProvider;
4042
import org.springframework.ai.retry.RetryUtils;
4143
import org.springframework.http.HttpHeaders;
@@ -107,12 +109,11 @@ public static Builder builder() {
107109
* @param responseErrorHandler Response error handler.
108110
* @param anthropicBetaFeatures Anthropic beta features.
109111
*/
110-
private AnthropicApi(String baseUrl, String completionsPath, String anthropicApiKey, String anthropicVersion,
112+
private AnthropicApi(String baseUrl, String completionsPath, ApiKey anthropicApiKey, String anthropicVersion,
111113
RestClient.Builder restClientBuilder, WebClient.Builder webClientBuilder,
112114
ResponseErrorHandler responseErrorHandler, String anthropicBetaFeatures) {
113115

114116
Consumer<HttpHeaders> jsonContentHeaders = headers -> {
115-
headers.add(HEADER_X_API_KEY, anthropicApiKey);
116117
headers.add(HEADER_ANTHROPIC_VERSION, anthropicVersion);
117118
headers.add(HEADER_ANTHROPIC_BETA, anthropicBetaFeatures);
118119
headers.setContentType(MediaType.APPLICATION_JSON);
@@ -124,6 +125,12 @@ private AnthropicApi(String baseUrl, String completionsPath, String anthropicApi
124125
.baseUrl(baseUrl)
125126
.defaultHeaders(jsonContentHeaders)
126127
.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+
})
127134
.build();
128135

129136
this.webClient = webClientBuilder.clone()
@@ -133,6 +140,12 @@ private AnthropicApi(String baseUrl, String completionsPath, String anthropicApi
133140
resp -> resp.bodyToMono(String.class)
134141
.flatMap(it -> Mono.error(new RuntimeException(
135142
"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+
})
136149
.build();
137150
}
138151

@@ -1339,7 +1352,7 @@ public static class Builder {
13391352

13401353
private String completionsPath = DEFAULT_MESSAGE_COMPLETIONS_PATH;
13411354

1342-
private String apiKey;
1355+
private ApiKey apiKey;
13431356

13441357
private String anthropicVersion = DEFAULT_ANTHROPIC_VERSION;
13451358

@@ -1363,12 +1376,18 @@ public Builder completionsPath(String completionsPath) {
13631376
return this;
13641377
}
13651378

1366-
public Builder apiKey(String apiKey) {
1379+
public Builder apiKey(ApiKey apiKey) {
13671380
Assert.notNull(apiKey, "apiKey cannot be null");
13681381
this.apiKey = apiKey;
13691382
return this;
13701383
}
13711384

1385+
public Builder apiKey(String simpleApiKey) {
1386+
Assert.notNull(simpleApiKey, "simpleApiKey cannot be null");
1387+
this.apiKey = new SimpleApiKey(simpleApiKey);
1388+
return this;
1389+
}
1390+
13721391
public Builder anthropicVersion(String anthropicVersion) {
13731392
Assert.notNull(anthropicVersion, "anthropicVersion cannot be null");
13741393
this.anthropicVersion = anthropicVersion;

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

Lines changed: 235 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,235 @@
1+
/*
2+
* Copyright 2023-2025 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package org.springframework.ai.anthropic.api;
18+
19+
import static org.assertj.core.api.Assertions.assertThat;
20+
import static org.assertj.core.api.Assertions.assertThatThrownBy;
21+
import static org.mockito.Mockito.mock;
22+
23+
import java.io.IOException;
24+
import java.util.LinkedList;
25+
import java.util.List;
26+
import java.util.Objects;
27+
import java.util.Queue;
28+
29+
import org.junit.jupiter.api.AfterEach;
30+
import org.junit.jupiter.api.BeforeEach;
31+
import org.junit.jupiter.api.Nested;
32+
import org.junit.jupiter.api.Test;
33+
import org.springframework.ai.model.ApiKey;
34+
import org.springframework.ai.model.SimpleApiKey;
35+
import org.springframework.http.HttpHeaders;
36+
import org.springframework.http.HttpStatus;
37+
import org.springframework.http.MediaType;
38+
import org.springframework.http.ResponseEntity;
39+
import org.springframework.web.client.ResponseErrorHandler;
40+
import org.springframework.web.client.RestClient;
41+
import org.springframework.web.reactive.function.client.WebClient;
42+
43+
import okhttp3.mockwebserver.MockResponse;
44+
import okhttp3.mockwebserver.MockWebServer;
45+
import okhttp3.mockwebserver.RecordedRequest;
46+
47+
public class AnthropicApiBuilderTests {
48+
49+
private static final ApiKey TEST_API_KEY = new SimpleApiKey("test-api-key");
50+
51+
private static final String TEST_BASE_URL = "https://test.anthropic.com";
52+
53+
private static final String TEST_COMPLETIONS_PATH = "/test/completions";
54+
55+
@Test
56+
void testMinimalBuilder() {
57+
AnthropicApi api = AnthropicApi.builder().apiKey(TEST_API_KEY).build();
58+
59+
assertThat(api).isNotNull();
60+
}
61+
62+
@Test
63+
void testFullBuilder() {
64+
RestClient.Builder restClientBuilder = RestClient.builder();
65+
WebClient.Builder webClientBuilder = WebClient.builder();
66+
ResponseErrorHandler errorHandler = mock(ResponseErrorHandler.class);
67+
68+
AnthropicApi api = AnthropicApi.builder()
69+
.apiKey(TEST_API_KEY)
70+
.baseUrl(TEST_BASE_URL)
71+
.completionsPath(TEST_COMPLETIONS_PATH)
72+
.restClientBuilder(restClientBuilder)
73+
.webClientBuilder(webClientBuilder)
74+
.responseErrorHandler(errorHandler)
75+
.build();
76+
77+
assertThat(api).isNotNull();
78+
}
79+
80+
@Test
81+
void testMissingApiKey() {
82+
assertThatThrownBy(() -> AnthropicApi.builder().build()).isInstanceOf(IllegalArgumentException.class)
83+
.hasMessageContaining("apiKey must be set");
84+
}
85+
86+
@Test
87+
void testInvalidBaseUrl() {
88+
assertThatThrownBy(() -> AnthropicApi.builder().baseUrl("").build())
89+
.isInstanceOf(IllegalArgumentException.class)
90+
.hasMessageContaining("baseUrl cannot be null or empty");
91+
92+
assertThatThrownBy(() -> AnthropicApi.builder().baseUrl(null).build())
93+
.isInstanceOf(IllegalArgumentException.class)
94+
.hasMessageContaining("baseUrl cannot be null or empty");
95+
}
96+
97+
@Test
98+
void testInvalidCompletionsPath() {
99+
assertThatThrownBy(() -> AnthropicApi.builder().completionsPath("").build())
100+
.isInstanceOf(IllegalArgumentException.class)
101+
.hasMessageContaining("completionsPath cannot be null or empty");
102+
103+
assertThatThrownBy(() -> AnthropicApi.builder().completionsPath(null).build())
104+
.isInstanceOf(IllegalArgumentException.class)
105+
.hasMessageContaining("completionsPath cannot be null or empty");
106+
}
107+
108+
@Test
109+
void testInvalidRestClientBuilder() {
110+
assertThatThrownBy(() -> AnthropicApi.builder().restClientBuilder(null).build())
111+
.isInstanceOf(IllegalArgumentException.class)
112+
.hasMessageContaining("restClientBuilder cannot be null");
113+
}
114+
115+
@Test
116+
void testInvalidWebClientBuilder() {
117+
assertThatThrownBy(() -> AnthropicApi.builder().webClientBuilder(null).build())
118+
.isInstanceOf(IllegalArgumentException.class)
119+
.hasMessageContaining("webClientBuilder cannot be null");
120+
}
121+
122+
@Test
123+
void testInvalidResponseErrorHandler() {
124+
assertThatThrownBy(() -> AnthropicApi.builder().responseErrorHandler(null).build())
125+
.isInstanceOf(IllegalArgumentException.class)
126+
.hasMessageContaining("responseErrorHandler cannot be null");
127+
}
128+
129+
@Nested
130+
class MockRequests {
131+
132+
MockWebServer mockWebServer;
133+
134+
@BeforeEach
135+
void setUp() throws IOException {
136+
mockWebServer = new MockWebServer();
137+
mockWebServer.start();
138+
}
139+
140+
@AfterEach
141+
void tearDown() throws IOException {
142+
mockWebServer.shutdown();
143+
}
144+
145+
@Test
146+
void dynamicApiKeyRestClient() throws InterruptedException {
147+
Queue<ApiKey> apiKeys = new LinkedList<>(List.of(new SimpleApiKey("key1"), new SimpleApiKey("key2")));
148+
AnthropicApi api = AnthropicApi.builder()
149+
.apiKey(() -> Objects.requireNonNull(apiKeys.poll()).getValue())
150+
.baseUrl(mockWebServer.url("/").toString())
151+
.build();
152+
153+
MockResponse mockResponse = new MockResponse().setResponseCode(200)
154+
.addHeader(HttpHeaders.CONTENT_TYPE, MediaType.APPLICATION_JSON_VALUE)
155+
.setBody("""
156+
{
157+
"id": "msg_1nZdL29xx5MUA1yADyHTEsnR8uuvGzszyY",
158+
"type": "message",
159+
"role": "assistant",
160+
"content": [],
161+
"model": "claude-opus-3-latest",
162+
"stop_reason": null,
163+
"stop_sequence": null,
164+
"usage": {
165+
"input_tokens": 25,
166+
"output_tokens": 1
167+
}
168+
}
169+
""");
170+
mockWebServer.enqueue(mockResponse);
171+
mockWebServer.enqueue(mockResponse);
172+
173+
AnthropicApi.AnthropicMessage chatCompletionMessage = new AnthropicApi.AnthropicMessage(
174+
List.of(new AnthropicApi.ContentBlock("Hello world")), AnthropicApi.Role.USER);
175+
AnthropicApi.ChatCompletionRequest request = AnthropicApi.ChatCompletionRequest.builder()
176+
.model(AnthropicApi.ChatModel.CLAUDE_3_OPUS)
177+
.temperature(0.8)
178+
.messages(List.of(chatCompletionMessage))
179+
.build();
180+
ResponseEntity<AnthropicApi.ChatCompletionResponse> response = api.chatCompletionEntity(request);
181+
assertThat(response.getStatusCode()).isEqualTo(HttpStatus.OK);
182+
RecordedRequest recordedRequest = mockWebServer.takeRequest();
183+
assertThat(recordedRequest.getHeader(HttpHeaders.AUTHORIZATION)).isNull();
184+
assertThat(recordedRequest.getHeader("x-api-key")).isEqualTo("key1");
185+
186+
response = api.chatCompletionEntity(request);
187+
assertThat(response.getStatusCode()).isEqualTo(HttpStatus.OK);
188+
189+
recordedRequest = mockWebServer.takeRequest();
190+
assertThat(recordedRequest.getHeader(HttpHeaders.AUTHORIZATION)).isNull();
191+
assertThat(recordedRequest.getHeader("x-api-key")).isEqualTo("key2");
192+
}
193+
194+
@Test
195+
void dynamicApiKeyWebClient() throws InterruptedException {
196+
Queue<ApiKey> apiKeys = new LinkedList<>(List.of(new SimpleApiKey("key1"), new SimpleApiKey("key2")));
197+
AnthropicApi api = AnthropicApi.builder()
198+
.apiKey(() -> Objects.requireNonNull(apiKeys.poll()).getValue())
199+
.baseUrl(mockWebServer.url("/").toString())
200+
.build();
201+
202+
MockResponse mockResponse = new MockResponse().setResponseCode(200)
203+
.addHeader(HttpHeaders.CONTENT_TYPE, MediaType.TEXT_EVENT_STREAM_VALUE)
204+
.setBody(
205+
"""
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+
""");
208+
mockWebServer.enqueue(mockResponse);
209+
mockWebServer.enqueue(mockResponse);
210+
211+
AnthropicApi.AnthropicMessage chatCompletionMessage = new AnthropicApi.AnthropicMessage(
212+
List.of(new AnthropicApi.ContentBlock("Hello world")), AnthropicApi.Role.USER);
213+
AnthropicApi.ChatCompletionRequest request = AnthropicApi.ChatCompletionRequest.builder()
214+
.model(AnthropicApi.ChatModel.CLAUDE_3_OPUS)
215+
.temperature(0.8)
216+
.messages(List.of(chatCompletionMessage))
217+
.stream(true)
218+
.build();
219+
List<AnthropicApi.ChatCompletionResponse> response = api.chatCompletionStream(request)
220+
.collectList()
221+
.block();
222+
RecordedRequest recordedRequest = mockWebServer.takeRequest();
223+
assertThat(recordedRequest.getHeader(HttpHeaders.AUTHORIZATION)).isNull();
224+
assertThat(recordedRequest.getHeader("x-api-key")).isEqualTo("key1");
225+
226+
response = api.chatCompletionStream(request).collectList().block();
227+
228+
recordedRequest = mockWebServer.takeRequest();
229+
assertThat(recordedRequest.getHeader(HttpHeaders.AUTHORIZATION)).isNull();
230+
assertThat(recordedRequest.getHeader("x-api-key")).isEqualTo("key2");
231+
}
232+
233+
}
234+
235+
}

0 commit comments

Comments
 (0)