AstrBot/astrbot/core/knowledge_base/manager.py

"""知识库管理器

该模块提供知识库的CRUD操作和文档上传处理流程。
"""

import uuid
from pathlib import Path
from typing import Optional

import aiofiles
from sqlalchemy import func, select, update

from astrbot.core.db import BaseDatabase
from astrbot.core.db.vec_db.base import BaseVecDB
from astrbot.core.knowledge_base.chunking.base import BaseChunker
from astrbot.core.knowledge_base.models import KBChunk, KBDocument, KnowledgeBase
from astrbot.core.knowledge_base.parsers.base import BaseParser


class KBManager:
    """知识库管理器

    职责:
    - 知识库的 CRUD 操作
    - 文档上传与解析
    - 文档块生成与存储
    - 多媒体资源管理
    """

    def __init__(
        self,
        db: BaseDatabase,
        vec_db: BaseVecDB,
        storage_path: str,
        parsers: dict[str, BaseParser],
        chunker: BaseChunker,
        provider_manager=None,
    ):
        self.db = db
        self.vec_db = vec_db
        self.storage_path = Path(storage_path)
        self.media_path = self.storage_path / "media"
        self.files_path = self.storage_path / "files"
        self.parsers = parsers
        self.chunker = chunker
        self.provider_manager = provider_manager

        # 确保目录存在
        self.media_path.mkdir(parents=True, exist_ok=True)
        self.files_path.mkdir(parents=True, exist_ok=True)

    # ===== 知识库操作 =====

    async def create_kb(
        self,
        kb_name: str,
        description: Optional[str] = None,
        emoji: Optional[str] = None,
        embedding_provider_id: Optional[str] = None,
        rerank_provider_id: Optional[str] = None,
        chunk_size: Optional[int] = None,
        chunk_overlap: Optional[int] = None,
        top_k_dense: Optional[int] = None,
        top_k_sparse: Optional[int] = None,
        top_m_final: Optional[int] = None,
        enable_rerank: Optional[bool] = None,
    ) -> KnowledgeBase:
        """创建知识库

        Args:
            enable_rerank: 是否启用重排序。
                - 如果明确传入 True/False，则使用该值
                - 如果为 None，则根据是否有可用的 rerank provider 自动决定
        """
        # 智能决定 enable_rerank 的默认值
        if enable_rerank is None:
            # 检查是否有可用的 rerank provider
            has_rerank_provider = (
                self.provider_manager
                and hasattr(self.provider_manager, 'rerank_provider_insts')
                and len(self.provider_manager.rerank_provider_insts) > 0
            )
            enable_rerank = has_rerank_provider

        kb = KnowledgeBase(
            kb_name=kb_name,
            description=description,
            emoji=emoji or "📚",
            embedding_provider_id=embedding_provider_id,
            rerank_provider_id=rerank_provider_id,
            chunk_size=chunk_size if chunk_size is not None else 512,
            chunk_overlap=chunk_overlap if chunk_overlap is not None else 50,
            top_k_dense=top_k_dense if top_k_dense is not None else 50,
            top_k_sparse=top_k_sparse if top_k_sparse is not None else 50,
            top_m_final=top_m_final if top_m_final is not None else 5,
            enable_rerank=enable_rerank,
        )
        async with self.db.get_db() as session:
            session.add(kb)
            await session.commit()
            await session.refresh(kb)
        return kb

    async def get_kb(self, kb_id: str) -> Optional[KnowledgeBase]:
        """获取知识库"""
        async with self.db.get_db() as session:
            stmt = select(KnowledgeBase).where(KnowledgeBase.kb_id == kb_id)
            result = await session.execute(stmt)
            return result.scalar_one_or_none()

    async def list_kbs(self, offset: int = 0, limit: int = 100) -> list[KnowledgeBase]:
        """列出所有知识库"""
        async with self.db.get_db() as session:
            stmt = (
                select(KnowledgeBase)
                .offset(offset)
                .limit(limit)
                .order_by(KnowledgeBase.created_at.desc())
            )
            result = await session.execute(stmt)
            return list(result.scalars().all())

    async def update_kb(
        self,
        kb_id: str,
        kb_name: Optional[str] = None,
        description: Optional[str] = None,
        emoji: Optional[str] = None,
        embedding_provider_id: Optional[str] = None,
        rerank_provider_id: Optional[str] = None,
        chunk_size: Optional[int] = None,
        chunk_overlap: Optional[int] = None,
        top_k_dense: Optional[int] = None,
        top_k_sparse: Optional[int] = None,
        top_m_final: Optional[int] = None,
        enable_rerank: Optional[bool] = None,
    ) -> Optional[KnowledgeBase]:
        """更新知识库"""
        async with self.db.get_db() as session:
            stmt = select(KnowledgeBase).where(KnowledgeBase.kb_id == kb_id)
            result = await session.execute(stmt)
            kb = result.scalar_one_or_none()
            if not kb:
                return None

            if kb_name is not None:
                kb.kb_name = kb_name
            if description is not None:
                kb.description = description
            if emoji is not None:
                kb.emoji = emoji
            if embedding_provider_id is not None:
                kb.embedding_provider_id = embedding_provider_id
            if rerank_provider_id is not None:
                kb.rerank_provider_id = rerank_provider_id
            if chunk_size is not None:
                kb.chunk_size = chunk_size
            if chunk_overlap is not None:
                kb.chunk_overlap = chunk_overlap
            if top_k_dense is not None:
                kb.top_k_dense = top_k_dense
            if top_k_sparse is not None:
                kb.top_k_sparse = top_k_sparse
            if top_m_final is not None:
                kb.top_m_final = top_m_final
            if enable_rerank is not None:
                kb.enable_rerank = enable_rerank

            await session.commit()
            await session.refresh(kb)
            return kb

    async def delete_kb(self, kb_id: str) -> bool:
        """删除知识库(级联删除所有文档和资源)"""
        # 1. 获取所有文档
        from astrbot.core.knowledge_base.manager_ops import KBManagerOps

        ops = KBManagerOps(self)
        docs = await ops.list_documents(kb_id)

        # 2. 删除所有文档(包括文件和向量)
        for doc in docs:
            await ops.delete_document(doc.doc_id)

        # 3. 删除知识库记录
        async with self.db.get_db() as session:
            stmt = select(KnowledgeBase).where(KnowledgeBase.kb_id == kb_id)
            result = await session.execute(stmt)
            kb = result.scalar_one_or_none()
            if not kb:
                return False

            await session.delete(kb)
            await session.commit()

        return True

    # ===== 文档上传 =====

    async def upload_document(
        self,
        kb_id: str,
        file_name: str,
        file_content: bytes,
        file_type: str,
    ) -> KBDocument:
        """上传并处理文档（带原子性保证和失败清理）

        流程:
        1. 保存原始文件
        2. 解析文档内容
        3. 提取多媒体资源
        4. 分块处理
        5. 生成向量并存储
        6. 保存元数据（事务）
        7. 更新统计
        """
        doc_id = str(uuid.uuid4())
        file_path = None
        media_paths = []
        vec_doc_ids = []

        try:
            # 1. 保存原始文件
            file_path = self.files_path / kb_id / f"{doc_id}.{file_type}"
            file_path.parent.mkdir(parents=True, exist_ok=True)

            async with aiofiles.open(file_path, "wb") as f:
                await f.write(file_content)

            # 2. 解析文档
            parser = self.parsers.get(file_type)
            if not parser:
                raise ValueError(f"不支持的文件类型: {file_type}")

            parse_result = await parser.parse(file_content, file_name)
            text_content = parse_result.text
            media_items = parse_result.media

            # 3. 保存多媒体资源
            from astrbot.core.knowledge_base.manager_ops import KBManagerOps

            ops = KBManagerOps(self)
            saved_media = []
            for media_item in media_items:
                media = await ops._save_media(
                    kb_id=kb_id,
                    doc_id=doc_id,
                    media_type=media_item.media_type,
                    file_name=media_item.file_name,
                    content=media_item.content,
                    mime_type=media_item.mime_type,
                )
                saved_media.append(media)
                media_paths.append(Path(media.file_path))

            # 4. 文档分块
            chunks_text = await self.chunker.chunk(text_content)

            # 5. 生成向量并存储
            saved_chunks = []
            for idx, chunk_text in enumerate(chunks_text):
                # 存储到向量数据库
                vec_doc_id = await self.vec_db.insert(
                    content=chunk_text,
                    metadata={
                        "kb_id": kb_id,
                        "doc_id": doc_id,
                        "chunk_index": idx,
                    },
                )
                vec_doc_ids.append(str(vec_doc_id))

                # 保存块元数据
                chunk = KBChunk(
                    doc_id=doc_id,
                    kb_id=kb_id,
                    chunk_index=idx,
                    content=chunk_text,
                    char_count=len(chunk_text),
                    vec_doc_id=str(vec_doc_id),
                )
                saved_chunks.append(chunk)

            # 6. 保存文档元数据（事务）
            doc = KBDocument(
                doc_id=doc_id,
                kb_id=kb_id,
                doc_name=file_name,
                file_type=file_type,
                file_size=len(file_content),
                file_path=str(file_path),
                chunk_count=len(saved_chunks),
                media_count=len(saved_media),
            )

            async with self.db.get_db() as session:
                async with session.begin():
                    session.add(doc)
                    for chunk in saved_chunks:
                        session.add(chunk)
                    for media in saved_media:
                        session.add(media)
                    await session.commit()

                await session.refresh(doc)

            # 7. 更新知识库统计
            await self._update_kb_stats(kb_id)

            return doc

        except Exception as e:
            # 失败清理：删除已创建的资源
            from astrbot.core import logger

            logger.error(f"文档上传失败，开始清理资源: {e}")

            # 清理向量数据库
            for vec_id in vec_doc_ids:
                try:
                    await self.vec_db.delete(vec_id)
                except Exception as ve:
                    logger.warning(f"清理向量失败 {vec_id}: {ve}")

            # 清理多媒体文件
            for media_path in media_paths:
                try:
                    if media_path.exists():
                        media_path.unlink()
                except Exception as me:
                    logger.warning(f"清理多媒体文件失败 {media_path}: {me}")

            # 清理文档文件
            if file_path and file_path.exists():
                try:
                    file_path.unlink()
                except Exception as fe:
                    logger.warning(f"清理文档文件失败 {file_path}: {fe}")

            # 重新抛出原始异常
            raise

    # ===== 统计更新 =====

    async def _update_kb_stats(self, kb_id: str):
        """更新知识库统计信息（事务中执行）"""
        async with self.db.get_db() as session:
            async with session.begin():
                # 统计文档数（在事务中查询）
                doc_count = (
                    await session.scalar(
                        select(func.count(KBDocument.id)).where(
                            KBDocument.kb_id == kb_id
                        )
                    )
                    or 0
                )

                # 统计块数（在事务中查询）
                chunk_count = (
                    await session.scalar(
                        select(func.count(KBChunk.id)).where(KBChunk.kb_id == kb_id)
                    )
                    or 0
                )

                # 更新知识库（在同一事务中）
                await session.execute(
                    update(KnowledgeBase)
                    .where(KnowledgeBase.kb_id == kb_id)
                    .values(doc_count=doc_count, chunk_count=chunk_count)
                )

                await session.commit()