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

"""简单查看数据库内容"""
import os
import sys
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()

# 切换到正确的工作目录
os.chdir(r"d:\agentlearning\lqq-agent-study\02_RAG_study")

from langchain_community.embeddings import DashScopeEmbeddings
from langchain_community.vectorstores import Chroma

# 初始化
embedding_model = DashScopeEmbeddings(
    model="text-embedding-v3",
    dashscope_api_key=os.getenv("DASHSCOPE_API_KEY", "")
)

# 加载数据库
db_path = "./car_info_knowledge_db"
print(f"加载数据库: {db_path}")

if not os.path.exists(db_path):
    print(f"数据库不存在！")
    sys.exit(1)

vectorstore = Chroma(
    persist_directory=db_path,
    embedding_function=embedding_model
)

# 获取所有数据
result = vectorstore.get()
documents = result.get('documents', [])
metadatas = result.get('metadatas', [])

print(f"\n总块数: {len(documents)}")

if documents:
    # 统计
    sizes = [len(doc) for doc in documents]
    print(f"总字符数: {sum(sizes):,}")
    print(f"平均块大小: {sum(sizes)/len(sizes):.1f}")
    print(f"最大块: {max(sizes)}")
    print(f"最小块: {min(sizes)}")
    
    # 显示前10个块
    print("\n" + "=" * 80)
    print("前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}】")
        print(f"页码: {metadata.get('page', 'N/A')}")
        print(f"大小: {len(doc)} 字符")
        print(f"内容: {doc[:150].replace(chr(10), ' ')}...")
        print("-" * 80)
    
    # 测试检索
    print("\n\n测试检索:")
    query = "HDC系统"
    print(f"查询: {query}")
    
    results = vectorstore.similarity_search_with_score(query, k=3)
    
    for i, (doc, score) in enumerate(results):
        print(f"\n结果 #{i+1} (相似度: {score:.4f})")
        print(f"内容: {doc.page_content[:150]}...")