Add observability support to VectorStore

spring-ai-core/src/main/java/org/springframework/ai/observation/conventions/AiObservationAttributes.java

@@ -132,7 +132,11 @@ public enum AiObservationAttributes {
 	/**
 	 * The full response received from the model.
 	 */
-	COMPLETION("gen_ai.completion");
+	COMPLETION("gen_ai.completion"),
+	/**
+	 * The name of the operation or command being executed.
+	 */
+	DB_OPERATION_NAME("db.operation.name"),;
 
 	private final String value;
 

spring-ai-core/src/main/java/org/springframework/ai/observation/conventions/VectorStoreProvider.java

@@ -0,0 +1,39 @@
+/*
+* Copyright 2024 - 2024 the original author or authors.
+*
+* Licensed under the Apache License, Version 2.0 (the "License");
+* you may not use this file except in compliance with the License.
+* You may obtain a copy of the License at
+*
+* https://www.apache.org/licenses/LICENSE-2.0
+*
+* Unless required by applicable law or agreed to in writing, software
+* distributed under the License is distributed on an "AS IS" BASIS,
+* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+* See the License for the specific language governing permissions and
+* limitations under the License.
+*/
+package org.springframework.ai.observation.conventions;
+
+/**
+ * @author Christian Tzolov
+ * @since 1.0.0
+ */
+public enum VectorStoreProvider {
+
+	// @formatter:off
+        PG_VECTOR("pg_vector"),
+        SIMPLE_VECTOR_STORE("simple_vector_store");
+
+        // @formatter:on
+	private final String value;
+
+	VectorStoreProvider(String value) {
+		this.value = value;
+	}
+
+	public String value() {
+		return this.value;
+	}
+
+}

spring-ai-core/src/main/java/org/springframework/ai/observation/conventions/VectorStoreSimilarityMetric.java

@@ -0,0 +1,54 @@
+/*
+* Copyright 2024 - 2024 the original author or authors.
+*
+* Licensed under the Apache License, Version 2.0 (the "License");
+* you may not use this file except in compliance with the License.
+* You may obtain a copy of the License at
+*
+* https://www.apache.org/licenses/LICENSE-2.0
+*
+* Unless required by applicable law or agreed to in writing, software
+* distributed under the License is distributed on an "AS IS" BASIS,
+* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+* See the License for the specific language governing permissions and
+* limitations under the License.
+*/
+package org.springframework.ai.observation.conventions;
+
+/**
+ * @author Christian Tzolov
+ * @since 1.0.0
+ */
+public enum VectorStoreSimilarityMetric {
+
+	// @formatter:off
+
+        /**
+         *  The cosine metric.
+         */
+        COSINE("cosine"),
+        /**
+         * The euclidean distance metric.
+         */
+        EUCLIDEAN("euclidean"),
+        /**
+         * The manhattan distance metric.
+         */
+        MANHATTAN("manhattan"),
+        /**
+         * The dot product metric.
+         */
+        DOT("dot");
+
+        // @formatter:on
+	private final String value;
+
+	VectorStoreSimilarityMetric(String value) {
+		this.value = value;
+	}
+
+	public String value() {
+		return this.value;
+	}
+
+}

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

@@ -15,16 +15,6 @@
  */
 package org.springframework.ai.vectorstore;
 
-import com.fasterxml.jackson.core.JsonProcessingException;
-import com.fasterxml.jackson.core.type.TypeReference;
-import com.fasterxml.jackson.databind.ObjectMapper;
-import com.fasterxml.jackson.databind.ObjectWriter;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-import org.springframework.ai.document.Document;
-import org.springframework.ai.embedding.EmbeddingModel;
-import org.springframework.core.io.Resource;
-
 import java.io.File;
 import java.io.FileOutputStream;
 import java.io.IOException;
@@ -40,6 +30,21 @@
 import java.util.Optional;
 import java.util.concurrent.ConcurrentHashMap;
 
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.springframework.ai.document.Document;
+import org.springframework.ai.embedding.EmbeddingModel;
+import org.springframework.ai.observation.conventions.VectorStoreProvider;
+import org.springframework.ai.observation.conventions.VectorStoreSimilarityMetric;
+import org.springframework.ai.vectorstore.observation.AbstractObservationVectorStore;
+import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext;
+import org.springframework.core.io.Resource;
+
+import com.fasterxml.jackson.core.JsonProcessingException;
+import com.fasterxml.jackson.core.type.TypeReference;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.ObjectWriter;
+
 /**
  * SimpleVectorStore is a simple implementation of the VectorStore interface.
  *
@@ -55,7 +60,7 @@
  * @author Mark Pollack
  * @author Christian Tzolov
  */
-public class SimpleVectorStore implements VectorStore {
+public class SimpleVectorStore extends AbstractObservationVectorStore {
 
 	private static final Logger logger = LoggerFactory.getLogger(SimpleVectorStore.class);
 
@@ -69,7 +74,7 @@ public SimpleVectorStore(EmbeddingModel embeddingModel) {
 	}
 
 	@Override
-	public void add(List<Document> documents) {
+	public void doAdd(List<Document> documents) {
 		for (Document document : documents) {
 			logger.info("Calling EmbeddingModel for document id = {}", document.getId());
 			float[] embedding = this.embeddingModel.embed(document);
@@ -79,15 +84,15 @@ public void add(List<Document> documents) {
 	}
 
 	@Override
-	public Optional<Boolean> delete(List<String> idList) {
+	public Optional<Boolean> doDelete(List<String> idList) {
 		for (String id : idList) {
 			this.store.remove(id);
 		}
 		return Optional.of(true);
 	}
 
 	@Override
-	public List<Document> similaritySearch(SearchRequest request) {
+	public List<Document> doSimilaritySearch(SearchRequest request) {
 		if (request.getFilterExpression() != null) {
 			throw new UnsupportedOperationException(
 					"The [" + this.getClass() + "] doesn't support metadata filtering!");
@@ -247,4 +252,13 @@ public static float norm(float[] vector) {
 
 	}
 
+	@Override
+	public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName) {
+
+		return VectorStoreObservationContext.builder(VectorStoreProvider.SIMPLE_VECTOR_STORE.value(), operationName)
+			.withDimensions(this.embeddingModel.dimensions())
+			.withCollectionName("in-memory-map")
+			.withSimilarityMetric(VectorStoreSimilarityMetric.COSINE.value());
+	}
+
 }

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

@@ -31,6 +31,10 @@
  */
 public interface VectorStore extends DocumentWriter {
 
+	default String getName() {
+		return this.getClass().getSimpleName();
+	}
+
 	/**
 	 * Adds list of {@link Document}s to the vector store.
 	 * @param documents the list of documents to store. Throws an exception if the

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

@@ -0,0 +1,107 @@
+/*
+* Copyright 2024 - 2024 the original author or authors.
+*
+* Licensed under the Apache License, Version 2.0 (the "License");
+* you may not use this file except in compliance with the License.
+* You may obtain a copy of the License at
+*
+* https://www.apache.org/licenses/LICENSE-2.0
+*
+* Unless required by applicable law or agreed to in writing, software
+* distributed under the License is distributed on an "AS IS" BASIS,
+* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+* See the License for the specific language governing permissions and
+* limitations under the License.
+*/
+package org.springframework.ai.vectorstore.observation;
+
+import java.util.List;
+import java.util.Optional;
+
+import org.springframework.ai.document.Document;
+import org.springframework.ai.vectorstore.SearchRequest;
+import org.springframework.ai.vectorstore.VectorStore;
+import org.springframework.lang.Nullable;
+
+import io.micrometer.observation.ObservationRegistry;
+
+/**
+ * @author Christian Tzolov
+ * @since 1.0.0
+ */
+public abstract class AbstractObservationVectorStore implements VectorStore {
+
+	private static final VectorStoreObservationConvention DEFAULT_OBSERVATION_CONVENTION = new DefaultVectorStoreObservationConvention();
+
+	private final ObservationRegistry observationRegistry;
+
+	@Nullable
+	private final VectorStoreObservationConvention customObservationConvention;
+
+	public AbstractObservationVectorStore() {
+		this(ObservationRegistry.NOOP, null);
+	}
+
+	public AbstractObservationVectorStore(ObservationRegistry observationRegistry) {
+		this(observationRegistry, null);
+	}
+
+	public AbstractObservationVectorStore(ObservationRegistry observationRegistry,
+			VectorStoreObservationConvention customSearchObservationConvention) {
+		this.observationRegistry = observationRegistry;
+		this.customObservationConvention = customSearchObservationConvention;
+	}
+
+	@Override
+	public void add(List<Document> documents) {
+
+		VectorStoreObservationContext observationContext = this
+			.createObservationContextBuilder(VectorStoreObservationContext.Operation.ADD.value())
+			.build();
+
+		VectorStoreObservationDocumentation.AI_VECTOR_STORE
+			.observation(this.customObservationConvention, DEFAULT_OBSERVATION_CONVENTION, () -> observationContext,
+					observationRegistry)
+			.observe(() -> this.doAdd(documents));
+	}
+
+	@Override
+	public Optional<Boolean> delete(List<String> deleteDocIds) {
+
+		VectorStoreObservationContext observationContext = this
+			.createObservationContextBuilder(VectorStoreObservationContext.Operation.DELETE.value())
+			.build();
+
+		return VectorStoreObservationDocumentation.AI_VECTOR_STORE
+			.observation(this.customObservationConvention, DEFAULT_OBSERVATION_CONVENTION, () -> observationContext,
+					this.observationRegistry)
+			.observe(() -> this.doDelete(deleteDocIds));
+	}
+
+	@Override
+	public List<Document> similaritySearch(SearchRequest request) {
+
+		VectorStoreObservationContext searchObservationContext = this
+			.createObservationContextBuilder(VectorStoreObservationContext.Operation.QUERY.value())
+			.withQueryRequest(request)
+			.build();
+
+		return VectorStoreObservationDocumentation.AI_VECTOR_STORE
+			.observation(this.customObservationConvention, DEFAULT_OBSERVATION_CONVENTION,
+					() -> searchObservationContext, this.observationRegistry)
+			.observe(() -> {
+				var documents = this.doSimilaritySearch(request);
+				searchObservationContext.setQueryResponse(documents);
+				return documents;
+			});
+	}
+
+	abstract public void doAdd(List<Document> documents);
+
+	abstract public Optional<Boolean> doDelete(List<String> idList);
+
+	abstract public List<Document> doSimilaritySearch(SearchRequest request);
+
+	abstract public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName);
+
+}