RAG 知识库接口服务实现

一、介绍

根据前一章节知识库的测试案例，本章将会把测试案例封装成接口形式提供给外部使用。

二、技术方案

知识库的上传和使用流程清晰，但需要明确选择哪个知识库，并记录相关信息。为此，我们采用 Redis 列表来高效记录知识库的选择。对于公司级的大型知识库，则需使用 MySQL 数据库进行持久化存储，以满足更高的数据管理需求。

三、具体实现

1. 项目结构

image-20250403110356890

2. 引入组件

<!-- lxf-rag-app pom.xml -->
<dependency>
    <groupId>org.redisson</groupId>
    <artifactId>redisson-spring-boot-starter</artifactId>
    <version>3.44.0</version>
</dependency>

<!-- lxf-rag-trigger pom.xml -->
<dependency>
    <groupId>org.springframework.ai</groupId>
    <artifactId>spring-ai-tika-document-reader</artifactId>
</dependency>
<dependency>
    <groupId>org.springframework.ai</groupId>
    <artifactId>spring-ai-starter-vector-store-pgvector</artifactId>
</dependency>
<dependency>
    <groupId>org.springframework.ai</groupId>
    <artifactId>spring-ai-ollama</artifactId>
</dependency>

本章节使用的是 redis 来存储知识库，所以在这里需要引入操作 redis 的客户端 Redisson 组件。

3. 配置链接

redis:
  sdk:
    config:
      host: 192.168.1.107
      port: 16379
      pool-size: 10
      min-idle-size: 5
      idle-timeout: 30000
      connect-timeout: 5000
      retry-attempts: 3
      retry-interval: 1000
      ping-interval: 60000
      keep-alive: true

这里配置你的 Redis 链接地址。IP、端口，设置为你自己的端口，也可以自己配置密码。

4. 上传知识

@Slf4j
@RestController
@CrossOrigin("*")
@RequestMapping("/api/v1/rag/")
public class RAGController implements IRAGService {

    @Resource
    private OllamaChatModel ollamaChatModel;
    @Resource
    private TokenTextSplitter tokenTextSplitter;
    @Resource
    private SimpleVectorStore simpleVectorStore;
    @Resource
    private PgVectorStore pgVectorStore;
    @Resource
    private RedissonClient redissonClient;

    @GetMapping("/query_rag_tag_list")
    @Override
    public Response<List<String>> queryRagTagList() {
        RList<String> ragTagList = redissonClient.getList("ragTag");
        return Response.<List<String>>builder()
                .code("1000")
                .info("调用成功")
                .data(ragTagList)
                .build();
    }

    @GetMapping(value = "file/upload", headers = "content-type=multipart/form-data")
    @Override
    public Response<String> uploadFile(@RequestParam String ragTag, @RequestParam("file") List<MultipartFile> files) {
        log.info("上传知识库开始 {}", ragTag);
        for (MultipartFile file : files) {
            TikaDocumentReader documentReader = new TikaDocumentReader(file.getResource());
            List<Document> documents = documentReader.get();
            List<Document> documentSplitterList = tokenTextSplitter.apply(documents);

            // 添加知识库标签
            documents.forEach(doc -> doc.getMetadata().put("cactusli", ragTag));
            documentSplitterList.forEach(doc -> doc.getMetadata().put("cactusli", ragTag));

            pgVectorStore.accept(documentSplitterList);

            // 添加知识库记录
            RList<String> elements = redissonClient.getList("ragTag");
            if (!elements.contains(ragTag)){
                elements.add(ragTag);
            }
        }
        log.info("上传知识库完成 {}", ragTag);
        return Response.<String>builder().code("1000").info("调用成功").build();
    }

}

此Controller 提供了上传文件的接口，之后此接口会对接前端页面。

5. 使用知识

@RestController
@CrossOrigin({"*"})
@RequestMapping("/api/v1/ollama/")
public class OllamaController implements IAiService {

    @Resource
    private OllamaChatModel chatModel;

    @Resource
    private PgVectorStore pgVectorStore;

    @GetMapping(value = "generate_stream_rag")
    @Override
    public Flux<ChatResponse> generateStreamRag(@RequestParam String model, @RequestParam String ragTag, @RequestParam String message) {
        String SYSTEM_PROMPT = """
                Use the information from the DOCUMENTS section to provide accurate answers but act as if you knew this information innately.
                If unsure, simply state that you don't know.
                Another thing you need to note is that your reply must be in Chinese!
                DOCUMENTS:
                    {documents}
                """;
        // 指定文档搜索
        SearchRequest request = SearchRequest.builder()
                .query(message)
                .topK(5)
                .filterExpression("cactusli == '" + ragTag + "'")
                .build();

        List<Document> documents = pgVectorStore.similaritySearch(request);
        String documentCollectors = documents.stream().map(Document::getText).collect(Collectors.joining());
        Message ragMessage = new SystemPromptTemplate(SYSTEM_PROMPT).createMessage(Map.of("documents", documentCollectors));

        // 组装消息
        ArrayList<Message> messages = new ArrayList<>();
        messages.add(new UserMessage(message));
        messages.add(ragMessage);

        return chatModel.stream(new Prompt(
                messages,
                OllamaOptions.builder()
                        .model(model)
                        .build())
        );
    }
}

OllamaController 提供了 generateStreamRag 使用知识接口,这个过程就相当于把用户描述的信息，增加了更多的数据内容，也更加准确了。