Skip to content

Commit f54b3d5

Browse files
committed
Improve ChromaVectorStore with builder pattern and package organization
- Introduce builder pattern for ChromaVectorStore configuration - Move Chroma-related classes to org.springframework.ai.chroma package - Add proper null-safety annotations and validation - Deprecate direct constructors in favor of builder - Improve JavaDoc documentation - Clean up code organization and imports
1 parent 771c4a3 commit f54b3d5

File tree

14 files changed

+214
-69
lines changed

14 files changed

+214
-69
lines changed

spring-ai-core/src/main/java/org/springframework/ai/vectorstore/observation/AbstractObservationVectorStore.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,7 @@ public abstract class AbstractObservationVectorStore implements VectorStore {
4848
* @param customObservationConvention the custom observation convention to use
4949
*/
5050
public AbstractObservationVectorStore(ObservationRegistry observationRegistry,
51-
VectorStoreObservationConvention customObservationConvention) {
51+
@Nullable VectorStoreObservationConvention customObservationConvention) {
5252
this.observationRegistry = observationRegistry;
5353
this.customObservationConvention = customObservationConvention;
5454
}

spring-ai-core/src/main/java/org/springframework/ai/vectorstore/observation/VectorStoreObservationContext.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -194,7 +194,7 @@ public Builder withDimensions(Integer dimensions) {
194194
return this;
195195
}
196196

197-
public Builder withFieldName(String fieldName) {
197+
public Builder withFieldName(@Nullable String fieldName) {
198198
this.context.setFieldName(fieldName);
199199
return this;
200200
}

spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/chroma/ChromaVectorStoreAutoConfiguration.java

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -19,11 +19,11 @@
1919
import com.fasterxml.jackson.databind.ObjectMapper;
2020
import io.micrometer.observation.ObservationRegistry;
2121

22-
import org.springframework.ai.chroma.ChromaApi;
22+
import org.springframework.ai.chroma.api.ChromaApi;
2323
import org.springframework.ai.embedding.BatchingStrategy;
2424
import org.springframework.ai.embedding.EmbeddingModel;
2525
import org.springframework.ai.embedding.TokenCountBatchingStrategy;
26-
import org.springframework.ai.vectorstore.ChromaVectorStore;
26+
import org.springframework.ai.chroma.vectorstore.ChromaVectorStore;
2727
import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention;
2828
import org.springframework.beans.factory.ObjectProvider;
2929
import org.springframework.boot.autoconfigure.AutoConfiguration;
@@ -86,9 +86,13 @@ public ChromaVectorStore vectorStore(EmbeddingModel embeddingModel, ChromaApi ch
8686
ChromaVectorStoreProperties storeProperties, ObjectProvider<ObservationRegistry> observationRegistry,
8787
ObjectProvider<VectorStoreObservationConvention> customObservationConvention,
8888
BatchingStrategy chromaBatchingStrategy) {
89-
return new ChromaVectorStore(embeddingModel, chromaApi, storeProperties.getCollectionName(),
90-
storeProperties.isInitializeSchema(), observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP),
91-
customObservationConvention.getIfAvailable(() -> null), chromaBatchingStrategy);
89+
return ChromaVectorStore.builder(embeddingModel, chromaApi)
90+
.collectionName(storeProperties.getCollectionName())
91+
.initializeSchema(storeProperties.isInitializeSchema())
92+
.observationRegistry(observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP))
93+
.customObservationConvention(customObservationConvention.getIfAvailable(() -> null))
94+
.batchingStrategy(chromaBatchingStrategy)
95+
.build();
9296
}
9397

