23.02.2024

Spring AI with Neo4j Vector Index

Cypher statements based on results from the Neo4j Vector store

SpringAI

At the moment still in pre-release phase, there is a new Spring module called SpringAI. It offers support for different "AI-vendors" like OpenAI, Bedrock, Ollama, etc. for creating embeddings or calling the chat completion feature. (Yes, also image creation, but this is not part of this post.) On the other end there are multiple store implementations to persist the embedding in various stores (PgVector, Redis, Neo4j etc.).

The idea of those store modules is to have a unified idea of storing Documents representing some content along with metadata into the store and adding the embeddings. Also, the stores offer a database specific implementation of a similarity search. The result of the search is -described on a high abstraction level- the documents having an embedding "close" to the given term provided with the similarity search invocation.

Since the project is very actively developed, it is best to check out the current documentation for more information.

Neo4j Vector Store

In this post, we will have a closer look into the usage of the Neo4j Vector Store implementation created by Michael Simons and me. As mentioned in the beginning, each store module unifies the behaviour defined by the VectorStore interface, so does the Neo4j implementation.

The store itself is configurable for your needs. By default, it will store the Documents as nodes labelled Document and store the embeddings on the embedding property. For customisation, it is possible to define label, property name for the embedding, index name, embedding dimension, or distance type for the index. When doing the first steps with SpringAI this is not needed but might come in handy if you have an already existing dataset within the database and you don't want to be forced into the default naming schema.

After initializing the Neo4j Vector Store, it will ensure that a vector index for the embedding exists. At the moment of writing, Neo4j supports embedding dimensions up to 2048. As a consequence retrieving embeddings e.g. via llama2 does not work because they have an embedding dimension of 4096.

Adding an embedding to a node will set the embedding by calling db.create.setNodeVectorProperty(node, <embedding_property>, <embedding>). This way, it is guaranteed that the value will get indexed.

Retrieving the Documents via similarity search will invoke db.index.vector.queryNodes(<index_name>, <amount_of_nearest_neighbours>, <search_term_as_embedding>), and return the matching nodes alongside their ranking.

Example

This example will not only demonstrate how to use the functionality mentioned above, but also how to create some Documents, and use them after a successful similarity search with the chat completion to get the information, you are looking for. The application should be a simple Spring Boot web application that accepts a user prompt and responds with the equivalent Cypher statement. We won't strive for perfection here but give some idea how you could also create something similar for your purposes.

Setup

A Neo4j instance version 5.15 or newer is required to follow along. If you don't know how to get one, there are several options to get one running:

• Docker image
• Neo4j Desktop
• Neo4j Aura
• Neo4j Server binaries

We will create a Spring Boot application at start.spring.io with the -at the moment- current version 3.2.3 and select only the web dependency (pre-configured application). Because the SpringAI project is still in pre-release phase, it is only available in Spring's repositories. After downloading and extracting the application, the additional repositories needs to get declared in the pom.xml.

<repositories>  
    <repository>  
        <id>spring-milestones</id>  
        <name>Spring Milestones</name>  
        <url>https://repo.spring.io/milestone</url>  
        <snapshots>  
            <enabled>false</enabled>  
        </snapshots>  
    </repository>  
    <repository>  
        <id>spring-snapshots</id>  
        <name>Spring Snapshots</name>  
        <url>https://repo.spring.io/snapshot</url>  
        <releases>  
            <enabled>false</enabled>  
        </releases>  
    </repository>  
</repositories>

There is a BOM (bill of materials) available to define the overall version of SpringAI modules.

<dependencyManagement>  
    <dependencies>  
        <dependency>  
            <groupId>org.springframework.ai</groupId>  
            <artifactId>spring-ai-bom</artifactId>  
            <version>${spring-ai.version}</version>  
            <type>pom</type>  
            <scope>import</scope>  
        </dependency>  
    </dependencies>  
</dependencyManagement>

In this example, we will use the version 0.8.0-SNAPSHOT.

Now we can add the needed modules of SpringAI for this example:

<dependency>  
    <groupId>org.springframework.ai</groupId>  
    <artifactId>spring-ai-neo4j-store-spring-boot-starter</artifactId>  
