rag/knowledge/lazy_enhance.py

"""
懒加载增强模块（Phase 4）

按需调用 LLM/VLM 生成表格摘要和图片描述
"""

import hashlib
import logging
from pathlib import Path

logger = logging.getLogger(__name__)

# 缓存目录（扁平化）
VLM_CACHE_DIR = Path(".data/cache/vlm")
LLM_CACHE_DIR = Path(".data/cache/llm")


def compute_file_hash(file_path: str) -> str:
    """计算文件哈希"""
    try:
        with open(file_path, 'rb') as f:
            return hashlib.md5(f.read()).hexdigest()
    except Exception as e:
        logger.warning(f"计算文件哈希失败: {e}")
        return hashlib.md5(file_path.encode()).hexdigest()


def _get_embedding_model():
    """从 RAGEngine 获取 embedding 模型（KnowledgeBaseManager 上没有此属性）"""
    try:
        from core.engine import get_engine
        engine = get_engine()
        if not engine._initialized:
            engine.initialize()
        return engine.embedding_model
    except Exception as e:
        logger.warning(f"获取 embedding 模型失败: {e}")
        return None


async def lazy_vlm_description(chunk_id: str, image_path: str, kb_name: str, metadata: dict = None, defer_chromadb: bool = False) -> str:
    """
    懒加载 VLM 描述

    触发条件：图片切片被检索命中

    Args:
        chunk_id: 切片 ID
        image_path: 图片路径（相对路径或绝对路径）
        kb_name: 知识库名称
        metadata: 图片元数据（包含 section、page、caption、上下文等）
        defer_chromadb: 为 True 时跳过 ChromaDB 更新（仅写文件缓存），避免后台线程写锁竞争

    Returns:
        VLM 生成的图片描述
    """
    import os
    from knowledge.manager import get_kb_manager

    # 构建完整图片路径
    if not os.path.isabs(image_path):
        full_image_path = os.path.join('.data/images', image_path)
    else:
        full_image_path = image_path

    # 1. 检查缓存（空缓存视为无效，需重新生成）
    img_hash = compute_file_hash(full_image_path)
    cache_file = VLM_CACHE_DIR / f"{img_hash}.txt"
    if cache_file.exists():
        cached = cache_file.read_text(encoding='utf-8')
        if len(cached.strip()) >= 5:
            logger.info(f"VLM 缓存命中: {image_path}")
            return cached
        else:
            logger.warning(f"VLM 缓存内容过短({len(cached.strip())}字符)，删除并重新生成: {image_path}")
            try:
                cache_file.unlink()
            except OSError:
                pass

    # 2. 调用 VLM（传入元数据）
    logger.info(f"VLM 懒加载: {image_path}")
    kb_manager = get_kb_manager()
    description = kb_manager._generate_image_description(full_image_path, metadata=metadata)

    # 3. 空描述保护：VLM 返回内容过短时不写入缓存和向量库
    if not description or len(description.strip()) < 5:
        logger.warning(f"VLM 返回描述过短({len(description.strip()) if description else 0}字符)，跳过缓存和向量库更新: {image_path}")
        return description or ''

    # 4. 写入缓存
    VLM_CACHE_DIR.mkdir(parents=True, exist_ok=True)
    cache_file.write_text(description, encoding='utf-8')

    # 5. 更新向量库（metadata + embedding），需校验 chunk_id 非空
    # defer_chromadb=True 时跳过（后台线程只写缓存，避免 SQLite 写锁竞争）
    if defer_chromadb:
        logger.info(f"延迟 ChromaDB 更新(仅写缓存): {chunk_id}")
        return description

    if not chunk_id:
        logger.warning("chunk_id 为空，跳过向量库更新")
        return description

    try:
        collection = kb_manager.get_collection(kb_name)
        result = collection.get(ids=[chunk_id], include=['metadatas'])
        if result['metadatas']:
            # 更新 metadata
            new_metadata = {
                **result['metadatas'][0],
                'has_vlm_desc': True,
                'vlm_desc': description
            }

            # 更新 embedding（使用 VLM 描述重新计算向量）
            # 这样 VLM 描述中的关键词（如"发电量"）才能参与相似度检索
            embedding_model = _get_embedding_model()
            if embedding_model:
                new_vector = embedding_model.encode(description).tolist()
                if isinstance(new_vector[0], list):
                    new_vector = new_vector[0]

                collection.update(
                    ids=[chunk_id],
                    metadatas=[new_metadata],
                    embeddings=[new_vector],
                    documents=[description]  # 同时更新 document 字段
                )
                logger.info(f"已更新向量库(embedding+metadata): {chunk_id}")
            else:
                # 无 embedding 模型时只更新 metadata
                collection.update(
                    ids=[chunk_id],
                    metadatas=[new_metadata]
                )
                logger.info(f"已更新向量库(仅metadata,无embedding模型): {chunk_id}")
    except Exception as e:
        logger.warning(f"更新向量库失败: {e}")

    return description


