Skip to content

Commit 9c7af56

Browse files
committed
Increase the timeout/retry for auto-conf BaseOllamaIT
Add ModelManagementOptions configuration to BaseOllamaIT with 2 max retries and 10 minute timeout
1 parent b6222d4 commit 9c7af56

File tree

1 file changed

+9
-1
lines changed
  • spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/ollama

1 file changed

+9
-1
lines changed

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

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,10 @@
1818

1919
import org.testcontainers.ollama.OllamaContainer;
2020

21+
import java.time.Duration;
22+
2123
import org.springframework.ai.ollama.api.OllamaApi;
24+
import org.springframework.ai.ollama.management.ModelManagementOptions;
2225
import org.springframework.ai.ollama.management.OllamaModelManager;
2326
import org.springframework.ai.ollama.management.PullModelStrategy;
2427

@@ -48,7 +51,12 @@ public static String buildConnectionWithModel(String model) {
4851
if (useTestcontainers) {
4952
baseUrl = ollamaContainer.getEndpoint();
5053
}
51-
var ollamaModelManager = new OllamaModelManager(new OllamaApi(baseUrl));
54+
55+
var modelManagementOptions = ModelManagementOptions.builder()
56+
.withMaxRetries(2)
57+
.withTimeout(Duration.ofMinutes(10))
58+
.build();
59+
var ollamaModelManager = new OllamaModelManager(new OllamaApi(baseUrl), modelManagementOptions);
5260
ollamaModelManager.pullModel(model, PullModelStrategy.WHEN_MISSING);
5361
return baseUrl;
5462
}

0 commit comments

Comments
 (0)