</dependency>  
<dependency>  
    <groupId>org.springframework.ai</groupId>  
    <artifactId>spring-ai-openai-spring-boot-starter</artifactId>  
</dependency>  
<dependency>  
    <groupId>org.springframework.ai</groupId>  
    <artifactId>spring-ai-pdf-document-reader</artifactId>  
</dependency>

As you can see, the store and AI vendor specific dependencies are offered as a Spring Boot starter. This means that we don't need to do any manual bean configuration but just provide the needed properties. If we are ok with the default settings of the store and OpenAI module (we are), there are just some mandatory properties left, that should be placed in the application.properties file:

spring.ai.openai.api-key=${OPENAI_API_KEY}  
  
spring.neo4j.uri=<neo4j_URI> 
spring.neo4j.authentication.username=<neo4j_username>  
spring.neo4j.authentication.password=<neo4j_password>

The OpenAI API token and the coordinates of the Neo4j database to store and query the Documents are required.

The third dependency is a PDF to Document converter module that is part of SpringAI to simplify the tokenisation of PDF documents into Document objects.

The application infrastructure is set up and we can start coding.

Application

The application should not only return Cypher statements for textual input but also trigger the creation of the documents from a given PDF file.

Controller

Let's start with the controller class CypherPromptController that we declare as a @Controller stereotype and add the @ReponseBody annotation to be able to return basic strings. The actual functionality of invoking the vector store and AI services will be done in a service class called CypherPromptService. This needs to get autowired into our controller. There are two endpoint that needs to be provided, one for the Document and embedding creation and one for the actual prompt to Cypher conversion.

@Controller  
@ResponseBody  
public class CypherPromptController {
  private final CypherPromptService promptService;  
  
  public CypherPromptController(CypherPromptService promptService) {  
    this.promptService = promptService;  
  }
  @GetMapping("/")  
  public String convertToCypher(@RequestParam("m") String message) {  
    return this.promptService.convertToCypher(message);  
  }  
  
  @GetMapping("/create")  
  public String createDocuments(@RequestParam("filePath") String filePath) {  
    this.promptService.createEmbeddings(filePath);  
    return "done";  
  }
}

In a more strict REST sense the GET mapping for the document creation is up for a debate, but in the context of this conceptual code, I can live with it ;) Also, I assume that the application gets deployed somewhere (e.g. locally) to have access to the PDF file for parsing.

Service

The service provides the actual business logic for fulfilling our requirements. It has dependencies to the (Neo4j) vector store and the (OpenAI) chat client implementation.

For processing the PDF file that can be found at the given path, it uses the PagePdfDocumentReader to represent each page of a PDF as a single Document. When adding them to the vector store by invoking vectorStore.add(parsedDocuments) it will automatically get enhanced with the embeddings generated by OpenAI before getting persisted to Neo4j.

Prompting for Cypher statements is done with a Prompt containing two parts:

1. A system message that defines the basis for the communication with OpenAI Chat
2. The actual user prompt containing the message For the first one, a prepared message with a placeholder {documents} will be used. The replacement for this will be the content of the top 4 results of a vectorStore.similaritySearch(message).

Telling the chat how to behave when a message comes in and restricting it strictly to use only the provided documents' content should only return Cypher. "Correct Cypher?" you may ask. "What it thinks correct Cypher is!" I would answer.

The whole service class looks like this:

@Service  
public class CypherPromptService {  
  
  private final VectorStore vectorStore;  
  
  private final ChatClient chatClient;  
  
  private static final String SYSTEM_PROMPT_TEMPLATE = """  
    You are an assistant that gives out Cypher code snippets.
    Use the information from the DOCUMENTS section only to provide accurate answers.
    Return just the code snippet without formatting. No descriptive text.
    Don't use any learned knowledge that is not within the DOCUMENTS section.

    DOCUMENTS:  
    {documents}""";  
  
  public CypherPromptService(VectorStore vectorStore, ChatClient chatClient) {  
   this.vectorStore = vectorStore;  
   this.chatClient = chatClient;  
  }  
    
