Extend Ollama's model auto-pull timeout for ITs

spring-projects · Oct 24, 2024 · b66ffb4 · b66ffb4
1 parent 89e583c
commit b66ffb4
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/BaseOllamaIT.java b/models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/BaseOllamaIT.java
@@ -1,6 +1,9 @@
 package org.springframework.ai.ollama;
 
+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;
 import org.springframework.util.StringUtils;
@@ -53,7 +56,11 @@ public static OllamaApi buildOllamaApiWithModel(String model) {
 	}
 
 	public static void ensureModelIsPresent(OllamaApi ollamaApi, String model) {
-		var ollamaModelManager = new OllamaModelManager(ollamaApi);
+		var modelManagementOptions = ModelManagementOptions.builder()
+			.withMaxRetries(2)
+			.withTimeout(Duration.ofMinutes(10))
+			.build();
+		var ollamaModelManager = new OllamaModelManager(ollamaApi, modelManagementOptions);
 		ollamaModelManager.pullModel(model, PullModelStrategy.WHEN_MISSING);
 	}