GenAI Starter Kit: Everything You Need to Build an Application with Spring AI in Java

- 8 minutes read - 1616 words
Photo credit

There are so many options when it comes to languages, frameworks, and tools for building generative AI (GenAI) applications. When you are just getting started, these decisions and figuring out how to integrate everything can be overwhelming.

My team has been working on some pre-packaged solutions to make this process easier by providing starter kit projects with a few key technologies. One of those is the topic of today’s post - building a GenAI application with Spring AI in Java.

Let’s get started!

What is Spring AI?

Spring AI is a framework for building generative AI applications. It provides a set of tools and utilities for working with generative AI models and architectures, such as large language models (LLMs) and retrieval augmented generation (RAG). There are also other options for GenAI in Java, but Spring AI is one of a handful of Neo4j’s initial integrations across various communities.

What’s in this project?

All of the code and background information is available in the spring-ai-starter-kit Github repository. We will be walking through some of the code in this post, but you can always refer back to the repository.

We will use Neo4j to store structured data (such as entities and relationships), as well as the unstructured text data with the related vector embeddings. We can execute similarity searches on the vectors, then run retrieval queries to pull additional related entities. For today, we’ll pick the OpenAI model, but you can swap that out for any of the other Spring AI-supported LLMs.

If you take a look at the project’s pom.xml file, you will see the four dependencies we have included for this project:

  • Spring Web (for creating a REST API)

  • OpenAI (or other LLM model, such as Mistral, Ollama, etc.)

  • Neo4j Vector Database (for storing and querying vectors)

  • Spring Data Neo4j (for working with Neo4j in Spring applications)

Data set

The data set we are using is the SEC filings provided by EDGAR. It contains a collection of documents for company and individual filings. The documents are in plain text and contain financial information about various companies. Our teams have curated this data into a knowledge graph with information on companies, financial forms, and managers. It includes the following forms:

  1. Form 10-K: the annual report that publicly traded companies must file with the SEC. It provides a comprehensive summary of a company’s financial performance.

  2. Form 13: filed by institutional managers who manage $100 million or more in assets.

The starter kit application defaults to using a pre-loaded version of the data set in a public Neo4j Aura cloud database. This removes the headaches of data import and formatting and such. As an alternative, you can load your own knowledge graph from scratch or find out more about the data set by looking at our sec-edgar-notebooks repository.

With all of the background information out of the way, let’s check out the code!

Project setup

Go ahead and clone the project repository, and then open it in your favorite IDE. Looking at the pom.xml file, you should see that the milestone repository is included, along with our four dependencies from earlier. Since Spring AI is not a general-availability release yet, the milestone repository is included.

To use OpenAI’s LLM, you will need to request an API key by signing up at OpenAI. Once you have that (plus your database, if you created your own), you can set up the config in the application.properties file. Here’s an example of what that might look like:

spring.ai.openai.api-key=<YOUR API KEY HERE>
spring.neo4j.uri=neo4j+s://9fcf58c6.databases.neo4j.io
spring.neo4j.authentication.username=public
spring.neo4j.authentication.password=read_only
spring.data.neo4j.database=neo4j

The database credentials are for the default, pre-loaded instance. So, if you are using your own database, you will need to update the .uri, .username, and .password properties to match yours.

Next, let’s review the code for our application!

Embedding and vector store setup

In the SpringAiApplication class, we’ll need to set up a couple of Spring Beans for the OpenAI client and the Neo4j vector store that will allow us to access necessary components wherever we need them in our application.

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

The EmbeddingClient bean creates a client for the OpenAI API and passes in our API key from the properties file. Then, the Neo4jVectorStore bean configures Neo4j as the store for embeddings (vectors). It pulls in the driver, which is autoconfigured from our database credentials, as well as the embedding client. The vector store configuration needs customized to specify our vector index name (the Spring AI default is spring-ai-document-index), the label for the nodes that will store the embeddings (default looks for Document entities), and the name of property containing the embedding (default is embedding).

Application model

Next, we have some standard domain classes that map the application entities to our database model. There is a Chunk class that represents a document chunk node, and also classes for Form, Company, and Manager that represent the corresponding entities in the database. The Chunk entity has the embedding (vector) associated with it, which we’ll use for similarity searches.

These entities are standard Spring Data Neo4j code, so I won’t show the code here. However, full code for each class is available in the project’s Github repository.

The project’s repository interface allows the application to interact with the database. There is one defined query method that finds related entities (Form, Company, Manager) for the similar chunks.

Next, the controller class is where all the pieces come together to handle user requests and generate responses. This class will contain the logic for taking a question from the user and calling the Neo4jVectorStore to calculate and return the most similar documents. We can then pass those similar chunks into a Neo4j query to retrieve connected entities, providing additional context in the prompt for the LLM. It will use all the information provided to respond with a more detailed answer.

Controller

Our controller class needs to inject the Neo4jVectorStore bean, OpenAiChatClient bean, and ChunkRepository interface to run similarity search, call the LLM, and query the database, respectively.

@RestController
@RequestMapping("/api")
public class ChunkController {
    private final OpenAiChatClient client;
    private final Neo4jVectorStore vectorStore;
    private final ChunkRepository repo;

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

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

The last piece is to define a method that will be called when a user makes a GET request to the /chat endpoint. This method will take a question as a query parameter and pass that to the vector store’s similaritySearch() method to find similar document chunks.

Then, the similar Chunk nodes are mapped to Document entities because Spring AI expects a general document type. The Neo4jVectorStore class contains methods to convert Document to a custom record, as well as the reverse for record to Document conversion.

Back in our controller method, we now have similar document chunks, but chunks of text may not give enough detail for a helpful answer. So now we need to run the repository query in Neo4j to retrieve the related forms, companies, and managers for those chunks. This is the retrieval augmented generation (RAG) piece of the application.

After the similarity search returns similar document chunks, we call the getRelatedEntitiesForSimilarChunks() method in the repository (passing in the list of similar document ids) to find the related entities for those chunks.

The next block of code is the prompt template with the text to send to the LLM, along with the user’s question and graph results containing related entities. Finally, we call the template’s create() method to generate the response from the LLM and return the contents key for the response string.

Let’s test it out!

Running the application

To run our starter kit application, you can use the ./mvnw spring-boot:run command in the terminal. Once the application is running, you can make a GET request to the /api/chat endpoint with a question about the EDGAR data as a query parameter. A couple of examples are included next.

curl "http://localhost:8080/api/chat?question=How%20many%20forms%20are%20there%3F"
curl "http://localhost:8080/api/chat?question=Which%20companies%20are%20in%20healthcare%3F"
curl "http://localhost:8080/api/chat?question=Which%20managers%20own%20stock%20in%20more%20than%20one%20company%3F"

Feel free to play around with some SEC-related questions or tweak the application or code to see how it responds. You can also check the console output to see the data being passed back and forth between the application and the LLM.

Wrapping Up!

In today’s post, we checked out the Spring AI Neo4j starter kit to help you get started building GenAI applications in Java. We used Spring AI to extend the richness of the well-established Spring ecosystem, giving us the ability to write a GenAI app in a JVM language (Java, for this post).

While Spring AI supports a variety of LLM models and vector stores, we chose the OpenAI model and Neo4j database. Neo4j provides the ability to store relationships and standard structured data alongside the unstructured text data and vector embeddings. We used the OpenAI model to generate responses to user questions based on the similarity search results and related entities from the database.

I hope this post helps to get you started with GenAI and beyond. Happy coding!

Resources