From cde8f7446ca733aeeccc17e868bc93f6c65614e8 Mon Sep 17 00:00:00 2001 From: Soby Chacko Date: Mon, 9 Dec 2024 19:03:18 -0500 Subject: [PATCH] Add builder pattern and refactor Neo4jVectorStore The Neo4j vector store implementation has been enhanced with a builder pattern to be more intuitive than using ctors and follows spring ai builder conventions. Current constructors have been deprecated to maintain backward compatibility for one releaes cycle. The change includes: * Move classes to dedicated neo4j package for better organization * Add comprehensive builder pattern implementation with validation * Improve documentation with detailed usage examples * Deprecate but maintain old configuration approach for compatibility * Update integration tests to demonstrate new builder pattern * Enhance code readability and maintainability --- .../ROOT/pages/api/vectordbs/neo4j.adoc | 252 +++++++----- .../Neo4jVectorStoreAutoConfiguration.java | 31 +- .../neo4j/Neo4jVectorStoreProperties.java | 2 +- .../{ => neo4j}/Neo4jVectorStore.java | 372 ++++++++++++++++-- .../Neo4jVectorFilterExpressionConverter.java | 3 +- .../vectorstore/{ => neo4j}/Neo4jImage.java | 2 +- .../{ => neo4j}/Neo4jVectorStoreIT.java | 11 +- .../Neo4jVectorStoreObservationIT.java | 14 +- ...jVectorFilterExpressionConverterTests.java | 6 +- 9 files changed, 525 insertions(+), 168 deletions(-) rename vector-stores/spring-ai-neo4j-store/src/main/java/org/springframework/ai/vectorstore/{ => neo4j}/Neo4jVectorStore.java (55%) rename vector-stores/spring-ai-neo4j-store/src/main/java/org/springframework/ai/vectorstore/{ => neo4j}/filter/Neo4jVectorFilterExpressionConverter.java (96%) rename vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/{ => neo4j}/Neo4jImage.java (94%) rename vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/{ => neo4j}/Neo4jVectorStoreIT.java (97%) rename vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/{ => neo4j}/Neo4jVectorStoreObservationIT.java (94%) rename vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/{ => neo4j}/filter/Neo4jVectorFilterExpressionConverterTests.java (95%) diff --git a/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs/neo4j.adoc b/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs/neo4j.adoc index 51cd2d94f..a5f3ad472 100644 --- a/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs/neo4j.adoc +++ b/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs/neo4j.adoc @@ -20,77 +20,12 @@ Those indexes are powered by Lucene using a Hierarchical Navigable Small World G ** link:https://neo4j.com/deployment-center/[Neo4j Server] instance * If required, an API key for the xref:api/embeddings.adoc#available-implementations[EmbeddingModel] to generate the embeddings stored by the `Neo4jVectorStore`. -== Dependencies - -Add the Neo4j Vector Store dependency to your project: - -[source,xml] ----- - - org.springframework.ai - spring-ai-neo4j-store - ----- - -or to your Gradle `build.gradle` build file. - -[source,groovy] ----- -dependencies { - implementation 'org.springframework.ai:spring-ai-neo4j-store' -} ----- - - - -TIP: Refer to the xref:getting-started.adoc#dependency-management[Dependency Management] section to add the Spring AI BOM to your build file. - - -The vector store implementation can initialize the requisite schema for you, but you must opt-in by specifying the `initializeSchema` boolean in the appropriate constructor or by setting `...initialize-schema=true` in the `application.properties` file. - -NOTE: this is a breaking change! In earlier versions of Spring AI, this schema initialization happened by default. - - -== Configuration - -To connect to Neo4j and use the `Neo4jVectorStore`, you need to provide access details for your instance. -A simple configuration can either be provided via Spring Boot's _application.properties_, - -[source,properties] ----- -spring.neo4j.uri= -spring.neo4j.authentication.username= -spring.neo4j.authentication.password= -# API key if needed, e.g. OpenAI -spring.ai.openai.api.key= ----- - -environment variables, - -[source,bash] ----- -export SPRING_NEO4J_URI= -export SPRING_NEO4J_AUTHENTICATION_USERNAME= -export SPRING_NEO4J_AUTHENTICATION_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 API key as an environment variable but keep the rest in the plain _application.properties_ 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`. - -NOTE: Besides _application.properties_ and environment variables, Spring Boot offers https://docs.spring.io/spring-boot/docs/current/reference/html/features.html#features.external-config[additional configuration options]. - -Spring Boot's auto-configuration feature for the Neo4j Driver will create a bean instance that will be used by the `Neo4jVectorStore`. - == Auto-configuration Spring AI provides Spring Boot auto-configuration for the Neo4j Vector Store. To enable it, add the following dependency to your project's Maven `pom.xml` file: -[source, xml] +[source,xml] ---- org.springframework.ai @@ -113,21 +48,114 @@ Please have a look at the list of xref:#_neo4jvectorstore_properties[configurati TIP: Refer to the xref:getting-started.adoc#repositories[Repositories] section to add Milestone and/or Snapshot Repositories to your build file. +The vector store implementation can initialize the requisite schema for you, but you must opt-in by specifying the `initializeSchema` boolean in the appropriate constructor or by setting `...initialize-schema=true` in the `application.properties` file. + +NOTE: this is a breaking change! In earlier versions of Spring AI, this schema initialization happened by default. + Additionally, you will need a configured `EmbeddingModel` bean. Refer to the xref:api/embeddings.adoc#available-implementations[EmbeddingModel] section for more information. -Here is an example of the needed bean: +Now you can auto-wire the `Neo4jVectorStore` as a vector store in your application. [source,java] ---- -@Bean -public EmbeddingModel embeddingModel() { - // Can be any other Embeddingmodel implementation. - return new OpenAiEmbeddingModel(new OpenAiApi(System.getenv("SPRING_AI_OPENAI_API_KEY"))); +@Autowired VectorStore vectorStore; + +// ... + +List documents = List.of( + new Document("Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!!", Map.of("meta1", "meta1")), + 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 Neo4j +vectorStore.add(documents); + +// Retrieve documents similar to a query +List results = vectorStore.similaritySearch(SearchRequest.query("Spring").withTopK(5)); +---- + +[[neo4jvector-properties]] +=== Configuration Properties + +To connect to Neo4j and use the `Neo4jVectorStore`, you need to provide access details for your instance. +A simple configuration can be provided via Spring Boot's `application.yml`: + +[source,yaml] +---- +spring: + neo4j: + uri: + authentication: + username: + password: + ai: + vectorstore: + neo4j: + initialize-schema: true + database-name: neo4j + index-name: custom-index + dimensions: 1536 + distance-type: cosine + batching-strategy: TOKEN_COUNT # Optional: Controls how documents are batched for embedding +---- + +The Spring Boot properties starting with `spring.neo4j.*` are used to configure the Neo4j client: + +[cols="2,5,1",stripes=even] +|=== +|Property | Description | Default Value + +| `spring.neo4j.uri` | URI for connecting to the Neo4j instance | `neo4j://localhost:7687` +| `spring.neo4j.authentication.username` | Username for authentication with Neo4j | `neo4j` +| `spring.neo4j.authentication.password` | Password for authentication with Neo4j | - +|=== + +Properties starting with `spring.ai.vectorstore.neo4j.*` are used to configure the `Neo4jVectorStore`: + +[cols="2,5,1",stripes=even] +|=== +|Property | Description | Default Value + +|`spring.ai.vectorstore.neo4j.initialize-schema`| Whether to initialize the required schema | `false` +|`spring.ai.vectorstore.neo4j.database-name` | The name of the Neo4j database to use | `neo4j` +|`spring.ai.vectorstore.neo4j.index-name` | The name of the index to store the vectors | `spring-ai-document-index` +|`spring.ai.vectorstore.neo4j.dimensions` | The number of dimensions in the vector | `1536` +|`spring.ai.vectorstore.neo4j.distance-type` | The distance function to use | `cosine` +|`spring.ai.vectorstore.neo4j.label` | The label used for document nodes | `Document` +|`spring.ai.vectorstore.neo4j.embedding-property` | The property name used to store embeddings | `embedding` +|`spring.ai.vectorstore.neo4j.batching-strategy` | Strategy for batching documents when calculating embeddings. Options are `TOKEN_COUNT` or `FIXED_SIZE` | `TOKEN_COUNT` +|=== + +The following distance functions are available: + +* `cosine` - Default, suitable for most use cases. Measures cosine similarity between vectors. +* `euclidean` - Euclidean distance between vectors. Lower values indicate higher similarity. + +== Manual Configuration + +Instead of using the Spring Boot auto-configuration, you can manually configure the Neo4j vector store. For this you need to add the `spring-ai-neo4j-store` to your project: + +[source,xml] +---- + + org.springframework.ai + spring-ai-neo4j-store + +---- + +or to your Gradle `build.gradle` build file. + +[source,groovy] +---- +dependencies { + implementation 'org.springframework.ai:spring-ai-neo4j-store' } ---- -In cases where the Spring Boot auto-configured Neo4j `Driver` bean is not what you want or need, you can still define your own bean. -Please read the https://neo4j.com/docs/java-manual/current/client-applications/[Neo4j Java Driver reference] for more in-depth information about the configuration of a custom driver. +TIP: Refer to the xref:getting-started.adoc#dependency-management[Dependency Management] section to add the Spring AI BOM to your build file. + +Create a Neo4j `Driver` bean. +Read the link:https://neo4j.com/docs/java-manual/current/client-applications/[Neo4j Documentation] for more in-depth information about the configuration of a custom driver. [source,java] ---- @@ -138,9 +166,34 @@ public Driver driver() { } ---- -Now you can auto-wire the `Neo4jVectorStore` as a vector store in your application. +Then create the `Neo4jVectorStore` bean using the builder pattern: -== Metadata filtering +[source,java] +---- +@Bean +public VectorStore vectorStore(Driver driver, EmbeddingModel embeddingModel) { + return Neo4jVectorStore.builder() + .driver(driver) + .embeddingModel(embeddingModel) + .databaseName("neo4j") // Optional: defaults to "neo4j" + .distanceType(Neo4jDistanceType.COSINE) // Optional: defaults to COSINE + .dimensions(1536) // Optional: defaults to 1536 + .label("Document") // Optional: defaults to "Document" + .embeddingProperty("embedding") // Optional: defaults to "embedding" + .indexName("custom-index") // Optional: defaults to "spring-ai-document-index" + .initializeSchema(true) // Optional: defaults to false + .batchingStrategy(new TokenCountBatchingStrategy()) // Optional: defaults to TokenCountBatchingStrategy + .build(); +} + +// This can be any EmbeddingModel implementation +@Bean +public EmbeddingModel embeddingModel() { + return new OpenAiEmbeddingModel(new OpenAiApi(System.getenv("OPENAI_API_KEY"))); +} +---- + +== Metadata Filtering You can leverage the generic, portable xref:api/vectordbs.adoc#metadata-filters[metadata filters] with Neo4j store as well. @@ -149,11 +202,11 @@ For example, you can use either the text expression language: [source,java] ---- vectorStore.similaritySearch( - SearchRequest.defaults() - .withQuery("The World") - .withTopK(TOP_K) - .withSimilarityThreshold(SIMILARITY_THRESHOLD) - .withFilterExpression("author in ['john', 'jill'] && 'article_type' == 'blog'")); + SearchRequest.defaults() + .withQuery("The World") + .withTopK(TOP_K) + .withSimilarityThreshold(SIMILARITY_THRESHOLD) + .withFilterExpression("author in ['john', 'jill'] && 'article_type' == 'blog'")); ---- or programmatically using the `Filter.Expression` DSL: @@ -163,41 +216,26 @@ or programmatically using the `Filter.Expression` DSL: FilterExpressionBuilder b = new FilterExpressionBuilder(); vectorStore.similaritySearch(SearchRequest.defaults() - .withQuery("The World") - .withTopK(TOP_K) - .withSimilarityThreshold(SIMILARITY_THRESHOLD) - .withFilterExpression(b.and( - b.in("author", "john", "jill"), - b.eq("article_type", "blog")).build())); + .withQuery("The World") + .withTopK(TOP_K) + .withSimilarityThreshold(SIMILARITY_THRESHOLD) + .withFilterExpression(b.and( + b.in("author", "john", "jill"), + b.eq("article_type", "blog")).build())); ---- NOTE: Those (portable) filter expressions get automatically converted into the proprietary Neo4j `WHERE` link:https://neo4j.com/developer/cypher/filtering-query-results/[filter expressions]. For example, this portable filter expression: -```sql +[source,sql] +---- author in ['john', 'jill'] && 'article_type' == 'blog' -``` +---- is converted into the proprietary Neo4j filter format: -``` +[source,text] +---- node.`metadata.author` IN ["john","jill"] AND node.`metadata.'article_type'` = "blog" -``` - -== Neo4jVectorStore properties - -You can use the following properties in your Spring Boot configuration to customize the Neo4j vector store. - -[stripes=even] -|=== -|Property|Default value - -|`spring.ai.vectorstore.neo4j.database-name`|neo4j -|`spring.ai.vectorstore.neo4j.initialize-schema`|false -|`spring.ai.vectorstore.neo4j.embedding-dimension`|1536 -|`spring.ai.vectorstore.neo4j.distance-type`|cosine -|`spring.ai.vectorstore.neo4j.label`|Document -|`spring.ai.vectorstore.neo4j.embedding-property`|embedding -|`spring.ai.vectorstore.neo4j.index-name`|spring-ai-document-index -|=== +---- diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/neo4j/Neo4jVectorStoreAutoConfiguration.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/neo4j/Neo4jVectorStoreAutoConfiguration.java index 0c65588cb..aac5576d1 100644 --- a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/neo4j/Neo4jVectorStoreAutoConfiguration.java +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/neo4j/Neo4jVectorStoreAutoConfiguration.java @@ -22,7 +22,7 @@ import org.neo4j.driver.Driver; import org.springframework.ai.embedding.BatchingStrategy; import org.springframework.ai.embedding.EmbeddingModel; import org.springframework.ai.embedding.TokenCountBatchingStrategy; -import org.springframework.ai.vectorstore.Neo4jVectorStore; +import org.springframework.ai.vectorstore.neo4j.Neo4jVectorStore; import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention; import org.springframework.beans.factory.ObjectProvider; import org.springframework.boot.autoconfigure.AutoConfiguration; @@ -57,20 +57,23 @@ public class Neo4jVectorStoreAutoConfiguration { Neo4jVectorStoreProperties properties, ObjectProvider observationRegistry, ObjectProvider customObservationConvention, BatchingStrategy batchingStrategy) { - Neo4jVectorStore.Neo4jVectorStoreConfig config = Neo4jVectorStore.Neo4jVectorStoreConfig.builder() - .withDatabaseName(properties.getDatabaseName()) - .withEmbeddingDimension(properties.getEmbeddingDimension()) - .withDistanceType(properties.getDistanceType()) - .withLabel(properties.getLabel()) - .withEmbeddingProperty(properties.getEmbeddingProperty()) - .withIndexName(properties.getIndexName()) - .withIdProperty(properties.getIdProperty()) - .withConstraintName(properties.getConstraintName()) - .build(); - return new Neo4jVectorStore(driver, embeddingModel, config, properties.isInitializeSchema(), - observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP), - customObservationConvention.getIfAvailable(() -> null), batchingStrategy); + return Neo4jVectorStore.builder() + .driver(driver) + .embeddingModel(embeddingModel) + .initializeSchema(properties.isInitializeSchema()) + .observationRegistry(observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP)) + .customObservationConvention(customObservationConvention.getIfAvailable(() -> null)) + .batchingStrategy(batchingStrategy) + .databaseName(properties.getDatabaseName()) + .embeddingDimension(properties.getEmbeddingDimension()) + .distanceType(properties.getDistanceType()) + .label(properties.getLabel()) + .embeddingProperty(properties.getEmbeddingProperty()) + .indexName(properties.getIndexName()) + .idProperty(properties.getIdProperty()) + .constraintName(properties.getConstraintName()) + .build(); } } diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/neo4j/Neo4jVectorStoreProperties.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/neo4j/Neo4jVectorStoreProperties.java index b6a4f1547..52606c754 100644 --- a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/neo4j/Neo4jVectorStoreProperties.java +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/neo4j/Neo4jVectorStoreProperties.java @@ -17,7 +17,7 @@ package org.springframework.ai.autoconfigure.vectorstore.neo4j; import org.springframework.ai.autoconfigure.vectorstore.CommonVectorStoreProperties; -import org.springframework.ai.vectorstore.Neo4jVectorStore; +import org.springframework.ai.vectorstore.neo4j.Neo4jVectorStore; import org.springframework.boot.context.properties.ConfigurationProperties; /** diff --git a/vector-stores/spring-ai-neo4j-store/src/main/java/org/springframework/ai/vectorstore/Neo4jVectorStore.java b/vector-stores/spring-ai-neo4j-store/src/main/java/org/springframework/ai/vectorstore/neo4j/Neo4jVectorStore.java similarity index 55% rename from vector-stores/spring-ai-neo4j-store/src/main/java/org/springframework/ai/vectorstore/Neo4jVectorStore.java rename to vector-stores/spring-ai-neo4j-store/src/main/java/org/springframework/ai/vectorstore/neo4j/Neo4jVectorStore.java index 989084504..f8e9edbbe 100644 --- a/vector-stores/spring-ai-neo4j-store/src/main/java/org/springframework/ai/vectorstore/Neo4jVectorStore.java +++ b/vector-stores/spring-ai-neo4j-store/src/main/java/org/springframework/ai/vectorstore/neo4j/Neo4jVectorStore.java @@ -14,7 +14,7 @@ * limitations under the License. */ -package org.springframework.ai.vectorstore; +package org.springframework.ai.vectorstore.neo4j; import java.util.HashMap; import java.util.List; @@ -36,15 +36,99 @@ import org.springframework.ai.embedding.EmbeddingOptionsBuilder; import org.springframework.ai.embedding.TokenCountBatchingStrategy; import org.springframework.ai.observation.conventions.VectorStoreProvider; import org.springframework.ai.observation.conventions.VectorStoreSimilarityMetric; -import org.springframework.ai.vectorstore.filter.Neo4jVectorFilterExpressionConverter; +import org.springframework.ai.vectorstore.AbstractVectorStoreBuilder; +import org.springframework.ai.vectorstore.SearchRequest; +import org.springframework.ai.vectorstore.neo4j.filter.Neo4jVectorFilterExpressionConverter; import org.springframework.ai.vectorstore.observation.AbstractObservationVectorStore; import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext; import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention; import org.springframework.beans.factory.InitializingBean; import org.springframework.util.Assert; +import org.springframework.util.StringUtils; /** - * A vector store implementation that stores and retrieves vectors in a Neo4j database. + * Neo4j-based vector store implementation using Neo4j's vector search capabilities. + * + *

