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

astrbot/core/knowledge_base/chunking/fixed_size.py

@@ -18,6 +18,7 @@ class FixedSizeChunker(BaseChunker):
         Args:
             chunk_size: 块的大小(字符数)
             chunk_overlap: 块之间的重叠字符数
+
         """
         self.chunk_size = chunk_size
         self.chunk_overlap = chunk_overlap
@@ -32,6 +33,7 @@ class FixedSizeChunker(BaseChunker):
 
         Returns:
             list[str]: 分块后的文本列表
+
         """
         chunk_size = kwargs.get("chunk_size", self.chunk_size)
         chunk_overlap = kwargs.get("chunk_overlap", self.chunk_overlap)

astrbot/core/knowledge_base/chunking/recursive.py

@@ -1,4 +1,5 @@
 from collections.abc import Callable
+
 from .base import BaseChunker
 
 
@@ -11,8 +12,7 @@ class RecursiveCharacterChunker(BaseChunker):
         is_separator_regex: bool = False,
         separators: list[str] | None = None,
     ):
-        """
-        初始化递归字符文本分割器
+        """初始化递归字符文本分割器
 
         Args:
             chunk_size: 每个文本块的最大大小
@@ -20,6 +20,7 @@ class RecursiveCharacterChunker(BaseChunker):
             length_function: 计算文本长度的函数
             is_separator_regex: 分隔符是否为正则表达式
             separators: 用于分割文本的分隔符列表，按优先级排序
+
         """
         self.chunk_size = chunk_size
         self.chunk_overlap = chunk_overlap
@@ -39,8 +40,7 @@ class RecursiveCharacterChunker(BaseChunker):
         ]
 
     async def chunk(self, text: str, **kwargs) -> list[str]:
-        """
-        递归地将文本分割成块
+        """递归地将文本分割成块
 
         Args:
             text: 要分割的文本
@@ -49,6 +49,7 @@ class RecursiveCharacterChunker(BaseChunker):
 
         Returns:
             分割后的文本块列表
+
         """
         if not text:
             return []
@@ -90,7 +91,7 @@ class RecursiveCharacterChunker(BaseChunker):
                                     combined_text,
                                     chunk_size=chunk_size,
                                     chunk_overlap=overlap,
-                                )
+                                ),
                             )
                             current_chunk = []
                             current_chunk_length = 0
@@ -98,8 +99,10 @@ class RecursiveCharacterChunker(BaseChunker):
                         # 递归分割过大的部分
                         final_chunks.extend(
                             await self.chunk(
-                                split, chunk_size=chunk_size, chunk_overlap=overlap
-                            )
+                                split,
+                                chunk_size=chunk_size,
+                                chunk_overlap=overlap,
+                            ),
                         )
                     # 如果添加这部分会使当前块超过chunk_size
                     elif current_chunk_length + split_length > chunk_size:
@@ -132,16 +135,19 @@ class RecursiveCharacterChunker(BaseChunker):
         return [text]
 
     def _split_by_character(
-        self, text: str, chunk_size: int | None = None, overlap: int | None = None
+        self,
+        text: str,
+        chunk_size: int | None = None,
+        overlap: int | None = None,
     ) -> list[str]:
-        """
-        按字符级别分割文本
+        """按字符级别分割文本
 
         Args:
             text: 要分割的文本
 
         Returns:
             分割后的文本块列表
