Ollama: Added missing fields in API

Fixes gh-553 Signed-off-by: Thomas Vitale <[email protected]>
spring-projects · Apr 6, 2024 · 4e473aa · 4e473aa
1 parent 11d0578
commit 4e473aa
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 8 deletions.
diff --git a/...s/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/OllamaEmbeddingClient.java b/...s/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/OllamaEmbeddingClient.java
@@ -135,7 +135,8 @@ OllamaApi.EmbeddingRequest ollamaEmbeddingRequest(String inputContent, Embedding
 			throw new IllegalArgumentException("Model is not set!");
 		}
 		String model = mergedOptions.getModel();
-		return new EmbeddingRequest(model, inputContent, OllamaOptions.filterNonSupportedFields(mergedOptions.toMap()));
+		return new EmbeddingRequest(model, inputContent, null,
+				OllamaOptions.filterNonSupportedFields(mergedOptions.toMap()));
 	}
 
 }
diff --git a/models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/api/OllamaApi.java b/models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/api/OllamaApi.java
@@ -150,7 +150,9 @@ public record GenerateRequest(
 			@JsonProperty("template") String template,
 			@JsonProperty("context") List<Integer> context,
 			@JsonProperty("stream") Boolean stream,
-			@JsonProperty("raw") Boolean raw) {
+			@JsonProperty("raw") Boolean raw,
+			@JsonProperty("images") List<String> images,
+			@JsonProperty("keep_alive") String keepAlive) {
 
 		/**
 		 * Short cut constructor to create a CompletionRequest without options.
@@ -159,7 +161,7 @@ public record GenerateRequest(
 		 * @param stream Whether to stream the response.
 		 */
 		public GenerateRequest(String model, String prompt, Boolean stream) {
-			this(model, prompt, null, null, null, null, null, stream, null);
+			this(model, prompt, null, null, null, null, null, stream, null, null, null);
 		}
 
 		/**
@@ -170,7 +172,7 @@ public GenerateRequest(String model, String prompt, Boolean stream) {
 		 * @param stream Whether to stream the response.
 		 */
 		public GenerateRequest(String model, String prompt, boolean enableJsonFormat, Boolean stream) {
-			this(model, prompt, (enableJsonFormat) ? "json" : null, null, null, null, null, stream, null);
+			this(model, prompt, (enableJsonFormat) ? "json" : null, null, null, null, null, stream, null, null, null);
 		}
 
 		/**
@@ -192,6 +194,8 @@ public static class Builder {
 			private List<Integer> context;
 			private Boolean stream;
 			private Boolean raw;
+			private List<String> images;
+			private String keepAlive;
 
 			public Builder(String prompt) {
 				this.prompt = prompt;
@@ -242,8 +246,18 @@ public Builder withRaw(Boolean raw) {
 				return this;
 			}
 
+			public Builder withImages(List<String> images) {
+				this.images = images;
+				return this;
+			}
+
+			public Builder withKeepAlive(String keepAlive) {
+				this.keepAlive = keepAlive;
+				return this;
+			}
+
 			public GenerateRequest build() {
-				return new GenerateRequest(model, prompt, format, options, system, template, context, stream, raw);
+				return new GenerateRequest(model, prompt, format, options, system, template, context, stream, raw, images, keepAlive);
 			}
 
 		}
@@ -462,14 +476,14 @@ public Builder withTemplate(String template) {
 			}
 
 			public Builder withOptions(Map<String, Object> options) {
-				Objects.requireNonNullElse(options, "The options can not be null.");
+				Objects.requireNonNull(options, "The options can not be null.");
 
 				this.options = OllamaOptions.filterNonSupportedFields(options);
 				return this;
 			}
 
 			public Builder withOptions(OllamaOptions options) {
-				Objects.requireNonNullElse(options, "The options can not be null.");
+				Objects.requireNonNull(options, "The options can not be null.");
 				this.options = OllamaOptions.filterNonSupportedFields(options.toMap());
 				return this;
 			}
@@ -574,6 +588,7 @@ public Flux<ChatResponse> streamingChat(ChatRequest chatRequest) {
 	public record EmbeddingRequest(
 			@JsonProperty("model") String model,
 			@JsonProperty("prompt") String prompt,
+			@JsonProperty("keep_alive") Duration keepAlive,
 			@JsonProperty("options") Map<String, Object> options) {
 
 		/**
@@ -582,7 +597,7 @@ public record EmbeddingRequest(
 		 * @param prompt The text to generate embeddings for.
 		 */
 		public EmbeddingRequest(String model, String prompt) {
-			this(model, prompt, null);
+			this(model, prompt, null, null);
 		}
 	}