From ef2e39a8adad01257ca05b0abf93a086589313a5 Mon Sep 17 00:00:00 2001 From: Thomas Vitale Date: Mon, 21 Oct 2024 17:15:01 +0200 Subject: [PATCH] Normalize Ollama model names in auto-pull feature In order to support edge cases due to different naming formats, this PR introduced an explicit normalization logic to ensure the correct matching when checking for the availability of a certain model. Integration tests have been added to cover the different scenarios, including models from Ollama and from Hugging Face. Also fix the container creation on the useTestcontainers flag (christian) Signed-off-by: Thomas Vitale --- .../ollama/management/OllamaModelManager.java | 19 ++++- .../ai/ollama/BaseOllamaIT.java | 8 +- .../management/OllamaModelManagerIT.java | 83 +++++++++++++++---- 3 files changed, 90 insertions(+), 20 deletions(-) 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