9498
static class PropertiesChromaConnectionDetails implements ChromaConnectionDetails {

spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/chroma/ChromaVectorStoreProperties.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717
package org.springframework.ai.autoconfigure.vectorstore.chroma;
1818

1919
import org.springframework.ai.autoconfigure.vectorstore.CommonVectorStoreProperties;
20-
import org.springframework.ai.vectorstore.ChromaVectorStore;
20+
import org.springframework.ai.chroma.vectorstore.ChromaVectorStore;
2121
import org.springframework.boot.context.properties.ConfigurationProperties;
2222

2323
/**
Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414
* limitations under the License.
1515
*/
1616

17-
package org.springframework.ai.chroma;
17+
package org.springframework.ai.chroma.api;
1818

1919
import java.util.ArrayList;
2020
import java.util.HashMap;
@@ -28,7 +28,7 @@
2828
import com.fasterxml.jackson.core.JsonProcessingException;
2929
import com.fasterxml.jackson.databind.ObjectMapper;
3030

31-
import org.springframework.ai.chroma.ChromaApi.QueryRequest.Include;
31+
import org.springframework.ai.chroma.api.ChromaApi.QueryRequest.Include;
3232
import org.springframework.http.HttpHeaders;
3333
import org.springframework.http.MediaType;
3434
import org.springframework.http.client.SimpleClientHttpRequestFactory;
@@ -49,10 +49,10 @@
4949
public class ChromaApi {
5050

5151
// Regular expression pattern that looks for a message inside the ValueError(...).
52-
private static Pattern VALUE_ERROR_PATTERN = Pattern.compile("ValueError\\('([^']*)'\\)");
52+
private static final Pattern VALUE_ERROR_PATTERN = Pattern.compile("ValueError\\('([^']*)'\\)");
5353

5454
// Regular expression pattern that looks for a message.
55-
private static Pattern MESSAGE_ERROR_PATTERN = Pattern.compile("\"message\":\"(.*?)\"");
55+
private static final Pattern MESSAGE_ERROR_PATTERN = Pattern.compile("\"message\":\"(.*?)\"");
5656

5757
private final ObjectMapper objectMapper;
5858

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414
* limitations under the License.
1515
*/
1616

17-
package org.springframework.ai.vectorstore;
17+
package org.springframework.ai.chroma.vectorstore;
1818

1919
import org.springframework.ai.vectorstore.filter.Filter;
2020
import org.springframework.ai.vectorstore.filter.converter.PineconeFilterExpressionConverter;
Lines changed: 122 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414
* limitations under the License.
1515
*/
1616

17-
package org.springframework.ai.vectorstore;
17+
package org.springframework.ai.chroma.vectorstore;
1818

1919
import java.util.ArrayList;
2020
import java.util.HashMap;
@@ -27,10 +27,10 @@
2727
import com.fasterxml.jackson.databind.json.JsonMapper;
2828
import io.micrometer.observation.ObservationRegistry;
2929

30-
import org.springframework.ai.chroma.ChromaApi;
31-
import org.springframework.ai.chroma.ChromaApi.AddEmbeddingsRequest;
32-
import org.springframework.ai.chroma.ChromaApi.DeleteEmbeddingsRequest;
33-
import org.springframework.ai.chroma.ChromaApi.Embedding;
30+
import org.springframework.ai.chroma.api.ChromaApi;
31+
import org.springframework.ai.chroma.api.ChromaApi.AddEmbeddingsRequest;
32+
import org.springframework.ai.chroma.api.ChromaApi.DeleteEmbeddingsRequest;
33+
import org.springframework.ai.chroma.api.ChromaApi.Embedding;
3434
import org.springframework.ai.document.Document;
3535
import org.springframework.ai.document.DocumentMetadata;
3636
import org.springframework.ai.embedding.BatchingStrategy;
@@ -39,12 +39,15 @@
3939
import org.springframework.ai.embedding.TokenCountBatchingStrategy;
4040
import org.springframework.ai.observation.conventions.VectorStoreProvider;
4141
import org.springframework.ai.util.JacksonUtils;
42+
import org.springframework.ai.vectorstore.SearchRequest;
43+
import org.springframework.ai.vectorstore.VectorStore;
4244
import org.springframework.ai.vectorstore.filter.FilterExpressionConverter;
4345
import org.springframework.ai.vectorstore.observation.AbstractObservationVectorStore;
4446
import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext;
4547
import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention;
4648
import org.springframework.beans.factory.InitializingBean;
4749
import org.springframework.lang.NonNull;
50+
import org.springframework.lang.Nullable;
4851
import org.springframework.util.Assert;
4952
import org.springframework.util.CollectionUtils;
5053

@@ -73,6 +76,7 @@ public class ChromaVectorStore extends AbstractObservationVectorStore implements
7376

7477
private FilterExpressionConverter filterExpressionConverter;
7578

79+
@Nullable
7680
private String collectionId;
7781

7882
private final boolean initializeSchema;
@@ -83,19 +87,31 @@ public class ChromaVectorStore extends AbstractObservationVectorStore implements
8387

8488
private boolean initialized = false;
8589

90+
/**
91+
* @deprecated in favor of the builder method
92+
*/
93+
@Deprecated(forRemoval = true)
8694
public ChromaVectorStore(EmbeddingModel embeddingModel, ChromaApi chromaApi, boolean initializeSchema) {
8795
this(embeddingModel, chromaApi, DEFAULT_COLLECTION_NAME, initializeSchema);
8896
}
8997

98+
/**
99+
* @deprecated in favor of the builder method
100+
*/
101+
@Deprecated(forRemoval = true)
90102
public ChromaVectorStore(EmbeddingModel embeddingModel, ChromaApi chromaApi, String collectionName,
91103
boolean initializeSchema) {
92104
this(embeddingModel, chromaApi, collectionName, initializeSchema, ObservationRegistry.NOOP, null,
93105
new TokenCountBatchingStrategy());
94106
}
95107

108+
/**
109+
* @deprecated in favor of the builder method
110+
*/
111+
@Deprecated(forRemoval = true)
96112
public ChromaVectorStore(EmbeddingModel embeddingModel, ChromaApi chromaApi, String collectionName,
97113
boolean initializeSchema, ObservationRegistry observationRegistry,
98-
VectorStoreObservationConvention customObservationConvention, BatchingStrategy batchingStrategy) {
114+
@Nullable VectorStoreObservationConvention customObservationConvention, BatchingStrategy batchingStrategy) {
99115

100116
super(observationRegistry, customObservationConvention);
101117

@@ -108,6 +124,9 @@ public ChromaVectorStore(EmbeddingModel embeddingModel, ChromaApi chromaApi, Str
108124
this.objectMapper = JsonMapper.builder().addModules(JacksonUtils.instantiateAvailableModules()).build();
109125
}
110126

127+
/**
128+
* @param builder {@link Builder} for chroma vector store
129+
*/
111130
private ChromaVectorStore(Builder builder) {
112131
super(builder.observationRegistry, builder.customObservationConvention);
113132
this.embeddingModel = builder.embeddingModel;
@@ -128,9 +147,27 @@ private ChromaVectorStore(Builder builder) {
128147
}
129148
}
130149

131-
public void setFilterExpressionConverter(FilterExpressionConverter filterExpressionConverter) {
132-
Assert.notNull(filterExpressionConverter, "FilterExpressionConverter should not be null.");
133-
this.filterExpressionConverter = filterExpressionConverter;
150+
@Override
151+
public void afterPropertiesSet() throws Exception {
152+
if (!this.initialized) {
153+
var collection = this.chromaApi.getCollection(this.collectionName);
154+
if (collection == null) {
155+
if (this.initializeSchema) {
156+
collection = this.chromaApi
157+
.createCollection(new ChromaApi.CreateCollectionRequest(this.collectionName));
158+
}
159+
else {
160+
throw new RuntimeException("Collection " + this.collectionName
161+
+ " doesn't exist and won't be created as the initializeSchema is set to false.");
162+
}
163+
}
164+
this.collectionId = collection.id();
165+
this.initialized = true;
166+
}
167+
}
168+
169+
public static Builder builder(EmbeddingModel embeddingModel, ChromaApi chromaApi) {
170+
return new Builder(embeddingModel, chromaApi);
134171
}
135172

136173
@Override
@@ -217,36 +254,34 @@ private Map<String, Object> jsonToMap(String jsonText) {
217254
}
218255
}
219256

257+
/**
258+
* @deprecated not used currently anywhere
259+
*/
260+
@Deprecated(forRemoval = true)
220261
public String getCollectionName() {
221262
return this.collectionName;
222263
}
223264

265+
/**
266+
* @deprecated only used in tests
267+
*/
268+
@Deprecated(forRemoval = true)
269+
@Nullable
224270
public String getCollectionId() {
225271
return this.collectionId;
226272
}
227273

228-
@Override
229-
public void afterPropertiesSet() throws Exception {
230-
if (!this.initialized) {
231-
var collection = this.chromaApi.getCollection(this.collectionName);
232-
if (collection == null) {
233-
if (this.initializeSchema) {
234-
collection = this.chromaApi
235-
.createCollection(new ChromaApi.CreateCollectionRequest(this.collectionName));
236-
}
237-
else {
238-
throw new RuntimeException("Collection " + this.collectionName
239-
+ " doesn't exist and won't be created as the initializeSchema is set to false.");
240-
}
241-
}
242-
this.collectionId = collection.id();
243-
this.initialized = true;
244-
}
274+
/**
275+
* @deprecated in favor the builder method
276+
*/
277+
@Deprecated(forRemoval = true)
278+
public void setFilterExpressionConverter(FilterExpressionConverter filterExpressionConverter) {
279+
Assert.notNull(filterExpressionConverter, "FilterExpressionConverter should not be null.");
280+
this.filterExpressionConverter = filterExpressionConverter;
245281
}
246282

247283
@Override
248-
public @NonNull VectorStoreObservationContext.Builder createObservationContextBuilder(
249-
@NonNull String operationName) {
284+
public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName) {
250285
return VectorStoreObservationContext.builder(VectorStoreProvider.CHROMA.value(), operationName)
251286
.withDimensions(this.embeddingModel.dimensions())
252287
.withCollectionName(this.collectionName + ":" + this.collectionId);
@@ -264,6 +299,7 @@ public static class Builder {
264299

265300
private ObservationRegistry observationRegistry = ObservationRegistry.NOOP;
266301

302+
@Nullable
267303
private VectorStoreObservationConvention customObservationConvention = null;
268304

269305
private BatchingStrategy batchingStrategy = new TokenCountBatchingStrategy();
@@ -272,46 +308,103 @@ public static class Builder {
272308

273309
private boolean initializeImmediately = false;
274310

311+
/**
312+
* Constructor with required parameters.
313+
* @param embeddingModel The {@link EmbeddingModel} to use
314+
* @param chromaApi The {@link ChromaApi} instance
315+
* @throws IllegalArgumentException if any required parameter is null
316+
*/
275317
public Builder(EmbeddingModel embeddingModel, ChromaApi chromaApi) {
318+
Assert.notNull(embeddingModel, "embeddingModel must not be null");
319+
Assert.notNull(chromaApi, "chromaApi must not be null");
276320
this.embeddingModel = embeddingModel;
277321
this.chromaApi = chromaApi;
278322
}
279323

324+
/**
325+
* Sets the collection name.
326+
* @param collectionName the name of the collection
327+
* @return the builder instance
328+
* @throws IllegalArgumentException if collectionName is null or empty
329+
*/
280330
public Builder collectionName(String collectionName) {
331+
Assert.hasText(collectionName, "collectionName must not be null or empty");
281332
this.collectionName = collectionName;
282333
return this;
283334
}
284335

336+
/**
337+
* Sets whether to initialize the schema.
338+
* @param initializeSchema true to initialize schema, false otherwise
339+
* @return the builder instance
340+
*/
285341
public Builder initializeSchema(boolean initializeSchema) {
286342
this.initializeSchema = initializeSchema;
287343
return this;
288344
}
289345

346+
/**
347+
* Sets the {@link ObservationRegistry}
348+
* @param observationRegistry the observation registry to use
349+
* @return the builder instance
350+
* @throws IllegalArgumentException if observationRegistry is null
351+
*/
290352
public Builder observationRegistry(ObservationRegistry observationRegistry) {
353+
Assert.notNull(observationRegistry, "observationRegistry must not be null");
291354
this.observationRegistry = observationRegistry;
292355
return this;
293356
}
294357

295-
public Builder customObservationConvention(VectorStoreObservationConvention convention) {
358+
/**
359+
* Sets the custom observation convention
360+
* {@link VectorStoreObservationConvention}.
361+
* @param convention the custom observation convention to use
362+
* @return the builder instance
363+
*/
364+
public Builder customObservationConvention(@Nullable VectorStoreObservationConvention convention) {
296365
this.customObservationConvention = convention;
297366
return this;
298367
}
299368

369+
/**
370+
* Sets the batching strategy.
371+
* @param batchingStrategy the batching strategy to use
372+
* @return the builder instance
373+
* @throws IllegalArgumentException if batchingStrategy is null
374+
*/
300375
public Builder batchingStrategy(BatchingStrategy batchingStrategy) {
376+
Assert.notNull(batchingStrategy, "batchingStrategy must not be null");
301377
this.batchingStrategy = batchingStrategy;
302378
return this;
303379
}
304380

381+
/**
382+
* Sets the filter expression converter.
383+
* @param converter the filter expression converter to use
384+
* @return the builder instance
385+
* @throws IllegalArgumentException if converter is null
386+
*/
305387
public Builder filterExpressionConverter(FilterExpressionConverter converter) {
388+
Assert.notNull(converter, "filterExpressionConverter must not be null");
306389
this.filterExpressionConverter = converter;
307390
return this;
308391
}
309392

393+
/**
394+
* Sets whether to initialize immediately.
395+
* @param initialize true to initialize immediately, false otherwise
396+
* @return the builder instance
397+
*/
310398
public Builder initializeImmediately(boolean initialize) {
311399
this.initializeImmediately = initialize;
312400
return this;
313401
}
314402

403+
/**
404+
* Builds the {@link ChromaVectorStore} instance.
405+
* @return a new ChromaVectorStore instance
406+
* @throws IllegalStateException if the builder is in an invalid state
407+
*/
315408
public ChromaVectorStore build() {
316409
return new ChromaVectorStore(this);
317410
}

0 commit comments

Comments
 (0)