AstrBot 4.3.3__py3-none-any.whl → 4.5.0__py3-none-any.whl

astrbot/core/knowledge_base/kb_mgr.py

@@ -0,0 +1,287 @@
+import traceback
+from pathlib import Path
+from astrbot.core import logger
+from astrbot.core.provider.manager import ProviderManager
+
+from .retrieval.manager import RetrievalManager, RetrievalResult
+from .retrieval.sparse_retriever import SparseRetriever
+from .retrieval.rank_fusion import RankFusion
+from .kb_db_sqlite import KBSQLiteDatabase
+
+# from .chunking.fixed_size import FixedSizeChunker
+from .chunking.recursive import RecursiveCharacterChunker
+from .kb_helper import KBHelper
+
+from .models import KnowledgeBase
+
+
+FILES_PATH = "data/knowledge_base"
+DB_PATH = Path(FILES_PATH) / "kb.db"
+"""Knowledge Base storage root directory"""
+CHUNKER = RecursiveCharacterChunker()
+
+
+class KnowledgeBaseManager:
+    kb_db: KBSQLiteDatabase
+    retrieval_manager: RetrievalManager
+
+    def __init__(
+        self,
+        provider_manager: ProviderManager,
+    ):
+        Path(DB_PATH).parent.mkdir(parents=True, exist_ok=True)
+        self.provider_manager = provider_manager
+        self._session_deleted_callback_registered = False
+
+        self.kb_insts: dict[str, KBHelper] = {}
+
+    async def initialize(self):
+        """初始化知识库模块"""
+        try:
+            logger.info("正在初始化知识库模块...")
+
+            # 初始化数据库
+            await self._init_kb_database()
+
+            # 初始化检索管理器
+            sparse_retriever = SparseRetriever(self.kb_db)
+            rank_fusion = RankFusion(self.kb_db)
+            self.retrieval_manager = RetrievalManager(
+                sparse_retriever=sparse_retriever,
+                rank_fusion=rank_fusion,
+                kb_db=self.kb_db,
+            )
+            await self.load_kbs()
+
+        except ImportError as e:
+            logger.error(f"知识库模块导入失败: {e}")
+            logger.warning("请确保已安装所需依赖: pypdf, aiofiles, Pillow, rank-bm25")
+        except Exception as e:
+            logger.error(f"知识库模块初始化失败: {e}")
+            logger.error(traceback.format_exc())
+
+    async def _init_kb_database(self):
+        self.kb_db = KBSQLiteDatabase(DB_PATH.as_posix())
+        await self.kb_db.initialize()
+        await self.kb_db.migrate_to_v1()
+        logger.info(f"KnowledgeBase database initialized: {DB_PATH}")
+
+    async def load_kbs(self):
+        """加载所有知识库实例"""
+        kb_records = await self.kb_db.list_kbs()
+        for record in kb_records:
+            kb_helper = KBHelper(
+                kb_db=self.kb_db,
+                kb=record,
+                provider_manager=self.provider_manager,
+                kb_root_dir=FILES_PATH,
+                chunker=CHUNKER,
+            )
+            await kb_helper.initialize()
+            self.kb_insts[record.kb_id] = kb_helper
+
+    async def create_kb(
+        self,
+        kb_name: str,
+        description: str | None = None,
+        emoji: str | None = None,
+        embedding_provider_id: str | None = None,
+        rerank_provider_id: str | None = None,
+        chunk_size: int | None = None,
+        chunk_overlap: int | None = None,
+        top_k_dense: int | None = None,
+        top_k_sparse: int | None = None,
+        top_m_final: int | None = None,
+    ) -> KBHelper:
+        """创建新的知识库实例"""
+        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,
+        )
+        async with self.kb_db.get_db() as session:
+            session.add(kb)
+            await session.commit()
+            await session.refresh(kb)
+
+            kb_helper = KBHelper(
+                kb_db=self.kb_db,
+                kb=kb,
+                provider_manager=self.provider_manager,
+                kb_root_dir=FILES_PATH,
+                chunker=CHUNKER,
+            )
+            await kb_helper.initialize()
+        self.kb_insts[kb.kb_id] = kb_helper
+        return kb_helper
+
+    async def get_kb(self, kb_id: str) -> KBHelper | None:
+        """获取知识库实例"""
+        if kb_id in self.kb_insts:
+            return self.kb_insts[kb_id]
+
+    async def get_kb_by_name(self, kb_name: str) -> KBHelper | None:
+        """通过名称获取知识库实例"""
+        for kb_helper in self.kb_insts.values():
+            if kb_helper.kb.kb_name == kb_name:
+                return kb_helper
+        return None
+
+    async def delete_kb(self, kb_id: str) -> bool:
+        """删除知识库实例"""
+        kb_helper = await self.get_kb(kb_id)
+        if not kb_helper:
+            return False
+
+        await kb_helper.delete_vec_db()
+        async with self.kb_db.get_db() as session:
+            await session.delete(kb_helper.kb)
+            await session.commit()
+
+        self.kb_insts.pop(kb_id, None)
+        return True
+
+    async def list_kbs(self) -> list[KnowledgeBase]:
+        """列出所有知识库实例"""
+        kbs = [kb_helper.kb for kb_helper in self.kb_insts.values()]
+        return kbs
+
+    async def update_kb(
+        self,
+        kb_id: str,
+        kb_name: str,
+        description: str | None = None,
+        emoji: str | None = None,
+        embedding_provider_id: str | None = None,
+        rerank_provider_id: str | None = None,
+        chunk_size: int | None = None,
+        chunk_overlap: int | None = None,
+        top_k_dense: int | None = None,
+        top_k_sparse: int | None = None,
+        top_m_final: int | None = None,
+    ) -> KBHelper | None:
+        """更新知识库实例"""
+        kb_helper = await self.get_kb(kb_id)
+        if not kb_helper:
+            return None
+
+        kb = kb_helper.kb
+        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
+        kb.rerank_provider_id = rerank_provider_id  # 允许设置为 None
+        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
+        async with self.kb_db.get_db() as session:
+            session.add(kb)
+            await session.commit()
+            await session.refresh(kb)
+
+        return kb_helper
+
+    async def retrieve(
+        self,
+        query: str,
+        kb_names: list[str],
+        top_k_fusion: int = 20,
+        top_m_final: int = 5,
+    ) -> dict | None:
+        """从指定知识库中检索相关内容"""
+        kb_ids = []
+        kb_id_helper_map = {}
+        for kb_name in kb_names:
+            if kb_helper := await self.get_kb_by_name(kb_name):
+                kb_ids.append(kb_helper.kb.kb_id)
+                kb_id_helper_map[kb_helper.kb.kb_id] = kb_helper
+
+        if not kb_ids:
+            return {}
+
+        results = await self.retrieval_manager.retrieve(
+            query=query,
+            kb_ids=kb_ids,
+            kb_id_helper_map=kb_id_helper_map,
+            top_k_fusion=top_k_fusion,
+            top_m_final=top_m_final,
+        )
+        if not results:
+            return None
+
+        context_text = self._format_context(results)
+
+        results_dict = [
+            {
+                "chunk_id": r.chunk_id,
+                "doc_id": r.doc_id,
+                "kb_id": r.kb_id,
+                "kb_name": r.kb_name,
+                "doc_name": r.doc_name,
+                "chunk_index": r.metadata.get("chunk_index", 0),
+                "content": r.content,
+                "score": r.score,
+                "char_count": r.metadata.get("char_count", 0),
+            }
+            for r in results
+        ]
+
+        return {
+            "context_text": context_text,
+            "results": results_dict,
+        }
+
+    def _format_context(self, results: list[RetrievalResult]) -> str:
+        """格式化知识上下文
+
+        Args:
+            results: 检索结果列表
+
+        Returns:
+            str: 格式化的上下文文本
+        """
+        lines = ["以下是相关的知识库内容,请参考这些信息回答用户的问题:\n"]
+
+        for i, result in enumerate(results, 1):
+            lines.append(f"【知识 {i}】")
+            lines.append(f"来源: {result.kb_name} / {result.doc_name}")
+            lines.append(f"内容: {result.content}")
+            lines.append(f"相关度: {result.score:.2f}")
+            lines.append("")
+
+        return "\n".join(lines)
+
+    async def terminate(self):
+        """终止所有知识库实例,关闭数据库连接"""
+        for kb_id, kb_helper in self.kb_insts.items():
+            try:
+                await kb_helper.terminate()
+            except Exception as e:
+                logger.error(f"关闭知识库 {kb_id} 失败: {e}")
+
+        self.kb_insts.clear()
+
+        # 关闭元数据数据库
+        if hasattr(self, "kb_db") and self.kb_db:
+            try:
+                await self.kb_db.close()
+            except Exception as e:
+                logger.error(f"关闭知识库元数据数据库失败: {e}")

