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

"""
直接显示向量数据库内容（非交互式）
"""

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

# 加载环境变量
load_dotenv()


def show_database(persist_directory: str):
    """显示数据库内容"""
    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
    )
    
    # 获取所有文档
    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)
    max_chars = max(len(doc) for doc in documents)
    min_chars = min(len(doc) for doc in documents)
    
    print(f"   总字符数: {total_chars:,}")
    print(f"   平均块大小: {avg_chars:.1f} 字符")
    print(f"   最大块: {max_chars} 字符")
    print(f"   最小块: {min_chars} 字符")
    
    # 显示前20个块
    print(f"\n📄 前20个块预览:")
    print("=" * 80)
    
    for i in range(min(20, len(documents))):
        doc = documents[i]
        metadata = metadatas[i] if i < len(metadatas) else {}
        
        print(f"\n【块 #{i+1}】")
        print(f"ID: {ids[i] if i < len(ids) else 'N/A'}")
        print(f"页码: {metadata.get('page', 'N/A')}")
        print(f"大小: {len(doc)} 字符")
        print(f"内容: {doc[:150].replace(chr(10), ' ')}...")
        print("-" * 80)


if __name__ == "__main__":
    # 检查API Key
    if not os.getenv("DASHSCOPE_API_KEY"):
        print("❌ 请先配置 DASHSCOPE_API_KEY")
        exit(1)
    
    # 显示数据库内容
    show_database("./car_info_knowledge_db")
    
    # 测试检索
    print("\n\n" + "=" * 80)
    print("测试检索功能")
    print("=" * 80)
    
    embedding_model = DashScopeEmbeddings(
        model=os.getenv("EMBEDDING_MODEL", "text-embedding-v3"),
        dashscope_api_key=os.getenv("DASHSCOPE_API_KEY", "")
    )
    
    vectorstore = Chroma(
        persist_directory="./car_info_knowledge_db",
        embedding_function=embedding_model
    )
    
    query = "HDC系统在什么速度下会激活？"
    print(f"\n查询: {query}")
    print("-" * 80)
    
    results = vectorstore.similarity_search_with_score(query, k=3)
    
    for i, (doc, score) in enumerate(results):
        print(f"\n结果 #{i+1}")
        print(f"相似度: {score:.4f}")
        print(f"页码: {doc.metadata.get('page', 'N/A')}")
        print(f"内容: {doc.page_content[:200]}...")
        print("-" * 80)