使用Java实现RAG(检索增强生成)的完整指南

rag:
  ocr:
    tesseract:
      use: true   # 启用 Tesseract OCR（PDF 图片文字识别）
  vector:
    milvus:
      use: true   # 启用 Milvus 向量数据库

@Test
public void loadPdfDocuments() {
    PdfboxLoader loader = PdfboxLoader.builder()
        .filePath("./files/pdf/en/Transformer.pdf")
        .build();
    loader.setExtractImages(false);  // 不提取图片，只提取文本

    List<Document> documents = loader.load();

    System.out.println("总页数：" + documents.size());
    // 每个 Document 对应 PDF 的一页
}

ApachePoiDocxLoader loader = ApachePoiDocxLoader.builder()
    .filePath("./files/docx/en/Transformer.docx")
    .build();

List<Document> documents = loader.load();

@Test
public void splitDocuments() {
    List<Document> documents = loader.load();
    System.out.println("切分前：" + documents.size() + " 页");

    StanfordNLPTextSplitter splitter = StanfordNLPTextSplitter.builder()
        .chunkSize(1000)    // 每块最多 1000 字符
        .chunkOverlap(100)  // 相邻块重叠 100 字符，保证上下文连贯
        .build();

    List<Document> splits = splitter.splitDocument(documents);
    System.out.println("切分后：" + splits.size() + " 块");
}

@Test
public void embedAndStore() {
    // ... 加载、切分文档 ...

    VectorStore vectorStore = Milvus.fromDocuments(
        splits,
        OllamaEmbeddings.builder()
            .model("nomic-embed-text")  // 本地 Embedding 模型，免费
            .vectorSize(768)            // 向量维度
            .build(),
        "MyKnowledgeBase"               // Milvus collection 名称
    );

    System.out.println("向量化完成！");
}

docker run -d --name milvus \
  -p 19530:19530 \
  milvusdb/milvus:latest standalone

@Test
public void retrieveAndAsk() {
    // 假设文档已经存入 Milvus...
    BaseRetriever retriever = vectorStore.asRetriever();

    BaseRunnable<StringPromptValue, ?> prompt = PromptTemplate.fromTemplate(
        """
        请根据以下文档内容回答问题。如果文档中没有相关信息，请说"文档中未找到相关信息"。
        
        文档内容：
        ${context}
        
        问题：${question}
        
        回答：
        """
    );

    Function<Object, String> formatDocs = input -> {
        List<Document> docs = (List<Document>) input;
        StringBuilder sb = new StringBuilder();
        for (Document doc : docs) {
            sb.append(doc.getPageContent()).append("\n\n");
        }
        return sb.toString();
    };

    FlowInstance ragChain = chainActor.builder()
        .next(retriever)   // 向量检索：输入问题，返回相关文档列表
        .next(formatDocs)  // 将文档列表拼接为字符串
        .next(input -> Map.of(
            "context",  input,
            "question", ContextBus.get().getFlowParam()  // 获取原始问题
        ))
        .next(prompt)
        .next(llm)
        .next(new StrOutputParser())
        .build();

    ChatGeneration result = chainActor.invoke(
        ragChain,
        "Transformer 模型中的注意力机制是如何工作的？"
    );

    System.out.println(result.getText());
}

@Test
public void documentSummary() {
    // 加载 PDF
    List<Document> documents = loader.load();
    String content = documents.stream()
        .map(Document::getPageContent)
        .collect(Collectors.joining("\n"));

    // 长文档截取首尾
    String textToSummarize = content.length() < 2000 ? content
        : content.substring(0, 1000) + "\n...\n" + content.substring(content.length() - 1000);

    FlowInstance chain = chainActor.builder()
        .next(PromptTemplate.fromTemplate("请对以下内容摘要（100字以内）：\n\n${text}"))
        .next(ChatOllama.builder().model("qwen2.5:0.5b").build())
        .next(new StrOutputParser())
        .build();

    ChatGeneration result = chainActor.invoke(chain, Map.of("text", textToSummarize));
    System.out.println(result.getText());
}