rag/knowledge/document.py

"""
知识库管理器 - 文档管理 Mixin

包含文档级别的管理方法
"""

import os
import logging
from datetime import datetime
from typing import List, Dict, Optional

from .base import _get_doc_type

logger = logging.getLogger(__name__)


class DocumentMixin:
    """文档管理方法"""

    def get_document_count(self, kb_name: str) -> int:
        """获取向量库中的文档数量"""
        collection = self.get_collection(kb_name)
        return collection.count() if collection else 0

    def list_documents(self, kb_name: str) -> List[dict]:
        """列出向量库中的文档"""
        collection = self.get_collection(kb_name)
        if not collection:
            return []

        result = collection.get()

        from collections import Counter
        file_chunks = Counter()

        for meta in result.get('metadatas', []):
            source = meta.get('source', 'unknown')
            file_chunks[source] += 1

        return [
            {"source": source, "chunks": count}
            for source, count in file_chunks.items()
        ]

    def delete_document(self, kb_name: str, filename: str) -> int:
        """从向量库删除文档"""
        collection = self.get_collection(kb_name)
        if not collection:
            return 0

        result = collection.get(where={"source": filename})

        if not result['ids']:
            return 0

        collection.delete(ids=result['ids'])
        deleted = len(result['ids'])

        logger.info(f"从 {kb_name} 删除文档: {filename}, 片段数: {deleted}")
        return deleted

    def deprecate_document(
        self,
        kb_name: str,
        filename: str,
        reason: str = "制度废止",
        deprecated_by: str = ""
    ) -> Dict:
        """软删除文档 - 将chunks状态标记为deprecated"""
        collection = self.get_collection(kb_name)
        if not collection:
            return {"success": False, "error": "向量库不存在"}

        result = collection.get(where={"source": filename})

        if not result['ids']:
            return {"success": False, "error": "文档不存在"}

        deprecated_date = datetime.now().isoformat()
        updated_metadatas = [
            {
                **m,
                "status": "deprecated",
                "deprecated_date": deprecated_date,
                "deprecated_reason": reason,
                "deprecated_by": deprecated_by
            }
            for m in result['metadatas']
        ]

        collection.update(
            ids=result['ids'],
            metadatas=updated_metadatas
        )

        self.rebuild_bm25_index(kb_name)

        logger.info(f"软删除文档: {kb_name}/{filename}, chunks: {len(result['ids'])}, 原因: {reason}")

        return {
            "success": True,
            "deprecated_chunks": len(result['ids']),
            "document_id": filename,
            "collection": kb_name,
            "deprecated_date": deprecated_date
        }

    def restore_document(self, kb_name: str, filename: str) -> Dict:
        """恢复已废止的文档"""
        collection = self.get_collection(kb_name)
        if not collection:
            return {"success": False, "error": "向量库不存在"}

        result = collection.get(
            where={
                "$and": [
                    {"source": filename},
                    {"status": "deprecated"}
                ]
            }
        )

        if not result['ids']:
            return {"success": False, "error": "未找到已废止的文档"}

        updated_metadatas = [
            {
                **m,
                "status": "active",
                "deprecated_date": None,
                "deprecated_reason": None
            }
            for m in result['metadatas']
        ]

        collection.update(
            ids=result['ids'],
            metadatas=updated_metadatas
        )

        self.rebuild_bm25_index(kb_name)

        logger.info(f"恢复文档: {kb_name}/{filename}, chunks: {len(result['ids'])}")

        return {
            "success": True,
            "restored_chunks": len(result['ids']),
            "document_id": filename,
            "collection": kb_name
        }

    def get_document_chunks(
        self,
        kb_name: str,
        filename: str,
        status: str = None
    ) -> List[Dict]:
        """获取文档的chunks列表"""
        collection = self.get_collection(kb_name)
        if not collection:
            return []

        where_filter = {"source": filename}
        if status:
            where_filter["status"] = status

        result = collection.get(where=where_filter)

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

    def get_document_info(self, kb_name: str, filename: str) -> Optional[Dict]:
        """获取文档基本信息"""
        collection = self.get_collection(kb_name)
        if not collection:
            return None

        result = collection.get(where={"source": filename})

        if not result['ids']:
            return None

        status_counts = {}
        for meta in result['metadatas']:
            status = meta.get("status", "active")
            status_counts[status] = status_counts.get(status, 0) + 1

        main_status = "active"
        if status_counts.get("deprecated", 0) > status_counts.get("active", 0):
            main_status = "deprecated"
        elif status_counts.get("superseded", 0) > 0:
            main_status = "superseded"

        first_meta = result['metadatas'][0] if result['metadatas'] else {}

        return {
            "document_id": filename,
            "collection": kb_name,
            "total_chunks": len(result['ids']),
            "status": main_status,
            "status_counts": status_counts,
            "version": first_meta.get("version", "v1"),
            "effective_date": first_meta.get("effective_date"),
            "deprecated_date": first_meta.get("deprecated_date"),
            "deprecated_reason": first_meta.get("deprecated_reason"),
            "security_level": first_meta.get("security_level", "public")
        }

    def list_documents_by_status(
        self,
        kb_name: str,
        status: str = None
    ) -> List[Dict]:
        """按状态列出文档"""
        collection = self.get_collection(kb_name)
        if not collection:
            return []

        result = collection.get()

        doc_info = {}
        for meta in result.get('metadatas', []):
            source = meta.get('source', 'unknown')
            chunk_status = meta.get('status', 'active')

            if source not in doc_info:
                doc_info[source] = {
                    "source": source,
                    "chunks": 0,
                    "status_counts": {},
                    "collection": kb_name
                }

            doc_info[source]["chunks"] += 1
            doc_info[source]["status_counts"][chunk_status] = \
                doc_info[source]["status_counts"].get(chunk_status, 0) + 1

        result_list = []
        for doc in doc_info.values():
            counts = doc["status_counts"]
            if counts.get("deprecated", 0) > counts.get("active", 0):
                doc["status"] = "deprecated"
            elif counts.get("superseded", 0) > 0:
                doc["status"] = "superseded"
            else:
                doc["status"] = "active"

            if status is None or doc["status"] == status:
                result_list.append(doc)

        return result_list

    # add_file_to_kb 方法较长，暂时保留在 manager.py 中
    # 后续可以拆分到单独的 document_processing.py