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

Refactor VectorStoreObservationContext builder methods

Browse Source 
This commit deprecates builder methods with the "with" prefix in the VectorStoreObservationContext class
and introduces new methods without the prefix for a cleaner API. For example:

withCollectionName() → collectionName()
withDimensions() → dimensions()
withNamespace() → namespace()

The old methods are marked as deprecated for maintaining backward compatibility.

All vector store implementations have been updated to use the new method names.
This commit is contained in:
Ilayaperumal Gopinathan
2024-12-18 14:34:15 +00:00
committed by Soby Chacko
parent 8ce8374825
commit 26fab03c2c
24 changed files with 137 additions and 77 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -246,9 +246,9 @@ public class SimpleVectorStore extends AbstractObservationVectorStore {
public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName) {
return VectorStoreObservationContext.builder(VectorStoreProvider.SIMPLE.value(), operationName)
.withDimensions(this.embeddingModel.dimensions())
.withCollectionName("in-memory-map")
.withSimilarityMetric(VectorStoreSimilarityMetric.COSINE.value());
.dimensions(this.embeddingModel.dimensions())
.collectionName("in-memory-map")
.similarityMetric(VectorStoreSimilarityMetric.COSINE.value());
}
public static final class EmbeddingMath {

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

@@ -111,7 +111,7 @@ public abstract class AbstractObservationVectorStore implements VectorStore {
VectorStoreObservationContext searchObservationContext = this
.createObservationContextBuilder(VectorStoreObservationContext.Operation.QUERY.value())
.withQueryRequest(request)
.queryRequest(request)
.build();
return VectorStoreObservationDocumentation.AI_VECTOR_STORE

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

@@ -30,6 +30,7 @@ import org.springframework.util.Assert;
*
* @author Christian Tzolov
* @author Thomas Vitale
* @author Ilayaperumal Gopinathan
* @since 1.0.0
*/
public class VectorStoreObservationContext extends Observation.Context {
@@ -184,36 +185,99 @@ public class VectorStoreObservationContext extends Observation.Context {
this.context = new VectorStoreObservationContext(databaseSystem, operationName);
}
public Builder collectionName(String collectionName) {
this.context.setCollectionName(collectionName);
return this;
}
public Builder dimensions(Integer dimensions) {
this.context.setDimensions(dimensions);
return this;
}
public Builder fieldName(@Nullable String fieldName) {
this.context.setFieldName(fieldName);
return this;
}
public Builder namespace(String namespace) {
this.context.setNamespace(namespace);
return this;
}
public Builder queryRequest(SearchRequest request) {
this.context.setQueryRequest(request);
return this;
}
public Builder queryResponse(List<Document> documents) {
this.context.setQueryResponse(documents);
return this;
}
public Builder similarityMetric(String similarityMetric) {
this.context.setSimilarityMetric(similarityMetric);
return this;
}
/**
* @deprecated use {@link #collectionName(String)} instead.
*/
@Deprecated(forRemoval = true, since = "1.0.0-M5")
public Builder withCollectionName(String collectionName) {
this.context.setCollectionName(collectionName);
return this;
}
/**
* @deprecated use {@link #dimensions(Integer)} instead.
*/
@Deprecated(forRemoval = true, since = "1.0.0-M5")
public Builder withDimensions(Integer dimensions) {
this.context.setDimensions(dimensions);
return this;
}
/**
* @deprecated use {@link #fieldName(String)} instead.
*/
@Deprecated(forRemoval = true, since = "1.0.0-M5")
public Builder withFieldName(@Nullable String fieldName) {
this.context.setFieldName(fieldName);
return this;
}
/**
* @deprecated use {@link #namespace(String)} instead.
*/
@Deprecated(forRemoval = true, since = "1.0.0-M5")
public Builder withNamespace(String namespace) {
this.context.setNamespace(namespace);
return this;
}
/**
* @deprecated use {@link #queryRequest(SearchRequest)} instead.
*/
@Deprecated(forRemoval = true, since = "1.0.0-M5")
public Builder withQueryRequest(SearchRequest request) {
this.context.setQueryRequest(request);
return this;
}
/**
* @deprecated use {@link #queryResponse(List)} instead.
*/
@Deprecated(forRemoval = true, since = "1.0.0-M5")
public Builder withQueryResponse(List<Document> documents) {
this.context.setQueryResponse(documents);
return this;
}
/**
* @deprecated use {@link #similarityMetric(String)} instead.
*/
@Deprecated(forRemoval = true, since = "1.0.0-M5")
public Builder withSimilarityMetric(String similarityMetric) {
this.context.setSimilarityMetric(similarityMetric);
return this;

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

@@ -78,12 +78,12 @@ class DefaultVectorStoreObservationConventionTests {
void shouldHaveOptionalKeyValues() {
VectorStoreObservationContext observationContext = VectorStoreObservationContext
.builder("my-database", VectorStoreObservationContext.Operation.QUERY)
.withCollectionName("COLLECTION_NAME")
.withDimensions(696)
.withFieldName("FIELD_NAME")
.withNamespace("NAMESPACE")
.withSimilarityMetric("SIMILARITY_METRIC")
.withQueryRequest(SearchRequest.query("VDB QUERY").withFilterExpression("country == 'UK' && year >= 2020"))
.collectionName("COLLECTION_NAME")
.dimensions(696)
.fieldName("FIELD_NAME")
.namespace("NAMESPACE")
.similarityMetric("SIMILARITY_METRIC")
.queryRequest(SearchRequest.query("VDB QUERY").withFilterExpression("country == 'UK' && year >= 2020"))
.build();
List<Document> queryResponseDocs = List.of(new Document("doc1"), new Document("doc2"));

2
spring-ai-core/src/test/java/org/springframework/ai/vectorstore/observation/VectorStoreQueryResponseObservationHandlerTests.java
View File

@@ -44,7 +44,7 @@ class VectorStoreQueryResponseObservationHandlerTests {
void whenCompletionWithTextThenSpanEvent() {
var observationContext = VectorStoreObservationContext
.builder("db", VectorStoreObservationContext.Operation.ADD)
.withQueryResponse(List.of(new Document("hello"), new Document("other-side")))
.queryResponse(List.of(new Document("hello"), new Document("other-side")))
.build();
var sdkTracer = SdkTracerProvider.builder().build().get("test");
var otelTracer = new OtelTracer(sdkTracer, new OtelCurrentTraceContext(), null);

9
vector-stores/spring-ai-azure-cosmos-db-store/src/main/java/org/springframework/ai/vectorstore/CosmosDBVectorStore.java
View File

@@ -18,7 +18,6 @@ package org.springframework.ai.vectorstore;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Optional;
import java.util.stream.Collectors;
@@ -354,10 +353,10 @@ public class CosmosDBVectorStore extends AbstractObservationVectorStore implemen
@Override
public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName) {
return VectorStoreObservationContext.builder(VectorStoreProvider.COSMOSDB.value(), operationName)
.withCollectionName(this.container.getId())
.withDimensions(this.embeddingModel.dimensions())
.withNamespace(this.container.getDatabase().getId())
.withSimilarityMetric("cosine");
.collectionName(this.container.getId())
.dimensions(this.embeddingModel.dimensions())
.namespace(this.container.getDatabase().getId())
.similarityMetric("cosine");
}
}

7
vector-stores/spring-ai-azure-store/src/main/java/org/springframework/ai/vectorstore/azure/AzureVectorStore.java
View File

@@ -59,7 +59,6 @@ import org.springframework.ai.vectorstore.SearchRequest;
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.VectorStoreObservationContext.Builder;
import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.util.Assert;
@@ -390,9 +389,9 @@ public class AzureVectorStore extends AbstractObservationVectorStore implements
public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName) {
return VectorStoreObservationContext.builder(VectorStoreProvider.AZURE.value(), operationName)
.withCollectionName(this.indexName)
.withDimensions(this.embeddingModel.dimensions())
.withSimilarityMetric(this.initializeSchema ? VectorStoreSimilarityMetric.COSINE.value() : null);
.collectionName(this.indexName)
.dimensions(this.embeddingModel.dimensions())
.similarityMetric(this.initializeSchema ? VectorStoreSimilarityMetric.COSINE.value() : null);
}
public record MetadataField(String name, SearchFieldDataType fieldType) {

8
vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/vectorstore/cassandra/CassandraVectorStore.java
View File

@@ -502,10 +502,10 @@ public class CassandraVectorStore extends AbstractObservationVectorStore impleme
@Override
public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName) {
return VectorStoreObservationContext.builder(VectorStoreProvider.CASSANDRA.value(), operationName)
.withCollectionName(this.schema.table())
.withDimensions(this.embeddingModel.dimensions())
.withNamespace(this.schema.keyspace())
.withSimilarityMetric(getSimilarityMetric());
.collectionName(this.schema.table())
.dimensions(this.embeddingModel.dimensions())
.namespace(this.schema.keyspace())
.similarityMetric(getSimilarityMetric());
}
private String getSimilarityMetric() {

4
vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/chroma/vectorstore/ChromaVectorStore.java
View File

@@ -272,8 +272,8 @@ public class ChromaVectorStore extends AbstractObservationVectorStore implements
@Override
public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName) {
return VectorStoreObservationContext.builder(VectorStoreProvider.CHROMA.value(), operationName)
.withDimensions(this.embeddingModel.dimensions())
.withCollectionName(this.collectionName + ":" + this.collectionId);
.dimensions(this.embeddingModel.dimensions())
.collectionName(this.collectionName + ":" + this.collectionId);
}
public static class ChromaBuilder extends AbstractVectorStoreBuilder<ChromaBuilder> {

6
vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/elasticsearch/ElasticsearchVectorStore.java
View File

@@ -362,9 +362,9 @@ public class ElasticsearchVectorStore extends AbstractObservationVectorStore imp
@Override
public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName) {
return VectorStoreObservationContext.builder(VectorStoreProvider.ELASTICSEARCH.value(), operationName)
.withCollectionName(this.options.getIndexName())
.withDimensions(this.embeddingModel.dimensions())
.withSimilarityMetric(getSimilarityMetric());
.collectionName(this.options.getIndexName())
.dimensions(this.embeddingModel.dimensions())
.similarityMetric(getSimilarityMetric());
}
private String getSimilarityMetric() {

6
vector-stores/spring-ai-gemfire-store/src/main/java/org/springframework/ai/vectorstore/GemFireVectorStore.java
View File

@@ -340,9 +340,9 @@ public class GemFireVectorStore extends AbstractObservationVectorStore implement
@Override
public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName) {
return VectorStoreObservationContext.builder(VectorStoreProvider.GEMFIRE.value(), operationName)
.withCollectionName(this.indexName)
.withDimensions(this.embeddingModel.dimensions())
.withFieldName(EMBEDDINGS);
.collectionName(this.indexName)
.dimensions(this.embeddingModel.dimensions())
.fieldName(EMBEDDINGS);
}
public static class CreateRequest {

7
vector-stores/spring-ai-hanadb-store/src/main/java/org/springframework/ai/vectorstore/HanaCloudVectorStore.java
View File

@@ -36,7 +36,6 @@ import org.springframework.ai.observation.conventions.VectorStoreSimilarityMetri
import org.springframework.ai.util.JacksonUtils;
import org.springframework.ai.vectorstore.observation.AbstractObservationVectorStore;
import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext;
import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext.Builder;
import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention;
/**
@@ -175,9 +174,9 @@ public class HanaCloudVectorStore extends AbstractObservationVectorStore {
public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName) {
return VectorStoreObservationContext.builder(VectorStoreProvider.HANA.value(), operationName)
.withDimensions(this.embeddingModel.dimensions())
.withCollectionName(this.config.getTableName())
.withSimilarityMetric(VectorStoreSimilarityMetric.COSINE.value());
.dimensions(this.embeddingModel.dimensions())
.collectionName(this.config.getTableName())
.similarityMetric(VectorStoreSimilarityMetric.COSINE.value());
}
}

8
vector-stores/spring-ai-mariadb-store/src/main/java/org/springframework/ai/vectorstore/mariadb/MariaDBVectorStore.java
View File

@@ -504,10 +504,10 @@ public class MariaDBVectorStore extends AbstractObservationVectorStore implement
public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName) {
return VectorStoreObservationContext.builder(VectorStoreProvider.MARIADB.value(), operationName)
.withCollectionName(this.vectorTableName)
.withDimensions(this.embeddingDimensions())
.withNamespace(this.schemaName)
.withSimilarityMetric(getSimilarityMetric());
.collectionName(this.vectorTableName)
.dimensions(this.embeddingDimensions())
.namespace(this.schemaName)
.similarityMetric(getSimilarityMetric());
}
private String getSimilarityMetric() {

8
vector-stores/spring-ai-milvus-store/src/main/java/org/springframework/ai/milvus/vectorstore/MilvusVectorStore.java
View File

@@ -564,10 +564,10 @@ public class MilvusVectorStore extends AbstractObservationVectorStore implements
String operationName) {
return VectorStoreObservationContext.builder(VectorStoreProvider.MILVUS.value(), operationName)
.withCollectionName(this.collectionName)
.withDimensions(this.embeddingModel.dimensions())
.withSimilarityMetric(getSimilarityMetric())
.withNamespace(this.databaseName);
.collectionName(this.collectionName)
.dimensions(this.embeddingModel.dimensions())
.similarityMetric(getSimilarityMetric())
.namespace(this.databaseName);
}
private String getSimilarityMetric() {

6
vector-stores/spring-ai-mongodb-atlas-store/src/main/java/org/springframework/ai/vectorstore/mongodb/atlas/MongoDBAtlasVectorStore.java
View File

@@ -345,9 +345,9 @@ public class MongoDBAtlasVectorStore extends AbstractObservationVectorStore impl
public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName) {
return VectorStoreObservationContext.builder(VectorStoreProvider.MONGODB.value(), operationName)
.withCollectionName(this.collectionName)
.withDimensions(this.embeddingModel.dimensions())
.withFieldName(this.pathName);
.collectionName(this.collectionName)
.dimensions(this.embeddingModel.dimensions())
.fieldName(this.pathName);
}
/**

6
vector-stores/spring-ai-neo4j-store/src/main/java/org/springframework/ai/vectorstore/neo4j/Neo4jVectorStore.java
View File

@@ -373,9 +373,9 @@ public class Neo4jVectorStore extends AbstractObservationVectorStore implements
public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName) {
return VectorStoreObservationContext.builder(VectorStoreProvider.NEO4J.value(), operationName)
.withCollectionName(this.indexName)
.withDimensions(this.embeddingModel.dimensions())
.withSimilarityMetric(getSimilarityMetric());
.collectionName(this.indexName)
.dimensions(this.embeddingModel.dimensions())
.similarityMetric(getSimilarityMetric());
}
private String getSimilarityMetric() {

6
vector-stores/spring-ai-opensearch-store/src/main/java/org/springframework/ai/vectorstore/opensearch/OpenSearchVectorStore.java
View File

@@ -418,9 +418,9 @@ public class OpenSearchVectorStore extends AbstractObservationVectorStore implem
@Override
public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName) {
return VectorStoreObservationContext.builder(VectorStoreProvider.OPENSEARCH.value(), operationName)
.withCollectionName(this.index)
.withDimensions(this.embeddingModel.dimensions())
.withSimilarityMetric(getSimilarityFunction());
.collectionName(this.index)
.dimensions(this.embeddingModel.dimensions())
.similarityMetric(getSimilarityFunction());
}
private String getSimilarityFunction() {

7
vector-stores/spring-ai-oracle-store/src/main/java/org/springframework/ai/vectorstore/OracleVectorStore.java
View File

@@ -49,7 +49,6 @@ import org.springframework.ai.observation.conventions.VectorStoreSimilarityMetri
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.VectorStoreObservationContext.Builder;
import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.jdbc.core.BatchPreparedStatementSetter;
@@ -551,9 +550,9 @@ public class OracleVectorStore extends AbstractObservationVectorStore implements
@Override
public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName) {
return VectorStoreObservationContext.builder(VectorStoreProvider.ORACLE.value(), operationName)
.withDimensions(this.embeddingModel.dimensions())
.withCollectionName(this.getTableName())
.withSimilarityMetric(getSimilarityMetric());
.dimensions(this.embeddingModel.dimensions())
.collectionName(this.getTableName())
.similarityMetric(getSimilarityMetric());
}
private String getSimilarityMetric() {

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

@@ -494,10 +494,10 @@ public class PgVectorStore extends AbstractObservationVectorStore implements Ini
public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName) {
return VectorStoreObservationContext.builder(VectorStoreProvider.PG_VECTOR.value(), operationName)
.withCollectionName(this.vectorTableName)
.withDimensions(this.embeddingDimensions())
.withNamespace(this.schemaName)
.withSimilarityMetric(getSimilarityMetric());
.collectionName(this.vectorTableName)
.dimensions(this.embeddingDimensions())
.namespace(this.schemaName)
.similarityMetric(getSimilarityMetric());
}
private String getSimilarityMetric() {

8
vector-stores/spring-ai-pinecone-store/src/main/java/org/springframework/ai/vectorstore/PineconeVectorStore.java
View File

@@ -287,10 +287,10 @@ public class PineconeVectorStore extends AbstractObservationVectorStore {
public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName) {
return VectorStoreObservationContext.builder(VectorStoreProvider.PINECONE.value(), operationName)
.withCollectionName(this.pineconeIndexName)
.withDimensions(this.embeddingModel.dimensions())
.withNamespace(this.pineconeNamespace)
.withFieldName(this.pineconeContentFieldName);
.collectionName(this.pineconeIndexName)
.dimensions(this.embeddingModel.dimensions())
.namespace(this.pineconeNamespace)
.fieldName(this.pineconeContentFieldName);
}
/**

4
vector-stores/spring-ai-qdrant-store/src/main/java/org/springframework/ai/vectorstore/qdrant/QdrantVectorStore.java
View File

@@ -359,8 +359,8 @@ public class QdrantVectorStore extends AbstractObservationVectorStore implements
public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName) {
return VectorStoreObservationContext.builder(VectorStoreProvider.QDRANT.value(), operationName)
.withDimensions(this.embeddingModel.dimensions())
.withCollectionName(this.collectionName);
.dimensions(this.embeddingModel.dimensions())
.collectionName(this.collectionName);
}

8
vector-stores/spring-ai-redis-store/src/main/java/org/springframework/ai/vectorstore/redis/RedisVectorStore.java
View File

@@ -454,10 +454,10 @@ public class RedisVectorStore extends AbstractObservationVectorStore implements
public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName) {
return VectorStoreObservationContext.builder(VectorStoreProvider.REDIS.value(), operationName)
.withCollectionName(this.indexName)
.withDimensions(this.embeddingModel.dimensions())
.withFieldName(this.embeddingFieldName)
.withSimilarityMetric(VectorStoreSimilarityMetric.COSINE.value());
.collectionName(this.indexName)
.dimensions(this.embeddingModel.dimensions())
.fieldName(this.embeddingFieldName)
.similarityMetric(VectorStoreSimilarityMetric.COSINE.value());
}

8
vector-stores/spring-ai-typesense-store/src/main/java/org/springframework/ai/vectorstore/typesense/TypesenseVectorStore.java
View File

@@ -443,10 +443,10 @@ public class TypesenseVectorStore extends AbstractObservationVectorStore impleme
public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName) {
return VectorStoreObservationContext.builder(VectorStoreProvider.TYPESENSE.value(), operationName)
.withDimensions(this.embeddingModel.dimensions())
.withCollectionName(this.collectionName)
.withFieldName(EMBEDDING_FIELD_NAME)
.withSimilarityMetric(VectorStoreSimilarityMetric.COSINE.value());
.dimensions(this.embeddingModel.dimensions())
.collectionName(this.collectionName)
.fieldName(EMBEDDING_FIELD_NAME)
.similarityMetric(VectorStoreSimilarityMetric.COSINE.value());
}
public static final class TypesenseBuilder extends AbstractVectorStoreBuilder<TypesenseBuilder> {

4
vector-stores/spring-ai-weaviate-store/src/main/java/org/springframework/ai/vectorstore/weaviate/WeaviateVectorStore.java
View File

@@ -448,8 +448,8 @@ public class WeaviateVectorStore extends AbstractObservationVectorStore {
public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName) {
return VectorStoreObservationContext.builder(VectorStoreProvider.WEAVIATE.value(), operationName)
.withDimensions(this.embeddingModel.dimensions())
.withCollectionName(this.weaviateObjectClass);
.dimensions(this.embeddingModel.dimensions())
.collectionName(this.weaviateObjectClass);
}
/**