lqq-agent-study/02_RAG_study/02_RAG_task.py

"""
RAG系统实战 - 投资方法论知识库
功能：加载PDF文档、数据清洗、分块、向量化、存储到向量库、检索生成回答
"""

import os
import re
from typing import List
from dotenv import load_dotenv

# LangChain 核心组件
from langchain_community.document_loaders import PyMuPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_community.vectorstores import Chroma
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from langchain_core.output_parsers import StrOutputParser
from langchain_core.documents import Document

# 加载环境变量
load_dotenv()


class RAGSystem:
    """RAG系统类，封装完整的检索增强生成流程"""
    
    def __init__(self, pdf_path: str, persist_directory: str = "./knowledge_db"):
        """
        初始化RAG系统
        
        Args:
            pdf_path: PDF文件路径
            persist_directory: 向量数据库持久化目录
        """
        self.pdf_path = pdf_path
        self.persist_directory = persist_directory
        self.pages = None
        self.split_docs = None
        self.vectorstore = None
        self.retriever = None
        
        # 初始化Embedding模型（使用阿里通义）
        self.embedding_model = DashScopeEmbeddings(
            model=os.getenv("EMBEDDING_MODEL", "text-embedding-v3"),
            dashscope_api_key=os.getenv("DASHSCOPE_API_KEY", "")
        )
        
        # 初始化大模型（使用DeepSeek）
        self.llm = ChatOpenAI(
            model=os.getenv("MODEL_NAME", "deepseek-chat"),
            openai_api_key=os.getenv("OPENAI_API_KEY", ""),
            openai_api_base=os.getenv("OPENAI_API_BASE", "https://api.deepseek.com/v1"),
            temperature=0.7
        )
        
        log_info("RAG系统初始化完成")
        log_info(f"使用模型: {os.getenv('MODEL_NAME', 'deepseek-chat')}")
        log_info(f"API地址: {os.getenv('OPENAI_API_BASE', 'https://api.deepseek.com/v1')}")
    
    def load_pdf(self) -> List[Document]:
        """
        步骤1：加载PDF文档
        
        Returns:
            文档页面列表
        """
        log_info(f"正在加载PDF文档: {self.pdf_path}")
        
        # 检查文件是否存在
        if not os.path.exists(self.pdf_path):
            raise FileNotFoundError(f"PDF文件不存在: {self.pdf_path}")
        
        # 使用PyMuPDF加载PDF
        loader = PyMuPDFLoader(self.pdf_path)
        self.pages = loader.load()
        
        log_info(f"PDF加载完成，共 {len(self.pages)} 页")
        
        # 检查是否有实际内容
        total_content = sum(len(page.page_content) for page in self.pages)
        
        if total_content == 0:
            log_error("PDF文件没有可提取的文本内容！")
            log_error("可能原因：")
            log_error("  1. PDF是扫描版或图片型PDF")
            log_error("  2. PDF有加密保护")
            log_error("  3. PDF格式特殊")
            log_error("\n解决方案：")
            log_error("  1. 使用OCR工具将PDF转换为文本")
            log_error("  2. 尝试其他PDF文件")
            raise ValueError("PDF文件无可提取文本内容，请使用文本型PDF或进行OCR处理")
        
        # 打印统计信息
        non_empty_pages = sum(1 for page in self.pages if len(page.page_content) > 0)
        log_info(f"有内容的页数: {non_empty_pages}/{len(self.pages)}")
        log_info(f"总字符数: {total_content}")
        
        # 打印第一页预览
        if self.pages and len(self.pages[0].page_content) > 0:
            preview = self.pages[0].page_content[:500]
            log_info(f"第一页内容预览:\n{preview}...")
        
        return self.pages
    
    def clean_text(self, text: str) -> str:
        """
        步骤2：清洗PDF解析出的文本，去除常见噪声
        
        Args:
            text: 原始文本
            
        Returns:
            清洗后的文本
        """
        if not text or not text.strip():
            return ""
        
        # 保存原始长度用于对比
        original_len = len(text)
        
        # 删除连续的换行符（保留最多2个）
        text = re.sub(r'\n{3,}', '\n\n', text)
        
        # 删除项目符号
        text = text.replace('•', '').replace('·', '')
        
        # 合并多余空格（但不删除换行符周围的空格）
        text = re.sub(r'[^\S\n]+', ' ', text)
        
        # 删除特殊控制字符
        text = re.sub(r'[\x00-\x08\x0b\x0c\x0e-\x1f\x7f-\x9f]', '', text)
        
        # 如果清洗后文本为空或太短，返回原始文本
        cleaned_text = text.strip()
        if len(cleaned_text) < 10:  # 如果清洗后少于10个字符，返回原始文本
            log_info(f"警告: 清洗后文本过短({len(cleaned_text)}字符)，保留原始文本({original_len}字符)")
            return text.strip()
        
        return cleaned_text
    
    def clean_all_pages(self) -> List[Document]:
        """
        清洗所有页面
        
        Returns:
            清洗后的文档列表
        """
        if not self.pages:
            raise ValueError("请先调用 load_pdf() 加载文档")
        
        log_info("开始清洗文档...")
        
        # 清洗每一页的内容
        for page in self.pages:
            page.page_content = self.clean_text(page.page_content)
        
        log_info(f"文档清洗完成，共处理 {len(self.pages)} 页")
        
        return self.pages
    
    def split_documents(
        self, 
        chunk_size: int = 100, 
        chunk_overlap: int = 20
    ) -> List[Document]:
        """
        步骤3：文档分块（使用递归字符分割器）
        
        Args:
            chunk_size: 每个块的最大字符数
            chunk_overlap: 相邻块之间的重叠字符数
            
        Returns:
            分块后的文档列表
        """
        if not self.pages:
            raise ValueError("请先调用 load_pdf() 加载文档")
        
        log_info(f"开始分块，chunk_size={chunk_size}, chunk_overlap={chunk_overlap}")
        
        # 创建递归字符分割器
        text_splitter = RecursiveCharacterTextSplitter(
            # 分隔符优先级：段落 → 换行 → 句号 → 空格 → 硬切
            separators=["\n\n", "\n", "。", "！", "？", "；", " ", ""],
            chunk_size=chunk_size,
            chunk_overlap=chunk_overlap,
            length_function=len,
            is_separator_regex=False
        )
        
        # 执行分块
        self.split_docs = text_splitter.split_documents(self.pages)
        
        # 过滤空块
        self.split_docs = [
            doc for doc in self.split_docs 
            if doc.page_content and doc.page_content.strip()
        ]
        
        # 统计信息
        total_chars = sum(len(doc.page_content) for doc in self.split_docs)
        log_info(f"分块完成，共 {len(self.split_docs)} 个块，总字符数: {total_chars}")
        
        # 打印前3个块预览
        for i, doc in enumerate(self.split_docs[:3]):
            log_info(f"块{i+1} 预览: {doc.page_content[:100]}...")
        
        return self.split_docs
    
    def create_vectorstore(self) -> Chroma:
        """
        步骤4：创建向量数据库并存储文档
        
        Returns:
            向量数据库实例
        """
        if not self.split_docs:
            raise ValueError("请先调用 split_documents() 分块文档")
        
        log_info("开始创建向量数据库...")
        
        # 创建向量数据库并持久化
        self.vectorstore = Chroma.from_documents(
            documents=self.split_docs,
            embedding=self.embedding_model,
            collection_metadata={"hnsw:space": "cosine"},  # 使用余弦相似度
            persist_directory=self.persist_directory
        )
        
        log_info(f"向量数据库创建完成，保存至: {self.persist_directory}")
        
        return self.vectorstore
    
    def load_vectorstore(self) -> Chroma:
        """
        加载已存在的向量数据库
        
        Returns:
            向量数据库实例
        """
        if not os.path.exists(self.persist_directory):
            raise FileNotFoundError(f"向量数据库不存在: {self.persist_directory}")
        
        log_info(f"加载向量数据库: {self.persist_directory}")
        
        self.vectorstore = Chroma(
            persist_directory=self.persist_directory,
            embedding_function=self.embedding_model
        )
        
        log_info("向量数据库加载完成")
        
        return self.vectorstore
    
    def create_retriever(self, k: int = 3):
        """
        步骤5：创建检索器
        
        Args:
            k: 返回的最相关文档数量
            
        Returns:
            检索器实例
        """
        if not self.vectorstore:
            raise ValueError("请先创建或加载向量数据库")
        
        log_info(f"创建检索器，返回Top-{k}文档")
        
        self.retriever = self.vectorstore.as_retriever(
            search_kwargs={"k": k}
        )
        
        return self.retriever
    
    def retrieve(self, query: str, k: int = 3) -> List[Document]:
        """
        步骤6：检索相关文档
        
        Args:
            query: 用户问题
            k: 返回的文档数量
            
        Returns:
            相关文档列表
        """
        if not self.retriever:
            raise ValueError("请先创建检索器")
        
        log_info(f"检索问题: {query}")
        
        # 执行检索
        relevant_docs = self.retriever.invoke(query)
        
        log_info(f"检索完成，找到 {len(relevant_docs)} 个相关文档")
        
        # 打印检索结果
        for i, doc in enumerate(relevant_docs):
            log_info(f"文档{i+1}: {doc.page_content[:150]}...")
        
        return relevant_docs
    
    def generate_answer(self, query: str, relevant_docs: List[Document]) -> str:
        """
        步骤7：基于检索结果生成回答
        
        Args:
            query: 用户问题
            relevant_docs: 相关文档列表
            
        Returns:
            生成的回答
        """
        log_info("开始生成回答...")
        
        # 拼接上下文
        context = "\n\n---\n\n".join([doc.page_content for doc in relevant_docs])
        
        # 构建Prompt模板
        prompt = ChatPromptTemplate.from_template("""
你是一个专业的投资知识库助手。请根据以下检索到的上下文回答用户问题。

**规则：**
- 只基于提供的上下文回答，不要编造
- 如果上下文中没有相关信息，直接说「根据现有资料，我找不到这个问题的答案」
- 回答要简洁直接，引用原文时用引号
- 回答时请标注信息来源的页码

**检索到的上下文：**
{context}

**用户问题：**
{question}
""")
        
        # 构建Chain
        chain = prompt | self.llm | StrOutputParser()
        
        # 生成回答
        answer = chain.invoke({
            "context": context,
            "question": query
        })
        
        log_info("回答生成完成")
        
        return answer
    
    def query(self, question: str, k: int = 3) -> str:
        """
        完整查询流程：检索 + 生成
        
        Args:
            question: 用户问题
            k: 检索文档数量
            
        Returns:
            生成的回答
        """
        # 检索相关文档
        relevant_docs = self.retrieve(question, k)
        
        # 生成回答
        answer = self.generate_answer(question, relevant_docs)
        
        return answer
    
    def build_knowledge_base(self):
        """
        构建完整的知识库：加载 → 清洗 → 分块 → 向量化 → 存储
        """
        log_info("=" * 50)
        log_info("开始构建知识库...")
        log_info("=" * 50)
        
        # 步骤1：加载PDF
        self.load_pdf()
        
        # 步骤2：清洗数据
        self.clean_all_pages()
        
        # 步骤3：分块
        self.split_documents()
        
        # 步骤4：向量化并存储
        self.create_vectorstore()
        
        # 步骤5：创建检索器
        self.create_retriever()
        
        log_info("=" * 50)
        log_info("知识库构建完成！")
        log_info("=" * 50)


