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

"""
RAG系统测试脚本
用于验证系统各个组件是否正常工作
"""

import os
import sys
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()


def test_api_key():
    """测试API Key配置"""
    print("\n=== 测试1: API Key配置 ===")
    
    # 检查DeepSeek API Key
    openai_api_key = os.getenv("OPENAI_API_KEY")
    dashscope_api_key = os.getenv("DASHSCOPE_API_KEY")
    
    openai_ok = True
    dashscope_ok = True
    
    if not openai_api_key:
        print("❌ 未配置 OPENAI_API_KEY (DeepSeek)")
        print("   请在 .env 文件中添加: OPENAI_API_KEY=your-deepseek-api-key")
        openai_ok = False
    else:
        print(f"✅ OPENAI_API_KEY 已配置 (长度: {len(openai_api_key)})")
    
    if not dashscope_api_key:
        print("❌ 未配置 DASHSCOPE_API_KEY (用于Embedding)")
        print("   请在 .env 文件中添加: DASHSCOPE_API_KEY=your-dashscope-api-key")
        dashscope_ok = False
    else:
        print(f"✅ DASHSCOPE_API_KEY 已配置 (长度: {len(dashscope_api_key)})")
    
    # 打印模型配置
    model_name = os.getenv("MODEL_NAME", "deepseek-chat")
    api_base = os.getenv("OPENAI_API_BASE", "https://api.deepseek.com/v1")
    embedding_model = os.getenv("EMBEDDING_MODEL", "text-embedding-v3")
    
    print(f"\n模型配置:")
    print(f"   LLM模型: {model_name}")
    print(f"   API地址: {api_base}")
    print(f"   Embedding模型: {embedding_model}")
    
    return openai_ok and dashscope_ok


def test_pdf_file():
    """测试PDF文件是否存在"""
    print("\n=== 测试2: PDF文件检查 ===")
    
    pdf_path = r"D:\investment\疯狂的里海 · 投资方法论 — 基于 367 篇投资周记提炼.pdf"
    
    if os.path.exists(pdf_path):
        file_size = os.path.getsize(pdf_path) / (1024 * 1024)  # MB
        print(f"✅ PDF文件存在")
        print(f"   路径: {pdf_path}")
        print(f"   大小: {file_size:.2f} MB")
        return True
    else:
        print(f"❌ PDF文件不存在")
        print(f"   路径: {pdf_path}")
        print("   请检查文件路径是否正确")
        return False


def test_dependencies():
    """测试依赖包是否安装"""
    print("\n=== 测试3: 依赖包检查 ===")
    
    required_packages = {
        'langchain': 'langchain',
        'langchain_community': 'langchain-community',
        'langchain_core': 'langchain-core',
        'langchain_text_splitters': 'langchain-text-splitters',
        'fitz': 'pymupdf',  # PyMuPDF
        'chromadb': 'chromadb',
        'dashscope': 'dashscope',
    }
    
    all_installed = True
    
    for module_name, package_name in required_packages.items():
        try:
            __import__(module_name)
            print(f"✅ {package_name} 已安装")
        except ImportError:
            print(f"❌ {package_name} 未安装")
            print(f"   安装命令: pip install {package_name}")
            all_installed = False
    
    return all_installed


def test_embedding_model():
    """测试Embedding模型"""
    print("\n=== 测试4: Embedding模型 ===")
    
    try:
        from langchain_community.embeddings import DashScopeEmbeddings
        
        api_key = os.getenv("DASHSCOPE_API_KEY")
        if not api_key:
            print("⚠️ 跳过测试（API Key未配置）")
            return None
        
        # 创建Embedding模型
        embedding_model = DashScopeEmbeddings(
            model=os.getenv("EMBEDDING_MODEL", "text-embedding-v3"),
            dashscope_api_key=api_key
        )
        
        # 测试向量化
        test_text = "这是一个测试文本"
        embedding = embedding_model.embed_query(test_text)
        
        print(f"✅ Embedding模型正常工作")
        print(f"   使用模型: {os.getenv('EMBEDDING_MODEL', 'text-embedding-v3')}")
        print(f"   向量维度: {len(embedding)}")
        print(f"   前5个值: {embedding[:5]}")
        return True
        
    except Exception as e:
        print(f"❌ Embedding模型测试失败")
        print(f"   错误: {str(e)}")
        return False


def test_llm_model():
    """测试大语言模型"""
    print("\n=== 测试5: 大语言模型 ===")
    
    try:
        from langchain_openai import ChatOpenAI
        
        api_key = os.getenv("OPENAI_API_KEY")
        api_base = os.getenv("OPENAI_API_BASE", "https://api.deepseek.com/v1")
        model_name = os.getenv("MODEL_NAME", "deepseek-chat")
        
        if not api_key:
            print("⚠️ 跳过测试（API Key未配置）")
            return None
        
        # 创建LLM（使用DeepSeek）
        llm = ChatOpenAI(
            model=model_name,
            openai_api_key=api_key,
            openai_api_base=api_base,
            temperature=0.7
        )
        
        # 测试调用
        test_message = "你好，请回复'测试成功'"
        response = llm.invoke(test_message)
        
        print(f"✅ 大语言模型正常工作")
        print(f"   使用模型: {model_name}")
        print(f"   API地址: {api_base}")
        print(f"   测试问题: {test_message}")
        print(f"   模型回复: {response.content}")
        return True
        
    except Exception as e:
        print(f"❌ 大语言模型测试失败")
        print(f"   错误: {str(e)}")
        return False


def test_rag_system():
    """测试完整的RAG系统"""
    print("\n=== 测试6: RAG系统完整性 ===")
    
    try:
        # 添加路径
        sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
        
        # 导入RAG系统
        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
        
        print("✅ RAG系统模块加载成功")
        return True
        
    except Exception as e:
        print(f"❌ RAG系统测试失败")
        print(f"   错误: {str(e)}")
        return False


def main():
    """运行所有测试"""
    print("=" * 60)
    print("RAG系统诊断测试")
    print("=" * 60)
    
    # 运行测试
    results = {
        "API Key配置": test_api_key(),
        "PDF文件检查": test_pdf_file(),
        "依赖包检查": test_dependencies(),
    }
    
    # 如果API Key配置正确，运行额外测试
    if results["API Key配置"]:
        results["Embedding模型"] = test_embedding_model()
        results["大语言模型"] = test_llm_model()
    else:
        results["Embedding模型"] = None
        results["大语言模型"] = None
    
    results["RAG系统"] = test_rag_system()
    
    # 总结
    print("\n" + "=" * 60)
    print("测试总结")
    print("=" * 60)
    
    passed = sum(1 for v in results.values() if v is True)
    failed = sum(1 for v in results.values() if v is False)
    skipped = sum(1 for v in results.values() if v is None)
    
    for test_name, result in results.items():
        if result is True:
            status = "✅ 通过"
        elif result is False:
            status = "❌ 失败"
        else:
            status = "⚠️ 跳过"
        print(f"{test_name}: {status}")
    
    print(f"\n总计: {passed} 通过, {failed} 失败, {skipped} 跳过")
    
    if failed == 0 and passed > 0:
        print("\n🎉 所有测试通过！系统准备就绪。")
        print("\n运行以下命令启动系统:")
        print("  python 02_RAG_task.py")
    elif failed > 0:
        print("\n⚠️ 部分测试失败，请检查上述错误信息。")
    else:
        print("\n⚠️ 请先配置 API Key 和安装依赖包。")


if __name__ == "__main__":
    main()