+
         """
         chunk_size = chunk_size or self.chunk_size
         overlap = overlap or self.chunk_overlap

astrbot/core/knowledge_base/kb_db_sqlite.py

@@ -1,18 +1,18 @@
 from contextlib import asynccontextmanager
 from pathlib import Path
 
-from sqlmodel import col, desc
-from sqlalchemy import text, func, select, update, delete
+from sqlalchemy import delete, func, select, text, update
 from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
+from sqlmodel import col, desc
 
 from astrbot.core import logger
+from astrbot.core.db.vec_db.faiss_impl import FaissVecDB
 from astrbot.core.knowledge_base.models import (
     BaseKBModel,
     KBDocument,
     KBMedia,
     KnowledgeBase,
 )
-from astrbot.core.db.vec_db.faiss_impl import FaissVecDB
 
 
 class KBSQLiteDatabase:
@@ -21,6 +21,7 @@ class KBSQLiteDatabase:
 
         Args:
             db_path: 数据库文件路径, 默认为 data/knowledge_base/kb.db
+
         """
         self.db_path = db_path
         self.DATABASE_URL = f"sqlite+aiosqlite:///{db_path}"
@@ -85,77 +86,77 @@ class KBSQLiteDatabase:
                 await session.execute(
                     text(
                         "CREATE INDEX IF NOT EXISTS idx_kb_kb_id "
-                        "ON knowledge_bases(kb_id)"
-                    )
+                        "ON knowledge_bases(kb_id)",
+                    ),
                 )
                 await session.execute(
                     text(
                         "CREATE INDEX IF NOT EXISTS idx_kb_name "
-                        "ON knowledge_bases(kb_name)"
-                    )
+                        "ON knowledge_bases(kb_name)",
+                    ),
                 )
                 await session.execute(
                     text(
                         "CREATE INDEX IF NOT EXISTS idx_kb_created_at "
-                        "ON knowledge_bases(created_at)"
-                    )
+                        "ON knowledge_bases(created_at)",
+                    ),
                 )
 
                 # 创建文档表索引
                 await session.execute(
                     text(
                         "CREATE INDEX IF NOT EXISTS idx_doc_doc_id "
-                        "ON kb_documents(doc_id)"
-                    )
+                        "ON kb_documents(doc_id)",
+                    ),
                 )
                 await session.execute(
                     text(
                         "CREATE INDEX IF NOT EXISTS idx_doc_kb_id "
-                        "ON kb_documents(kb_id)"
-                    )
+                        "ON kb_documents(kb_id)",
+                    ),
                 )
                 await session.execute(
                     text(
                         "CREATE INDEX IF NOT EXISTS idx_doc_name "
-                        "ON kb_documents(doc_name)"
-                    )
+                        "ON kb_documents(doc_name)",
+                    ),
                 )
                 await session.execute(
                     text(
                         "CREATE INDEX IF NOT EXISTS idx_doc_type "
-                        "ON kb_documents(file_type)"
-                    )
+                        "ON kb_documents(file_type)",
+                    ),
                 )
                 await session.execute(
                     text(
                         "CREATE INDEX IF NOT EXISTS idx_doc_created_at "
-                        "ON kb_documents(created_at)"
-                    )
+                        "ON kb_documents(created_at)",
+                    ),
                 )
 
                 # 创建多媒体表索引
                 await session.execute(
                     text(
                         "CREATE INDEX IF NOT EXISTS idx_media_media_id "
-                        "ON kb_media(media_id)"
-                    )
+                        "ON kb_media(media_id)",
+                    ),
                 )
                 await session.execute(
                     text(
                         "CREATE INDEX IF NOT EXISTS idx_media_doc_id "
-                        "ON kb_media(doc_id)"
-                    )
+                        "ON kb_media(doc_id)",
+                    ),
                 )
                 await session.execute(
                     text(
-                        "CREATE INDEX IF NOT EXISTS idx_media_kb_id ON kb_media(kb_id)"
-                    )
+                        "CREATE INDEX IF NOT EXISTS idx_media_kb_id ON kb_media(kb_id)",
+                    ),
                 )
                 await session.execute(
                     text(
                         "CREATE INDEX IF NOT EXISTS idx_media_type "
-                        "ON kb_media(media_type)"
-                    )
+                        "ON kb_media(media_type)",
+                    ),
                 )
 
                 await session.commit()
