Spring AI – Using vector database to implement retrieval AI dialogue

Spring AI – Using vector database to implement retrieval AI dialogue

Spring AI is not limited to the unified encapsulation of large language model dialogue APIs, it can also implement some functions of LangChain in a simple way. This article will lead readers to implement a simple search-based AI dialogue interface.

1. Demand background

In some scenarios, we want the AI ​​to respond based on the data we provide. Because the conversation has a maximum Token limit, we cannot directly send all the data to AI in many scenarios. On the one hand, when the amount of data is large, the Token limit will be exceeded. On the other hand, if the Token limit is not exceeded, There will also be unnecessary conversation fee overhead under restrictions. Therefore, how we can make AI respond better based on the data we provide while spending the least amount of money is a very critical issue. To solve this problem, we can use data vectorization to solve it.

2. Implementation principle

Store our personal data in a vector database. Then, before the user initiates a conversation with the AI, a set of similar documents is first retrieved from the vector database. These documents are then used as context for the user's question and sent to the AI ​​model along with the user's conversation, allowing for accurate responses. This approach is called Retrieval Augmented Generation (RAG).

Step one: Data vectorization

We have many ways to vectorize data, the simplest is by calling a third-party API. Taking OpenAI's API as an example, it provides the https://api.openai.com/v1/embeddings interface. By requesting this interface, you can obtain the vectorized data of a certain text. For details, please refer to the official API introduction: Create embeddings. In Spring AI, we do not have to call this interface to manually perform vectorization processing. Spring AI will automatically call it when storing it in the vector database.

Step 2: Vector storage and retrieval

There is a VectorStore abstract interface in Spring AI, which defines the interaction between Spring AI and the vector database. We can use this interface to operate the vector database through simple vector database configuration.

public interface VectorStore {

    void add(List<Document> documents);

    Optional<Boolean> delete(List<String> idList);

    List<Document> similaritySearch(String query);

    List<Document> similaritySearch(SearchRequest request);
} 

