From bf3a8765d31b6ec491bf7ce34e3cee52631717d9 Mon Sep 17 00:00:00 2001
From: sblashuk <siarhei.blashuk@outlook.com>
Date: Fri, 1 Mar 2024 17:33:00 +0100
Subject: [PATCH] #311 Create ollama enum with supported models POC

---
 .../ai/ollama/api/OllamaModel.java            | 67 +++++++++++++++++++
 .../ai/ollama/api/OllamaOptions.java          |  5 +-
 2 files changed, 71 insertions(+), 1 deletion(-)
 create mode 100644 models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/api/OllamaModel.java

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
new file mode 100644
index 00000000000..e455ddf0d9f
--- /dev/null
+++ b/models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/api/OllamaModel.java
@@ -0,0 +1,67 @@
+package org.springframework.ai.ollama.api;
+
+/**
+ * Helper class for common Ollama models.
+ *
+ * @author Siarhei Blashuk
+ * @since 0.8.1
+ */
+public enum OllamaModel {
+
+	/**
+	 * Llama 2 is a collection of language models ranging from 7B to 70B parameters.
+	 */
+	LLAMA2("llama2"),
+	/**
+	 * The 7B parameters model
+	 */
+	MISTRAL("mistral"),
+	/**
+	 * The 2.7B uncensored Dolphin model
+	 */
+	DOLPHIN_PHI("dolphin-phi"),
+	/**
+	 * The Phi-2 2.7B language model
+	 */
+	PHI("phi"),
+	/**
+	 * A fine-tuned Mistral model
+	 */
+	NEURAL_CHAT("neural-chat"),
+	/**
+	 * Starling-7B model
+	 */
+	STARLING_LM("starling-lm"),
+	/**
+	 * Code Llama is based on Llama 2 model
+	 */
+	CODELLAMA("codellama"),
+	/**
+	 * Orca Mini is based on Llama and Llama 2 ranging from 3 billion parameters to 70
+	 * billion
+	 */
+	ORCA_MINI("orca-mini"),
+	/**
+	 * Llava is a Large Language and Vision Assistant model
+	 */
+	LLAVA("llava"),
+	/**
+	 * Gemma is a lightweight model with 2 billion and 7 billion
+	 */
+	GEMMA("gemma"),
+	/**
+	 * Uncensored Llama 2 model
+	 */
+	LLAMA2_UNCENSORED("llama2-uncensored");
+
+	private final String id;
+
+	OllamaModel(String id) {
+		this.id = id;
+	}
+
+	public String id() {
+		return 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 a6dcc8b67bd..842351e3480 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
@@ -44,7 +44,7 @@
 @JsonInclude(Include.NON_NULL)
 public class OllamaOptions implements ChatOptions, EmbeddingOptions {
 
-	public static final String DEFAULT_MODEL = "mistral";
+	public static final String DEFAULT_MODEL = OllamaModel.MISTRAL.id();
 
 	// @formatter:off
 	/**
@@ -244,6 +244,9 @@ public class OllamaOptions implements ChatOptions, EmbeddingOptions {
 	 */
 	@JsonProperty("model") private String model;
 
+	/**
+	 * @param model The ollama model names to use. See the {@link OllamaModel} for the common models.
+	 */
 	public OllamaOptions withModel(String model) {
 		this.model = model;
 		return this;