@@ -208,7 +209,10 @@ class KBSQLiteDatabase:
             return result.scalar_one_or_none()
 
     async def list_documents_by_kb(
-        self, kb_id: str, offset: int = 0, limit: int = 100
+        self,
+        kb_id: str,
+        offset: int = 0,
+        limit: int = 100,
     ) -> list[KBDocument]:
         """列出知识库的所有文档"""
         async with self.get_db() as session:
@@ -226,7 +230,7 @@ class KBSQLiteDatabase:
         """统计知识库的文档数量"""
         async with self.get_db() as session:
             stmt = select(func.count(col(KBDocument.id))).where(
-                col(KBDocument.kb_id) == kb_id
+                col(KBDocument.kb_id) == kb_id,
             )
             result = await session.execute(stmt)
             return result.scalar() or 0
@@ -252,12 +256,11 @@ class KBSQLiteDatabase:
     async def delete_document_by_id(self, doc_id: str, vec_db: FaissVecDB):
         """删除单个文档及其相关数据"""
         # 在知识库表中删除
-        async with self.get_db() as session:
-            async with session.begin():
-                # 删除文档记录
-                delete_stmt = delete(KBDocument).where(col(KBDocument.doc_id) == doc_id)
-                await session.execute(delete_stmt)
-                await session.commit()
+        async with self.get_db() as session, session.begin():
+            # 删除文档记录
+            delete_stmt = delete(KBDocument).where(col(KBDocument.doc_id) == doc_id)
+            await session.execute(delete_stmt)
+            await session.commit()
 
         # 在 vec db 中删除相关向量
         await vec_db.delete_documents(metadata_filters={"kb_doc_id": doc_id})
@@ -282,18 +285,17 @@ class KBSQLiteDatabase:
         """更新知识库统计信息"""
         chunk_cnt = await vec_db.count_documents()
 
