From fc1f92d11cacc79bdf8d41eb7f941a64fdf67482 Mon Sep 17 00:00:00 2001 From: Soby Chacko Date: Mon, 9 Dec 2024 11:13:26 -0500 Subject: [PATCH] Add builder pattern and refactor Elasticsearch store package The changes introduce a fluent builder pattern for ElasticsearchVectorStore configuration, making it easier to create and customize instances with optional parameters. All Elasticsearch-related classes are moved to a dedicated elasticsearch package for better organization. Key changes: * Add ElasticsearchVectorStore.builder() with comprehensive options * Move classes to org.springframework.ai.vectorstore.elasticsearch package * Deprecate old constructors in favor of builder pattern * Add support for configurable batching strategies * Enhance documentation with usage examples and best practices --- .../pages/api/vectordbs/elasticsearch.adoc | 80 +++---- ...ticsearchVectorStoreAutoConfiguration.java | 16 +- .../ElasticsearchVectorStoreProperties.java | 2 +- ...csearchVectorStoreAutoConfigurationIT.java | 4 +- ...archAiSearchFilterExpressionConverter.java | 2 +- .../ElasticsearchVectorStore.java | 217 ++++++++++++++++-- .../ElasticsearchVectorStoreOptions.java | 2 +- .../SimilarityFunction.java | 2 +- ...AiSearchFilterExpressionConverterTest.java | 2 +- .../ElasticsearchImage.java | 2 +- .../ElasticsearchVectorStoreIT.java | 23 +- ...ElasticsearchVectorStoreObservationIT.java | 16 +- 12 files changed, 285 insertions(+), 83 deletions(-) rename vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/{ => elasticsearch}/ElasticsearchAiSearchFilterExpressionConverter.java (98%) rename vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/{ => elasticsearch}/ElasticsearchVectorStore.java (62%) rename vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/{ => elasticsearch}/ElasticsearchVectorStoreOptions.java (96%) rename vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/{ => elasticsearch}/SimilarityFunction.java (94%) rename vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/{ => elasticsearch}/ElasticsearchAiSearchFilterExpressionConverterTest.java (98%) rename vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/{ => elasticsearch}/ElasticsearchImage.java (94%) rename vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/{ => elasticsearch}/ElasticsearchVectorStoreIT.java (96%) rename vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/{ => elasticsearch}/ElasticsearchVectorStoreObservationIT.java (94%) diff --git a/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs/elasticsearch.adoc b/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs/elasticsearch.adoc index fab97b3f0..34835e6d9 100644 --- a/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs/elasticsearch.adoc +++ b/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs/elasticsearch.adoc @@ -76,14 +76,11 @@ Alternatively you can opt-out the initialization and create the index manually u NOTE: this is a breaking change! In earlier versions of Spring AI, this schema initialization happened by default. - - Please have a look at the list of <> for the vector store to learn about the default values and configuration options. These properties can be also set by configuring the `ElasticsearchVectorStoreOptions` bean. Additionally, you will need a configured `EmbeddingModel` bean. Refer to the xref:api/embeddings.adoc#available-implementations[EmbeddingModel] section for more information. - Now you can auto-wire the `ElasticsearchVectorStore` as a vector store in your application. [source,java] @@ -97,7 +94,7 @@ List documents = List.of( new Document("The World is Big and Salvation Lurks Around the Corner"), new Document("You walk forward facing the past and you turn back toward the future.", Map.of("meta2", "meta2"))); -// Add the documents to Qdrant +// Add the documents to Elasticsearch vectorStore.add(documents); // Retrieve documents similar to a query @@ -117,34 +114,19 @@ spring: uris: username: password: -# API key if needed, e.g. OpenAI ai: - openai: - api: - key: + vectorstore: + elasticsearch: + initialize-schema: true + index-name: custom-index + dimensions: 1536 + similarity: cosine + batching-strategy: TOKEN_COUNT # Optional: Controls how documents are batched for embedding ---- -environment variables, - -[source,bash] ----- -export SPRING_ELASTICSEARCH_URIS= -export SPRING_ELASTICSEARCH_USERNAME= -export SPRING_ELASTICSEARCH_PASSWORD= -# API key if needed, e.g. OpenAI -export SPRING_AI_OPENAI_API_KEY= ----- - -or can be a mix of those. -For example, if you want to store your password as an environment variable but keep the rest in the plain `application.yml` file. - -NOTE: If you choose to create a shell script for ease in future work, be sure to run it prior to starting your application by "sourcing" the file, i.e. `source .sh`. - -Spring Boot's auto-configuration feature for the Elasticsearch RestClient will create a bean instance that will be used by the `ElasticsearchVectorStore`. - The Spring Boot properties starting with `spring.elasticsearch.*` are used to configure the Elasticsearch client: -[stripes=even] +[cols="2,5,1",stripes=even] |=== |Property | Description | Default Value @@ -160,23 +142,24 @@ The Spring Boot properties starting with `spring.elasticsearch.*` are used to co | `spring.elasticsearch.socket-timeout` | Socket timeout used when communicating with Elasticsearch. | `30s` |=== -Properties starting with the `spring.ai.vectorstore.elasticsearch.*` prefix are used to configure `ElasticsearchVectorStore`. +Properties starting with `spring.ai.vectorstore.elasticsearch.*` are used to configure the `ElasticsearchVectorStore`: -[stripes=even] +[cols="2,5,1",stripes=even] |=== |Property | Description | Default Value -|`spring.ai.vectorstore.elasticsearch.initialize-schema`| Whether to initialize the required schema | `false` -|`spring.ai.vectorstore.elasticsearch.index-name` | The name of the index to store the vectors. | spring-ai-document-index -|`spring.ai.vectorstore.elasticsearch.dimensions` | The number of dimensions in the vector. | 1536 -|`spring.ai.vectorstore.elasticsearch.similarity` | The similarity function to use. | `cosine` +|`spring.ai.vectorstore.elasticsearch.initialize-schema`| Whether to initialize the required schema | `false` +|`spring.ai.vectorstore.elasticsearch.index-name` | The name of the index to store the vectors | `spring-ai-document-index` +|`spring.ai.vectorstore.elasticsearch.dimensions` | The number of dimensions in the vector | `1536` +|`spring.ai.vectorstore.elasticsearch.similarity` | The similarity function to use | `cosine` +|`spring.ai.vectorstore.elasticsearch.batching-strategy` | Strategy for batching documents when calculating embeddings. Options are `TOKEN_COUNT` or `FIXED_SIZE` | `TOKEN_COUNT` |=== The following similarity functions are available: -* cosine -* l2_norm -* dot_product +* `cosine` - Default, suitable for most use cases. Measures cosine similarity between vectors. +* `l2_norm` - Euclidean distance between vectors. Lower values indicate higher similarity. +* `dot_product` - Best performance for normalized vectors (e.g., OpenAI embeddings). More details about each in the https://www.elastic.co/guide/en/elasticsearch/reference/master/dense-vector.html#dense-vector-params[Elasticsearch Documentation] on dense vectors. @@ -206,7 +189,7 @@ vectorStore.similaritySearch(SearchRequest.defaults() .withTopK(TOP_K) .withSimilarityThreshold(SIMILARITY_THRESHOLD) .withFilterExpression(b.and( - b.in("john", "jill"), + b.in("author", "john", "jill"), b.eq("article_type", "blog")).build())); ---- @@ -247,7 +230,6 @@ dependencies { } ---- - Create an Elasticsearch `RestClient` bean. Read the link:https://www.elastic.co/guide/en/elasticsearch/client/java-api-client/current/java-rest-low-usage-initialization.html[Elasticsearch Documentation] for more in-depth information about the configuration of a custom RestClient. @@ -255,7 +237,7 @@ Read the link:https://www.elastic.co/guide/en/elasticsearch/client/java-api-clie ---- @Bean public RestClient restClient() { - RestClient.builder(new HttpHost("", 9200, "http")) + return RestClient.builder(new HttpHost("", 9200, "http")) .setDefaultHeaders(new Header[]{ new BasicHeader("Authorization", "Basic ") }) @@ -263,19 +245,29 @@ public RestClient restClient() { } ---- -and then create the `ElasticsearchVectorStore` bean: +Then create the `ElasticsearchVectorStore` bean using the builder pattern: [source,java] ---- @Bean -public ElasticsearchVectorStore vectorStore(EmbeddingModel embeddingModel, RestClient restClient) { - return new ElasticsearchVectorStore( restClient, embeddingModel); +public VectorStore vectorStore(RestClient restClient, EmbeddingModel embeddingModel) { + ElasticsearchVectorStoreOptions options = new ElasticsearchVectorStoreOptions(); + options.setIndexName("custom-index"); // Optional: defaults to "spring-ai-document-index" + options.setSimilarity(COSINE); // Optional: defaults to COSINE + options.setDimensions(1536); // Optional: defaults to model dimensions or 1536 + + return ElasticsearchVectorStore.builder() + .restClient(restClient) + .embeddingModel(embeddingModel) + .options(options) // Optional: use custom options + .initializeSchema(true) // Optional: defaults to false + .batchingStrategy(new TokenCountBatchingStrategy()) // Optional: defaults to TokenCountBatchingStrategy + .build(); } -// This can be any EmbeddingModel implementation. +// This can be any EmbeddingModel implementation @Bean public EmbeddingModel embeddingModel() { return new OpenAiEmbeddingModel(new OpenAiApi(System.getenv("OPENAI_API_KEY"))); } ---- - diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/elasticsearch/ElasticsearchVectorStoreAutoConfiguration.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/elasticsearch/ElasticsearchVectorStoreAutoConfiguration.java index d5e14d229..33d0c2810 100644 --- a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/elasticsearch/ElasticsearchVectorStoreAutoConfiguration.java +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/elasticsearch/ElasticsearchVectorStoreAutoConfiguration.java @@ -22,8 +22,8 @@ import org.elasticsearch.client.RestClient; import org.springframework.ai.embedding.BatchingStrategy; import org.springframework.ai.embedding.EmbeddingModel; import org.springframework.ai.embedding.TokenCountBatchingStrategy; -import org.springframework.ai.vectorstore.ElasticsearchVectorStore; -import org.springframework.ai.vectorstore.ElasticsearchVectorStoreOptions; +import org.springframework.ai.vectorstore.elasticsearch.ElasticsearchVectorStore; +import org.springframework.ai.vectorstore.elasticsearch.ElasticsearchVectorStoreOptions; import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention; import org.springframework.beans.factory.ObjectProvider; import org.springframework.boot.autoconfigure.AutoConfiguration; @@ -73,9 +73,15 @@ public class ElasticsearchVectorStoreAutoConfiguration { elasticsearchVectorStoreOptions.setSimilarity(properties.getSimilarity()); } - return new ElasticsearchVectorStore(elasticsearchVectorStoreOptions, restClient, embeddingModel, - properties.isInitializeSchema(), observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP), - customObservationConvention.getIfAvailable(() -> null), batchingStrategy); + return ElasticsearchVectorStore.builder() + .restClient(restClient) + .options(elasticsearchVectorStoreOptions) + .embeddingModel(embeddingModel) + .initializeSchema(properties.isInitializeSchema()) + .observationRegistry(observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP)) + .customObservationConvention(customObservationConvention.getIfAvailable(() -> null)) + .batchingStrategy(batchingStrategy) + .build(); } } diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/elasticsearch/ElasticsearchVectorStoreProperties.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/elasticsearch/ElasticsearchVectorStoreProperties.java index bbcf95cb5..876c13ab5 100644 --- a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/elasticsearch/ElasticsearchVectorStoreProperties.java +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/elasticsearch/ElasticsearchVectorStoreProperties.java @@ -17,7 +17,7 @@ package org.springframework.ai.autoconfigure.vectorstore.elasticsearch; import org.springframework.ai.autoconfigure.vectorstore.CommonVectorStoreProperties; -import org.springframework.ai.vectorstore.SimilarityFunction; +import org.springframework.ai.vectorstore.elasticsearch.SimilarityFunction; import org.springframework.boot.context.properties.ConfigurationProperties; /** diff --git a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/elasticsearch/ElasticsearchVectorStoreAutoConfigurationIT.java b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/elasticsearch/ElasticsearchVectorStoreAutoConfigurationIT.java index c9da9584b..f517a0306 100644 --- a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/elasticsearch/ElasticsearchVectorStoreAutoConfigurationIT.java +++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/elasticsearch/ElasticsearchVectorStoreAutoConfigurationIT.java @@ -33,9 +33,9 @@ import org.springframework.ai.autoconfigure.openai.OpenAiAutoConfiguration; import org.springframework.ai.autoconfigure.retry.SpringAiRetryAutoConfiguration; import org.springframework.ai.document.Document; import org.springframework.ai.observation.conventions.VectorStoreProvider; -import org.springframework.ai.vectorstore.ElasticsearchVectorStore; +import org.springframework.ai.vectorstore.elasticsearch.ElasticsearchVectorStore; import org.springframework.ai.vectorstore.SearchRequest; -import org.springframework.ai.vectorstore.SimilarityFunction; +import org.springframework.ai.vectorstore.elasticsearch.SimilarityFunction; import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext; import org.springframework.boot.autoconfigure.AutoConfigurations; import org.springframework.boot.autoconfigure.elasticsearch.ElasticsearchRestClientAutoConfiguration; diff --git a/vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/ElasticsearchAiSearchFilterExpressionConverter.java b/vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/elasticsearch/ElasticsearchAiSearchFilterExpressionConverter.java similarity index 98% rename from vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/ElasticsearchAiSearchFilterExpressionConverter.java rename to vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/elasticsearch/ElasticsearchAiSearchFilterExpressionConverter.java index 4a7dc588f..eb8b505a8 100644 --- a/vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/ElasticsearchAiSearchFilterExpressionConverter.java +++ b/vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/elasticsearch/ElasticsearchAiSearchFilterExpressionConverter.java @@ -14,7 +14,7 @@ * limitations under the License. */ -package org.springframework.ai.vectorstore; +package org.springframework.ai.vectorstore.elasticsearch; import java.text.ParseException; import java.text.SimpleDateFormat; diff --git a/vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/ElasticsearchVectorStore.java b/vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/elasticsearch/ElasticsearchVectorStore.java similarity index 62% rename from vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/ElasticsearchVectorStore.java rename to vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/elasticsearch/ElasticsearchVectorStore.java index 2e4cef70c..9bfe2e9b5 100644 --- a/vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/ElasticsearchVectorStore.java +++ b/vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/elasticsearch/ElasticsearchVectorStore.java @@ -14,7 +14,7 @@ * limitations under the License. */ -package org.springframework.ai.vectorstore; +package org.springframework.ai.vectorstore.elasticsearch; import java.io.IOException; import java.util.List; @@ -48,23 +48,100 @@ import org.springframework.ai.embedding.TokenCountBatchingStrategy; import org.springframework.ai.model.EmbeddingUtils; import org.springframework.ai.observation.conventions.VectorStoreProvider; import org.springframework.ai.observation.conventions.VectorStoreSimilarityMetric; +import org.springframework.ai.vectorstore.AbstractVectorStoreBuilder; +import org.springframework.ai.vectorstore.SearchRequest; import org.springframework.ai.vectorstore.filter.Filter; 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; /** - * The ElasticsearchVectorStore class implements the VectorStore interface and provides - * functionality for managing and querying documents in Elasticsearch. It uses an - * embedding model to generate vector representations of the documents and performs - * similarity searches based on these vectors. + * Elasticsearch-based vector store implementation using the dense_vector field type. * - * The ElasticsearchVectorStore class requires a RestClient and an EmbeddingModel to be - * instantiated. It also supports optional initialization of the Elasticsearch schema. + *

