From 8ed2af41154ac877de3c649255a15214f152e2a4 Mon Sep 17 00:00:00 2001
From: sblashuk <siarhei.blashuk@outlook.com>
Date: Fri, 1 Mar 2024 17:33:00 +0100
Subject: [PATCH] Add Ollama enum with supported models and their ids

 - Add missing license header
---
 .../ai/ollama/api/OllamaModel.java            | 92 +++++++++++++++++++
 .../ai/ollama/api/OllamaOptions.java          |  5 +-
 2 files changed, 96 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 000000000..bdb7ee1b0
--- /dev/null
+++ b/models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/api/OllamaModel.java
@@ -0,0 +1,92 @@
+/*
+ * 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.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 this.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 786dfc4ba..ac060d9f2 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
@@ -43,7 +43,7 @@ import org.springframework.ai.embedding.EmbeddingOptions;
 @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
 	/**
@@ -243,6 +243,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;