async def lazy_table_summary(chunk_id: str, table_md: str, kb_name: str, defer_chromadb: bool = False) -> str:
    """
    懒加载表格摘要

    触发条件：表格切片被检索命中且相关性 > 0.7

    Args:
        chunk_id: 切片 ID
        table_md: 表格 Markdown 内容
        kb_name: 知识库名称
        defer_chromadb: 为 True 时跳过 ChromaDB 更新（仅写文件缓存），避免后台线程写锁竞争

    Returns:
        LLM 生成的表格摘要
    """
    from knowledge.manager import get_kb_manager

    # 1. 检查缓存（空缓存视为无效）
    table_hash = hashlib.md5(table_md.encode()).hexdigest()
    cache_file = LLM_CACHE_DIR / f"{table_hash}.txt"
    if cache_file.exists():
        cached = cache_file.read_text(encoding='utf-8')
        if len(cached.strip()) >= 5:
            logger.info(f"LLM 缓存命中: {chunk_id}")
            return cached
        else:
            logger.warning(f"LLM 缓存内容过短({len(cached.strip())}字符)，删除并重新生成: {chunk_id}")
            try:
                cache_file.unlink()
            except OSError:
                pass

    # 2. 调用 LLM
    logger.info(f"LLM 懒加载: {chunk_id}")
    kb_manager = get_kb_manager()
    summary = kb_manager._generate_table_summary(table_md, None)

    # 空摘要保护
    if not summary or len(summary.strip()) < 5:
        logger.warning(f"LLM 返回摘要过短，跳过缓存和向量库更新: {chunk_id}")
        return summary or ''

    # 3. 写入缓存
    LLM_CACHE_DIR.mkdir(parents=True, exist_ok=True)
    cache_file.write_text(summary, encoding='utf-8')

    # 4. 更新向量库，需校验 chunk_id 非空
    # defer_chromadb=True 时跳过（后台线程只写缓存，避免 SQLite 写锁竞争）
    if defer_chromadb:
        logger.info(f"延迟 ChromaDB 更新(仅写缓存): {chunk_id}")
        return summary

    if not chunk_id:
        logger.warning("chunk_id 为空，跳过表格向量库更新")
        return summary
    try:
        collection = kb_manager.get_collection(kb_name)
        result = collection.get(ids=[chunk_id], include=['metadatas'])
        if result['metadatas']:
            # 新增摘要切片（需要 embedding 模型）
            embedding_model = _get_embedding_model()
            if embedding_model:
                vector = embedding_model.encode(summary).tolist()
                if isinstance(vector[0], list):
                    vector = vector[0]

                collection.add(
                    ids=[f"{chunk_id}_summary"],
                    embeddings=[vector],
                    documents=[summary],
                    metadatas=[{
                        **result['metadatas'][0],
                        'is_summary': True,
                        'original_doc_id': chunk_id
                    }]
                )
                logger.info(f"已新增摘要切片(embedding): {chunk_id}_summary")
            else:
                logger.info(f"跳过摘要切片(无embedding模型): {chunk_id}")
            # 更新原切片标记（不依赖 embedding 模型）
            collection.update(
                ids=[chunk_id],
                metadatas=[{**result['metadatas'][0], 'has_summary': True}]
            )
    except Exception as e:
        logger.warning(f"更新向量库失败: {e}")

    return summary


