1414 * limitations under the License.
1515 */
1616
17- package org .springframework .ai .vectorstore ;
17+ package org .springframework .ai .chroma . vectorstore ;
1818
1919import java .util .ArrayList ;
2020import java .util .HashMap ;
2727import com .fasterxml .jackson .databind .json .JsonMapper ;
2828import 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 ;
3434import org .springframework .ai .document .Document ;
3535import org .springframework .ai .document .DocumentMetadata ;
3636import org .springframework .ai .embedding .BatchingStrategy ;
3939import org .springframework .ai .embedding .TokenCountBatchingStrategy ;
4040import org .springframework .ai .observation .conventions .VectorStoreProvider ;
4141import org .springframework .ai .util .JacksonUtils ;
42+ import org .springframework .ai .vectorstore .SearchRequest ;
43+ import org .springframework .ai .vectorstore .VectorStore ;
4244import org .springframework .ai .vectorstore .filter .FilterExpressionConverter ;
4345import org .springframework .ai .vectorstore .observation .AbstractObservationVectorStore ;
4446import org .springframework .ai .vectorstore .observation .VectorStoreObservationContext ;
4547import org .springframework .ai .vectorstore .observation .VectorStoreObservationConvention ;
4648import org .springframework .beans .factory .InitializingBean ;
4749import org .springframework .lang .NonNull ;
50+ import org .springframework .lang .Nullable ;
4851import org .springframework .util .Assert ;
4952import 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