Skip to content

Commit 13f86f1

Browse files
ilayaperumalgilayaperumalg
committed
Refactor VectorStoreObservationContext builder methods
 - Deprecate the builder methods which has the prefix `with`
1 parent 99a44fb commit 13f86f1

File tree

24 files changed

+137
-78
lines changed

24 files changed

+137
-78
lines changed

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

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,6 @@
4343
import org.slf4j.LoggerFactory;
4444

4545
import org.springframework.ai.document.Document;
46-
import org.springframework.ai.document.DocumentMetadata;
4746
import org.springframework.ai.embedding.EmbeddingModel;
4847
import org.springframework.ai.observation.conventions.VectorStoreProvider;
4948
import org.springframework.ai.observation.conventions.VectorStoreSimilarityMetric;
@@ -231,9 +230,9 @@ private float[] getUserQueryEmbedding(String query) {
231230
public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName) {
232231

233232
return VectorStoreObservationContext.builder(VectorStoreProvider.SIMPLE.value(), operationName)
234-
.withDimensions(this.embeddingModel.dimensions())
235-
.withCollectionName("in-memory-map")
236-
.withSimilarityMetric(VectorStoreSimilarityMetric.COSINE.value());
233+
.dimensions(this.embeddingModel.dimensions())
234+
.collectionName("in-memory-map")
235+
.similarityMetric(VectorStoreSimilarityMetric.COSINE.value());
237236
}
238237

239238
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)