+ * The store uses an Elasticsearch index to persist vector embeddings along with their + * associated document content and metadata. The implementation leverages Elasticsearch's + * k-NN search capabilities for efficient similarity search operations. + *

+ * + *

+ * Features: + *

+ *
    + *
  • Automatic schema initialization with configurable index creation
    • + *
  • Support for multiple similarity functions: Cosine, L2 Norm, and Dot Product
    • + *
  • Metadata filtering using Elasticsearch query strings
    • + *
  • Configurable similarity thresholds for search results
    • + *
  • Batch processing support with configurable strategies
    • + *
  • Observation and metrics support through Micrometer
    • + *
+ * + *

+ * Basic usage example: + *

+ * 
{@code
+ * ElasticsearchVectorStore vectorStore = ElasticsearchVectorStore.builder()
+ *     .restClient(restClient)
+ *     .embeddingModel(embeddingModel)
+ *     .initializeSchema(true)
+ *     .build();
+ *
+ * // Add documents
+ * vectorStore.add(List.of(
+ *     new Document("content1", Map.of("key1", "value1")),
+ *     new Document("content2", Map.of("key2", "value2"))
+ * ));
+ *
+ * // Search with filters
+ * List results = vectorStore.similaritySearch(
+ *     SearchRequest.query("search text")
+ *         .withTopK(5)
+ *         .withSimilarityThreshold(0.7)
+ *         .withFilterExpression("key1 == 'value1'")
+ * );
+ * }
+ * + *

