From 6307377ef5ab2f70e92e5182220a4af185d9127d Mon Sep 17 00:00:00 2001
From: Soby Chacko <soby.chacko@broadcom.com>
Date: Fri, 6 Dec 2024 18:13:58 -0500
Subject: [PATCH] Fix null issues with custom observation convention

 - Relax null constraints for custom observation convention in AbstractVectorStoreBuilder
---
 .../ai/vectorstore/AbstractVectorStoreBuilder.java    |  1 -
 .../chroma/ChromaVectorStoreAutoConfiguration.java    | 11 ++++-------
 2 files changed, 4 insertions(+), 8 deletions(-)

diff --git a/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/AbstractVectorStoreBuilder.java b/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/AbstractVectorStoreBuilder.java
index 65912b3a4..0affbdcc8 100644
--- a/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/AbstractVectorStoreBuilder.java
+++ b/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/AbstractVectorStoreBuilder.java
@@ -72,7 +72,6 @@ public abstract class AbstractVectorStoreBuilder<T extends AbstractVectorStoreBu
 
 	@Override
 	public T customObservationConvention(VectorStoreObservationConvention convention) {
-		Assert.notNull(convention, "VectorStoreObservationConvention must not be null");
 		this.customObservationConvention = convention;
 		return self();
 	}
diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/chroma/ChromaVectorStoreAutoConfiguration.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/chroma/ChromaVectorStoreAutoConfiguration.java
index 3a1313f48..9f0b8493c 100644
--- a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/chroma/ChromaVectorStoreAutoConfiguration.java
+++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/chroma/ChromaVectorStoreAutoConfiguration.java
@@ -86,17 +86,14 @@ public class ChromaVectorStoreAutoConfiguration {
 			ChromaVectorStoreProperties storeProperties, ObjectProvider<ObservationRegistry> observationRegistry,
 			ObjectProvider<VectorStoreObservationConvention> customObservationConvention,
 			BatchingStrategy chromaBatchingStrategy) {
-		ChromaVectorStore.ChromaBuilder chromaBuilder = ChromaVectorStore.builder(chromaApi)
+		return ChromaVectorStore.builder(chromaApi)
 			.embeddingModel(embeddingModel)
 			.collectionName(storeProperties.getCollectionName())
 			.initializeSchema(storeProperties.isInitializeSchema())
 			.observationRegistry(observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP))
-			.batchingStrategy(chromaBatchingStrategy);
-		VectorStoreObservationConvention observationConvention = customObservationConvention.getIfAvailable();
-		if (observationConvention != null) {
-			chromaBuilder.customObservationConvention(observationConvention);
-		}
-		return chromaBuilder.build();
+			.customObservationConvention(customObservationConvention.getIfAvailable(() -> null))
+			.batchingStrategy(chromaBatchingStrategy)
+			.build();
 	}
 
 	static class PropertiesChromaConnectionDetails implements ChromaConnectionDetails {