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

feat(ollama): add retry template integration to OllamaChatModel

Browse Source 
* Update tests that are supposed to fail to not use retry
* Upgrade ot use Ollama 0.6.7

Signed-off-by: Alexandros Pappas <alexandros.pappas@yiluhub.com>
This commit is contained in:
Alexandros Pappas
2024-12-02 16:10:36 +01:00
committed by Mark Pollack
parent 6ef15bdd74
commit cfbefee6a2
14 changed files with 340 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

29
models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/OllamaChatModel.java
View File

@@ -65,8 +65,13 @@ import org.springframework.ai.ollama.api.common.OllamaApiConstants;
import org.springframework.ai.ollama.management.ModelManagementOptions;
import org.springframework.ai.ollama.management.OllamaModelManager;
import org.springframework.ai.ollama.management.PullModelStrategy;
import org.springframework.ai.tool.definition.ToolDefinition;
import org.springframework.ai.util.json.JsonParser;
import org.springframework.ai.retry.RetryUtils;
import org.springframework.retry.support.RetryTemplate;
import org.springframework.util.Assert;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;
@@ -129,27 +134,32 @@ public class OllamaChatModel implements ChatModel {
private ChatModelObservationConvention observationConvention = DEFAULT_OBSERVATION_CONVENTION;
private final RetryTemplate retryTemplate;
public OllamaChatModel(OllamaApi ollamaApi, OllamaOptions defaultOptions, ToolCallingManager toolCallingManager,
ObservationRegistry observationRegistry, ModelManagementOptions modelManagementOptions) {
this(ollamaApi, defaultOptions, toolCallingManager, observationRegistry, modelManagementOptions,
new DefaultToolExecutionEligibilityPredicate());
new DefaultToolExecutionEligibilityPredicate(), RetryUtils.DEFAULT_RETRY_TEMPLATE);
}
public OllamaChatModel(OllamaApi ollamaApi, OllamaOptions defaultOptions, ToolCallingManager toolCallingManager,
ObservationRegistry observationRegistry, ModelManagementOptions modelManagementOptions,
ToolExecutionEligibilityPredicate toolExecutionEligibilityPredicate) {
ToolExecutionEligibilityPredicate toolExecutionEligibilityPredicate, RetryTemplate retryTemplate) {
Assert.notNull(ollamaApi, "ollamaApi must not be null");
Assert.notNull(defaultOptions, "defaultOptions must not be null");
Assert.notNull(toolCallingManager, "toolCallingManager must not be null");
Assert.notNull(observationRegistry, "observationRegistry must not be null");
Assert.notNull(modelManagementOptions, "modelManagementOptions must not be null");
Assert.notNull(toolExecutionEligibilityPredicate, "toolExecutionEligibilityPredicate must not be null");
Assert.notNull(retryTemplate, "retryTemplate must not be null");
this.chatApi = ollamaApi;
this.defaultOptions = defaultOptions;
this.toolCallingManager = toolCallingManager;
this.observationRegistry = observationRegistry;
this.modelManager = new OllamaModelManager(this.chatApi, modelManagementOptions);
this.toolExecutionEligibilityPredicate = toolExecutionEligibilityPredicate;
this.retryTemplate = retryTemplate;
initializeModel(defaultOptions.getModel(), modelManagementOptions.pullModelStrategy());
}
@@ -237,7 +247,7 @@ public class OllamaChatModel implements ChatModel {
this.observationRegistry)
.observe(() -> {
OllamaApi.ChatResponse ollamaResponse = this.chatApi.chat(request);
OllamaApi.ChatResponse ollamaResponse = this.retryTemplate.execute(ctx -> this.chatApi.chat(request));
List<AssistantMessage.ToolCall> toolCalls = ollamaResponse.message().toolCalls() == null ? List.of()
: ollamaResponse.message()
@@ -540,6 +550,8 @@ public class OllamaChatModel implements ChatModel {
private ModelManagementOptions modelManagementOptions = ModelManagementOptions.defaults();
private RetryTemplate retryTemplate = RetryUtils.DEFAULT_RETRY_TEMPLATE;
private Builder() {
}
@@ -574,13 +586,20 @@ public class OllamaChatModel implements ChatModel {
return this;
}
public Builder retryTemplate(RetryTemplate retryTemplate) {
this.retryTemplate = retryTemplate;
return this;
}
public OllamaChatModel build() {
if (this.toolCallingManager != null) {
return new OllamaChatModel(this.ollamaApi, this.defaultOptions, this.toolCallingManager,
this.observationRegistry, this.modelManagementOptions, this.toolExecutionEligibilityPredicate);
this.observationRegistry, this.modelManagementOptions, this.toolExecutionEligibilityPredicate,
this.retryTemplate);
}
return new OllamaChatModel(this.ollamaApi, this.defaultOptions, DEFAULT_TOOL_CALLING_MANAGER,
this.observationRegistry, this.modelManagementOptions, this.toolExecutionEligibilityPredicate);
this.observationRegistry, this.modelManagementOptions, this.toolExecutionEligibilityPredicate,
this.retryTemplate);
}
}

