Spring AI

The Spring AI project aims to be something like LangChain for the Java world.

Combining Spring AI with Spring Data Neo4j makes it possible to build on existing domain models, turn them into a graph, and enrich them with vector embeddings.

Installation

If you are using the Spring Initializr, you can add the Spring AI model of your choosing and Neo4j Vector Database as dependencies to your project, and it will pull in all the related pieces.

If you add the dependencies manually, Spring AI is not yet a general release library, so you will need to add the dependencies, as well as the snapshot and milestone repositories to your pom.xml file, as outlined in the Getting Started documentation.

The next piece you will need to add is the configuration for the Neo4j Vector Database, which also needs to be accompanied by custom driver and embedding client configuration. Probably the simplest approach is to create Spring Beans for each of these components in the main application class. This is mentioned in the documentation, but the full code is not provided. An example of how to configure the Neo4j Vector Database is shown below (note that you will need to alter the vectorStore bean to match your specific configuration):

@Bean
public Driver driver() {
  return GraphDatabase.driver(System.getenv("SPRING_NEO4J_URI"),
      AuthTokens.basic(System.getenv("SPRING_NEO4J_AUTHENTICATION_USERNAME"),
          System.getenv("SPRING_NEO4J_AUTHENTICATION_PASSWORD")));
}

@Bean
public EmbeddingClient embeddingClient() {
  return new OpenAiEmbeddingClient(new OpenAiApi(System.getenv("SPRING_AI_OPENAI_API_KEY")));
}

@Bean
public Neo4jVectorStore vectorStore(Driver driver, EmbeddingClient embeddingClient) {

  return new Neo4jVectorStore(driver, embeddingClient,
      Neo4jVectorStore.Neo4jVectorStoreConfig.builder()
          .withIndexName("form_10k_chunks")
          .withLabel("Chunk")
          .withEmbeddingProperty("textEmbedding")
          .build());
}

Now that we have the vector store configured, we can use Spring AI for retrieval augmented generation (RAG) via a three-step process.

Call the vector similarity search method to retrieve the most similar documents.
Use the similar documents as input to a retrieval query that pulls related entities from the graph.
Provide the similar documents (with their related graph entities) as input to a prompt that the LLM will use to generate a response.

The code snippet below demonstrates how to use Spring AI for RAG with Neo4j Vector Store:

@GetMapping("/chat")
String getGeneratedResponse(@RequestParam String question) {
    List<Document> results = vectorStore.similaritySearch(SearchRequest.query(question));

    System.out.println("Id list to graph: " + results.stream()
            .map(Document::getId)
            .collect(Collectors.toList()));

    List<Chunk> docList = repo.getRelatedEntitiesForSimilarChunks(results.stream()
            .map(Document::getId)
            .collect(Collectors.toList()));

    var template = new PromptTemplate("""
            You are a helpful question-answering agent. Your task is to analyze
            and synthesize information from the top result from a similarity search
            and relevant data from a graph database.
            Given the user's query: {question}, provide a meaningful and efficient answer based
            on the insights derived from the following data:
                            
            {graph_result}
            """,
            Map.of("question", question,
                    "graph_result", docList.stream().map(chunk -> chunk.toString()).collect(Collectors.joining("\n"))));
    System.out.println("----- PROMPT -----");
    System.out.println(template.render());

    return client.call(template.create().getContents());
}