Add support for ApiKey for Anthropic and use it dynamically for every request

Fixes #3365

 - Set ApiKey as late as possible

Signed-off-by: Filip Hrisafov <filip.hrisafov@gmail.com>
(cherry picked from commit 0a1cf81cfc)
This commit is contained in:
Filip Hrisafov
2025-05-29 01:18:25 +02:00
committed by Spring Builds
parent 5ecfcce163
commit ae9284cb01
2 changed files with 377 additions and 6 deletions

View File

@@ -34,8 +34,10 @@ import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import org.springframework.ai.anthropic.api.StreamHelper.ChatCompletionResponseBuilder;
import org.springframework.ai.model.ApiKey;
import org.springframework.ai.model.ChatModelDescription;
import org.springframework.ai.model.ModelOptionsUtils;
import org.springframework.ai.model.SimpleApiKey;
import org.springframework.ai.observation.conventions.AiProvider;
import org.springframework.ai.retry.RetryUtils;
import org.springframework.http.HttpHeaders;
@@ -60,6 +62,7 @@ import org.springframework.web.reactive.function.client.WebClient;
* @author Alexandros Pappas
* @author Jonghoon Park
* @author Claudio Silva Junior
* @author Filip Hrisafov
* @since 1.0.0
*/
public final class AnthropicApi {
@@ -96,6 +99,8 @@ public final class AnthropicApi {
private final WebClient webClient;
private final ApiKey apiKey;
/**
* Create a new client api.
* @param baseUrl api base URL.
@@ -107,18 +112,18 @@ public final class AnthropicApi {
* @param responseErrorHandler Response error handler.
* @param anthropicBetaFeatures Anthropic beta features.
*/
private AnthropicApi(String baseUrl, String completionsPath, String anthropicApiKey, String anthropicVersion,
private AnthropicApi(String baseUrl, String completionsPath, ApiKey anthropicApiKey, String anthropicVersion,
RestClient.Builder restClientBuilder, WebClient.Builder webClientBuilder,
ResponseErrorHandler responseErrorHandler, String anthropicBetaFeatures) {
Consumer<HttpHeaders> jsonContentHeaders = headers -> {
headers.add(HEADER_X_API_KEY, anthropicApiKey);
headers.add(HEADER_ANTHROPIC_VERSION, anthropicVersion);
headers.add(HEADER_ANTHROPIC_BETA, anthropicBetaFeatures);
headers.setContentType(MediaType.APPLICATION_JSON);
};
this.completionsPath = completionsPath;
this.apiKey = anthropicApiKey;
this.restClient = restClientBuilder.clone()
.baseUrl(baseUrl)
@@ -160,12 +165,17 @@ public final class AnthropicApi {
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(ChatCompletionResponse.class);
// @formatter:on
}
/**
@@ -196,9 +206,13 @@ public final class AnthropicApi {
AtomicReference<ChatCompletionResponseBuilder> chatCompletionReference = new AtomicReference<>();
// @formatter:off
return this.webClient.post()
.uri(this.completionsPath)
.headers(headers -> headers.addAll(additionalHttpHeader))
.headers(headers -> {
headers.addAll(additionalHttpHeader);
addDefaultHeadersIfMissing(headers);
}) // @formatter:off
.body(Mono.just(chatRequest), ChatCompletionRequest.class)
.retrieve()
.bodyToFlux(String.class)
@@ -232,6 +246,15 @@ public final class AnthropicApi {
.filter(chatCompletionResponse -> chatCompletionResponse.type() != null);
}
private void addDefaultHeadersIfMissing(HttpHeaders headers) {
if (!headers.containsKey(HEADER_X_API_KEY)) {
String apiKeyValue = this.apiKey.getValue();
if (StringUtils.hasText(apiKeyValue)) {
headers.add(HEADER_X_API_KEY, apiKeyValue);
}
}
}
/**
* Check the <a href="https://docs.anthropic.com/claude/docs/models-overview">Models
* overview</a> and <a href=
@@ -1349,7 +1372,7 @@ public final class AnthropicApi {
private String completionsPath = DEFAULT_MESSAGE_COMPLETIONS_PATH;
private String apiKey;
private ApiKey apiKey;
private String anthropicVersion = DEFAULT_ANTHROPIC_VERSION;
@@ -1373,12 +1396,18 @@ public final class AnthropicApi {
return this;
}
public Builder apiKey(String apiKey) {
public Builder apiKey(ApiKey apiKey) {
Assert.notNull(apiKey, "apiKey cannot be null");
this.apiKey = apiKey;
return this;
}
public Builder apiKey(String simpleApiKey) {
Assert.notNull(simpleApiKey, "simpleApiKey cannot be null");
this.apiKey = new SimpleApiKey(simpleApiKey);
return this;
}
public Builder anthropicVersion(String anthropicVersion) {
Assert.notNull(anthropicVersion, "anthropicVersion cannot be null");
this.anthropicVersion = anthropicVersion;

View File

@@ -0,0 +1,342 @@
/*
* Copyright 2023-2025 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.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.anthropic.api;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
import static org.mockito.Mockito.mock;
import java.io.IOException;
import java.util.LinkedList;
import java.util.List;
import java.util.Objects;
import java.util.Queue;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Nested;
import org.junit.jupiter.api.Test;
import org.springframework.ai.model.ApiKey;
import org.springframework.ai.model.SimpleApiKey;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.util.LinkedMultiValueMap;
import org.springframework.util.MultiValueMap;
import org.springframework.web.client.ResponseErrorHandler;
import org.springframework.web.client.RestClient;
import org.springframework.web.reactive.function.client.WebClient;
import okhttp3.mockwebserver.MockResponse;
import okhttp3.mockwebserver.MockWebServer;
import okhttp3.mockwebserver.RecordedRequest;
import org.opentest4j.AssertionFailedError;
/**
* @author Filip Hrisafov
*/
public class AnthropicApiBuilderTests {
private static final ApiKey TEST_API_KEY = new SimpleApiKey("test-api-key");
private static final String TEST_BASE_URL = "https://test.anthropic.com";
private static final String TEST_COMPLETIONS_PATH = "/test/completions";
@Test
void testMinimalBuilder() {
AnthropicApi api = AnthropicApi.builder().apiKey(TEST_API_KEY).build();
assertThat(api).isNotNull();
}
@Test
void testFullBuilder() {
RestClient.Builder restClientBuilder = RestClient.builder();
WebClient.Builder webClientBuilder = WebClient.builder();
ResponseErrorHandler errorHandler = mock(ResponseErrorHandler.class);
AnthropicApi api = AnthropicApi.builder()
.apiKey(TEST_API_KEY)
.baseUrl(TEST_BASE_URL)
.completionsPath(TEST_COMPLETIONS_PATH)
.restClientBuilder(restClientBuilder)
.webClientBuilder(webClientBuilder)
.responseErrorHandler(errorHandler)
.build();
assertThat(api).isNotNull();
}
@Test
void testMissingApiKey() {
assertThatThrownBy(() -> AnthropicApi.builder().build()).isInstanceOf(IllegalArgumentException.class)
.hasMessageContaining("apiKey must be set");
}
@Test
void testInvalidBaseUrl() {
assertThatThrownBy(() -> AnthropicApi.builder().baseUrl("").build())
.isInstanceOf(IllegalArgumentException.class)
.hasMessageContaining("baseUrl cannot be null or empty");
assertThatThrownBy(() -> AnthropicApi.builder().baseUrl(null).build())
.isInstanceOf(IllegalArgumentException.class)
.hasMessageContaining("baseUrl cannot be null or empty");
}
@Test
void testInvalidCompletionsPath() {
assertThatThrownBy(() -> AnthropicApi.builder().completionsPath("").build())
.isInstanceOf(IllegalArgumentException.class)
.hasMessageContaining("completionsPath cannot be null or empty");
assertThatThrownBy(() -> AnthropicApi.builder().completionsPath(null).build())
.isInstanceOf(IllegalArgumentException.class)
.hasMessageContaining("completionsPath cannot be null or empty");
}
@Test
void testInvalidRestClientBuilder() {
assertThatThrownBy(() -> AnthropicApi.builder().restClientBuilder(null).build())
.isInstanceOf(IllegalArgumentException.class)
.hasMessageContaining("restClientBuilder cannot be null");
}
@Test
void testInvalidWebClientBuilder() {
assertThatThrownBy(() -> AnthropicApi.builder().webClientBuilder(null).build())
.isInstanceOf(IllegalArgumentException.class)
.hasMessageContaining("webClientBuilder cannot be null");
}
@Test
void testInvalidResponseErrorHandler() {
assertThatThrownBy(() -> AnthropicApi.builder().responseErrorHandler(null).build())
.isInstanceOf(IllegalArgumentException.class)
.hasMessageContaining("responseErrorHandler cannot be null");
}
@Nested
class MockRequests {
MockWebServer mockWebServer;
@BeforeEach
void setUp() throws IOException {
mockWebServer = new MockWebServer();
mockWebServer.start();
}
@AfterEach
void tearDown() throws IOException {
mockWebServer.shutdown();
}
@Test
void dynamicApiKeyRestClient() throws InterruptedException {
Queue<ApiKey> apiKeys = new LinkedList<>(List.of(new SimpleApiKey("key1"), new SimpleApiKey("key2")));
AnthropicApi api = AnthropicApi.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("""
{
"id": "msg_1nZdL29xx5MUA1yADyHTEsnR8uuvGzszyY",
"type": "message",
"role": "assistant",
"content": [],
"model": "claude-opus-3-latest",
"stop_reason": null,
"stop_sequence": null,
"usage": {
"input_tokens": 25,
"output_tokens": 1
}
}
""");
mockWebServer.enqueue(mockResponse);
mockWebServer.enqueue(mockResponse);
AnthropicApi.AnthropicMessage chatCompletionMessage = new AnthropicApi.AnthropicMessage(
List.of(new AnthropicApi.ContentBlock("Hello world")), AnthropicApi.Role.USER);
AnthropicApi.ChatCompletionRequest request = AnthropicApi.ChatCompletionRequest.builder()
.model(AnthropicApi.ChatModel.CLAUDE_3_OPUS)
.temperature(0.8)
.messages(List.of(chatCompletionMessage))
.build();
ResponseEntity<AnthropicApi.ChatCompletionResponse> response = api.chatCompletionEntity(request);
assertThat(response.getStatusCode()).isEqualTo(HttpStatus.OK);
RecordedRequest recordedRequest = mockWebServer.takeRequest();
assertThat(recordedRequest.getHeader(HttpHeaders.AUTHORIZATION)).isNull();
assertThat(recordedRequest.getHeader("x-api-key")).isEqualTo("key1");
response = api.chatCompletionEntity(request);
assertThat(response.getStatusCode()).isEqualTo(HttpStatus.OK);
recordedRequest = mockWebServer.takeRequest();
assertThat(recordedRequest.getHeader(HttpHeaders.AUTHORIZATION)).isNull();
assertThat(recordedRequest.getHeader("x-api-key")).isEqualTo("key2");
}
@Test
void dynamicApiKeyRestClientWithAdditionalApiKeyHeader() throws InterruptedException {
AnthropicApi api = AnthropicApi.builder().apiKey(() -> {
throw new AssertionFailedError("Should not be called, API key is provided in headers");
}).baseUrl(mockWebServer.url("/").toString()).build();
MockResponse mockResponse = new MockResponse().setResponseCode(200)
.addHeader(HttpHeaders.CONTENT_TYPE, MediaType.APPLICATION_JSON_VALUE)
.setBody("""
{
"id": "msg_1nZdL29xx5MUA1yADyHTEsnR8uuvGzszyY",
"type": "message",
"role": "assistant",
"content": [],
"model": "claude-opus-3-latest",
"stop_reason": null,
"stop_sequence": null,
"usage": {
"input_tokens": 25,
"output_tokens": 1
}
}
""");
mockWebServer.enqueue(mockResponse);
AnthropicApi.AnthropicMessage chatCompletionMessage = new AnthropicApi.AnthropicMessage(
List.of(new AnthropicApi.ContentBlock("Hello world")), AnthropicApi.Role.USER);
AnthropicApi.ChatCompletionRequest request = AnthropicApi.ChatCompletionRequest.builder()
.model(AnthropicApi.ChatModel.CLAUDE_3_OPUS)
.temperature(0.8)
.messages(List.of(chatCompletionMessage))
.build();
MultiValueMap<String, String> additionalHeaders = new LinkedMultiValueMap<>();
additionalHeaders.add("x-api-key", "additional-key");
ResponseEntity<AnthropicApi.ChatCompletionResponse> response = api.chatCompletionEntity(request,
additionalHeaders);
assertThat(response.getStatusCode()).isEqualTo(HttpStatus.OK);
RecordedRequest recordedRequest = mockWebServer.takeRequest();
assertThat(recordedRequest.getHeader(HttpHeaders.AUTHORIZATION)).isNull();
assertThat(recordedRequest.getHeader("x-api-key")).isEqualTo("additional-key");
}
@Test
void dynamicApiKeyWebClient() throws InterruptedException {
Queue<ApiKey> apiKeys = new LinkedList<>(List.of(new SimpleApiKey("key1"), new SimpleApiKey("key2")));
AnthropicApi api = AnthropicApi.builder()
.apiKey(() -> Objects.requireNonNull(apiKeys.poll()).getValue())
.baseUrl(mockWebServer.url("/").toString())
.build();
MockResponse mockResponse = new MockResponse().setResponseCode(200)
.addHeader(HttpHeaders.CONTENT_TYPE, MediaType.TEXT_EVENT_STREAM_VALUE)
.setBody("""
{
"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
}
}
}
""".replace("\n", ""));
mockWebServer.enqueue(mockResponse);
mockWebServer.enqueue(mockResponse);
AnthropicApi.AnthropicMessage chatCompletionMessage = new AnthropicApi.AnthropicMessage(
List.of(new AnthropicApi.ContentBlock("Hello world")), AnthropicApi.Role.USER);
AnthropicApi.ChatCompletionRequest request = AnthropicApi.ChatCompletionRequest.builder()
.model(AnthropicApi.ChatModel.CLAUDE_3_OPUS)
.temperature(0.8)
.messages(List.of(chatCompletionMessage))
.stream(true)
.build();
api.chatCompletionStream(request).collectList().block();
RecordedRequest recordedRequest = mockWebServer.takeRequest();
assertThat(recordedRequest.getHeader(HttpHeaders.AUTHORIZATION)).isNull();
assertThat(recordedRequest.getHeader("x-api-key")).isEqualTo("key1");
api.chatCompletionStream(request).collectList().block();
recordedRequest = mockWebServer.takeRequest();
assertThat(recordedRequest.getHeader(HttpHeaders.AUTHORIZATION)).isNull();
assertThat(recordedRequest.getHeader("x-api-key")).isEqualTo("key2");
}
@Test
void dynamicApiKeyWebClientWithAdditionalApiKey() throws InterruptedException {
Queue<ApiKey> apiKeys = new LinkedList<>(List.of(new SimpleApiKey("key1"), new SimpleApiKey("key2")));
AnthropicApi api = AnthropicApi.builder()
.apiKey(() -> Objects.requireNonNull(apiKeys.poll()).getValue())
.baseUrl(mockWebServer.url("/").toString())
.build();
MockResponse mockResponse = new MockResponse().setResponseCode(200)
.addHeader(HttpHeaders.CONTENT_TYPE, MediaType.TEXT_EVENT_STREAM_VALUE)
.setBody("""
{
"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
}
}
}
""".replace("\n", ""));
mockWebServer.enqueue(mockResponse);
AnthropicApi.AnthropicMessage chatCompletionMessage = new AnthropicApi.AnthropicMessage(
List.of(new AnthropicApi.ContentBlock("Hello world")), AnthropicApi.Role.USER);
AnthropicApi.ChatCompletionRequest request = AnthropicApi.ChatCompletionRequest.builder()
.model(AnthropicApi.ChatModel.CLAUDE_3_OPUS)
.temperature(0.8)
.messages(List.of(chatCompletionMessage))
.stream(true)
.build();
MultiValueMap<String, String> additionalHeaders = new LinkedMultiValueMap<>();
additionalHeaders.add("x-api-key", "additional-key");
api.chatCompletionStream(request, additionalHeaders).collectList().block();
RecordedRequest recordedRequest = mockWebServer.takeRequest();
assertThat(recordedRequest.getHeader(HttpHeaders.AUTHORIZATION)).isNull();
assertThat(recordedRequest.getHeader("x-api-key")).isEqualTo("additional-key");
}
}
}