+ * Advanced configuration example: + *

+ * 
{@code
+ * ElasticsearchVectorStoreOptions options = new ElasticsearchVectorStoreOptions();
+ * options.setIndexName("custom_vectors");
+ * options.setSimilarity(SimilarityFunction.dot_product);
+ * options.setDimensions(1536);
+ *
+ * ElasticsearchVectorStore vectorStore = ElasticsearchVectorStore.builder()
+ *     .restClient(restClient)
+ *     .embeddingModel(embeddingModel)
+ *     .options(options)
+ *     .initializeSchema(true)
+ *     .batchingStrategy(new TokenCountBatchingStrategy())
+ *     .build();
+ * }
+ * + *

+ * Requirements: + *

+ *
    + *
  • Elasticsearch 8.0 or later
    • + *
  • Index mapping with id (string), content (text), metadata (object), and embedding + * (dense_vector) fields
    • + *
+ * + *

+ * Similarity Functions: + *

+ *
    + *
  • cosine: Default, suitable for most use cases. Measures cosine similarity between + * vectors.
    • + *
  • l2_norm: Euclidean distance between vectors. Lower values indicate higher + * similarity.
    • + *
  • dot_product: Best performance for normalized vectors (e.g., OpenAI + * embeddings).
    • + *
* * @author Jemin Huh * @author Wei Jiang @@ -83,8 +160,6 @@ public class ElasticsearchVectorStore extends AbstractObservationVectorStore imp SimilarityFunction.cosine, VectorStoreSimilarityMetric.COSINE, SimilarityFunction.l2_norm, VectorStoreSimilarityMetric.EUCLIDEAN, SimilarityFunction.dot_product, VectorStoreSimilarityMetric.DOT); - private final EmbeddingModel embeddingModel; - private final ElasticsearchClient elasticsearchClient; private final ElasticsearchVectorStoreOptions options; @@ -95,34 +170,47 @@ public class ElasticsearchVectorStore extends AbstractObservationVectorStore imp private final BatchingStrategy batchingStrategy; + @Deprecated(since = "1.0.0-M5", forRemoval = true) public ElasticsearchVectorStore(RestClient restClient, EmbeddingModel embeddingModel, boolean initializeSchema) { this(new ElasticsearchVectorStoreOptions(), restClient, embeddingModel, initializeSchema); } + @Deprecated(since = "1.0.0-M5", forRemoval = true) public ElasticsearchVectorStore(ElasticsearchVectorStoreOptions options, RestClient restClient, EmbeddingModel embeddingModel, boolean initializeSchema) { this(options, restClient, embeddingModel, initializeSchema, ObservationRegistry.NOOP, null, new TokenCountBatchingStrategy()); } + @Deprecated(since = "1.0.0-M5", forRemoval = true) public ElasticsearchVectorStore(ElasticsearchVectorStoreOptions options, RestClient restClient, EmbeddingModel embeddingModel, boolean initializeSchema, ObservationRegistry observationRegistry, VectorStoreObservationConvention customObservationConvention, BatchingStrategy batchingStrategy) { - super(observationRegistry, customObservationConvention); + this(builder().restClient(restClient) + .options(options) + .embeddingModel(embeddingModel) + .initializeSchema(initializeSchema) + .observationRegistry(observationRegistry) + .customObservationConvention(customObservationConvention) + .batchingStrategy(batchingStrategy)); + } + + protected ElasticsearchVectorStore(ElasticsearchBuilder builder) { + super(builder); + + Assert.notNull(builder.restClient, "RestClient must not be null"); + + this.initializeSchema = builder.initializeSchema; + this.options = builder.options; + this.filterExpressionConverter = builder.filterExpressionConverter; + this.batchingStrategy = builder.batchingStrategy; - this.initializeSchema = initializeSchema; - Objects.requireNonNull(embeddingModel, "RestClient must not be null"); - Objects.requireNonNull(embeddingModel, "EmbeddingModel must not be null"); String version = Version.VERSION == null ? "Unknown" : Version.VERSION.toString(); - this.elasticsearchClient = new ElasticsearchClient(new RestClientTransport(restClient, + this.elasticsearchClient = new ElasticsearchClient(new RestClientTransport(builder.restClient, new JacksonJsonpMapper( new ObjectMapper().configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false)))) .withTransportOptions(t -> t.addHeader("user-agent", "spring-ai elastic-java/" + version)); - this.embeddingModel = embeddingModel; - this.options = options; - this.filterExpressionConverter = new ElasticsearchAiSearchFilterExpressionConverter(); - this.batchingStrategy = batchingStrategy; } @Override @@ -297,4 +385,95 @@ public class ElasticsearchVectorStore extends AbstractObservationVectorStore imp public record ElasticSearchDocument(String id, String content, Map metadata, float[] embedding) { } + /** + * Creates a new builder instance for ElasticsearchVectorStore. + * @return a new ElasticsearchBuilder instance + */ + public static ElasticsearchBuilder builder() { + return new ElasticsearchBuilder(); + } + + public static class ElasticsearchBuilder extends AbstractVectorStoreBuilder { + + private RestClient restClient; + + private ElasticsearchVectorStoreOptions options = new ElasticsearchVectorStoreOptions(); + + private boolean initializeSchema = false; + + private BatchingStrategy batchingStrategy = new TokenCountBatchingStrategy(); + + private FilterExpressionConverter filterExpressionConverter = new ElasticsearchAiSearchFilterExpressionConverter(); + + /** + * Sets the Elasticsearch REST client. + * @param restClient the Elasticsearch REST client + * @return the builder instance + * @throws IllegalArgumentException if restClient is null + */ + public ElasticsearchBuilder restClient(RestClient restClient) { + Assert.notNull(restClient, "RestClient must not be null"); + this.restClient = restClient; + return this; + } + + /** + * Sets the Elasticsearch vector store options. + * @param options the vector store options to use + * @return the builder instance + * @throws IllegalArgumentException if options is null + */ + public ElasticsearchBuilder options(ElasticsearchVectorStoreOptions options) { + Assert.notNull(options, "options must not be null"); + this.options = options; + return this; + } + + /** + * Sets whether to initialize the schema. + * @param initializeSchema true to initialize schema, false otherwise + * @return the builder instance + */ + public ElasticsearchBuilder initializeSchema(boolean initializeSchema) { + this.initializeSchema = initializeSchema; + return this; + } + + /** + * Sets the batching strategy for vector operations. + * @param batchingStrategy the batching strategy to use + * @return the builder instance + * @throws IllegalArgumentException if batchingStrategy is null + */ + public ElasticsearchBuilder batchingStrategy(BatchingStrategy batchingStrategy) { + Assert.notNull(batchingStrategy, "batchingStrategy must not be null"); + this.batchingStrategy = batchingStrategy; + return this; + } + + /** + * Sets the filter expression converter. + * @param converter the filter expression converter to use + * @return the builder instance + * @throws IllegalArgumentException if converter is null + */ + public ElasticsearchBuilder filterExpressionConverter(FilterExpressionConverter converter) { + Assert.notNull(converter, "filterExpressionConverter must not be null"); + this.filterExpressionConverter = converter; + return this; + } + + /** + * Builds the ElasticsearchVectorStore instance. + * @return a new ElasticsearchVectorStore instance + * @throws IllegalStateException if the builder is in an invalid state + */ + @Override + public ElasticsearchVectorStore build() { + validate(); + return new ElasticsearchVectorStore(this); + } + + } + } diff --git a/vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/ElasticsearchVectorStoreOptions.java b/vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/elasticsearch/ElasticsearchVectorStoreOptions.java similarity index 96% rename from vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/ElasticsearchVectorStoreOptions.java rename to vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/elasticsearch/ElasticsearchVectorStoreOptions.java index 8aee18ac0..2958b6e66 100644 --- a/vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/ElasticsearchVectorStoreOptions.java +++ b/vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/elasticsearch/ElasticsearchVectorStoreOptions.java @@ -14,7 +14,7 @@ * limitations under the License. */ -package org.springframework.ai.vectorstore; +package org.springframework.ai.vectorstore.elasticsearch; /** * Provided Elasticsearch vector option configuration. diff --git a/vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/SimilarityFunction.java b/vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/elasticsearch/SimilarityFunction.java similarity index 94% rename from vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/SimilarityFunction.java rename to vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/elasticsearch/SimilarityFunction.java index b28e7313d..f551b271c 100644 --- a/vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/SimilarityFunction.java +++ b/vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/elasticsearch/SimilarityFunction.java @@ -14,7 +14,7 @@ * limitations under the License. */ -package org.springframework.ai.vectorstore; +package org.springframework.ai.vectorstore.elasticsearch; /** * https://www.elastic.co/guide/en/elasticsearch/reference/master/dense-vector.html diff --git a/vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/ElasticsearchAiSearchFilterExpressionConverterTest.java b/vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/elasticsearch/ElasticsearchAiSearchFilterExpressionConverterTest.java similarity index 98% rename from vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/ElasticsearchAiSearchFilterExpressionConverterTest.java rename to vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/elasticsearch/ElasticsearchAiSearchFilterExpressionConverterTest.java index 7a6f737ee..8a366533a 100644 --- a/vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/ElasticsearchAiSearchFilterExpressionConverterTest.java +++ b/vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/elasticsearch/ElasticsearchAiSearchFilterExpressionConverterTest.java @@ -14,7 +14,7 @@ * limitations under the License. */ -package org.springframework.ai.vectorstore; +package org.springframework.ai.vectorstore.elasticsearch; import java.util.Date; import java.util.List; diff --git a/vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/ElasticsearchImage.java b/vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/elasticsearch/ElasticsearchImage.java similarity index 94% rename from vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/ElasticsearchImage.java rename to vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/elasticsearch/ElasticsearchImage.java index 6cc156f35..4af1decdc 100644 --- a/vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/ElasticsearchImage.java +++ b/vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/elasticsearch/ElasticsearchImage.java @@ -14,7 +14,7 @@ * limitations under the License. */ -package org.springframework.ai.vectorstore; +package org.springframework.ai.vectorstore.elasticsearch; import org.testcontainers.utility.DockerImageName; diff --git a/vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/ElasticsearchVectorStoreIT.java b/vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/elasticsearch/ElasticsearchVectorStoreIT.java similarity index 96% rename from vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/ElasticsearchVectorStoreIT.java rename to vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/elasticsearch/ElasticsearchVectorStoreIT.java index 2372a5d9c..4f1de5fac 100644 --- a/vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/ElasticsearchVectorStoreIT.java +++ b/vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/elasticsearch/ElasticsearchVectorStoreIT.java @@ -14,7 +14,7 @@ * limitations under the License. */ -package org.springframework.ai.vectorstore; +package org.springframework.ai.vectorstore.elasticsearch; import java.io.IOException; import java.nio.charset.StandardCharsets; @@ -51,6 +51,7 @@ import org.springframework.ai.document.Document; import org.springframework.ai.embedding.EmbeddingModel; import org.springframework.ai.openai.OpenAiEmbeddingModel; import org.springframework.ai.openai.api.OpenAiApi; +import org.springframework.ai.vectorstore.SearchRequest; import org.springframework.boot.SpringBootConfiguration; import org.springframework.boot.autoconfigure.EnableAutoConfiguration; import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration; @@ -376,7 +377,11 @@ class ElasticsearchVectorStoreIT { @Bean("vectorStore_cosine") public ElasticsearchVectorStore vectorStoreDefault(EmbeddingModel embeddingModel, RestClient restClient) { - return new ElasticsearchVectorStore(restClient, embeddingModel, true); + return ElasticsearchVectorStore.builder() + .restClient(restClient) + .embeddingModel(embeddingModel) + .initializeSchema(true) + .build(); } @Bean("vectorStore_l2_norm") @@ -384,7 +389,12 @@ class ElasticsearchVectorStoreIT { ElasticsearchVectorStoreOptions options = new ElasticsearchVectorStoreOptions(); options.setIndexName("index_l2"); options.setSimilarity(SimilarityFunction.l2_norm); - return new ElasticsearchVectorStore(options, restClient, embeddingModel, true); + return ElasticsearchVectorStore.builder() + .restClient(restClient) + .embeddingModel(embeddingModel) + .initializeSchema(true) + .options(options) + .build(); } @Bean("vectorStore_dot_product") @@ -392,7 +402,12 @@ class ElasticsearchVectorStoreIT { ElasticsearchVectorStoreOptions options = new ElasticsearchVectorStoreOptions(); options.setIndexName("index_dot_product"); options.setSimilarity(SimilarityFunction.dot_product); - return new ElasticsearchVectorStore(options, restClient, embeddingModel, true); + return ElasticsearchVectorStore.builder() + .restClient(restClient) + .embeddingModel(embeddingModel) + .initializeSchema(true) + .options(options) + .build(); } @Bean diff --git a/vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/ElasticsearchVectorStoreObservationIT.java b/vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/elasticsearch/ElasticsearchVectorStoreObservationIT.java similarity index 94% rename from vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/ElasticsearchVectorStoreObservationIT.java rename to vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/elasticsearch/ElasticsearchVectorStoreObservationIT.java index 1d074fba0..ae19e2ff0 100644 --- a/vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/ElasticsearchVectorStoreObservationIT.java +++ b/vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/elasticsearch/ElasticsearchVectorStoreObservationIT.java @@ -14,7 +14,7 @@ * limitations under the License. */ -package org.springframework.ai.vectorstore; +package org.springframework.ai.vectorstore.elasticsearch; import java.io.IOException; import java.nio.charset.StandardCharsets; @@ -51,6 +51,8 @@ import org.springframework.ai.observation.conventions.VectorStoreProvider; import org.springframework.ai.observation.conventions.VectorStoreSimilarityMetric; import org.springframework.ai.openai.OpenAiEmbeddingModel; import org.springframework.ai.openai.api.OpenAiApi; +import org.springframework.ai.vectorstore.SearchRequest; +import org.springframework.ai.vectorstore.VectorStore; import org.springframework.ai.vectorstore.observation.DefaultVectorStoreObservationConvention; import org.springframework.ai.vectorstore.observation.VectorStoreObservationDocumentation.HighCardinalityKeyNames; import org.springframework.ai.vectorstore.observation.VectorStoreObservationDocumentation.LowCardinalityKeyNames; @@ -67,6 +69,7 @@ import static org.hamcrest.Matchers.greaterThan; /** * @author Christian Tzolov * @author Thomas Vitale + * @author Soby Chacko */ @Testcontainers @EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+") @@ -205,8 +208,15 @@ public class ElasticsearchVectorStoreObservationIT { @Bean public ElasticsearchVectorStore vectorStoreDefault(EmbeddingModel embeddingModel, RestClient restClient, ObservationRegistry observationRegistry) { - return new ElasticsearchVectorStore(new ElasticsearchVectorStoreOptions(), restClient, embeddingModel, true, - observationRegistry, null, new TokenCountBatchingStrategy()); + return ElasticsearchVectorStore.builder() + .restClient(restClient) + .embeddingModel(embeddingModel) + .initializeSchema(true) + .options(new ElasticsearchVectorStoreOptions()) + .observationRegistry(observationRegistry) + .customObservationConvention(null) + .batchingStrategy(new TokenCountBatchingStrategy()) + .build(); } @Bean