rag/cleanup_orphans.py

"""
孤儿文件清理工具

清理 .data/images/ 和 .data/cache/vlm/ 中不再被任何 ChromaDB 切片引用的文件。

用法:
    # 预览模式（不删除，只显示孤儿文件）
    python cleanup_orphans.py

    # 实际删除
    python cleanup_orphans.py --force

    # 仅清理特定知识库
    python cleanup_orphans.py --collections public_kb test_kb
"""

import argparse
import hashlib
import logging
import os
import sys
from pathlib import Path

logging.basicConfig(level=logging.INFO, format='%(levelname)s: %(message)s')
logger = logging.getLogger(__name__)

IMAGES_DIR = Path(".data/images")
VLM_CACHE_DIR = Path(".data/cache/vlm")


def compute_file_hash(file_path: str) -> str:
    """计算文件 MD5"""
    with open(file_path, 'rb') as f:
        return hashlib.md5(f.read()).hexdigest()


def collect_referenced_images(manager, collections=None) -> dict:
    """
    从 ChromaDB 收集所有被引用的图片路径。

    Returns:
        {image_filename: set of chunk_ids referencing it}
    """
    referenced = {}

    if collections:
        kb_names = collections
    else:
        kb_names = [c.name if hasattr(c, 'name') else str(c)
                     for c in manager.list_collections()]

    for kb_name in kb_names:
        try:
            col = manager.get_collection(kb_name)
        except Exception as e:
            logger.warning(f"无法获取 {kb_name}: {e}")
            continue

        # 查找所有带 image_path 的切片
        results = col.get(include=['metadatas'])
        if not results['ids']:
            continue

        for chunk_id, meta in zip(results['ids'], results['metadatas']):
            image_path = meta.get('image_path', '')
            if not image_path:
                continue

            # image_path 可能是: "185a7a75d246.png" 或相对路径
            filename = os.path.basename(image_path)
            if filename not in referenced:
                referenced[filename] = set()
            referenced[filename].add(f"{kb_name}/{chunk_id}")

    return referenced


def find_orphan_images(referenced: dict) -> list:
    """
    查找 .data/images/ 中不再被引用的图片文件。

    Returns:
        [(filepath, filename, size_bytes)]
    """
    orphans = []
    if not IMAGES_DIR.exists():
        return orphans

    for f in IMAGES_DIR.iterdir():
        if not f.is_file():
            continue
        if f.name not in referenced:
            orphans.append((str(f), f.name, f.stat().st_size))

    return orphans


def find_orphan_vlm_caches(referenced: dict) -> list:
    """
    查找 .data/cache/vlm/ 中对应的图片已不存在的缓存文件。

    缓存文件以图片 MD5 命名，如果图片被删了，缓存也应该是孤儿。

    Returns:
        [(filepath, filename, size_bytes)]
    """
    orphans = []
    if not VLM_CACHE_DIR.exists():
        return orphans

    # 构建 referenced 中所有图片的 MD5 集合
    referenced_hashes = set()
    for filename in referenced:
        full_path = IMAGES_DIR / filename
        if full_path.exists():
            try:
                img_hash = compute_file_hash(str(full_path))
                referenced_hashes.add(img_hash)
            except Exception:
                pass

    for f in VLM_CACHE_DIR.iterdir():
        if not f.is_file() or not f.suffix == '.txt':
            continue
        cache_hash = f.stem  # 文件名就是 MD5
        if cache_hash not in referenced_hashes:
            orphans.append((str(f), f.name, f.stat().st_size))

    return orphans


def delete_files(file_list: list, dry_run: bool = True) -> int:
    """删除文件列表，返回删除数量"""
    deleted = 0
    for filepath, filename, size in file_list:
        if dry_run:
            logger.info(f"  [DRY-RUN] 将删除: {filename} ({size/1024:.1f} KB)")
        else:
            try:
                os.remove(filepath)
                logger.info(f"  已删除: {filename} ({size/1024:.1f} KB)")
                deleted += 1
            except OSError as e:
                logger.warning(f"  删除失败: {filename} - {e}")
    return deleted


def main():
    parser = argparse.ArgumentParser(description="孤儿文件清理工具")
    parser.add_argument("--force", action="store_true",
                        help="实际删除文件（默认仅预览）")
    parser.add_argument("--collections", nargs="+",
                        help="仅检查指定知识库（默认检查全部）")
    args = parser.parse_args()

    os.chdir(os.path.dirname(os.path.abspath(__file__)))

    mode = "删除" if args.force else "预览（不删除）"
    print("=" * 60)
    print(f"孤儿文件清理 - {mode}")
    print("=" * 60)

    # 1. 收集引用
    print("\n[1/4] 扫描 ChromaDB 引用...")
    sys.path.insert(0, '.')
    from knowledge.manager import get_kb_manager
    manager = get_kb_manager()
    referenced = collect_referenced_images(manager, args.collections)
    print(f"  被引用的图片: {len(referenced)} 个")

    # 2. 查找孤儿图片
    print("\n[2/4] 查找孤儿图片...")
    orphan_images = find_orphan_images(referenced)
    total_img_size = sum(s for _, _, s in orphan_images)
    print(f"  孤儿图片: {len(orphan_images)} 个 ({total_img_size/1024:.1f} KB)")

    # 3. 查找孤儿 VLM 缓存
    print("\n[3/4] 查找孤儿 VLM 缓存...")
    orphan_caches = find_orphan_vlm_caches(referenced)
    total_cache_size = sum(s for _, _, s in orphan_caches)
    print(f"  孤儿缓存: {len(orphan_caches)} 个 ({total_cache_size/1024:.1f} KB)")

    # 4. 清理
    print("\n[4/4] 清理...")
    if not orphan_images and not orphan_caches:
        print("  没有需要清理的文件")
    else:
        if not args.force:
            print("  预览模式，以下文件将被删除：")

        img_deleted = delete_files(orphan_images, dry_run=not args.force)
        cache_deleted = delete_files(orphan_caches, dry_run=not args.force)

        if args.force:
            print(f"\n  删除了 {img_deleted} 个图片 + {cache_deleted} 个缓存")
            freed = total_img_size + total_cache_size
            print(f"  释放空间: {freed/1024:.1f} KB")
        else:
            print(f"\n  共 {len(orphan_images) + len(orphan_caches)} 个文件待清理")
            print(f"  释放空间: {(total_img_size + total_cache_size)/1024:.1f} KB")
            print("  使用 --force 参数执行实际删除")

    print("=" * 60)


if __name__ == '__main__':
    main()