From 5e8cecd8b196985284ed5b59be2b10dca3f8dd8f Mon Sep 17 00:00:00 2001 From: Thomas Vitale Date: Sat, 19 Oct 2024 15:34:54 +0200 Subject: [PATCH] Enhance Ollama model auto-pull feature * Fix configuration inheritance issue when default value is not specified. * Make it possible to enable the auto-pull feature only for specific model types (e.g. for chat models only). * Add the possibility to list explicit models to auto-pull at startup time. Update Ollama model defaults and add new embedding model * Change default chat model to Mistral * Change default embedding model to mxbai-embed-large * Add MXBAI_EMBED_LARGE to OllamaModel enum * Remove DEFAULT_MODEL constant from OllamaOptions * Update relevant classes to use new defaults Signed-off-by: Thomas Vitale Co-authored-by:Christian Tzolov --- .../ai/ollama/OllamaChatModel.java | 5 +- .../ai/ollama/OllamaEmbeddingModel.java | 5 +- .../ai/ollama/api/OllamaModel.java | 7 ++- .../ai/ollama/api/OllamaOptions.java | 4 +- .../management/ModelManagementOptions.java | 6 ++- .../ollama/management/OllamaModelManager.java | 4 ++ .../ai/ollama/BaseOllamaIT.java | 9 +++- .../management/OllamaModelManagerIT.java | 19 +++++++ .../ROOT/pages/api/chat/ollama-chat.adoc | 32 +++++++++++- .../api/embeddings/ollama-embeddings.adoc | 32 +++++++++++- .../ollama/OllamaAutoConfiguration.java | 17 +++++-- .../ollama/OllamaChatProperties.java | 3 +- .../ollama/OllamaEmbeddingProperties.java | 3 +- .../OllamaInitializationProperties.java | 50 +++++++++++++++++++ 14 files changed, 177 insertions(+), 19 deletions(-) diff --git a/models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/OllamaChatModel.java b/models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/OllamaChatModel.java index 750546d88..2e9532461 100644 --- a/models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/OllamaChatModel.java +++ b/models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/OllamaChatModel.java @@ -43,6 +43,7 @@ import org.springframework.ai.model.function.FunctionCallback; import org.springframework.ai.model.function.FunctionCallbackContext; import org.springframework.ai.model.function.FunctionCallingOptions; import org.springframework.ai.ollama.api.OllamaApi; +import org.springframework.ai.ollama.api.OllamaModel; import org.springframework.ai.ollama.api.OllamaApi.ChatRequest; import org.springframework.ai.ollama.api.OllamaApi.Message.Role; import org.springframework.ai.ollama.api.OllamaApi.Message.ToolCall; @@ -379,7 +380,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) { - if (!PullModelStrategy.NEVER.equals(pullModelStrategy)) { + if (pullModelStrategy != null && !PullModelStrategy.NEVER.equals(pullModelStrategy)) { this.modelManager.pullModel(model, pullModelStrategy); } } @@ -397,7 +398,7 @@ public class OllamaChatModel extends AbstractToolCallSupport implements ChatMode private OllamaApi ollamaApi; - private OllamaOptions defaultOptions = OllamaOptions.create().withModel(OllamaOptions.DEFAULT_MODEL); + private OllamaOptions defaultOptions = OllamaOptions.create().withModel(OllamaModel.MISTRAL.id()); private FunctionCallbackContext functionCallbackContext; diff --git a/models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/OllamaEmbeddingModel.java b/models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/OllamaEmbeddingModel.java index 6da6714ee..534d1ba33 100644 --- a/models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/OllamaEmbeddingModel.java +++ b/models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/OllamaEmbeddingModel.java @@ -30,6 +30,7 @@ import org.springframework.ai.embedding.observation.EmbeddingModelObservationCon import org.springframework.ai.embedding.observation.EmbeddingModelObservationDocumentation; import org.springframework.ai.model.ModelOptionsUtils; import org.springframework.ai.ollama.api.OllamaApi; +import org.springframework.ai.ollama.api.OllamaModel; import org.springframework.ai.ollama.api.OllamaApi.EmbeddingsResponse; import org.springframework.ai.ollama.management.ModelManagementOptions; import org.springframework.ai.ollama.management.OllamaModelManager; @@ -163,7 +164,7 @@ public class OllamaEmbeddingModel extends AbstractEmbeddingModel { * Pull the given model into Ollama based on the specified strategy. */ private void initializeModelIfEnabled(String model, PullModelStrategy pullModelStrategy) { - if (!PullModelStrategy.NEVER.equals(pullModelStrategy)) { + if (pullModelStrategy != null && !PullModelStrategy.NEVER.equals(pullModelStrategy)) { this.modelManager.pullModel(model, pullModelStrategy); } } @@ -212,7 +213,7 @@ public class OllamaEmbeddingModel extends AbstractEmbeddingModel { private OllamaApi ollamaApi; - private OllamaOptions defaultOptions = OllamaOptions.create().withModel(OllamaOptions.DEFAULT_MODEL); + private OllamaOptions defaultOptions = OllamaOptions.create().withModel(OllamaModel.MXBAI_EMBED_LARGE.id()); private ObservationRegistry observationRegistry = ObservationRegistry.NOOP; diff --git a/models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/api/OllamaModel.java b/models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/api/OllamaModel.java index 5687e2aeb..a70765249 100644 --- a/models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/api/OllamaModel.java +++ b/models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/api/OllamaModel.java @@ -121,7 +121,12 @@ public enum OllamaModel implements ChatModelDescription { /** * A high-performing open embedding model with a large token context window. */ - NOMIC_EMBED_TEXT("nomic-embed-text"); + NOMIC_EMBED_TEXT("nomic-embed-text"), + + /** + * State-of-the-art large embedding model from mixedbread.ai + */ + MXBAI_EMBED_LARGE("mxbai-embed-large"); private final String id; diff --git a/models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/api/OllamaOptions.java b/models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/api/OllamaOptions.java index b25b53102..f7a780668 100644 --- a/models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/api/OllamaOptions.java +++ b/models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/api/OllamaOptions.java @@ -51,8 +51,6 @@ import com.fasterxml.jackson.annotation.JsonProperty; @JsonInclude(Include.NON_NULL) public class OllamaOptions implements FunctionCallingOptions, ChatOptions, EmbeddingOptions { - public static final String DEFAULT_MODEL = OllamaModel.MISTRAL.id(); - private static final List NON_SUPPORTED_FIELDS = List.of("model", "format", "keep_alive", "truncate"); // Following fields are options which must be set when the model is loaded into @@ -309,7 +307,7 @@ public class OllamaOptions implements FunctionCallingOptions, ChatOptions, Embed * Strategy for pulling models at run-time. */ @JsonIgnore - private PullModelStrategy pullModelStrategy = PullModelStrategy.NEVER; + private PullModelStrategy pullModelStrategy; public static OllamaOptions builder() { return new OllamaOptions(); diff --git a/models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/management/ModelManagementOptions.java b/models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/management/ModelManagementOptions.java index 92676d6e8..b49d0978c 100644 --- a/models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/management/ModelManagementOptions.java +++ b/models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/management/ModelManagementOptions.java @@ -16,6 +16,7 @@ package org.springframework.ai.ollama.management; import java.time.Duration; +import java.util.List; /** * Options for managing models in Ollama. @@ -23,8 +24,9 @@ import java.time.Duration; * @author Thomas Vitale * @since 1.0.0 */ -public record ModelManagementOptions(PullModelStrategy pullModelStrategy, Duration timeout, Integer maxRetries) { +public record ModelManagementOptions(PullModelStrategy pullModelStrategy, List additionalModels, + Duration timeout, Integer maxRetries) { public static ModelManagementOptions defaults() { - return new ModelManagementOptions(PullModelStrategy.NEVER, Duration.ofMinutes(5), 0); + return new ModelManagementOptions(PullModelStrategy.NEVER, List.of(), Duration.ofMinutes(5), 0); } } 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 01d444b94..31e2249b8 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 @@ -48,6 +48,10 @@ public class OllamaModelManager { public OllamaModelManager(OllamaApi ollamaApi, ModelManagementOptions options) { this.ollamaApi = ollamaApi; this.options = options; + + if (!CollectionUtils.isEmpty(options.additionalModels())) { + options.additionalModels().forEach(this::pullModel); + } } public boolean isModelAvailable(String 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 f3d4ef07a..edb438072 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 @@ -3,6 +3,7 @@ package org.springframework.ai.ollama; import org.springframework.ai.ollama.api.OllamaApi; import org.springframework.ai.ollama.management.OllamaModelManager; import org.springframework.ai.ollama.management.PullModelStrategy; +import org.springframework.util.StringUtils; import org.testcontainers.ollama.OllamaContainer; public class BaseOllamaIT { @@ -31,6 +32,10 @@ public class BaseOllamaIT { return false; } + public static OllamaApi buildOllamaApi() { + return buildOllamaApiWithModel(null); + } + public static OllamaApi buildOllamaApiWithModel(String model) { var baseUrl = "http://localhost:11434"; if (useTestcontainers) { @@ -38,7 +43,9 @@ public class BaseOllamaIT { } var ollamaApi = new OllamaApi(baseUrl); - ensureModelIsPresent(ollamaApi, model); + if (StringUtils.hasText(model)) { + ensureModelIsPresent(ollamaApi, model); + } return ollamaApi; } 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 dac65820f..52d19e6d8 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 @@ -23,6 +23,8 @@ import org.springframework.ai.ollama.api.OllamaModel; import org.testcontainers.junit.jupiter.Testcontainers; import java.io.IOException; +import java.time.Duration; +import java.util.List; import static org.assertj.core.api.Assertions.assertThat; @@ -81,4 +83,21 @@ class OllamaModelManagerIT extends BaseOllamaIT { assertThat(isModelAvailable).isFalse(); } + @Test + public void pullAdditionalModels() { + var model = "all-minilm"; + var isModelAvailable = modelManager.isModelAvailable(model); + assertThat(isModelAvailable).isFalse(); + + new OllamaModelManager(buildOllamaApi(), + new ModelManagementOptions(PullModelStrategy.WHEN_MISSING, List.of(model), Duration.ofMinutes(5), 0)); + + isModelAvailable = modelManager.isModelAvailable(model); + assertThat(isModelAvailable).isTrue(); + + modelManager.deleteModel(model); + isModelAvailable = modelManager.isModelAvailable(model); + assertThat(isModelAvailable).isFalse(); + } + } diff --git a/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/chat/ollama-chat.adoc b/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/chat/ollama-chat.adoc index 15df628ad..e78b8439d 100644 --- a/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/chat/ollama-chat.adoc +++ b/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/chat/ollama-chat.adoc @@ -72,6 +72,8 @@ Here are the properties for initializing the Ollama integration and xref:auto-pu | spring.ai.ollama.init.pull-model-strategy | Whether to pull models at startup-time and how. | `never` | spring.ai.ollama.init.timeout | How long to wait for a model to be pulled. | `5m` | spring.ai.ollama.init.max-retries | Maximum number of retries for the model pull operation. | `0` +| spring.ai.ollama.init.chat.include | Include this type of models in the initialization task. | `true` +| spring.ai.ollama.init.chat.additional-models | Additional models to initialize besides the ones configured via default properties. | `[]` |==== === Chat Properties @@ -188,6 +190,34 @@ spring: CAUTION: The application will not complete its initialization until all the models become available in Ollama. Depending on the model size and the speed of the Internet connection, your application might be slow at starting up. +You can also initialize additional models at startup time, useful for those models used dynamically at runtime. + +[source,yaml] +---- +spring: + ai: + ollama: + init: + pull-model-strategy: always + chat: + additional-models: + - llama3.2 + - qwen2.5 +---- + +If you want to apply the pulling strategy only to other types of models, you can exclude the chat models from the initialization task. + +[source,yaml] +---- +spring: + ai: + ollama: + init: + pull-model-strategy: always + chat: + include: false +---- + === Pulling models at runtime To enable auto-pulling of models at runtime, you can configure the `pullModelStrategy` option in your `OllamaOptions`: @@ -205,7 +235,7 @@ ChatResponse response = chatModel.call(new Prompt( You can also configure this option using the following property: `spring.ai.ollama.chat.options.pull-model-strategy=always`. -CAUTION: The time to process an incoming request might incur unexpected delays, waiting for the needed model to become available in Ollama. Depending on the model size and the speed of the Internet connection, your application might be slow at processing requests. +CAUTION: The time to process an incoming request might incur unexpected delays, waiting for the needed model to become available in Ollama. Depending on the model size and the speed of the Internet connection, your application might be slow at processing requests. You might want to initialize these models at startup time instead, using the `spring.ai.ollama.init.chat.additional-models` property. == Function Calling diff --git a/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/embeddings/ollama-embeddings.adoc b/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/embeddings/ollama-embeddings.adoc index 872837360..e414df625 100644 --- a/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/embeddings/ollama-embeddings.adoc +++ b/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/embeddings/ollama-embeddings.adoc @@ -76,6 +76,8 @@ Here are the properties for initializing the Ollama integration and xref:auto-pu | spring.ai.ollama.init.pull-model-strategy | Whether to pull models at startup-time and how. | `never` | spring.ai.ollama.init.timeout | How long to wait for a model to be pulled. | `5m` | spring.ai.ollama.init.max-retries | Maximum number of retries for the model pull operation. | `0` +| spring.ai.ollama.init.embedding.include | Include this type of models in the initialization task. | `true` +| spring.ai.ollama.init.embedding.additional-models | Additional models to initialize besides the ones configured via default properties. | `[]` |==== === Embedding Properties @@ -190,6 +192,34 @@ spring: CAUTION: The application will not complete its initialization until all the models become available in Ollama. Depending on the model size and the speed of the Internet connection, your application might be slow at starting up. +You can also initialize additional models at startup time, useful for those models used dynamically at runtime. + +[source,yaml] +---- +spring: + ai: + ollama: + init: + pull-model-strategy: always + embedding: + additional-models: + - mxbai-embed-large + - nomic-embed-text +---- + +If you want to apply the pulling strategy only to other types of models, you can exclude the embedding models from the initialization task. + +[source,yaml] +---- +spring: + ai: + ollama: + init: + pull-model-strategy: always + embedding: + include: false +---- + === Pulling models at runtime To enable auto-pulling of models at runtime, you can configure the `pullModelStrategy` option in your `OllamaOptions`: @@ -206,7 +236,7 @@ EmbeddingResponse embeddingResponse = embeddingModel You can also configure this option using the following property: `spring.ai.ollama.embedding.options.pull-model-strategy=always`. -CAUTION: The time to process an incoming request might incur unexpected delays, waiting for the needed model to become available in Ollama. Depending on the model size and the speed of the Internet connection, your application might be slow at processing requests. +CAUTION: The time to process an incoming request might incur unexpected delays, waiting for the needed model to become available in Ollama. Depending on the model size and the speed of the Internet connection, your application might be slow at processing requests. You might want to initialize these models at startup time instead, using the `spring.ai.ollama.init.embedding.additional-models` property. == Sample Controller diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/ollama/OllamaAutoConfiguration.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/ollama/OllamaAutoConfiguration.java index 46f6a292d..453c23726 100644 --- a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/ollama/OllamaAutoConfiguration.java +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/ollama/OllamaAutoConfiguration.java @@ -25,6 +25,7 @@ import org.springframework.ai.ollama.OllamaChatModel; import org.springframework.ai.ollama.OllamaEmbeddingModel; import org.springframework.ai.ollama.api.OllamaApi; import org.springframework.ai.ollama.management.ModelManagementOptions; +import org.springframework.ai.ollama.management.PullModelStrategy; import org.springframework.beans.factory.ObjectProvider; import org.springframework.boot.autoconfigure.AutoConfiguration; import org.springframework.boot.autoconfigure.ImportAutoConfiguration; @@ -80,14 +81,18 @@ public class OllamaAutoConfiguration { OllamaInitializationProperties initProperties, List toolFunctionCallbacks, FunctionCallbackContext functionCallbackContext, ObjectProvider observationRegistry, ObjectProvider observationConvention) { + var chatModelPullStrategy = initProperties.getChat().isInclude() ? initProperties.getPullModelStrategy() + : PullModelStrategy.NEVER; + var chatModel = OllamaChatModel.builder() .withOllamaApi(ollamaApi) .withDefaultOptions(properties.getOptions()) .withFunctionCallbackContext(functionCallbackContext) .withToolFunctionCallbacks(toolFunctionCallbacks) .withObservationRegistry(observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP)) - .withModelManagementOptions(new ModelManagementOptions(initProperties.getPullModelStrategy(), - initProperties.getTimeout(), initProperties.getMaxRetries())) + .withModelManagementOptions( + new ModelManagementOptions(chatModelPullStrategy, initProperties.getChat().getAdditionalModels(), + initProperties.getTimeout(), initProperties.getMaxRetries())) .build(); observationConvention.ifAvailable(chatModel::setObservationConvention); @@ -102,12 +107,16 @@ public class OllamaAutoConfiguration { public OllamaEmbeddingModel ollamaEmbeddingModel(OllamaApi ollamaApi, OllamaEmbeddingProperties properties, OllamaInitializationProperties initProperties, ObjectProvider observationRegistry, ObjectProvider observationConvention) { + var embeddingModelPullStrategy = initProperties.getEmbedding().isInclude() + ? initProperties.getPullModelStrategy() : PullModelStrategy.NEVER; + var embeddingModel = OllamaEmbeddingModel.builder() .withOllamaApi(ollamaApi) .withDefaultOptions(properties.getOptions()) .withObservationRegistry(observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP)) - .withModelManagementOptions(new ModelManagementOptions(initProperties.getPullModelStrategy(), - initProperties.getTimeout(), initProperties.getMaxRetries())) + .withModelManagementOptions(new ModelManagementOptions(embeddingModelPullStrategy, + initProperties.getEmbedding().getAdditionalModels(), initProperties.getTimeout(), + initProperties.getMaxRetries())) .build(); observationConvention.ifAvailable(embeddingModel::setObservationConvention); diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/ollama/OllamaChatProperties.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/ollama/OllamaChatProperties.java index 2439c83e8..c106c8c35 100644 --- a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/ollama/OllamaChatProperties.java +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/ollama/OllamaChatProperties.java @@ -15,6 +15,7 @@ */ package org.springframework.ai.autoconfigure.ollama; +import org.springframework.ai.ollama.api.OllamaModel; import org.springframework.ai.ollama.api.OllamaOptions; import org.springframework.boot.context.properties.ConfigurationProperties; import org.springframework.boot.context.properties.NestedConfigurationProperty; @@ -41,7 +42,7 @@ public class OllamaChatProperties { * generative's defaults. */ @NestedConfigurationProperty - private OllamaOptions options = OllamaOptions.create().withModel(OllamaOptions.DEFAULT_MODEL); + private OllamaOptions options = OllamaOptions.create().withModel(OllamaModel.MISTRAL.id()); public String getModel() { return this.options.getModel(); diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/ollama/OllamaEmbeddingProperties.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/ollama/OllamaEmbeddingProperties.java index a2368cd2e..9b21a92d5 100644 --- a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/ollama/OllamaEmbeddingProperties.java +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/ollama/OllamaEmbeddingProperties.java @@ -15,6 +15,7 @@ */ package org.springframework.ai.autoconfigure.ollama; +import org.springframework.ai.ollama.api.OllamaModel; import org.springframework.ai.ollama.api.OllamaOptions; import org.springframework.boot.context.properties.ConfigurationProperties; import org.springframework.boot.context.properties.NestedConfigurationProperty; @@ -41,7 +42,7 @@ public class OllamaEmbeddingProperties { * generative's defaults. */ @NestedConfigurationProperty - private OllamaOptions options = OllamaOptions.create().withModel(OllamaOptions.DEFAULT_MODEL); + private OllamaOptions options = OllamaOptions.create().withModel(OllamaModel.MXBAI_EMBED_LARGE.id()); public String getModel() { return this.options.getModel(); diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/ollama/OllamaInitializationProperties.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/ollama/OllamaInitializationProperties.java index 572b1b442..b884404be 100644 --- a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/ollama/OllamaInitializationProperties.java +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/ollama/OllamaInitializationProperties.java @@ -19,6 +19,7 @@ import org.springframework.ai.ollama.management.PullModelStrategy; import org.springframework.boot.context.properties.ConfigurationProperties; import java.time.Duration; +import java.util.List; /** * Ollama initialization configuration properties. @@ -36,6 +37,16 @@ public class OllamaInitializationProperties { */ private PullModelStrategy pullModelStrategy = PullModelStrategy.NEVER; + /** + * Chat models initialization settings. + */ + private final ModelTypeInit chat = new ModelTypeInit(); + + /** + * Embedding models initialization settings. + */ + private final ModelTypeInit embedding = new ModelTypeInit(); + /** * How long to wait for a model to be pulled. */ @@ -54,6 +65,14 @@ public class OllamaInitializationProperties { this.pullModelStrategy = pullModelStrategy; } + public ModelTypeInit getChat() { + return chat; + } + + public ModelTypeInit getEmbedding() { + return embedding; + } + public Duration getTimeout() { return timeout; } @@ -70,4 +89,35 @@ public class OllamaInitializationProperties { this.maxRetries = maxRetries; } + public static class ModelTypeInit { + + /** + * Include this type of models in the initialization task. + */ + private boolean include = true; + + /** + * Additional models to initialize besides the ones configured via default + * properties. + */ + private List additionalModels = List.of(); + + public boolean isInclude() { + return include; + } + + public void setInclude(boolean include) { + this.include = include; + } + + public List getAdditionalModels() { + return additionalModels; + } + + public void setAdditionalModels(List additionalModels) { + this.additionalModels = additionalModels; + } + + } + }