Skip to content

Resolve OpenAI ApiKey for every request #3364

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 1 commit into from
Closed

Resolve OpenAI ApiKey for every request #3364

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 10 additions & 4 deletions models/spring-ai-openai/src/main/java/org/springframework/ai/openai/api/OpenAiApi.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -128,22 +128,28 @@ 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);
};
this.restClient = restClientBuilder.clone()
.baseUrl(baseUrl)
.defaultHeaders(finalHeaders)
.defaultStatusHandler(responseErrorHandler)
.defaultRequest(requestHeadersSpec -> {
if (!(apiKey instanceof NoopApiKey)) {
requestHeadersSpec.header(HttpHeaders.AUTHORIZATION, "Bearer " + apiKey.getValue());
}
})
.build();

this.webClient = webClientBuilder.clone()
.baseUrl(baseUrl)
.defaultHeaders(finalHeaders)
.defaultRequest(requestHeadersSpec -> {
if (!(apiKey instanceof NoopApiKey)) {
requestHeadersSpec.header(HttpHeaders.AUTHORIZATION, "Bearer " + apiKey.getValue());
}
})
.build(); // @formatter:on
}

Expand Down
20 changes: 15 additions & 5 deletions models/spring-ai-openai/src/main/java/org/springframework/ai/openai/api/OpenAiAudioApi.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -71,20 +71,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() {
Expand Down
12 changes: 8 additions & 4 deletions models/spring-ai-openai/src/main/java/org/springframework/ai/openai/api/OpenAiImageApi.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -62,15 +63,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

Expand Down
22 changes: 15 additions & 7 deletions ...spring-ai-openai/src/main/java/org/springframework/ai/openai/api/OpenAiModerationApi.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -64,13 +65,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) {
Expand Down
139 changes: 139 additions & 0 deletions ...ring-ai-openai/src/test/java/org/springframework/ai/openai/api/OpenAiApiBuilderTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,10 +16,27 @@

package org.springframework.ai.openai.api;

import java.io.IOException;
import java.util.LinkedList;
import java.util.List;
import java.util.Objects;
import java.util.Queue;

import okhttp3.mockwebserver.MockResponse;
import okhttp3.mockwebserver.MockWebServer;
import okhttp3.mockwebserver.RecordedRequest;

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;
Expand Down Expand Up @@ -142,4 +159,126 @@ void testInvalidResponseErrorHandler() {
.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")));
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("""
{
"id": "chatcmpl-12345",
"object": "chat.completion",
"created": 1677858242,
"model": "gpt-3.5-turbo",
"choices": [
{
"index": 0,
"message": {
"role": "assistant",
"content": "Hello world"
},
"finish_reason": "stop"
}
],
"usage": {
"prompt_tokens": 10,
"completion_tokens": 5,
"total_tokens": 15
}
}
""");
mockWebServer.enqueue(mockResponse);
mockWebServer.enqueue(mockResponse);

OpenAiApi.ChatCompletionMessage chatCompletionMessage = new OpenAiApi.ChatCompletionMessage("Hello world",
OpenAiApi.ChatCompletionMessage.Role.USER);
OpenAiApi.ChatCompletionRequest request = new OpenAiApi.ChatCompletionRequest(
List.of(chatCompletionMessage), "gpt-3.5-turbo", 0.8, false);
ResponseEntity<OpenAiApi.ChatCompletion> response = api.chatCompletionEntity(request);
assertThat(response.getStatusCode()).isEqualTo(HttpStatus.OK);
RecordedRequest recordedRequest = mockWebServer.takeRequest();
assertThat(recordedRequest.getHeader(HttpHeaders.AUTHORIZATION)).isEqualTo("Bearer key1");

response = api.chatCompletionEntity(request);
assertThat(response.getStatusCode()).isEqualTo(HttpStatus.OK);

recordedRequest = mockWebServer.takeRequest();
assertThat(recordedRequest.getHeader(HttpHeaders.AUTHORIZATION)).isEqualTo("Bearer key2");
}

@Test
void dynamicApiKeyWebClient() 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("""
{
"id": "chatcmpl-12345",
"object": "chat.completion",
"created": 1677858242,
"model": "gpt-3.5-turbo",
"choices": [
{
"index": 0,
"message": {
"role": "assistant",
"content": "Hello world"
},
"finish_reason": "stop"
}
],
"usage": {
"prompt_tokens": 10,
"completion_tokens": 5,
"total_tokens": 15
}
}
""".replace("\n", ""));
mockWebServer.enqueue(mockResponse);
mockWebServer.enqueue(mockResponse);

OpenAiApi.ChatCompletionMessage chatCompletionMessage = new OpenAiApi.ChatCompletionMessage("Hello world",
OpenAiApi.ChatCompletionMessage.Role.USER);
OpenAiApi.ChatCompletionRequest request = new OpenAiApi.ChatCompletionRequest(
List.of(chatCompletionMessage), "gpt-3.5-turbo", 0.8, true);
List<OpenAiApi.ChatCompletionChunk> response = api.chatCompletionStream(request).collectList().block();
assertThat(response).hasSize(1);
RecordedRequest recordedRequest = mockWebServer.takeRequest();
assertThat(recordedRequest.getHeader(HttpHeaders.AUTHORIZATION)).isEqualTo("Bearer key1");

response = api.chatCompletionStream(request).collectList().block();
assertThat(response).hasSize(1);

recordedRequest = mockWebServer.takeRequest();
assertThat(recordedRequest.getHeader(HttpHeaders.AUTHORIZATION)).isEqualTo("Bearer key2");
}

}

}
Loading