From a6fba7a3b6306b64c48cc881ecdc89096d3dd98f Mon Sep 17 00:00:00 2001 From: Gerrit Meier Date: Thu, 21 Sep 2023 13:15:16 +0200 Subject: [PATCH] GH-26 - Add support for Neo4j vector index. This commit brings support for Neo4j graph database in general, and uses the vector index functionality available since version 5.11. Aligned with the existing PgVector store and its tests. The module creates indexes, if needed, for the vector search and the identifier of the document object. Add neo4j to vectordb docs page Co-authored-by: Michael Simons --- pom.xml | 1 + .../modules/ROOT/pages/api/vectordbs.adoc | 5 +- vector-stores/spring-ai-neo4j-store/pom.xml | 78 ++++ .../ai/vectorstore/Neo4jVectorStore.java | 358 ++++++++++++++++++ .../ai/vectorstore/Neo4jVectorStoreIT.java | 194 ++++++++++ 5 files changed, 634 insertions(+), 2 deletions(-) create mode 100644 vector-stores/spring-ai-neo4j-store/pom.xml create mode 100644 vector-stores/spring-ai-neo4j-store/src/main/java/org/springframework/ai/vectorstore/Neo4jVectorStore.java create mode 100644 vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/Neo4jVectorStoreIT.java diff --git a/pom.xml b/pom.xml index 3cef21997..76eaa8555 100644 --- a/pom.xml +++ b/pom.xml @@ -21,6 +21,7 @@ spring-ai-docs vector-stores/spring-ai-pgvector-store vector-stores/spring-ai-milvus-store + vector-stores/spring-ai-neo4j-store diff --git a/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs.adoc b/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs.adoc index 5290956d3..b6dad236a 100644 --- a/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs.adoc +++ b/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs.adoc @@ -53,8 +53,9 @@ The `VectorStore` implementations supported by Spring AI are: * InMemoryVectorStore * SimplePersistentVectorStore -* PgVector - A Vector Store build on https://github.com/pgvector/pgvector[PostgreSQL/PGVector]. -* Milvus - A Vector Store build on https://milvus.io/[Milvus] +* PgVector - The Vector Store https://github.com/pgvector/pgvector[PostgreSQL/PGVector]. +* Milvus - The Vector Store https://milvus.io/[Milvus] +* Neo4j - The Vector Store https://neo4j.com/[Neo4j] More are implementations are coming, with Pinecone being the next implementation. diff --git a/vector-stores/spring-ai-neo4j-store/pom.xml b/vector-stores/spring-ai-neo4j-store/pom.xml new file mode 100644 index 000000000..1515c505a --- /dev/null +++ b/vector-stores/spring-ai-neo4j-store/pom.xml @@ -0,0 +1,78 @@ + + + 4.0.0 + + org.springframework.experimental.ai + spring-ai + 0.2.0-SNAPSHOT + ../../pom.xml + + spring-ai-neo4j-store + jar + Spring AI Vector Store - neo4j + Spring AI Neo4j Vector Store + https://github.com/spring-projects-experimental/spring-ai + + + https://github.com/spring-projects-experimental/spring-ai + git://github.com/spring-projects-experimental/spring-ai.git + git@github.com:spring-projects-experimental/spring-ai.git + + + + 0.2.0-SNAPSHOT + + 1.19.0 + + + + + org.springframework.experimental.ai + spring-ai-core + ${spring-ai.version} + + + + org.neo4j.driver + neo4j-java-driver + + + + org.neo4j + neo4j-cypher-dsl-schema-name-support + 2023.7.0 + + + + + org.springframework.experimental.ai + spring-ai-openai-spring-boot-starter + ${spring-ai.version} + test + + + + org.springframework.boot + spring-boot-starter-test + test + + + + org.testcontainers + neo4j + ${testcontainers.version} + test + + + + org.testcontainers + junit-jupiter + ${testcontainers.version} + test + + + + + + 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/Neo4jVectorStore.java new file mode 100644 index 000000000..b82359d5a --- /dev/null +++ b/vector-stores/spring-ai-neo4j-store/src/main/java/org/springframework/ai/vectorstore/Neo4jVectorStore.java @@ -0,0 +1,358 @@ +/* + * Copyright 2023-2023 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.ai.vectorstore; + +import org.neo4j.cypherdsl.support.schema_name.SchemaNames; +import org.neo4j.driver.Driver; +import org.neo4j.driver.SessionConfig; +import org.neo4j.driver.Values; +import org.springframework.ai.document.Document; +import org.springframework.ai.embedding.EmbeddingClient; +import org.springframework.beans.factory.InitializingBean; +import org.springframework.util.Assert; + +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.function.Predicate; + +/** + * @author Gerrit Meier + * @author Michael Simons + */ +public class Neo4jVectorStore implements VectorStore, InitializingBean { + + /** + * An enum to configure the distance function used in the Neo4j vector index. + */ + public enum Neo4jDistanceType { + + COSINE("cosine"), EUCLIDEAN("euclidean"); + + public final String name; + + Neo4jDistanceType(String name) { + this.name = name; + } + + } + + /** + * Configuration for the Neo4j vector store. + */ + public static final class Neo4jVectorStoreConfig { + + private final SessionConfig sessionConfig; + + private final int embeddingDimension; + + private final Neo4jDistanceType distanceType; + + private final String label; + + private final String embeddingProperty; + + private final String quotedLabel; + + /** + * Start building a new configuration. + * @return The entry point for creating a new configuration. + */ + public static Builder builder() { + + return new Builder(); + } + + /** + * {@return the default config} + */ + public static Neo4jVectorStoreConfig defaultConfig() { + + return builder().build(); + } + + private Neo4jVectorStoreConfig(Builder builder) { + + this.sessionConfig = Optional.ofNullable(builder.databaseName) + .filter(Predicate.not(String::isBlank)) + .map(SessionConfig::forDatabase) + .orElseGet(SessionConfig::defaultConfig); + this.embeddingDimension = builder.embeddingDimension; + this.distanceType = builder.distanceType; + this.label = builder.label; + this.embeddingProperty = builder.embeddingProperty; + this.quotedLabel = SchemaNames.sanitize(this.label).orElseThrow(); + } + + public static class Builder { + + private String databaseName; + + private int embeddingDimension = DEFAULT_EMBEDDING_DIMENSION; + + private Neo4jDistanceType distanceType = Neo4jDistanceType.COSINE; + + private String label = DEFAULT_LABEL; + + private String embeddingProperty = DEFAULT_EMBEDDING_PROPERTY; + + private Builder() { + } + + /** + * Configures the Neo4j database name to use. Leave {@literal null} or blank + * to use the default database. + * @param databaseName the database name to use + * @return this builder + */ + public Builder withDatabaseName(String databaseName) { + this.databaseName = databaseName; + return this; + } + + /** + * Configures the size of the embedding. Defaults to {@literal 1536}, inline + * with OpenAIs embeddings. + * @param newEmbeddingDimension The dimension of the embedding + * @return this builder + */ + public Builder withEmbeddingDimension(int newEmbeddingDimension) { + + Assert.isTrue(newEmbeddingDimension >= 1 && newEmbeddingDimension <= 2048, + "Dimension has to be withing the boundaries 1 and 2048 (inclusively)"); + + this.embeddingDimension = newEmbeddingDimension; + return this; + } + + /** + * Configures the distance type to store in the index and to use in queries. + * @param newDistanceType The distance type, must not be {@literal null} + * @return this builder + */ + public Builder withDistanceType(Neo4jDistanceType newDistanceType) { + + Assert.notNull(newDistanceType, "Distance type may not be null"); + + this.distanceType = newDistanceType; + return this; + } + + /** + * Configures the node label to use for storing documents. Defaults to + * {@literal Document}. + * @param newLabel The label used on the nodes representing the document + * @return this builder + */ + public Builder withLabel(String newLabel) { + + Assert.hasText(newLabel, "Node label may not be null or blank"); + + this.label = newLabel; + return this; + } + + /** + * Configures the property of the node to use for storing embedding. Defaults + * to {@literal embedding}. + * @param newEmbeddingProperty The property of the nodes for storing the + * embedding + * @return this builder + */ + public Builder withEmbeddingProperty(String newEmbeddingProperty) { + + Assert.hasText(newEmbeddingProperty, "Embedding property may not be null or blank"); + + this.embeddingProperty = newEmbeddingProperty; + return this; + } + + /** + * {@return the immutable configuration} + */ + public Neo4jVectorStoreConfig build() { + + return new Neo4jVectorStoreConfig(this); + } + + } + + } + + private static final int DEFAULT_EMBEDDING_DIMENSION = 1536; + + private static final String DEFAULT_LABEL = "Document"; + + private static final String INDEX_NAME = "spring-ai-document-index"; + + private static final String DEFAULT_EMBEDDING_PROPERTY = "embedding"; + + private final Driver driver; + + private final EmbeddingClient embeddingClient; + + private final Neo4jVectorStoreConfig config; + + public Neo4jVectorStore(Driver driver, EmbeddingClient embeddingClient, Neo4jVectorStoreConfig config) { + + Assert.notNull(driver, "Neo4j driver must not be null"); + Assert.notNull(embeddingClient, "Embedding client must not be null"); + + this.driver = driver; + this.embeddingClient = embeddingClient; + + this.config = config; + } + + @Override + public void add(List documents) { + + var rows = documents.stream().map(this::documentToRecord).toList(); + + try (var session = this.driver.session()) { + var statement = """ + UNWIND $rows AS row + MERGE (u:%s {id: row.id}) + ON CREATE + SET u += row.properties + ON MATCH + SET u = {} + SET u.id = row.id, + u += row.properties + WITH row, u + CALL db.create.setVectorProperty(u, $embeddingProperty, row.embedding) + YIELD node + RETURN count(node) + """.formatted(this.config.quotedLabel); + session.run(statement, Map.of("rows", rows, "embeddingProperty", this.config.embeddingProperty)).consume(); + } + } + + @Override + public Optional delete(List idList) { + + try (var session = this.driver.session(this.config.sessionConfig)) { + + var summary = session.run(""" + MATCH (n:%s) WHERE n.id IN $ids + CALL { WITH n DETACH DELETE n } IN TRANSACTIONS OF $transactionSize ROWS + """.formatted(this.config.quotedLabel), Map.of("ids", idList, "transactionSize", 10_000)) + .consume(); + return Optional.of(idList.size() == summary.counters().nodesDeleted()); + } + } + + @Override + public List similaritySearch(String query) { + return this.similaritySearch(query, 5); + } + + @Override + public List similaritySearch(String query, int k) { + return this.similaritySearch(query, k, 0); + } + + @Override + public List similaritySearch(String query, int k, double threshold) { + + Assert.isTrue(k > 0, "The number of documents to returned must be greater than zero"); + Assert.isTrue(threshold >= 0 && threshold <= 1, + "The similarity score is bounded between 0 and 1; least to most similar respectively."); + + var embedding = Values.value(toFloatArray(this.embeddingClient.embed(query))); + try (var session = this.driver.session(this.config.sessionConfig)) { + return session + .run(""" + CALL db.index.vector.queryNodes($indexName, $numberOfNearestNeighbours, $embeddingValue) + YIELD node, score + WHERE score >= $threshold + RETURN node, score + """, Map.of("indexName", INDEX_NAME, "numberOfNearestNeighbours", k, "embeddingValue", + embedding, "threshold", threshold)) + .list(Neo4jVectorStore::recordToDocument); + } + } + + @Override + public void afterPropertiesSet() { + + try (var session = this.driver.session(this.config.sessionConfig)) { + + session + .run("CREATE CONSTRAINT %s_unique_idx IF NOT EXISTS FOR (n:%s) REQUIRE n.id IS UNIQUE".formatted( + SchemaNames.sanitize(this.config.label + "_unique_idx").orElseThrow(), this.config.quotedLabel)) + .consume(); + + var vectorIndexExists = session + .run("SHOW INDEXES YIELD name WHERE name = $name RETURN count(*) > 0", Map.of("name", INDEX_NAME)) + .single() + .get(0) + .asBoolean(); + + if (!vectorIndexExists) { + var statement = "CALL db.index.vector.createNodeIndex($indexName, $label, $embeddingProperty, $embeddingDimension, $distanceType)"; + session.run(statement, + Map.of("indexName", INDEX_NAME, "label", this.config.label, "embeddingProperty", + this.config.embeddingProperty, "embeddingDimension", this.config.embeddingDimension, + "distanceType", this.config.distanceType.name)) + .consume(); + session.run("CALL db.awaitIndexes()").consume(); + } + } + } + + private Map documentToRecord(Document document) { + var embedding = this.embeddingClient.embed(document); + document.setEmbedding(embedding); + + var row = new HashMap(); + + row.put("id", document.getId()); + + var properties = new HashMap(); + properties.put("text", document.getText()); + + document.getMetadata().forEach((k, v) -> properties.put("metadata." + k, Values.value(v))); + row.put("properties", properties); + + row.put(DEFAULT_EMBEDDING_PROPERTY, Values.value(toFloatArray(embedding))); + return row; + } + + private static float[] toFloatArray(List embeddingDouble) { + float[] embeddingFloat = new float[embeddingDouble.size()]; + int i = 0; + for (Double d : embeddingDouble) { + embeddingFloat[i++] = d.floatValue(); + } + return embeddingFloat; + } + + private static Document recordToDocument(org.neo4j.driver.Record neoRecord) { + var node = neoRecord.get("node").asNode(); + var metaData = new HashMap(); + node.keys().forEach(key -> { + if (key.startsWith("metadata.")) { + metaData.put(key.substring(key.indexOf(".") + 1), node.get(key).asObject()); + } + }); + + return new Document(node.get("id").asString(), node.get("text").asString(), Map.copyOf(metaData)); + } + +} 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/Neo4jVectorStoreIT.java new file mode 100644 index 000000000..adecbe7fe --- /dev/null +++ b/vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/Neo4jVectorStoreIT.java @@ -0,0 +1,194 @@ +package org.springframework.ai.vectorstore; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.neo4j.driver.AuthTokens; +import org.neo4j.driver.Driver; +import org.neo4j.driver.GraphDatabase; +import org.springframework.ai.autoconfigure.openai.OpenAiAutoConfiguration; +import org.springframework.ai.document.Document; +import org.springframework.ai.embedding.EmbeddingClient; +import org.springframework.boot.SpringBootConfiguration; +import org.springframework.boot.autoconfigure.AutoConfigurations; +import org.springframework.boot.autoconfigure.EnableAutoConfiguration; +import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration; +import org.springframework.boot.test.context.runner.ApplicationContextRunner; +import org.springframework.context.annotation.Bean; +import org.testcontainers.containers.Neo4jContainer; +import org.testcontainers.junit.jupiter.Container; +import org.testcontainers.junit.jupiter.Testcontainers; +import org.testcontainers.utility.DockerImageName; + +import java.util.Collections; +import java.util.List; +import java.util.UUID; +import java.util.stream.Collectors; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * @author Gerrit Meier + * @author Michael Simons + */ +@Testcontainers +class Neo4jVectorStoreIT { + + // Neo4j 5.12 has a bug wrt checking limits, so either 5.11 or anything higher than + // 5.12 works + @Container + static Neo4jContainer neo4jContainer = new Neo4jContainer<>(DockerImageName.parse("neo4j:5.11")) + .withRandomPassword(); + + List documents = List.of( + new Document("Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!!", + Collections.singletonMap("meta1", "meta1")), + new Document("Hello World Hello World Hello World Hello World Hello World Hello World Hello World"), + new Document( + "Great Depression Great Depression Great Depression Great Depression Great Depression Great Depression", + Collections.singletonMap("meta2", "meta2"))); + + private final ApplicationContextRunner contextRunner = new ApplicationContextRunner() + .withUserConfiguration(TestApplication.class) + .withPropertyValues("spring.ai.openai.apiKey=" + System.getenv("OPENAI_API_KEY")); + + @BeforeEach + void cleanDatabase() { + this.contextRunner + .run(context -> context.getBean(Driver.class).executableQuery("MATCH (n) DETACH DELETE n").execute()); + } + + @Test + void addAndSearchTest() { + this.contextRunner.withConfiguration(AutoConfigurations.of(OpenAiAutoConfiguration.class)).run(context -> { + + VectorStore vectorStore = context.getBean(VectorStore.class); + + vectorStore.add(this.documents); + + List results = vectorStore.similaritySearch("Great", 1); + + assertThat(results).hasSize(1); + Document resultDoc = results.get(0); + assertThat(resultDoc.getId()).isEqualTo(this.documents.get(2).getId()); + assertThat(resultDoc.getText()).isEqualTo( + "Great Depression Great Depression Great Depression Great Depression Great Depression Great Depression"); + assertThat(resultDoc.getMetadata()).isEqualTo(Collections.singletonMap("meta2", "meta2")); + + // Remove all documents from the store + vectorStore.delete(this.documents.stream().map(Document::getId).collect(Collectors.toList())); + + List results2 = vectorStore.similaritySearch("Great", 1); + assertThat(results2).isEmpty(); + }); + } + + @Test + void documentUpdateTest() { + + this.contextRunner.withConfiguration(AutoConfigurations.of(OpenAiAutoConfiguration.class)).run(context -> { + + VectorStore vectorStore = context.getBean(VectorStore.class); + + Document document = new Document(UUID.randomUUID().toString(), "Spring AI rocks!!", + Collections.singletonMap("meta1", "meta1")); + + vectorStore.add(List.of(document)); + + List results = vectorStore.similaritySearch("Spring", 5); + + assertThat(results).hasSize(1); + Document resultDoc = results.get(0); + assertThat(resultDoc.getId()).isEqualTo(document.getId()); + assertThat(resultDoc.getText()).isEqualTo("Spring AI rocks!!"); + assertThat(resultDoc.getMetadata()).isEqualTo(Collections.singletonMap("meta1", "meta1")); + + Document sameIdDocument = new Document(document.getId(), + "The World is Big and Salvation Lurks Around the Corner", + Collections.singletonMap("meta2", "meta2")); + + vectorStore.add(List.of(sameIdDocument)); + + results = vectorStore.similaritySearch("FooBar", 5); + + assertThat(results).hasSize(1); + resultDoc = results.get(0); + assertThat(resultDoc.getId()).isEqualTo(document.getId()); + assertThat(resultDoc.getText()).isEqualTo("The World is Big and Salvation Lurks Around the Corner"); + assertThat(resultDoc.getMetadata()).isEqualTo(Collections.singletonMap("meta2", "meta2")); + + }); + } + + @Test + void searchThresholdTest() { + + this.contextRunner.withConfiguration(AutoConfigurations.of(OpenAiAutoConfiguration.class)).run(context -> { + + VectorStore vectorStore = context.getBean(VectorStore.class); + + vectorStore.add(this.documents); + + assertThat(vectorStore.similaritySearch("Great", 5, 0)).hasSize(3); + + List results = vectorStore.similaritySearch("Great", 5, 0.89); + + assertThat(results).hasSize(1); + Document resultDoc = results.get(0); + assertThat(resultDoc.getId()).isEqualTo(this.documents.get(2).getId()); + assertThat(resultDoc.getText()).isEqualTo( + "Great Depression Great Depression Great Depression Great Depression Great Depression Great Depression"); + assertThat(resultDoc.getMetadata()).isEqualTo(Collections.singletonMap("meta2", "meta2")); + + }); + } + + @Test + void ensureVectorIndexGetsCreated() { + this.contextRunner.run(context -> { + assertThat(context.getBean(Driver.class) + .executableQuery( + "SHOW indexes yield name, type WHERE name = 'spring-ai-document-index' AND type = 'VECTOR' return count(*) > 0") + .execute() + .records() + .get(0) // get first record + .get(0) + .asBoolean()) // get returned result + .isTrue(); + }); + } + + @Test + void ensureIdIndexGetsCreated() { + this.contextRunner.run(context -> { + assertThat(context.getBean(Driver.class) + .executableQuery( + "SHOW indexes yield labelsOrTypes, properties, type WHERE any(x in labelsOrTypes where x = 'Document') AND any(x in properties where x = 'id') AND type = 'RANGE' return count(*) > 0") + .execute() + .records() + .get(0) // get first record + .get(0) + .asBoolean()) // get returned result + .isTrue(); + }); + } + + @SpringBootConfiguration + @EnableAutoConfiguration(exclude = { DataSourceAutoConfiguration.class }) + public static class TestApplication { + + @Bean + public VectorStore vectorStore(Driver driver, EmbeddingClient embeddingClient) { + + return new Neo4jVectorStore(driver, embeddingClient, + Neo4jVectorStore.Neo4jVectorStoreConfig.defaultConfig()); + } + + @Bean + public Driver driver() { + return GraphDatabase.driver(neo4jContainer.getBoltUrl(), + AuthTokens.basic("neo4j", neo4jContainer.getAdminPassword())); + } + + } + +}