async def enhance_retrieved_chunks(contexts: list, query: str, kb_name: str, defer_chromadb: bool = False):
    """
    检索后增强：按需调用 LLM/VLM

    Args:
        contexts: 检索上下文列表
        query: 用户查询
        kb_name: 知识库名称
        defer_chromadb: 为 True 时后台线程只写文件缓存，不更新 ChromaDB（避免写锁竞争）
    """
    import re

    for ctx in contexts:
        try:
            meta = ctx.get('meta', {})
            chunk_type = meta.get('chunk_type', 'text')
            image_path = meta.get('image_path', '')

            # 图片切片：懒加载 VLM 描述
            if chunk_type in ('image', 'chart') and not meta.get('has_vlm_desc'):
                if image_path:
                    # 从 doc 字段中提取图号（上下文可能包含"见图2.5"等）
                    doc_text = ctx.get('doc', '')

                    # 提取图号（从前文/后文中）
                    figure_number = ""
                    fig_match = re.search(r'[见如]?图\s*(\d+\.?\d*)', doc_text)
                    if fig_match:
                        figure_number = fig_match.group(1)

                    # 如果 doc 中没有，尝试从 section 中提取
                    section = meta.get('section') or meta.get('section_path', '')
                    if not figure_number and section:
                        fig_match = re.search(r'[见如]?图\s*(\d+\.?\d*)', section)
                        if fig_match:
                            figure_number = fig_match.group(1)

                    # 传入图片元数据，增强 VLM 描述
                    image_metadata = {
                        'section': section,
                        'page': meta.get('page'),
                        'caption': meta.get('caption', ''),
                        'source': meta.get('source', ''),
                        'figure_number': figure_number,
                        'doc_text': doc_text
                    }
                    vlm_desc = await lazy_vlm_description(
                        meta.get('chunk_id', ''),
                        image_path,
                        kb_name,
                        metadata=image_metadata,
                        defer_chromadb=defer_chromadb
                    )
                    if vlm_desc:
                        ctx['doc'] = vlm_desc
                        ctx['vlm_enhanced'] = True

            # 表格切片：同时处理摘要和关联图片的 VLM 描述
            elif chunk_type == 'table':
                doc_text = ctx.get('doc', '')

                # 1. 懒加载表格摘要（高分切片）
                if not meta.get('has_summary'):
                    score = ctx.get('score', 0)
                    if score > 0.7:
                        summary = await lazy_table_summary(
                            meta.get('chunk_id', ''),
                            doc_text,
                            kb_name,
                            defer_chromadb=defer_chromadb
                        )
                        if summary:
                            ctx['summary'] = summary
                            ctx['llm_enhanced'] = True

                # 2. 表格有关联图片时，懒加载 VLM 描述
                if image_path and not meta.get('has_vlm_desc'):
                    # 提取表号（如 "表2.2"、"见表2.1"）
                    table_number = ""
                    table_match = re.search(r'[见如]?表\s*(\d+\.?\d*)', doc_text)
                    if table_match:
                        table_number = table_match.group(1)

                    section = meta.get('section') or meta.get('section_path', '')
                    if not table_number and section:
                        table_match = re.search(r'[见如]?表\s*(\d+\.?\d*)', section)
                        if table_match:
                            table_number = table_match.group(1)

                    table_image_metadata = {
                        'section': section,
                        'page': meta.get('page'),
                        'caption': meta.get('caption', ''),
                        'source': meta.get('source', ''),
                        'table_number': table_number,
                        'figure_number': table_number,
                        'doc_text': doc_text,
                        'is_table': True
                    }
                    vlm_desc = await lazy_vlm_description(
                        meta.get('chunk_id', ''),
                        image_path,
                        kb_name,
                        metadata=table_image_metadata,
                        defer_chromadb=defer_chromadb
                    )
                    if vlm_desc:
                        ctx['image_description'] = vlm_desc
                        ctx['vlm_enhanced'] = True

        except Exception as e:
            chunk_id = ctx.get('meta', {}).get('chunk_id', '?')
            logger.warning(f"增强切片失败(chunk_id={chunk_id}): {e}")