---

Building a RAG Chatbot with Spring Using Yellowbrick as a Vector Store

Overview

If you are a Java developer, you have probably used libraries from the Spring ecosystem. Java developers love Spring because it simplifies the development of complex applications. Its comprehensive ecosystem includes many specialized frameworks that reduce the complexity of integrating third-party APIs, databases, messaging systems, and more.

Spring AI is the latest project within the Spring ecosystem, designed to make it easy for Java developers to integrate AI capabilities into their applications. The framework abstracts the complexities of connecting to various AI-provider APIs by using familiar Spring abstractions.

Spring AI supports common tasks, such as implementing chatbots and virtual assistants, content generation and summarization, sentiment analysis, customer feedback processing, image and text recognition, and decision support.

In this tutorial, we will explain how to create a chatbot by using Spring AI. First, we will build a simple chatbot by using OpenAI, then enhance it with Retrieval-Augmented Generation (RAG) for improved contextual accuracy.

The RAG technique improves accuracy by augmenting the chatbot's knowledge with external data, such as a product manual, beyond the model's training data.

The Process

The process begins when the model receives a query, such as a user question. Use a pre-trained language model to encode this query into a vector (a mathematical representation).

Next, the system searches for relevant information within an external knowledge base by using the encoded query. The retrieved documents are then appended to the original query, providing additional context.

A vector database is ideal for storing these mathematical representations. While Yellowbrick does not natively support vector data types, it can store vector data efficiently for retrieval.

Spring abstracts the logic for interfacing with the vector store, allowing developers to switch databases without significant code changes. We will explain this by first using pgVector (a PostgreSQL extension) and then switching to Yellowbrick.

Create the Chatbot

All code for the following steps is available on GitHub.

Prerequisites

• JDK 17 or later
• OpenAI developer key

Getting Started

Please perform the following steps:

1. Visit start.spring.io and create a new project. Select Java 17 and Maven. Add dependencies for Spring Web and OpenAI. Decompress the generated project and import it into your IDE.

2. Add your OpenAI key to application.properties:

spring.ai.openai.api-key=<YOUR_OPENAI_KEY>
spring.main.web-application-type=none

3. Update your Spring application class:

@Bean
public CommandLineRunner runner(ChatClient.Builder builder) {
    return args -> {
        ChatClient chatClient = builder.build();
        String response = chatClient.prompt("Tell me a joke").call().content();
        System.out.println(response);
    };
}

4. Run the application:

./mvnw spring-boot:run

Add pgVector as a Vector Store

1. Start a pgVector Docker container:

docker run -it --rm -p 5432:5432 -e POSTGRES_USER=postgres -e POSTGRES_PASSWORD=postgres pgvector/pgvector:pg17

2. Add the pgVector dependency to pom.xml:

<dependency>
  <groupId>org.springframework.ai</groupId>
  <artifactId>spring-ai-pgvector-store-spring-boot-starter</artifactId>
</dependency>

3. To read PDF files, add:

<dependency>
  <groupId>org.springframework.ai</groupId>
  <artifactId>spring-ai-pdf-document-reader</artifactId>
</dependency>

4. Configure the database connection in application.properties:

spring.application.name=springai
spring.datasource.url=jdbc:postgresql://localhost:5432/postgres
spring.datasource.username=postgres
spring.datasource.password=postgres
spring.ai.vectorstore.pgvector.initialize-schema=true
spring.ai.vectorstore.pgvector.max-document-batch-size=1000

Add Content to the Vector Store

We will use an Employee Handbook PDF for this example. Place it in src/main/resources/data.

Create a DataLoaderService class:

@Service
public class DataLoaderService {

  @Value("classpath:/data/Employee_Handbook.pdf")
  private Resource pdfResource;

  @Autowired
  private VectorStore vectorStore;

