Skip to content

Commit b66ffb4

Browse files
committed
Extend Ollama's model auto-pull timeout for ITs
1 parent 89e583c commit b66ffb4

File tree

1 file changed

+8
-1
lines changed

1 file changed

+8
-1
lines changed

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

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,9 @@
11
package org.springframework.ai.ollama;
22

3+
import java.time.Duration;
4+
35
import org.springframework.ai.ollama.api.OllamaApi;
6+
import org.springframework.ai.ollama.management.ModelManagementOptions;
47
import org.springframework.ai.ollama.management.OllamaModelManager;
58
import org.springframework.ai.ollama.management.PullModelStrategy;
69
import org.springframework.util.StringUtils;
@@ -53,7 +56,11 @@ public static OllamaApi buildOllamaApiWithModel(String model) {
5356
}
5457

5558
public static void ensureModelIsPresent(OllamaApi ollamaApi, String model) {
56-
var ollamaModelManager = new OllamaModelManager(ollamaApi);
59+
var modelManagementOptions = ModelManagementOptions.builder()
60+
.withMaxRetries(2)
61+
.withTimeout(Duration.ofMinutes(10))
62+
.build();
63+
var ollamaModelManager = new OllamaModelManager(ollamaApi, modelManagementOptions);
5764
ollamaModelManager.pullModel(model, PullModelStrategy.WHEN_MISSING);
5865
}
5966

0 commit comments

Comments
 (0)