Skip to content

Commit f32b122

Browse files
leehautilayaperumalg
leehaut
authored and
ilayaperumalg
committed
Ollama Embedding keep_alive format
 - Use String instead of Duration for keep_alive embedding property - Update tests Signed-off-by: lance <[email protected]>
1 parent 3f79373 commit f32b122

File tree

4 files changed

+32
-68
lines changed

4 files changed

+32
-68
lines changed

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

Lines changed: 2 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -16,12 +16,9 @@
1616

1717
package org.springframework.ai.ollama;
1818

19-
import java.time.Duration;
2019
import java.util.List;
2120
import java.util.Optional;
2221
import java.util.concurrent.atomic.AtomicInteger;
23-
import java.util.regex.Matcher;
24-
import java.util.regex.Pattern;
2522

2623
import io.micrometer.observation.ObservationRegistry;
2724

@@ -178,8 +175,8 @@ OllamaApi.EmbeddingsRequest ollamaEmbeddingRequest(EmbeddingRequest embeddingReq
178175
}
179176

180177
return new OllamaApi.EmbeddingsRequest(requestOptions.getModel(), embeddingRequest.getInstructions(),
181-
DurationParser.parse(requestOptions.getKeepAlive()),
182-
OllamaEmbeddingOptions.filterNonSupportedFields(requestOptions.toMap()), requestOptions.getTruncate());
178+
requestOptions.getKeepAlive(), OllamaEmbeddingOptions.filterNonSupportedFields(requestOptions.toMap()),
179+
requestOptions.getTruncate());
183180
}
184181

