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 3419ee28332..0f30db20799 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 @@ -30,118 +30,159 @@ public enum OllamaModel implements ChatModelDescription { /** * Llama 2 is a collection of language models ranging from 7B to 70B parameters. */ - LLAMA2("llama2"), + LLAMA2("llama2", null), /** * Llama 3 is a collection of language models ranging from 8B and 70B parameters. */ - LLAMA3("llama3"), + LLAMA3("llama3", null), /** * The 8B language model from Meta. */ - LLAMA3_1("llama3.1"), + LLAMA3_1("llama3.1", null), /** * The Llama 3.2 3B language model from Meta. */ - LLAMA3_2("llama3.2"), + LLAMA3_2("llama3.2", null), /** * The Llama 3.2 1B language model from Meta. */ - LLAMA3_2_1B("llama3.2:1b"), + LLAMA3_2_1B("llama3.2", "1b"), /** * The 7B parameters model */ - MISTRAL("mistral"), + MISTRAL("mistral", null), /** * A 12B model with 128k context length, built by Mistral AI in collaboration with * NVIDIA. */ - MISTRAL_NEMO("mistral-nemo"), + MISTRAL_NEMO("mistral-nemo", null), /** * A small vision language model designed to run efficiently on edge devices. */ - MOONDREAM("moondream"), + MOONDREAM("moondream", null), /** * The 2.7B uncensored Dolphin model */ - DOLPHIN_PHI("dolphin-phi"), + DOLPHIN_PHI("dolphin-phi", null), /** * The Phi-2 2.7B language model */ - PHI("phi"), + PHI("phi", null), /** * The Phi-3 3.8B language model */ - PHI3("phi3"), + PHI3("phi3", null), /** * A fine-tuned Mistral model */ - NEURAL_CHAT("neural-chat"), + NEURAL_CHAT("neural-chat", null), /** * Starling-7B model */ - STARLING_LM("starling-lm"), + STARLING_LM("starling-lm", null), /** * Code Llama is based on Llama 2 model */ - CODELLAMA("codellama"), + CODELLAMA("codellama", null), /** * Orca Mini is based on Llama and Llama 2 ranging from 3 billion parameters to 70 * billion */ - ORCA_MINI("orca-mini"), + ORCA_MINI("orca-mini", null), /** * Llava is a Large Language and Vision Assistant model */ - LLAVA("llava"), + LLAVA("llava", null), /** * Gemma is a lightweight model with 2 billion and 7 billion */ - GEMMA("gemma"), + GEMMA("gemma", null), /** * Uncensored Llama 2 model */ - LLAMA2_UNCENSORED("llama2-uncensored"), + LLAMA2_UNCENSORED("llama2-uncensored", null), /** * A high-performing open embedding model with a large token context window. */ - NOMIC_EMBED_TEXT("nomic-embed-text"), + NOMIC_EMBED_TEXT("nomic-embed-text", null), /** * State-of-the-art large embedding model from mixedbread.ai */ - MXBAI_EMBED_LARGE("mxbai-embed-large"); + MXBAI_EMBED_LARGE("mxbai-embed-large", null), + /** + * Qwen2.5 models are pretrained on Alibaba's latest large-scale dataset, encompassing + * up to 18 trillion tokens. The model supports up to 128K tokens and has multilingual + * support. + */ + QWEN2_5("qwen2.5", null), QWEN2_5_7b("qwen2.5", "7b"), QWEN2_5_14b("qwen2.5", "14b"), QWEN2_5_32b("qwen2.5", "32b"), + QWEN2_5_72b("qwen2.5", "72b"), + /** + * The latest series of Code-Specific Qwen models, with significant improvements in + * code generation, code reasoning, and code fixing. + */ + QWEN2_5_CODER("qwen2.5-coder", null), QWEN2_5_CODER_3b("qwen2.5-coder", "3b"), + QWEN2_5_CODER_7b("qwen2.5-coder", "7b"), QWEN2_5_CODER_14b("qwen2.5-coder", "14b"), + QWEN2_5_CODER_32b("qwen2.5-coder", "32b"), + /** + * Llama 3.2 Vision is a collection of instruction-tuned image reasoning generative + * models in 11B and 90B sizes. + */ + LLAMA3_2_VISION("llama3.2-vision", null), LLAMA3_2_VISION_11b("llama3.2-vision", "11b"), + LLAMA3_2_VISION_90b("llama3.2-vision", "90b"); + /** + * model id + */ private final String id; - OllamaModel(String id) { + /** + * model size + */ + private final String size; + + OllamaModel(String id, String size) { this.id = id; + this.size = size; } public String id() { return this.id; } + /** + * for example: qwen2.5:7b,llama3.2:1b + * @return model id with size + */ @Override public String getName() { - return this.id; + return this.id + (this.size == null || this.size.isEmpty() ? "" : ":" + this.size); + } + + /** + * @return model size + */ + @Override + public String getModelSize() { + return this.size; } } diff --git a/models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/OllamaModelTests.java b/models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/OllamaModelTests.java new file mode 100644 index 00000000000..ba90bcce4e5 --- /dev/null +++ b/models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/OllamaModelTests.java @@ -0,0 +1,32 @@ +/* + * Copyright 2023-2024 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.ai.ollama; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.springframework.ai.ollama.api.OllamaModel; + +class OllamaModelTests extends BaseOllamaIT { + + @Test + void autoPullModelTest() { + Assertions.assertEquals(OllamaModel.LLAMA3_2_VISION.getName(), "llama3.2-vision"); + Assertions.assertEquals(OllamaModel.LLAMA3_2_VISION_11b.getName(), "llama3.2-vision:11b"); + + } + +} diff --git a/spring-ai-core/src/main/java/org/springframework/ai/model/ModelDescription.java b/spring-ai-core/src/main/java/org/springframework/ai/model/ModelDescription.java index e1d9ec2b64d..394af8ae272 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/model/ModelDescription.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/model/ModelDescription.java @@ -47,4 +47,12 @@ default String getVersion() { return ""; } + /** + * Returns the size of the model + * @return + */ + default String getModelSize() { + return null; + } + }