Spring/spring-ai
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Refactor and centralize Retry logic:

Browse Source 
- Establish a new "spring-ai-retry" project, implementing a default HTTP error handler,
   RetryTemplate, and handling both Transient and Non-Transient Exceptions.
 - Streamline existing clients (e.g., OpenAI and MistralAI) to utilize "spring-ai-retry."
 - Integrate retry auto-configuration with customizable properties, extending it to OpenAI and MistralAI Auto-Configs.
 - Allow configuration of RetryTemplate and ResponseErrorHandler for various clients, including OpenAIChatClient,
   OpenAiEmbeddingClient, OpenAiAudioTranscriptionCline, OpenAiImageClient, MistralAiChatClient, and MistralAiEmbeddingClient.
 - Add tests for default RestTemplate and ResponseErrorHandler configurations in OpenAI and MistralAI.
 - Introduce new retry auto-config properties: "onClientErrors" and "onHttpCodes".
 - Implement tests for retry auto-config properties.
 - Generate missing license headers.
This commit is contained in:
Christian Tzolov
2024-03-06 11:53:24 +01:00
committed by Mark Pollack
parent 78f73d17be
commit 1e3eaec7b9
57 changed files with 1415 additions and 389 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
models/spring-ai-mistral-ai/pom.xml
View File

@@ -29,24 +29,13 @@
<version>${project.parent.version}</version>
</dependency>
<dependency>
<groupId>org.springframework</groupId>
<artifactId>spring-web</artifactId>
<version>${spring-framework.version}</version>
</dependency>
<dependency>
<groupId>org.springframework.retry</groupId>
<artifactId>spring-retry</artifactId>
<version>2.0.4</version>
</dependency>
<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-retry</artifactId>
<version>${project.parent.version}</version>
</dependency>
<!-- Spring Framework -->
<dependency>
<groupId>org.springframework</groupId>
<artifactId>spring-webflux</artifactId>
</dependency>
<dependency>
<groupId>org.springframework</groupId>
<artifactId>spring-context-support</artifactId>

57
models/spring-ai-mistral-ai/src/main/java/org/springframework/ai/mistralai/MistralAiChatClient.java
View File

