diff --git a/spring-ai-core/src/main/java/org/springframework/ai/observation/conventions/AiObservationAttributes.java b/spring-ai-core/src/main/java/org/springframework/ai/observation/conventions/AiObservationAttributes.java index 29566449380..9e7aec175a9 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/observation/conventions/AiObservationAttributes.java +++ b/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; diff --git a/spring-ai-core/src/main/java/org/springframework/ai/observation/conventions/VectorStoreProvider.java b/spring-ai-core/src/main/java/org/springframework/ai/observation/conventions/VectorStoreProvider.java new file mode 100644 index 00000000000..1ec2ae534df --- /dev/null +++ b/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; + } + +} diff --git a/spring-ai-core/src/main/java/org/springframework/ai/observation/conventions/VectorStoreSimilarityMetric.java b/spring-ai-core/src/main/java/org/springframework/ai/observation/conventions/VectorStoreSimilarityMetric.java new file mode 100644 index 00000000000..6b95fa941d9 --- /dev/null +++ b/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; + } + +} diff --git a/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/SimpleVectorStore.java b/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/SimpleVectorStore.java index 7026fa8675b..b68e2bb8fd6 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/SimpleVectorStore.java +++ b/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 documents) { + public void doAdd(List documents) { for (Document document : documents) { logger.info("Calling EmbeddingModel for document id = {}", document.getId()); float[] embedding = this.embeddingModel.embed(document); @@ -79,7 +84,7 @@ public void add(List documents) { } @Override - public Optional delete(List idList) { + public Optional doDelete(List idList) { for (String id : idList) { this.store.remove(id); } @@ -87,7 +92,7 @@ public Optional delete(List idList) { } @Override - public List similaritySearch(SearchRequest request) { + public List 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()); + } + } diff --git a/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/VectorStore.java b/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/VectorStore.java index 8c4ff0c2a9c..dadcdda3b18 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/VectorStore.java +++ b/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 diff --git a/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/observation/AbstractObservationVectorStore.java b/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/observation/AbstractObservationVectorStore.java new file mode 100644 index 00000000000..c0e09886734 --- /dev/null +++ b/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 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 delete(List 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 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 documents); + + abstract public Optional doDelete(List idList); + + abstract public List doSimilaritySearch(SearchRequest request); + + abstract public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName); + +} \ No newline at end of file diff --git a/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/observation/DefaultVectorStoreObservationConvention.java b/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/observation/DefaultVectorStoreObservationConvention.java new file mode 100644 index 00000000000..69c4192739c --- /dev/null +++ b/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/observation/DefaultVectorStoreObservationConvention.java @@ -0,0 +1,171 @@ +/* +* 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 org.springframework.ai.vectorstore.observation.VectorStoreObservationDocumentation.HighCardinalityKeyNames; +import org.springframework.ai.vectorstore.observation.VectorStoreObservationDocumentation.LowCardinalityKeyNames; +import org.springframework.lang.Nullable; +import org.springframework.util.StringUtils; + +import io.micrometer.common.KeyValue; +import io.micrometer.common.KeyValues; + +/** + * @author Christian Tzolov + * @since 1.0.0 + */ + +public class DefaultVectorStoreObservationConvention implements VectorStoreObservationConvention { + + public static final String DEFAULT_NAME = "spring.ai.vector.store"; + + private static final String VECTOR_STORE_SPRING_AI_KIND = "vector_store"; + + private static final KeyValue DIMENSIONS_NONE = KeyValue.of(HighCardinalityKeyNames.DIMENSIONS, + KeyValue.NONE_VALUE); + + private static final KeyValue QUERY_NONE = KeyValue.of(HighCardinalityKeyNames.QUERY, KeyValue.NONE_VALUE); + + private static final KeyValue METADATA_FILTER_NONE = KeyValue.of(HighCardinalityKeyNames.QUERY_METADATA_FILTER, + KeyValue.NONE_VALUE); + + private static final KeyValue TOP_K_NONE = KeyValue.of(HighCardinalityKeyNames.TOP_K, KeyValue.NONE_VALUE); + + private static final KeyValue SIMILARITY_METRIC_NONE = KeyValue.of(HighCardinalityKeyNames.SIMILARITY_METRIC, + KeyValue.NONE_VALUE); + + private static final KeyValue COLLECTION_NAME_NONE = KeyValue.of(HighCardinalityKeyNames.COLLECTION_NAME, + KeyValue.NONE_VALUE); + + private static final KeyValue NAMESPACE_NONE = KeyValue.of(HighCardinalityKeyNames.NAMESPACE, KeyValue.NONE_VALUE); + + private static final KeyValue FIELD_NAME_NONE = KeyValue.of(HighCardinalityKeyNames.FIELD_NAME, + KeyValue.NONE_VALUE); + + private static final KeyValue INDEX_NAME_NONE = KeyValue.of(HighCardinalityKeyNames.INDEX_NAME, + KeyValue.NONE_VALUE); + + private final String name; + + public DefaultVectorStoreObservationConvention() { + this(DEFAULT_NAME); + } + + public DefaultVectorStoreObservationConvention(String name) { + this.name = name; + } + + @Override + public String getName() { + return this.name; + } + + @Override + @Nullable + public String getContextualName(VectorStoreObservationContext context) { + return "%s %s %s".formatted(VECTOR_STORE_SPRING_AI_KIND, context.getDatabaseSystem(), + context.getOperationName()); + } + + @Override + public KeyValues getLowCardinalityKeyValues(VectorStoreObservationContext context) { + return KeyValues.of(springAiKind(context), dbSystem(context), dbOperationName(context)); + } + + @Override + public KeyValues getHighCardinalityKeyValues(VectorStoreObservationContext context) { + return KeyValues.of(query(context), metadataFilter(context), topK(context), dimensions(context), + similarityMetric(context), collectionName(context), namespace(context), fieldName(context), + indexName(context)); + } + + protected KeyValue springAiKind(VectorStoreObservationContext context) { + return KeyValue.of(LowCardinalityKeyNames.SPRING_AI_KIND, VECTOR_STORE_SPRING_AI_KIND); + } + + protected KeyValue dbSystem(VectorStoreObservationContext context) { + return KeyValue.of(LowCardinalityKeyNames.DB_SYSTEM, context.getDatabaseSystem()); + } + + protected KeyValue dbOperationName(VectorStoreObservationContext context) { + return KeyValue.of(LowCardinalityKeyNames.DB_OPERATION_NAME, context.getOperationName()); + } + + protected KeyValue dimensions(VectorStoreObservationContext context) { + if (context.getDimensions() > 0) { + return KeyValue.of(HighCardinalityKeyNames.DIMENSIONS, "" + context.getDimensions()); + } + return DIMENSIONS_NONE; + } + + protected KeyValue query(VectorStoreObservationContext context) { + if (context.getQueryRequest() != null && StringUtils.hasText(context.getQueryRequest().getQuery())) { + return KeyValue.of(HighCardinalityKeyNames.QUERY, "" + context.getQueryRequest().getQuery()); + } + return QUERY_NONE; + } + + protected KeyValue metadataFilter(VectorStoreObservationContext context) { + if (context.getQueryRequest() != null && context.getQueryRequest().getFilterExpression() != null) { + return KeyValue.of(HighCardinalityKeyNames.QUERY_METADATA_FILTER, + "" + context.getQueryRequest().getFilterExpression().toString()); + } + return METADATA_FILTER_NONE; + } + + protected KeyValue topK(VectorStoreObservationContext context) { + if (context.getQueryRequest() != null && context.getQueryRequest().getTopK() > 0) { + return KeyValue.of(HighCardinalityKeyNames.TOP_K, "" + context.getQueryRequest().getTopK()); + } + return TOP_K_NONE; + } + + protected KeyValue similarityMetric(VectorStoreObservationContext context) { + if (StringUtils.hasText(context.getSimilarityMetric())) { + return KeyValue.of(HighCardinalityKeyNames.SIMILARITY_METRIC, context.getSimilarityMetric()); + } + return SIMILARITY_METRIC_NONE; + } + + protected KeyValue collectionName(VectorStoreObservationContext context) { + if (StringUtils.hasText(context.getCollectionName())) { + return KeyValue.of(HighCardinalityKeyNames.COLLECTION_NAME, context.getCollectionName()); + } + return COLLECTION_NAME_NONE; + } + + protected KeyValue namespace(VectorStoreObservationContext context) { + if (StringUtils.hasText(context.getNamespace())) { + return KeyValue.of(HighCardinalityKeyNames.NAMESPACE, context.getNamespace()); + } + return NAMESPACE_NONE; + } + + protected KeyValue fieldName(VectorStoreObservationContext context) { + if (StringUtils.hasText(context.getFieldName())) { + return KeyValue.of(HighCardinalityKeyNames.FIELD_NAME, context.getFieldName()); + } + return FIELD_NAME_NONE; + } + + protected KeyValue indexName(VectorStoreObservationContext context) { + if (StringUtils.hasText(context.getIndexName())) { + return KeyValue.of(HighCardinalityKeyNames.INDEX_NAME, context.getIndexName()); + } + return INDEX_NAME_NONE; + } + +} \ No newline at end of file diff --git a/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/observation/VectorStoreObservationContext.java b/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/observation/VectorStoreObservationContext.java new file mode 100644 index 00000000000..4017cc9bc7c --- /dev/null +++ b/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/observation/VectorStoreObservationContext.java @@ -0,0 +1,223 @@ +/* +* 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 org.springframework.ai.document.Document; +import org.springframework.ai.vectorstore.SearchRequest; +import org.springframework.util.Assert; + +import io.micrometer.observation.Observation; + +/** + * @author Christian Tzolov + * @since 1.0.0 + */ + +public class VectorStoreObservationContext extends Observation.Context { + + public enum Operation { + + /** + * VectorStore delete operation. + */ + ADD("add"), + /** + * VectorStore add operation. + */ + DELETE("delete"), + /** + * VectorStore similarity search operation. + */ + QUERY("query"); + + public final String value; + + Operation(String value) { + this.value = value; + } + + public String value() { + return this.value; + } + + } + + // SEARCH + private SearchRequest queryRequest; + + private List queryResponse; + + // COMMON + private final String databaseSystem; + + private int dimensions = -1; + + private String similarityMetric = ""; + + private String collectionName = ""; + + private String namespace = ""; + + private String fieldName = ""; + + private String indexName = ""; + + private final String operationName; + + public VectorStoreObservationContext(String databaseSystem, String operationName) { + Assert.hasText(databaseSystem, "databaseSystem cannot be null or empty"); + Assert.hasText(operationName, "operationName cannot be null or empty"); + this.databaseSystem = databaseSystem; + this.operationName = operationName; + } + + public SearchRequest getQueryRequest() { + return this.queryRequest; + } + + public void setQueryRequest(SearchRequest request) { + this.queryRequest = request; + } + + public List getQueryResponse() { + return this.queryResponse; + } + + public void setQueryResponse(List documents) { + this.queryResponse = documents; + } + + public String getDatabaseSystem() { + return this.databaseSystem; + } + + public int getDimensions() { + return this.dimensions; + } + + public void setDimensions(int dimensions) { + this.dimensions = dimensions; + } + + public String getSimilarityMetric() { + return this.similarityMetric; + } + + public void setSimilarityMetric(String similarityMetric) { + this.similarityMetric = similarityMetric; + } + + public String getCollectionName() { + return this.collectionName; + } + + public void setCollectionName(String collectionName) { + this.collectionName = collectionName; + } + + public String getNamespace() { + return this.namespace; + } + + public void setNamespace(String namespace) { + this.namespace = namespace; + } + + public String getFieldName() { + return this.fieldName; + } + + public void setFieldName(String fieldName) { + this.fieldName = fieldName; + } + + public String getIndexName() { + return this.indexName; + } + + public void setIndexName(String indexName) { + this.indexName = indexName; + } + + public String getOperationName() { + return this.operationName; + } + + public static Builder builder(String databaseSystem, String operationName) { + return new Builder(databaseSystem, operationName); + } + + public static Builder builder(String databaseSystem, Operation operation) { + return builder(databaseSystem, operation.value); + } + + public static class Builder { + + private VectorStoreObservationContext context; + + public Builder(String databaseSystem, String operationName) { + this.context = new VectorStoreObservationContext(databaseSystem, operationName); + } + + public Builder withQueryRequest(SearchRequest request) { + this.context.setQueryRequest(request); + return this; + } + + public Builder withQueryResponse(List documents) { + this.context.setQueryResponse(documents); + return this; + } + + public Builder withDimensions(int dimensions) { + this.context.setDimensions(dimensions); + return this; + } + + public Builder withSimilarityMetric(String similarityMetric) { + this.context.setSimilarityMetric(similarityMetric); + return this; + } + + public Builder withCollectionName(String collectionName) { + this.context.setCollectionName(collectionName); + return this; + } + + public Builder withNamespace(String namespace) { + this.context.setNamespace(namespace); + return this; + } + + public Builder withFieldName(String fieldName) { + this.context.setFieldName(fieldName); + return this; + } + + public Builder withIndexName(String indexName) { + this.context.setIndexName(indexName); + return this; + } + + public VectorStoreObservationContext build() { + return this.context; + } + + } + +} \ No newline at end of file diff --git a/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/observation/VectorStoreObservationConvention.java b/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/observation/VectorStoreObservationConvention.java new file mode 100644 index 00000000000..9bf80d838d2 --- /dev/null +++ b/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/observation/VectorStoreObservationConvention.java @@ -0,0 +1,33 @@ +/* +* 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 io.micrometer.observation.Observation; +import io.micrometer.observation.ObservationConvention; + +/** + * @author Christian Tzolov + * @since 1.0.0 + */ + +public interface VectorStoreObservationConvention extends ObservationConvention { + + @Override + default boolean supportsContext(Observation.Context context) { + return context instanceof VectorStoreObservationContext; + } + +} \ No newline at end of file diff --git a/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/observation/VectorStoreObservationDocumentation.java b/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/observation/VectorStoreObservationDocumentation.java new file mode 100644 index 00000000000..6ae8a15174e --- /dev/null +++ b/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/observation/VectorStoreObservationDocumentation.java @@ -0,0 +1,179 @@ +/* +* 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 org.springframework.ai.observation.conventions.AiObservationAttributes; + +import io.micrometer.common.docs.KeyName; +import io.micrometer.observation.Observation; +import io.micrometer.observation.ObservationConvention; +import io.micrometer.observation.docs.ObservationDocumentation; + +/** + * @author Christian Tzolov + * @since 1.0.0 + */ +public enum VectorStoreObservationDocumentation implements ObservationDocumentation { + + /** + * Vector Store observations for clients. + */ + AI_VECTOR_STORE { + @Override + public Class> getDefaultConvention() { + return DefaultVectorStoreObservationConvention.class; + } + + @Override + public KeyName[] getLowCardinalityKeyNames() { + return LowCardinalityKeyNames.values(); + } + + @Override + public KeyName[] getHighCardinalityKeyNames() { + return HighCardinalityKeyNames.values(); + } + }; + + public enum LowCardinalityKeyNames implements KeyName { + + /** + * Spring AI kind. + */ + SPRING_AI_KIND { + @Override + public String asString() { + return "spring.ai.kind"; + } + }, + /** + * The name of the operation or command being executed. + */ + DB_OPERATION_NAME { + @Override + public String asString() { + return AiObservationAttributes.DB_OPERATION_NAME.value(); + } + }, + /** + * The database management system (DBMS) product as identified by the client + * instrumentation. + */ + DB_SYSTEM { + @Override + public String asString() { + return "db.system"; + } + }; + + } + + public enum HighCardinalityKeyNames implements KeyName { + + /** + * Similarity search response content. + */ + QUERY_RESPONSE { + @Override + public String asString() { + return "db.vector.query.response.documents"; + } + }, + /** + * The database query being executed. + */ + QUERY { + @Override + public String asString() { + return "db.vector.query.content"; + } + }, + /** + * The metadata filters used in the query. + */ + QUERY_METADATA_FILTER { + @Override + public String asString() { + return "db.vector.query.filter"; + } + }, + /** + * The metric used in similarity search. + */ + SIMILARITY_METRIC { + @Override + public String asString() { + return "db.vector.similarity_metric"; + } + }, + /** + * The top-k most similar vectors returned by a query. + */ + TOP_K { + @Override + public String asString() { + return "db.vector.query.top_k"; + } + }, + /** + * The dimension of the vector. + */ + DIMENSIONS { + @Override + public String asString() { + return "db.vector.dimension_count"; + } + }, + /** + * The name field as of the vector (e.g. a field name). + */ + FIELD_NAME { + @Override + public String asString() { + return "db.vector.name"; + } + }, + /** + * The name of a collection (table, container) within the database. + */ + COLLECTION_NAME { + @Override + public String asString() { + return "db.collection.name"; + } + }, + /** + * The namespace of the database. + */ + NAMESPACE { + @Override + public String asString() { + return "db.namespace"; + } + }, + /** + * The index name used in the query. + */ + INDEX_NAME { + @Override + public String asString() { + return "db.index.name"; + } + } + + } + +} \ No newline at end of file diff --git a/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/observation/VectorStoreQueryResponseObservationFilter.java b/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/observation/VectorStoreQueryResponseObservationFilter.java new file mode 100644 index 00000000000..0855d278a21 --- /dev/null +++ b/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/observation/VectorStoreQueryResponseObservationFilter.java @@ -0,0 +1,55 @@ +/* + * Copyright 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.StringJoiner; + +import org.springframework.util.CollectionUtils; + +import io.micrometer.observation.Observation; +import io.micrometer.observation.ObservationFilter; + +/** + * An {@link ObservationFilter} to include the Vector Store search response content in the + * observation. + * + * @author Christian Tzolov + * @since 1.0.0 + */ +public class VectorStoreQueryResponseObservationFilter implements ObservationFilter { + + @Override + public Observation.Context map(Observation.Context context) { + + if (!(context instanceof VectorStoreObservationContext observationContext)) { + return context; + } + + if (CollectionUtils.isEmpty(observationContext.getQueryResponse())) { + return observationContext; + } + + StringJoiner joiner = new StringJoiner(", ", "[", "]"); + observationContext.getQueryResponse().forEach(document -> joiner.add("\"" + document.getContent() + "\"")); + + observationContext + .addHighCardinalityKeyValue(VectorStoreObservationDocumentation.HighCardinalityKeyNames.QUERY_RESPONSE + .withValue(joiner.toString())); + + return observationContext; + } + +} diff --git a/spring-ai-core/src/test/java/org/springframework/ai/vectorstore/observation/DefaultVectorStoreObservationConventionTests.java b/spring-ai-core/src/test/java/org/springframework/ai/vectorstore/observation/DefaultVectorStoreObservationConventionTests.java new file mode 100644 index 00000000000..d92366849c3 --- /dev/null +++ b/spring-ai-core/src/test/java/org/springframework/ai/vectorstore/observation/DefaultVectorStoreObservationConventionTests.java @@ -0,0 +1,129 @@ +/* + * Copyright 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 static org.assertj.core.api.Assertions.assertThat; + +import java.util.List; + +import org.junit.jupiter.api.Test; +import org.springframework.ai.document.Document; +import org.springframework.ai.vectorstore.SearchRequest; +import org.springframework.ai.vectorstore.observation.VectorStoreObservationDocumentation.HighCardinalityKeyNames; +import org.springframework.ai.vectorstore.observation.VectorStoreObservationDocumentation.LowCardinalityKeyNames; + +import io.micrometer.common.KeyValue; +import io.micrometer.observation.Observation; + +/** + * Unit tests for {@link DefaultVectorStoreObservationConvention}. + * + * @author Christian Tzolov + */ +class DefaultVectorStoreObservationConventionTests { + + private final DefaultVectorStoreObservationConvention observationConvention = new DefaultVectorStoreObservationConvention(); + + @Test + void shouldHaveName() { + assertThat(this.observationConvention.getName()) + .isEqualTo(DefaultVectorStoreObservationConvention.DEFAULT_NAME); + } + + @Test + void shouldHaveContextualName() { + VectorStoreObservationContext observationContext = VectorStoreObservationContext + .builder("my-database", VectorStoreObservationContext.Operation.QUERY) + .build(); + assertThat(this.observationConvention.getContextualName(observationContext)) + .isEqualTo("vector_store my-database query"); + } + + @Test + void supportsOnlyVectorStoreObservationContext() { + VectorStoreObservationContext observationContext = VectorStoreObservationContext + .builder("my-database", VectorStoreObservationContext.Operation.QUERY) + .build(); + assertThat(this.observationConvention.supportsContext(observationContext)).isTrue(); + assertThat(this.observationConvention.supportsContext(new Observation.Context())).isFalse(); + } + + @Test + void shouldHaveRequiredKeyValues() { + VectorStoreObservationContext observationContext = VectorStoreObservationContext + .builder("my_database", VectorStoreObservationContext.Operation.QUERY) + .build(); + assertThat(this.observationConvention.getLowCardinalityKeyValues(observationContext)).contains( + KeyValue.of(LowCardinalityKeyNames.DB_OPERATION_NAME.asString(), "query"), + KeyValue.of(LowCardinalityKeyNames.DB_SYSTEM.asString(), "my_database")); + } + + @Test + void shouldHaveOptionalKeyValues() { + + VectorStoreObservationContext observationContext = VectorStoreObservationContext + .builder("my-database", VectorStoreObservationContext.Operation.QUERY) + .withCollectionName("COLLECTION_NAME") + .withDimensions(696) + .withFieldName("FIELD_NAME") + .withIndexName("INDEX_NAME") + .withNamespace("NAMESPACE") + .withSimilarityMetric("SIMILARITY_METRIC") + .withQueryRequest(SearchRequest.query("VDB QUERY").withFilterExpression("country == 'UK' && year >= 2020")) + .build(); + + List queryResponseDocs = List.of(new Document("doc1"), new Document("doc2")); + + observationContext.setQueryResponse(queryResponseDocs); + + assertThat(this.observationConvention.getLowCardinalityKeyValues(observationContext)) + .contains(KeyValue.of(LowCardinalityKeyNames.DB_OPERATION_NAME.asString(), + VectorStoreObservationContext.Operation.QUERY.value)); + + // Optional, filter only added content + assertThat(this.observationConvention.getHighCardinalityKeyValues(observationContext)) + .doesNotContain(KeyValue.of(HighCardinalityKeyNames.QUERY_RESPONSE, "[doc1,doc2]")); + + assertThat(this.observationConvention.getHighCardinalityKeyValues(observationContext)).contains( + KeyValue.of(HighCardinalityKeyNames.COLLECTION_NAME.asString(), "COLLECTION_NAME"), + KeyValue.of(HighCardinalityKeyNames.DIMENSIONS.asString(), "696"), + KeyValue.of(HighCardinalityKeyNames.FIELD_NAME.asString(), "FIELD_NAME"), + KeyValue.of(HighCardinalityKeyNames.INDEX_NAME.asString(), "INDEX_NAME"), + KeyValue.of(HighCardinalityKeyNames.NAMESPACE.asString(), "NAMESPACE"), + KeyValue.of(HighCardinalityKeyNames.SIMILARITY_METRIC.asString(), "SIMILARITY_METRIC"), + KeyValue.of(HighCardinalityKeyNames.QUERY.asString(), "VDB QUERY"), + KeyValue.of(HighCardinalityKeyNames.QUERY_METADATA_FILTER.asString(), + "Expression[type=AND, left=Expression[type=EQ, left=Key[key=country], right=Value[value=UK]], right=Expression[type=GTE, left=Key[key=year], right=Value[value=2020]]]")); + } + + @Test + void shouldHaveMissingKeyValues() { + VectorStoreObservationContext observationContext = VectorStoreObservationContext + .builder("my-database", VectorStoreObservationContext.Operation.QUERY) + .build(); + + assertThat(this.observationConvention.getHighCardinalityKeyValues(observationContext)).contains( + KeyValue.of(HighCardinalityKeyNames.COLLECTION_NAME.asString(), KeyValue.NONE_VALUE), + KeyValue.of(HighCardinalityKeyNames.DIMENSIONS.asString(), KeyValue.NONE_VALUE), + KeyValue.of(HighCardinalityKeyNames.FIELD_NAME.asString(), KeyValue.NONE_VALUE), + KeyValue.of(HighCardinalityKeyNames.INDEX_NAME.asString(), KeyValue.NONE_VALUE), + KeyValue.of(HighCardinalityKeyNames.NAMESPACE.asString(), KeyValue.NONE_VALUE), + KeyValue.of(HighCardinalityKeyNames.SIMILARITY_METRIC.asString(), KeyValue.NONE_VALUE), + KeyValue.of(HighCardinalityKeyNames.QUERY.asString(), KeyValue.NONE_VALUE), + KeyValue.of(HighCardinalityKeyNames.QUERY_METADATA_FILTER.asString(), KeyValue.NONE_VALUE)); + } + +} diff --git a/spring-ai-core/src/test/java/org/springframework/ai/vectorstore/observation/VectorStoreObservationContextTests.java b/spring-ai-core/src/test/java/org/springframework/ai/vectorstore/observation/VectorStoreObservationContextTests.java new file mode 100644 index 00000000000..06d54315135 --- /dev/null +++ b/spring-ai-core/src/test/java/org/springframework/ai/vectorstore/observation/VectorStoreObservationContextTests.java @@ -0,0 +1,52 @@ +/* + * Copyright 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 static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +import org.junit.jupiter.api.Test; + +/** + * Unit tests for {@link VectorStoreObservationContext}. + * + * @author Christian Tzolov + */ +class VectorStoreObservationContextTests { + + @Test + void whenMandatoryFieldsThenReturn() { + var observationContext = VectorStoreObservationContext + .builder("db", VectorStoreObservationContext.Operation.ADD) + .build(); + assertThat(observationContext).isNotNull(); + } + + @Test + void whenDbSystemIsNullThenThrow() { + assertThatThrownBy(() -> VectorStoreObservationContext.builder(null, "delete").build()) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("databaseSystem cannot be null or empty"); + } + + @Test + void whenOperationNameIsNullThenThrow() { + assertThatThrownBy(() -> VectorStoreObservationContext.builder("Db", "").build()) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("operationName cannot be null or empty"); + } + +} diff --git a/spring-ai-core/src/test/java/org/springframework/ai/vectorstore/observation/VectorStoreQueryResponseObservationFilterTests.java b/spring-ai-core/src/test/java/org/springframework/ai/vectorstore/observation/VectorStoreQueryResponseObservationFilterTests.java new file mode 100644 index 00000000000..08ab42835c8 --- /dev/null +++ b/spring-ai-core/src/test/java/org/springframework/ai/vectorstore/observation/VectorStoreQueryResponseObservationFilterTests.java @@ -0,0 +1,71 @@ +/* + * Copyright 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 static org.assertj.core.api.Assertions.assertThat; + +import java.util.List; + +import org.junit.jupiter.api.Test; +import org.springframework.ai.document.Document; +import org.springframework.ai.vectorstore.observation.VectorStoreObservationDocumentation.HighCardinalityKeyNames; + +import io.micrometer.common.KeyValue; +import io.micrometer.observation.Observation; + +/** + * Unit tests for {@link VectorStoreQueryResponseObservationFilter}. + * + * @author Christian Tzolov + */ +class VectorStoreQueryResponseObservationFilterTests { + + private final VectorStoreQueryResponseObservationFilter observationFilter = new VectorStoreQueryResponseObservationFilter(); + + @Test + void whenNotSupportedObservationContextThenReturnOriginalContext() { + var expectedContext = new Observation.Context(); + var actualContext = observationFilter.map(expectedContext); + + assertThat(actualContext).isEqualTo(expectedContext); + } + + @Test + void whenEmptyQueryResponseThenReturnOriginalContext() { + var expectedContext = VectorStoreObservationContext.builder("db", VectorStoreObservationContext.Operation.ADD) + .build(); + + var actualContext = observationFilter.map(expectedContext); + + assertThat(actualContext).isEqualTo(expectedContext); + } + + @Test + void whenNonEmptyQueryResponseThenAugmentContext() { + var expectedContext = VectorStoreObservationContext.builder("db", VectorStoreObservationContext.Operation.ADD) + .build(); + + List queryResponseDocs = List.of(new Document("doc1"), new Document("doc2")); + + expectedContext.setQueryResponse(queryResponseDocs); + + var augmentedContext = observationFilter.map(expectedContext); + + assertThat(augmentedContext.getHighCardinalityKeyValues()) + .contains(KeyValue.of(HighCardinalityKeyNames.QUERY_RESPONSE.asString(), "[\"doc1\", \"doc2\"]")); + } + +} diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/observation/VectorStoreObservationAutoConfiguration.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/observation/VectorStoreObservationAutoConfiguration.java new file mode 100644 index 00000000000..46feafc9448 --- /dev/null +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/observation/VectorStoreObservationAutoConfiguration.java @@ -0,0 +1,53 @@ +/* + * Copyright 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.autoconfigure.vectorstore.observation; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import org.springframework.ai.vectorstore.VectorStore; +import org.springframework.ai.vectorstore.observation.VectorStoreQueryResponseObservationFilter; +import org.springframework.boot.autoconfigure.AutoConfiguration; +import org.springframework.boot.autoconfigure.condition.ConditionalOnClass; +import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean; +import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty; +import org.springframework.boot.context.properties.EnableConfigurationProperties; +import org.springframework.context.annotation.Bean; + +/** + * Auto-configuration for Spring AI vector store observations. + * + * @author Christian Tzolov + * @since 1.0.0 + */ +@AutoConfiguration( + afterName = { "org.springframework.boot.actuate.autoconfigure.observation.ObservationAutoConfiguration" }) +@ConditionalOnClass(VectorStore.class) +@EnableConfigurationProperties({ VectorStoreObservationProperties.class }) +public class VectorStoreObservationAutoConfiguration { + + private static final Logger logger = LoggerFactory.getLogger(VectorStoreObservationAutoConfiguration.class); + + @Bean + @ConditionalOnMissingBean + @ConditionalOnProperty(prefix = VectorStoreObservationProperties.CONFIG_PREFIX, name = "include-query-response", + havingValue = "true") + VectorStoreQueryResponseObservationFilter vectorStoreQueryResponseContentObservationFilter() { + logger.warn( + "You have enabled the inclusion of the query response content in the observations, with the risk of exposing sensitive or private information. Please, be careful!"); + return new VectorStoreQueryResponseObservationFilter(); + } + +} diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/observation/VectorStoreObservationProperties.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/observation/VectorStoreObservationProperties.java new file mode 100644 index 00000000000..33589abad26 --- /dev/null +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/observation/VectorStoreObservationProperties.java @@ -0,0 +1,44 @@ +/* + * Copyright 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.autoconfigure.vectorstore.observation; + +import org.springframework.boot.context.properties.ConfigurationProperties; + +/** + * Configuration properties for vector store observations. + * + * @author Christian Tzolov + * @since 1.0.0 + */ +@ConfigurationProperties(VectorStoreObservationProperties.CONFIG_PREFIX) +public class VectorStoreObservationProperties { + + public static final String CONFIG_PREFIX = "spring.ai.vectorstore.observations"; + + /** + * Whether to include the search response content in the observations. + */ + private boolean includeQueryResponse = false; + + public boolean isIncludeQueryResponse() { + return this.includeQueryResponse; + } + + public void setIncludeQueryResponse(boolean includeQueryResponse) { + this.includeQueryResponse = includeQueryResponse; + } + +} diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/observation/package-info.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/observation/package-info.java new file mode 100644 index 00000000000..347dd6a3ef9 --- /dev/null +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/observation/package-info.java @@ -0,0 +1,22 @@ +/* + * Copyright 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. + */ + +@NonNullApi +@NonNullFields +package org.springframework.ai.autoconfigure.vectorstore.observation; + +import org.springframework.lang.NonNullApi; +import org.springframework.lang.NonNullFields; diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/pgvector/PgVectorStoreAutoConfiguration.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/pgvector/PgVectorStoreAutoConfiguration.java index 7a81d9a5456..4c8746527b0 100644 --- a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/pgvector/PgVectorStoreAutoConfiguration.java +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/pgvector/PgVectorStoreAutoConfiguration.java @@ -19,6 +19,8 @@ import org.springframework.ai.embedding.EmbeddingModel; import org.springframework.ai.vectorstore.PgVectorStore; +import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention; +import org.springframework.beans.factory.ObjectProvider; import org.springframework.boot.autoconfigure.AutoConfiguration; import org.springframework.boot.autoconfigure.condition.ConditionalOnClass; import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean; @@ -27,6 +29,8 @@ import org.springframework.context.annotation.Bean; import org.springframework.jdbc.core.JdbcTemplate; +import io.micrometer.observation.ObservationRegistry; + /** * @author Christian Tzolov * @author Josh Long @@ -39,7 +43,9 @@ public class PgVectorStoreAutoConfiguration { @Bean @ConditionalOnMissingBean public PgVectorStore vectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel, - PgVectorStoreProperties properties) { + PgVectorStoreProperties properties, ObjectProvider observationRegistry, + ObjectProvider customSearchObservationConvention) { + var initializeSchema = properties.isInitializeSchema(); return new PgVectorStore.Builder(jdbcTemplate, embeddingModel).withSchemaName(properties.getSchemaName()) @@ -50,6 +56,8 @@ public PgVectorStore vectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel embed .withRemoveExistingVectorStoreTable(properties.isRemoveExistingVectorStoreTable()) .withIndexType(properties.getIndexType()) .withInitializeSchema(initializeSchema) + .withObservationRegistry(observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP)) + .withSearchObservationConvention(customSearchObservationConvention.getIfAvailable(() -> null)) .build(); } diff --git a/spring-ai-spring-boot-autoconfigure/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports b/spring-ai-spring-boot-autoconfigure/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports index 1c74d1a0a7f..9fe46cc2609 100644 --- a/spring-ai-spring-boot-autoconfigure/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports +++ b/spring-ai-spring-boot-autoconfigure/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports @@ -46,3 +46,4 @@ org.springframework.ai.autoconfigure.qianfan.QianFanAutoConfiguration org.springframework.ai.autoconfigure.minimax.MiniMaxAutoConfiguration org.springframework.ai.autoconfigure.vertexai.embedding.VertexAiEmbeddingAutoConfiguration org.springframework.ai.autoconfigure.chat.memory.cassandra.CassandraChatMemoryAutoConfiguration +org.springframework.ai.autoconfigure.vectorstore.observation.VectorStoreObservationAutoConfiguration diff --git a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/observation/VectorStoreObservationAutoConfigurationTests.java b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/observation/VectorStoreObservationAutoConfigurationTests.java new file mode 100644 index 00000000000..24386fbc0e6 --- /dev/null +++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/observation/VectorStoreObservationAutoConfigurationTests.java @@ -0,0 +1,50 @@ +/* + * Copyright 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.autoconfigure.vectorstore.observation; + +import static org.assertj.core.api.Assertions.assertThat; + +import org.junit.jupiter.api.Test; +import org.springframework.ai.vectorstore.observation.VectorStoreQueryResponseObservationFilter; +import org.springframework.boot.autoconfigure.AutoConfigurations; +import org.springframework.boot.test.context.runner.ApplicationContextRunner; + +/** + * Unit tests for {@link VectorStoreObservationAutoConfiguration}. + * + * @author Christian Tzolov + */ +class VectorStoreObservationAutoConfigurationTests { + + private final ApplicationContextRunner contextRunner = new ApplicationContextRunner() + .withConfiguration(AutoConfigurations.of(VectorStoreObservationAutoConfiguration.class)); + + @Test + void queryResponseFilterDefault() { + contextRunner.run(context -> { + assertThat(context).doesNotHaveBean(VectorStoreQueryResponseObservationFilter.class); + }); + } + + @Test + void queryResponseFilterFilterEnabled() { + contextRunner.withPropertyValues("spring.ai.vectorstore.observations.include-query-response=true") + .run(context -> { + assertThat(context).hasSingleBean(VectorStoreQueryResponseObservationFilter.class); + }); + } + +} diff --git a/vector-stores/spring-ai-pgvector-store/pom.xml b/vector-stores/spring-ai-pgvector-store/pom.xml index 0da1ef57527..2cd142626d2 100644 --- a/vector-stores/spring-ai-pgvector-store/pom.xml +++ b/vector-stores/spring-ai-pgvector-store/pom.xml @@ -88,6 +88,12 @@ test + + io.micrometer + micrometer-observation-test + test + + diff --git a/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/PgVectorStore.java b/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/PgVectorStore.java index aa7966e767f..c1d67e16410 100644 --- a/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/PgVectorStore.java +++ b/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/PgVectorStore.java @@ -28,7 +28,12 @@ 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.filter.FilterExpressionConverter; +import org.springframework.ai.vectorstore.observation.AbstractObservationVectorStore; +import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext; +import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention; import org.springframework.beans.factory.InitializingBean; import org.springframework.jdbc.core.BatchPreparedStatementSetter; import org.springframework.jdbc.core.JdbcTemplate; @@ -42,6 +47,8 @@ import com.fasterxml.jackson.databind.ObjectMapper; import com.pgvector.PGvector; +import io.micrometer.observation.ObservationRegistry; + /** * Uses the "vector_store" table to store the Spring AI vector data. The table and the * vector index will be auto-created if not available. @@ -50,7 +57,7 @@ * @author Josh Long * @author Muthukumaran Navaneethakrishnan */ -public class PgVectorStore implements VectorStore, InitializingBean { +public class PgVectorStore extends AbstractObservationVectorStore implements InitializingBean { private static final Logger logger = LoggerFactory.getLogger(PgVectorStore.class); @@ -124,6 +131,19 @@ private PgVectorStore(String schemaName, String vectorTableName, boolean vectorT JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel, int dimensions, PgDistanceType distanceType, boolean removeExistingVectorStoreTable, PgIndexType createIndexMethod, boolean initializeSchema) { + this(schemaName, vectorTableName, vectorTableValidationsEnabled, jdbcTemplate, embeddingModel, dimensions, + distanceType, removeExistingVectorStoreTable, createIndexMethod, initializeSchema, + ObservationRegistry.NOOP, null); + } + + private PgVectorStore(String schemaName, String vectorTableName, boolean vectorTableValidationsEnabled, + JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel, int dimensions, PgDistanceType distanceType, + boolean removeExistingVectorStoreTable, PgIndexType createIndexMethod, boolean initializeSchema, + ObservationRegistry observationRegistry, + VectorStoreObservationConvention customSearchObservationConvention) { + + super(observationRegistry, customSearchObservationConvention); + this.vectorTableName = (null == vectorTableName || vectorTableName.isEmpty()) ? DEFAULT_TABLE_NAME : vectorTableName.trim(); logger.info("Using the vector table name: {}", @@ -150,7 +170,7 @@ public PgDistanceType getDistanceType() { } @Override - public void add(List documents) { + public void doAdd(List documents) { int size = documents.size(); @@ -205,7 +225,7 @@ private float[] toFloatArray(List embedding) { } @Override - public Optional delete(List idList) { + public Optional doDelete(List idList) { int updateCount = 0; for (String id : idList) { int count = jdbcTemplate.update("DELETE FROM " + getFullyQualifiedTableName() + " WHERE id = ?", @@ -217,7 +237,7 @@ public Optional delete(List idList) { } @Override - public List similaritySearch(SearchRequest request) { + public List doSimilaritySearch(SearchRequest request) { String nativeFilterExpression = (request.getFilterExpression() != null) ? this.filterExpressionConverter.convertExpression(request.getFilterExpression()) : ""; @@ -484,6 +504,11 @@ public static class Builder { private boolean initializeSchema; + private ObservationRegistry observationRegistry = ObservationRegistry.NOOP; + + @Nullable + private VectorStoreObservationConvention searchObservationConvention; + // Builder constructor with mandatory parameters public Builder(JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel) { if (jdbcTemplate == null || embeddingModel == null) { @@ -533,12 +558,46 @@ public Builder withInitializeSchema(boolean initializeSchema) { return this; } + public Builder withObservationRegistry(ObservationRegistry observationRegistry) { + this.observationRegistry = observationRegistry; + return this; + } + + public Builder withSearchObservationConvention( + VectorStoreObservationConvention customSearchObservationConvention) { + this.searchObservationConvention = customSearchObservationConvention; + return this; + } + public PgVectorStore build() { return new PgVectorStore(schemaName, vectorTableName, vectorTableValidationsEnabled, jdbcTemplate, embeddingModel, dimensions, distanceType, removeExistingVectorStoreTable, indexType, - initializeSchema); + initializeSchema, observationRegistry, searchObservationConvention); } } + @Override + public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName) { + + return VectorStoreObservationContext.builder(VectorStoreProvider.PG_VECTOR.value(), operationName) + .withDimensions(this.embeddingDimensions()) + .withCollectionName(this.vectorTableName) + .withNamespace(this.schemaName) + .withSimilarityMetric(getSimilarityMetric()) + .withIndexName(this.createIndexMethod.name()); + } + + private static Map SIMILARITY_TYPE_MAPPING = Map.of( + PgDistanceType.COSINE_DISTANCE, VectorStoreSimilarityMetric.COSINE, PgDistanceType.EUCLIDEAN_DISTANCE, + VectorStoreSimilarityMetric.EUCLIDEAN, PgDistanceType.NEGATIVE_INNER_PRODUCT, + VectorStoreSimilarityMetric.DOT); + + private String getSimilarityMetric() { + if (!SIMILARITY_TYPE_MAPPING.containsKey(this.getDistanceType())) { + return this.getDistanceType().name(); + } + return SIMILARITY_TYPE_MAPPING.get(this.distanceType).value(); + } + } \ No newline at end of file diff --git a/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorObservationIT.java b/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorObservationIT.java new file mode 100644 index 00000000000..59590735fd1 --- /dev/null +++ b/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorObservationIT.java @@ -0,0 +1,202 @@ +/* + * Copyright 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; + +import static org.assertj.core.api.Assertions.assertThat; + +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.util.List; +import java.util.Map; + +import javax.sql.DataSource; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable; +import org.springframework.ai.document.Document; +import org.springframework.ai.embedding.EmbeddingModel; +import org.springframework.ai.openai.OpenAiChatModel; +import org.springframework.ai.openai.OpenAiEmbeddingModel; +import org.springframework.ai.openai.api.OpenAiApi; +import org.springframework.ai.vectorstore.PgVectorStore.PgIndexType; +import org.springframework.ai.vectorstore.observation.DefaultVectorStoreObservationConvention; +import org.springframework.ai.vectorstore.observation.VectorStoreObservationDocumentation.HighCardinalityKeyNames; +import org.springframework.ai.vectorstore.observation.VectorStoreObservationDocumentation.LowCardinalityKeyNames; +import org.springframework.boot.SpringBootConfiguration; +import org.springframework.boot.autoconfigure.EnableAutoConfiguration; +import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration; +import org.springframework.boot.autoconfigure.jdbc.DataSourceProperties; +import org.springframework.boot.context.properties.ConfigurationProperties; +import org.springframework.boot.test.context.runner.ApplicationContextRunner; +import org.springframework.context.annotation.Bean; +import org.springframework.context.annotation.Primary; +import org.springframework.core.io.DefaultResourceLoader; +import org.springframework.jdbc.core.JdbcTemplate; +import org.testcontainers.containers.PostgreSQLContainer; +import org.testcontainers.junit.jupiter.Container; +import org.testcontainers.junit.jupiter.Testcontainers; + +import com.zaxxer.hikari.HikariDataSource; + +import io.micrometer.observation.ObservationRegistry; +import io.micrometer.observation.tck.TestObservationRegistry; +import io.micrometer.observation.tck.TestObservationRegistryAssert; + +/** + * Integration tests for observation instrumentation in {@link OpenAiChatModel}. + * + * @author Christian Tzolov + */ +@EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+") +@Testcontainers +public class PgVectorObservationIT { + + @Container + @SuppressWarnings("resource") + static PostgreSQLContainer postgresContainer = new PostgreSQLContainer<>("pgvector/pgvector:pg16") + .withUsername("postgres") + .withPassword("postgres"); + + List documents = List.of( + new Document(getText("classpath:/test/data/spring.ai.txt"), Map.of("meta1", "meta1")), + new Document(getText("classpath:/test/data/time.shelter.txt")), + new Document(getText("classpath:/test/data/great.depression.txt"), Map.of("meta2", "meta2"))); + + public static String getText(String uri) { + var resource = new DefaultResourceLoader().getResource(uri); + try { + return resource.getContentAsString(StandardCharsets.UTF_8); + } + catch (IOException e) { + throw new RuntimeException(e); + } + } + + private final ApplicationContextRunner contextRunner = new ApplicationContextRunner() + .withUserConfiguration(Config.class) + .withPropertyValues("test.spring.ai.vectorstore.pgvector.distanceType=COSINE_DISTANCE", + + // JdbcTemplate configuration + String.format("app.datasource.url=jdbc:postgresql://%s:%d/%s", postgresContainer.getHost(), + postgresContainer.getMappedPort(5432), "postgres"), + "app.datasource.username=postgres", "app.datasource.password=postgres", + "app.datasource.type=com.zaxxer.hikari.HikariDataSource"); + + @Test + void observationVectorStoreAddAndQueryOperations() { + + contextRunner.run(context -> { + + VectorStore vectorStore = context.getBean(VectorStore.class); + + TestObservationRegistry observationRegistry = context.getBean(TestObservationRegistry.class); + + vectorStore.add(documents); + + TestObservationRegistryAssert.assertThat(observationRegistry) + .doesNotHaveAnyRemainingCurrentObservation() + .hasObservationWithNameEqualTo(DefaultVectorStoreObservationConvention.DEFAULT_NAME) + .that() + .hasContextualNameEqualTo("vector_store pg_vector add") + .hasLowCardinalityKeyValue(LowCardinalityKeyNames.DB_OPERATION_NAME.asString(), "add") + .hasLowCardinalityKeyValue(LowCardinalityKeyNames.DB_SYSTEM.asString(), "pg_vector") + .hasLowCardinalityKeyValue(LowCardinalityKeyNames.SPRING_AI_KIND.asString(), "vector_store") + .hasHighCardinalityKeyValue(HighCardinalityKeyNames.QUERY.asString(), "none") + .hasHighCardinalityKeyValue(HighCardinalityKeyNames.DIMENSIONS.asString(), "1536") + .hasHighCardinalityKeyValue(HighCardinalityKeyNames.COLLECTION_NAME.asString(), "vector_store") + .hasHighCardinalityKeyValue(HighCardinalityKeyNames.NAMESPACE.asString(), "public") + .hasHighCardinalityKeyValue(HighCardinalityKeyNames.FIELD_NAME.asString(), "none") + .hasHighCardinalityKeyValue(HighCardinalityKeyNames.SIMILARITY_METRIC.asString(), "cosine") + .hasHighCardinalityKeyValue(HighCardinalityKeyNames.TOP_K.asString(), "none") + + .hasBeenStarted() + .hasBeenStopped(); + + observationRegistry.clear(); + + List results = vectorStore + .similaritySearch(SearchRequest.query("What is Great Depression").withTopK(1)); + + assertThat(results).isNotEmpty(); + + TestObservationRegistryAssert.assertThat(observationRegistry) + .doesNotHaveAnyRemainingCurrentObservation() + .hasObservationWithNameEqualTo(DefaultVectorStoreObservationConvention.DEFAULT_NAME) + .that() + .hasContextualNameEqualTo("vector_store pg_vector query") + .hasLowCardinalityKeyValue(LowCardinalityKeyNames.DB_OPERATION_NAME.asString(), "query") + .hasLowCardinalityKeyValue(LowCardinalityKeyNames.DB_SYSTEM.asString(), "pg_vector") + .hasLowCardinalityKeyValue(LowCardinalityKeyNames.SPRING_AI_KIND.asString(), "vector_store") + + .hasHighCardinalityKeyValue(HighCardinalityKeyNames.QUERY.asString(), "What is Great Depression") + .hasHighCardinalityKeyValue(HighCardinalityKeyNames.DIMENSIONS.asString(), "1536") + .hasHighCardinalityKeyValue(HighCardinalityKeyNames.COLLECTION_NAME.asString(), "vector_store") + .hasHighCardinalityKeyValue(HighCardinalityKeyNames.NAMESPACE.asString(), "public") + .hasHighCardinalityKeyValue(HighCardinalityKeyNames.FIELD_NAME.asString(), "none") + .hasHighCardinalityKeyValue(HighCardinalityKeyNames.SIMILARITY_METRIC.asString(), "cosine") + .hasHighCardinalityKeyValue(HighCardinalityKeyNames.TOP_K.asString(), "1") + + .hasBeenStarted() + .hasBeenStopped(); + + }); + } + + @SpringBootConfiguration + @EnableAutoConfiguration(exclude = { DataSourceAutoConfiguration.class }) + static class Config { + + @Bean + public TestObservationRegistry observationRegistry() { + return TestObservationRegistry.create(); + } + + @Bean + public VectorStore vectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel, + ObservationRegistry observationRegistry) { + return new PgVectorStore.Builder(jdbcTemplate, embeddingModel) + .withDistanceType(PgVectorStore.PgDistanceType.COSINE_DISTANCE) + .withIndexType(PgIndexType.HNSW) + .withObservationRegistry(observationRegistry) + .withInitializeSchema(true) + .build(); + } + + @Bean + public JdbcTemplate myJdbcTemplate(DataSource dataSource) { + return new JdbcTemplate(dataSource); + } + + @Bean + @Primary + @ConfigurationProperties("app.datasource") + public DataSourceProperties dataSourceProperties() { + return new DataSourceProperties(); + } + + @Bean + public HikariDataSource dataSource(DataSourceProperties dataSourceProperties) { + return dataSourceProperties.initializeDataSourceBuilder().type(HikariDataSource.class).build(); + } + + @Bean + public EmbeddingModel embeddingModel() { + return new OpenAiEmbeddingModel(new OpenAiApi(System.getenv("OPENAI_API_KEY"))); + } + + } + +}