rag/core/agentic_media.py

"""
Agentic RAG - 富媒体处理 Mixin

包含图表查找、图片提取、富媒体附加等方法
"""

import re
import json
import logging

from .agentic_base import logger

logger = logging.getLogger(__name__)


class RichMediaMixin:
    """富媒体处理方法"""

    def _find_figure(self, query: str, contexts: list, source: str = None) -> dict:
        """精确查找图表，带 fallback"""
        patterns = [
            r'图\s*(\d+[\.\-]\d+)',
            r'Fig\.?\s*(\d+[\.\-]\d+)',
            r'Figure\s*(\d+[\.\-]\d+)',
        ]

        target_figure = None
        for pattern in patterns:
            match = re.search(pattern, query, re.IGNORECASE)
            if match:
                target_figure = match.group(1).replace('-', '.')
                break

        if not target_figure:
            return {"found": False}

        # 从 contexts 中查找
        for ctx in contexts:
            meta = ctx.get('meta', {})
            fig_num = meta.get('figure_number', '')
            if fig_num == target_figure:
                if not source or meta.get('source') == source:
                    return {
                        "found": True,
                        "chunk_id": meta.get('chunk_id'),
                        "source": meta.get('source'),
                        "page": meta.get('page'),
                        "caption": meta.get('caption'),
                        "image_path": meta.get('image_path'),
                    }

        # Fallback: 直接查向量库
        try:
            from knowledge.manager import get_kb_manager
            kb_mgr = get_kb_manager()
            coll = kb_mgr.get_collection('public_kb')

            if coll:
                where_conditions = [{'chunk_type': {'$in': ['image', 'chart']}}]
                if source:
                    where_conditions.append({'source': source})

                result = coll.get(
                    where={'$and': where_conditions} if len(where_conditions) > 1 else where_conditions[0],
                    include=['metadatas', 'documents']
                )

                for meta, doc in zip(result.get('metadatas', []), result.get('documents', [])):
                    if meta.get('figure_number') == target_figure:
                        return {
                            "found": True,
                            "chunk_id": meta.get('chunk_id'),
                            "source": meta.get('source'),
                            "page": meta.get('page'),
                            "caption": meta.get('caption'),
                            "image_path": meta.get('image_path'),
                        }
                    caption = meta.get('caption', '') or (doc if doc else '')
                    if f"图{target_figure}" in caption or f"图 {target_figure}" in caption:
                        return {
                            "found": True,
                            "chunk_id": meta.get('chunk_id'),
                            "source": meta.get('source'),
                            "page": meta.get('page'),
                            "caption": meta.get('caption'),
                            "image_path": meta.get('image_path'),
                        }
        except Exception as e:
            logger.warning(f"_find_figure fallback 查询失败: {e}")

        return {"found": False}

    def _get_images_for_source(self, source: str, collections: list = None) -> list:
        """直接从向量库获取指定文件的所有图片"""
        try:
            from knowledge.manager import get_kb_manager
            kb_mgr = get_kb_manager()
        except ImportError:
            return []

        images = []
        seen_ids = set()

        target_collections = collections or ['public_kb']

        for kb_name in target_collections:
            try:
                coll = kb_mgr.get_collection(kb_name)
                if not coll:
                    continue

                result = coll.get(
                    where={'source': source},
                    include=['metadatas']
                )

                for meta in result.get('metadatas', []):
                    images_json = meta.get('images_json')
                    if images_json:
                        try:
                            imgs = json.loads(images_json)
                            for img in imgs:
                                img_id = img.get('id')
                                if img_id and img_id not in seen_ids:
                                    seen_ids.add(img_id)
                                    images.append({
                                        "id": img_id,
                                        "caption": img.get("caption", ""),
                                        "url": f"/images/{img_id}",
                                        "page": img.get("page") or meta.get("page"),
                                        "source": source,
                                        "width": img.get("width"),
                                        "height": img.get("height")
                                    })
                        except (json.JSONDecodeError, TypeError):
                            pass
            except Exception as e:
                logger.warning(f"从 {kb_name} 获取图片失败: {e}")
                continue

        return images

    def _extract_rich_media(self, contexts: list, sources_filter: list = None, max_images: int = 10,
                           max_tables: int = 5) -> dict:
        """从检索结果中提取富媒体（图片、表格）"""
        images = []
        tables = []
        seen_image_ids = set()
        seen_table_ids = set()

        for ctx in contexts:
            meta = ctx.get('meta', {})
            source = meta.get('source', '')

            # 过滤来源
            if sources_filter and source not in sources_filter:
                continue

            # 提取图片
            images_json = meta.get('images_json')
            if images_json:
                try:
                    imgs = json.loads(images_json)
                    for img in imgs:
                        img_id = img.get('id')
                        if img_id and img_id not in seen_image_ids:
                            seen_image_ids.add(img_id)
                            images.append({
                                "id": img_id,
                                "caption": img.get("caption", ""),
                                "url": f"/images/{img_id}",
                                "page": img.get("page") or meta.get("page"),
                                "source": source,
                                "type": img.get("type", "image")
                            })
                except (json.JSONDecodeError, TypeError):
                    pass

            # 提取表格
            table_json = meta.get('table_json')
            if table_json:
                try:
                    tbl = json.loads(table_json)
                    tbl_id = tbl.get('id') or meta.get('chunk_id')
                    if tbl_id and tbl_id not in seen_table_ids:
                        seen_table_ids.add(tbl_id)
                        tables.append({
                            "id": tbl_id,
                            "caption": tbl.get("caption", ""),
                            "markdown": tbl.get("markdown", ""),
                            "page": meta.get("page"),
                            "source": source
                        })
                except (json.JSONDecodeError, TypeError):
                    pass

        return {
            "images": images[:max_images],
            "tables": tables[:max_tables]
        }