astrbot/core/knowledge_base/models.py

@@ -0,0 +1,114 @@
+import uuid
+from datetime import datetime, timezone
+
+from sqlmodel import Field, SQLModel, Text, UniqueConstraint, MetaData
+
+
+class BaseKBModel(SQLModel, table=False):
+    metadata = MetaData()
+
+
+class KnowledgeBase(BaseKBModel, table=True):
+    """知识库表
+
+    存储知识库的基本信息和统计数据。
+    """
+
+    __tablename__ = "knowledge_bases"  # type: ignore
+
+    id: int | None = Field(
+        primary_key=True, sa_column_kwargs={"autoincrement": True}, default=None
+    )
+    kb_id: str = Field(
+        max_length=36,
+        nullable=False,
+        unique=True,
+        default_factory=lambda: str(uuid.uuid4()),
+        index=True,
+    )
+    kb_name: str = Field(max_length=100, nullable=False)
+    description: str | None = Field(default=None, sa_type=Text)
+    emoji: str | None = Field(default="📚", max_length=10)
+    embedding_provider_id: str | None = Field(default=None, max_length=100)
+    rerank_provider_id: str | None = Field(default=None, max_length=100)
+    # 分块配置参数
+    chunk_size: int | None = Field(default=512, nullable=True)
+    chunk_overlap: int | None = Field(default=50, nullable=True)
+    # 检索配置参数
+    top_k_dense: int | None = Field(default=50, nullable=True)
+    top_k_sparse: int | None = Field(default=50, nullable=True)
+    top_m_final: int | None = Field(default=5, nullable=True)
+    created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
+    updated_at: datetime = Field(
+        default_factory=lambda: datetime.now(timezone.utc),
+        sa_column_kwargs={"onupdate": datetime.now(timezone.utc)},
+    )
+    doc_count: int = Field(default=0, nullable=False)
+    chunk_count: int = Field(default=0, nullable=False)
+
+    __table_args__ = (
+        UniqueConstraint(
+            "kb_name",
+            name="uix_kb_name",
+        ),
+    )
+
+
+class KBDocument(BaseKBModel, table=True):
+    """文档表
+
+    存储上传到知识库的文档元数据。
+    """
+
+    __tablename__ = "kb_documents"  # type: ignore
+
+    id: int | None = Field(
+        primary_key=True, sa_column_kwargs={"autoincrement": True}, default=None
+    )
+    doc_id: str = Field(
+        max_length=36,
+        nullable=False,
+        unique=True,
+        default_factory=lambda: str(uuid.uuid4()),
+        index=True,
+    )
+    kb_id: str = Field(max_length=36, nullable=False, index=True)
+    doc_name: str = Field(max_length=255, nullable=False)
+    file_type: str = Field(max_length=20, nullable=False)
+    file_size: int = Field(nullable=False)
+    file_path: str = Field(max_length=512, nullable=False)
+    chunk_count: int = Field(default=0, nullable=False)
+    media_count: int = Field(default=0, nullable=False)
+    created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
+    updated_at: datetime = Field(
+        default_factory=lambda: datetime.now(timezone.utc),
+        sa_column_kwargs={"onupdate": datetime.now(timezone.utc)},
+    )
+
+
+class KBMedia(BaseKBModel, table=True):
+    """多媒体资源表
+
+    存储从文档中提取的图片、视频等多媒体资源。
+    """
+
+    __tablename__ = "kb_media"  # type: ignore
+
+    id: int | None = Field(
+        primary_key=True, sa_column_kwargs={"autoincrement": True}, default=None
+    )
+    media_id: str = Field(
+        max_length=36,
+        nullable=False,
+        unique=True,
+        default_factory=lambda: str(uuid.uuid4()),
+        index=True,
+    )
+    doc_id: str = Field(max_length=36, nullable=False, index=True)
+    kb_id: str = Field(max_length=36, nullable=False, index=True)
+    media_type: str = Field(max_length=20, nullable=False)
+    file_name: str = Field(max_length=255, nullable=False)
+    file_path: str = Field(max_length=512, nullable=False)
+    file_size: int = Field(nullable=False)
+    mime_type: str = Field(max_length=100, nullable=False)
+    created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))

