Skip to content

Commit 26fab03

Browse files
ilayaperumalgsobychacko
ilayaperumalg
authored and
sobychacko
committed
Refactor VectorStoreObservationContext builder methods
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.
1 parent 8ce8374 commit 26fab03

File tree

24 files changed

+137
-77
lines changed

24 files changed

+137
-77
lines changed

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -246,9 +246,9 @@ private float[] getUserQueryEmbedding(String query) {
246246
public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName) {
247247

248248
return VectorStoreObservationContext.builder(VectorStoreProvider.SIMPLE.value(), operationName)
249-
.withDimensions(this.embeddingModel.dimensions())
250-
.withCollectionName("in-memory-map")
251-
.withSimilarityMetric(VectorStoreSimilarityMetric.COSINE.value());
249+
.dimensions(this.embeddingModel.dimensions())
250+
.collectionName("in-memory-map")
251+
.similarityMetric(VectorStoreSimilarityMetric.COSINE.value());
252252
}
253253

254254
public static final class EmbeddingMath {

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -111,7 +111,7 @@ public List<Document> similaritySearch(SearchRequest request) {
111111

112112
VectorStoreObservationContext searchObservationContext = this
113113
.createObservationContextBuilder(VectorStoreObservationContext.Operation.QUERY.value())
114-
.withQueryRequest(request)
114+
.queryRequest(request)
115115
.build();
116116

117117
return VectorStoreObservationDocumentation.AI_VECTOR_STORE

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

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@
3030
*
3131
* @author Christian Tzolov
3232
* @author Thomas Vitale
33+
* @author Ilayaperumal Gopinathan
3334
* @since 1.0.0
3435
*/
3536
public class VectorStoreObservationContext extends Observation.Context {
@@ -184,36 +185,99 @@ public Builder(String databaseSystem, String operationName) {
184185
this.context = new VectorStoreObservationContext(databaseSystem, operationName);
185186
}
186187

188+
public Builder collectionName(String collectionName) {
189+
this.context.setCollectionName(collectionName);
190+
return this;
191+
}
192+
193+
public Builder dimensions(Integer dimensions) {
194+
this.context.setDimensions(dimensions);
195+
return this;
196+
}
197+
198+
public Builder fieldName(@Nullable String fieldName) {
199+
this.context.setFieldName(fieldName);
200+
return this;
201+
}
202+
203+
public Builder namespace(String namespace) {
204+
this.context.setNamespace(namespace);
205+
return this;
206+
}
207+
208+
public Builder queryRequest(SearchRequest request) {
209+
this.context.setQueryRequest(request);
210+
return this;
211+
}
212+
213+
public Builder queryResponse(List<Document> documents) {
214+
this.context.setQueryResponse(documents);
215+
return this;
216+
}
217+
218+
public Builder similarityMetric(String similarityMetric) {
219+
this.context.setSimilarityMetric(similarityMetric);
220+
return this;
221+
}
222+
223+
/**
224+
* @deprecated use {@link #collectionName(String)} instead.
225+
*/
226+
@Deprecated(forRemoval = true, since = "1.0.0-M5")
187227
public Builder withCollectionName(String collectionName) {
188228
this.context.setCollectionName(collectionName);
189229
return this;
190230
}
191231

232+
/**
233+
* @deprecated use {@link #dimensions(Integer)} instead.
234+
*/
235+
@Deprecated(forRemoval = true, since = "1.0.0-M5")
192236
public Builder withDimensions(Integer dimensions) {
193237
this.context.setDimensions(dimensions);
194238
return this;
195239
}
196240

241+
/**
242+
* @deprecated use {@link #fieldName(String)} instead.
243+
*/
244+
@Deprecated(forRemoval = true, since = "1.0.0-M5")
197245
public Builder withFieldName(@Nullable String fieldName) {
198246
this.context.setFieldName(fieldName);
199247
return this;
200248
}
201249

250+
/**
251+
* @deprecated use {@link #namespace(String)} instead.
252+
*/
253+
@Deprecated(forRemoval = true, since = "1.0.0-M5")
202254
public Builder withNamespace(String namespace) {
203255
this.context.setNamespace(namespace);
204256
return this;
205257
}
206258

259+
/**
260+
* @deprecated use {@link #queryRequest(SearchRequest)} instead.
261+
*/
262+
@Deprecated(forRemoval = true, since = "1.0.0-M5")
207263
public Builder withQueryRequest(SearchRequest request) {
208264
this.context.setQueryRequest(request);
209265
return this;
210266
}
211267

268+
/**
269+
* @deprecated use {@link #queryResponse(List)} instead.
270+
*/
271+
@Deprecated(forRemoval = true, since = "1.0.0-M5")
212272
public Builder withQueryResponse(List<Document> documents) {
213273
this.context.setQueryResponse(documents);
214274
return this;
215275
}
216276

277+
/**
278+
* @deprecated use {@link #similarityMetric(String)} instead.
279+
*/
280+
@Deprecated(forRemoval = true, since = "1.0.0-M5")
217281
public Builder withSimilarityMetric(String similarityMetric) {
218282
this.context.setSimilarityMetric(similarityMetric);
219283
return this;

‎spring-ai-core/src/test/java/org/springframework/ai/vectorstore/observation/DefaultVectorStoreObservationConventionTests.java‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -78,12 +78,12 @@ void shouldHaveRequiredKeyValues() {
7878
void shouldHaveOptionalKeyValues() {
7979
VectorStoreObservationContext observationContext = VectorStoreObservationContext
8080
.builder("my-database", VectorStoreObservationContext.Operation.QUERY)
81-
.withCollectionName("COLLECTION_NAME")
82-
.withDimensions(696)
83-
.withFieldName("FIELD_NAME")
84-
.withNamespace("NAMESPACE")
85-
.withSimilarityMetric("SIMILARITY_METRIC")
86-
.withQueryRequest(SearchRequest.query("VDB QUERY").withFilterExpression("country == 'UK' && year >= 2020"))
81+
.collectionName("COLLECTION_NAME")
82+
.dimensions(696)
83+
.fieldName("FIELD_NAME")
84+
.namespace("NAMESPACE")
85+
.similarityMetric("SIMILARITY_METRIC")
86+
.queryRequest(SearchRequest.query("VDB QUERY").withFilterExpression("country == 'UK' && year >= 2020"))
8787
.build();
8888

8989
List<Document> queryResponseDocs = List.of(new Document("doc1"), new Document("doc2"));

‎spring-ai-core/src/test/java/org/springframework/ai/vectorstore/observation/VectorStoreQueryResponseObservationHandlerTests.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,7 @@ class VectorStoreQueryResponseObservationHandlerTests {
4444
void whenCompletionWithTextThenSpanEvent() {
4545
var observationContext = VectorStoreObservationContext
4646
.builder("db", VectorStoreObservationContext.Operation.ADD)
47-
.withQueryResponse(List.of(new Document("hello"), new Document("other-side")))
47+
.queryResponse(List.of(new Document("hello"), new Document("other-side")))
4848
.build();
4949
var sdkTracer = SdkTracerProvider.builder().build().get("test");
5050
var otelTracer = new OtelTracer(sdkTracer, new OtelCurrentTraceContext(), null);

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

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,6 @@
1818

1919
import java.util.ArrayList;
2020
import java.util.Collections;
21-
import java.util.HashMap;
2221
import java.util.List;
2322
import java.util.Optional;
2423
import java.util.stream.Collectors;
@@ -354,10 +353,10 @@ public List<Document> doSimilaritySearch(SearchRequest request) {
354353
@Override
355354
public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName) {
356355
return VectorStoreObservationContext.builder(VectorStoreProvider.COSMOSDB.value(), operationName)
357-
.withCollectionName(this.container.getId())
358-
.withDimensions(this.embeddingModel.dimensions())
359-
.withNamespace(this.container.getDatabase().getId())
360-
.withSimilarityMetric("cosine");
356+
.collectionName(this.container.getId())
357+
.dimensions(this.embeddingModel.dimensions())
358+
.namespace(this.container.getDatabase().getId())
359+
.similarityMetric("cosine");
361360
}
362361

363362
}

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

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,6 @@
5959
import org.springframework.ai.vectorstore.filter.FilterExpressionConverter;
6060
import org.springframework.ai.vectorstore.observation.AbstractObservationVectorStore;
6161
import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext;
62-
import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext.Builder;
6362
import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention;
6463
import org.springframework.beans.factory.InitializingBean;
6564
import org.springframework.util.Assert;
@@ -390,9 +389,9 @@ public void afterPropertiesSet() throws Exception {
390389
public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName) {
391390

392391
return VectorStoreObservationContext.builder(VectorStoreProvider.AZURE.value(), operationName)
393-
.withCollectionName(this.indexName)
394-
.withDimensions(this.embeddingModel.dimensions())
395-
.withSimilarityMetric(this.initializeSchema ? VectorStoreSimilarityMetric.COSINE.value() : null);
392+
.collectionName(this.indexName)
393+
.dimensions(this.embeddingModel.dimensions())
394+
.similarityMetric(this.initializeSchema ? VectorStoreSimilarityMetric.COSINE.value() : null);
396395
}
397396

398397
public record MetadataField(String name, SearchFieldDataType fieldType) {

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

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -502,10 +502,10 @@ private String getDocumentId(Row row) {
502502
@Override
503503
public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName) {
504504
return VectorStoreObservationContext.builder(VectorStoreProvider.CASSANDRA.value(), operationName)
505-
.withCollectionName(this.schema.table())
506-
.withDimensions(this.embeddingModel.dimensions())
507-
.withNamespace(this.schema.keyspace())
508-
.withSimilarityMetric(getSimilarityMetric());
505+
.collectionName(this.schema.table())
506+
.dimensions(this.embeddingModel.dimensions())
507+
.namespace(this.schema.keyspace())
508+
.similarityMetric(getSimilarityMetric());
509509
}
510510

511511
private String getSimilarityMetric() {

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -272,8 +272,8 @@ public void setFilterExpressionConverter(FilterExpressionConverter filterExpress
272272
@Override
273273
public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName) {
274274
return VectorStoreObservationContext.builder(VectorStoreProvider.CHROMA.value(), operationName)
275-
.withDimensions(this.embeddingModel.dimensions())
276-
.withCollectionName(this.collectionName + ":" + this.collectionId);
275+
.dimensions(this.embeddingModel.dimensions())
276+
.collectionName(this.collectionName + ":" + this.collectionId);
277277
}
278278

279279
public static class ChromaBuilder extends AbstractVectorStoreBuilder<ChromaBuilder> {

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -362,9 +362,9 @@ public void afterPropertiesSet() {
362362
@Override
363363
public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName) {
364364
return VectorStoreObservationContext.builder(VectorStoreProvider.ELASTICSEARCH.value(), operationName)
365-
.withCollectionName(this.options.getIndexName())
366-
.withDimensions(this.embeddingModel.dimensions())
367-
.withSimilarityMetric(getSimilarityMetric());
365+
.collectionName(this.options.getIndexName())
366+
.dimensions(this.embeddingModel.dimensions())
367+
.similarityMetric(getSimilarityMetric());
368368
}
369369

370370
private String getSimilarityMetric() {

0 commit comments

Comments
 (0)