Spring/spring-ai
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Add observability support to VectorStore

Browse Source 
Implementation:
- Introduce AbstractObservationVectorStore with instrumentation for add, delete, and similaritySearch methods
- Create VectorStoreObservationContext to capture operation details
- Implement DefaultVectorStoreObservationConvention for naming and tagging
- Add VectorStoreObservationDocumentation for defining observation keys
- Create VectorStoreObservationAutoConfiguration for auto-configuring observations
- Add VectorStoreObservationProperties to control optional observation content filters
- Update VectorStore interface with getName() method
- Modify PgVectorStore and SimpleVectorStore to extend AbstractObservationVectorStore
- Add vector_store Spring AI kind

Filters:
- Implement VectorStoreQueryResponseObservationFilter
- Add VectorStoreDeleteRequestContentObservationFilter and VectorStoreAddRequestContentObservationFilter

Enhancements:
- Update PgVectorStoreAutoConfiguration to support observations
- Add observation support to PgVectorStore's Builder
- Add VectorStoreObservationContext.Operation enum with ADD, DELETE, and QUERY options

Tests:
- Add tests for VectorStore context, convention, and filters
- Add VectorStoreObservationAutoConfiguration tests
- Add PgVectorObservationIT

Resolves #1205
This commit is contained in:
Christian Tzolov
2024-08-11 08:30:02 +02:00
committed by Mark Pollack
parent 72369d515f
commit bc3f9acb86
23 changed files with 1620 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -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;

39
spring-ai-core/src/main/java/org/springframework/ai/observation/conventions/VectorStoreProvider.java Normal file
View File

@@ -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;
}
}

54
spring-ai-core/src/main/java/org/springframework/ai/observation/conventions/VectorStoreSimilarityMetric.java Normal file
View File

@@ -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;
}
}

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

@@ -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;
@@ -32,6 +22,8 @@ import java.io.OutputStream;
import java.io.OutputStreamWriter;
import java.io.Writer;
import java.nio.charset.StandardCharsets;
import java.nio.file.FileAlreadyExistsException;
import java.nio.file.Files;
import java.util.Comparator;
import java.util.HashMap;
import java.util.List;
@@ -40,6 +32,21 @@ import java.util.Objects;
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 +62,7 @@ import java.util.concurrent.ConcurrentHashMap;
* @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 +76,7 @@ public class SimpleVectorStore implements VectorStore {
}
@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,7 +86,7 @@ public class SimpleVectorStore implements VectorStore {
}
@Override
public Optional<Boolean> delete(List<String> idList) {
public Optional<Boolean> doDelete(List<String> idList) {
for (String id : idList) {
this.store.remove(id);
}
@@ -87,7 +94,7 @@ public class SimpleVectorStore implements VectorStore {
}
@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!");
@@ -114,7 +121,15 @@ public class SimpleVectorStore implements VectorStore {
try {
if (!file.exists()) {
logger.info("Creating new vector store file: {}", file);
file.createNewFile();
try {
Files.createFile(file.toPath());
}
catch (FileAlreadyExistsException e) {
throw new RuntimeException("File already exists: " + file, e);
}
catch (IOException e) {
throw new RuntimeException("Failed to create new file: " + file + ". Reason: " + e.getMessage(), e);
}
}
else {
logger.info("Overwriting existing vector store file: {}", file);
@@ -247,4 +262,13 @@ public class SimpleVectorStore implements VectorStore {
}
@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());
}
}

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

@@ -31,6 +31,10 @@ import org.springframework.ai.document.DocumentWriter;
*/
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

107
spring-ai-core/src/main/java/org/springframework/ai/vectorstore/observation/AbstractObservationVectorStore.java Normal file
View File

@@ -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;
});
}
public abstract void doAdd(List<Document> documents);
public abstract Optional<Boolean> doDelete(List<String> idList);
public abstract List<Document> doSimilaritySearch(SearchRequest request);
public abstract VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName);
}

171
spring-ai-core/src/main/java/org/springframework/ai/vectorstore/observation/DefaultVectorStoreObservationConvention.java Normal file
View File

@@ -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(), 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() {
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;
}
}

223
spring-ai-core/src/main/java/org/springframework/ai/vectorstore/observation/VectorStoreObservationContext.java Normal file
View File

@@ -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<Document> 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<Document> getQueryResponse() {
return this.queryResponse;
}
public void setQueryResponse(List<Document> 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<Document> 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;
}
}
}

33
spring-ai-core/src/main/java/org/springframework/ai/vectorstore/observation/VectorStoreObservationConvention.java Normal file
View File

@@ -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<VectorStoreObservationContext> {
@Override
default boolean supportsContext(Observation.Context context) {
return context instanceof VectorStoreObservationContext;
}
}

179
spring-ai-core/src/main/java/org/springframework/ai/vectorstore/observation/VectorStoreObservationDocumentation.java Normal file
View File

@@ -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<? extends ObservationConvention<? extends Observation.Context>> 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";
}
}
}
}

55
spring-ai-core/src/main/java/org/springframework/ai/vectorstore/observation/VectorStoreQueryResponseObservationFilter.java Normal file
View File

@@ -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;
}
}

129
spring-ai-core/src/test/java/org/springframework/ai/vectorstore/observation/DefaultVectorStoreObservationConventionTests.java Normal file
View File

@@ -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<Document> 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));
}
}

52
spring-ai-core/src/test/java/org/springframework/ai/vectorstore/observation/VectorStoreObservationContextTests.java Normal file
View File

@@ -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");
}
}

71
spring-ai-core/src/test/java/org/springframework/ai/vectorstore/observation/VectorStoreQueryResponseObservationFilterTests.java Normal file
View File