-        async with self.get_db() as session:
-            async with session.begin():
-                update_stmt = (
-                    update(KnowledgeBase)
-                    .where(col(KnowledgeBase.kb_id) == kb_id)
-                    .values(
-                        doc_count=select(func.count(col(KBDocument.id)))
-                        .where(col(KBDocument.kb_id) == kb_id)
-                        .scalar_subquery(),
-                        chunk_count=chunk_cnt,
-                    )
+        async with self.get_db() as session, session.begin():
+            update_stmt = (
+                update(KnowledgeBase)
+                .where(col(KnowledgeBase.kb_id) == kb_id)
+                .values(
+                    doc_count=select(func.count(col(KBDocument.id)))
+                    .where(col(KBDocument.kb_id) == kb_id)
+                    .scalar_subquery(),
+                    chunk_count=chunk_cnt,
                 )
+            )
 
-                await session.execute(update_stmt)
-                await session.commit()
+            await session.execute(update_stmt)
+            await session.commit()

astrbot/core/knowledge_base/kb_helper.py

@@ -1,16 +1,19 @@
-import uuid
-import aiofiles
 import json
+import uuid
 from pathlib import Path
-from .models import KnowledgeBase, KBDocument, KBMedia
-from .kb_db_sqlite import KBSQLiteDatabase
+
+import aiofiles
+
+from astrbot.core import logger
 from astrbot.core.db.vec_db.base import BaseVecDB
 from astrbot.core.db.vec_db.faiss_impl.vec_db import FaissVecDB
-from astrbot.core.provider.provider import EmbeddingProvider, RerankProvider
 from astrbot.core.provider.manager import ProviderManager
-from .parsers.util import select_parser
+from astrbot.core.provider.provider import EmbeddingProvider, RerankProvider
+
 from .chunking.base import BaseChunker
-from astrbot.core import logger
+from .kb_db_sqlite import KBSQLiteDatabase
+from .models import KBDocument, KBMedia, KnowledgeBase
+from .parsers.util import select_parser
 
 
 class KBHelper:
@@ -45,11 +48,11 @@ class KBHelper:
         if not self.kb.embedding_provider_id:
             raise ValueError(f"知识库 {self.kb.kb_name} 未配置 Embedding Provider")
         ep: EmbeddingProvider = await self.prov_mgr.get_provider_by_id(
-            self.kb.embedding_provider_id
+            self.kb.embedding_provider_id,
         )  # type: ignore
         if not ep:
             raise ValueError(
-                f"无法找到 ID 为 {self.kb.embedding_provider_id} 的 Embedding Provider"
+                f"无法找到 ID 为 {self.kb.embedding_provider_id} 的 Embedding Provider",
             )
         return ep
 
@@ -57,11 +60,11 @@ class KBHelper:
         if not self.kb.rerank_provider_id:
             return None
         rp: RerankProvider = await self.prov_mgr.get_provider_by_id(
-            self.kb.rerank_provider_id
+            self.kb.rerank_provider_id,
         )  # type: ignore
         if not rp:
             raise ValueError(
-                f"无法找到 ID 为 {self.kb.rerank_provider_id} 的 Rerank Provider"
+                f"无法找到 ID 为 {self.kb.rerank_provider_id} 的 Rerank Provider",
             )
         return rp
 
@@ -122,6 +125,7 @@ class KBHelper:
                 - stage: 当前阶段 ('parsing', 'chunking', 'embedding')
                 - current: 当前进度
                 - total: 总数
+
         """
         await self._ensure_vec_db()
         doc_id = str(uuid.uuid4())
@@ -162,7 +166,9 @@ class KBHelper:
                 await progress_callback("chunking", 0, 100)
 
             chunks_text = await self.chunker.chunk(
-                text_content, chunk_size=chunk_size, chunk_overlap=chunk_overlap
+                text_content,
+                chunk_size=chunk_size,
+                chunk_overlap=chunk_overlap,
             )
             contents = []
             metadatas = []
@@ -173,7 +179,7 @@ class KBHelper:
                         "kb_id": self.kb.kb_id,
                         "kb_doc_id": doc_id,
                         "chunk_index": idx,
-                    }
+                    },
                 )
 
             if progress_callback:
@@ -234,7 +240,9 @@ class KBHelper:
             raise e
 
     async def list_documents(
-        self, offset: int = 0, limit: int = 100
+        self,
+        offset: int = 0,
+        limit: int = 100,
     ) -> list[KBDocument]:
         """列出知识库的所有文档"""
         docs = await self.kb_db.list_documents_by_kb(self.kb.kb_id, offset, limit)
@@ -288,12 +296,17 @@ class KBHelper:
             await session.refresh(doc)
 
     async def get_chunks_by_doc_id(
-        self, doc_id: str, offset: int = 0, limit: int = 100
+        self,
+        doc_id: str,
+        offset: int = 0,
+        limit: int = 100,
     ) -> list[dict]:
         """获取文档的所有块及其元数据"""
         vec_db: FaissVecDB = self.vec_db  # type: ignore
         chunks = await vec_db.document_storage.get_documents(
-            metadata_filters={"kb_doc_id": doc_id}, offset=offset, limit=limit
+            metadata_filters={"kb_doc_id": doc_id},
+            offset=offset,
+            limit=limit,
         )
         result = []
         for chunk in chunks:
@@ -306,7 +319,7 @@ class KBHelper:
                     "chunk_index": chunk_md["chunk_index"],
                     "content": chunk["text"],
                     "char_count": len(chunk["text"]),
-                }
+                },
             )
         return result
 

astrbot/core/knowledge_base/kb_mgr.py

@@ -1,19 +1,17 @@
 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_db_sqlite import KBSQLiteDatabase
 from .kb_helper import KBHelper
-
 from .models import KnowledgeBase
-
+from .retrieval.manager import RetrievalManager, RetrievalResult
+from .retrieval.rank_fusion import RankFusion
+from .retrieval.sparse_retriever import SparseRetriever
 
 FILES_PATH = "data/knowledge_base"
 DB_PATH = Path(FILES_PATH) / "kb.db"
@@ -257,6 +255,7 @@ class KnowledgeBaseManager:
 
         Returns:
             str: 格式化的上下文文本
+
         """
         lines = ["以下是相关的知识库内容,请参考这些信息回答用户的问题:\n"]
 

astrbot/core/knowledge_base/models.py

@@ -1,7 +1,7 @@
 import uuid
 from datetime import datetime, timezone
 
-from sqlmodel import Field, SQLModel, Text, UniqueConstraint, MetaData
+from sqlmodel import Field, MetaData, SQLModel, Text, UniqueConstraint
 
 
 class BaseKBModel(SQLModel, table=False):
@@ -17,7 +17,9 @@ class KnowledgeBase(BaseKBModel, table=True):
     __tablename__ = "knowledge_bases"  # type: ignore
 
     id: int | None = Field(
-        primary_key=True, sa_column_kwargs={"autoincrement": True}, default=None
+        primary_key=True,
+        sa_column_kwargs={"autoincrement": True},
+        default=None,
     )
     kb_id: str = Field(
         max_length=36,
@@ -63,7 +65,9 @@ class KBDocument(BaseKBModel, table=True):
     __tablename__ = "kb_documents"  # type: ignore
 
     id: int | None = Field(
-        primary_key=True, sa_column_kwargs={"autoincrement": True}, default=None
+        primary_key=True,
+        sa_column_kwargs={"autoincrement": True},
+        default=None,
     )
     doc_id: str = Field(
         max_length=36,
@@ -95,7 +99,9 @@ class KBMedia(BaseKBModel, table=True):
     __tablename__ = "kb_media"  # type: ignore
 
     id: int | None = Field(
-        primary_key=True, sa_column_kwargs={"autoincrement": True}, default=None
+        primary_key=True,
+        sa_column_kwargs={"autoincrement": True},
+        default=None,
     )
     media_id: str = Field(
         max_length=36,

astrbot/core/knowledge_base/parsers/__init__.py

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

astrbot/core/knowledge_base/parsers/base.py

@@ -47,4 +47,5 @@ class BaseParser(ABC):
 
         Returns:
             ParseResult: 解析结果
+
         """

astrbot/core/knowledge_base/parsers/markitdown_parser.py

@@ -1,11 +1,12 @@
 import io
 import os
 
+from markitdown_no_magika import MarkItDown, StreamInfo
+
 from astrbot.core.knowledge_base.parsers.base import (
     BaseParser,
     ParseResult,
 )
-from markitdown_no_magika import MarkItDown, StreamInfo
 
 
 class MarkitdownParser(BaseParser):

astrbot/core/knowledge_base/parsers/pdf_parser.py

@@ -29,6 +29,7 @@ class PDFParser(BaseParser):
 
         Returns:
             ParseResult: 包含文本和图片的解析结果
+
         """
         pdf_file = io.BytesIO(file_content)
         reader = PdfReader(pdf_file)
@@ -87,7 +88,7 @@ class PDFParser(BaseParser):
                                 file_name=f"page_{page_num}_img_{image_counter}.{ext}",
                                 content=image_data,
                                 mime_type=mime_type,
-                            )
+                            ),
                         )
                     except Exception:
                         # 单个图片提取失败不影响整体

astrbot/core/knowledge_base/parsers/text_parser.py

@@ -26,6 +26,7 @@ class TextParser(BaseParser):
 
         Raises:
             ValueError: 如果无法解码文件
+
         """
         # 尝试多种编码
         for encoding in ["utf-8", "gbk", "gb2312", "gb18030"]:

astrbot/core/knowledge_base/parsers/util.py

@@ -6,7 +6,7 @@ async def select_parser(ext: str) -> BaseParser:
         from .markitdown_parser import MarkitdownParser
 
         return MarkitdownParser()
-    elif ext == ".pdf":
+    if ext == ".pdf":
         from .pdf_parser import PDFParser
 
         return PDFParser()

astrbot/core/knowledge_base/retrieval/__init__.py

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