@@ -15,7 +15,6 @@
*/
package org.springframework.ai.mistralai;
import java.time.Duration;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
@@ -41,10 +40,8 @@ import org.springframework.ai.mistralai.api.MistralAiApi.ChatCompletionRequest;
import org.springframework.ai.model.ModelOptionsUtils;
import org.springframework.ai.model.function.AbstractFunctionCallSupport;
import org.springframework.ai.model.function.FunctionCallbackContext;
import org.springframework.ai.retry.RetryUtils;
import org.springframework.http.ResponseEntity;
import org.springframework.retry.RetryCallback;
import org.springframework.retry.RetryContext;
import org.springframework.retry.RetryListener;
import org.springframework.retry.support.RetryTemplate;
import org.springframework.util.Assert;
import org.springframework.util.CollectionUtils;
@@ -70,17 +67,7 @@ public class MistralAiChatClient extends
*/
private final MistralAiApi mistralAiApi;
private final RetryTemplate retryTemplate = RetryTemplate.builder()
.maxAttempts(10)
.retryOn(MistralAiApi.MistralAiApiException.class)
.exponentialBackoff(Duration.ofMillis(2000), 5, Duration.ofMillis(3 * 60000))
.withListener(new RetryListener() {
public <T extends Object, E extends Throwable> void onError(RetryContext context,
RetryCallback<T, E> callback, Throwable throwable) {
log.warn("Retry error. Retry count:" + context.getRetryCount(), throwable);
};
})
.build();
private final RetryTemplate retryTemplate;
public MistralAiChatClient(MistralAiApi mistralAiApi) {
this(mistralAiApi,
@@ -93,46 +80,50 @@ public class MistralAiChatClient extends
}
public MistralAiChatClient(MistralAiApi mistralAiApi, MistralAiChatOptions options) {
this(mistralAiApi, options, null);
this(mistralAiApi, options, null, RetryUtils.DEFAULT_RETRY_TEMPLATE);
}
public MistralAiChatClient(MistralAiApi mistralAiApi, MistralAiChatOptions options,
FunctionCallbackContext functionCallbackContext) {
FunctionCallbackContext functionCallbackContext, RetryTemplate retryTemplate) {
super(functionCallbackContext);
Assert.notNull(mistralAiApi, "MistralAiApi must not be null");
Assert.notNull(options, "Options must not be null");
Assert.notNull(retryTemplate, "RetryTemplate must not be null");
this.mistralAiApi = mistralAiApi;
this.defaultOptions = options;
this.retryTemplate = retryTemplate;
}
@Override
public ChatResponse call(Prompt prompt) {
// return retryTemplate.execute(ctx -> {
var request = createRequest(prompt, false);
// var completionEntity = this.mistralAiApi.chatCompletionEntity(request);
ResponseEntity<ChatCompletion> completionEntity = this.callWithFunctionSupport(request);
return retryTemplate.execute(ctx -> {
var chatCompletion = completionEntity.getBody();
if (chatCompletion == null) {
log.warn("No chat completion returned for prompt: {}", prompt);
return new ChatResponse(List.of());
}
ResponseEntity<ChatCompletion> completionEntity = this.callWithFunctionSupport(request);
List<Generation> generations = chatCompletion.choices()
.stream()
.map(choice -> new Generation(choice.message().content(), Map.of("role", choice.message().role().name()))
.withGenerationMetadata(ChatGenerationMetadata.from(choice.finishReason().name(), null)))
.toList();
var chatCompletion = completionEntity.getBody();
if (chatCompletion == null) {
log.warn("No chat completion returned for prompt: {}", prompt);
return new ChatResponse(List.of());
}
return new ChatResponse(generations);
// });
List<Generation> generations = chatCompletion.choices()
.stream()
.map(choice -> new Generation(choice.message().content(),
Map.of("role", choice.message().role().name()))
.withGenerationMetadata(ChatGenerationMetadata.from(choice.finishReason().name(), null)))
.toList();
return new ChatResponse(generations);
});
}
@Override
public Flux<ChatResponse> stream(Prompt prompt) {
var request = createRequest(prompt, true);
return retryTemplate.execute(ctx -> {
var request = createRequest(prompt, true);
var completionChunks = this.mistralAiApi.chatCompletionStream(request);

39
models/spring-ai-mistral-ai/src/main/java/org/springframework/ai/mistralai/MistralAiEmbeddingClient.java
View File

@@ -15,23 +15,25 @@
*/
package org.springframework.ai.mistralai;
import java.util.List;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.MetadataMode;
import org.springframework.ai.embedding.*;
import org.springframework.ai.embedding.AbstractEmbeddingClient;
import org.springframework.ai.embedding.Embedding;
import org.springframework.ai.embedding.EmbeddingOptions;
import org.springframework.ai.embedding.EmbeddingRequest;
import org.springframework.ai.embedding.EmbeddingResponse;
import org.springframework.ai.embedding.EmbeddingResponseMetadata;
import org.springframework.ai.mistralai.api.MistralAiApi;
import org.springframework.ai.mistralai.api.MistralAiApi.MistralAiApiException;
import org.springframework.ai.model.ModelOptionsUtils;
import org.springframework.retry.RetryCallback;
import org.springframework.retry.RetryContext;
import org.springframework.retry.RetryListener;
import org.springframework.ai.retry.RetryUtils;
import org.springframework.retry.support.RetryTemplate;
import org.springframework.util.Assert;
import java.time.Duration;
import java.util.List;
/**
* @author Ricken Bazolo
* @since 0.8.1
@@ -46,17 +48,7 @@ public class MistralAiEmbeddingClient extends AbstractEmbeddingClient {
private final MistralAiApi mistralAiApi;
private final RetryTemplate retryTemplate = RetryTemplate.builder()
.maxAttempts(10)
.retryOn(MistralAiApiException.class)
.exponentialBackoff(Duration.ofMillis(2000), 5, Duration.ofMillis(3 * 60000))
.withListener(new RetryListener() {
public <T extends Object, E extends Throwable> void onError(RetryContext context,
RetryCallback<T, E> callback, Throwable throwable) {
log.warn("Retry error. Retry count:" + context.getRetryCount(), throwable);
};
})
.build();
private final RetryTemplate retryTemplate;
public MistralAiEmbeddingClient(MistralAiApi mistralAiApi) {
this(mistralAiApi, MetadataMode.EMBED);
@@ -64,22 +56,25 @@ public class MistralAiEmbeddingClient extends AbstractEmbeddingClient {
public MistralAiEmbeddingClient(MistralAiApi mistralAiApi, MetadataMode metadataMode) {
this(mistralAiApi, metadataMode,
MistralAiEmbeddingOptions.builder().withModel(MistralAiApi.EmbeddingModel.EMBED.getValue()).build());
MistralAiEmbeddingOptions.builder().withModel(MistralAiApi.EmbeddingModel.EMBED.getValue()).build(),
RetryUtils.DEFAULT_RETRY_TEMPLATE);
}
public MistralAiEmbeddingClient(MistralAiApi mistralAiApi, MistralAiEmbeddingOptions options) {
this(mistralAiApi, MetadataMode.EMBED, options);
this(mistralAiApi, MetadataMode.EMBED, options, RetryUtils.DEFAULT_RETRY_TEMPLATE);
}
public MistralAiEmbeddingClient(MistralAiApi mistralAiApi, MetadataMode metadataMode,
MistralAiEmbeddingOptions options) {
MistralAiEmbeddingOptions options, RetryTemplate retryTemplate) {
Assert.notNull(mistralAiApi, "MistralAiApi must not be null");
Assert.notNull(metadataMode, "metadataMode must not be null");
Assert.notNull(options, "options must not be null");
Assert.notNull(retryTemplate, "retryTemplate must not be null");
this.mistralAiApi = mistralAiApi;
this.metadataMode = metadataMode;
this.defaultOptions = options;
this.retryTemplate = retryTemplate;
}
@Override

67
models/spring-ai-mistral-ai/src/main/java/org/springframework/ai/mistralai/api/MistralAiApi.java
View File

@@ -15,8 +15,6 @@
*/
package org.springframework.ai.mistralai.api;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.Map;
import java.util.function.Consumer;
@@ -25,22 +23,18 @@ import java.util.function.Predicate;
import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.annotation.JsonInclude.Include;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.databind.DeserializationFeature;
import com.fasterxml.jackson.databind.ObjectMapper;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import org.springframework.ai.model.ModelOptionsUtils;
import org.springframework.ai.retry.RetryUtils;
import org.springframework.boot.context.properties.bind.ConstructorBinding;
import org.springframework.core.ParameterizedTypeReference;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.http.client.ClientHttpResponse;
import org.springframework.lang.NonNull;
import org.springframework.util.Assert;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StreamUtils;
import org.springframework.web.client.ResponseErrorHandler;
import org.springframework.web.client.RestClient;
import org.springframework.web.reactive.function.client.WebClient;
@@ -70,8 +64,6 @@ public class MistralAiApi {
private WebClient webClient;
private final ObjectMapper objectMapper;
/**
* Create a new client api with DEFAULT_BASE_URL
* @param mistralAiApiKey Mistral api Key.
@@ -86,7 +78,7 @@ public class MistralAiApi {
* @param mistralAiApiKey Mistral api Key.
*/
public MistralAiApi(String baseUrl, String mistralAiApiKey) {
this(baseUrl, mistralAiApiKey, RestClient.builder());
this(baseUrl, mistralAiApiKey, RestClient.builder(), RetryUtils.DEFAULT_RESPONSE_ERROR_HANDLER);
}
/**
@@ -94,69 +86,24 @@ public class MistralAiApi {
* @param baseUrl api base URL.
* @param mistralAiApiKey Mistral api Key.
* @param restClientBuilder RestClient builder.
* @param responseErrorHandler Response error handler.
*/
public MistralAiApi(String baseUrl, String mistralAiApiKey, RestClient.Builder restClientBuilder) {
this.objectMapper = new ObjectMapper().configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false);
public MistralAiApi(String baseUrl, String mistralAiApiKey, RestClient.Builder restClientBuilder,
ResponseErrorHandler responseErrorHandler) {
Consumer<HttpHeaders> jsonContentHeaders = headers -> {
headers.setBearerAuth(mistralAiApiKey);
headers.setContentType(MediaType.APPLICATION_JSON);
};
var responseErrorHandler = new ResponseErrorHandler() {
@Override
public boolean hasError(@NonNull ClientHttpResponse response) throws IOException {
return response.getStatusCode().isError();
}
@Override
public void handleError(@NonNull ClientHttpResponse response) throws IOException {
if (response.getStatusCode().isError()) {
String error = StreamUtils.copyToString(response.getBody(), StandardCharsets.UTF_8);
String message = String.format("%s - %s", response.getStatusCode().value(), error);
if (response.getStatusCode().is4xxClientError()) {
throw new MistralAiApiClientErrorException(message);
}
throw new MistralAiApiException(message);
}
}
};
this.restClient = restClientBuilder.baseUrl(baseUrl)
.defaultHeaders(jsonContentHeaders)
.defaultStatusHandler(responseErrorHandler)
.build();
this.webClient = WebClient.builder().baseUrl(baseUrl).defaultHeaders(jsonContentHeaders).build();
}
public static class MistralAiApiException extends RuntimeException {
public MistralAiApiException(String message) {
super(message);
}
public MistralAiApiException(String message, Throwable t) {
super(message, t);
}
}
/**
* Thrown on 4xx client errors, such as 401 - Incorrect API key provided, 401 - You
* must be a member of an organization to use the API, 429 - Rate limit reached for
* requests, 429 - You exceeded your current quota , please check your plan and
* billing details.
*/
public static class MistralAiApiClientErrorException extends RuntimeException {
public MistralAiApiClientErrorException(String message) {
super(message);
}
}
/**
* Represents a tool the model may call. Currently, only functions are supported as a
* tool.
@@ -594,7 +541,7 @@ public class MistralAiApi {
// anticipation of future changes. Based on:
// https://github.com/mistralai/client-python/blob/main/src/mistralai/models/chat_completion.py
@JsonProperty("error") ERROR,
@JsonProperty("tool_calls") TOOL_CALLS
// @formatter:on

2
models/spring-ai-mistral-ai/src/test/java/org/springframework/ai/mistralai/MistralChatCompletionRequestTest.java → models/spring-ai-mistral-ai/src/test/java/org/springframework/ai/mistralai/MistralAiChatCompletionRequestTest.java
View File

@@ -30,7 +30,7 @@ import static org.assertj.core.api.Assertions.assertThat;
*/
@SpringBootTest
@EnabledIfEnvironmentVariable(named = "MISTRAL_AI_API_KEY", matches = ".+")
public class MistralChatCompletionRequestTest {
public class MistralAiChatCompletionRequestTest {
MistralAiChatClient chatClient = new MistralAiChatClient(new MistralAiApi("test"));

2
models/spring-ai-mistral-ai/src/test/java/org/springframework/ai/mistralai/MistralEmbeddingIT.java → models/spring-ai-mistral-ai/src/test/java/org/springframework/ai/mistralai/MistralAiEmbeddingIT.java
View File

@@ -28,7 +28,7 @@ import static org.assertj.core.api.Assertions.assertThat;
@SpringBootTest
@EnabledIfEnvironmentVariable(named = "MISTRAL_AI_API_KEY", matches = ".+")
class MistralEmbeddingIT {
class MistralAiEmbeddingIT {
@Autowired
private MistralAiEmbeddingClient mistralAiEmbeddingClient;

192
models/spring-ai-mistral-ai/src/test/java/org/springframework/ai/mistralai/MistralAiRetryTests.java Normal file
View File

@@ -0,0 +1,192 @@
/*
* Copyright 2023 - 2024 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.mistralai;
import java.util.List;
import java.util.Optional;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.mockito.Mock;
import org.mockito.junit.jupiter.MockitoExtension;
import reactor.core.publisher.Flux;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.document.MetadataMode;
import org.springframework.ai.mistralai.api.MistralAiApi;
import org.springframework.ai.mistralai.api.MistralAiApi.ChatCompletion;
import org.springframework.ai.mistralai.api.MistralAiApi.ChatCompletionChunk;
import org.springframework.ai.mistralai.api.MistralAiApi.ChatCompletionFinishReason;
import org.springframework.ai.mistralai.api.MistralAiApi.ChatCompletionMessage;
import org.springframework.ai.mistralai.api.MistralAiApi.ChatCompletionMessage.Role;
import org.springframework.ai.mistralai.api.MistralAiApi.ChatCompletionRequest;
import org.springframework.ai.mistralai.api.MistralAiApi.Embedding;
import org.springframework.ai.mistralai.api.MistralAiApi.EmbeddingList;
import org.springframework.ai.mistralai.api.MistralAiApi.EmbeddingRequest;
import org.springframework.ai.retry.RetryUtils;
import org.springframework.ai.retry.TransientAiException;
import org.springframework.http.ResponseEntity;
import org.springframework.retry.RetryCallback;
import org.springframework.retry.RetryContext;
import org.springframework.retry.RetryListener;
import org.springframework.retry.support.RetryTemplate;
import static org.assertj.core.api.Assertions.assertThat;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.mockito.ArgumentMatchers.isA;
import static org.mockito.Mockito.when;
/**
* @author Christian Tzolov
*/
@SuppressWarnings("unchecked")
@ExtendWith(MockitoExtension.class)
public class MistralAiRetryTests {
private class TestRetryListener implements RetryListener {
int onErrorRetryCount = 0;
int onSuccessRetryCount = 0;
@Override
public <T, E extends Throwable> void onSuccess(RetryContext context, RetryCallback<T, E> callback, T result) {
onSuccessRetryCount = context.getRetryCount();
}
@Override
public <T, E extends Throwable> void onError(RetryContext context, RetryCallback<T, E> callback,
Throwable throwable) {
onErrorRetryCount = context.getRetryCount();
}
}
private TestRetryListener retryListener;
private RetryTemplate retryTemplate;
private @Mock MistralAiApi mistralAiApi;
private MistralAiChatClient chatClient;
private MistralAiEmbeddingClient embeddingClient;
@BeforeEach
public void beforeEach() {
retryTemplate = RetryUtils.DEFAULT_RETRY_TEMPLATE;
retryListener = new TestRetryListener();
retryTemplate.registerListener(retryListener);
chatClient = new MistralAiChatClient(mistralAiApi,
MistralAiChatOptions.builder()
.withTemperature(0.7f)
.withTopP(1f)
.withSafePrompt(false)
.withModel(MistralAiApi.ChatModel.TINY.getValue())
.build(),
null, retryTemplate);
embeddingClient = new MistralAiEmbeddingClient(mistralAiApi, MetadataMode.EMBED,
MistralAiEmbeddingOptions.builder().withModel(MistralAiApi.EmbeddingModel.EMBED.getValue()).build(),
retryTemplate);
}
@Test
public void mistralAiChatTransientError() {
var choice = new ChatCompletion.Choice(0, new ChatCompletionMessage("Response", Role.ASSISTANT),
ChatCompletionFinishReason.STOP);
ChatCompletion expectedChatCompletion = new ChatCompletion("id", "chat.completion", 789l, "model",
List.of(choice), new MistralAiApi.Usage(10, 10, 10));
when(mistralAiApi.chatCompletionEntity(isA(ChatCompletionRequest.class)))
.thenThrow(new TransientAiException("Transient Error 1"))
.thenThrow(new TransientAiException("Transient Error 2"))
.thenReturn(ResponseEntity.of(Optional.of(expectedChatCompletion)));
var result = chatClient.call(new Prompt("text"));
assertThat(result).isNotNull();
assertThat(result.getResult().getOutput().getContent()).isSameAs("Response");
assertThat(retryListener.onSuccessRetryCount).isEqualTo(2);
assertThat(retryListener.onErrorRetryCount).isEqualTo(2);
}
@Test
public void mistralAiChatNonTransientError() {
when(mistralAiApi.chatCompletionEntity(isA(ChatCompletionRequest.class)))
.thenThrow(new RuntimeException("Non Transient Error"));
assertThrows(RuntimeException.class, () -> chatClient.call(new Prompt("text")));
}
@Test
public void mistralAiChatStreamTransientError() {
var choice = new ChatCompletionChunk.ChunkChoice(0, new ChatCompletionMessage("Response", Role.ASSISTANT),
ChatCompletionFinishReason.STOP);
ChatCompletionChunk expectedChatCompletion = new ChatCompletionChunk("id", "chat.completion.chunk", 789l,
"model", List.of(choice));
when(mistralAiApi.chatCompletionStream(isA(ChatCompletionRequest.class)))
.thenThrow(new TransientAiException("Transient Error 1"))
.thenThrow(new TransientAiException("Transient Error 2"))
.thenReturn(Flux.just(expectedChatCompletion));
var result = chatClient.stream(new Prompt("text"));
assertThat(result).isNotNull();
assertThat(result.collectList().block().get(0).getResult().getOutput().getContent()).isSameAs("Response");
assertThat(retryListener.onSuccessRetryCount).isEqualTo(2);
assertThat(retryListener.onErrorRetryCount).isEqualTo(2);
}
@Test
public void mistralAiChatStreamNonTransientError() {
when(mistralAiApi.chatCompletionStream(isA(ChatCompletionRequest.class)))
.thenThrow(new RuntimeException("Non Transient Error"));
assertThrows(RuntimeException.class, () -> chatClient.stream(new Prompt("text")));
}
@Test
public void mistralAiEmbeddingTransientError() {
EmbeddingList<Embedding> expectedEmbeddings = new EmbeddingList<>("list",
List.of(new Embedding(0, List.of(9.9, 8.8))), "model", new MistralAiApi.Usage(10, 10, 10));
when(mistralAiApi.embeddings(isA(EmbeddingRequest.class)))
.thenThrow(new TransientAiException("Transient Error 1"))
.thenThrow(new TransientAiException("Transient Error 2"))
.thenReturn(ResponseEntity.of(Optional.of(expectedEmbeddings)));
var result = embeddingClient
.call(new org.springframework.ai.embedding.EmbeddingRequest(List.of("text1", "text2"), null));
assertThat(result).isNotNull();
assertThat(result.getResult().getOutput()).isEqualTo(List.of(9.9, 8.8));
assertThat(retryListener.onSuccessRetryCount).isEqualTo(2);
assertThat(retryListener.onErrorRetryCount).isEqualTo(2);
}
@Test
public void mistralAiEmbeddingNonTransientError() {
when(mistralAiApi.embeddings(isA(EmbeddingRequest.class)))
.thenThrow(new RuntimeException("Non Transient Error"));
assertThrows(RuntimeException.class, () -> embeddingClient
.call(new org.springframework.ai.embedding.EmbeddingRequest(List.of("text1", "text2"), null)));
}
}

11
models/spring-ai-openai/pom.xml
View File

@@ -29,9 +29,9 @@
</dependency>
<dependency>
<groupId>org.springframework.retry</groupId>
<artifactId>spring-retry</artifactId>
<version>2.0.4</version>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-retry</artifactId>
<version>${project.parent.version}</version>
</dependency>
<!-- NOTE: Required only by the @ConstructorBinding. -->
@@ -57,11 +57,6 @@
<version>${victools.version}</version>
</dependency>
<dependency>
<groupId>org.springframework</groupId>
<artifactId>spring-webflux</artifactId>
</dependency>
<dependency>
<groupId>org.springframework</groupId>
<artifactId>spring-context-support</artifactId>

18
models/spring-ai-openai/src/main/java/org/springframework/ai/openai/OpenAiAudioTranscriptionClient.java
View File

@@ -31,8 +31,6 @@
package org.springframework.ai.openai;
import java.time.Duration;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -40,12 +38,12 @@ import org.springframework.ai.chat.metadata.RateLimit;
import org.springframework.ai.model.ModelClient;
import org.springframework.ai.openai.api.OpenAiAudioApi;
import org.springframework.ai.openai.api.OpenAiAudioApi.StructuredResponse;
import org.springframework.ai.openai.api.common.OpenAiApiException;
import org.springframework.ai.openai.audio.transcription.AudioTranscription;
import org.springframework.ai.openai.audio.transcription.AudioTranscriptionPrompt;
import org.springframework.ai.openai.audio.transcription.AudioTranscriptionResponse;
import org.springframework.ai.openai.metadata.audio.OpenAiAudioTranscriptionResponseMetadata;
import org.springframework.ai.openai.metadata.support.OpenAiResponseHeaderExtractor;
import org.springframework.ai.retry.RetryUtils;
import org.springframework.core.io.Resource;
import org.springframework.http.ResponseEntity;
import org.springframework.retry.support.RetryTemplate;
@@ -66,11 +64,7 @@ public class OpenAiAudioTranscriptionClient
private final OpenAiAudioTranscriptionOptions defaultOptions;
public final RetryTemplate retryTemplate = RetryTemplate.builder()
.maxAttempts(10)
.retryOn(OpenAiApiException.class)
.exponentialBackoff(Duration.ofMillis(2000), 5, Duration.ofMillis(3 * 60000))
.build();
public final RetryTemplate retryTemplate;
private final OpenAiAudioApi audioApi;
@@ -80,14 +74,18 @@ public class OpenAiAudioTranscriptionClient
.withModel(OpenAiAudioApi.WhisperModel.WHISPER_1.getValue())
.withResponseFormat(OpenAiAudioApi.TranscriptResponseFormat.JSON)
.withTemperature(0.7f)
.build());
.build(),
RetryUtils.DEFAULT_RETRY_TEMPLATE);
}
public OpenAiAudioTranscriptionClient(OpenAiAudioApi audioApi, OpenAiAudioTranscriptionOptions options) {
public OpenAiAudioTranscriptionClient(OpenAiAudioApi audioApi, OpenAiAudioTranscriptionOptions options,
RetryTemplate retryTemplate) {
Assert.notNull(audioApi, "OpenAiAudioApi must not be null");
Assert.notNull(options, "OpenAiTranscriptionOptions must not be null");
Assert.notNull(retryTemplate, "RetryTemplate must not be null");
this.audioApi = audioApi;
this.defaultOptions = options;
this.retryTemplate = retryTemplate;
}
public String call(Resource audioResource) {

51
models/spring-ai-openai/src/main/java/org/springframework/ai/openai/OpenAiChatClient.java
View File

@@ -15,7 +15,6 @@
*/
package org.springframework.ai.openai;
import java.time.Duration;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
@@ -43,14 +42,11 @@ import org.springframework.ai.openai.api.OpenAiApi.ChatCompletion;
import org.springframework.ai.openai.api.OpenAiApi.ChatCompletionMessage;
import org.springframework.ai.openai.api.OpenAiApi.ChatCompletionMessage.Role;
import org.springframework.ai.openai.api.OpenAiApi.ChatCompletionMessage.ToolCall;
import org.springframework.ai.openai.api.common.OpenAiApiException;
import org.springframework.ai.openai.api.OpenAiApi.ChatCompletionRequest;
import org.springframework.ai.openai.metadata.OpenAiChatResponseMetadata;
import org.springframework.ai.openai.metadata.support.OpenAiResponseHeaderExtractor;
import org.springframework.ai.retry.RetryUtils;
import org.springframework.http.ResponseEntity;
import org.springframework.retry.RetryCallback;
import org.springframework.retry.RetryContext;
import org.springframework.retry.RetryListener;
import org.springframework.retry.support.RetryTemplate;
import org.springframework.util.Assert;
import org.springframework.util.CollectionUtils;
@@ -73,7 +69,7 @@ public class OpenAiChatClient extends
AbstractFunctionCallSupport<ChatCompletionMessage, OpenAiApi.ChatCompletionRequest, ResponseEntity<ChatCompletion>>
implements ChatClient, StreamingChatClient {
private final Logger logger = LoggerFactory.getLogger(getClass());
private static final Logger logger = LoggerFactory.getLogger(OpenAiChatClient.class);
/**
* The default options used for the chat completion requests.
@@ -83,48 +79,59 @@ public class OpenAiChatClient extends
/**
* The retry template used to retry the OpenAI API calls.
*/
public final RetryTemplate retryTemplate = RetryTemplate.builder()
.maxAttempts(10)
.retryOn(OpenAiApiException.class)
.exponentialBackoff(Duration.ofMillis(2000), 5, Duration.ofMillis(3 * 60000))
.withListener(new RetryListener() {
@Override
public <T extends Object, E extends Throwable> void onError(RetryContext context,
RetryCallback<T, E> callback, Throwable throwable) {
logger.warn("Retry error. Retry count:" + context.getRetryCount(), throwable);
};
})
.build();
public final RetryTemplate retryTemplate;
/**
* Low-level access to the OpenAI API.
*/
private final OpenAiApi openAiApi;
/**
* Creates an instance of the OpenAiChatClient.
* @param openAiApi The OpenAiApi instance to be used for interacting with the OpenAI
* Chat API.
* @throws IllegalArgumentException if openAiApi is null
*/
public OpenAiChatClient(OpenAiApi openAiApi) {
this(openAiApi,
OpenAiChatOptions.builder().withModel(OpenAiApi.DEFAULT_CHAT_MODEL).withTemperature(0.7f).build());
}
/**
* Initializes an instance of the OpenAiChatClient.
* @param openAiApi The OpenAiApi instance to be used for interacting with the OpenAI
* Chat API.
* @param options The OpenAiChatOptions to configure the chat client.
*/
public OpenAiChatClient(OpenAiApi openAiApi, OpenAiChatOptions options) {
this(openAiApi, options, null);
this(openAiApi, options, null, RetryUtils.DEFAULT_RETRY_TEMPLATE);
}
/**
* Initializes a new instance of the OpenAiChatClient.
* @param openAiApi The OpenAiApi instance to be used for interacting with the OpenAI
* Chat API.
* @param options The OpenAiChatOptions to configure the chat client.
* @param functionCallbackContext The function callback context.
* @param retryTemplate The retry template.
*/
public OpenAiChatClient(OpenAiApi openAiApi, OpenAiChatOptions options,
FunctionCallbackContext functionCallbackContext) {
FunctionCallbackContext functionCallbackContext, RetryTemplate retryTemplate) {
super(functionCallbackContext);
Assert.notNull(openAiApi, "OpenAiApi must not be null");
Assert.notNull(options, "Options must not be null");
Assert.notNull(retryTemplate, "RetryTemplate must not be null");
this.openAiApi = openAiApi;
this.defaultOptions = options;
this.retryTemplate = retryTemplate;
}
@Override
public ChatResponse call(Prompt prompt) {
return this.retryTemplate.execute(ctx -> {
ChatCompletionRequest request = createRequest(prompt, false);
ChatCompletionRequest request = createRequest(prompt, false);
return this.retryTemplate.execute(ctx -> {
ResponseEntity<ChatCompletion> completionEntity = this.callWithFunctionSupport(request);

26
models/spring-ai-openai/src/main/java/org/springframework/ai/openai/OpenAiEmbeddingClient.java
View File

@@ -15,7 +15,6 @@
*/
package org.springframework.ai.openai;
import java.time.Duration;
import java.util.List;
import org.slf4j.Logger;
@@ -33,10 +32,7 @@ import org.springframework.ai.model.ModelOptionsUtils;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.ai.openai.api.OpenAiApi.EmbeddingList;
import org.springframework.ai.openai.api.OpenAiApi.Usage;
import org.springframework.ai.openai.api.common.OpenAiApiException;
import org.springframework.retry.RetryCallback;
import org.springframework.retry.RetryContext;
import org.springframework.retry.RetryListener;
import org.springframework.ai.retry.RetryUtils;
import org.springframework.retry.support.RetryTemplate;
import org.springframework.util.Assert;
@@ -51,17 +47,7 @@ public class OpenAiEmbeddingClient extends AbstractEmbeddingClient {
private final OpenAiEmbeddingOptions defaultOptions;
private final RetryTemplate retryTemplate = RetryTemplate.builder()
.maxAttempts(10)
.retryOn(OpenAiApiException.class)
.exponentialBackoff(Duration.ofMillis(2000), 5, Duration.ofMillis(3 * 60000))
.withListener(new RetryListener() {
public <T extends Object, E extends Throwable> void onError(RetryContext context,
RetryCallback<T, E> callback, Throwable throwable) {
logger.warn("Retry error. Retry count:" + context.getRetryCount(), throwable);
};
})
.build();
private final RetryTemplate retryTemplate;
private final OpenAiApi openAiApi;
@@ -73,17 +59,21 @@ public class OpenAiEmbeddingClient extends AbstractEmbeddingClient {
public OpenAiEmbeddingClient(OpenAiApi openAiApi, MetadataMode metadataMode) {
this(openAiApi, metadataMode,
OpenAiEmbeddingOptions.builder().withModel(OpenAiApi.DEFAULT_EMBEDDING_MODEL).build());
OpenAiEmbeddingOptions.builder().withModel(OpenAiApi.DEFAULT_EMBEDDING_MODEL).build(),
RetryUtils.DEFAULT_RETRY_TEMPLATE);
}
public OpenAiEmbeddingClient(OpenAiApi openAiApi, MetadataMode metadataMode, OpenAiEmbeddingOptions options) {
public OpenAiEmbeddingClient(OpenAiApi openAiApi, MetadataMode metadataMode, OpenAiEmbeddingOptions options,
RetryTemplate retryTemplate) {
Assert.notNull(openAiApi, "OpenAiService must not be null");
Assert.notNull(metadataMode, "metadataMode must not be null");
Assert.notNull(options, "options must not be null");
Assert.notNull(retryTemplate, "retryTemplate must not be null");
this.openAiApi = openAiApi;
this.metadataMode = metadataMode;
this.defaultOptions = options;
this.retryTemplate = retryTemplate;
}
@Override

34
models/spring-ai-openai/src/main/java/org/springframework/ai/openai/OpenAiImageClient.java
View File

@@ -15,7 +15,6 @@
*/
package org.springframework.ai.openai;
import java.time.Duration;
import java.util.List;
import org.slf4j.Logger;
@@ -30,13 +29,10 @@ import org.springframework.ai.image.ImageResponse;
import org.springframework.ai.image.ImageResponseMetadata;
import org.springframework.ai.model.ModelOptionsUtils;
import org.springframework.ai.openai.api.OpenAiImageApi;
import org.springframework.ai.openai.api.common.OpenAiApiException;
import org.springframework.ai.openai.metadata.OpenAiImageGenerationMetadata;
import org.springframework.ai.openai.metadata.OpenAiImageResponseMetadata;
import org.springframework.ai.retry.RetryUtils;
import org.springframework.http.ResponseEntity;
import org.springframework.retry.RetryCallback;
import org.springframework.retry.RetryContext;
import org.springframework.retry.RetryListener;
import org.springframework.retry.support.RetryTemplate;
import org.springframework.util.Assert;
@@ -50,38 +46,32 @@ import org.springframework.util.Assert;
*/
public class OpenAiImageClient implements ImageClient {
private final Logger logger = LoggerFactory.getLogger(getClass());
private final static Logger logger = LoggerFactory.getLogger(OpenAiImageClient.class);
private OpenAiImageOptions defaultOptions;
private final OpenAiImageApi openAiImageApi;
public final RetryTemplate retryTemplate = RetryTemplate.builder()
.maxAttempts(10)
.retryOn(OpenAiApiException.class)
.exponentialBackoff(Duration.ofMillis(2000), 5, Duration.ofMillis(3 * 60000))
.withListener(new RetryListener() {
public <T extends Object, E extends Throwable> void onError(RetryContext context,
RetryCallback<T, E> callback, Throwable throwable) {
logger.warn("Retry error. Retry count:" + context.getRetryCount(), throwable);
};
})
.build();
public final RetryTemplate retryTemplate;
public OpenAiImageClient(OpenAiImageApi openAiImageApi) {
this(openAiImageApi, OpenAiImageOptions.builder().build(), RetryUtils.DEFAULT_RETRY_TEMPLATE);
}
public OpenAiImageClient(OpenAiImageApi openAiImageApi, OpenAiImageOptions defaultOptions,
RetryTemplate retryTemplate) {
Assert.notNull(openAiImageApi, "OpenAiImageApi must not be null");
Assert.notNull(defaultOptions, "defaultOptions must not be null");
Assert.notNull(retryTemplate, "retryTemplate must not be null");
this.openAiImageApi = openAiImageApi;
this.defaultOptions = defaultOptions;
this.retryTemplate = retryTemplate;
}
public OpenAiImageOptions getDefaultOptions() {
return this.defaultOptions;
}
public OpenAiImageClient withDefaultOptions(OpenAiImageOptions defaultOptions) {
this.defaultOptions = defaultOptions;
return this;
}
@Override
public ImageResponse call(ImagePrompt imagePrompt) {
return this.retryTemplate.execute(ctx -> {

28
models/spring-ai-openai/src/main/java/org/springframework/ai/openai/api/common/ApiUtils.java → models/spring-ai-openai/src/main/java/org/springframework/ai/openai/api/ApiUtils.java
View File

@@ -13,18 +13,12 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.openai.api.common;
package org.springframework.ai.openai.api;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.function.Consumer;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.http.client.ClientHttpResponse;
import org.springframework.lang.NonNull;
import org.springframework.util.StreamUtils;
import org.springframework.web.client.ResponseErrorHandler;
/**
* @author Christian Tzolov
@@ -40,24 +34,4 @@ public class ApiUtils {
};
};
public static final ResponseErrorHandler DEFAULT_RESPONSE_ERROR_HANDLER = new ResponseErrorHandler() {
@Override
public boolean hasError(@NonNull ClientHttpResponse response) throws IOException {
return response.getStatusCode().isError();
}
@Override
public void handleError(@NonNull ClientHttpResponse response) throws IOException {
if (response.getStatusCode().isError()) {
String error = StreamUtils.copyToString(response.getBody(), StandardCharsets.UTF_8);
String message = String.format("%s - %s", response.getStatusCode().value(), error);
if (response.getStatusCode().is4xxClientError()) {
throw new OpenAiApiClientErrorException(message);
}
throw new OpenAiApiException(message);
}
}
};
}

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

@@ -26,14 +26,13 @@ import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import org.springframework.ai.model.ModelOptionsUtils;
import org.springframework.ai.openai.api.common.ApiUtils;
import org.springframework.ai.retry.RetryUtils;
import org.springframework.boot.context.properties.bind.ConstructorBinding;
import org.springframework.core.ParameterizedTypeReference;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.util.Assert;
import org.springframework.util.CollectionUtils;
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;
@@ -53,8 +52,6 @@ public class OpenAiApi {
private final RestClient restClient;
private final RestClient multipartRestClient;
private final WebClient webClient;
/**
@@ -84,20 +81,23 @@ public class OpenAiApi {
* @param restClientBuilder RestClient builder.
*/
public OpenAiApi(String baseUrl, String openAiToken, RestClient.Builder restClientBuilder) {
this(baseUrl, openAiToken, restClientBuilder, RetryUtils.DEFAULT_RESPONSE_ERROR_HANDLER);
}
/**
* Create a new chat completion api.
*
* @param baseUrl api base URL.
* @param openAiToken OpenAI apiKey.
* @param restClientBuilder RestClient builder.
* @param responseErrorHandler Response error handler.
*/
public OpenAiApi(String baseUrl, String openAiToken, RestClient.Builder restClientBuilder, ResponseErrorHandler responseErrorHandler) {
this.restClient = restClientBuilder
.baseUrl(baseUrl)
.defaultHeaders(ApiUtils.getJsonContentHeaders(openAiToken))
.defaultStatusHandler(ApiUtils.DEFAULT_RESPONSE_ERROR_HANDLER)
.build();
this.multipartRestClient = restClientBuilder
.baseUrl(baseUrl)
.defaultHeaders(multipartFormDataHeaders -> {
multipartFormDataHeaders.setBearerAuth(openAiToken);
multipartFormDataHeaders.setContentType(MediaType.MULTIPART_FORM_DATA);
})
.defaultStatusHandler(ApiUtils.DEFAULT_RESPONSE_ERROR_HANDLER)
.defaultStatusHandler(RetryUtils.DEFAULT_RESPONSE_ERROR_HANDLER)
.build();
this.webClient = WebClient.builder()

11
models/spring-ai-openai/src/main/java/org/springframework/ai/openai/api/OpenAiAudioApi.java
View File

@@ -21,12 +21,13 @@ import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.annotation.JsonInclude.Include;
import com.fasterxml.jackson.annotation.JsonProperty;
import org.springframework.ai.openai.api.common.ApiUtils;
import org.springframework.ai.retry.RetryUtils;
import org.springframework.core.io.ByteArrayResource;
import org.springframework.http.ResponseEntity;
import org.springframework.util.Assert;
import org.springframework.util.LinkedMultiValueMap;
import org.springframework.util.MultiValueMap;
import org.springframework.web.client.ResponseErrorHandler;
import org.springframework.web.client.RestClient;
/**
@@ -45,7 +46,7 @@ public class OpenAiAudioApi {
* @param openAiToken OpenAI apiKey.
*/
public OpenAiAudioApi(String openAiToken) {
this(ApiUtils.DEFAULT_BASE_URL, openAiToken, RestClient.builder());
this(ApiUtils.DEFAULT_BASE_URL, openAiToken, RestClient.builder(), RetryUtils.DEFAULT_RESPONSE_ERROR_HANDLER);
}
/**
@@ -53,12 +54,14 @@ public class OpenAiAudioApi {
* @param baseUrl api base URL.
* @param openAiToken OpenAI apiKey.
* @param restClientBuilder RestClient builder.
* @param responseErrorHandler Response error handler.
*/
public OpenAiAudioApi(String baseUrl, String openAiToken, RestClient.Builder restClientBuilder) {
public OpenAiAudioApi(String baseUrl, String openAiToken, RestClient.Builder restClientBuilder,
ResponseErrorHandler responseErrorHandler) {
this.restClient = restClientBuilder.baseUrl(baseUrl).defaultHeaders(headers -> {
headers.setBearerAuth(openAiToken);
}).defaultStatusHandler(ApiUtils.DEFAULT_RESPONSE_ERROR_HANDLER).build();
}).defaultStatusHandler(responseErrorHandler).build();
}
/**

23
models/spring-ai-openai/src/main/java/org/springframework/ai/openai/api/OpenAiImageApi.java
View File

@@ -20,9 +20,10 @@ import java.util.List;
import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.annotation.JsonProperty;
import org.springframework.ai.openai.api.common.ApiUtils;
import org.springframework.ai.retry.RetryUtils;
import org.springframework.http.ResponseEntity;
import org.springframework.util.Assert;
import org.springframework.web.client.ResponseErrorHandler;
import org.springframework.web.client.RestClient;
/**
@@ -44,11 +45,29 @@ public class OpenAiImageApi {
this(ApiUtils.DEFAULT_BASE_URL, openAiToken, RestClient.builder());
}
/**
* Create a new OpenAI Image API with the provided base URL.
* @param baseUrl the base URL for the OpenAI API.
* @param openAiToken OpenAI apiKey.
* @param restClientBuilder the rest client builder to use.
*/
public OpenAiImageApi(String baseUrl, String openAiToken, RestClient.Builder restClientBuilder) {
this(baseUrl, openAiToken, restClientBuilder, RetryUtils.DEFAULT_RESPONSE_ERROR_HANDLER);
}
/**
* Create a new OpenAI Image API with the provided base URL.
* @param baseUrl the base URL for the OpenAI API.
* @param openAiToken OpenAI apiKey.
* @param restClientBuilder the rest client builder to use.
* @param responseErrorHandler the response error handler to use.
*/
public OpenAiImageApi(String baseUrl, String openAiToken, RestClient.Builder restClientBuilder,
ResponseErrorHandler responseErrorHandler) {
this.restClient = restClientBuilder.baseUrl(baseUrl)
.defaultHeaders(ApiUtils.getJsonContentHeaders(openAiToken))
.defaultStatusHandler(ApiUtils.DEFAULT_RESPONSE_ERROR_HANDLER)
.defaultStatusHandler(responseErrorHandler)
.build();
}

2
models/spring-ai-openai/src/test/java/org/springframework/ai/openai/ChatCompletionRequestTests.java
View File

@@ -22,7 +22,7 @@ import org.junit.jupiter.api.Test;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.model.function.FunctionCallbackWrapper;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.ai.openai.chat.api.tool.MockWeatherService;
import org.springframework.ai.openai.api.tool.MockWeatherService;
import static org.assertj.core.api.Assertions.assertThat;

2
models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/api/OpenAiApiIT.java → models/spring-ai-openai/src/test/java/org/springframework/ai/openai/api/OpenAiApiIT.java
View File

@@ -13,7 +13,7 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.openai.chat.api;
package org.springframework.ai.openai.api;
import java.util.List;

3
models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/api/RestClientBuilderTests.java → models/spring-ai-openai/src/test/java/org/springframework/ai/openai/api/RestClientBuilderTests.java
View File

@@ -13,9 +13,8 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.openai.chat.api;
package org.springframework.ai.openai.api;
import org.junit.jupiter.api.Disabled;
import org.junit.jupiter.api.Test;
import org.springframework.http.client.SimpleClientHttpRequestFactory;

2
models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/api/tool/MockWeatherService.java → models/spring-ai-openai/src/test/java/org/springframework/ai/openai/api/tool/MockWeatherService.java
View File

@@ -13,7 +13,7 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.openai.chat.api.tool;
package org.springframework.ai.openai.api.tool;
import java.util.function.Function;

3
models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/api/tool/OpenAiApiToolFunctionCallIT.java → models/spring-ai-openai/src/test/java/org/springframework/ai/openai/api/tool/OpenAiApiToolFunctionCallIT.java
View File

@@ -13,7 +13,8 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.openai.chat.api.tool;
package org.springframework.ai.openai.api.tool;
import java.util.ArrayList;
import java.util.List;

3
models/spring-ai-openai/src/test/java/org/springframework/ai/openai/audio/transcription/OpenAiTranscriptionClientWithTranscriptionResponseMetadataTests.java
View File

@@ -26,6 +26,7 @@ import org.springframework.ai.openai.api.OpenAiAudioApi;
import org.springframework.ai.openai.metadata.audio.OpenAiAudioTranscriptionMetadata;
import org.springframework.ai.openai.metadata.audio.OpenAiAudioTranscriptionResponseMetadata;
import org.springframework.ai.openai.metadata.support.OpenAiApiResponseHeaders;
import org.springframework.ai.retry.RetryUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.boot.test.autoconfigure.web.client.RestClientTest;
@@ -151,7 +152,7 @@ public class OpenAiTranscriptionClientWithTranscriptionResponseMetadataTests {
@Bean
public OpenAiAudioApi chatCompletionApi(RestClient.Builder builder) {
return new OpenAiAudioApi("", TEST_API_KEY, builder);
return new OpenAiAudioApi("", TEST_API_KEY, builder, RetryUtils.DEFAULT_RESPONSE_ERROR_HANDLER);
}
@Bean

2
models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/OpenAiChatClientIT.java
View File

@@ -37,7 +37,7 @@ import org.springframework.ai.chat.prompt.SystemPromptTemplate;
import org.springframework.ai.model.function.FunctionCallbackWrapper;
import org.springframework.ai.openai.OpenAiChatOptions;
import org.springframework.ai.openai.OpenAiTestConfiguration;
import org.springframework.ai.openai.chat.api.tool.MockWeatherService;
import org.springframework.ai.openai.api.tool.MockWeatherService;
import org.springframework.ai.openai.testutils.AbstractIT;
import org.springframework.ai.parser.BeanOutputParser;
import org.springframework.ai.parser.ListOutputParser;

272
models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/OpenAiRetryTests.java Normal file
View File

@@ -0,0 +1,272 @@
/*
* Copyright 2023 - 2024 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.openai.chat;
import java.util.List;
import java.util.Optional;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.mockito.Mock;
import org.mockito.junit.jupiter.MockitoExtension;
import reactor.core.publisher.Flux;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.document.MetadataMode;
import org.springframework.ai.image.ImageMessage;
import org.springframework.ai.image.ImagePrompt;
import org.springframework.ai.openai.OpenAiAudioTranscriptionClient;
import org.springframework.ai.openai.OpenAiAudioTranscriptionOptions;
import org.springframework.ai.openai.OpenAiChatClient;
import org.springframework.ai.openai.OpenAiChatOptions;
import org.springframework.ai.openai.OpenAiEmbeddingClient;
import org.springframework.ai.openai.OpenAiEmbeddingOptions;
import org.springframework.ai.openai.OpenAiImageClient;
import org.springframework.ai.openai.OpenAiImageOptions;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.ai.openai.api.OpenAiApi.ChatCompletion;
import org.springframework.ai.openai.api.OpenAiApi.ChatCompletionChunk;
import org.springframework.ai.openai.api.OpenAiApi.ChatCompletionFinishReason;
import org.springframework.ai.openai.api.OpenAiApi.ChatCompletionMessage;
import org.springframework.ai.openai.api.OpenAiApi.ChatCompletionMessage.Role;
import org.springframework.ai.openai.api.OpenAiApi.ChatCompletionRequest;
import org.springframework.ai.openai.api.OpenAiApi.Embedding;
import org.springframework.ai.openai.api.OpenAiApi.EmbeddingList;
import org.springframework.ai.openai.api.OpenAiApi.EmbeddingRequest;
import org.springframework.ai.openai.api.OpenAiAudioApi;
import org.springframework.ai.openai.api.OpenAiAudioApi.StructuredResponse;
import org.springframework.ai.openai.api.OpenAiAudioApi.TranscriptResponseFormat;
import org.springframework.ai.openai.api.OpenAiAudioApi.TranscriptionRequest;
import org.springframework.ai.openai.api.OpenAiImageApi;
import org.springframework.ai.openai.api.OpenAiImageApi.Data;
import org.springframework.ai.openai.api.OpenAiImageApi.OpenAiImageRequest;
import org.springframework.ai.openai.api.OpenAiImageApi.OpenAiImageResponse;
import org.springframework.ai.openai.audio.transcription.AudioTranscriptionPrompt;
import org.springframework.ai.openai.audio.transcription.AudioTranscriptionResponse;
import org.springframework.ai.retry.RetryUtils;
import org.springframework.ai.retry.TransientAiException;
import org.springframework.core.io.ClassPathResource;
import org.springframework.http.ResponseEntity;
import org.springframework.retry.RetryCallback;
import org.springframework.retry.RetryContext;
import org.springframework.retry.RetryListener;
import org.springframework.retry.support.RetryTemplate;
import static org.assertj.core.api.Assertions.assertThat;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.mockito.ArgumentMatchers.isA;
import static org.mockito.Mockito.when;
/**
* @author Christian Tzolov
*/
@SuppressWarnings("unchecked")
@ExtendWith(MockitoExtension.class)
public class OpenAiRetryTests {
private class TestRetryListener implements RetryListener {
int onErrorRetryCount = 0;
int onSuccessRetryCount = 0;
@Override
public <T, E extends Throwable> void onSuccess(RetryContext context, RetryCallback<T, E> callback, T result) {
onSuccessRetryCount = context.getRetryCount();
}
@Override
public <T, E extends Throwable> void onError(RetryContext context, RetryCallback<T, E> callback,
Throwable throwable) {
onErrorRetryCount = context.getRetryCount();
}
}
private TestRetryListener retryListener;
private RetryTemplate retryTemplate;
private @Mock OpenAiApi openAiApi;
private @Mock OpenAiAudioApi openAiAudioApi;
private @Mock OpenAiImageApi openAiImageApi;
private OpenAiChatClient chatClient;
private OpenAiEmbeddingClient embeddingClient;
private OpenAiAudioTranscriptionClient audioTranscriptionClient;
private OpenAiImageClient imageClient;
@BeforeEach
public void beforeEach() {
retryTemplate = RetryUtils.DEFAULT_RETRY_TEMPLATE;
retryListener = new TestRetryListener();
retryTemplate.registerListener(retryListener);
chatClient = new OpenAiChatClient(openAiApi, OpenAiChatOptions.builder().build(), null, retryTemplate);
embeddingClient = new OpenAiEmbeddingClient(openAiApi, MetadataMode.EMBED,
OpenAiEmbeddingOptions.builder().build(), retryTemplate);
audioTranscriptionClient = new OpenAiAudioTranscriptionClient(openAiAudioApi,
OpenAiAudioTranscriptionOptions.builder()
.withModel("model")
.withResponseFormat(TranscriptResponseFormat.JSON)
.build(),
retryTemplate);
imageClient = new OpenAiImageClient(openAiImageApi, OpenAiImageOptions.builder().build(), retryTemplate);
}
@Test
public void openAiChatTransientError() {
var choice = new ChatCompletion.Choice(ChatCompletionFinishReason.STOP, 0,
new ChatCompletionMessage("Response", Role.ASSISTANT), null);
ChatCompletion expectedChatCompletion = new ChatCompletion("id", List.of(choice), 666l, "model", null, null,
new OpenAiApi.Usage(10, 10, 10));
when(openAiApi.chatCompletionEntity(isA(ChatCompletionRequest.class)))
.thenThrow(new TransientAiException("Transient Error 1"))
.thenThrow(new TransientAiException("Transient Error 2"))
.thenReturn(ResponseEntity.of(Optional.of(expectedChatCompletion)));
var result = chatClient.call(new Prompt("text"));
assertThat(result).isNotNull();
assertThat(result.getResult().getOutput().getContent()).isSameAs("Response");
assertThat(retryListener.onSuccessRetryCount).isEqualTo(2);
assertThat(retryListener.onErrorRetryCount).isEqualTo(2);
}
@Test
public void openAiChatNonTransientError() {
when(openAiApi.chatCompletionEntity(isA(ChatCompletionRequest.class)))
.thenThrow(new RuntimeException("Non Transient Error"));
assertThrows(RuntimeException.class, () -> chatClient.call(new Prompt("text")));
}
@Test
public void openAiChatStreamTransientError() {
var choice = new ChatCompletionChunk.ChunkChoice(ChatCompletionFinishReason.STOP, 0,
new ChatCompletionMessage("Response", Role.ASSISTANT), null);
ChatCompletionChunk expectedChatCompletion = new ChatCompletionChunk("id", List.of(choice), 666l, "model", null,
null);
when(openAiApi.chatCompletionStream(isA(ChatCompletionRequest.class)))
.thenThrow(new TransientAiException("Transient Error 1"))
.thenThrow(new TransientAiException("Transient Error 2"))
.thenReturn(Flux.just(expectedChatCompletion));
var result = chatClient.stream(new Prompt("text"));
assertThat(result).isNotNull();
assertThat(result.collectList().block().get(0).getResult().getOutput().getContent()).isSameAs("Response");
assertThat(retryListener.onSuccessRetryCount).isEqualTo(2);
assertThat(retryListener.onErrorRetryCount).isEqualTo(2);
}
@Test
public void openAiChatStreamNonTransientError() {
when(openAiApi.chatCompletionStream(isA(ChatCompletionRequest.class)))
.thenThrow(new RuntimeException("Non Transient Error"));
assertThrows(RuntimeException.class, () -> chatClient.stream(new Prompt("text")));
}
@Test
public void openAiEmbeddingTransientError() {
EmbeddingList<Embedding> expectedEmbeddings = new EmbeddingList<>("list",
List.of(new Embedding(0, List.of(9.9, 8.8))), "model", new OpenAiApi.Usage(10, 10, 10));
when(openAiApi.embeddings(isA(EmbeddingRequest.class))).thenThrow(new TransientAiException("Transient Error 1"))
.thenThrow(new TransientAiException("Transient Error 2"))
.thenReturn(ResponseEntity.of(Optional.of(expectedEmbeddings)));
var result = embeddingClient
.call(new org.springframework.ai.embedding.EmbeddingRequest(List.of("text1", "text2"), null));
assertThat(result).isNotNull();
assertThat(result.getResult().getOutput()).isEqualTo(List.of(9.9, 8.8));
assertThat(retryListener.onSuccessRetryCount).isEqualTo(2);
assertThat(retryListener.onErrorRetryCount).isEqualTo(2);
}
@Test
public void openAiEmbeddingNonTransientError() {
when(openAiApi.embeddings(isA(EmbeddingRequest.class)))
.thenThrow(new RuntimeException("Non Transient Error"));
assertThrows(RuntimeException.class, () -> embeddingClient
.call(new org.springframework.ai.embedding.EmbeddingRequest(List.of("text1", "text2"), null)));
}
@Test
public void openAiAudioTranscriptionTransientError() {
var expectedResponse = new StructuredResponse("nl", 6.7f, "Transcription Text", List.of(), List.of());
when(openAiAudioApi.createTranscription(isA(TranscriptionRequest.class), isA(Class.class)))
.thenThrow(new TransientAiException("Transient Error 1"))
.thenThrow(new TransientAiException("Transient Error 2"))
.thenReturn(ResponseEntity.of(Optional.of(expectedResponse)));
AudioTranscriptionResponse result = audioTranscriptionClient
.call(new AudioTranscriptionPrompt(new ClassPathResource("speech/jfk.flac")));
assertThat(result).isNotNull();
assertThat(result.getResult().getOutput()).isEqualTo(expectedResponse.text());
assertThat(retryListener.onSuccessRetryCount).isEqualTo(2);
assertThat(retryListener.onErrorRetryCount).isEqualTo(2);
}
@Test
public void openAiAudioTranscriptionNonTransientError() {
when(openAiAudioApi.createTranscription(isA(TranscriptionRequest.class), isA(Class.class)))
.thenThrow(new RuntimeException("Transient Error 1"));
assertThrows(RuntimeException.class, () -> audioTranscriptionClient
.call(new AudioTranscriptionPrompt(new ClassPathResource("speech/jfk.flac"))));
}
@Test
public void openAiImageTransientError() {
var expectedResponse = new OpenAiImageResponse(678l, List.of(new Data("url678", "b64", "prompt")));
when(openAiImageApi.createImage(isA(OpenAiImageRequest.class)))
.thenThrow(new TransientAiException("Transient Error 1"))
.thenThrow(new TransientAiException("Transient Error 2"))
.thenReturn(ResponseEntity.of(Optional.of(expectedResponse)));
var result = imageClient.call(new ImagePrompt(List.of(new ImageMessage("Image Message"))));
assertThat(result).isNotNull();
assertThat(result.getResult().getOutput().getUrl()).isEqualTo("url678");
assertThat(retryListener.onSuccessRetryCount).isEqualTo(2);
assertThat(retryListener.onErrorRetryCount).isEqualTo(2);
}
@Test
public void openAiImageNonTransientError() {
when(openAiImageApi.createImage(isA(OpenAiImageRequest.class)))
.thenThrow(new RuntimeException("Transient Error 1"));
assertThrows(RuntimeException.class,
() -> imageClient.call(new ImagePrompt(List.of(new ImageMessage("Image Message")))));
}
}

7
models/spring-ai-stability-ai/pom.xml
View File

@@ -30,12 +30,11 @@
</dependency>
<dependency>
<groupId>org.springframework</groupId>
<artifactId>spring-web</artifactId>
<version>${spring-framework.version}</version>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-retry</artifactId>
<version>${project.parent.version}</version>
</dependency>
<!-- Spring Framework -->
<dependency>
<groupId>org.springframework</groupId>

30
models/spring-ai-stability-ai/src/main/java/org/springframework/ai/stabilityai/StabilityAiImageClient.java
View File

@@ -15,17 +15,24 @@
*/
package org.springframework.ai.stabilityai;
import java.util.List;
import java.util.stream.Collectors;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.image.*;
import org.springframework.ai.image.Image;
import org.springframework.ai.image.ImageClient;
import org.springframework.ai.image.ImageGeneration;
import org.springframework.ai.image.ImageOptions;
import org.springframework.ai.image.ImagePrompt;
import org.springframework.ai.image.ImageResponse;
import org.springframework.ai.image.ImageResponseMetadata;
import org.springframework.ai.model.ModelOptionsUtils;
import org.springframework.ai.stabilityai.api.StabilityAiApi;
import org.springframework.ai.stabilityai.api.StabilityAiImageOptions;
import org.springframework.util.Assert;
import java.util.List;
import java.util.stream.Collectors;
/**
* StabilityAiImageClient is a class that implements the ImageClient interface. It
* provides a client for calling the StabilityAI image generation API.
@@ -50,7 +57,7 @@ public class StabilityAiImageClient implements ImageClient {
}
public StabilityAiImageOptions getOptions() {
return options;
return this.options;
}
/**
@@ -159,17 +166,4 @@ public class StabilityAiImageClient implements ImageClient {
return builder.build();
}
private ImagePrompt createUpdatedPrompt(ImagePrompt prompt) {
ImageOptions runtimeImageModelOptions = prompt.getOptions();
ImageOptionsBuilder imageOptionsBuilder = ImageOptionsBuilder.builder();
if (runtimeImageModelOptions != null) {
if (runtimeImageModelOptions.getModel() != null) {
imageOptionsBuilder.withModel(runtimeImageModelOptions.getModel());
}
}
ImageOptions updatedImageModelOptions = imageOptionsBuilder.build();
return new ImagePrompt(prompt.getInstructions(), updatedImageModelOptions);
}
}

11
models/spring-ai-stability-ai/src/main/java/org/springframework/ai/stabilityai/StabilityAiImageGenerationMetadata.java
View File

@@ -35,16 +35,17 @@ public class StabilityAiImageGenerationMetadata implements ImageGenerationMetada
}
public String getFinishReason() {
return finishReason;
return this.finishReason;
}
public Long getSeed() {
return seed;
return this.seed;
}
@Override
public String toString() {
return "StabilityAiImageGenerationMetadata{" + "finishReason='" + finishReason + '\'' + ", seed=" + seed + '}';
return "StabilityAiImageGenerationMetadata{" + "finishReason='" + this.finishReason + '\'' + ", seed="
+ this.seed + '}';
}
@Override
@@ -53,12 +54,12 @@ public class StabilityAiImageGenerationMetadata implements ImageGenerationMetada
return true;
if (!(o instanceof StabilityAiImageGenerationMetadata that))
return false;
return Objects.equals(finishReason, that.finishReason) && Objects.equals(seed, that.seed);
return Objects.equals(this.finishReason, that.finishReason) && Objects.equals(this.seed, that.seed);
}
@Override
public int hashCode() {
return Objects.hash(finishReason, seed);
return Objects.hash(this.finishReason, this.seed);
}
}

22
models/spring-ai-stability-ai/src/main/java/org/springframework/ai/stabilityai/StyleEnum.java
View File

@@ -20,11 +20,25 @@ package org.springframework.ai.stabilityai;
*/
public enum StyleEnum {
THREE_D_MODEL("3d-model"), ANALOG_FILM("analog-film"), ANIME("anime"), CINEMATIC("cinematic"),
COMIC_BOOK("comic-book"), DIGITAL_ART("digital-art"), ENHANCE("enhance"), FANTASY_ART("fantasy-art"),
ISOMETRIC("isometric"), LINE_ART("line-art"), LOW_POLY("low-poly"), MODELING_COMPOUND("modeling-compound"),
NEON_PUNK("neon-punk"), ORIGAMI("origami"), PHOTOGRAPHIC("photographic"), PIXEL_ART("pixel-art"),
// @formatter:off
THREE_D_MODEL("3d-model"),
ANALOG_FILM("analog-film"),
ANIME("anime"),
CINEMATIC("cinematic"),
COMIC_BOOK("comic-book"),
DIGITAL_ART("digital-art"),
ENHANCE("enhance"),
FANTASY_ART("fantasy-art"),
ISOMETRIC("isometric"),
LINE_ART("line-art"),
LOW_POLY("low-poly"),
MODELING_COMPOUND("modeling-compound"),
NEON_PUNK("neon-punk"),
ORIGAMI("origami"),
PHOTOGRAPHIC("photographic"),
PIXEL_ART("pixel-art"),
TILE_TEXTURE("tile-texture");
// @formatter:on
private final String text;

45
models/spring-ai-stability-ai/src/main/java/org/springframework/ai/stabilityai/api/StabilityAiApi.java
View File

@@ -15,20 +15,18 @@
*/
package org.springframework.ai.stabilityai.api;
import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.http.client.ClientHttpResponse;
import org.springframework.util.Assert;
import org.springframework.web.client.ResponseErrorHandler;
import org.springframework.web.client.RestClient;
import java.io.IOException;
import java.util.List;
import java.util.function.Consumer;
import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.annotation.JsonProperty;
import org.springframework.ai.retry.RetryUtils;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.util.Assert;
import org.springframework.web.client.RestClient;
/**
* Represents the StabilityAI API.
*/
@@ -80,35 +78,12 @@ public class StabilityAiApi {
headers.setContentType(MediaType.APPLICATION_JSON);
};
ResponseErrorHandler responseErrorHandler = new ResponseErrorHandler() {
@Override
public boolean hasError(ClientHttpResponse response) throws IOException {
return response.getStatusCode().isError();
}
@Override
public void handleError(ClientHttpResponse response) throws IOException {
if (response.getStatusCode().isError()) {
throw new RuntimeException(String.format("%s - %s", response.getStatusCode().value(),
new ObjectMapper().readValue(response.getBody(), ResponseError.class)));
}
}
};
this.restClient = restClientBuilder.baseUrl(baseUrl)
.defaultHeaders(jsonContentHeaders)
.defaultStatusHandler(responseErrorHandler)
.defaultStatusHandler(RetryUtils.DEFAULT_RESPONSE_ERROR_HANDLER)
.build();
}
@JsonInclude(JsonInclude.Include.NON_NULL)
public record ResponseError(@JsonProperty("id") String id, @JsonProperty("name") String name,
@JsonProperty("message") String message
) {
}
@JsonInclude(JsonInclude.Include.NON_NULL)
public record GenerateImageRequest(@JsonProperty("text_prompts") List<TextPrompts> textPrompts,
@JsonProperty("height") Integer height, @JsonProperty("width") Integer width,