diff --git a/models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/management/OllamaModelManager.java b/models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/management/OllamaModelManager.java index 31e2249b8..ebc736c82 100644 --- a/models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/management/OllamaModelManager.java +++ b/models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/management/OllamaModelManager.java @@ -21,6 +21,7 @@ import org.springframework.ai.ollama.api.OllamaApi; import org.springframework.ai.ollama.api.OllamaApi.DeleteModelRequest; import org.springframework.ai.ollama.api.OllamaApi.ListModelResponse; import org.springframework.ai.ollama.api.OllamaApi.PullModelRequest; +import org.springframework.util.Assert; import org.springframework.util.CollectionUtils; import reactor.util.retry.Retry; @@ -55,14 +56,28 @@ public class OllamaModelManager { } public boolean isModelAvailable(String modelName) { + Assert.hasText(modelName, "modelName must not be empty"); ListModelResponse listModelResponse = ollamaApi.listModels(); if (!CollectionUtils.isEmpty(listModelResponse.models())) { - // Not an equality check to support the implicit ":latest" tag. - return listModelResponse.models().stream().anyMatch(m -> m.name().contains(modelName)); + var normalizedModelName = normalizeModelName(modelName); + return listModelResponse.models().stream().anyMatch(m -> m.name().equals(normalizedModelName)); } return false; } + /** + * If the name follows the format ":", leave it as is. If the name + * follows the format "" and doesn't include any ":" sign, then add ":latest" + * as a suffix. + */ + private String normalizeModelName(String modelName) { + var modelNameWithoutSpaces = modelName.trim(); + if (modelNameWithoutSpaces.contains(":")) { + return modelNameWithoutSpaces; + } + return modelNameWithoutSpaces + ":latest"; + } + public void deleteModel(String modelName) { logger.info("Start deletion of model: {}", modelName); if (!isModelAvailable(modelName)) { 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 index edb438072..635c0a16f 100644 --- 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 @@ -11,11 +11,13 @@ public class BaseOllamaIT { // Toggle for running tests locally on native Ollama for a faster feedback loop. private static final boolean useTestcontainers = true; - public static final OllamaContainer ollamaContainer; + public static OllamaContainer ollamaContainer; static { - ollamaContainer = new OllamaContainer(OllamaImage.DEFAULT_IMAGE).withReuse(true); - ollamaContainer.start(); + if (useTestcontainers) { + ollamaContainer = new OllamaContainer(OllamaImage.DEFAULT_IMAGE).withReuse(true); + ollamaContainer.start(); + } } /** diff --git a/models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/management/OllamaModelManagerIT.java b/models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/management/OllamaModelManagerIT.java index 52d19e6d8..ab99833ca 100644 --- a/models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/management/OllamaModelManagerIT.java +++ b/models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/management/OllamaModelManagerIT.java @@ -63,24 +63,77 @@ class OllamaModelManagerIT extends BaseOllamaIT { } @Test - public void pullAndDeleteModel() { - var model = "all-minilm"; - modelManager.pullModel(model, PullModelStrategy.WHEN_MISSING); - var isModelAvailable = modelManager.isModelAvailable(model); - assertThat(isModelAvailable).isTrue(); + public void pullAndDeleteModelFromOllama() { + // Pull model with explicit version. + var modelWithExplicitVersion = "all-minilm:33m"; + modelManager.deleteModel(modelWithExplicitVersion); + modelManager.pullModel(modelWithExplicitVersion, PullModelStrategy.WHEN_MISSING); + var isModelWithExplicitVersionAvailable = modelManager.isModelAvailable(modelWithExplicitVersion); + assertThat(isModelWithExplicitVersionAvailable).isTrue(); - modelManager.deleteModel(model); - isModelAvailable = modelManager.isModelAvailable(model); - assertThat(isModelAvailable).isFalse(); + // Pull same model without version, which should pull the "latest" version. + var modelWithoutVersion = "all-minilm"; + modelManager.deleteModel(modelWithoutVersion); + var isModelWithoutVersionAvailable = modelManager.isModelAvailable(modelWithoutVersion); + assertThat(isModelWithoutVersionAvailable).isFalse(); + isModelWithExplicitVersionAvailable = modelManager.isModelAvailable(modelWithExplicitVersion); + assertThat(isModelWithExplicitVersionAvailable).isTrue(); - model = "all-minilm:latest"; - modelManager.pullModel(model, PullModelStrategy.WHEN_MISSING); - isModelAvailable = modelManager.isModelAvailable(model); - assertThat(isModelAvailable).isTrue(); + modelManager.pullModel(modelWithoutVersion, PullModelStrategy.WHEN_MISSING); + isModelWithoutVersionAvailable = modelManager.isModelAvailable(modelWithoutVersion); + assertThat(isModelWithoutVersionAvailable).isTrue(); - modelManager.deleteModel(model); - isModelAvailable = modelManager.isModelAvailable(model); - assertThat(isModelAvailable).isFalse(); + // Pull model with ":latest" suffix, with has the same effect as pulling the model + // without version. + var modelWithLatestVersion = "all-minilm:latest"; + var isModelWithLatestVersionAvailable = modelManager.isModelAvailable(modelWithLatestVersion); + assertThat(isModelWithLatestVersionAvailable).isTrue(); + + // Final clean-up. + modelManager.deleteModel(modelWithExplicitVersion); + isModelWithExplicitVersionAvailable = modelManager.isModelAvailable(modelWithExplicitVersion); + assertThat(isModelWithExplicitVersionAvailable).isFalse(); + + modelManager.deleteModel(modelWithLatestVersion); + isModelWithLatestVersionAvailable = modelManager.isModelAvailable(modelWithLatestVersion); + assertThat(isModelWithLatestVersionAvailable).isFalse(); + } + + @Test + public void pullAndDeleteModelFromHuggingFace() { + // Pull model with explicit version. + var modelWithExplicitVersion = "hf.co/SanctumAI/Llama-3.2-1B-Instruct-GGUF:Q3_K_S"; + modelManager.deleteModel(modelWithExplicitVersion); + modelManager.pullModel(modelWithExplicitVersion, PullModelStrategy.WHEN_MISSING); + var isModelWithExplicitVersionAvailable = modelManager.isModelAvailable(modelWithExplicitVersion); + assertThat(isModelWithExplicitVersionAvailable).isTrue(); + + // Pull same model without version, which should pull the "latest" version. + var modelWithoutVersion = "hf.co/SanctumAI/Llama-3.2-1B-Instruct-GGUF"; + modelManager.deleteModel(modelWithoutVersion); + var isModelWithoutVersionAvailable = modelManager.isModelAvailable(modelWithoutVersion); + assertThat(isModelWithoutVersionAvailable).isFalse(); + isModelWithExplicitVersionAvailable = modelManager.isModelAvailable(modelWithExplicitVersion); + assertThat(isModelWithExplicitVersionAvailable).isTrue(); + + modelManager.pullModel(modelWithoutVersion, PullModelStrategy.WHEN_MISSING); + isModelWithoutVersionAvailable = modelManager.isModelAvailable(modelWithoutVersion); + assertThat(isModelWithoutVersionAvailable).isTrue(); + + // Pull model with ":latest" suffix, with has the same effect as pulling the model + // without version. + var modelWithLatestVersion = "hf.co/SanctumAI/Llama-3.2-1B-Instruct-GGUF:latest"; + var isModelWithLatestVersionAvailable = modelManager.isModelAvailable(modelWithLatestVersion); + assertThat(isModelWithLatestVersionAvailable).isTrue(); + + // Final clean-up. + modelManager.deleteModel(modelWithExplicitVersion); + isModelWithExplicitVersionAvailable = modelManager.isModelAvailable(modelWithExplicitVersion); + assertThat(isModelWithExplicitVersionAvailable).isFalse(); + + modelManager.deleteModel(modelWithLatestVersion); + isModelWithLatestVersionAvailable = modelManager.isModelAvailable(modelWithLatestVersion); + assertThat(isModelWithLatestVersionAvailable).isFalse(); } @Test