185182
/**
@@ -200,37 +197,6 @@ public void setObservationConvention(EmbeddingModelObservationConvention observa
200197
this.observationConvention = observationConvention;
201198
}
202199

203-
public static class DurationParser {
204-
205-
private static final Pattern PATTERN = Pattern.compile("(-?\\d+)(ms|s|m|h)");
206-
207-
public static Duration parse(String input) {
208-
209-
if (!StringUtils.hasText(input)) {
210-
return null;
211-
}
212-
213-
Matcher matcher = PATTERN.matcher(input);
214-
215-
if (matcher.matches()) {
216-
long value = Long.parseLong(matcher.group(1));
217-
String unit = matcher.group(2);
218-
219-
return switch (unit) {
220-
case "ms" -> Duration.ofMillis(value);
221-
case "s" -> Duration.ofSeconds(value);
222-
case "m" -> Duration.ofMinutes(value);
223-
case "h" -> Duration.ofHours(value);
224-
default -> throw new IllegalArgumentException("Unsupported time unit: " + unit);
225-
};
226-
}
227-
else {
228-
throw new IllegalArgumentException("Invalid duration format: " + input);
229-
}
230-
}
231-
232-
}
233-
234200
public static final class Builder {
235201

236202
private OllamaApi ollamaApi;

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -620,7 +620,7 @@ public Duration getEvalDuration() {
620620
public record EmbeddingsRequest(
621621
@JsonProperty("model") String model,
622622
@JsonProperty("input") List<String> input,
623-
@JsonProperty("keep_alive") Duration keepAlive,
623+
@JsonProperty("keep_alive") String keepAlive,
624624
@JsonProperty("options") Map<String, Object> options,
625625
@JsonProperty("truncate") Boolean truncate) {
626626

‎models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/OllamaEmbeddingModelTests.java‎

Lines changed: 11 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,6 @@
1616

1717
package org.springframework.ai.ollama;
1818

19-
import java.time.Duration;
2019
import java.util.List;
2120
import java.util.Map;
2221

@@ -46,7 +45,7 @@
4645
* @since 1.0.0
4746
*/
4847
@ExtendWith(MockitoExtension.class)
49-
public class OllamaEmbeddingModelTests {
48+
class OllamaEmbeddingModelTests {
5049

5150
@Mock
5251
OllamaApi ollamaApi;
@@ -55,7 +54,7 @@ public class OllamaEmbeddingModelTests {
5554
ArgumentCaptor<EmbeddingsRequest> embeddingsRequestCaptor;
5655

5756
@Test
58-
public void options() {
57+
void options() {
5958

6059
given(this.ollamaApi.embed(this.embeddingsRequestCaptor.capture()))
6160
.willReturn(new EmbeddingsResponse("RESPONSE_MODEL_NAME",
@@ -109,7 +108,7 @@ public void options() {
109108
}
110109

111110
@Test
112-
public void singleInputEmbedding() {
111+
void singleInputEmbedding() {
113112
given(this.ollamaApi.embed(this.embeddingsRequestCaptor.capture()))
114113
.willReturn(new EmbeddingsResponse("TEST_MODEL", List.of(new float[] { 0.1f, 0.2f, 0.3f }), 10L, 5L, 1));
115114

@@ -131,7 +130,7 @@ public void singleInputEmbedding() {
131130
}
132131

133132
@Test
134-
public void embeddingWithNullOptions() {
133+
void embeddingWithNullOptions() {
135134
given(this.ollamaApi.embed(this.embeddingsRequestCaptor.capture()))
136135
.willReturn(new EmbeddingsResponse("NULL_OPTIONS_MODEL", List.of(new float[] { 0.5f }), 5L, 2L, 1));
137136

@@ -150,7 +149,7 @@ public void embeddingWithNullOptions() {
150149
}
151150

152151
@Test
153-
public void embeddingWithMultipleLargeInputs() {
152+
void embeddingWithMultipleLargeInputs() {
154153
List<String> largeInputs = List.of(
155154
"This is a very long text input that might be used for document embedding scenarios",
156155
"Another substantial piece of text content that could represent a paragraph or section",
@@ -179,7 +178,7 @@ public void embeddingWithMultipleLargeInputs() {
179178
}
180179

181180
@Test
182-
public void embeddingWithCustomKeepAliveFormats() {
181+
void embeddingWithCustomKeepAliveFormats() {
183182
given(this.ollamaApi.embed(this.embeddingsRequestCaptor.capture()))
184183
.willReturn(new EmbeddingsResponse("KEEPALIVE_MODEL", List.of(new float[] { 1.0f }), 5L, 2L, 1));
185184

@@ -192,17 +191,17 @@ public void embeddingWithCustomKeepAliveFormats() {
192191
var secondsOptions = OllamaEmbeddingOptions.builder().model("KEEPALIVE_MODEL").keepAlive("300s").build();
193192

194193
embeddingModel.call(new EmbeddingRequest(List.of("Keep alive seconds"), secondsOptions));
195-
assertThat(this.embeddingsRequestCaptor.getValue().keepAlive()).isEqualTo(Duration.ofSeconds(300));
194+
assertThat(this.embeddingsRequestCaptor.getValue().keepAlive()).isEqualTo("300s");
196195

197196
// Test with hours format
198197
var hoursOptions = OllamaEmbeddingOptions.builder().model("KEEPALIVE_MODEL").keepAlive("2h").build();
199198

200199
embeddingModel.call(new EmbeddingRequest(List.of("Keep alive hours"), hoursOptions));
201-
assertThat(this.embeddingsRequestCaptor.getValue().keepAlive()).isEqualTo(Duration.ofHours(2));
200+
assertThat(this.embeddingsRequestCaptor.getValue().keepAlive()).isEqualTo("2h");
202201
}
203202

204203
@Test
205-
public void embeddingResponseMetadata() {
204+
void embeddingResponseMetadata() {
206205
given(this.ollamaApi.embed(this.embeddingsRequestCaptor.capture()))
207206
.willReturn(new EmbeddingsResponse("METADATA_MODEL", List.of(new float[] { 0.1f, 0.2f }), 100L, 50L, 25));
208207

@@ -220,7 +219,7 @@ public void embeddingResponseMetadata() {
220219
}
221220

222221
@Test
223-
public void embeddingWithZeroLengthVectors() {
222+
void embeddingWithZeroLengthVectors() {
224223
given(this.ollamaApi.embed(this.embeddingsRequestCaptor.capture()))
225224
.willReturn(new EmbeddingsResponse("ZERO_MODEL", List.of(new float[] {}), 0L, 0L, 1));
226225

@@ -237,7 +236,7 @@ public void embeddingWithZeroLengthVectors() {
237236
}
238237

239238
@Test
240-
public void builderValidation() {
239+
void builderValidation() {
241240
// Test that builder requires ollamaApi
242241
assertThatThrownBy(() -> OllamaEmbeddingModel.builder().build()).isInstanceOf(IllegalArgumentException.class);
243242

‎models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/OllamaEmbeddingRequestTests.java‎

Lines changed: 18 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,6 @@
1616

1717
package org.springframework.ai.ollama;
1818

19-
import java.time.Duration;
2019
import java.util.Arrays;
2120
import java.util.Collections;
2221
import java.util.List;
@@ -35,12 +34,12 @@
3534
* @author Thomas Vitale
3635
* @author Jonghoon Park
3736
*/
38-
public class OllamaEmbeddingRequestTests {
37+
class OllamaEmbeddingRequestTests {
3938

4039
private OllamaEmbeddingModel embeddingModel;
4140

4241
@BeforeEach
43-
public void setUp() {
42+
void setUp() {
4443
this.embeddingModel = OllamaEmbeddingModel.builder()
4544
.ollamaApi(OllamaApi.builder().build())
4645
.defaultOptions(
@@ -49,7 +48,7 @@ public void setUp() {
4948
}
5049

5150
@Test
52-
public void ollamaEmbeddingRequestDefaultOptions() {
51+
void ollamaEmbeddingRequestDefaultOptions() {
5352
var embeddingRequest = this.embeddingModel.buildEmbeddingRequest(new EmbeddingRequest(List.of("Hello"), null));
5453
var ollamaRequest = this.embeddingModel.ollamaEmbeddingRequest(embeddingRequest);
5554

@@ -58,7 +57,7 @@ public void ollamaEmbeddingRequestDefaultOptions() {
5857
}
5958

6059
@Test
61-
public void ollamaEmbeddingRequestRequestOptions() {
60+
void ollamaEmbeddingRequestRequestOptions() {
6261
var promptOptions = OllamaEmbeddingOptions.builder()//
6362
.model("PROMPT_MODEL")//
6463
.build();
@@ -72,18 +71,18 @@ public void ollamaEmbeddingRequestRequestOptions() {
7271
}
7372

7473
@Test
75-
public void ollamaEmbeddingRequestWithNegativeKeepAlive() {
74+
void ollamaEmbeddingRequestWithNegativeKeepAlive() {
7675
var promptOptions = OllamaEmbeddingOptions.builder().model("PROMPT_MODEL").keepAlive("-1m").build();
7776

7877
var embeddingRequest = this.embeddingModel
7978
.buildEmbeddingRequest(new EmbeddingRequest(List.of("Hello"), promptOptions));
8079
var ollamaRequest = this.embeddingModel.ollamaEmbeddingRequest(embeddingRequest);
8180

82-
assertThat(ollamaRequest.keepAlive()).isEqualTo(Duration.ofMinutes(-1));
81+
assertThat(ollamaRequest.keepAlive()).isEqualTo("-1m");
8382
}
8483

8584
@Test
86-
public void ollamaEmbeddingRequestWithEmptyInput() {
85+
void ollamaEmbeddingRequestWithEmptyInput() {
8786
var embeddingRequest = this.embeddingModel
8887
.buildEmbeddingRequest(new EmbeddingRequest(Collections.emptyList(), null));
8988
var ollamaRequest = this.embeddingModel.ollamaEmbeddingRequest(embeddingRequest);
@@ -93,7 +92,7 @@ public void ollamaEmbeddingRequestWithEmptyInput() {
9392
}
9493

9594
@Test
96-
public void ollamaEmbeddingRequestWithMultipleInputs() {
95+
void ollamaEmbeddingRequestWithMultipleInputs() {
9796
List<String> inputs = Arrays.asList("Hello", "World", "How are you?");
9897
var embeddingRequest = this.embeddingModel.buildEmbeddingRequest(new EmbeddingRequest(inputs, null));
9998
var ollamaRequest = this.embeddingModel.ollamaEmbeddingRequest(embeddingRequest);
@@ -103,7 +102,7 @@ public void ollamaEmbeddingRequestWithMultipleInputs() {
103102
}
104103

105104
@Test
106-
public void ollamaEmbeddingRequestOptionsOverrideDefaults() {
105+
void ollamaEmbeddingRequestOptionsOverrideDefaults() {
107106
var requestOptions = OllamaEmbeddingOptions.builder().model("OVERRIDE_MODEL").build();
108107

109108
var embeddingRequest = this.embeddingModel
@@ -115,24 +114,24 @@ public void ollamaEmbeddingRequestOptionsOverrideDefaults() {
115114
}
116115

117116
@Test
118-
public void ollamaEmbeddingRequestWithDifferentKeepAliveFormats() {
117+
void ollamaEmbeddingRequestWithDifferentKeepAliveFormats() {
119118
// Test seconds format
120119
var optionsSeconds = OllamaEmbeddingOptions.builder().keepAlive("30s").build();
121120
var requestSeconds = this.embeddingModel
122121
.buildEmbeddingRequest(new EmbeddingRequest(List.of("Test"), optionsSeconds));
123122
var ollamaRequestSeconds = this.embeddingModel.ollamaEmbeddingRequest(requestSeconds);
124-
assertThat(ollamaRequestSeconds.keepAlive()).isEqualTo(Duration.ofSeconds(30));
123+
assertThat(ollamaRequestSeconds.keepAlive()).isEqualTo("30s");
125124

126125
// Test hours format
127126
var optionsHours = OllamaEmbeddingOptions.builder().keepAlive("2h").build();
128127
var requestHours = this.embeddingModel
129128
.buildEmbeddingRequest(new EmbeddingRequest(List.of("Test"), optionsHours));
130129
var ollamaRequestHours = this.embeddingModel.ollamaEmbeddingRequest(requestHours);
131-
assertThat(ollamaRequestHours.keepAlive()).isEqualTo(Duration.ofHours(2));
130+
assertThat(ollamaRequestHours.keepAlive()).isEqualTo("2h");
132131
}
133132

134133
@Test
135-
public void ollamaEmbeddingRequestWithMinimalDefaults() {
134+
void ollamaEmbeddingRequestWithMinimalDefaults() {
136135
// Create model with minimal defaults
137136
var minimalModel = OllamaEmbeddingModel.builder()
138137
.ollamaApi(OllamaApi.builder().build())
@@ -151,7 +150,7 @@ public void ollamaEmbeddingRequestWithMinimalDefaults() {
151150
}
152151

153152
@Test
154-
public void ollamaEmbeddingRequestPreservesInputOrder() {
153+
void ollamaEmbeddingRequestPreservesInputOrder() {
155154
List<String> orderedInputs = Arrays.asList("First", "Second", "Third", "Fourth");
156155
var embeddingRequest = this.embeddingModel.buildEmbeddingRequest(new EmbeddingRequest(orderedInputs, null));
157156
var ollamaRequest = this.embeddingModel.ollamaEmbeddingRequest(embeddingRequest);
@@ -160,7 +159,7 @@ public void ollamaEmbeddingRequestPreservesInputOrder() {
160159
}
161160

162161
@Test
163-
public void ollamaEmbeddingRequestWithWhitespaceInputs() {
162+
void ollamaEmbeddingRequestWithWhitespaceInputs() {
164163
List<String> inputs = Arrays.asList("", " ", "\t\n", "normal text", " spaced ");
165164
var embeddingRequest = this.embeddingModel.buildEmbeddingRequest(new EmbeddingRequest(inputs, null));
166165
var ollamaRequest = this.embeddingModel.ollamaEmbeddingRequest(embeddingRequest);
@@ -170,7 +169,7 @@ public void ollamaEmbeddingRequestWithWhitespaceInputs() {
170169
}
171170

172171
@Test
173-
public void ollamaEmbeddingRequestWithNullInput() {
172+
void ollamaEmbeddingRequestWithNullInput() {
174173
// Test behavior when input list contains null values
175174
List<String> inputsWithNull = Arrays.asList("Hello", null, "World");
176175
var embeddingRequest = this.embeddingModel.buildEmbeddingRequest(new EmbeddingRequest(inputsWithNull, null));
@@ -181,7 +180,7 @@ public void ollamaEmbeddingRequestWithNullInput() {
181180
}
182181

183182
@Test
184-
public void ollamaEmbeddingRequestPartialOptionsOverride() {
183+
void ollamaEmbeddingRequestPartialOptionsOverride() {
185184
// Test that only specified options are overridden, others remain default
186185
var requestOptions = OllamaEmbeddingOptions.builder()
187186
.model("PARTIAL_OVERRIDE_MODEL")
@@ -199,7 +198,7 @@ public void ollamaEmbeddingRequestPartialOptionsOverride() {
199198
}
200199

201200
@Test
202-
public void ollamaEmbeddingRequestWithEmptyStringInput() {
201+
void ollamaEmbeddingRequestWithEmptyStringInput() {
203202
// Test with list containing only empty string
204203
var embeddingRequest = this.embeddingModel.buildEmbeddingRequest(new EmbeddingRequest(List.of(""), null));
205204
var ollamaRequest = this.embeddingModel.ollamaEmbeddingRequest(embeddingRequest);

0 commit comments

Comments
 (0)