From 9c7af5641f13f65cc9237152a771de9161106466 Mon Sep 17 00:00:00 2001 From: Christian Tzolov Date: Fri, 25 Oct 2024 15:43:17 +0200 Subject: [PATCH] Increase the timeout/retry for auto-conf BaseOllamaIT Add ModelManagementOptions configuration to BaseOllamaIT with 2 max retries and 10 minute timeout --- .../ai/autoconfigure/ollama/BaseOllamaIT.java | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/ollama/BaseOllamaIT.java b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/ollama/BaseOllamaIT.java index 3a096dc09..3af059f33 100644 --- a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/ollama/BaseOllamaIT.java +++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/ollama/BaseOllamaIT.java @@ -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; }