check_db.py

from langchain_openai import OpenAIEmbeddings
from langchain_community.vectorstores import Chroma
from dotenv import load_dotenv

load_dotenv()

# 加载向量数据库
embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
vector_db = Chroma(persist_directory="./chroma_db", embedding_function=embeddings)

# 检查数据库内容
print("=== 检查向量数据库 ===")
print(f"总文档数: {vector_db._collection.count()}")

# 获取所有文档的元数据
all_docs = vector_db.get()
sources = set()
for metadata in all_docs["metadatas"]:
    if metadata and "source" in metadata:
        sources.add(metadata["source"])

print(f"\n发现的文档来源:")
for source in sorted(sources):
    print(f"  - {source}")

# 测试检索World Bank相关内容
print("\n=== 测试检索 'World Bank infrastructure' ===")
results = vector_db.similarity_search("World Bank infrastructure", k=5)
for i, doc in enumerate(results):
    source = doc.metadata.get("source", "Unknown")
    print(f"\n结果 {i + 1}: {source}")
    print(f"内容预览: {doc.page_content[:200]}...")