package xyz.wbsite.ai;

import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.splitter.DocumentSplitters;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.embedding.onnx.bgesmallenv15q.BgeSmallEnV15QuantizedEmbeddingModel;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;

import java.util.List;

/**
 * 主函数入口
 */
public class Naive_RAG_Example {

    public static void main(String[] args) {
        // 加载文档
        Document document = Helper.getDocument();

        //现在，我们需要将此文档拆分为更小的部分，也称为“块”
        //这种方法允许我们仅向LLM发送相关段以响应用户查询，
        //而不是整个文档。例如如果用户询问取消政策，
        //我们将仅识别并发送与取消相关的片段。
        //一个好的起点是使用递归文档拆分器，最初尝试
        //按段落分割。如果一个段落太大而无法放入单个片段中，
        //拆分器将递归地按换行符、句子和单词进行划分，
        //如有必要，确保每段文本都适合一个片段。
        DocumentSplitter splitter = DocumentSplitters.recursive(300, 0);
        List<TextSegment> segments = splitter.split(document);

        //现在，我们需要嵌入（也称为“矢量化”）这些片段。
        //执行相似性搜索需要嵌入。
        //对于这个例子，我们将使用本地进程内嵌入模型，但您可以选择任何支持的模型。
        //Langchain4j目前支持10多个流行的嵌入模型提供者。
        EmbeddingModel embeddingModel = new BgeSmallEnV15QuantizedEmbeddingModel();
        List<Embedding> embeddings = embeddingModel.embedAll(segments).content();

        //接下来，我们将把这些嵌入存储在嵌入存储中（也称为“向量数据库”）。
        //此存储将用于在每次与LLM交互时搜索相关细分市场。
        //为简单起见，此示例使用内存中的嵌入存储，但您可以从任何支持的存储中进行选择。
        //Langchain4j目前支持超过15个流行的嵌入商店。
        InMemoryEmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();
        embeddingStore.addAll(embeddings, segments);

        // 我们还可以使用EmbeddingStoreIngestor将上面的手动步骤隐藏在更简单的API后面。
        // 请参阅_01_Advanced_RAG_with_Query_Compression_example中使用嵌入式存储器的示例。

        //内容检索器负责根据用户查询检索相关内容。
        //目前，它能够检索文本段，但未来的增强功能将包括支持
        //其他模态，如图像、音频等。
        ContentRetriever contentRetriever = EmbeddingStoreContentRetriever.builder()
                .embeddingStore(embeddingStore)
                .embeddingModel(embeddingModel)
                .maxResults(2) // 在每次交互中，我们将检索2个最相关的片段
                .minScore(0.1) // 我们希望检索至少与用户查询有些相似的段
                .build();

        //我们可以选择使用聊天存储器，与LLM进行来回对话
        //并允许它记住之前的交互。
        //目前，LangChain4j提供了两种聊天内存实现：
        //MessageWindowChatMemory和TokenWindowChatMemory。
        ChatMemory chatMemory = MessageWindowChatMemory.withMaxMessages(10);

        //最后一步是构建我们的人工智能服务，
        //配置它以使用我们上面创建的组件。
        Assistant assistant = AiServices.builder(Assistant.class)
                .chatLanguageModel(Helper.getChatModel())
                .contentRetriever(contentRetriever)
                .chatMemory(chatMemory)
                .build();

        String chat = assistant.chat("你是谁？");
        System.out.println(chat);
    }


    // 创建一个助手接口
    interface Assistant {

        String chat(String userMessage);
    }
}