Skip to content

Commit

Permalink
Add observability support to existing vector stores
Browse files Browse the repository at this point in the history
 Add observability support to:
 - Cassandra
 - Chroma
 - Elasticsearch
 - Milvus
 - Neo4j
 - OpenSearch
 - Qdrant
 - Redis
 - Typesense
 - Weaviate
 - Pinecone
 - Oracle
 - Gemifire
 - MongoDB
 - HanaDB

 Add autoconfiguration obsrvability for the above vector stores.
 Add integration tests for all vector stores.
  • Loading branch information
tzolov authored and markpollack committed Aug 20, 2024
1 parent 3b7522b commit 93fa2bf
Show file tree
Hide file tree
Showing 98 changed files with 4,643 additions and 300 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,23 @@ public enum VectorStoreProvider {

// @formatter:off
PG_VECTOR("pg_vector"),
SIMPLE_VECTOR_STORE("simple_vector_store");
AZURE("azure"),
CASSANDRA("cassandra"),
CHROMA("chroma"),
ELASTICSEARCH("elasticsearch"),
MILVUS("milvus"),
NEO4J("neo4j"),
OPENSEARCH("opensearch"),
QDRANT("qdrant"),
REDIS("redis"),
TYPESENSE("typesense"),
WEAVIATE("weaviate"),
PINECONE("pinecone"),
ORACLE("oracle"),
MONGODB("mongodb"),
GEMFIRE("gemfire"),
HANA("hana"),
SIMPLE("simple");

// @formatter:on
private final String value;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -40,13 +40,16 @@
import org.springframework.ai.observation.conventions.VectorStoreSimilarityMetric;
import org.springframework.ai.vectorstore.observation.AbstractObservationVectorStore;
import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext;
import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention;
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;

import io.micrometer.observation.ObservationRegistry;

/**
* SimpleVectorStore is a simple implementation of the VectorStore interface.
*
Expand All @@ -71,6 +74,14 @@ public class SimpleVectorStore extends AbstractObservationVectorStore {
protected EmbeddingModel embeddingModel;

public SimpleVectorStore(EmbeddingModel embeddingModel) {
this(embeddingModel, ObservationRegistry.NOOP, null);
}

public SimpleVectorStore(EmbeddingModel embeddingModel, ObservationRegistry observationRegistry,
VectorStoreObservationConvention customObservationConvention) {

super(observationRegistry, customObservationConvention);

Objects.requireNonNull(embeddingModel, "EmbeddingModel must not be null");
this.embeddingModel = embeddingModel;
}
Expand Down Expand Up @@ -265,7 +276,7 @@ public static float norm(float[] vector) {
@Override
public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName) {

return VectorStoreObservationContext.builder(VectorStoreProvider.SIMPLE_VECTOR_STORE.value(), operationName)
return VectorStoreObservationContext.builder(VectorStoreProvider.SIMPLE.value(), operationName)
.withDimensions(this.embeddingModel.dimensions())
.withCollectionName("in-memory-map")
.withSimilarityMetric(VectorStoreSimilarityMetric.COSINE.value());
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -38,18 +38,10 @@ public abstract class AbstractObservationVectorStore implements VectorStore {
@Nullable
private final VectorStoreObservationConvention customObservationConvention;

public AbstractObservationVectorStore() {
this(ObservationRegistry.NOOP, null);
}

public AbstractObservationVectorStore(ObservationRegistry observationRegistry) {
this(observationRegistry, null);
}

public AbstractObservationVectorStore(ObservationRegistry observationRegistry,
VectorStoreObservationConvention customSearchObservationConvention) {
VectorStoreObservationConvention customObservationConvention) {
this.observationRegistry = observationRegistry;
this.customObservationConvention = customSearchObservationConvention;
this.customObservationConvention = customObservationConvention;
}

@Override
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,9 @@ public class DefaultVectorStoreObservationConvention implements VectorStoreObser

private static final KeyValue TOP_K_NONE = KeyValue.of(HighCardinalityKeyNames.TOP_K, KeyValue.NONE_VALUE);

private static final KeyValue SIMILARITY_THRESHOLD_NONE = KeyValue.of(HighCardinalityKeyNames.SIMILARITY_THRESHOLD,
KeyValue.NONE_VALUE);

private static final KeyValue SIMILARITY_METRIC_NONE = KeyValue.of(HighCardinalityKeyNames.SIMILARITY_METRIC,
KeyValue.NONE_VALUE);

Expand Down Expand Up @@ -89,7 +92,7 @@ public KeyValues getLowCardinalityKeyValues(VectorStoreObservationContext contex
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));
indexName(context), similarityThreshold(context));
}

protected KeyValue springAiKind() {
Expand Down Expand Up @@ -133,6 +136,14 @@ protected KeyValue topK(VectorStoreObservationContext context) {
return TOP_K_NONE;
}

protected KeyValue similarityThreshold(VectorStoreObservationContext context) {
if (context.getQueryRequest() != null && context.getQueryRequest().getSimilarityThreshold() >= 0) {
return KeyValue.of(HighCardinalityKeyNames.SIMILARITY_THRESHOLD,
"" + context.getQueryRequest().getSimilarityThreshold());
}
return SIMILARITY_THRESHOLD_NONE;
}

protected KeyValue similarityMetric(VectorStoreObservationContext context) {
if (StringUtils.hasText(context.getSimilarityMetric())) {
return KeyValue.of(HighCardinalityKeyNames.SIMILARITY_METRIC, context.getSimilarityMetric());
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,6 @@
* @author Christian Tzolov
* @since 1.0.0
*/

public class VectorStoreObservationContext extends Observation.Context {

public enum Operation {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -128,6 +128,17 @@ public String asString() {
return "db.vector.query.top_k";
}
},
/**
* Similarity threshold that accepts all search scores. A threshold value of 0.0
* means any similarity is accepted or disable the similarity threshold filtering.
* A threshold value of 1.0 means an exact match is required.
*/
SIMILARITY_THRESHOLD {
@Override
public String asString() {
return "db.vector.query.similarity_threshold";
}
},
/**
* The dimension of the vector.
*/
Expand Down
6 changes: 6 additions & 0 deletions spring-ai-spring-boot-autoconfigure/pom.xml
Original file line number Diff line number Diff line change
Expand Up @@ -519,6 +519,12 @@
<scope>test</scope>
</dependency>

<dependency>
<groupId>io.micrometer</groupId>
<artifactId>micrometer-observation-test</artifactId>
<scope>test</scope>
</dependency>

</dependencies>

</project>
Original file line number Diff line number Diff line change
Expand Up @@ -19,8 +19,14 @@
import com.azure.search.documents.indexes.SearchIndexClient;
import com.azure.search.documents.indexes.SearchIndexClientBuilder;

import io.micrometer.observation.ObservationRegistry;

import java.util.List;

import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.vectorstore.azure.AzureVectorStore;
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;
Expand Down Expand Up @@ -48,9 +54,12 @@ public SearchIndexClient searchIndexClient(AzureVectorStoreProperties properties
@Bean
@ConditionalOnMissingBean
public AzureVectorStore vectorStore(SearchIndexClient searchIndexClient, EmbeddingModel embeddingModel,
AzureVectorStoreProperties properties) {
AzureVectorStoreProperties properties, ObjectProvider<ObservationRegistry> observationRegistry,
ObjectProvider<VectorStoreObservationConvention> customObservationConvention) {

var vectorStore = new AzureVectorStore(searchIndexClient, embeddingModel, properties.isInitializeSchema());
var vectorStore = new AzureVectorStore(searchIndexClient, embeddingModel, properties.isInitializeSchema(),
List.of(), observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP),
customObservationConvention.getIfAvailable(() -> null));

vectorStore.setIndexName(properties.getIndexName());

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -20,9 +20,13 @@
import com.datastax.oss.driver.api.core.CqlSession;
import com.datastax.oss.driver.api.core.config.DefaultDriverOption;

import io.micrometer.observation.ObservationRegistry;

import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.vectorstore.CassandraVectorStore;
import org.springframework.ai.vectorstore.CassandraVectorStoreConfig;
import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.boot.autoconfigure.AutoConfiguration;
import org.springframework.boot.autoconfigure.cassandra.CassandraAutoConfiguration;
import org.springframework.boot.autoconfigure.cassandra.DriverConfigLoaderBuilderCustomizer;
Expand All @@ -33,6 +37,7 @@

/**
* @author Mick Semb Wever
* @author Christian Tzolov
* @since 1.0.0
*/
@AutoConfiguration(after = CassandraAutoConfiguration.class)
Expand All @@ -43,7 +48,8 @@ public class CassandraVectorStoreAutoConfiguration {
@Bean
@ConditionalOnMissingBean
public CassandraVectorStore vectorStore(EmbeddingModel embeddingModel, CassandraVectorStoreProperties properties,
CqlSession cqlSession) {
CqlSession cqlSession, ObjectProvider<ObservationRegistry> observationRegistry,
ObjectProvider<VectorStoreObservationConvention> customObservationConvention) {

var builder = CassandraVectorStoreConfig.builder().withCqlSession(cqlSession);

Expand All @@ -61,7 +67,9 @@ public CassandraVectorStore vectorStore(EmbeddingModel embeddingModel, Cassandra
builder = builder.returnEmbeddings();
}

return CassandraVectorStore.create(builder.build(), embeddingModel);
return new CassandraVectorStore(builder.build(), embeddingModel,
observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP),
customObservationConvention.getIfAvailable(() -> null));
}

@Bean
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,8 @@
import org.springframework.ai.chroma.ChromaApi;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.vectorstore.ChromaVectorStore;
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;
Expand All @@ -29,6 +31,8 @@

import com.fasterxml.jackson.databind.ObjectMapper;

import io.micrometer.observation.ObservationRegistry;

/**
* @author Christian Tzolov
* @author Eddú Meléndez
Expand Down Expand Up @@ -72,9 +76,11 @@ else if (StringUtils.hasText(apiProperties.getUsername()) && StringUtils.hasText
@Bean
@ConditionalOnMissingBean
public ChromaVectorStore vectorStore(EmbeddingModel embeddingModel, ChromaApi chromaApi,
ChromaVectorStoreProperties storeProperties) {
ChromaVectorStoreProperties storeProperties, ObjectProvider<ObservationRegistry> observationRegistry,
ObjectProvider<VectorStoreObservationConvention> customObservationConvention) {
return new ChromaVectorStore(embeddingModel, chromaApi, storeProperties.getCollectionName(),
storeProperties.isInitializeSchema());
storeProperties.isInitializeSchema(), observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP),
customObservationConvention.getIfAvailable(() -> null));
}

static class PropertiesChromaConnectionDetails implements ChromaConnectionDetails {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,8 @@
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.vectorstore.ElasticsearchVectorStore;
import org.springframework.ai.vectorstore.ElasticsearchVectorStoreOptions;
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;
Expand All @@ -28,10 +30,13 @@
import org.springframework.context.annotation.Bean;
import org.springframework.util.StringUtils;

import io.micrometer.observation.ObservationRegistry;

/**
* @author Eddú Meléndez
* @author Wei Jiang
* @author Josh Long
* @author Christian Tzolov
* @since 1.0.0
*/

Expand All @@ -43,7 +48,8 @@ class ElasticsearchVectorStoreAutoConfiguration {
@Bean
@ConditionalOnMissingBean
ElasticsearchVectorStore vectorStore(ElasticsearchVectorStoreProperties properties, RestClient restClient,
EmbeddingModel embeddingModel) {
EmbeddingModel embeddingModel, ObjectProvider<ObservationRegistry> observationRegistry,
ObjectProvider<VectorStoreObservationConvention> customObservationConvention) {
ElasticsearchVectorStoreOptions elasticsearchVectorStoreOptions = new ElasticsearchVectorStoreOptions();

if (StringUtils.hasText(properties.getIndexName())) {
Expand All @@ -57,7 +63,8 @@ ElasticsearchVectorStore vectorStore(ElasticsearchVectorStoreProperties properti
}

return new ElasticsearchVectorStore(elasticsearchVectorStoreOptions, restClient, embeddingModel,
properties.isInitializeSchema());
properties.isInitializeSchema(), observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP),
customObservationConvention.getIfAvailable(() -> null));
}

}
Original file line number Diff line number Diff line change
Expand Up @@ -19,15 +19,20 @@
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.vectorstore.GemFireVectorStore;
import org.springframework.ai.vectorstore.GemFireVectorStoreConfig;
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;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;

import io.micrometer.observation.ObservationRegistry;

/**
* @author Geet Rawat
* @author Christian Tzolov
*/
@AutoConfiguration
@ConditionalOnClass({ GemFireVectorStore.class, EmbeddingModel.class })
Expand All @@ -45,7 +50,8 @@ GemFireVectorStoreAutoConfiguration.PropertiesGemFireConnectionDetails gemfireCo
@Bean
@ConditionalOnMissingBean
public GemFireVectorStore gemfireVectorStore(EmbeddingModel embeddingModel, GemFireVectorStoreProperties properties,
GemFireConnectionDetails gemFireConnectionDetails) {
GemFireConnectionDetails gemFireConnectionDetails, ObjectProvider<ObservationRegistry> observationRegistry,
ObjectProvider<VectorStoreObservationConvention> customObservationConvention) {
var config = new GemFireVectorStoreConfig();

config.setHost(gemFireConnectionDetails.getHost())
Expand All @@ -57,7 +63,9 @@ public GemFireVectorStore gemfireVectorStore(EmbeddingModel embeddingModel, GemF
.setVectorSimilarityFunction(properties.getVectorSimilarityFunction())
.setFields(properties.getFields())
.setSslEnabled(properties.isSslEnabled());
return new GemFireVectorStore(config, embeddingModel, properties.isInitializeSchema());
return new GemFireVectorStore(config, embeddingModel, properties.isInitializeSchema(),
observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP),
customObservationConvention.getIfAvailable(() -> null));
}

private static class PropertiesGemFireConnectionDetails implements GemFireConnectionDetails {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -22,15 +22,20 @@
import org.springframework.ai.vectorstore.HanaCloudVectorStoreConfig;
import org.springframework.ai.vectorstore.HanaVectorEntity;
import org.springframework.ai.vectorstore.HanaVectorRepository;
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;
import org.springframework.boot.autoconfigure.data.jpa.JpaRepositoriesAutoConfiguration;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;

import io.micrometer.observation.ObservationRegistry;

/**
* @author Rahul Mittal
* @author Christian Tzolov
* @since 1.0.0
*/
@AutoConfiguration(after = { JpaRepositoriesAutoConfiguration.class })
Expand All @@ -41,13 +46,17 @@ public class HanaCloudVectorStoreAutoConfiguration {
@Bean
@ConditionalOnMissingBean
public HanaCloudVectorStore vectorStore(HanaVectorRepository<? extends HanaVectorEntity> repository,
EmbeddingModel embeddingModel, HanaCloudVectorStoreProperties properties) {
EmbeddingModel embeddingModel, HanaCloudVectorStoreProperties properties,
ObjectProvider<ObservationRegistry> observationRegistry,
ObjectProvider<VectorStoreObservationConvention> customObservationConvention) {

return new HanaCloudVectorStore(repository, embeddingModel,
HanaCloudVectorStoreConfig.builder()
.tableName(properties.getTableName())
.topK(properties.getTopK())
.build());
.build(),
observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP),
customObservationConvention.getIfAvailable(() -> null));
}

}
Loading

0 comments on commit 93fa2bf

Please sign in to comment.