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

Increase the timeout/retry for auto-conf BaseOllamaIT

Browse Source 
Add ModelManagementOptions configuration to BaseOllamaIT
  with 2 max retries and 10 minute timeout
This commit is contained in:
Christian Tzolov
2024-10-25 15:43:17 +02:00
parent b6222d42e8
commit 9c7af5641f
1 changed files with 9 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/ollama/BaseOllamaIT.java
View File

@@ -18,7 +18,10 @@ package org.springframework.ai.autoconfigure.ollama;
import org.testcontainers.ollama.OllamaContainer;
import java.time.Duration;
import org.springframework.ai.ollama.api.OllamaApi;
import org.springframework.ai.ollama.management.ModelManagementOptions;
import org.springframework.ai.ollama.management.OllamaModelManager;
import org.springframework.ai.ollama.management.PullModelStrategy;
@@ -48,7 +51,12 @@ public class BaseOllamaIT {
if (useTestcontainers) {
baseUrl = ollamaContainer.getEndpoint();
}
var ollamaModelManager = new OllamaModelManager(new OllamaApi(baseUrl));
var modelManagementOptions = ModelManagementOptions.builder()
.withMaxRetries(2)
.withTimeout(Duration.ofMinutes(10))
.build();
var ollamaModelManager = new OllamaModelManager(new OllamaApi(baseUrl), modelManagementOptions);
ollamaModelManager.pullModel(model, PullModelStrategy.WHEN_MISSING);
return baseUrl;
}