Skip to content

Commit 4e473aa

Browse files
ThomasVitaletzolov
authored andcommitted
Ollama: Added missing fields in API
Fixes gh-553 Signed-off-by: Thomas Vitale <ThomasVitale@users.noreply.github.com>
1 parent 11d0578 commit 4e473aa

File tree

2 files changed

+24
-8
lines changed

2 files changed

+24
-8
lines changed

models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/OllamaEmbeddingClient.java

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -135,7 +135,8 @@ OllamaApi.EmbeddingRequest ollamaEmbeddingRequest(String inputContent, Embedding
135135
throw new IllegalArgumentException("Model is not set!");
136136
}
137137
String model = mergedOptions.getModel();
138-
return new EmbeddingRequest(model, inputContent, OllamaOptions.filterNonSupportedFields(mergedOptions.toMap()));
138+
return new EmbeddingRequest(model, inputContent, null,
139+
OllamaOptions.filterNonSupportedFields(mergedOptions.toMap()));
139140
}
140141

141142
}

models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/api/OllamaApi.java

Lines changed: 22 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -150,7 +150,9 @@ public record GenerateRequest(
150150
@JsonProperty("template") String template,
151151
@JsonProperty("context") List<Integer> context,
152152
@JsonProperty("stream") Boolean stream,
153-
@JsonProperty("raw") Boolean raw) {
153+
@JsonProperty("raw") Boolean raw,
154+
@JsonProperty("images") List<String> images,
155+
@JsonProperty("keep_alive") String keepAlive) {
154156

155157
/**
156158
* Short cut constructor to create a CompletionRequest without options.
@@ -159,7 +161,7 @@ public record GenerateRequest(
159161
* @param stream Whether to stream the response.
160162
*/
161163
public GenerateRequest(String model, String prompt, Boolean stream) {
162-
this(model, prompt, null, null, null, null, null, stream, null);
164+
this(model, prompt, null, null, null, null, null, stream, null, null, null);
163165
}
164166

165167
/**
@@ -170,7 +172,7 @@ public GenerateRequest(String model, String prompt, Boolean stream) {
170172
* @param stream Whether to stream the response.
171173
*/
172174
public GenerateRequest(String model, String prompt, boolean enableJsonFormat, Boolean stream) {
173-
this(model, prompt, (enableJsonFormat) ? "json" : null, null, null, null, null, stream, null);
175+
this(model, prompt, (enableJsonFormat) ? "json" : null, null, null, null, null, stream, null, null, null);
174176
}
175177

176178
/**
@@ -192,6 +194,8 @@ public static class Builder {
192194
private List<Integer> context;
193195
private Boolean stream;
194196
private Boolean raw;
197+
private List<String> images;
198+
private String keepAlive;
195199

196200
public Builder(String prompt) {
197201
this.prompt = prompt;
@@ -242,8 +246,18 @@ public Builder withRaw(Boolean raw) {
242246
return this;
243247
}
244248

249+
public Builder withImages(List<String> images) {
250+
this.images = images;
251+
return this;
252+
}
253+
254+
public Builder withKeepAlive(String keepAlive) {
255+
this.keepAlive = keepAlive;
256+
return this;
257+
}
258+
245259
public GenerateRequest build() {
246-
return new GenerateRequest(model, prompt, format, options, system, template, context, stream, raw);
260+
return new GenerateRequest(model, prompt, format, options, system, template, context, stream, raw, images, keepAlive);
247261
}
248262

249263
}
@@ -462,14 +476,14 @@ public Builder withTemplate(String template) {
462476
}
463477

464478
public Builder withOptions(Map<String, Object> options) {
465-
Objects.requireNonNullElse(options, "The options can not be null.");
479+
Objects.requireNonNull(options, "The options can not be null.");
466480

467481
this.options = OllamaOptions.filterNonSupportedFields(options);
468482
return this;
469483
}
470484

471485
public Builder withOptions(OllamaOptions options) {
472-
Objects.requireNonNullElse(options, "The options can not be null.");
486+
Objects.requireNonNull(options, "The options can not be null.");
473487
this.options = OllamaOptions.filterNonSupportedFields(options.toMap());
474488
return this;
475489
}
@@ -574,6 +588,7 @@ public Flux<ChatResponse> streamingChat(ChatRequest chatRequest) {
574588
public record EmbeddingRequest(
575589
@JsonProperty("model") String model,
576590
@JsonProperty("prompt") String prompt,
591+
@JsonProperty("keep_alive") Duration keepAlive,
577592
@JsonProperty("options") Map<String, Object> options) {
578593

579594
/**
@@ -582,7 +597,7 @@ public record EmbeddingRequest(
582597
* @param prompt The text to generate embeddings for.
583598
*/
584599
public EmbeddingRequest(String model, String prompt) {
585-
this(model, prompt, null);
600+
this(model, prompt, null, null);
586601
}
587602
}
588603

0 commit comments

Comments
 (0)