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

Normalize Ollama model names in auto-pull feature

Browse Source 
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 <ThomasVitale@users.noreply.github.com>
This commit is contained in:
Thomas Vitale
2024-10-21 17:15:01 +02:00
committed by Christian Tzolov
parent d5bc9c998c
commit ef2e39a8ad
3 changed files with 90 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/management/OllamaModelManager.java
View File

@@ -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 "<string>:<string>", leave it as is. If the name
* follows the format "<string>" 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)) {

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

@@ -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();
}
}
/**

83
models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/management/OllamaModelManagerIT.java
View File

@@ -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