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

Consolidate Ollama auto-pull logic

Browse Source 
Consolidate the Ollama auto-pull logic at startup time, supporting the auto-pull for the default models specified via configuration properties and for optional models specified for initialization.

Signed-off-by: Thomas Vitale <ThomasVitale@users.noreply.github.com>
This commit is contained in:
Thomas Vitale
2024-10-21 19:56:01 +02:00
committed by Christian Tzolov
parent 1cadc49a9b
commit d5bc9c998c
13 changed files with 140 additions and 155 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/OllamaChatModel.java
View File

@@ -97,7 +97,7 @@ public class OllamaChatModel extends AbstractToolCallSupport implements ChatMode
this.defaultOptions = defaultOptions;
this.observationRegistry = observationRegistry;
this.modelManager = new OllamaModelManager(chatApi, modelManagementOptions);
initializeModelIfEnabled(defaultOptions.getModel(), modelManagementOptions.pullModelStrategy());
initializeModel(defaultOptions.getModel(), modelManagementOptions.pullModelStrategy());
}
public static Builder builder() {
@@ -302,11 +302,6 @@ public class OllamaChatModel extends AbstractToolCallSupport implements ChatMode
}
OllamaOptions mergedOptions = ModelOptionsUtils.merge(runtimeOptions, this.defaultOptions, OllamaOptions.class);
mergedOptions.setPullModelStrategy(this.defaultOptions.getPullModelStrategy());
if (runtimeOptions != null && runtimeOptions.getPullModelStrategy() != null) {
mergedOptions.setPullModelStrategy(runtimeOptions.getPullModelStrategy());
}
// Override the model.
if (!StringUtils.hasText(mergedOptions.getModel())) {
throw new IllegalArgumentException("Model is not set!");
@@ -331,8 +326,6 @@ public class OllamaChatModel extends AbstractToolCallSupport implements ChatMode
requestBuilder.withTools(this.getFunctionTools(functionsForThisRequest));
}
initializeModelIfEnabled(mergedOptions.getModel(), mergedOptions.getPullModelStrategy());
return requestBuilder.build();
}
@@ -379,7 +372,7 @@ public class OllamaChatModel extends AbstractToolCallSupport implements ChatMode
/**
* Pull the given model into Ollama based on the specified strategy.
*/
private void initializeModelIfEnabled(String model, PullModelStrategy pullModelStrategy) {
private void initializeModel(String model, PullModelStrategy pullModelStrategy) {
if (pullModelStrategy != null && !PullModelStrategy.NEVER.equals(pullModelStrategy)) {
this.modelManager.pullModel(model, pullModelStrategy);
}

11
models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/OllamaEmbeddingModel.java
View File

@@ -77,7 +77,7 @@ public class OllamaEmbeddingModel extends AbstractEmbeddingModel {
this.observationRegistry = observationRegistry;
this.modelManager = new OllamaModelManager(ollamaApi, modelManagementOptions);
initializeModelIfEnabled(defaultOptions.getModel(), modelManagementOptions.pullModelStrategy());
initializeModel(defaultOptions.getModel(), modelManagementOptions.pullModelStrategy());
}
public static Builder builder() {
@@ -139,19 +139,12 @@ public class OllamaEmbeddingModel extends AbstractEmbeddingModel {
OllamaOptions mergedOptions = ModelOptionsUtils.merge(runtimeOptions, this.defaultOptions, OllamaOptions.class);
mergedOptions.setPullModelStrategy(this.defaultOptions.getPullModelStrategy());
if (runtimeOptions != null && runtimeOptions.getPullModelStrategy() != null) {
mergedOptions.setPullModelStrategy(runtimeOptions.getPullModelStrategy());
}
// Override the model.
if (!StringUtils.hasText(mergedOptions.getModel())) {
throw new IllegalArgumentException("Model is not set!");
}
String model = mergedOptions.getModel();
initializeModelIfEnabled(mergedOptions.getModel(), mergedOptions.getPullModelStrategy());
return new OllamaApi.EmbeddingsRequest(model, inputContent, DurationParser.parse(mergedOptions.getKeepAlive()),
OllamaOptions.filterNonSupportedFields(mergedOptions.toMap()), mergedOptions.getTruncate());
}
@@ -163,7 +156,7 @@ public class OllamaEmbeddingModel extends AbstractEmbeddingModel {
/**
* Pull the given model into Ollama based on the specified strategy.
*/
private void initializeModelIfEnabled(String model, PullModelStrategy pullModelStrategy) {
private void initializeModel(String model, PullModelStrategy pullModelStrategy) {
if (pullModelStrategy != null && !PullModelStrategy.NEVER.equals(pullModelStrategy)) {
this.modelManager.pullModel(model, pullModelStrategy);
}

28
models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/api/OllamaOptions.java
View File

@@ -28,7 +28,6 @@ import org.springframework.ai.embedding.EmbeddingOptions;
import org.springframework.ai.model.ModelOptionsUtils;
import org.springframework.ai.model.function.FunctionCallback;
import org.springframework.ai.model.function.FunctionCallingOptions;
import org.springframework.ai.ollama.management.PullModelStrategy;
import org.springframework.boot.context.properties.NestedConfigurationProperty;
import org.springframework.util.Assert;
@@ -303,12 +302,6 @@ public class OllamaOptions implements FunctionCallingOptions, ChatOptions, Embed
@JsonIgnore
private Map<String, Object> toolContext;
/**
* Strategy for pulling models at run-time.
*/
@JsonIgnore
private PullModelStrategy pullModelStrategy;
public static OllamaOptions builder() {
return new OllamaOptions();
}
@@ -521,11 +514,6 @@ public class OllamaOptions implements FunctionCallingOptions, ChatOptions, Embed
return this;
}
public OllamaOptions withPullModelStrategy(PullModelStrategy pullModelStrategy) {
this.pullModelStrategy = pullModelStrategy;
return this;
}
// -------------------
// Getters and Setters
// -------------------
@@ -866,14 +854,6 @@ public class OllamaOptions implements FunctionCallingOptions, ChatOptions, Embed
this.toolContext = toolContext;
}
public PullModelStrategy getPullModelStrategy() {
return this.pullModelStrategy;
}
public void setPullModelStrategy(PullModelStrategy pullModelStrategy) {
this.pullModelStrategy = pullModelStrategy;
}
/**
* Convert the {@link OllamaOptions} object to a {@link Map} of key/value pairs.
* @return The {@link Map} of key/value pairs.
@@ -944,8 +924,7 @@ public class OllamaOptions implements FunctionCallingOptions, ChatOptions, Embed
.withFunctions(fromOptions.getFunctions())
.withProxyToolCalls(fromOptions.getProxyToolCalls())
.withFunctionCallbacks(fromOptions.getFunctionCallbacks())
.withToolContext(fromOptions.getToolContext())
.withPullModelStrategy(fromOptions.getPullModelStrategy());
.withToolContext(fromOptions.getToolContext());
}
// @formatter:on
@@ -975,8 +954,7 @@ public class OllamaOptions implements FunctionCallingOptions, ChatOptions, Embed
&& Objects.equals(penalizeNewline, that.penalizeNewline) && Objects.equals(stop, that.stop)
&& Objects.equals(functionCallbacks, that.functionCallbacks)
&& Objects.equals(proxyToolCalls, that.proxyToolCalls) && Objects.equals(functions, that.functions)
&& Objects.equals(toolContext, that.toolContext)
&& Objects.equals(pullModelStrategy, that.pullModelStrategy);
&& Objects.equals(toolContext, that.toolContext);
}
@Override
@@ -987,7 +965,7 @@ public class OllamaOptions implements FunctionCallingOptions, ChatOptions, Embed
this.topP, tfsZ, this.typicalP, this.repeatLastN, this.temperature, this.repeatPenalty,
this.presencePenalty, this.frequencyPenalty, this.mirostat, this.mirostatTau, this.mirostatEta,
this.penalizeNewline, this.stop, this.functionCallbacks, this.functions, this.proxyToolCalls,
this.toolContext, this.pullModelStrategy);
this.toolContext);
}
}

42
models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/management/ModelManagementOptions.java
View File

@@ -26,7 +26,49 @@ import java.util.List;
*/
public record ModelManagementOptions(PullModelStrategy pullModelStrategy, List<String> additionalModels,
Duration timeout, Integer maxRetries) {
public static ModelManagementOptions defaults() {
return new ModelManagementOptions(PullModelStrategy.NEVER, List.of(), Duration.ofMinutes(5), 0);
}
public static Builder builder() {
return new Builder();
}
public static class Builder {
private PullModelStrategy pullModelStrategy = PullModelStrategy.NEVER;
private List<String> additionalModels = List.of();
private Duration timeout = Duration.ofMinutes(5);
private Integer maxRetries = 0;
public Builder withPullModelStrategy(PullModelStrategy pullModelStrategy) {
this.pullModelStrategy = pullModelStrategy;
return this;
}
public Builder withAdditionalModels(List<String> additionalModels) {
this.additionalModels = additionalModels;
return this;
}
public Builder withTimeout(Duration timeout) {
this.timeout = timeout;
return this;
}
public Builder withMaxRetries(Integer maxRetries) {
this.maxRetries = maxRetries;
return this;
}
public ModelManagementOptions build() {
return new ModelManagementOptions(pullModelStrategy, additionalModels, timeout, maxRetries);
}
}
}

19
models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/OllamaChatModelIT.java
View File

@@ -33,6 +33,7 @@ import org.springframework.ai.converter.ListOutputConverter;
import org.springframework.ai.converter.MapOutputConverter;
import org.springframework.ai.ollama.api.OllamaApi;
import org.springframework.ai.ollama.api.OllamaModel;
import org.springframework.ai.ollama.management.ModelManagementOptions;
import org.springframework.ai.ollama.management.OllamaModelManager;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.ai.ollama.management.PullModelStrategy;
@@ -56,6 +57,8 @@ class OllamaChatModelIT extends BaseOllamaIT {
private static final String MODEL = OllamaModel.LLAMA3_2.getName();
private static final String ADDITIONAL_MODEL = "tinyllama";
@Autowired
private OllamaChatModel chatModel;
@@ -65,23 +68,17 @@ class OllamaChatModelIT extends BaseOllamaIT {
@Test
void autoPullModelTest() {
var modelManager = new OllamaModelManager(ollamaApi);
var model = "tinyllama";
modelManager.deleteModel(model);
assertThat(modelManager.isModelAvailable(model)).isFalse();
assertThat(modelManager.isModelAvailable(ADDITIONAL_MODEL)).isTrue();
String joke = ChatClient.create(chatModel)
.prompt("Tell me a joke")
.options(OllamaOptions.builder()
.withModel(model)
.withPullModelStrategy(PullModelStrategy.WHEN_MISSING)
.build())
.options(OllamaOptions.builder().withModel(ADDITIONAL_MODEL).build())
.call()
.content();
assertThat(joke).isNotEmpty();
assertThat(modelManager.isModelAvailable(model)).isTrue();
modelManager.deleteModel(model);
modelManager.deleteModel(ADDITIONAL_MODEL);
}
@Test
@@ -249,6 +246,10 @@ class OllamaChatModelIT extends BaseOllamaIT {
return OllamaChatModel.builder()
.withOllamaApi(ollamaApi)
.withDefaultOptions(OllamaOptions.create().withModel(MODEL).withTemperature(0.9))
.withModelManagementOptions(ModelManagementOptions.builder()
.withPullModelStrategy(PullModelStrategy.WHEN_MISSING)
.withAdditionalModels(List.of(ADDITIONAL_MODEL))
.build())
.build();
}

24
models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/OllamaEmbeddingModelIT.java
View File

@@ -21,6 +21,7 @@ import org.springframework.ai.embedding.EmbeddingRequest;
import org.springframework.ai.embedding.EmbeddingResponse;
import org.springframework.ai.ollama.api.OllamaApi;
import org.springframework.ai.ollama.api.OllamaModel;
import org.springframework.ai.ollama.management.ModelManagementOptions;
import org.springframework.ai.ollama.management.OllamaModelManager;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.ai.ollama.management.PullModelStrategy;
@@ -41,6 +42,8 @@ class OllamaEmbeddingModelIT extends BaseOllamaIT {
private static final String MODEL = OllamaModel.NOMIC_EMBED_TEXT.getName();
private static final String ADDITIONAL_MODEL = "all-minilm";
@Autowired
private OllamaEmbeddingModel embeddingModel;
@@ -65,36 +68,29 @@ class OllamaEmbeddingModelIT extends BaseOllamaIT {
}
@Test
void autoPullModel() {
void autoPullModelAtStartupTime() {
var model = "all-minilm";
assertThat(embeddingModel).isNotNull();
var modelManager = new OllamaModelManager(ollamaApi);
modelManager.deleteModel(model);
assertThat(modelManager.isModelAvailable(model)).isFalse();
assertThat(modelManager.isModelAvailable(ADDITIONAL_MODEL)).isTrue();
EmbeddingResponse embeddingResponse = embeddingModel
.call(new EmbeddingRequest(List.of("Hello World", "Something else"),
OllamaOptions.builder()
.withModel(model)
.withPullModelStrategy(PullModelStrategy.WHEN_MISSING)
.withTruncate(false)
.build()));
assertThat(modelManager.isModelAvailable(model)).isTrue();
OllamaOptions.builder().withModel(model).withTruncate(false).build()));
assertThat(embeddingResponse.getResults()).hasSize(2);
assertThat(embeddingResponse.getResults().get(0).getIndex()).isEqualTo(0);
assertThat(embeddingResponse.getResults().get(0).getOutput()).isNotEmpty();
assertThat(embeddingResponse.getResults().get(1).getIndex()).isEqualTo(1);
assertThat(embeddingResponse.getResults().get(1).getOutput()).isNotEmpty();
assertThat(embeddingResponse.getMetadata().getModel()).contains(model);
assertThat(embeddingResponse.getMetadata().getModel()).contains(ADDITIONAL_MODEL);
assertThat(embeddingResponse.getMetadata().getUsage().getPromptTokens()).isEqualTo(4);
assertThat(embeddingResponse.getMetadata().getUsage().getTotalTokens()).isEqualTo(4);
assertThat(embeddingModel.dimensions()).isEqualTo(768);
modelManager.deleteModel(model);
modelManager.deleteModel(ADDITIONAL_MODEL);
}
@SpringBootConfiguration
@@ -110,6 +106,10 @@ class OllamaEmbeddingModelIT extends BaseOllamaIT {
return OllamaEmbeddingModel.builder()
.withOllamaApi(ollamaApi)
.withDefaultOptions(OllamaOptions.create().withModel(MODEL))
.withModelManagementOptions(ModelManagementOptions.builder()
.withPullModelStrategy(PullModelStrategy.WHEN_MISSING)
.withAdditionalModels(List.of(ADDITIONAL_MODEL))
.build())
.build();
}

2
models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/OllamaImage.java
View File

@@ -22,6 +22,6 @@ import org.testcontainers.utility.DockerImageName;
*/
public class OllamaImage {
public static final DockerImageName DEFAULT_IMAGE = DockerImageName.parse("ollama/ollama:0.3.13");
public static final DockerImageName DEFAULT_IMAGE = DockerImageName.parse("ollama/ollama:0.3.14");
}

14
models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/api/OllamaApiModelsIT.java
View File

@@ -15,6 +15,11 @@
*/
package org.springframework.ai.ollama.api;
import static org.assertj.core.api.Assertions.assertThat;
import java.io.IOException;
import java.time.Duration;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.DisabledIf;
@@ -22,11 +27,6 @@ import org.springframework.ai.ollama.BaseOllamaIT;
import org.springframework.http.HttpStatus;
import org.testcontainers.junit.jupiter.Testcontainers;
import java.io.IOException;
import java.time.Duration;
import static org.assertj.core.api.Assertions.assertThat;
/**
* Integration tests for the Ollama APIs to manage models.
*
@@ -36,7 +36,7 @@ import static org.assertj.core.api.Assertions.assertThat;
@DisabledIf("isDisabled")
public class OllamaApiModelsIT extends BaseOllamaIT {
private static final String MODEL = OllamaModel.NOMIC_EMBED_TEXT.getName();
private static final String MODEL = "all-minilm";
static OllamaApi ollamaApi;
@@ -60,7 +60,7 @@ public class OllamaApiModelsIT extends BaseOllamaIT {
var showModelResponse = ollamaApi.showModel(showModelRequest);
assertThat(showModelResponse).isNotNull();
assertThat(showModelResponse.details().family()).isEqualTo("nomic-bert");
assertThat(showModelResponse.details().family()).isEqualTo("bert");
}
@Test