Skip to content

Commit ef2e39a

Browse files
ThomasVitaletzolov
ThomasVitale
authored and
tzolov
committed
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 <[email protected]>
1 parent d5bc9c9 commit ef2e39a

File tree

3 files changed

+92
-22
lines changed

3 files changed

+92
-22
lines changed

‎models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/management/OllamaModelManager.java‎

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
import org.springframework.ai.ollama.api.OllamaApi.DeleteModelRequest;
2222
import org.springframework.ai.ollama.api.OllamaApi.ListModelResponse;
2323
import org.springframework.ai.ollama.api.OllamaApi.PullModelRequest;
24+
import org.springframework.util.Assert;
2425
import org.springframework.util.CollectionUtils;
2526
import reactor.util.retry.Retry;
2627

@@ -55,14 +56,28 @@ public OllamaModelManager(OllamaApi ollamaApi, ModelManagementOptions options) {
5556
}
5657

5758
public boolean isModelAvailable(String modelName) {
59+
Assert.hasText(modelName, "modelName must not be empty");
5860
ListModelResponse listModelResponse = ollamaApi.listModels();
5961
if (!CollectionUtils.isEmpty(listModelResponse.models())) {
60-
// Not an equality check to support the implicit ":latest" tag.
61-
return listModelResponse.models().stream().anyMatch(m -> m.name().contains(modelName));
62+
var normalizedModelName = normalizeModelName(modelName);
63+
return listModelResponse.models().stream().anyMatch(m -> m.name().equals(normalizedModelName));
6264
}
6365
return false;
6466
}
6567