@@ -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<Document> 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\"]"));
}
}

53
spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/observation/VectorStoreObservationAutoConfiguration.java Normal file
View File

@@ -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();
}
}

44
spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/observation/VectorStoreObservationProperties.java Normal file
View File

@@ -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;
}
}

22
spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/observation/package-info.java Normal file
View File

@@ -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;

10
spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/pgvector/PgVectorStoreAutoConfiguration.java
View File

@@ -19,6 +19,8 @@ import javax.sql.DataSource;
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.boot.context.properties.EnableConfigurationProperties
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> observationRegistry,
ObjectProvider<VectorStoreObservationConvention> customSearchObservationConvention) {
var initializeSchema = properties.isInitializeSchema();
return new PgVectorStore.Builder(jdbcTemplate, embeddingModel).withSchemaName(properties.getSchemaName())
@@ -50,6 +56,8 @@ public class PgVectorStoreAutoConfiguration {
.withRemoveExistingVectorStoreTable(properties.isRemoveExistingVectorStoreTable())
.withIndexType(properties.getIndexType())
.withInitializeSchema(initializeSchema)
.withObservationRegistry(observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP))
.withSearchObservationConvention(customSearchObservationConvention.getIfAvailable(() -> null))
.build();
}

1
spring-ai-spring-boot-autoconfigure/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports
View File

@@ -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

50
spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/observation/VectorStoreObservationAutoConfigurationTests.java Normal file
View File

@@ -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);
});
}
}

6
vector-stores/spring-ai-pgvector-store/pom.xml
View File

@@ -88,6 +88,12 @@
<scope>test</scope>
</dependency>
<dependency>
<groupId>io.micrometer</groupId>
<artifactId>micrometer-observation-test</artifactId>
<scope>test</scope>
</dependency>
</dependencies>
</project>

92
vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/PgVectorStore.java
View File

@@ -28,7 +28,12 @@ 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.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.core.JsonProcessingException;
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 @@ import com.pgvector.PGvector;
* @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);
@@ -58,13 +65,13 @@ public class PgVectorStore implements VectorStore, InitializingBean {
public static final int INVALID_EMBEDDING_DIMENSION = -1;
public final static String DEFAULT_TABLE_NAME = "vector_store";
public static final String DEFAULT_TABLE_NAME = "vector_store";
public final static String DEFAULT_VECTOR_INDEX_NAME = "spring_ai_vector_index";
public static final String DEFAULT_VECTOR_INDEX_NAME = "spring_ai_vector_index";
public final static String DEFAULT_SCHEMA_NAME = "public";
public static final String DEFAULT_SCHEMA_NAME = "public";
public final static boolean DEFAULT_SCHEMA_VALIDATION = false;
public static final boolean DEFAULT_SCHEMA_VALIDATION = false;
public final FilterExpressionConverter filterExpressionConverter = new PgVectorFilterExpressionConverter();
@@ -124,10 +131,23 @@ public class PgVectorStore implements VectorStore, InitializingBean {
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: {}",
this.vectorTableName + " is empty" + (null == vectorTableName || vectorTableName.isEmpty()));
logger.info("Using the vector table name: {}. Is empty: {}", this.vectorTableName,
(vectorTableName == null || vectorTableName.isEmpty()));
this.vectorIndexName = this.vectorTableName.equals(DEFAULT_TABLE_NAME) ? DEFAULT_VECTOR_INDEX_NAME
: this.vectorTableName + "_index";
@@ -150,7 +170,7 @@ public class PgVectorStore implements VectorStore, InitializingBean {
}
@Override
public void add(List<Document> documents) {
public void doAdd(List<Document> documents) {
int size = documents.size();
@@ -195,17 +215,8 @@ public class PgVectorStore implements VectorStore, InitializingBean {
}
}
private float[] toFloatArray(List<Float> embedding) {
float[] embeddingFloat = new float[embedding.size()];
int i = 0;
for (Float d : embedding) {
embeddingFloat[i++] = d.floatValue();
}
return embeddingFloat;
}
@Override
public Optional<Boolean> delete(List<String> idList) {
public Optional<Boolean> doDelete(List<String> idList) {
int updateCount = 0;
for (String id : idList) {
int count = jdbcTemplate.update("DELETE FROM " + getFullyQualifiedTableName() + " WHERE id = ?",
@@ -217,7 +228,7 @@ public class PgVectorStore implements VectorStore, InitializingBean {
}
@Override
public List<Document> similaritySearch(SearchRequest request) {
public List<Document> doSimilaritySearch(SearchRequest request) {
String nativeFilterExpression = (request.getFilterExpression() != null)
? this.filterExpressionConverter.convertExpression(request.getFilterExpression()) : "";
@@ -276,7 +287,7 @@ public class PgVectorStore implements VectorStore, InitializingBean {
}
if (!this.initializeSchema) {
logger.debug("Skipping the schema initialization for the table: " + this.getFullyQualifiedTableName());
logger.debug("Skipping the schema initialization for the table: {}", this.getFullyQualifiedTableName());
return;
}
@@ -484,6 +495,11 @@ public class PgVectorStore implements VectorStore, InitializingBean {
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 +549,46 @@ public class PgVectorStore implements VectorStore, InitializingBean {
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<PgDistanceType, VectorStoreSimilarityMetric> 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();
}
}

203
vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorObservationIT.java Normal file
View File

@@ -0,0 +1,203 @@
/*
* 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 instruAbstractObservationVectorStorementation 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<Document> 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<Document> 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")));
}
}
}