astrbot/core/knowledge_base/parsers/__init__.py

@@ -0,0 +1,15 @@
+"""
+文档解析器模块
+"""
+
+from .base import BaseParser, MediaItem, ParseResult
+from .text_parser import TextParser
+from .pdf_parser import PDFParser
+
+__all__ = [
+    "BaseParser",
+    "MediaItem",
+    "ParseResult",
+    "TextParser",
+    "PDFParser",
+]

astrbot/core/knowledge_base/parsers/base.py

@@ -0,0 +1,50 @@
+"""文档解析器基类和数据结构
+
+定义了文档解析器的抽象接口和相关数据类。
+"""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+
+
+@dataclass
+class MediaItem:
+    """多媒体项
+
+    表示从文档中提取的多媒体资源。
+    """
+
+    media_type: str  # image, video
+    file_name: str
+    content: bytes
+    mime_type: str
+
+
+@dataclass
+class ParseResult:
+    """解析结果
+
+    包含解析后的文本内容和提取的多媒体资源。
+    """
+
+    text: str
+    media: list[MediaItem]
+
+
+class BaseParser(ABC):
+    """文档解析器基类
+
+    所有文档解析器都应该继承此类并实现 parse 方法。
+    """
+
+    @abstractmethod
+    async def parse(self, file_content: bytes, file_name: str) -> ParseResult:
+        """解析文档
+
+        Args:
+            file_content: 文件内容
+            file_name: 文件名
+
+        Returns:
+            ParseResult: 解析结果
+        """

