Skip to content

Commit 057e26b

Browse files
committed
Add observability support to existing vector stores
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.
1 parent 10e1e13 commit 057e26b

File tree

97 files changed

+4600
-298
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

97 files changed

+4600
-298
lines changed

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

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,22 @@ public enum VectorStoreProvider {
2323

2424
// @formatter:off
2525
PG_VECTOR("pg_vector"),
26+
AZURE_VECTOR_STORE("azure_vector_store"),
27+
CASSANDRA_VECTOR_STORE("cassandra_vector_store"),
28+
CHROMA_VECTOR_STORE("chroma_vector_store"),
29+
ELASTICSEARCH_VECTOR_STORE("elasticsearch_vector_store"),
30+
MILVUS_VECTOR_STORE("milvus_vector_store"),
31+
NEO4J_VECTOR_STORE("neo4j_vector_store"),
32+
OPENSEARCH_VECTOR_STORE("opensearch_vector_store"),
33+
QDRANT_VECTOR_STORE("qdrant_vector_store"),
34+
REDIS_VECTOR_STORE("redis_vector_store"),
35+
TYPESENSE_VECTOR_STORE("typesense_vector_store"),
36+
WEAVIATE_VECTOR_STORE("weaviate_vector_store"),
37+
PINECONE_VECTOR_STORE("pinecone_vector_store"),
38+
ORACLE_VECTOR_STORE("oracle_vector_store"),
39+
MONGODB_VECTOR_STORE("mongodb_vector_store"),
40+
GEMFIRE_VECTOR_STORE("gemfire_vector_store"),
41+
HANA_VECTOR_STORE("hana_vector_store"),
2642
SIMPLE_VECTOR_STORE("simple_vector_store");
2743

2844
// @formatter:on

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

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,13 +40,16 @@
4040
import org.springframework.ai.observation.conventions.VectorStoreSimilarityMetric;
4141
import org.springframework.ai.vectorstore.observation.AbstractObservationVectorStore;
4242
import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext;
43+
import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention;
4344
import org.springframework.core.io.Resource;
4445

4546
import com.fasterxml.jackson.core.JsonProcessingException;
4647
import com.fasterxml.jackson.core.type.TypeReference;
4748
import com.fasterxml.jackson.databind.ObjectMapper;
4849
import com.fasterxml.jackson.databind.ObjectWriter;
4950

51+
import io.micrometer.observation.ObservationRegistry;
52+
5053
/**
5154
* SimpleVectorStore is a simple implementation of the VectorStore interface.
5255
*
@@ -71,6 +74,14 @@ public class SimpleVectorStore extends AbstractObservationVectorStore {
7174
protected EmbeddingModel embeddingModel;
7275

7376
public SimpleVectorStore(EmbeddingModel embeddingModel) {
77+
this(embeddingModel, ObservationRegistry.NOOP, null);
78+
}
79+
80+
public SimpleVectorStore(EmbeddingModel embeddingModel, ObservationRegistry observationRegistry,
81+
VectorStoreObservationConvention customObservationConvention) {
82+
83+
super(observationRegistry, customObservationConvention);
84+
7485
Objects.requireNonNull(embeddingModel, "EmbeddingModel must not be null");
7586
this.embeddingModel = embeddingModel;
7687
}

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

Lines changed: 2 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -38,18 +38,10 @@ public abstract class AbstractObservationVectorStore implements VectorStore {
3838
@Nullable
3939
private final VectorStoreObservationConvention customObservationConvention;
4040

41-
public AbstractObservationVectorStore() {
42-
this(ObservationRegistry.NOOP, null);
43-
}
44-
45-
public AbstractObservationVectorStore(ObservationRegistry observationRegistry) {
46-
this(observationRegistry, null);
47-
}
48-
4941
public AbstractObservationVectorStore(ObservationRegistry observationRegistry,
50-
VectorStoreObservationConvention customSearchObservationConvention) {
42+
VectorStoreObservationConvention customObservationConvention) {
5143
this.observationRegistry = observationRegistry;
52-
this.customObservationConvention = customSearchObservationConvention;
44+
this.customObservationConvention = customObservationConvention;
5345
}
5446

5547
@Override

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

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,9 @@ public class DefaultVectorStoreObservationConvention implements VectorStoreObser
4444

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

47+
private static final KeyValue SIMILARITY_THRESHOLD_NONE = KeyValue.of(HighCardinalityKeyNames.SIMILARITY_THRESHOLD,
48+
KeyValue.NONE_VALUE);
49+
4750
private static final KeyValue SIMILARITY_METRIC_NONE = KeyValue.of(HighCardinalityKeyNames.SIMILARITY_METRIC,
4851
KeyValue.NONE_VALUE);
4952

@@ -89,7 +92,7 @@ public KeyValues getLowCardinalityKeyValues(VectorStoreObservationContext contex
8992
public KeyValues getHighCardinalityKeyValues(VectorStoreObservationContext context) {
9093
return KeyValues.of(query(context), metadataFilter(context), topK(context), dimensions(context),
9194
similarityMetric(context), collectionName(context), namespace(context), fieldName(context),
92-
indexName(context));
95+
indexName(context), similarityThreshold(context));
9396
}
9497

9598
protected KeyValue springAiKind() {
@@ -133,6 +136,14 @@ protected KeyValue topK(VectorStoreObservationContext context) {
133136
return TOP_K_NONE;
134137
}
135138

139+
protected KeyValue similarityThreshold(VectorStoreObservationContext context) {
140+
if (context.getQueryRequest() != null && context.getQueryRequest().getSimilarityThreshold() >= 0) {
141+
return KeyValue.of(HighCardinalityKeyNames.SIMILARITY_THRESHOLD,
142+
"" + context.getQueryRequest().getSimilarityThreshold());
143+
}
144+
return SIMILARITY_THRESHOLD_NONE;
145+
}
146+
136147
protected KeyValue similarityMetric(VectorStoreObservationContext context) {
137148
if (StringUtils.hasText(context.getSimilarityMetric())) {
138149
return KeyValue.of(HighCardinalityKeyNames.SIMILARITY_METRIC, context.getSimilarityMetric());

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

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,6 @@
2727
* @author Christian Tzolov
2828
* @since 1.0.0
2929
*/
30-
3130
public class VectorStoreObservationContext extends Observation.Context {
3231

3332
public enum Operation {

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

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -128,6 +128,17 @@ public String asString() {
128128
return "db.vector.query.top_k";
129129
}
130130
},
131+
/**
132+
* Similarity threshold that accepts all search scores. A threshold value of 0.0
133+
* means any similarity is accepted or disable the similarity threshold filtering.
134+
* A threshold value of 1.0 means an exact match is required.
135+
*/
136+
SIMILARITY_THRESHOLD {
137+
@Override
138+
public String asString() {
139+
return "db.vector.query.similarity_threshold";
140+
}
141+
},
131142
/**
132143
* The dimension of the vector.
133144
*/

spring-ai-spring-boot-autoconfigure/pom.xml

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -513,6 +513,12 @@
513513
<scope>test</scope>
514514
</dependency>
515515

516+
<dependency>
517+
<groupId>io.micrometer</groupId>
518+
<artifactId>micrometer-observation-test</artifactId>
519+
<scope>test</scope>
520+
</dependency>
521+
516522
</dependencies>
517523

518524
</project>

spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/azure/AzureVectorStoreAutoConfiguration.java

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,14 @@
1919
import com.azure.search.documents.indexes.SearchIndexClient;
2020
import com.azure.search.documents.indexes.SearchIndexClientBuilder;
2121

22+
import io.micrometer.observation.ObservationRegistry;
23+
24+
import java.util.List;
25+
2226
import org.springframework.ai.embedding.EmbeddingModel;
2327
import org.springframework.ai.vectorstore.azure.AzureVectorStore;
28+
import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention;
29+
import org.springframework.beans.factory.ObjectProvider;
2430
import org.springframework.boot.autoconfigure.AutoConfiguration;
2531
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
2632
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
@@ -48,9 +54,12 @@ public SearchIndexClient searchIndexClient(AzureVectorStoreProperties properties
4854
@Bean
4955
@ConditionalOnMissingBean
5056
public AzureVectorStore vectorStore(SearchIndexClient searchIndexClient, EmbeddingModel embeddingModel,
51-
AzureVectorStoreProperties properties) {
57+
AzureVectorStoreProperties properties, ObjectProvider<ObservationRegistry> observationRegistry,
58+
ObjectProvider<VectorStoreObservationConvention> customObservationConvention) {
5259

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

5564
vectorStore.setIndexName(properties.getIndexName());
5665

spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/cassandra/CassandraVectorStoreAutoConfiguration.java

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,9 +20,13 @@
2020
import com.datastax.oss.driver.api.core.CqlSession;
2121
import com.datastax.oss.driver.api.core.config.DefaultDriverOption;
2222

23+
import io.micrometer.observation.ObservationRegistry;
24+
2325
import org.springframework.ai.embedding.EmbeddingModel;
2426
import org.springframework.ai.vectorstore.CassandraVectorStore;
2527
import org.springframework.ai.vectorstore.CassandraVectorStoreConfig;
28+
import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention;
29+
import org.springframework.beans.factory.ObjectProvider;
2630
import org.springframework.boot.autoconfigure.AutoConfiguration;
2731
import org.springframework.boot.autoconfigure.cassandra.CassandraAutoConfiguration;
2832
import org.springframework.boot.autoconfigure.cassandra.DriverConfigLoaderBuilderCustomizer;
@@ -33,6 +37,7 @@
3337

3438
/**
3539
* @author Mick Semb Wever
40+
* @author Christian Tzolov
3641
* @since 1.0.0
3742
*/
3843
@AutoConfiguration(after = CassandraAutoConfiguration.class)
@@ -43,7 +48,8 @@ public class CassandraVectorStoreAutoConfiguration {
4348
@Bean
4449
@ConditionalOnMissingBean
4550
public CassandraVectorStore vectorStore(EmbeddingModel embeddingModel, CassandraVectorStoreProperties properties,
46-
CqlSession cqlSession) {
51+
CqlSession cqlSession, ObjectProvider<ObservationRegistry> observationRegistry,
52+
ObjectProvider<VectorStoreObservationConvention> customObservationConvention) {
4753

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

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

64-
return CassandraVectorStore.create(builder.build(), embeddingModel);
70+
return new CassandraVectorStore(builder.build(), embeddingModel,
71+
observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP),
72+
customObservationConvention.getIfAvailable(() -> null));
6573
}
6674

6775
@Bean

spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/chroma/ChromaVectorStoreAutoConfiguration.java

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,8 @@
1818
import org.springframework.ai.chroma.ChromaApi;
1919
import org.springframework.ai.embedding.EmbeddingModel;
2020
import org.springframework.ai.vectorstore.ChromaVectorStore;
21+
import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention;
22+
import org.springframework.beans.factory.ObjectProvider;
2123
import org.springframework.boot.autoconfigure.AutoConfiguration;
2224
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
2325
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
@@ -29,6 +31,8 @@
2931

3032
import com.fasterxml.jackson.databind.ObjectMapper;
3133

34+
import io.micrometer.observation.ObservationRegistry;
35+
3236
/**
3337
* @author Christian Tzolov
3438
* @author Eddú Meléndez
@@ -72,9 +76,11 @@ else if (StringUtils.hasText(apiProperties.getUsername()) && StringUtils.hasText
7276
@Bean
7377
@ConditionalOnMissingBean
7478
public ChromaVectorStore vectorStore(EmbeddingModel embeddingModel, ChromaApi chromaApi,
75-
ChromaVectorStoreProperties storeProperties) {
79+
ChromaVectorStoreProperties storeProperties, ObjectProvider<ObservationRegistry> observationRegistry,
80+
ObjectProvider<VectorStoreObservationConvention> customObservationConvention) {
7681
return new ChromaVectorStore(embeddingModel, chromaApi, storeProperties.getCollectionName(),
77-
storeProperties.isInitializeSchema());
82+
storeProperties.isInitializeSchema(), observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP),
83+
customObservationConvention.getIfAvailable(() -> null));
7884
}
7985

8086
static class PropertiesChromaConnectionDetails implements ChromaConnectionDetails {

0 commit comments

Comments
 (0)