rag:
  ocr:
    tesseract:
      use: true   # 啟用 Tesseract OCR（PDF 圖片文字識(shí)別）
  vector:
    milvus:
      use: true   # 啟用 Milvus 向量數(shù)據(jù)庫(kù)

@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("總頁(yè)數(shù)：" + documents.size());
    // 每個(gè) Document 對(duì)應(yīng) PDF 的一頁(yè)
}

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() + " 頁(yè)");

    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 模型，免費(fèi)
            .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() {
    // 假設(shè)文檔已經(jīng)存入 Milvus...
    BaseRetriever retriever = vectorStore.asRetriever();

    BaseRunnable<StringPromptValue, ?> prompt = PromptTemplate.fromTemplate(
        """
        請(qǐng)根據(jù)以下文檔內(nèi)容回答問題。如果文檔中沒有相關(guān)信息，請(qǐng)說"文檔中未找到相關(guān)信息"。
        
        文檔內(nèi)容：
        ${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)   // 向量檢索：輸入問題，返回相關(guān)文檔列表
        .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 模型中的注意力機(jī)制是如何工作的？"
    );

    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"));

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

    FlowInstance chain = chainActor.builder()
        .next(PromptTemplate.fromTemplate("請(qǐng)對(duì)以下內(nèi)容摘要（100字以內(nèi)）：\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());
}