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

"""
快速开始示例 - RAG系统
演示如何使用RAG系统进行知识库问答
"""

import os
import sys
from dotenv import load_dotenv

# 添加当前目录到系统路径
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

# 导入RAG系统（使用importlib处理数字开头的文件名）
import importlib.util
spec = importlib.util.spec_from_file_location("rag_task", "./02_RAG_task.py")
rag_module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(rag_module)
RAGSystem = rag_module.RAGSystem

# 加载环境变量
load_dotenv()


def example_1_build_and_query():
    """示例1：构建知识库并查询"""
    print("\n=== 示例1：构建知识库并查询 ===\n")
    
    # 创建RAG系统
    rag = RAGSystem(
        pdf_path=r"D:\investment\疯狂的里海 · 投资方法论 — 基于 367 篇投资周记提炼.pdf",
        persist_directory="./investment_knowledge_db"
    )
    
    # 构建知识库（首次运行需要，后续可跳过）
    rag.build_knowledge_base()
    
    # 提问
    questions = [
        "什么是价值投资？",
        "投资的核心原则是什么？",
        "如何判断一个公司是否值得投资？"
    ]
    
    for question in questions:
        print(f"\n问题: {question}")
        answer = rag.query(question)
        print(f"回答: {answer}")
        print("-" * 60)


def example_2_load_and_query():
    """示例2：加载已有知识库并查询"""
    print("\n=== 示例2：加载已有知识库并查询 ===\n")
    
    # 创建RAG系统
    rag = RAGSystem(
        pdf_path=r"D:\investment\疯狂的里海 · 投资方法论 — 基于 367 篇投资周记提炼.pdf",
        persist_directory="./investment_knowledge_db"
    )
    
    # 加载已有向量数据库
    rag.load_vectorstore()
    rag.create_retriever()
    
    # 单次查询
    question = "投资者应该如何控制风险？"
    print(f"问题: {question}")
    
    answer = rag.query(question)
    print(f"回答: {answer}")


def example_3_custom_retrieval():
    """示例3：自定义检索参数"""
    print("\n=== 示例3：自定义检索参数 ===\n")
    
    # 创建RAG系统
    rag = RAGSystem(
        pdf_path=r"D:\investment\疯狂的里海 · 投资方法论 — 基于 367 篇投资周记提炼.pdf",
        persist_directory="./investment_knowledge_db"
    )
    
    # 加载知识库
    rag.load_vectorstore()
    
    # 自定义检索器：返回更多相关文档
    rag.create_retriever(k=5)
    
    # 检索相关文档
    question = "投资周记中提到了哪些投资策略？"
    relevant_docs = rag.retrieve(question, k=5)
    
    print(f"问题: {question}")
    print(f"\n找到 {len(relevant_docs)} 个相关文档片段:\n")
    
    for i, doc in enumerate(relevant_docs, 1):
        print(f"--- 文档片段 {i} ---")
        print(f"内容: {doc.page_content[:200]}...")
        print(f"来源: {doc.metadata}")
        print()


def example_4_step_by_step():
    """示例4：分步骤执行RAG流程"""
    print("\n=== 示例4：分步骤执行RAG流程 ===\n")
    
    # 创建RAG系统
    rag = RAGSystem(
        pdf_path=r"D:\investment\疯狂的里海 · 投资方法论 — 基于 367 篇投资周记提炼.pdf",
        persist_directory="./custom_knowledge_db"
    )
    
    # 步骤1：加载PDF
    print("步骤1: 加载PDF文档...")
    pages = rag.load_pdf()
    print(f"  加载了 {len(pages)} 页\n")
    
    # 步骤2：清洗数据
    print("步骤2: 清洗数据...")
    rag.clean_all_pages()
    print("  数据清洗完成\n")
    
    # 步骤3：分块（自定义参数）
    print("步骤3: 文档分块...")
    docs = rag.split_documents(chunk_size=1000, chunk_overlap=200)
    print(f"  分成 {len(docs)} 个块\n")
    
    # 步骤4：创建向量库
    print("步骤4: 创建向量数据库...")
    vectorstore = rag.create_vectorstore()
    print("  向量数据库创建完成\n")
    
    # 步骤5：创建检索器
    print("步骤5: 创建检索器...")
    retriever = rag.create_retriever(k=3)
    print("  检索器创建完成\n")
    
    # 步骤6：检索
    print("步骤6: 检索相关文档...")
    question = "什么是安全边际？"
    relevant_docs = rag.retrieve(question)
    print(f"  找到 {len(relevant_docs)} 个相关文档\n")
    
    # 步骤7：生成答案
    print("步骤7: 生成答案...")
    answer = rag.generate_answer(question, relevant_docs)
    print(f"问题: {question}")
    print(f"回答: {answer}")


def main():
    """主函数"""
    # 检查API Key
    if not os.getenv("DASHSCOPE_API_KEY") or os.getenv("DASHSCOPE_API_KEY") == "your_dashscope_api_key_here":
        print("错误: 请先在 .env 文件中配置 DASHSCOPE_API_KEY")
        print("获取API Key: https://dashscope.console.aliyun.com/")
        return
    
    print("=" * 60)
    print("RAG系统快速开始示例")
    print("=" * 60)
    
    # 选择要运行的示例
    print("\n请选择示例:")
    print("1. 构建知识库并查询（首次运行）")
    print("2. 加载已有知识库并查询")
    print("3. 自定义检索参数")
    print("4. 分步骤执行RAG流程")
    print("0. 退出")
    
    choice = input("\n请输入选项 (0-4): ").strip()
    
    if choice == "1":
        example_1_build_and_query()
    elif choice == "2":
        example_2_load_and_query()
    elif choice == "3":
        example_3_custom_retrieval()
    elif choice == "4":
        example_4_step_by_step()
    elif choice == "0":
        print("退出程序")
    else:
        print("无效选项")


if __name__ == "__main__":
    main()