# 基于LangChain和Llama.cpp的本地RAG实现
# 安装依赖: pip install langchain llama-cpp-python chromadb pypdf sentence-transformers

import os
from langchain.document_loaders import PyPDFLoader, DirectoryLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import Chroma
from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate
from langchain.llms import LlamaCpp
from langchain.callbacks.manager import CallbackManager
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler


from llama_cpp_embeddings import LlamaCppEmbeddings


class LocalRAGSystem:
    def __init__(self, 
                 data_dir="./documents",
                 db_dir="./local_chroma_db",
                 model_path="./models/llama-2-7b-chat.Q4_K_M.gguf",
                 embedding_model_name="all-MiniLM-L6-v2"):
        """
        初始化本地RAG系统
        :param data_dir: 文档存放目录
        :param db_dir: 向量数据库存放目录
        :param model_path: 本地LLM模型路径(GGUF格式)
        :param embedding_model_name: 用于生成嵌入的模型名称
        """
        self.data_dir = data_dir
        self.db_dir = db_dir
        self.model_path = model_path
        self.embedding_model_name = embedding_model_name
        
        # 初始化组件
        self.embeddings = None
        self.vector_db = None
        self.llm = None
        self.qa_chain = None
        
        # 创建必要的目录
        os.makedirs(self.data_dir, exist_ok=True)
        os.makedirs(self.db_dir, exist_ok=True)

    def load_and_process_documents(self):
        """加载并处理文档"""
        # 加载PDF文档
        loader = DirectoryLoader(
            path=self.data_dir,
            glob="*.pdf",
            loader_cls=PyPDFLoader
        )
        documents = loader.load()
        print(f"加载了 {len(documents)} 个PDF文档")
        
        # 分割文档
        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=500,
            chunk_overlap=50,
            separators=["\n\n", "\n", ". ", " ", ""]
        )
        splits = text_splitter.split_documents(documents)
        print(f"文档分割完成，得到 {len(splits)} 个片段")
        
        return splits

    def initialize_embeddings(self,embeddings_type : str):
        """初始化嵌入模型"""
        if embeddings_type == "HF":
            self.embeddings = HuggingFaceEmbeddings(
                model_name=self.embedding_model_name,
                model_kwargs={'device': 'cpu'},  # 可以改为'cuda'如果有GPU
                encode_kwargs={'normalize_embeddings': True}
            )
        if embeddings_type == "llama_cpp":
            self.embeddings =  LlamaCppEmbeddings(
                model_path=self.embedding_model_name,
                n_threads=4  # 根据 CPU 核心数调整（如 8 核心设为 4 或 8）
            )
        print(f"已加载嵌入模型: {self.embedding_model_name}")
        return self.embeddings

    def create_vector_database(self, documents, overwrite=False):
        """创建或加载向量数据库"""
        if not self.embeddings:
            self.initialize_embeddings()
            
        # 检查是否需要覆盖现有数据库
        if os.path.exists(self.db_dir) and not overwrite:
            print(f"加载现有向量数据库: {self.db_dir}")
            self.vector_db = Chroma(
                persist_directory=self.db_dir,
                embedding_function=self.embeddings
            )
        else:
            print(f"创建新的向量数据库: {self.db_dir}")
            self.vector_db = Chroma.from_documents(
                documents=documents,
                embedding=self.embeddings,
                persist_directory=self.db_dir
            )
            self.vector_db.persist()
        return self.vector_db

    def initialize_llm(self):
        """初始化本地LLM (使用Llama.cpp)"""
        # 检查模型文件是否存在
        if not os.path.exists(self.model_path):
            raise FileNotFoundError(f"模型文件未找到: {self.model_path}\n请下载GGUF格式的LLM模型并更新路径")
        
        # 回调管理器用于流式输出
        callback_manager = CallbackManager([StreamingStdOutCallbackHandler()])
        
        # 初始化Llama.cpp LLM
        self.llm = LlamaCpp(
            model_path=self.model_path,
            callback_manager=callback_manager,
            verbose=False,  # 设为True可查看调试信息
            n_ctx=2048,     # 上下文窗口大小
            n_threads=4,    # 线程数，根据CPU核心数调整
            n_gpu_layers=0  # GPU加速层数，设为0表示仅用CPU，有GPU可适当增加
        )
        print(f"已加载本地LLM模型: {os.path.basename(self.model_path)}")
        return self.llm

    def create_qa_chain(self):
        """创建检索增强问答链"""
        if not self.vector_db:
            raise ValueError("请先创建向量数据库")
            
        if not self.llm:
            self.initialize_llm()
        
        # 定义提示词模板
        prompt_template = """使用以下提供的上下文信息来回答用户的问题。
        如果你不知道答案，直接说不知道，不要编造。回答要简洁明了。
        
        上下文:
        {context}
        
        问题:
        {question}
        
        回答:
        """
        
        # 创建提示词
        prompt = PromptTemplate(
            template=prompt_template,
            input_variables=["context", "question"]
        )
        
        # 创建检索式问答链
        self.qa_chain = RetrievalQA.from_chain_type(
            llm=self.llm,
            chain_type="stuff",
            retriever=self.vector_db.as_retriever(
                search_kwargs={"k": 3}  # 检索最相关的3个文档片段
            ),
            chain_type_kwargs={"prompt": prompt},
            return_source_documents=True  # 返回用于回答的源文档
        )
        
        return self.qa_chain

    def answer_query(self, query):
        """回答用户查询"""
        if not self.qa_chain:
            self.create_qa_chain()
            
        print(f"\n查询: {query}")
        print("回答: ", end="")  # 保持与流式输出在同一行
        result = self.qa_chain({"query": query})
        
        # 输出源文档信息
        print("\n\n参考文档:")
        for i, doc in enumerate(result["source_documents"], 1):
            print(f"文档 {i}: {doc.metadata.get('source', '未知来源')}, 页码: {doc.metadata.get('page', '未知')}")
            
        return result

if __name__ == "__main__":
    # 初始化RAG系统
    rag_system = LocalRAGSystem(
        model_path="./models/llama-2-7b-chat.Q4_K_M.gguf",  # 替换为你的模型路径
        data_dir="./documents",                           # 替换为你的文档目录
        db_dir="./local_chroma_db"
    )
    
    # 处理文档并创建向量数据库 (首次运行需要，之后可注释掉)
    documents = rag_system.load_and_process_documents()
    rag_system.create_vector_database(documents)
    
    # 初始化LLM和问答链
    rag_system.initialize_llm()
    rag_system.create_qa_chain()
    
    # 示例查询
    while True:
        user_query = input("\n请输入你的问题 (输入 'exit' 退出): ")
        if user_query.lower() == 'exit':
            break
        rag_system.answer_query(user_query)
    