4141import  org .springframework .ai .vectorstore .filter .FilterExpressionConverter ;
4242import  org .springframework .ai .vectorstore .observation .AbstractObservationVectorStore ;
4343import  org .springframework .ai .vectorstore .observation .VectorStoreObservationContext ;
44- import  org .springframework .ai .vectorstore .observation .VectorStoreObservationContext .Builder ;
4544import  org .springframework .ai .vectorstore .observation .VectorStoreObservationConvention ;
4645import  org .springframework .beans .factory .InitializingBean ;
4746import  org .springframework .util .Assert ;
5857 * @author Christian Tzolov 
5958 * @author Fu Cheng 
6059 * @author Sebastien Deleuze 
61-  * 
60+  * @author Soby Chacko  
6261 */ 
6362public  class  ChromaVectorStore  extends  AbstractObservationVectorStore  implements  InitializingBean  {
6463
6564	public  static  final  String  DISTANCE_FIELD_NAME  = "distance" ;
6665
6766	public  static  final  String  DEFAULT_COLLECTION_NAME  = "SpringAiCollection" ;
6867
69- 	public  static  final  double  SIMILARITY_THRESHOLD_ALL  = 0.0 ;
70- 
71- 	public  static  final  int  DEFAULT_TOP_K  = 4 ;
72- 
7368	private  final  EmbeddingModel  embeddingModel ;
7469
7570	private  final  ChromaApi  chromaApi ;
@@ -86,6 +81,8 @@ public class ChromaVectorStore extends AbstractObservationVectorStore implements
8681
8782	private  final  ObjectMapper  objectMapper ;
8883
84+ 	private  boolean  initialized  = false ;
85+ 
8986	public  ChromaVectorStore (EmbeddingModel  embeddingModel , ChromaApi  chromaApi , boolean  initializeSchema ) {
9087		this (embeddingModel , chromaApi , DEFAULT_COLLECTION_NAME , initializeSchema );
9188	}
@@ -111,6 +108,26 @@ public ChromaVectorStore(EmbeddingModel embeddingModel, ChromaApi chromaApi, Str
111108		this .objectMapper  = JsonMapper .builder ().addModules (JacksonUtils .instantiateAvailableModules ()).build ();
112109	}
113110
111+ 	private  ChromaVectorStore (Builder  builder ) {
112+ 		super (builder .observationRegistry , builder .customObservationConvention );
113+ 		this .embeddingModel  = builder .embeddingModel ;
114+ 		this .chromaApi  = builder .chromaApi ;
115+ 		this .collectionName  = builder .collectionName ;
116+ 		this .initializeSchema  = builder .initializeSchema ;
117+ 		this .filterExpressionConverter  = builder .filterExpressionConverter ;
118+ 		this .batchingStrategy  = builder .batchingStrategy ;
119+ 		this .objectMapper  = JsonMapper .builder ().addModules (JacksonUtils .instantiateAvailableModules ()).build ();
120+ 
121+ 		if  (builder .initializeImmediately ) {
122+ 			try  {
123+ 				afterPropertiesSet ();
124+ 			}
125+ 			catch  (Exception  e ) {
126+ 				throw  new  IllegalStateException ("Failed to initialize ChromaVectorStore" , e );
127+ 			}
128+ 		}
129+ 	}
130+ 
114131	public  void  setFilterExpressionConverter (FilterExpressionConverter  filterExpressionConverter ) {
115132		Assert .notNull (filterExpressionConverter , "FilterExpressionConverter should not be null." );
116133		this .filterExpressionConverter  = filterExpressionConverter ;
@@ -207,26 +224,95 @@ public String getCollectionId() {
207224
208225	@ Override 
209226	public  void  afterPropertiesSet () throws  Exception  {
210- 		var  collection  = this .chromaApi .getCollection (this .collectionName );
211- 		if  (collection  == null ) {
212- 			if  (this .initializeSchema ) {
213- 				collection  = this .chromaApi 
214- 					.createCollection (new  ChromaApi .CreateCollectionRequest (this .collectionName ));
215- 			}
216- 			else  {
217- 				throw  new  RuntimeException ("Collection "  + this .collectionName 
218- 						+ " doesn't exist and won't be created as the initializeSchema is set to false." );
227+ 		if  (!this .initialized ) {
228+ 			var  collection  = this .chromaApi .getCollection (this .collectionName );
229+ 			if  (collection  == null ) {
230+ 				if  (this .initializeSchema ) {
231+ 					collection  = this .chromaApi 
232+ 						.createCollection (new  ChromaApi .CreateCollectionRequest (this .collectionName ));
233+ 				}
234+ 				else  {
235+ 					throw  new  RuntimeException ("Collection "  + this .collectionName 
236+ 							+ " doesn't exist and won't be created as the initializeSchema is set to false." );
237+ 				}
219238			}
239+ 			this .collectionId  = collection .id ();
240+ 			this .initialized  = true ;
220241		}
221- 		this .collectionId  = collection .id ();
222242	}
223243
224244	@ Override 
225- 	public  Builder  createObservationContextBuilder (String  operationName ) {
245+ 	public  VectorStoreObservationContext . Builder  createObservationContextBuilder (String  operationName ) {
226246		return  VectorStoreObservationContext .builder (VectorStoreProvider .CHROMA .value (), operationName )
227247			.withDimensions (this .embeddingModel .dimensions ())
228248			.withCollectionName (this .collectionName  + ":"  + this .collectionId )
229249			.withFieldName (this .initializeSchema  ? DISTANCE_FIELD_NAME  : null );
230250	}
231251
252+ 	public  static  class  Builder  {
253+ 
254+ 		private  final  EmbeddingModel  embeddingModel ;
255+ 
256+ 		private  final  ChromaApi  chromaApi ;
257+ 
258+ 		private  String  collectionName  = DEFAULT_COLLECTION_NAME ;
259+ 
260+ 		private  boolean  initializeSchema  = false ;
261+ 
262+ 		private  ObservationRegistry  observationRegistry  = ObservationRegistry .NOOP ;
263+ 
264+ 		private  VectorStoreObservationConvention  customObservationConvention  = null ;
265+ 
266+ 		private  BatchingStrategy  batchingStrategy  = new  TokenCountBatchingStrategy ();
267+ 
268+ 		private  FilterExpressionConverter  filterExpressionConverter  = new  ChromaFilterExpressionConverter ();
269+ 
270+ 		private  boolean  initializeImmediately  = false ;
271+ 
272+ 		public  Builder (EmbeddingModel  embeddingModel , ChromaApi  chromaApi ) {
273+ 			this .embeddingModel  = embeddingModel ;
274+ 			this .chromaApi  = chromaApi ;
275+ 		}
276+ 
277+ 		public  Builder  collectionName (String  collectionName ) {
278+ 			this .collectionName  = collectionName ;
279+ 			return  this ;
280+ 		}
281+ 
282+ 		public  Builder  initializeSchema (boolean  initializeSchema ) {
283+ 			this .initializeSchema  = initializeSchema ;
284+ 			return  this ;
285+ 		}
286+ 
287+ 		public  Builder  observationRegistry (ObservationRegistry  observationRegistry ) {
288+ 			this .observationRegistry  = observationRegistry ;
289+ 			return  this ;
290+ 		}
291+ 
292+ 		public  Builder  customObservationConvention (VectorStoreObservationConvention  convention ) {
293+ 			this .customObservationConvention  = convention ;
294+ 			return  this ;
295+ 		}
296+ 
297+ 		public  Builder  batchingStrategy (BatchingStrategy  batchingStrategy ) {
298+ 			this .batchingStrategy  = batchingStrategy ;
299+ 			return  this ;
300+ 		}
301+ 
302+ 		public  Builder  filterExpressionConverter (FilterExpressionConverter  converter ) {
303+ 			this .filterExpressionConverter  = converter ;
304+ 			return  this ;
305+ 		}
306+ 
307+ 		public  Builder  initializeImmediately (boolean  initialize ) {
308+ 			this .initializeImmediately  = initialize ;
309+ 			return  this ;
310+ 		}
311+ 
312+ 		public  ChromaVectorStore  build () {
313+ 			return  new  ChromaVectorStore (this );
314+ 		}
315+ 
316+ 	}
317+ 
232318}
0 commit comments