+ * The store uses Neo4j's vector search functionality to persist and query vector + * embeddings along with their associated document content and metadata. The + * implementation leverages Neo4j's HNSW (Hierarchical Navigable Small World) algorithm + * for efficient k-NN search operations. + *

+ * + *

+ * Features: + *

+ *
    + *
  • Automatic schema initialization with configurable index creation
    • + *
  • Support for multiple distance functions: Cosine and Euclidean
    • + *
  • Metadata filtering using Neo4j's WHERE clause expressions
    • + *
  • Configurable similarity thresholds for search results
    • + *
  • Batch processing support with configurable strategies
    • + *
  • Observation and metrics support through Micrometer
    • + *
+ * + *

+ * Basic usage example: + *

+ * 
{@code
+ * Neo4jVectorStore vectorStore = Neo4jVectorStore.builder()
+ *     .driver(driver)
+ *     .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
+ * Neo4jVectorStore vectorStore = Neo4jVectorStore.builder()
+ *     .driver(driver)
+ *     .embeddingModel(embeddingModel)
+ *     .databaseName("neo4j")
+ *     .distanceType(Neo4jDistanceType.COSINE)
+ *     .dimensions(1536)
+ *     .label("CustomDocument")
+ *     .embeddingProperty("vector")
+ *     .indexName("custom-vectors")
+ *     .initializeSchema(true)
+ *     .batchingStrategy(new TokenCountBatchingStrategy())
+ *     .build();
+ * }
+ * + *