def log_info(message: str):
    """打印日志信息"""
    print(f"[INFO] {message}")


def interactive_query(rag_system: RAGSystem):
    """
    交互式问答模式
    
    Args:
        rag_system: RAG系统实例
    """
    print("\n" + "=" * 60)
    print("RAG知识库问答系统 - 投资方法论")
    print("=" * 60)
    print("输入问题开始查询，输入 'quit' 或 'exit' 退出\n")
    
    while True:
        try:
            # 获取用户输入
            question = input("你的问题: ").strip()
            
            # 检查退出命令
            if question.lower() in ['quit', 'exit', 'q']:
                print("\n感谢使用，再见！")
                break
            
            # 跳过空问题
            if not question:
                print("请输入有效问题\n")
                continue
            
            # 执行查询
            print("\n正在检索并生成答案...\n")
            answer = rag_system.query(question)
            
            # 显示结果
            print("-" * 60)
            print(f"回答:\n{answer}")
            print("-" * 60)
            print()
            
        except KeyboardInterrupt:
            print("\n\n感谢使用，再见！")
            break
        except Exception as e:
            log_error(f"查询出错: {str(e)}")


def log_error(message: str):
    """打印错误日志"""
    print(f"[ERROR] {message}")


def main():
    """主函数"""
    # PDF文件路径
    pdf_path = r"D:\investment\data\疯狂的里海_投资体系框架.pdf"
    
    # 备用PDF文件路径（用于测试）
    backup_pdf_path = r"./car_info.pdf"
    
    # 向量数据库存储路径
    persist_directory = "D:\agentlearning\lqq-agent-study\investment_db"
    
    # 检查API Key配置
    openai_api_key = os.getenv("OPENAI_API_KEY")
    dashscope_api_key = os.getenv("DASHSCOPE_API_KEY")
    
    if not openai_api_key:
        print("警告: 未设置 OPENAI_API_KEY 环境变量")
        print("请在 .env 文件中添加: OPENAI_API_KEY=your-deepseek-api-key")
        print("或者访问 https://platform.deepseek.com/ 获取API Key")
        return
    
    if not dashscope_api_key:
        print("警告: 未设置 DASHSCOPE_API_KEY 环境变量")
        print("请在 .env 文件中添加: DASHSCOPE_API_KEY=your-dashscope-api-key")
        print("或者访问 https://dashscope.console.aliyun.com/ 获取API Key")
        return
    
    try:
        # 检查主PDF文件是否存在
        if not os.path.exists(pdf_path):
            print(f"警告: 指定的PDF文件不存在: {pdf_path}")
            
            # 尝试使用备用PDF文件
            if os.path.exists(backup_pdf_path):
                print(f"将使用备用PDF文件进行测试: {backup_pdf_path}")
                pdf_path = backup_pdf_path
                persist_directory = "./car_info_knowledge_db"
            else:
                print("备用PDF文件也不存在，请检查文件路径")
                return
        
        # 创建RAG系统实例
        rag = RAGSystem(
            pdf_path=pdf_path,
            persist_directory=persist_directory
        )
        
        # 检查向量数据库是否已存在
        if os.path.exists(persist_directory):
            print(f"检测到已有向量数据库: {persist_directory}")
            choice = input("是否重新构建知识库？(y/n): ").strip().lower()
            
            if choice == 'y':
                # 重新构建知识库
                rag.build_knowledge_base()
            else:
                # 加载已有向量数据库
                rag.load_vectorstore()
                rag.create_retriever()
        else:
            # 构建新知识库
            rag.build_knowledge_base()
        
        # 进入交互式问答
        interactive_query(rag)
        
    except ValueError as e:
        if "PDF文件无可提取文本内容" in str(e):
            print("\n" + "=" * 60)
            print("PDF文件处理失败")
            print("=" * 60)
            print("\n您的PDF文件是扫描版或图片型PDF，无法直接提取文本。")
            print("\n解决方案：")
            print("1. 使用OCR工具将PDF转换为文本")
            print("   推荐工具：Adobe Acrobat、福昕PDF、ABBYY FineReader")
            print("   在线工具：https://www.pdf2go.com/zh/ocr-pdf")
            print("\n2. 使用项目中的示例PDF文件进行测试:")
            print(f"   文件路径: {backup_pdf_path}")
            print("\n3. 寻找其他文本型PDF文件")
            print("\n提示：您可以修改代码中的pdf_path变量指向可用的PDF文件")
        else:
            log_error(str(e))
    except FileNotFoundError as e:
        log_error(str(e))
        log_error("请检查PDF文件路径是否正确")
    except Exception as e:
        log_error(f"系统运行出错: {str(e)}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    main()