Skip to content

Commit a6c5d6d

Browse files
committed
Elasticsearch vector store builder refactoring
1 parent a1a2e00 commit a6c5d6d

File tree

4 files changed

+21
-13
lines changed

4 files changed

+21
-13
lines changed

spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/elasticsearch/ElasticsearchVectorStoreAutoConfiguration.java

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -73,7 +73,8 @@ ElasticsearchVectorStore vectorStore(ElasticsearchVectorStoreProperties properti
7373
elasticsearchVectorStoreOptions.setSimilarity(properties.getSimilarity());
7474
}
7575

76-
return ElasticsearchVectorStore.builder(restClient)
76+
return ElasticsearchVectorStore.builder()
77+
.restClient(restClient)
7778
.options(elasticsearchVectorStoreOptions)
7879
.embeddingModel(embeddingModel)
7980
.initializeSchema(properties.isInitializeSchema())

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

Lines changed: 11 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -111,16 +111,20 @@ public ElasticsearchVectorStore(ElasticsearchVectorStoreOptions options, RestCli
111111
EmbeddingModel embeddingModel, boolean initializeSchema, ObservationRegistry observationRegistry,
112112
VectorStoreObservationConvention customObservationConvention, BatchingStrategy batchingStrategy) {
113113

114-
this(builder(restClient).options(options)
114+
this(builder().restClient(restClient)
115+
.options(options)
115116
.embeddingModel(embeddingModel)
116117
.initializeSchema(initializeSchema)
117118
.observationRegistry(observationRegistry)
118119
.customObservationConvention(customObservationConvention)
119120
.batchingStrategy(batchingStrategy));
120121
}
121122

122-
private ElasticsearchVectorStore(ElasticsearchBuilder builder) {
123+
protected ElasticsearchVectorStore(ElasticsearchBuilder builder) {
123124
super(builder);
125+
126+
Assert.notNull(builder.restClient, "RestClient must not be null");
127+
124128
this.initializeSchema = builder.initializeSchema;
125129
this.options = builder.options;
126130
this.filterExpressionConverter = builder.filterExpressionConverter;
@@ -307,16 +311,15 @@ public record ElasticSearchDocument(String id, String content, Map<String, Objec
307311

308312
/**
309313
* Creates a new builder instance for ElasticsearchVectorStore.
310-
* @param restClient the Elasticsearch REST client
311314
* @return a new ElasticsearchBuilder instance
312315
*/
313-
public static ElasticsearchBuilder builder(RestClient restClient) {
314-
return new ElasticsearchBuilder(restClient);
316+
public static ElasticsearchBuilder builder() {
317+
return new ElasticsearchBuilder();
315318
}
316319

317320
public static class ElasticsearchBuilder extends AbstractVectorStoreBuilder<ElasticsearchBuilder> {
318321

319-
private final RestClient restClient;
322+
private RestClient restClient;
320323

321324
private ElasticsearchVectorStoreOptions options = new ElasticsearchVectorStoreOptions();
322325

@@ -327,13 +330,13 @@ public static class ElasticsearchBuilder extends AbstractVectorStoreBuilder<Elas
327330
private FilterExpressionConverter filterExpressionConverter = new ElasticsearchAiSearchFilterExpressionConverter();
328331

329332
/**
330-
* Creates a new builder instance with the specified REST client.
331333
* @param restClient the Elasticsearch REST client
332334
* @throws IllegalArgumentException if restClient is null
333335
*/
334-
public ElasticsearchBuilder(RestClient restClient) {
336+
public ElasticsearchBuilder restClient(RestClient restClient) {
335337
Assert.notNull(restClient, "RestClient must not be null");
336338
this.restClient = restClient;
339+
return this;
337340
}
338341

339342
/**

vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/elasticsearch/vectorstore/ElasticsearchVectorStoreIT.java

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -377,7 +377,8 @@ public static class TestApplication {
377377

378378
@Bean("vectorStore_cosine")
379379
public ElasticsearchVectorStore vectorStoreDefault(EmbeddingModel embeddingModel, RestClient restClient) {
380-
return ElasticsearchVectorStore.builder(restClient)
380+
return ElasticsearchVectorStore.builder()
381+
.restClient(restClient)
381382
.embeddingModel(embeddingModel)
382383
.initializeSchema(true)
383384
.build();
@@ -388,7 +389,8 @@ public ElasticsearchVectorStore vectorStoreL2(EmbeddingModel embeddingModel, Res
388389
ElasticsearchVectorStoreOptions options = new ElasticsearchVectorStoreOptions();
389390
options.setIndexName("index_l2");
390391
options.setSimilarity(SimilarityFunction.l2_norm);
391-
return ElasticsearchVectorStore.builder(restClient)
392+
return ElasticsearchVectorStore.builder()
393+
.restClient(restClient)
392394
.embeddingModel(embeddingModel)
393395
.initializeSchema(true)
394396
.options(options)
@@ -400,7 +402,8 @@ public ElasticsearchVectorStore vectorStoreDotProduct(EmbeddingModel embeddingMo
400402
ElasticsearchVectorStoreOptions options = new ElasticsearchVectorStoreOptions();
401403
options.setIndexName("index_dot_product");
402404
options.setSimilarity(SimilarityFunction.dot_product);
403-
return ElasticsearchVectorStore.builder(restClient)
405+
return ElasticsearchVectorStore.builder()
406+
.restClient(restClient)
404407
.embeddingModel(embeddingModel)
405408
.initializeSchema(true)
406409
.options(options)

vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/elasticsearch/vectorstore/ElasticsearchVectorStoreObservationIT.java

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -208,7 +208,8 @@ public TestObservationRegistry observationRegistry() {
208208
@Bean
209209
public ElasticsearchVectorStore vectorStoreDefault(EmbeddingModel embeddingModel, RestClient restClient,
210210
ObservationRegistry observationRegistry) {
211-
return ElasticsearchVectorStore.builder(restClient)
211+
return ElasticsearchVectorStore.builder()
212+
.restClient(restClient)
212213
.embeddingModel(embeddingModel)
213214
.initializeSchema(true)
214215
.options(new ElasticsearchVectorStoreOptions())

0 commit comments

Comments
 (0)