102-Spring AI Alibaba RAG Pgvector 示例

本示例将引导您一步步构建一个基于 Spring AI Alibaba 的检索增强生成(RAG)应用,使用 PostgreSQL 的 pgvector 扩展作为向量数据库。该应用能够导入文档、处理文档内容、生成向量嵌入,并基于导入的文档内容回答用户问题。
1. 示例目标
我们将创建一个包含以下核心功能的 Web 应用:
- 文档导入 (
/ai/rag/importFileV2):支持上传各种格式的文档(如 PDF、Word、文本等),解析文档内容,将文本分块,生成向量嵌入并存储到 PostgreSQL 的 pgvector 中。 - 文档检索与问答 (
/ai/rag/searchV2):根据用户的问题,从向量数据库中检索相关文档片段,并基于检索到的内容生成回答。 - 文档管理 (
/ai/rag/deleteFilesV2):支持删除已导入的文档及其相关的向量数据。
2. 技术栈与核心依赖
- Spring Boot 3.x
- Spring AI Alibaba (用于对接阿里云 DashScope 通义大模型)
- PostgreSQL with pgvector (作为向量数据库)
- Spring AI Vector Store (用于向量存储和检索)
- Spring AI Document Readers (用于文档解析)
- Maven (项目构建工具)
在 pom.xml 中,你需要引入以下核心依赖:
<dependencies><!-- Spring Web 用于构建 RESTful API --><dependency><groupId>org.springframework.boot</groupId><artifactId>spring-boot-starter-web</artifactId></dependency><!-- Spring AI Alibaba 核心启动器,集成 DashScope --><dependency><groupId>com.alibaba.cloud.ai</groupId><artifactId>spring-ai-alibaba-starter-dashscope</artifactId></dependency><!-- PDF 文档读取器 --><dependency><groupId>org.springframework.ai</groupId><artifactId>spring-ai-pdf-document-reader</artifactId></dependency><!-- Pgvector 向量存储 --><dependency><groupId>org.springframework.ai</groupId><artifactId>spring-ai-pgvector-store</artifactId></dependency><!-- Pgvector 自动配置 --><dependency><groupId>org.springframework.ai</groupId><artifactId>spring-ai-autoconfigure-vector-store-pgvector</artifactId></dependency><!-- Tika 文档读取器,支持多种格式 --><dependency><groupId>org.springframework.ai</groupId><artifactId>spring-ai-tika-document-reader</artifactId></dependency>
</dependencies>3. 项目配置
在 src/main/resources/application.yml 文件中,配置你的 DashScope API Key 和 PostgreSQL 连接信息。
spring:application:name: rag-pgvector-exampledatasource:url: jdbc:postgresql://127.0.0.1:5432/postgresusername: postgrespassword: mysecretpasswordai:dashscope:api-key: ${DASH_SCOPE_API_KEY}vectorstore:pgvector:dimensions: 1536index-type: hnswdistance-type: cosine_distance重要提示:请将 DASH_SCOPE_API_KEY 环境变量设置为你从阿里云获取的有效 API Key。同时,确保 PostgreSQL 数据库已安装 pgvector 扩展。
4. 数据库准备
在使用本示例前,需要在 PostgreSQL 数据库中创建向量存储表。执行以下 SQL 命令:
CREATE EXTENSION IF NOT EXISTS "uuid-ossp";
CREATE TABLE IF NOT EXISTS vector_store (
id uuid DEFAULT uuid_generate_v4() PRIMARY KEY,
content text,
metadata json,
embedding vector(1536)
);CREATE INDEX ON vector_store USING HNSW (embedding vector_cosine_ops);5. 准备提示词模板文件
在 src/main/resources 目录下创建以下文件结构:
src/main/resources/
└── prompts/└── system-qa.st5.1 prompts/system-qa.st
系统提示词模板,用于定义 RAG 问答的上下文和回答规则。
Context information is below.
---------------------
{question_answer_context}
---------------------
Given the context and provided history information and not prior knowledge,
reply to the user comment. If the answer is not in the context, inform
the user that you can't answer the question.6. 编写 Java 代码
6.1 RagService.java
定义 RAG 服务的接口。
package com.alibaba.cloud.ai.example.rag;import org.springframework.ai.chat.model.ChatResponse;
import reactor.core.publisher.Flux;/*** Title Rag service.* Description Rag service.* */public interface RagService {void importDocuments();Flux<ChatResponse> retrieve(String message);
}6.2 RagExampleApplication.java
Spring Boot 应用程序的主类。
package com.alibaba.cloud.ai.example.rag;import org.springframework.boot.SpringApplication;
import org.springframework.boot.autoconfigure.SpringBootApplication;@SpringBootApplication
public class RagExampleApplication {public static void main(String[] args) {SpringApplication.run(RagExampleApplication.class, args);}}6.3 RagPgVectorController.java
实现 RAG 功能的 REST API 控制器。
package com.alibaba.cloud.ai.example.rag.controller;import com.alibaba.cloud.ai.advisor.RetrievalRerankAdvisor;
import com.alibaba.cloud.ai.model.RerankModel;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.SystemPromptTemplate;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.DocumentReader;
import org.springframework.ai.reader.pdf.PagePdfDocumentReader;
import org.springframework.ai.reader.tika.TikaDocumentReader;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.ai.vectorstore.filter.FilterExpressionBuilder;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.io.Resource;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.util.StringUtils;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;
import reactor.core.publisher.Flux;import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.UUID;@RestController
@RequestMapping("/ai")
public class RagPgVectorController {@Value("classpath:/prompts/system-qa.st")private Resource systemResource;@Value("classpath:/data/spring_ai_alibaba_quickstart.pdf")private Resource springAiResource;private final VectorStore vectorStore;private final ChatModel chatModel;private final RerankModel rerankModel;public RagPgVectorController(VectorStore vectorStore, ChatModel chatModel, RerankModel rerankModel) {this.vectorStore = vectorStore;this.chatModel = chatModel;this.rerankModel = rerankModel;}@GetMapping("/rag/importDocument")public void importDocument() {// 1. parse documentDocumentReader reader = new PagePdfDocumentReader(springAiResource);List<Document> documents = reader.get();// 1.2 use local file// FileSystemResource fileSystemResource = new FileSystemResource("D:\\file.pdf");// DocumentReader reader = new PagePdfDocumentReader(fileSystemResource);// 2. split trunksList<Document> splitDocuments = new TokenTextSplitter().apply(documents);// 3. create embedding and store to vector storevectorStore.add(splitDocuments);}/*** Receive any long text, split it and write it into a vector store*/@GetMapping("/rag/importText")public ResponseEntity<String> insertText(@RequestParam("text") String text) {// 1.parameter verificationif (!StringUtils.hasText(text)) {return ResponseEntity.badRequest().body("Please enter text");}// 2.parse documentList<Document> documents = List.of(new Document(text));// 3.Splitting TextList<Document> splitDocuments = new TokenTextSplitter().apply(documents);// 4.create embedding and store to vector storevectorStore.add(splitDocuments);// 5.return success promptString msg = String.format("successfully inserted %d text fragments into vector store", splitDocuments.size());return ResponseEntity.ok(msg);}/*** read and write multiple files and write it into a vector store* @param file* @return*/@PostMapping(value = "/rag/importFile", consumes = MediaType.MULTIPART_FORM_DATA_VALUE)public ResponseEntity<String> insertFiles( @RequestPart(value = "file", required = false) MultipartFile file) {// 1. file verificationif (file == null || file.isEmpty()) {return ResponseEntity.badRequest().body("必须上传非空的文件");}// 2. parse filesList<Document> docs = new TikaDocumentReader(file.getResource()).get();// 3. Splitting TextList<Document> splitDocs = new TokenTextSplitter().apply(docs);// 4. create embedding and store to vector storevectorStore.add(splitDocs);// 5.return success promptString msg = String.format("successfully inserted %d text fragments into vector store", splitDocs.size());return ResponseEntity.ok(msg);}@GetMapping(value = "/rag", produces = MediaType.TEXT_EVENT_STREAM_VALUE)public Flux<ChatResponse> generate(@RequestParam(value = "message",defaultValue = "how to get start with spring ai alibaba?") String message) throws IOException {SearchRequest searchRequest = SearchRequest.builder().topK(2).build();String promptTemplate = systemResource.getContentAsString(StandardCharsets.UTF_8);return ChatClient.builder(chatModel).defaultAdvisors(new RetrievalRerankAdvisor(vectorStore, rerankModel, searchRequest, new SystemPromptTemplate(promptTemplate), 0.1)).build().prompt().user(message).stream().chatResponse();}/*** read and write multiple files and write it into a vector store* @param file* @return*/@PostMapping(value = "/rag/importFileV2", consumes = MediaType.MULTIPART_FORM_DATA_VALUE)public ResponseEntity<String> importFileV2(@RequestPart(value = "file", required = false) MultipartFile file) {// 1. file verificationif (file == null || file.isEmpty()) {return ResponseEntity.badRequest().body("必须上传非空的文件");}// 2. parse filesList<Document> docs = new TikaDocumentReader(file.getResource()).get();// 3. Splitting TextList<Document> splitDocs = new TokenTextSplitter().apply(docs);String fileId = UUID.randomUUID().toString();for (Document doc : splitDocs) {doc.getMetadata().put("fileId", fileId);}// 4. create embedding and store to vector storevectorStore.add(splitDocs);// 5.return success promptString msg = String.format("successfully inserted %d text fragments into vector store," +" fileId: %s", splitDocs.size(), fileId);return ResponseEntity.ok(msg);}/*** search the vector store* @param message* @param fileId* @return* @throws IOException*/@GetMapping(value = "/rag/searchV2")public Flux<String> search(@RequestParam(value = "message",defaultValue = "what is blibaba?") String message,@RequestParam(value = "fileId", required = true)String fileId) throws IOException {FilterExpressionBuilder b = new FilterExpressionBuilder();Filter.Expression expression = b.eq("fileId", fileId).build();SearchRequest searchRequest = SearchRequest.builder().topK(1).filterExpression(expression).build();String promptTemplate = systemResource.getContentAsString(StandardCharsets.UTF_8);return ChatClient.builder(chatModel).defaultAdvisors(new RetrievalRerankAdvisor(vectorStore, rerankModel, searchRequest, new SystemPromptTemplate(promptTemplate), 0.1)).build().prompt().user(message).stream().content();}@PostMapping(value = "/rag/deleteFilesV2")public ResponseEntity<String> deleteFiles(@RequestParam(value = "fileId", required = false) String fileId) {FilterExpressionBuilder b = new FilterExpressionBuilder();Filter.Expression expression = b.eq("fileId", fileId).build();vectorStore.delete(expression);return ResponseEntity.ok("successfully deleted");}}7. 运行与测试
- 启动应用:运行你的 Spring Boot 主程序。
- 使用浏览器或 API 工具(如 Postman, curl)进行测试。
测试 1:导入文档
使用以下命令上传文档到向量数据库:
curl -X POST http://localhost:8080/ai/rag/importFileV2 \-F "file=@/path/to/your/file"预期响应:
successfully inserted XX text fragments into vector store, fileId: [UUID]
测试 2:基于文档内容进行问答
使用以下命令基于导入的文档内容提问:
curl -G 'http://localhost:8080/ai/rag/searchV2' \--data-urlencode 'messages=what is alibaba?' \--data-urlencode 'fileId={fileId}'将 {fileId} 替换为导入文档时返回的实际文件 ID。
预期响应:
基于导入的文档内容,系统会返回相关回答。如果文档中没有相关信息,系统会告知无法回答该问题。
测试 3:删除文档
使用以下命令删除已导入的文档:
curl -X DELETE 'http://localhost:8080/ai/rag/deleteFilesV2?fileId={fileId}'将 {fileId} 替换为要删除的文档 ID。
预期响应:
successfully deleted
8. 实现思路与扩展建议
实现思路
本示例的核心思想是"检索增强生成"。通过以下步骤实现:
- 文档处理:使用 TikaDocumentReader 解析各种格式的文档,使用 TokenTextSplitter 将文档分割成适当大小的文本块。
- 向量存储:将文本块转换为向量嵌入,并存储到 PostgreSQL 的 pgvector 中,同时为每个文本块添加元数据(如文件 ID)。
- 检索与重排序:根据用户问题检索相关文本块,使用 RerankModel 对检索结果进行重排序,提高相关性。
- 生成回答:将检索到的文本块作为上下文,使用大语言模型生成回答。
扩展建议
- 支持更多文档格式:通过集成更多文档解析器,支持更多格式的文档导入。
- 优化文本分割策略:根据不同类型的文档,采用更适合的文本分割策略,提高检索效果。
- 实现增量更新:支持对已导入文档的增量更新,避免重复处理。
- 添加用户权限管理:实现基于用户的文档访问权限控制,确保数据安全。
- 集成多模态能力:支持图像、音频等多模态内容的处理和检索。
- 实现对话历史:结合 Spring AI Alibaba 的对话记忆功能,实现多轮对话的上下文感知。