7
models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/api/OllamaApi.java
View File

@@ -51,6 +51,7 @@ import org.springframework.web.reactive.function.client.WebClient;
* @author Christian Tzolov
* @author Thomas Vitale
* @author Jonghoon Park
* @author Alexandros Pappas
* @since 0.8.0
*/
// @formatter:off
@@ -64,6 +65,9 @@ public final class OllamaApi {
private static final Log logger = LogFactory.getLog(OllamaApi.class);
private static final String DEFAULT_BASE_URL = "http://localhost:11434";
private final RestClient restClient;
private final WebClient webClient;
@@ -77,11 +81,13 @@ public final class OllamaApi {
*/
private OllamaApi(String baseUrl, RestClient.Builder restClientBuilder, WebClient.Builder webClientBuilder, ResponseErrorHandler responseErrorHandler) {
Consumer<HttpHeaders> defaultHeaders = headers -> {
headers.setContentType(MediaType.APPLICATION_JSON);
headers.setAccept(List.of(MediaType.APPLICATION_JSON));
};
this.restClient = restClientBuilder
.clone()
.baseUrl(baseUrl)
@@ -89,6 +95,7 @@ public final class OllamaApi {
.defaultStatusHandler(responseErrorHandler)
.build();
this.webClient = webClientBuilder
.clone()
.baseUrl(baseUrl)

2
models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/OllamaChatModelFunctionCallingIT.java
View File

@@ -36,6 +36,7 @@ import org.springframework.ai.ollama.api.OllamaApi;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.ai.ollama.api.tool.MockWeatherService;
import org.springframework.ai.tool.function.FunctionToolCallback;
import org.springframework.ai.retry.RetryUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.boot.test.context.SpringBootTest;
@@ -120,6 +121,7 @@ class OllamaChatModelFunctionCallingIT extends BaseOllamaIT {
return OllamaChatModel.builder()
.ollamaApi(ollamaApi)
.defaultOptions(OllamaOptions.builder().model(MODEL).temperature(0.9).build())
.retryTemplate(RetryUtils.DEFAULT_RETRY_TEMPLATE)
.build();
}

5
models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/OllamaChatModelIT.java
View File

@@ -52,8 +52,12 @@ import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.ai.ollama.management.ModelManagementOptions;
import org.springframework.ai.ollama.management.OllamaModelManager;
import org.springframework.ai.ollama.management.PullModelStrategy;
import org.springframework.ai.support.ToolCallbacks;
import org.springframework.ai.tool.annotation.Tool;
import org.springframework.ai.retry.RetryUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.boot.test.context.SpringBootTest;
@@ -371,6 +375,7 @@ class OllamaChatModelIT extends BaseOllamaIT {
.pullModelStrategy(PullModelStrategy.WHEN_MISSING)
.additionalModels(List.of(ADDITIONAL_MODEL))
.build())
.retryTemplate(RetryUtils.DEFAULT_RETRY_TEMPLATE)
.build();
}

21
models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/OllamaChatModelMultimodalIT.java
View File

@@ -16,6 +16,7 @@
package org.springframework.ai.ollama;
import java.time.Duration;
import java.util.List;
import org.junit.jupiter.api.Test;
@@ -27,11 +28,17 @@ import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.content.Media;
import org.springframework.ai.ollama.api.OllamaApi;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.ai.retry.RetryUtils;
import org.springframework.ai.retry.TransientAiException;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.context.annotation.Bean;
import org.springframework.core.io.ClassPathResource;
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.MimeTypeUtils;
import static org.assertj.core.api.Assertions.assertThat;
@@ -86,9 +93,23 @@ class OllamaChatModelMultimodalIT extends BaseOllamaIT {
@Bean
public OllamaChatModel ollamaChat(OllamaApi ollamaApi) {
RetryTemplate retryTemplate = RetryTemplate.builder()
.maxAttempts(1)
.retryOn(TransientAiException.class)
.fixedBackoff(Duration.ofSeconds(1))
.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();
return OllamaChatModel.builder()
.ollamaApi(ollamaApi)
.defaultOptions(OllamaOptions.builder().model(MODEL).temperature(0.9).build())
.retryTemplate(retryTemplate)
.build();
}

8
models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/OllamaChatModelObservationIT.java
View File

@@ -34,6 +34,7 @@ import org.springframework.ai.observation.conventions.AiProvider;
import org.springframework.ai.ollama.api.OllamaApi;
import org.springframework.ai.ollama.api.OllamaModel;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.ai.retry.RetryUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.boot.test.context.SpringBootTest;
@@ -47,6 +48,7 @@ import static org.springframework.ai.chat.observation.ChatModelObservationDocume
* Integration tests for observation instrumentation in {@link OllamaChatModel}.
*
* @author Thomas Vitale
* @author Alexandros Pappas
*/
@SpringBootTest(classes = OllamaChatModelObservationIT.Config.class)
public class OllamaChatModelObservationIT extends BaseOllamaIT {
@@ -169,7 +171,11 @@ public class OllamaChatModelObservationIT extends BaseOllamaIT {
@Bean
public OllamaChatModel openAiChatModel(OllamaApi ollamaApi, TestObservationRegistry observationRegistry) {
return OllamaChatModel.builder().ollamaApi(ollamaApi).observationRegistry(observationRegistry).build();
return OllamaChatModel.builder()
.ollamaApi(ollamaApi)
.observationRegistry(observationRegistry)
.retryTemplate(RetryUtils.DEFAULT_RETRY_TEMPLATE)
.build();
}
}

2
models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/OllamaChatModelTests.java
View File

@@ -35,6 +35,7 @@ import org.springframework.ai.ollama.api.OllamaApi;
import org.springframework.ai.ollama.api.OllamaModel;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.ai.ollama.management.ModelManagementOptions;
import org.springframework.ai.retry.RetryUtils;
import static org.assertj.core.api.Assertions.assertThat;
import static org.junit.jupiter.api.Assertions.*;
@@ -82,6 +83,7 @@ class OllamaChatModelTests {
() -> OllamaChatModel.builder()
.ollamaApi(this.ollamaApi)
.defaultOptions(OllamaOptions.builder().model(OllamaModel.LLAMA2).build())
.retryTemplate(RetryUtils.DEFAULT_RETRY_TEMPLATE)
.modelManagementOptions(null)
.build());
assertEquals("modelManagementOptions must not be null", exception.getMessage());

4
models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/OllamaChatRequestTests.java
View File

@@ -28,18 +28,21 @@ import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.ai.tool.definition.DefaultToolDefinition;
import org.springframework.ai.tool.definition.ToolDefinition;
import org.springframework.ai.retry.RetryUtils;
import static org.assertj.core.api.Assertions.assertThat;
/**
* @author Christian Tzolov
* @author Thomas Vitale
* @author Alexandros Pappas
*/
class OllamaChatRequestTests {
OllamaChatModel chatModel = OllamaChatModel.builder()
.ollamaApi(OllamaApi.builder().build())
.defaultOptions(OllamaOptions.builder().model("MODEL_NAME").topK(99).temperature(66.6).numGPU(1).build())
.retryTemplate(RetryUtils.DEFAULT_RETRY_TEMPLATE)
.build();
@Test
@@ -146,6 +149,7 @@ class OllamaChatRequestTests {
OllamaChatModel chatModel = OllamaChatModel.builder()
.ollamaApi(OllamaApi.builder().build())
.defaultOptions(OllamaOptions.builder().model("DEFAULT_OPTIONS_MODEL").build())
.retryTemplate(RetryUtils.DEFAULT_RETRY_TEMPLATE)
.build();
var prompt1 = chatModel.buildRequestPrompt(new Prompt("Test message content"));

2
models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/OllamaImage.java
View File

@@ -23,7 +23,7 @@ import org.testcontainers.utility.DockerImageName;
*/
public final class OllamaImage {
public static final DockerImageName DEFAULT_IMAGE = DockerImageName.parse("ollama/ollama:0.5.2");
public static final DockerImageName DEFAULT_IMAGE = DockerImageName.parse("ollama/ollama:0.6.7");
private OllamaImage() {

97
models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/OllamaRetryTests.java Normal file
View File

@@ -0,0 +1,97 @@
package org.springframework.ai.ollama;
import java.time.Instant;
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 org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.ollama.api.OllamaApi;
import org.springframework.ai.ollama.api.OllamaModel;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.ai.retry.RetryUtils;
import org.springframework.ai.retry.TransientAiException;
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.mockito.ArgumentMatchers.isA;
import static org.mockito.Mockito.when;
/**
* Tests for the OllamaRetryTests class.
*
* @author Alexandros Pappas
*/
@ExtendWith(MockitoExtension.class)
class OllamaRetryTests {
private static final String MODEL = OllamaModel.LLAMA3_2.getName();
private TestRetryListener retryListener;
private RetryTemplate retryTemplate;
@Mock
private OllamaApi ollamaApi;
private OllamaChatModel chatModel;
@BeforeEach
public void beforeEach() {
this.retryTemplate = RetryUtils.SHORT_RETRY_TEMPLATE;
this.retryListener = new TestRetryListener();
this.retryTemplate.registerListener(this.retryListener);
this.chatModel = OllamaChatModel.builder()
.ollamaApi(this.ollamaApi)
.defaultOptions(OllamaOptions.builder().model(MODEL).temperature(0.9).build())
.retryTemplate(this.retryTemplate)
.build();
}
@Test
void ollamaChatTransientError() {
String promptText = "What is the capital of Bulgaria and what is the size? What it the national anthem?";
var expectedChatResponse = new OllamaApi.ChatResponse("CHAT_COMPLETION_ID", Instant.now(),
OllamaApi.Message.builder(OllamaApi.Message.Role.ASSISTANT).content("Response").build(), null, true,
null, null, null, null, null, null);
when(this.ollamaApi.chat(isA(OllamaApi.ChatRequest.class)))
.thenThrow(new TransientAiException("Transient Error 1"))
.thenThrow(new TransientAiException("Transient Error 2"))
.thenReturn(expectedChatResponse);
var result = this.chatModel.call(new Prompt(promptText));
assertThat(result).isNotNull();
assertThat(result.getResult().getOutput().getText()).isSameAs("Response");
assertThat(this.retryListener.onSuccessRetryCount).isEqualTo(2);
assertThat(this.retryListener.onErrorRetryCount).isEqualTo(2);
}
private static 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) {
this.onSuccessRetryCount = context.getRetryCount();
}
@Override
public <T, E extends Throwable> void onError(RetryContext context, RetryCallback<T, E> callback,
Throwable throwable) {
this.onErrorRetryCount = context.getRetryCount();
}
}
}