astrbot/core/knowledge_base/parsers/markitdown_parser.py

@@ -0,0 +1,25 @@
+import io
+import os
+
+from astrbot.core.knowledge_base.parsers.base import (
+    BaseParser,
+    ParseResult,
+)
+from markitdown_no_magika import MarkItDown, StreamInfo
+
+
+class MarkitdownParser(BaseParser):
+    """解析 docx, xls, xlsx 格式"""
+
+    async def parse(self, file_content: bytes, file_name: str) -> ParseResult:
+        md = MarkItDown(enable_plugins=False)
+        bio = io.BytesIO(file_content)
+        stream_info = StreamInfo(
+            extension=os.path.splitext(file_name)[1].lower(),
+            filename=file_name,
+        )
+        result = md.convert(bio, stream_info=stream_info)
+        return ParseResult(
+            text=result.markdown,
+            media=[],
+        )

astrbot/core/knowledge_base/parsers/pdf_parser.py

@@ -0,0 +1,100 @@
+"""PDF 文件解析器
+
+支持解析 PDF 文件中的文本和图片资源。
+"""
+
+import io
+
+from pypdf import PdfReader
+
+from astrbot.core.knowledge_base.parsers.base import (
+    BaseParser,
+    MediaItem,
+    ParseResult,
+)
+
+
+class PDFParser(BaseParser):
+    """PDF 文档解析器
+
+    提取 PDF 中的文本内容和嵌入的图片资源。
+    """
+
+    async def parse(self, file_content: bytes, file_name: str) -> ParseResult:
+        """解析 PDF 文件
+
+        Args:
+            file_content: 文件内容
+            file_name: 文件名
+
+        Returns:
+            ParseResult: 包含文本和图片的解析结果
+        """
+        pdf_file = io.BytesIO(file_content)
+        reader = PdfReader(pdf_file)
+
+        text_parts = []
+        media_items = []
+
+        # 提取文本
+        for page in reader.pages:
+            text = page.extract_text()
+            if text:
+                text_parts.append(text)
+
+        # 提取图片
+        image_counter = 0
+        for page_num, page in enumerate(reader.pages):
+            try:
+                # 安全检查 Resources
+                if "/Resources" not in page:
+                    continue
+
+                resources = page["/Resources"]
+                if not resources or "/XObject" not in resources:  # type: ignore
+                    continue
+
+                xobjects = resources["/XObject"].get_object()  # type: ignore
+                if not xobjects:
+                    continue
+
+                for obj_name in xobjects:
+                    try:
+                        obj = xobjects[obj_name]
+
+                        if obj.get("/Subtype") != "/Image":
+                            continue
+
+                        # 提取图片数据
+                        image_data = obj.get_data()
+
+                        # 确定格式
+                        filter_type = obj.get("/Filter", "")
+                        if filter_type == "/DCTDecode":
+                            ext = "jpg"
+                            mime_type = "image/jpeg"
+                        elif filter_type == "/FlateDecode":
+                            ext = "png"
+                            mime_type = "image/png"
+                        else:
+                            ext = "png"
+                            mime_type = "image/png"
+
+                        image_counter += 1
+                        media_items.append(
+                            MediaItem(
+                                media_type="image",
+                                file_name=f"page_{page_num}_img_{image_counter}.{ext}",
+                                content=image_data,
+                                mime_type=mime_type,
+                            )
+                        )
+                    except Exception:
+                        # 单个图片提取失败不影响整体
+                        continue
+            except Exception:
+                # 页面处理失败不影响其他页面
+                continue
+
+        full_text = "\n\n".join(text_parts)
+        return ParseResult(text=full_text, media=media_items)

