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

"""
查看RAG知识库内容工具
功能：查看分块结果、向量数据库内容、检索测试
"""

import os
from dotenv import load_dotenv
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_community.vectorstores import Chroma

# 加载环境变量
load_dotenv()


def view_vectorstore(persist_directory: str = "./car_info_knowledge_db"):
    """
    查看向量数据库内容
    
    Args:
        persist_directory: 向量数据库路径
    """
    print("=" * 80)
    print(f"查看向量数据库: {persist_directory}")
    print("=" * 80)
    
    if not os.path.exists(persist_directory):
        print(f"❌ 数据库不存在: {persist_directory}")
        return
    
    # 初始化Embedding模型
    embedding_model = DashScopeEmbeddings(
        model=os.getenv("EMBEDDING_MODEL", "text-embedding-v3"),
        dashscope_api_key=os.getenv("DASHSCOPE_API_KEY", "")
    )
    
    # 加载向量数据库
    vectorstore = Chroma(
        persist_directory=persist_directory,
        embedding_function=embedding_model
    )
    
    # 获取所有文档
    try:
        # 使用get方法获取所有数据
        result = vectorstore.get()
        
        documents = result.get('documents', [])
        metadatas = result.get('metadatas', [])
        ids = result.get('ids', [])
        
        print(f"\n📊 数据库统计信息:")
        print(f"   总文档数: {len(documents)}")
        
        if len(documents) == 0:
            print("\n❌ 数据库为空")
            return
        
        # 统计信息
        total_chars = sum(len(doc) for doc in documents)
        avg_chars = total_chars / len(documents) if len(documents) > 0 else 0
        
        print(f"   总字符数: {total_chars}")
        print(f"   平均每块字符数: {avg_chars:.1f}")
        print(f"   最大块: {max(len(doc) for doc in documents)} 字符")
        print(f"   最小块: {min(len(doc) for doc in documents)} 字符")
        
        # 显示前N个块
        print(f"\n📄 前10个块的内容预览:")
        print("-" * 80)
        
        for i in range(min(10, len(documents))):
            doc = documents[i]
            metadata = metadatas[i] if i < len(metadatas) else {}
            
            print(f"\n块 #{i+1} (ID: {ids[i] if i < len(ids) else 'N/A'})")
            print(f"页码: {metadata.get('page', 'N/A')}")
            print(f"字符数: {len(doc)}")
            print(f"内容预览:")
            print(f"  {doc[:200]}...")
            print("-" * 80)
        
    except Exception as e:
        print(f"❌ 读取数据库出错: {e}")
        import traceback
        traceback.print_exc()


def test_search(persist_directory: str = "./car_info_knowledge_db", query: str = None):
    """
    测试检索功能
    
    Args:
        persist_directory: 向量数据库路径
        query: 查询文本
    """
    print("\n" + "=" * 80)
    print("测试检索功能")
    print("=" * 80)
    
    if not os.path.exists(persist_directory):
        print(f"❌ 数据库不存在: {persist_directory}")
        return
    
    # 初始化Embedding模型
    embedding_model = DashScopeEmbeddings(
        model=os.getenv("EMBEDDING_MODEL", "text-embedding-v3"),
        dashscope_api_key=os.getenv("DASHSCOPE_API_KEY", "")
    )
    
    # 加载向量数据库
    vectorstore = Chroma(
        persist_directory=persist_directory,
        embedding_function=embedding_model
    )
    
    # 默认查询
    if not query:
        query = "HDC系统在什么速度下会激活？"
    
    print(f"\n🔍 查询: {query}")
    print("-" * 80)
    
    try:
        # 执行相似度搜索
        results = vectorstore.similarity_search_with_score(query, k=3)
        
        print(f"\n找到 {len(results)} 个相关结果:\n")
        
        for i, (doc, score) in enumerate(results):
            print(f"结果 #{i+1}")
            print(f"相似度分数: {score:.4f} (越小越相似)")
            print(f"页码: {doc.metadata.get('page', 'N/A')}")
            print(f"字符数: {len(doc.page_content)}")
            print(f"内容:")
            print(f"  {doc.page_content[:300]}...")
            print("-" * 80)
            
    except Exception as e:
        print(f"❌ 检索出错: {e}")
        import traceback
        traceback.print_exc()


def interactive_query(persist_directory: str = "./car_info_knowledge_db"):
    """
    交互式查询模式
    """
    print("\n" + "=" * 80)
    print("交互式查询模式")
    print("=" * 80)
    print("输入问题进行检索，输入 'quit' 或 'exit' 退出\n")
    
    if not os.path.exists(persist_directory):
        print(f"❌ 数据库不存在: {persist_directory}")
        return
    
    # 初始化Embedding模型
    embedding_model = DashScopeEmbeddings(
        model=os.getenv("EMBEDDING_MODEL", "text-embedding-v3"),
        dashscope_api_key=os.getenv("DASHSCOPE_API_KEY", "")
    )
    
    # 加载向量数据库
    vectorstore = Chroma(
        persist_directory=persist_directory,
        embedding_function=embedding_model
    )
    
    while True:
        try:
            query = input("你的问题: ").strip()
            
            if query.lower() in ['quit', 'exit', 'q']:
                print("\n再见！")
                break
            
            if not query:
                continue
            
            # 执行检索
            results = vectorstore.similarity_search_with_score(query, k=3)
            
            print(f"\n找到 {len(results)} 个相关结果:\n")
            
            for i, (doc, score) in enumerate(results):
                print(f"结果 #{i+1} (相似度: {score:.4f})")
                print(f"内容: {doc.page_content[:200]}...")
                print("-" * 60)
            
        except KeyboardInterrupt:
            print("\n\n再见！")
            break
        except Exception as e:
            print(f"错误: {e}")


def main():
    """主函数"""
    print("\n" + "=" * 80)
    print("RAG知识库内容查看工具")
    print("=" * 80)
    
    # 检查API Key
    if not os.getenv("DASHSCOPE_API_KEY"):
        print("❌ 请先配置 DASHSCOPE_API_KEY")
        return
    
    # 查看数据库
    db_path = "./car_info_knowledge_db"
    
    while True:
        print("\n请选择功能:")
        print("1. 查看向量数据库内容")
        print("2. 测试检索功能")
        print("3. 交互式查询")
        print("4. 更换数据库路径")
        print("0. 退出")
        
        choice = input("\n请输入选项 (0-4): ").strip()
        
        if choice == "1":
            view_vectorstore(db_path)
        elif choice == "2":
            query = input("输入查询内容 (直接回车使用默认查询): ").strip()
            test_search(db_path, query if query else None)
        elif choice == "3":
            interactive_query(db_path)
        elif choice == "4":
            new_path = input(f"输入数据库路径 (当前: {db_path}): ").strip()
            if new_path:
                db_path = new_path
        elif choice == "0":
            print("\n再见！")
            break
        else:
            print("无效选项")


if __name__ == "__main__":
    main()