package com.iwe3.langchain4j;

import ai.djl.huggingface.tokenizers.HuggingFaceTokenizer;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.parser.apache.pdfbox.ApachePdfBoxDocumentParser;
import dev.langchain4j.data.document.splitter.DocumentByParagraphSplitter;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.output.TokenUsage;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import jakarta.annotation.Resource;
import org.junit.jupiter.api.Test;
import org.springframework.boot.test.context.SpringBootTest;

@SpringBootTest
class RagApplicationTest {

    @Resource
    private ChatModel chatModel;


    @Test
    public void testCountToken(){
        var res = chatModel.chat(UserMessage.from("你好，这是一段文字！"));
        var result = res.aiMessage().text();
        System.out.println("通过调用大模型返回结果：" + result);
        /*Token用量计算底层API*/
        var tu = res.tokenUsage();
        System.out.println("本地调用销毁的TOKEN：" + tu);
    }


    /**
     * 加载文档并存入向量数据库
     */
    @Test
    public void testReadDocumentAndStore() {
        // 使用FileSystemDocumentLoader读取指定目录下的知识库文档
        // 并使用默认的文档解析器TextDocumentParser对文档进行解析
        var document = FileSystemDocumentLoader.loadDocument("D:/files/电子科大.md");

        // 为了简单起见，我们暂时使用基于内存的向量存储
        InMemoryEmbeddingStore<TextSegment> db = new InMemoryEmbeddingStore<>();

        // /ingest
        // /1、分割文档，默认使用递归分割器，将文档分割为多个文本片段，每个片段包含不超过 300个token，并且有 30个token的重叠部分保证连贯性
        // /2、文本向量化：使用一个LangChain4j内置的轻量化向量模型对每个文本片段进行向量化
        // /3、将原始文本和向量存储到向量数据库中(InMemoryEmbeddingStore)
        EmbeddingStoreIngestor.ingest(document, db);

        // 查看向量数据库内容
        System.out.println(db);
    }



    /**
     * 解析PDF
     */
    @Test
    public void testParsePDF() {
        var document = FileSystemDocumentLoader.loadDocument(
                "D:/files/电子科大.pdf",
                new ApachePdfBoxDocumentParser()
        );
        System.out.println(document);
    }

    @Test
    public void testReadDocument() {
        // 使用FileSystemDocumentLoader读取指定目录下的知识库文档
        // 并使用默认的文档解析器TextDocumentParser对文档进行解析
        var document = FileSystemDocumentLoader
                .loadDocument("D:/files/test.txt");
        System.out.println(document.metadata());
        System.out.println(document.text());
    }

}