68+
/**
69+
* If the name follows the format "<string>:<string>", leave it as is. If the name
70+
* follows the format "<string>" and doesn't include any ":" sign, then add ":latest"
71+
* as a suffix.
72+
*/
73+
private String normalizeModelName(String modelName) {
74+
var modelNameWithoutSpaces = modelName.trim();
75+
if (modelNameWithoutSpaces.contains(":")) {
76+
return modelNameWithoutSpaces;
77+
}
78+
return modelNameWithoutSpaces + ":latest";
79+
}
80+
6681
public void deleteModel(String modelName) {
6782
logger.info("Start deletion of model: {}", modelName);
6883
if (!isModelAvailable(modelName)) {

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

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -11,11 +11,13 @@ public class BaseOllamaIT {
1111
// Toggle for running tests locally on native Ollama for a faster feedback loop.
1212
private static final boolean useTestcontainers = true;
1313

14-
public static final OllamaContainer ollamaContainer;
14+
public static OllamaContainer ollamaContainer;
1515

1616
static {
17-
ollamaContainer = new OllamaContainer(OllamaImage.DEFAULT_IMAGE).withReuse(true);
18-
ollamaContainer.start();
17+
if (useTestcontainers) {
18+
ollamaContainer = new OllamaContainer(OllamaImage.DEFAULT_IMAGE).withReuse(true);
19+
ollamaContainer.start();
20+
}
1921
}
2022

2123
/**

‎models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/management/OllamaModelManagerIT.java‎

Lines changed: 70 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -63,24 +63,77 @@ public void whenModelNotAvailableReturnFalse() {
6363
}
6464

6565
@Test
66-
public void pullAndDeleteModel() {
67-
var model = "all-minilm";
68-
modelManager.pullModel(model, PullModelStrategy.WHEN_MISSING);
69-
var isModelAvailable = modelManager.isModelAvailable(model);
70-
assertThat(isModelAvailable).isTrue();
71-
72-
modelManager.deleteModel(model);
73-
isModelAvailable = modelManager.isModelAvailable(model);
74-
assertThat(isModelAvailable).isFalse();
75-
76-
model = "all-minilm:latest";
77-
modelManager.pullModel(model, PullModelStrategy.WHEN_MISSING);
78-
isModelAvailable = modelManager.isModelAvailable(model);
79-
assertThat(isModelAvailable).isTrue();
66+
public void pullAndDeleteModelFromOllama() {
67+
// Pull model with explicit version.
68+
var modelWithExplicitVersion = "all-minilm:33m";
69+
modelManager.deleteModel(modelWithExplicitVersion);
70+
modelManager.pullModel(modelWithExplicitVersion, PullModelStrategy.WHEN_MISSING);
71+
var isModelWithExplicitVersionAvailable = modelManager.isModelAvailable(modelWithExplicitVersion);
72+
assertThat(isModelWithExplicitVersionAvailable).isTrue();
73+
74+
// Pull same model without version, which should pull the "latest" version.
75+
var modelWithoutVersion = "all-minilm";
76+
modelManager.deleteModel(modelWithoutVersion);
77+
var isModelWithoutVersionAvailable = modelManager.isModelAvailable(modelWithoutVersion);
78+
assertThat(isModelWithoutVersionAvailable).isFalse();
79+
isModelWithExplicitVersionAvailable = modelManager.isModelAvailable(modelWithExplicitVersion);
80+
assertThat(isModelWithExplicitVersionAvailable).isTrue();
81+
82+
modelManager.pullModel(modelWithoutVersion, PullModelStrategy.WHEN_MISSING);
83+
isModelWithoutVersionAvailable = modelManager.isModelAvailable(modelWithoutVersion);
84+
assertThat(isModelWithoutVersionAvailable).isTrue();
85+
86+
// Pull model with ":latest" suffix, with has the same effect as pulling the model
87+
// without version.
88+
var modelWithLatestVersion = "all-minilm:latest";
89+
var isModelWithLatestVersionAvailable = modelManager.isModelAvailable(modelWithLatestVersion);
90+
assertThat(isModelWithLatestVersionAvailable).isTrue();
91+
92+
// Final clean-up.
93+
modelManager.deleteModel(modelWithExplicitVersion);
94+
isModelWithExplicitVersionAvailable = modelManager.isModelAvailable(modelWithExplicitVersion);
95+
assertThat(isModelWithExplicitVersionAvailable).isFalse();
96+
97+
modelManager.deleteModel(modelWithLatestVersion);
98+
isModelWithLatestVersionAvailable = modelManager.isModelAvailable(modelWithLatestVersion);
99+
assertThat(isModelWithLatestVersionAvailable).isFalse();
100+
}
80101

81-
modelManager.deleteModel(model);
82-
isModelAvailable = modelManager.isModelAvailable(model);
83-
assertThat(isModelAvailable).isFalse();
102+
@Test
103+
public void pullAndDeleteModelFromHuggingFace() {
104+
// Pull model with explicit version.
105+
var modelWithExplicitVersion = "hf.co/SanctumAI/Llama-3.2-1B-Instruct-GGUF:Q3_K_S";
106+
modelManager.deleteModel(modelWithExplicitVersion);
107+
modelManager.pullModel(modelWithExplicitVersion, PullModelStrategy.WHEN_MISSING);
108+
var isModelWithExplicitVersionAvailable = modelManager.isModelAvailable(modelWithExplicitVersion);
109+
assertThat(isModelWithExplicitVersionAvailable).isTrue();
110+
111+
// Pull same model without version, which should pull the "latest" version.
112+
var modelWithoutVersion = "hf.co/SanctumAI/Llama-3.2-1B-Instruct-GGUF";
113+
modelManager.deleteModel(modelWithoutVersion);
114+
var isModelWithoutVersionAvailable = modelManager.isModelAvailable(modelWithoutVersion);
115+
assertThat(isModelWithoutVersionAvailable).isFalse();
116+
isModelWithExplicitVersionAvailable = modelManager.isModelAvailable(modelWithExplicitVersion);
117+
assertThat(isModelWithExplicitVersionAvailable).isTrue();
118+
119+
modelManager.pullModel(modelWithoutVersion, PullModelStrategy.WHEN_MISSING);
120+
isModelWithoutVersionAvailable = modelManager.isModelAvailable(modelWithoutVersion);
121+
assertThat(isModelWithoutVersionAvailable).isTrue();
122+
123+
// Pull model with ":latest" suffix, with has the same effect as pulling the model
124+
// without version.
125+
var modelWithLatestVersion = "hf.co/SanctumAI/Llama-3.2-1B-Instruct-GGUF:latest";
126+
var isModelWithLatestVersionAvailable = modelManager.isModelAvailable(modelWithLatestVersion);
127+
assertThat(isModelWithLatestVersionAvailable).isTrue();
128+
129+
// Final clean-up.
130+
modelManager.deleteModel(modelWithExplicitVersion);
131+
isModelWithExplicitVersionAvailable = modelManager.isModelAvailable(modelWithExplicitVersion);
132+
assertThat(isModelWithExplicitVersionAvailable).isFalse();
133+
134+
modelManager.deleteModel(modelWithLatestVersion);
135+
isModelWithLatestVersionAvailable = modelManager.isModelAvailable(modelWithLatestVersion);
136+
assertThat(isModelWithLatestVersionAvailable).isFalse();
84137
}
85138

86139
@Test

0 commit comments

Comments
 (0)