118 lines
4.7 KiB
Python
118 lines
4.7 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
用于检查chromadb数据库中所有文档的metadata字段
|
||
运行方式：项目根目录下执行
|
||
python utils/check_metadata.py
|
||
"""
|
||
import os
|
||
import sys
|
||
# 将上一级目录添加到sys.path
|
||
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||
from rag.vector_store import VectorStoreManager
|
||
|
||
|
||
def check_metadata():
|
||
"""检查chromadb数据库中所有文档的metadata字段"""
|
||
try:
|
||
# 初始化VectorStoreManager
|
||
print("正在初始化VectorStoreManager...")
|
||
vector_store_manager = VectorStoreManager()
|
||
|
||
# 获取所有文档的metadata
|
||
print("正在获取所有文档的metadata...")
|
||
results = vector_store_manager.collections['non_code'].get(include=['metadatas', 'documents'])
|
||
print(results)
|
||
input()
|
||
# print(vector_store_manager)
|
||
results = vector_store_manager.collections['code'].get(include=['metadatas', 'documents'])
|
||
|
||
processed_shas = {}
|
||
for metadata in results.get('metadatas', []):
|
||
if metadata and 'file_path' in metadata and 'file_blob_sha' in metadata:
|
||
file_path = metadata['file_path']
|
||
blob_sha = metadata['file_blob_sha']
|
||
processed_shas[file_path] = blob_sha
|
||
print(f"已处理文件SHA: {file_path} -> {blob_sha}")
|
||
# input("按任意键继续...")
|
||
collection = vector_store_manager.collections['code']
|
||
target_source = 'git_server_172_26_120_125'
|
||
result = collection.get(where={"repo_id": {"$contains": target_source}})
|
||
print(f"{target_source}查询结果: {str(result)[:50]}")
|
||
target_source = 'git_server_172_26_120_125_testrepo'
|
||
result = collection.get(where={"repo_id": {"$contains": target_source}})
|
||
print(f"{target_source}查询结果: {str(result)[:50]}")
|
||
result = collection.get(where={"repo_id": target_source})
|
||
print(f"{target_source}查询结果: {str(result)[:50]}")
|
||
exit(0)
|
||
|
||
# 提取数据
|
||
ids = results.get('ids', [])
|
||
metadatas = results.get('metadatas', [])
|
||
documents = results.get('documents', [])
|
||
|
||
print(f"共找到 {len(ids)} 个文档")
|
||
print("\n检查metadata中的file_path字段：")
|
||
print("-" * 80)
|
||
|
||
# 统计信息
|
||
total_docs = len(ids)
|
||
hello_docs = 0
|
||
hello_copy_docs = 0
|
||
hello_copy_copy_docs = 0
|
||
other_docs = 0
|
||
|
||
# 检查每个文档的metadata
|
||
for i, (doc_id, metadata, document) in enumerate(zip(ids, metadatas, documents)):
|
||
if metadata:
|
||
file_path = metadata.get('file_path', 'N/A')
|
||
func_name = metadata.get('func_name', 'N/A')
|
||
doc_id_meta = metadata.get('doc_id', 'N/A')
|
||
chunk_id = metadata.get('chunk_id', 'N/A')
|
||
|
||
if "check_metadata" in file_path:
|
||
print(f"文档ID: {doc_id}")
|
||
print(f" file_path: {file_path}")
|
||
print(f" func_name: {func_name}")
|
||
print(f" doc_id: {doc_id_meta}")
|
||
print(f" chunk_id: {chunk_id}")
|
||
for key, value in metadata.items():
|
||
if "mysql_sync" in str(value):
|
||
print(f" {key}: {value}")
|
||
print("+++++")
|
||
hello_copy_copy_docs += 1
|
||
|
||
# 统计file_path中的hello和hello_copy
|
||
# if 'hello' in file_path and 'hello_copy' not in file_path:
|
||
# hello_docs += 1
|
||
# print(" 状态: 仍为hello，未更新")
|
||
# elif 'hello_copy' in file_path:
|
||
# hello_copy_docs += 1
|
||
# print(" 状态: 已更新为hello_copy")
|
||
# else:
|
||
# other_docs += 1
|
||
# print(" 状态: 不包含hello或hello_copy")
|
||
else:
|
||
print(f"文档ID: {doc_id}")
|
||
print(" 无metadata")
|
||
other_docs += 1
|
||
|
||
# print("-" * 80)
|
||
|
||
# 打印统计结果
|
||
print("\n统计结果：")
|
||
print(f"总文档数: {total_docs}")
|
||
print(f"包含'hello_copy_copy'的文档数: {hello_copy_copy_docs}")
|
||
|
||
if hello_docs > 0:
|
||
print("\n警告: 仍有文档的file_path包含'hello'，未更新为'hello_copy'")
|
||
else:
|
||
print("\n所有文档的file_path已更新为'hello_copy'")
|
||
|
||
except Exception as e:
|
||
print(f"错误: {e}")
|
||
sys.exit(1)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
check_metadata()
|