+ * Requirements: + *

+ *
    + *
  • Neo4j 5.15 or later
    • + *
  • Node schema with id (string), text (string), metadata (object), and embedding + * (vector) properties
    • + *
+ * + *

+ * Distance Functions: + *

+ *
    + *
  • cosine: Default, suitable for most use cases. Measures cosine similarity between + * vectors.
    • + *
  • euclidean: Euclidean distance between vectors. Lower values indicate higher + * similarity.
    • + *
* * @author Gerrit Meier * @author Michael Simons @@ -52,6 +136,7 @@ import org.springframework.util.Assert; * @author Thomas Vitale * @author Soby Chacko * @author Jihoon Kim + * @since 1.0.0 */ public class Neo4jVectorStore extends AbstractObservationVectorStore implements InitializingBean { @@ -73,36 +158,77 @@ public class Neo4jVectorStore extends AbstractObservationVectorStore implements Neo4jDistanceType.COSINE, VectorStoreSimilarityMetric.COSINE, Neo4jDistanceType.EUCLIDEAN, VectorStoreSimilarityMetric.EUCLIDEAN); - private final Neo4jVectorFilterExpressionConverter filterExpressionConverter = new Neo4jVectorFilterExpressionConverter(); - private final Driver driver; - private final EmbeddingModel embeddingModel; + private final SessionConfig sessionConfig; - private final Neo4jVectorStoreConfig config; + private final int embeddingDimension; + + private final Neo4jDistanceType distanceType; + + private final String embeddingProperty; + + private final String label; + + private final String indexName; + + private final String indexNameNotSanitized; + + private final String idProperty; + + private final String constraintName; + + private final Neo4jVectorFilterExpressionConverter filterExpressionConverter = new Neo4jVectorFilterExpressionConverter(); private final boolean initializeSchema; private final BatchingStrategy batchingStrategy; + @Deprecated(since = "1.0.0-M5", forRemoval = true) public Neo4jVectorStore(Driver driver, EmbeddingModel embeddingModel, Neo4jVectorStoreConfig config, boolean initializeSchema) { this(driver, embeddingModel, config, initializeSchema, ObservationRegistry.NOOP, null, new TokenCountBatchingStrategy()); } + @Deprecated(since = "1.0.0-M5", forRemoval = true) public Neo4jVectorStore(Driver driver, EmbeddingModel embeddingModel, Neo4jVectorStoreConfig config, boolean initializeSchema, ObservationRegistry observationRegistry, VectorStoreObservationConvention customObservationConvention, BatchingStrategy batchingStrategy) { - super(observationRegistry, customObservationConvention); - this.initializeSchema = initializeSchema; - Assert.notNull(driver, "Neo4j driver must not be null"); - Assert.notNull(embeddingModel, "Embedding model must not be null"); - this.driver = driver; - this.embeddingModel = embeddingModel; - this.config = config; - this.batchingStrategy = batchingStrategy; + this(builder().driver(driver) + .embeddingModel(embeddingModel) + .sessionConfig(config.sessionConfig) + .embeddingDimension(config.embeddingDimension) + .distanceType(config.distanceType) + .embeddingProperty(config.embeddingProperty) + .label(config.label) + .indexName(config.indexName) + .idProperty(config.idProperty) + .constraintName(config.constraintName) + .initializeSchema(initializeSchema) + .observationRegistry(observationRegistry) + .customObservationConvention(customObservationConvention) + .batchingStrategy(batchingStrategy)); + } + + protected Neo4jVectorStore(Neo4jBuilder builder) { + super(builder); + + Assert.notNull(builder.driver, "Neo4j driver must not be null"); + + this.driver = builder.driver; + this.sessionConfig = builder.sessionConfig; + this.embeddingDimension = builder.embeddingDimension; + this.distanceType = builder.distanceType; + this.embeddingProperty = SchemaNames.sanitize(builder.embeddingProperty).orElseThrow(); + this.label = SchemaNames.sanitize(builder.label).orElseThrow(); + this.indexNameNotSanitized = builder.indexName; + this.indexName = SchemaNames.sanitize(builder.indexName, true).orElseThrow(); + this.idProperty = SchemaNames.sanitize(builder.idProperty).orElseThrow(); + this.constraintName = SchemaNames.sanitize(builder.constraintName).orElseThrow(); + this.initializeSchema = builder.initializeSchema; + this.batchingStrategy = builder.batchingStrategy; } @Override @@ -127,17 +253,17 @@ public class Neo4jVectorStore extends AbstractObservationVectorStore implements u += row.properties WITH row, u CALL db.create.setNodeVectorProperty(u, $embeddingProperty, row[$embeddingProperty]) - """.formatted(this.config.label, this.config.idProperty); - session.executeWrite( - tx -> tx.run(statement, Map.of("rows", rows, "embeddingProperty", this.config.embeddingProperty)) - .consume()); + """.formatted(this.label, this.idProperty); + session + .executeWrite(tx -> tx.run(statement, Map.of("rows", rows, "embeddingProperty", this.embeddingProperty)) + .consume()); } } @Override public Optional doDelete(List idList) { - try (var session = this.driver.session(this.config.sessionConfig)) { + try (var session = this.driver.session(this.sessionConfig)) { // Those queries with internal, cypher based transaction management cannot be // run with executeWrite @@ -145,7 +271,7 @@ public class Neo4jVectorStore extends AbstractObservationVectorStore implements .run(""" MATCH (n:%s) WHERE n.%s IN $ids CALL { WITH n DETACH DELETE n } IN TRANSACTIONS OF $transactionSize ROWS - """.formatted(this.config.label, this.config.idProperty), + """.formatted(this.label, this.idProperty), Map.of("ids", idList, "transactionSize", DEFAULT_TRANSACTION_SIZE)) .consume(); return Optional.of(idList.size() == summary.counters().nodesDeleted()); @@ -159,7 +285,7 @@ public class Neo4jVectorStore extends AbstractObservationVectorStore implements "The similarity score is bounded between 0 and 1; least to most similar respectively."); var embedding = Values.value(this.embeddingModel.embed(request.getQuery())); - try (var session = this.driver.session(this.config.sessionConfig)) { + try (var session = this.driver.session(this.sessionConfig)) { StringBuilder condition = new StringBuilder("score >= $threshold"); if (request.hasFilterExpression()) { condition.append(" AND ") @@ -172,8 +298,9 @@ public class Neo4jVectorStore extends AbstractObservationVectorStore implements RETURN node, score""".formatted(condition); return session.executeRead(tx -> tx - .run(query, Map.of("indexName", this.config.indexNameNotSanitized, "numberOfNearestNeighbours", - request.getTopK(), "embeddingValue", embedding, "threshold", request.getSimilarityThreshold())) + .run(query, + Map.of("indexName", this.indexNameNotSanitized, "numberOfNearestNeighbours", request.getTopK(), + "embeddingValue", embedding, "threshold", request.getSimilarityThreshold())) .list(this::recordToDocument)); } } @@ -185,11 +312,11 @@ public class Neo4jVectorStore extends AbstractObservationVectorStore implements return; } - try (var session = this.driver.session(this.config.sessionConfig)) { + try (var session = this.driver.session(this.sessionConfig)) { session.executeWriteWithoutResult(tx -> { tx.run("CREATE CONSTRAINT %s IF NOT EXISTS FOR (n:%s) REQUIRE n.%s IS UNIQUE" - .formatted(this.config.constraintName, this.config.label, this.config.idProperty)).consume(); + .formatted(this.constraintName, this.label, this.idProperty)).consume(); var statement = """ CREATE VECTOR INDEX %s IF NOT EXISTS FOR (n:%s) ON (n.%s) @@ -197,8 +324,8 @@ public class Neo4jVectorStore extends AbstractObservationVectorStore implements `vector.dimensions`: %d, `vector.similarity_function`: '%s' }} - """.formatted(this.config.indexName, this.config.label, this.config.embeddingProperty, - this.config.embeddingDimension, this.config.distanceType.name); + """.formatted(this.indexName, this.label, this.embeddingProperty, this.embeddingDimension, + this.distanceType.name); tx.run(statement).consume(); }); @@ -219,7 +346,7 @@ public class Neo4jVectorStore extends AbstractObservationVectorStore implements document.getMetadata().forEach((k, v) -> properties.put("metadata." + k, Values.value(v))); row.put("properties", properties); - row.put(this.config.embeddingProperty, Values.value(embedding)); + row.put(this.embeddingProperty, Values.value(embedding)); return row; } @@ -235,7 +362,7 @@ public class Neo4jVectorStore extends AbstractObservationVectorStore implements }); return Document.builder() - .id(node.get(this.config.idProperty).asString()) + .id(node.get(this.idProperty).asString()) .text(node.get("text").asString()) .metadata(Map.copyOf(metaData)) .score((double) score) @@ -246,16 +373,16 @@ public class Neo4jVectorStore extends AbstractObservationVectorStore implements public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName) { return VectorStoreObservationContext.builder(VectorStoreProvider.NEO4J.value(), operationName) - .withCollectionName(this.config.indexName) + .withCollectionName(this.indexName) .withDimensions(this.embeddingModel.dimensions()) .withSimilarityMetric(getSimilarityMetric()); } private String getSimilarityMetric() { - if (!SIMILARITY_TYPE_MAPPING.containsKey(this.config.distanceType)) { - return this.config.distanceType.name(); + if (!SIMILARITY_TYPE_MAPPING.containsKey(this.distanceType)) { + return this.distanceType.name(); } - return SIMILARITY_TYPE_MAPPING.get(this.config.distanceType).value(); + return SIMILARITY_TYPE_MAPPING.get(this.distanceType).value(); } /** @@ -273,9 +400,181 @@ public class Neo4jVectorStore extends AbstractObservationVectorStore implements } + public static Neo4jBuilder builder() { + return new Neo4jBuilder(); + } + + public static class Neo4jBuilder extends AbstractVectorStoreBuilder { + + private Driver driver; + + private SessionConfig sessionConfig = SessionConfig.defaultConfig(); + + private int embeddingDimension = DEFAULT_EMBEDDING_DIMENSION; + + private Neo4jDistanceType distanceType = Neo4jDistanceType.COSINE; + + private String label = DEFAULT_LABEL; + + private String embeddingProperty = DEFAULT_EMBEDDING_PROPERTY; + + private String indexName = DEFAULT_INDEX_NAME; + + private String idProperty = DEFAULT_ID_PROPERTY; + + private String constraintName = DEFAULT_CONSTRAINT_NAME; + + private boolean initializeSchema = false; + + private BatchingStrategy batchingStrategy = new TokenCountBatchingStrategy(); + + public Neo4jBuilder driver(Driver driver) { + Assert.notNull(driver, "Neo4j driver must not be null"); + this.driver = driver; + return this; + } + + /** + * Sets the database name. When provided and not blank, creates a session config + * for that database. + * @param databaseName the database name to use + * @return the builder instance + */ + public Neo4jBuilder databaseName(String databaseName) { + if (StringUtils.hasText(databaseName)) { + this.sessionConfig = SessionConfig.forDatabase(databaseName); + } + return this; + } + + /** + * Sets the session configuration directly. + * @param sessionConfig the session configuration to use + * @return the builder instance + */ + public Neo4jBuilder sessionConfig(SessionConfig sessionConfig) { + this.sessionConfig = sessionConfig; + return this; + } + + /** + * Sets the embedding dimension. Must be positive. + * @param dimension the dimension of the embedding + * @return the builder instance + * @throws IllegalArgumentException if dimension is less than 1 + */ + public Neo4jBuilder embeddingDimension(int dimension) { + Assert.isTrue(dimension >= 1, "Dimension has to be positive"); + this.embeddingDimension = dimension; + return this; + } + + /** + * Sets the distance type for index storage and queries. + * @param distanceType the distance type to use + * @return the builder instance + * @throws IllegalArgumentException if distanceType is null + */ + public Neo4jBuilder distanceType(Neo4jDistanceType distanceType) { + Assert.notNull(distanceType, "Distance type may not be null"); + this.distanceType = distanceType; + return this; + } + + /** + * Sets the label for document nodes. + * @param label the label to use + * @return the builder instance + */ + public Neo4jBuilder label(String label) { + if (StringUtils.hasText(label)) { + this.label = label; + } + return this; + } + + /** + * Sets the property name for storing embeddings. + * @param embeddingProperty the property name to use + * @return the builder instance + */ + public Neo4jBuilder embeddingProperty(String embeddingProperty) { + if (StringUtils.hasText(embeddingProperty)) { + this.embeddingProperty = embeddingProperty; + } + return this; + } + + /** + * Sets the name of the vector index. + * @param indexName the index name to use + * @return the builder instance + */ + public Neo4jBuilder indexName(String indexName) { + if (StringUtils.hasText(indexName)) { + this.indexName = indexName; + } + return this; + } + + /** + * Sets the property name for document IDs. + * @param idProperty the property name to use + * @return the builder instance + */ + public Neo4jBuilder idProperty(String idProperty) { + if (StringUtils.hasText(idProperty)) { + this.idProperty = idProperty; + } + return this; + } + + /** + * Sets the name of the unique constraint. + * @param constraintName the constraint name to use + * @return the builder instance + */ + public Neo4jBuilder constraintName(String constraintName) { + if (StringUtils.hasText(constraintName)) { + this.constraintName = constraintName; + } + return this; + } + + /** + * Sets whether to initialize the schema. + * @param initializeSchema true to initialize schema, false otherwise + * @return the builder instance + */ + public Neo4jBuilder initializeSchema(boolean initializeSchema) { + this.initializeSchema = initializeSchema; + return this; + } + + /** + * Sets the batching strategy. + * @param batchingStrategy the strategy to use + * @return the builder instance + * @throws IllegalArgumentException if batchingStrategy is null + */ + public Neo4jBuilder batchingStrategy(BatchingStrategy batchingStrategy) { + Assert.notNull(batchingStrategy, "BatchingStrategy must not be null"); + this.batchingStrategy = batchingStrategy; + return this; + } + + @Override + public Neo4jVectorStore build() { + validate(); + return new Neo4jVectorStore(this); + } + + } + /** * Configuration for the Neo4j vector store. */ + @Deprecated(since = "1.0.0-M5", forRemoval = true) public static final class Neo4jVectorStoreConfig { private final SessionConfig sessionConfig; @@ -317,19 +616,20 @@ public class Neo4jVectorStore extends AbstractObservationVectorStore implements * Start building a new configuration. * @return The entry point for creating a new configuration. */ + @Deprecated(since = "1.0.0-M5", forRemoval = true) public static Builder builder() { - return new Builder(); } /** * {@return the default config} */ + @Deprecated(since = "1.0.0-M5", forRemoval = true) public static Neo4jVectorStoreConfig defaultConfig() { - return builder().build(); } + @Deprecated(since = "1.0.0-M5", forRemoval = true) public static final class Builder { private String databaseName; diff --git a/vector-stores/spring-ai-neo4j-store/src/main/java/org/springframework/ai/vectorstore/filter/Neo4jVectorFilterExpressionConverter.java b/vector-stores/spring-ai-neo4j-store/src/main/java/org/springframework/ai/vectorstore/neo4j/filter/Neo4jVectorFilterExpressionConverter.java similarity index 96% rename from vector-stores/spring-ai-neo4j-store/src/main/java/org/springframework/ai/vectorstore/filter/Neo4jVectorFilterExpressionConverter.java rename to vector-stores/spring-ai-neo4j-store/src/main/java/org/springframework/ai/vectorstore/neo4j/filter/Neo4jVectorFilterExpressionConverter.java index 747c0807a..6786e0e3b 100644 --- a/vector-stores/spring-ai-neo4j-store/src/main/java/org/springframework/ai/vectorstore/filter/Neo4jVectorFilterExpressionConverter.java +++ b/vector-stores/spring-ai-neo4j-store/src/main/java/org/springframework/ai/vectorstore/neo4j/filter/Neo4jVectorFilterExpressionConverter.java @@ -14,8 +14,9 @@ * limitations under the License. */ -package org.springframework.ai.vectorstore.filter; +package org.springframework.ai.vectorstore.neo4j.filter; +import org.springframework.ai.vectorstore.filter.Filter; import org.springframework.ai.vectorstore.filter.Filter.Expression; import org.springframework.ai.vectorstore.filter.Filter.Group; import org.springframework.ai.vectorstore.filter.Filter.Key; diff --git a/vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/Neo4jImage.java b/vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/neo4j/Neo4jImage.java similarity index 94% rename from vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/Neo4jImage.java rename to vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/neo4j/Neo4jImage.java index e0133dbf1..17ab2e168 100644 --- a/vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/Neo4jImage.java +++ b/vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/neo4j/Neo4jImage.java @@ -14,7 +14,7 @@ * limitations under the License. */ -package org.springframework.ai.vectorstore; +package org.springframework.ai.vectorstore.neo4j; import org.testcontainers.utility.DockerImageName; diff --git a/vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/Neo4jVectorStoreIT.java b/vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/neo4j/Neo4jVectorStoreIT.java similarity index 97% rename from vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/Neo4jVectorStoreIT.java rename to vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/neo4j/Neo4jVectorStoreIT.java index b5bc62ef9..401b8a559 100644 --- a/vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/Neo4jVectorStoreIT.java +++ b/vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/neo4j/Neo4jVectorStoreIT.java @@ -14,7 +14,7 @@ * limitations under the License. */ -package org.springframework.ai.vectorstore; +package org.springframework.ai.vectorstore.neo4j; import java.util.Collections; import java.util.List; @@ -37,6 +37,8 @@ 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.ai.vectorstore.VectorStore; import org.springframework.ai.vectorstore.filter.FilterExpressionTextParser; import org.springframework.boot.SpringBootConfiguration; import org.springframework.boot.autoconfigure.EnableAutoConfiguration; @@ -289,8 +291,11 @@ class Neo4jVectorStoreIT { @Bean public VectorStore vectorStore(Driver driver, EmbeddingModel embeddingModel) { - return new Neo4jVectorStore(driver, embeddingModel, Neo4jVectorStore.Neo4jVectorStoreConfig.defaultConfig(), - true); + return Neo4jVectorStore.builder() + .driver(driver) + .embeddingModel(embeddingModel) + .initializeSchema(true) + .build(); } @Bean diff --git a/vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/Neo4jVectorStoreObservationIT.java b/vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/neo4j/Neo4jVectorStoreObservationIT.java similarity index 94% rename from vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/Neo4jVectorStoreObservationIT.java rename to vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/neo4j/Neo4jVectorStoreObservationIT.java index 1c841b1c1..0239f5454 100644 --- a/vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/Neo4jVectorStoreObservationIT.java +++ b/vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/neo4j/Neo4jVectorStoreObservationIT.java @@ -14,7 +14,7 @@ * limitations under the License. */ -package org.springframework.ai.vectorstore; +package org.springframework.ai.vectorstore.neo4j; import java.io.IOException; import java.nio.charset.StandardCharsets; @@ -43,6 +43,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; @@ -176,8 +178,14 @@ public class Neo4jVectorStoreObservationIT { public VectorStore vectorStore(Driver driver, EmbeddingModel embeddingModel, ObservationRegistry observationRegistry) { - return new Neo4jVectorStore(driver, embeddingModel, Neo4jVectorStore.Neo4jVectorStoreConfig.defaultConfig(), - true, observationRegistry, null, new TokenCountBatchingStrategy()); + return Neo4jVectorStore.builder() + .driver(driver) + .embeddingModel(embeddingModel) + .initializeSchema(true) + .observationRegistry(observationRegistry) + .customObservationConvention(null) + .batchingStrategy(new TokenCountBatchingStrategy()) + .build(); } @Bean diff --git a/vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/filter/Neo4jVectorFilterExpressionConverterTests.java b/vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/neo4j/filter/Neo4jVectorFilterExpressionConverterTests.java similarity index 95% rename from vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/filter/Neo4jVectorFilterExpressionConverterTests.java rename to vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/neo4j/filter/Neo4jVectorFilterExpressionConverterTests.java index 4eeaa22f1..aeb6c01f4 100644 --- a/vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/filter/Neo4jVectorFilterExpressionConverterTests.java +++ b/vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/neo4j/filter/Neo4jVectorFilterExpressionConverterTests.java @@ -14,7 +14,7 @@ * limitations under the License. */ -package org.springframework.ai.vectorstore.filter; +package org.springframework.ai.vectorstore.neo4j.filter; import java.util.List; @@ -24,6 +24,8 @@ import org.springframework.ai.vectorstore.filter.Filter.Expression; import org.springframework.ai.vectorstore.filter.Filter.Group; import org.springframework.ai.vectorstore.filter.Filter.Key; import org.springframework.ai.vectorstore.filter.Filter.Value; +import org.springframework.ai.vectorstore.filter.FilterExpressionConverter; +import org.springframework.ai.vectorstore.filter.FilterExpressionTextParser; import static org.assertj.core.api.Assertions.assertThat; import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.AND; @@ -130,7 +132,7 @@ public class Neo4jVectorFilterExpressionConverterTests { @Test public void testComplexIdentifiers2() { - Filter.Expression expr = new FilterExpressionTextParser() + Expression expr = new FilterExpressionTextParser() .parse("author in ['john', 'jill'] && 'article_type' == 'blog'"); String vectorExpr = this.converter.convertExpression(expr); assertThat(vectorExpr)