Skip to content
Open
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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 + (null == this.size ? "" : ":" + this.size);

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Missing Null Handling in getName:

The logic in getName assumes size will be null or non-empty but doesn't guard against unexpected input like an empty string.
Add validation to handle this edge case

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If there is no difference between id and name, then that's right

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
return this.id + (null == this.size ? "" : ":" + this.size);
return this.id + (this.size == null || this.size.isEmpty() ? "" : ":" + this.size);

}

/**
* @return model size
*/
@Override
public String getModelSize() {
return this.size;
}

}
Original file line number Diff line number Diff line change
@@ -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");

}

}
Original file line number Diff line number Diff line change
Expand Up @@ -47,4 +47,12 @@ default String getVersion() {
return "";
}

/**
* Returns the size of the model
* @return
*/
default String getModelSize() {
return null;
}

}