rag/knowledge/image_cleanup.py

"""
图片/VLM缓存孤儿文件清理模块

提供可被 document.py / collection.py 调用的清理函数，
也可被 cleanup_orphans.py 独立脚本使用。
"""

import hashlib
import logging
import os
from pathlib import Path

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) -> set:
    """
    从 ChromaDB 收集所有被引用的图片文件名。

    Args:
        manager: KnowledgeBaseManager 实例
        collections: 限定知识库列表，None 表示全部

    Returns:
        set of image filenames (e.g., {"185a7a75d246.png", ...})
    """
    referenced = set()

    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:
            continue

        results = col.get(include=['metadatas'])
        if not results['ids']:
            continue

        for meta in results['metadatas']:
            image_path = meta.get('image_path', '')
            if image_path:
                referenced.add(os.path.basename(image_path))

    return referenced


def cleanup_image_orphans(manager, collections=None, dry_run=False) -> dict:
    """
    清理不再被任何 ChromaDB 切片引用的图片和 VLM 缓存文件。

    Args:
        manager: KnowledgeBaseManager 实例
        collections: 限定知识库列表，None 表示全部
        dry_run: True 时只返回孤儿列表不实际删除

    Returns:
        {
            'orphan_images': [(filepath, filename, size_bytes)],
            'orphan_caches': [(filepath, filename, size_bytes)],
            'deleted_images': int,
            'deleted_caches': int,
            'freed_bytes': int
        }
    """
    result = {
        'orphan_images': [],
        'orphan_caches': [],
        'deleted_images': 0,
        'deleted_caches': 0,
        'freed_bytes': 0
    }

    # 1. 收集引用
    referenced = collect_referenced_images(manager, collections)

    # 2. 查找孤儿图片
    if IMAGES_DIR.exists():
        for f in IMAGES_DIR.iterdir():
            if f.is_file() and f.name not in referenced:
                result['orphan_images'].append((str(f), f.name, f.stat().st_size))

    # 3. 查找孤儿 VLM 缓存（图片已删除则缓存也应是孤儿）
    if VLM_CACHE_DIR.exists():
        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 f.is_file() and f.suffix == '.txt':
                cache_hash = f.stem
                if cache_hash not in referenced_hashes:
                    result['orphan_caches'].append((str(f), f.name, f.stat().st_size))

    # 4. 删除
    if not dry_run:
        for filepath, filename, size in result['orphan_images']:
            try:
                os.remove(filepath)
                result['deleted_images'] += 1
                result['freed_bytes'] += size
            except OSError as e:
                logger.warning(f"删除孤儿图片失败: {filename} - {e}")

        for filepath, filename, size in result['orphan_caches']:
            try:
                os.remove(filepath)
                result['deleted_caches'] += 1
                result['freed_bytes'] += size
            except OSError as e:
                logger.warning(f"删除孤儿缓存失败: {filename} - {e}")

        if result['deleted_images'] or result['deleted_caches']:
            logger.info(
                f"清理孤儿: {result['deleted_images']} 图片 + "
                f"{result['deleted_caches']} 缓存, "
                f"释放 {result['freed_bytes']/1024:.1f} KB"
            )

    return result