d67a8ff50b
- 修正 preview 接口路径（去掉多余的 /api 前缀） - 修正响应示例中 chunk id 格式（文件名_序号，非 UUID） - 补充 preview 响应中 status、version 字段 - 明确前端引用匹配方式为 chunk_id 查找而非数组下标 - fix_image_paths.py 改为遍历知识库子目录，避免生成空 chroma.sqlite3
140 lines
4.6 KiB
Python
140 lines
4.6 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""
|
|
修复向量库中的图片路径格式
|
|
|
|
问题：历史数据中 image_path 存储为 "images/xxx.jpg" 格式
|
|
修复：统一改为只存储文件名 "xxx.jpg"
|
|
|
|
使用方式：
|
|
python scripts/fix_image_paths.py [--dry-run]
|
|
"""
|
|
|
|
import argparse
|
|
import os
|
|
import sys
|
|
|
|
# 添加项目根目录到路径
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
import chromadb
|
|
from config import CHROMA_DB_PATH
|
|
|
|
|
|
def fix_image_paths(dry_run: bool = False):
|
|
"""
|
|
修复向量库中的图片路径格式
|
|
|
|
多向量库模式下，遍历 chroma/ 下的每个知识库子目录，
|
|
分别创建 PersistentClient 处理各自的集合。
|
|
|
|
Args:
|
|
dry_run: 仅预览，不实际修改
|
|
"""
|
|
# 扫描知识库子目录（每个子目录是一个独立的向量库）
|
|
kb_dirs = []
|
|
if os.path.exists(CHROMA_DB_PATH):
|
|
for item in os.listdir(CHROMA_DB_PATH):
|
|
item_path = os.path.join(CHROMA_DB_PATH, item)
|
|
if os.path.isdir(item_path) and not item.startswith('.'):
|
|
if os.path.exists(os.path.join(item_path, 'chroma.sqlite3')):
|
|
kb_dirs.append(item)
|
|
|
|
if not kb_dirs:
|
|
print("未找到任何知识库子目录")
|
|
return
|
|
|
|
print(f"发现 {len(kb_dirs)} 个知识库: {kb_dirs}\n")
|
|
|
|
for kb_name in kb_dirs:
|
|
kb_path = os.path.join(CHROMA_DB_PATH, kb_name)
|
|
print(f"处理知识库: {kb_name}")
|
|
client = chromadb.PersistentClient(path=kb_path)
|
|
|
|
# 获取所有集合
|
|
collections = client.list_collections()
|
|
print(f" 找到 {len(collections)} 个集合")
|
|
|
|
for collection in collections:
|
|
print(f"\n 处理集合: {collection.name}")
|
|
|
|
# 获取所有图片类型的切片
|
|
try:
|
|
results = collection.get(
|
|
where={"chunk_type": "image"},
|
|
include=["metadatas", "documents", "embeddings"]
|
|
)
|
|
|
|
if not results["ids"]:
|
|
print(f" 无图片类型切片")
|
|
continue
|
|
|
|
print(f" 找到 {len(results['ids'])} 个图片切片")
|
|
|
|
# 检查需要修复的路径
|
|
needs_fix = []
|
|
for i, (id, meta) in enumerate(zip(results["ids"], results["metadatas"])):
|
|
image_path = meta.get("image_path", "")
|
|
if image_path and ("/" in image_path or "\\" in image_path):
|
|
# 路径包含目录，需要修复
|
|
new_path = os.path.basename(image_path)
|
|
needs_fix.append({
|
|
"id": id,
|
|
"old_path": image_path,
|
|
"new_path": new_path,
|
|
"index": i
|
|
})
|
|
|
|
if not needs_fix:
|
|
print(f" 所有路径格式正确，无需修复")
|
|
continue
|
|
|
|
print(f" 需要修复 {len(needs_fix)} 条记录:")
|
|
for item in needs_fix[:5]: # 只显示前5条
|
|
print(f" {item['old_path']} -> {item['new_path']}")
|
|
if len(needs_fix) > 5:
|
|
print(f" ... 还有 {len(needs_fix) - 5} 条")
|
|
|
|
if dry_run:
|
|
print(f" [DRY-RUN] 跳过实际修改")
|
|
continue
|
|
|
|
# 执行修复
|
|
ids_to_update = []
|
|
metadatas_to_update = []
|
|
|
|
for item in needs_fix:
|
|
# 更新元数据
|
|
new_meta = results["metadatas"][item["index"]].copy()
|
|
new_meta["image_path"] = item["new_path"]
|
|
ids_to_update.append(item["id"])
|
|
metadatas_to_update.append(new_meta)
|
|
|
|
# 批量更新
|
|
collection.update(
|
|
ids=ids_to_update,
|
|
metadatas=metadatas_to_update
|
|
)
|
|
print(f" ✓ 已修复 {len(ids_to_update)} 条记录")
|
|
|
|
except Exception as e:
|
|
print(f" 处理集合 {collection.name} 时出错: {e}")
|
|
|
|
print()
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="修复向量库中的图片路径格式")
|
|
parser.add_argument("--dry-run", action="store_true", help="仅预览，不实际修改")
|
|
args = parser.parse_args()
|
|
|
|
if args.dry_run:
|
|
print("=== DRY-RUN 模式：仅预览，不修改数据 ===\n")
|
|
|
|
fix_image_paths(dry_run=args.dry_run)
|
|
|
|
print("\n完成!")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|