  public String convertToCypher(String message) {  
  
   var result = vectorStore.similaritySearch(message);  
   String documents = result.stream().map(Document::getContent).collect(Collectors.joining(System.lineSeparator()));  
   var systemMessage = new SystemPromptTemplate(SYSTEM_PROMPT_TEMPLATE)  
     .createMessage(Map.of("documents", documents));  
  
   var userMessage = new UserMessage(message);  
  
   var prompt = new Prompt(List.of(systemMessage, userMessage));  
   ChatResponse response = chatClient.call(prompt);  
  
   return response.getResult().getOutput().getContent();  
  }  
  
  public void createEmbeddings(String filePath) {  
   List<Document> parsedDocuments = new PagePdfDocumentReader(filePath).get();  
   parsedDocuments.forEach(document -> document.getMetadata().clear());  
  
   vectorStore.add(parsedDocuments);  
  }  
}

Run it

We need to get the current Cypher 5 manual as a PDF. This can be found herein the docs archive. For not having to write too much, I simply put it in the main/resources folder of my application before starting the application.

The application can be started within your IDE or by calling ./mvnw spring-boot:run on the command line.

Creation of Documents with embeddings

In this curl call we need to provide the file path that the application should be looking for the PDF to create the Documents. If you followed along until here, this should be just the file name because the file was placed in the resources folder.

curl 'http://localhost:8080/create?filePath=cypher.pdf'

Don't get confused by some exceptions regarding missing fonts or even the java.lang.IllegalArgumentException: Comparison method violates its general contract! in the application logs. It will take a while until all pages of the PDF got processed and saved with their embeddings in the database.

After this call returns, the dataset has been successfully created and the actual text to Cypher conversion can be invoked.

Conversion

Now the chat-based conversion can be invoked. Let's start with something simple to see if the interaction works and we get a correct response.

curl 'http://localhost:8080/?m=Create%20a%20node'

(For convenience I am using curl but it requires the conversion of special characters, like white spaces, quotes, etc. Use a browser or something like postman for your own sanity :) )

And the response is:

CREATE (n)

This looks promising. For the next example we will use the prompt: "Create a node labelled User with a property called 'nickname' having a value 'meistermeier'".

curl 'http://localhost:8080/?m=Create%20a%20node%20labelled%20User%20with%20a%20property%20called%20%27nickname%27%20having%20a%20value%20%27meistermeier%27'

returns

CREATE (:User {nickname: 'meistermeier'})

The last bit for this example would be to create a query to get to know how many User nodes exist in the database. "How many nodes labelled 'User' exists in the database?"

curl 'http://localhost:8080/?m=%22How%20many%20nodes%20labelled%20%27User%27%20exists%20in%20the%20database?%22'

will yield

MATCH (n:User) RETURN count(n)

Conclusion

The results for expected to be easily to produce are really interesting to observe, and in this context surprised me positively. "Create a node Create a node labelled User with a property called 'nickname' having a value 'meistermeier' linked to another new node labelled User with a property called 'nickname' having a value 'rotnroll666'"

curl 'http://localhost:8080/?m=%22Create%20a%20node%20Create%20a%20node%20labelled%20User%20with%20a%20property%20called%20%27nickname%27%20having%20a%20value%20%27meistermeier%27%20linked%20to%20another%20new%20node%20labelled%20User%20with%20a%20property%20called%20%27nickname%27%20having%20a%20value%20%27rotnroll666%27%22'

Also work partially by taking some relation type from the documentation.

CREATE (a:User {nickname: 'meistermeier'})-[:FRIEND]-(b:User {nickname: 'rotnroll666'})

But this is something we have to keep in mind when using a system with a limited understanding (the PDF) and giving it a incomplete (missing relationship type) information. After presenting the prompts that will create the Cypher statements it is now up for the reader to decide if this is a valuable way to create Cypher -or even execute it blindly with the driver in the database- or if it would be better in the long run to learn and understand Cypher.

If you want to give some feedback or leave some thoughts, feel free to ping/mention me on mastodon. Any questions regarding the code? You can find the sources of this example on GitHub meistermeier/spring-ai-neo4j-example.