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 <michael@simons.ac>
This commit is contained in:
Gerrit Meier
2023-09-21 13:15:16 +02:00
committed by Mark Pollack
parent 92773f17c0
commit a6fba7a3b6
5 changed files with 634 additions and 2 deletions

View File

@@ -21,6 +21,7 @@
<module>spring-ai-docs</module>
<module>vector-stores/spring-ai-pgvector-store</module>
<module>vector-stores/spring-ai-milvus-store</module>
<module>vector-stores/spring-ai-neo4j-store</module>
</modules>
<organization>

View File

@@ -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.

View File

@@ -0,0 +1,78 @@
<?xml version="1.0" encoding="UTF-8"?>
<project xmlns="http://maven.apache.org/POM/4.0.0"
xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/maven-v4_0_0.xsd">
<modelVersion>4.0.0</modelVersion>
<parent>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai</artifactId>
<version>0.2.0-SNAPSHOT</version>
<relativePath>../../pom.xml</relativePath>
</parent>
<artifactId>spring-ai-neo4j-store</artifactId>
<packaging>jar</packaging>
<name>Spring AI Vector Store - neo4j</name>
<description>Spring AI Neo4j Vector Store</description>
<url>https://github.com/spring-projects-experimental/spring-ai</url>
<scm>
<url>https://github.com/spring-projects-experimental/spring-ai</url>
<connection>git://github.com/spring-projects-experimental/spring-ai.git</connection>
<developerConnection>git@github.com:spring-projects-experimental/spring-ai.git</developerConnection>
</scm>
<properties>
<spring-ai.version>0.2.0-SNAPSHOT</spring-ai.version>
<!-- testing -->
<testcontainers.version>1.19.0</testcontainers.version>
</properties>
<dependencies>
<dependency>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai-core</artifactId>
<version>${spring-ai.version}</version>
</dependency>
<dependency>
<groupId>org.neo4j.driver</groupId>
<artifactId>neo4j-java-driver</artifactId>
</dependency>
<dependency>
<groupId>org.neo4j</groupId>
<artifactId>neo4j-cypher-dsl-schema-name-support</artifactId>
<version>2023.7.0</version>
</dependency>
<!-- TESTING -->
<dependency>
<groupId>org.springframework.experimental.ai</groupId>
<artifactId>spring-ai-openai-spring-boot-starter</artifactId>
<version>${spring-ai.version}</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.springframework.boot</groupId>
<artifactId>spring-boot-starter-test</artifactId>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.testcontainers</groupId>
<artifactId>neo4j</artifactId>
<version>${testcontainers.version}</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.testcontainers</groupId>
<artifactId>junit-jupiter</artifactId>
<version>${testcontainers.version}</version>
<scope>test</scope>
</dependency>
</dependencies>
</project>

View File

@@ -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<Document> 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<Boolean> delete(List<String> 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<Document> similaritySearch(String query) {
return this.similaritySearch(query, 5);
}
@Override
public List<Document> similaritySearch(String query, int k) {
return this.similaritySearch(query, k, 0);
}
@Override
public List<Document> 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<String, Object> documentToRecord(Document document) {
var embedding = this.embeddingClient.embed(document);
document.setEmbedding(embedding);
var row = new HashMap<String, Object>();
row.put("id", document.getId());
var properties = new HashMap<String, Object>();
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<Double> 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<String, Object>();
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));
}
}

View File

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