astrbot/core/knowledge_base/parsers/text_parser.py

@@ -0,0 +1,41 @@
+"""文本文件解析器
+
+支持解析 TXT 和 Markdown 文件。
+"""
+
+from astrbot.core.knowledge_base.parsers.base import BaseParser, ParseResult
+
+
+class TextParser(BaseParser):
+    """TXT/MD 文本解析器
+
+    支持多种字符编码的自动检测。
+    """
+
+    async def parse(self, file_content: bytes, file_name: str) -> ParseResult:
+        """解析文本文件
+
+        尝试使用多种编码解析文件内容。
+
+        Args:
+            file_content: 文件内容
+            file_name: 文件名
+
+        Returns:
+            ParseResult: 解析结果,不包含多媒体资源
+
+        Raises:
+            ValueError: 如果无法解码文件
+        """
+        # 尝试多种编码
+        for encoding in ["utf-8", "gbk", "gb2312", "gb18030"]:
+            try:
+                text = file_content.decode(encoding)
+                break
+            except UnicodeDecodeError:
+                continue
+        else:
+            raise ValueError(f"无法解码文件: {file_name}")
+
+        # 文本文件无多媒体资源
+        return ParseResult(text=text, media=[])

astrbot/core/knowledge_base/parsers/util.py

@@ -0,0 +1,13 @@
+from .base import BaseParser
+
+
+async def select_parser(ext: str) -> BaseParser:
+    if ext in {".md", ".txt", ".markdown", ".xlsx", ".docx", ".xls"}:
+        from .markitdown_parser import MarkitdownParser
+
+        return MarkitdownParser()
+    elif ext == ".pdf":
+        from .pdf_parser import PDFParser
+
+        return PDFParser()
+    raise ValueError(f"暂时不支持的文件格式: {ext}")

astrbot/core/knowledge_base/retrieval/__init__.py

@@ -0,0 +1,16 @@
+"""
+检索模块
+"""
+
+from .manager import RetrievalManager, RetrievalResult
+from .sparse_retriever import SparseRetriever, SparseResult
+from .rank_fusion import RankFusion, FusedResult
+
+__all__ = [
+    "RetrievalManager",
+    "RetrievalResult",
+    "SparseRetriever",
+    "SparseResult",
+    "RankFusion",
+    "FusedResult",
+]