Vector Database is a special type of database that plays an important role in artificial intelligence applications. In vector databases, query operations are different from traditional relational databases. They perform similarity searches rather than exact matches. When a vector is given as a query, the vector database returns vectors that are “similar” to the query vector. In this way, we can integrate personal data with AI models. `

Common vector databases include: Chroma, Milvus, Pgvector, Redis, Neo4j, etc.

3. Code implementation

This article will implement the RAG based on ChatGPT and the interface for uploading PDF files to be stored in the vector database. The vector database uses Pgvector. Pgvector is an extension based on PostgreSQL that can store and retrieve embeddings generated during machine learning.

The source code has been uploaded to GitHub: https://github.com/NingNing0111/vector-database-demo

Version Information

• JDK >= 17
• Spring Boot >= 3.2.2
• Spring AI = 0.8.0-SNAPSHOT

1. Install Pgvector

Pgvector will be installed using Docker. The docker-compose.yml file is as follows:

version: '3.7'
services:
  postgres:
    image: ankane/pgvector:v0.5.0
    restart: always
    environment:
      - POSTGRES_USER=postgres
      - POSTGRES_PASSWORD=postgres
      - POSTGRES_DB=vector_store
      - PGPASSWORD=postgres
    logging:
      options:
        max-size: 10m
        max-file: "3"
    ports:
      - '5432:5432'
    healthcheck:
      test: "pg_isready -U postgres -d vector_store"
      interval: 2s
      timeout: 20s
      retries: 10 

2. Create a Spring project and add dependencies

The creation process of the Spring project is briefly described. The core content of pom.xml is as follows:

 <properties>
        <java.version>17</java.version>
        <!--  Spring AIversion information  -->
        <spring-ai.version>0.8.0-SNAPSHOT</spring-ai.version>
    </properties>

    <dependencies>
        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-jdbc</artifactId>
        </dependency>

        <!-- useOpenAI -->
        <dependency>
            <groupId>org.springframework.ai</groupId>
            <artifactId>spring-ai-openai-spring-boot-starter</artifactId>
            <version>{spring-ai.version}</version>
        </dependency>
        <!-- usePGVectoras vector database -->
        <dependency>
            <groupId>org.springframework.ai</groupId>
            <artifactId>spring-ai-pgvector-store-spring-boot-starter</artifactId>
            <version>{spring-ai.version}</version>
        </dependency>
        <!-- introducePDFparser -->
        <dependency>
            <groupId>org.springframework.ai</groupId>
            <artifactId>spring-ai-pdf-document-reader</artifactId>
            <version>${spring-ai.version}</version>
        </dependency>
    </dependencies>

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

3. Configure API, Key, PGVector connection information

server:
  port: 8801

spring:
  ai:
    openai:
      base-url: https://api.example.com
      api-key: sk-aec103e6cfxxxxxxxxxxxxxxxxxxxxxxx71da57a

  datasource:
    username: postgres
    password: postgres
    url: jdbc:postgresql://localhost/vector_store 

4. Create VectorStore and text splitter TokenTextSplitter

Here I created an ApplicationConfig configuration class

package com.ningning0111.vectordatabasedemo.config;

import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.PgVectorStore;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.jdbc.core.JdbcTemplate;

@Configuration
public class ApplicationConfig {

    /**
     * Vector database search operation
     * @param embeddingClient
     * @param jdbcTemplate
     * @return
     */
    @Bean
    public VectorStore vectorStore(EmbeddingClient embeddingClient, JdbcTemplate jdbcTemplate){
        return new PgVectorStore(jdbcTemplate,embeddingClient);
    }

    /**
     * text splitter
     * @return
     */
    @Bean
    public TokenTextSplitter tokenTextSplitter() {
        return new TokenTextSplitter();
    }
} 

5. Build PDF storage service layer

Create a class named PdfStoreService under the service layer to store PDF files in the vector database.

package com.ningning0111.vectordatabasedemo.service;

import lombok.RequiredArgsConstructor;
import org.springframework.ai.reader.ExtractedTextFormatter;
import org.springframework.ai.reader.pdf.PagePdfDocumentReader;
import org.springframework.ai.reader.pdf.ParagraphPdfDocumentReader;
import org.springframework.ai.reader.pdf.config.PdfDocumentReaderConfig;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.core.io.DefaultResourceLoader;
import org.springframework.core.io.FileSystemResource;
import org.springframework.core.io.Resource;
import org.springframework.stereotype.Service;
import org.springframework.web.multipart.MultipartFile;

import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;

@Service
@RequiredArgsConstructor
public class PdfStoreService {

    private final DefaultResourceLoader resourceLoader;
    private final VectorStore vectorStore;
    private final TokenTextSplitter tokenTextSplitter;

    /**
     * according toPDFdivided by the number of pages
     * @param url
     */
    public void saveSourceByPage(String url){
        // Load resources，Need local path information
        Resource resource = resourceLoader.getResource(url);
        // loadPDFConfiguration object in file
        PdfDocumentReaderConfig loadConfig = PdfDocumentReaderConfig.builder()
                .withPageExtractedTextFormatter(
                        new ExtractedTextFormatter
                                .Builder()
                                .withNumberOfBottomTextLinesToDelete(3)
                                .withNumberOfTopPagesToSkipBeforeDelete(1)
                                .build()
                )
                .withPagesPerDocument(1)
                .build();

        PagePdfDocumentReader pagePdfDocumentReader = new PagePdfDocumentReader(resource, loadConfig);
        // Store in vector database
        vectorStore.accept(tokenTextSplitter.apply(pagePdfDocumentReader.get()));
    }

    /**
     * according toPDFDirectory（paragraph）divide
     * @param url
     */
    public void saveSourceByParagraph(String url){
        Resource resource = resourceLoader.getResource(url);

        PdfDocumentReaderConfig loadConfig = PdfDocumentReaderConfig.builder()
                .withPageExtractedTextFormatter(
                        new ExtractedTextFormatter
                                .Builder()
                                .withNumberOfBottomTextLinesToDelete(3)
                                .withNumberOfTopPagesToSkipBeforeDelete(1)
                                .build()
                )
                .withPagesPerDocument(1)
                .build();

        ParagraphPdfDocumentReader pdfReader = new ParagraphPdfDocumentReader(
                resource,
                loadConfig
        );
        vectorStore.accept(tokenTextSplitter.apply(pdfReader.get()));
    }

    /**
     * MultipartFileobject storage，usePagePdfDocumentReader
     * @param file
     */
    public void saveSource(MultipartFile file){
        try {
            // Get file name
            String fileName = file.getOriginalFilename();
            // Get file content type
            String contentType = file.getContentType();
            // Get file byte array
            byte[] bytes = file.getBytes();
            // Create a temporary file
            Path tempFile = Files.createTempFile("temp-", fileName);
            // Save file byte array to temporary file
            Files.write(tempFile, bytes);
            // Create FileSystemResource object
            Resource fileResource = new FileSystemResource(tempFile.toFile());
            PdfDocumentReaderConfig loadConfig = PdfDocumentReaderConfig.builder()
                    .withPageExtractedTextFormatter(
                            new ExtractedTextFormatter
                                    .Builder()
                                    .withNumberOfBottomTextLinesToDelete(3)
                                    .withNumberOfTopPagesToSkipBeforeDelete(1)
                                    .build()
                    )
                    .withPagesPerDocument(1)
                    .build();
            PagePdfDocumentReader pagePdfDocumentReader = new PagePdfDocumentReader(fileResource, loadConfig);
            vectorStore.accept(tokenTextSplitter.apply(pagePdfDocumentReader.get()));
        }catch (IOException e){
            e.printStackTrace();
        }

    }
} 

6. Build conversation service

Create the ChatService class, which provides two dialogue modes: normal dialogue mode without retrieval and dialogue mode for retrieval of vector database

package com.ningning0111.vectordatabasedemo.service;

import lombok.RequiredArgsConstructor;
import org.springframework.ai.chat.ChatClient;
import org.springframework.ai.chat.ChatResponse;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.SystemPromptTemplate;
import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.stereotype.Service;

import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

@Service
@RequiredArgsConstructor
public class ChatService {

    // System prompt word
    private final static String SYSTEM_PROMPT = """
            You need to use the document content to respond to user questions，And you need to act like you know this stuff innately，
            You cannot show in your reply that you are replying based on the content of the document given.，This point is very important。

            When the user asks a question that cannot be answered based on the document content or you are not clear about it，Just reply "I don't know"。

            The document content is as follows:
            {documents}

            """;

    private final ChatClient chatClient;
    private final VectorStore vectorStore;

    // simple conversation，Vector database is not searched
    public String simpleChat(String userMessage) {
        return chatClient.call(userMessage);
    }

    // Search through vector database
    public String chatByVectorStore(String message) {
        // Similarity search based on question text
        List<Document> listOfSimilarDocuments = vectorStore.similaritySearch(message);
        // WillDocumentfor each element in the listcontentThe content is obtained by splicingdocuments
        String documents = listOfSimilarDocuments.stream().map(Document::getContent).collect(Collectors.joining());
        // useSpring AI Built using templates providedSystemMessageobject
        Message systemMessage = new SystemPromptTemplate(SYSTEM_PROMPT).createMessage(Map.of("documents", documents));
        // ConstructUserMessageobject
        UserMessage userMessage = new UserMessage(message);
        // WillMessageThe list is sent toChatGPT
        ChatResponse rsp = chatClient.call(new Prompt(List.of(systemMessage, userMessage)));
        return rsp.getResult().getOutput().getContent();
    }
} 

7. Build Controller layer

ChatController provides a conversation interface:

package com.ningning0111.vectordatabasedemo.controller;

import com.ningning0111.vectordatabasedemo.service.ChatService;
import lombok.RequiredArgsConstructor;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

@RestController
@RequiredArgsConstructor
@RequestMapping("/api/v1/chat")
public class ChatController {

    private final ChatService chatService;

    @GetMapping("/simple")
    public String simpleChat(
            @RequestParam String message
    ){
        return chatService.simpleChat(message);
    }

    @GetMapping("/")
    public String chat(
            @RequestParam String message
    ){
        return chatService.chatByVectorStore(message);
    }
} 

PdfUploadController provides an interface for uploading files and saving them to the vector database

package com.ningning0111.vectordatabasedemo.controller;

import com.ningning0111.vectordatabasedemo.service.PdfStoreService;
import lombok.RequiredArgsConstructor;
import org.springframework.stereotype.Controller;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.multipart.MultipartFile;

@Controller
@RequestMapping("/api/v1/pdf")
@RequiredArgsConstructor
public class PdfUploadController {
    private final PdfStoreService pdfStoreService;

    @PostMapping("/upload")
    public void upload(
            @RequestParam MultipartFile file
    ){
        pdfStoreService.saveSource(file);
    }
} 

The source code has been uploaded to GitHub: https://github.com/NingNing0111/vector-database-demo