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

Extend Ollama's model auto-pull timeout for ITs

Browse Source
This commit is contained in:
Christian Tzolov
2024-10-24 12:51:35 +02:00
parent 89e583ca7f
commit b66ffb4d22
1 changed files with 8 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -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 class BaseOllamaIT {
}
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);
}