rag/knowledge/chunk.py

"""
知识库管理器 - 切片管理 Mixin

提供文档切片（Chunk）的 CRUD 操作，支持：
- 新增切片：手动添加文本切片到向量库
- 修改切片：更新切片内容和元数据
- 删除切片：从向量库移除切片
- 查询切片：分页获取切片列表

切片是向量检索的基本单位，每个切片包含：
- 文档内容（用于向量化）
- 元数据（来源、页码、状态等）
- 向量嵌入（由 embedding 模型生成）

主要方法：
- add_chunk: 新增切片
- update_chunk: 修改切片
- delete_chunk: 删除切片
- list_chunks: 查询切片列表
"""

import hashlib
import logging
from datetime import datetime
from typing import List, Dict, Optional, Tuple

logger = logging.getLogger(__name__)


class ChunkMixin:
    """
    切片管理 Mixin

    提供切片级别的 CRUD 操作，同时维护向量索引和 BM25 索引。

    依赖属性（需由主类提供）：
    - self.get_collection: 获取向量库集合的方法
    - self.get_bm25_index: 获取 BM25 索引的方法
    - self.save_bm25_index: 保存 BM25 索引的方法
    """

    def add_chunk(
        self,
        kb_name: str,
        content: str,
        metadata: dict = None
    ) -> str:
        """
        新增单个切片

        将文本内容切片添加到指定向量库，自动生成向量嵌入和 BM25 索引。

        Args:
            kb_name: 目标向量库名称
            content: 切片文本内容
            metadata: 可选的元数据字典，可包含：
                - source: 来源文件名
                - page: 页码
                - chunk_type: 切片类型 ('text' | 'table' | 'image')
                - status: 状态 ('active' | 'deprecated')
                - 其他自定义字段

        Returns:
            新创建的切片 ID，格式为: manual_{时间戳}_{内容哈希}

        Raises:
            ValueError: 向量库不存在
            Exception: 向量生成失败

        Example:
            >>> chunk_id = kb_manager.add_chunk(
            ...     "public_kb",
            ...     "这是要添加的文本内容",
            ...     {"source": "manual", "page": 1}
            ... )
        """
        collection = self.get_collection(kb_name)
        if not collection:
            raise ValueError(f"向量库 '{kb_name}' 不存在")

        content_hash = hashlib.md5(content.encode('utf-8')).hexdigest()[:8]
        chunk_id = f"manual_{datetime.now().strftime('%Y%m%d%H%M%S')}_{content_hash}"

        chunk_metadata = {
            "chunk_id": chunk_id,
            "chunk_type": "text",
            "collection": kb_name,
            "source": "manual",
            "status": "active",
            "version": "v1",
            "change_time": datetime.now().isoformat(),
        }
        if metadata:
            chunk_metadata.update(metadata)

        try:
            from core.engine import get_engine
            engine = get_engine()
            if not engine._initialized:
                engine.initialize()
            embedding = engine.embedding_model.encode(content).tolist()
        except Exception as e:
            logger.error(f"生成向量失败: {e}")
            raise

        collection.add(
            ids=[chunk_id],
            documents=[content],
            metadatas=[chunk_metadata],
            embeddings=[embedding]
        )

        bm25 = self.get_bm25_index(kb_name)
        if bm25:
            bm25.add_documents([chunk_id], [content], [chunk_metadata])
            self.save_bm25_index(kb_name)

        logger.info(f"新增切片: {chunk_id} -> {kb_name}")
        return chunk_id

    def update_chunk(
        self,
        kb_name: str,
        chunk_id: str,
        content: str = None,
        metadata: dict = None
    ) -> bool:
        """
        修改切片

        更新切片的内容和/或元数据。如果更新内容，会重新生成向量嵌入。

        Args:
            kb_name: 向量库名称
            chunk_id: 切片 ID
            content: 新的文本内容（None 表示不更新）
            metadata: 要更新的元数据字段（None 表示不更新）

        Returns:
            更新成功返回 True，切片不存在或更新失败返回 False

        Example:
            >>> kb_manager.update_chunk(
            ...     "public_kb",
            ...     "manual_20260517_abc123",
            ...     content="更新后的内容",
            ...     metadata={"status": "updated"}
            ... )
            True
        """
        collection = self.get_collection(kb_name)
        if not collection:
            return False

        existing = collection.get(ids=[chunk_id])
        if not existing['ids']:
            logger.warning(f"切片不存在: {chunk_id}")
            return False

        update_kwargs = {"ids": [chunk_id]}

        if content is not None:
            update_kwargs["documents"] = [content]
            try:
                from core.engine import get_engine
                engine = get_engine()
                if not engine._initialized:
                    engine.initialize()
                embedding = engine.embedding_model.encode(content).tolist()
                update_kwargs["embeddings"] = [embedding]
            except Exception as e:
                logger.error(f"生成向量失败: {e}")
                return False

        if metadata is not None:
            existing_metadata = existing['metadatas'][0] if existing['metadatas'] else {}
            existing_metadata.update(metadata)
            update_kwargs["metadatas"] = [existing_metadata]

        collection.update(**update_kwargs)

        if content is not None:
            bm25 = self.get_bm25_index(kb_name)
            if bm25:
                bm25.add_documents([chunk_id], [content], [metadata or {}])
                self.save_bm25_index(kb_name)

        logger.info(f"更新切片: {chunk_id}")
        return True

    def get_chunk(self, kb_name: str, chunk_id: str) -> Optional[Dict]:
        """
        获取单个切片信息

        Args:
            kb_name: 向量库名称
            chunk_id: 切片 ID

        Returns:
            切片信息字典，包含 id, document, metadata 等字段
            切片不存在时返回 None
        """
        collection = self.get_collection(kb_name)
        if not collection:
            return None

        result = collection.get(ids=[chunk_id], include=["documents", "metadatas"])
        if not result['ids']:
            return None

        return {
            'id': result['ids'][0],
            'document': result['documents'][0] if result['documents'] else '',
            'metadata': result['metadatas'][0] if result['metadatas'] else {}
        }

    def delete_chunk(self, kb_name: str, chunk_id: str) -> Tuple[bool, Optional[str]]:
        """
        删除切片

        从向量库中移除指定的切片。

        Args:
            kb_name: 向量库名称
            chunk_id: 切片 ID

        Returns:
            元组 (success, source_file)：
            - success: 删除是否成功
            - source_file: 被删除切片的来源文件名（用于清理哈希记录）

        Note:
            删除后 BM25 索引不会立即更新，需要手动调用 rebuild_bm25_index。
        """
        collection = self.get_collection(kb_name)
        if not collection:
            return False, None

        # 获取切片信息（删除前）
        existing = collection.get(ids=[chunk_id], include=["metadatas"])
        if not existing['ids']:
            logger.warning(f"切片不存在: {chunk_id}")
            return False, None

        source_file = None
        if existing['metadatas'] and existing['metadatas'][0]:
            source_file = existing['metadatas'][0].get('source')

        collection.delete(ids=[chunk_id])
        logger.info(f"删除切片: {chunk_id}")
        return True, source_file

    def delete_chunks_by_source(self, kb_name: str, source: str) -> int:
        """
        批量删除指定文件的所有切片

        Args:
            kb_name: 向量库名称
            source: 文件名（source 字段值）

        Returns:
            删除的切片数量

        Note:
            删除后会重建 BM25 索引。
        """
        collection = self.get_collection(kb_name)
        if not collection:
            return 0

        # 获取所有切片，找到匹配 source 的
        all_chunks = collection.get(include=["metadatas"])

        # 找到要删除的切片 ID
        chunk_ids_to_delete = []
        for i, chunk_id in enumerate(all_chunks['ids']):
            metadata = all_chunks['metadatas'][i] if all_chunks['metadatas'] else {}
            if metadata.get('source') == source:
                chunk_ids_to_delete.append(chunk_id)

        if not chunk_ids_to_delete:
            logger.warning(f"未找到文件 {source} 的切片")
            return 0

        # 批量删除
        collection.delete(ids=chunk_ids_to_delete)
        logger.info(f"批量删除切片: {source}, 共 {len(chunk_ids_to_delete)} 个")

        # 重建 BM25 索引
        try:
            self._bm25_indexes.pop(kb_name, None)
            bm25 = self.get_bm25_index(kb_name)
            if bm25:
                remaining = collection.get(include=["documents", "metadatas"])
                if remaining['ids']:
                    bm25.add_documents(
                        remaining['ids'],
                        remaining['documents'] or [],
                        remaining['metadatas'] or []
                    )
                    self.save_bm25_index(kb_name)
        except Exception as e:
            logger.warning(f"重建 BM25 索引失败: {e}")

        return len(chunk_ids_to_delete)

    def list_chunks(
        self,
        kb_name: str,
        document_id: str = None,
        limit: int = 100,
        offset: int = 0
    ) -> List[Dict]:
        """
        获取向量库中的切片列表

        支持分页和按文档过滤。

        Args:
            kb_name: 向量库名称
            document_id: 文档 ID（source 字段），用于过滤特定文档的切片
            limit: 返回数量限制（默认 100）
            offset: 偏移量（默认 0）

        Returns:
            切片字典列表，每个字典包含：
            - id: 切片 ID
            - document: 切片内容
            - metadata: 元数据字典
            - status: 状态
            - version: 版本号

        Example:
            >>> chunks = kb_manager.list_chunks("public_kb", limit=10)
            >>> for chunk in chunks:
            ...     print(chunk["id"], chunk["status"])
        """
        collection = self.get_collection(kb_name)
        if not collection:
            return []

        where_filter = {}
        if document_id:
            where_filter["source"] = document_id

        try:
            result = collection.get(
                where=where_filter if where_filter else None,
                limit=limit,
                offset=offset
            )
        except Exception as e:
            logger.warning(f"获取切片失败: {e}")
            return []

        return [
            {
                "id": id,
                "document": doc,
                "metadata": meta,
                "status": meta.get("status", "active") if meta else "active",
                "version": meta.get("version", "v1") if meta else "v1"
            }
            for id, doc, meta in zip(
                result['ids'],
                result.get('documents', []),
                result.get('metadatas', [])
            )
        ]