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())); + } + + } + +}