  @PostConstruct
  public void load() {
    PagePdfDocumentReader pdfReader = new PagePdfDocumentReader(pdfResource,
        PdfDocumentReaderConfig.builder()
            .withPageExtractedTextFormatter(ExtractedTextFormatter.builder()
                .withNumberOfBottomTextLinesToDelete(3)
                .withNumberOfTopPagesToSkipBeforeDelete(1)
                .build())
            .withPagesPerDocument(1)
            .build());

    var tokenTextSplitter = new TokenTextSplitter();
    vectorStore.accept(tokenTextSplitter.apply(pdfReader.get()));
  }
}

Query with Context from the Vector Store

Prompts

The prompt ensures contextually correct responses. It consists of the following:

• Context: Text snippets from the knowledge base.
• Query: The user's input.

Code

Please perform the following steps:

1. Create a ChatService:

@Service
public class ChatService {

  @Autowired
  @Qualifier("openAiChatModel")
  private ChatModel chatClient;

  @Autowired
  private VectorStore vectorStore;

  private final String PROMPT_BLUEPRINT = """
    Answer the query using the provided context:
    {context}
    Query:
    {query}
    If no context is available, respond with:
    "I'm sorry, I don't have the information you're looking for."
  """;

  public String chat(String query) {
    return chatClient.call(createPrompt(query, searchData(query)));
  }

  private String createPrompt(String query, List<Document> context) {
    PromptTemplate template = new PromptTemplate(PROMPT_BLUEPRINT);
    template.add("query", query);
    template.add("context", context);
    return template.render();
  }

  public List<Document> searchData(String query) {
    return vectorStore.similaritySearch(query);
  }
}

2. Add a controller:

@RestController
class AIController {

  private final ChatService chatService;

  AIController(ChatService chatService) {
    this.chatService = chatService;
  }

  @GetMapping("/chat")
  public Map<String, String> chat(@RequestParam String query) {
    return Map.of("answer", chatService.chat(query));
  }
}

3. Register a ChatClient bean:

@Bean
ChatClient chatClient(ChatClient.Builder builder) {
  return builder.build();
}

4. Remove spring.main.web-application-type=none from application.properties, then run the application.

5. Test with HTTPie, curl, or Postman:

http -v GET localhost:8080/chat?query='What is the vacation policy'

6. The following Json statement shows the expected response:

{
  "answer": "Vacations are provided as a benefit..."
}

Switch to Yellowbrick as a Vector Store

Steps

Please perform the following steps:

1. Clone the repository:

git clone https://github.com/YellowbrickData/yellowbrick-learn

2. Navigate to the vector store directory:

cd yellowbrick-learn/samples/rags/vectorstore

3. Build the artifact:

./mvnw install

4. Replace the pgVector dependency in pom.xml with:

<dependency>
  <groupId>com.yellowbrick.spring</groupId>
  <artifactId>vectorstore</artifactId>
  <version>0.0.1-SNAPSHOT</version>
</dependency>

5. Update application.properties:

spring.datasource.url=jdbc:postgresql://myinstance.aws.yellowbrickcloud.com:5432/<username>
spring.datasource.username=<username>
spring.datasource.password=<password>

6. Restart the application. Queries will now run on Yellowbrick.

Bonus: Switch to Amazon Bedrock

Steps

Please perform the following steps:

1. Enable the Titan model in your AWS account.

2. Replace the OpenAI dependency with Bedrock in pom.xml:

<dependency>
  <groupId>org.springframework.ai</groupId>
  <artifactId>spring-ai-bedrock-ai-spring-boot-starter</artifactId>
</dependency>

3. Update application.properties:

spring.ai.bedrock.titan.chat.enabled=true
spring.ai.bedrock.titan.embedding.enabled=true
spring.ai.bedrock.titan.embedding.model=amazon.titan-embed-text-v2:0
spring.ai.bedrock.aws.region=us-east-1
spring.ai.bedrock.aws.secret-key=<your_secret_key>
spring.ai.bedrock.aws.access-key=<your_access_key>

4. Rerun the application to see similar results.

References

1. Spring AI
2. Yellowbrick as a Vector Store (Python Edition)
3. What is Retrieval-Augmented Generation
4. Sample App on GitHub