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

astrbot/core/db/vec_db/faiss_impl/document_storage.py

@@ -1,12 +1,13 @@
-import os
 import json
-from datetime import datetime
+import os
 from contextlib import asynccontextmanager
+from datetime import datetime
 
-from sqlalchemy import Text, Column
+from sqlalchemy import Column, Text
 from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, create_async_engine
 from sqlalchemy.orm import sessionmaker
-from sqlmodel import Field, SQLModel, select, col, func, text, MetaData
+from sqlmodel import Field, MetaData, SQLModel, col, func, select, text
+
 from astrbot.core import logger
 
 
@@ -20,7 +21,9 @@ class Document(BaseDocModel, table=True):
     __tablename__ = "documents"  # type: ignore
 
     id: int | None = Field(
-        default=None, primary_key=True, sa_column_kwargs={"autoincrement": True}
+        default=None,
+        primary_key=True,
+        sa_column_kwargs={"autoincrement": True},
     )
     doc_id: str = Field(nullable=False)
     text: str = Field(nullable=False)
@@ -36,7 +39,8 @@ class DocumentStorage:
         self.engine: AsyncEngine | None = None
         self.async_session_maker: sessionmaker | None = None
         self.sqlite_init_path = os.path.join(
-            os.path.dirname(__file__), "sqlite_init.sql"
+            os.path.dirname(__file__),
+            "sqlite_init.sql",
         )
 
     async def initialize(self):
@@ -50,26 +54,26 @@ class DocumentStorage:
                 await conn.execute(
                     text(
                         "ALTER TABLE documents ADD COLUMN kb_doc_id TEXT "
-                        "GENERATED ALWAYS AS (json_extract(metadata, '$.kb_doc_id')) STORED"
-                    )
+                        "GENERATED ALWAYS AS (json_extract(metadata, '$.kb_doc_id')) STORED",
+                    ),
                 )
                 await conn.execute(
                     text(
                         "ALTER TABLE documents ADD COLUMN user_id TEXT "
-                        "GENERATED ALWAYS AS (json_extract(metadata, '$.user_id')) STORED"
-                    )
+                        "GENERATED ALWAYS AS (json_extract(metadata, '$.user_id')) STORED",
+                    ),
                 )
 
                 # Create indexes
                 await conn.execute(
                     text(
-                        "CREATE INDEX IF NOT EXISTS idx_documents_kb_doc_id ON documents(kb_doc_id)"
-                    )
+                        "CREATE INDEX IF NOT EXISTS idx_documents_kb_doc_id ON documents(kb_doc_id)",
+                    ),
                 )
                 await conn.execute(
                     text(
-                        "CREATE INDEX IF NOT EXISTS idx_documents_user_id ON documents(user_id)"
-                    )
+                        "CREATE INDEX IF NOT EXISTS idx_documents_user_id ON documents(user_id)",
+                    ),
                 )
             except BaseException:
                 pass
@@ -113,10 +117,11 @@ class DocumentStorage:
 
         Returns:
             list: The list of documents that match the filters.
+
         """
         if self.engine is None:
             logger.warning(
-                "Database connection is not initialized, returning empty result"
+                "Database connection is not initialized, returning empty result",
             )
             return []
 
@@ -125,7 +130,7 @@ class DocumentStorage:
 
             for key, val in metadata_filters.items():
                 query = query.where(
-                    text(f"json_extract(metadata, '$.{key}') = :filter_{key}")
+                    text(f"json_extract(metadata, '$.{key}') = :filter_{key}"),
                 ).params(**{f"filter_{key}": val})
 
             if ids is not None and len(ids) > 0:
@@ -153,24 +158,27 @@ class DocumentStorage:
 
         Returns:
             int: The integer ID of the inserted document.
+
         """
         assert self.engine is not None, "Database connection is not initialized."
 
-        async with self.get_session() as session:
-            async with session.begin():
-                document = Document(
-                    doc_id=doc_id,
-                    text=text,
-                    metadata_=json.dumps(metadata),
-                    created_at=datetime.now(),
-                    updated_at=datetime.now(),
-                )
-                session.add(document)
-                await session.flush()  # Flush to get the ID
-                return document.id  # type: ignore
+        async with self.get_session() as session, session.begin():
+            document = Document(
+                doc_id=doc_id,
+                text=text,
+                metadata_=json.dumps(metadata),
+                created_at=datetime.now(),
+                updated_at=datetime.now(),
+            )
+            session.add(document)
+            await session.flush()  # Flush to get the ID
+            return document.id  # type: ignore
 
     async def insert_documents_batch(
-        self, doc_ids: list[str], texts: list[str], metadatas: list[dict]
+        self,
+        doc_ids: list[str],
+        texts: list[str],
+        metadatas: list[dict],
     ) -> list[int]:
         """Batch insert documents and return their integer IDs.
 
@@ -181,44 +189,44 @@ class DocumentStorage:
 
         Returns:
             list[int]: List of integer IDs of the inserted documents.
+
         """
         assert self.engine is not None, "Database connection is not initialized."
 
-        async with self.get_session() as session:
-            async with session.begin():
-                import json
-
-                documents = []
-                for doc_id, text, metadata in zip(doc_ids, texts, metadatas):
-                    document = Document(
-                        doc_id=doc_id,
-                        text=text,
-                        metadata_=json.dumps(metadata),
-                        created_at=datetime.now(),
-                        updated_at=datetime.now(),
-                    )
-                    documents.append(document)
-                    session.add(document)
-
-                await session.flush()  # Flush to get all IDs
-                return [doc.id for doc in documents]  # type: ignore
+        async with self.get_session() as session, session.begin():
+            import json
+
+            documents = []
+            for doc_id, text, metadata in zip(doc_ids, texts, metadatas):
+                document = Document(
+                    doc_id=doc_id,
+                    text=text,
+                    metadata_=json.dumps(metadata),
+                    created_at=datetime.now(),
+                    updated_at=datetime.now(),
+                )
+                documents.append(document)
+                session.add(document)
+
+            await session.flush()  # Flush to get all IDs
+            return [doc.id for doc in documents]  # type: ignore
 
     async def delete_document_by_doc_id(self, doc_id: str):
         """Delete a document by its doc_id.
 
         Args:
             doc_id (str): The doc_id of the document to delete.
+
         """
         assert self.engine is not None, "Database connection is not initialized."
 
-        async with self.get_session() as session:
-            async with session.begin():
-                query = select(Document).where(col(Document.doc_id) == doc_id)
-                result = await session.execute(query)
-                document = result.scalar_one_or_none()
+        async with self.get_session() as session, session.begin():
+            query = select(Document).where(col(Document.doc_id) == doc_id)
+            result = await session.execute(query)
+            document = result.scalar_one_or_none()
 
-                if document:
-                    await session.delete(document)
+            if document:
+                await session.delete(document)
 
     async def get_document_by_doc_id(self, doc_id: str):
         """Retrieve a document by its doc_id.
@@ -228,6 +236,7 @@ class DocumentStorage:
 
         Returns:
             dict: The document data or None if not found.
+
         """
         assert self.engine is not None, "Database connection is not initialized."
 
@@ -246,46 +255,46 @@ class DocumentStorage:
         Args:
             doc_id (str): The doc_id.
             new_text (str): The new text to update the document with.
+
         """
         assert self.engine is not None, "Database connection is not initialized."
 
-        async with self.get_session() as session:
-            async with session.begin():
-                query = select(Document).where(col(Document.doc_id) == doc_id)
-                result = await session.execute(query)
-                document = result.scalar_one_or_none()
+        async with self.get_session() as session, session.begin():
+            query = select(Document).where(col(Document.doc_id) == doc_id)
+            result = await session.execute(query)
+            document = result.scalar_one_or_none()
 
-                if document:
-                    document.text = new_text
-                    document.updated_at = datetime.now()
-                    session.add(document)
+            if document:
+                document.text = new_text
+                document.updated_at = datetime.now()
+                session.add(document)
 
     async def delete_documents(self, metadata_filters: dict):
         """Delete documents by their metadata filters.
 
         Args:
             metadata_filters (dict): The metadata filters to apply.
+
         """
         if self.engine is None:
             logger.warning(
-                "Database connection is not initialized, skipping delete operation"
+                "Database connection is not initialized, skipping delete operation",
             )
             return
 
-        async with self.get_session() as session:
-            async with session.begin():
-                query = select(Document)
+        async with self.get_session() as session, session.begin():
+            query = select(Document)
 
-                for key, val in metadata_filters.items():
-                    query = query.where(
-                        text(f"json_extract(metadata, '$.{key}') = :filter_{key}")
-                    ).params(**{f"filter_{key}": val})
+            for key, val in metadata_filters.items():
+                query = query.where(
+                    text(f"json_extract(metadata, '$.{key}') = :filter_{key}"),
+                ).params(**{f"filter_{key}": val})
 
-                result = await session.execute(query)
-                documents = result.scalars().all()
+            result = await session.execute(query)
+            documents = result.scalars().all()
 
-                for doc in documents:
-                    await session.delete(doc)
+            for doc in documents:
+                await session.delete(doc)
 
     async def count_documents(self, metadata_filters: dict | None = None) -> int:
         """Count documents in the database.
@@ -295,6 +304,7 @@ class DocumentStorage:
 
         Returns:
             int: The count of documents.
+
         """
         if self.engine is None:
             logger.warning("Database connection is not initialized, returning 0")
@@ -306,7 +316,7 @@ class DocumentStorage:
             if metadata_filters:
                 for key, val in metadata_filters.items():
                     query = query.where(
-                        text(f"json_extract(metadata, '$.{key}') = :filter_{key}")
+                        text(f"json_extract(metadata, '$.{key}') = :filter_{key}"),
                     ).params(**{f"filter_{key}": val})
 
             result = await session.execute(query)
@@ -318,12 +328,13 @@ class DocumentStorage:
 
         Returns:
             list: A list of user IDs.
+
         """
         assert self.engine is not None, "Database connection is not initialized."
 
         async with self.get_session() as session:
             query = text(
-                "SELECT DISTINCT user_id FROM documents WHERE user_id IS NOT NULL"
+                "SELECT DISTINCT user_id FROM documents WHERE user_id IS NOT NULL",
             )
             result = await session.execute(query)
             rows = result.fetchall()
@@ -337,6 +348,7 @@ class DocumentStorage:
 
         Returns:
             dict: The converted dictionary.
+
         """
         return {
             "id": document.id,
@@ -361,6 +373,7 @@ class DocumentStorage:
             dict: The converted dictionary.
 
         Note: This method is kept for backward compatibility but is no longer used internally.
+
         """
         return {
             "id": row[0],

astrbot/core/db/vec_db/faiss_impl/embedding_storage.py

@@ -2,9 +2,10 @@ try:
     import faiss
 except ModuleNotFoundError:
     raise ImportError(
-        "faiss 未安装。请使用 'pip install faiss-cpu' 或 'pip install faiss-gpu' 安装。"
+        "faiss 未安装。请使用 'pip install faiss-cpu' 或 'pip install faiss-gpu' 安装。",
     )
 import os
+
 import numpy as np
 
 
@@ -27,11 +28,12 @@ class EmbeddingStorage:
             id (int): 向量的ID
         Raises:
             ValueError: 如果向量的维度与存储的维度不匹配
+
         """
         assert self.index is not None, "FAISS index is not initialized."
         if vector.shape[0] != self.dimension:
             raise ValueError(
-                f"向量维度不匹配, 期望: {self.dimension}, 实际: {vector.shape[0]}"
+                f"向量维度不匹配, 期望: {self.dimension}, 实际: {vector.shape[0]}",
             )
         self.index.add_with_ids(vector.reshape(1, -1), np.array([id]))
         await self.save_index()
@@ -44,11 +46,12 @@ class EmbeddingStorage:
             ids (list[int]): 向量的ID列表
         Raises:
             ValueError: 如果向量的维度与存储的维度不匹配
+
         """
         assert self.index is not None, "FAISS index is not initialized."
         if vectors.shape[1] != self.dimension:
             raise ValueError(
-                f"向量维度不匹配, 期望: {self.dimension}, 实际: {vectors.shape[1]}"
+                f"向量维度不匹配, 期望: {self.dimension}, 实际: {vectors.shape[1]}",
             )
         self.index.add_with_ids(vectors, np.array(ids))
         await self.save_index()
@@ -61,6 +64,7 @@ class EmbeddingStorage:
             k (int): 返回的最相似向量的数量
         Returns:
             tuple: (距离, 索引)
+
         """
         assert self.index is not None, "FAISS index is not initialized."
         faiss.normalize_L2(vector)
@@ -72,6 +76,7 @@ class EmbeddingStorage:
 
         Args:
             ids (list[int]): 要删除的向量ID列表
+
         """
         assert self.index is not None, "FAISS index is not initialized."
         id_array = np.array(ids, dtype=np.int64)
@@ -83,5 +88,6 @@ class EmbeddingStorage:
 
         Args:
             path (str): 保存索引的路径
+
         """
         faiss.write_index(self.index, self.path)

astrbot/core/db/vec_db/faiss_impl/vec_db.py

@@ -1,18 +1,18 @@
-import uuid
 import time
+import uuid
+
 import numpy as np
+
+from astrbot import logger
+from astrbot.core.provider.provider import EmbeddingProvider, RerankProvider
+
+from ..base import BaseVecDB, Result
 from .document_storage import DocumentStorage
 from .embedding_storage import EmbeddingStorage
-from ..base import Result, BaseVecDB
-from astrbot.core.provider.provider import EmbeddingProvider
-from astrbot.core.provider.provider import RerankProvider
-from astrbot import logger
 
 
 class FaissVecDB(BaseVecDB):
-    """
-    A class to represent a vector database.
-    """
+    """A class to represent a vector database."""
 
     def __init__(
         self,
@@ -26,7 +26,8 @@ class FaissVecDB(BaseVecDB):
         self.embedding_provider = embedding_provider
         self.document_storage = DocumentStorage(doc_store_path)
         self.embedding_storage = EmbeddingStorage(
-            embedding_provider.get_dim(), index_store_path
+            embedding_provider.get_dim(),
+            index_store_path,
         )
         self.embedding_provider = embedding_provider
         self.rerank_provider = rerank_provider
@@ -35,11 +36,12 @@ class FaissVecDB(BaseVecDB):
         await self.document_storage.initialize()
 
     async def insert(
-        self, content: str, metadata: dict | None = None, id: str | None = None
+        self,
+        content: str,
+        metadata: dict | None = None,
+        id: str | None = None,
     ) -> int:
-        """
-        插入一条文本和其对应向量，自动生成 ID 并保持一致性。
-        """
+        """插入一条文本和其对应向量，自动生成 ID 并保持一致性。"""
         metadata = metadata or {}
         str_id = id or str(uuid.uuid4())  # 使用 UUID 作为原始 ID
 
@@ -63,11 +65,11 @@ class FaissVecDB(BaseVecDB):
         max_retries: int = 3,
         progress_callback=None,
     ) -> list[int]:
-        """
-        批量插入文本和其对应向量，自动生成 ID 并保持一致性。
+        """批量插入文本和其对应向量，自动生成 ID 并保持一致性。
 
         Args:
             progress_callback: 进度回调函数，接收参数 (current, total)
+
         """
         metadatas = metadatas or [{} for _ in contents]
         ids = ids or [str(uuid.uuid4()) for _ in contents]
@@ -83,12 +85,14 @@ class FaissVecDB(BaseVecDB):
         )
         end = time.time()
         logger.debug(
-            f"Generated embeddings for {len(contents)} contents in {end - start:.2f} seconds."
+            f"Generated embeddings for {len(contents)} contents in {end - start:.2f} seconds.",
         )
 
         # 使用 DocumentStorage 的批量插入方法
         int_ids = await self.document_storage.insert_documents_batch(
-            ids, contents, metadatas
+            ids,
+            contents,
+            metadatas,
         )
 
         # 批量插入向量到 FAISS
@@ -104,8 +108,7 @@ class FaissVecDB(BaseVecDB):
         rerank: bool = False,
         metadata_filters: dict | None = None,
     ) -> list[Result]:
-        """
-        搜索最相似的文档。
+        """搜索最相似的文档。
 
         Args:
             query (str): 查询文本
@@ -116,6 +119,7 @@ class FaissVecDB(BaseVecDB):
 
         Returns:
             List[Result]: 查询结果
+
         """
         embedding = await self.embedding_provider.get_embedding(query)
         scores, indices = await self.embedding_storage.search(
@@ -128,7 +132,8 @@ class FaissVecDB(BaseVecDB):
         scores[0] = 1.0 - (scores[0] / 2.0)
         # NOTE: maybe the size is less than k.
         fetched_docs = await self.document_storage.get_documents(
-            metadata_filters=metadata_filters or {}, ids=indices[0]
+            metadata_filters=metadata_filters or {},
+            ids=indices[0],
         )
         if not fetched_docs:
             return []
@@ -149,7 +154,9 @@ class FaissVecDB(BaseVecDB):
             documents = [doc.data["text"] for doc in top_k_results]
             reranked_results = await self.rerank_provider.rerank(query, documents)
             reranked_results = sorted(
-                reranked_results, key=lambda x: x.relevance_score, reverse=True
+                reranked_results,
+                key=lambda x: x.relevance_score,
+                reverse=True,
             )
             top_k_results = [
                 top_k_results[reranked_result.index]
@@ -159,9 +166,7 @@ class FaissVecDB(BaseVecDB):
         return top_k_results
 
     async def delete(self, doc_id: str):
-        """
-        删除一条文档块（chunk）
-        """
+        """删除一条文档块（chunk）"""
         # 获得对应的 int id
         result = await self.document_storage.get_document_by_doc_id(doc_id)
         int_id = result["id"] if result else None
@@ -176,23 +181,23 @@ class FaissVecDB(BaseVecDB):
         await self.document_storage.close()
 
     async def count_documents(self, metadata_filter: dict | None = None) -> int:
-        """
-        计算文档数量
+        """计算文档数量
 
         Args:
             metadata_filter (dict | None): 元数据过滤器
+
         """
         count = await self.document_storage.count_documents(
-            metadata_filters=metadata_filter or {}
+            metadata_filters=metadata_filter or {},
         )
         return count
 
     async def delete_documents(self, metadata_filters: dict):
-        """
-        根据元数据过滤器删除文档
-        """
+        """根据元数据过滤器删除文档"""
         docs = await self.document_storage.get_documents(
-            metadata_filters=metadata_filters, offset=None, limit=None
+            metadata_filters=metadata_filters,
+            offset=None,
+            limit=None,
         )
         doc_ids: list[int] = [doc["id"] for doc in docs]
         await self.embedding_storage.delete(doc_ids)

astrbot/core/event_bus.py

@@ -1,5 +1,4 @@
-"""
-事件总线, 用于处理事件的分发和处理
+"""事件总线, 用于处理事件的分发和处理
 事件总线是一个异步队列, 用于接收各种消息事件, 并将其发送到Scheduler调度器进行处理
 其中包含了一个无限循环的调度函数, 用于从事件队列中获取新的事件, 并创建一个新的异步任务来执行管道调度器的处理逻辑
 
@@ -13,10 +12,12 @@ class:
 
 import asyncio
 from asyncio import Queue
-from astrbot.core.pipeline.scheduler import PipelineScheduler
+
 from astrbot.core import logger
-from .platform import AstrMessageEvent
 from astrbot.core.astrbot_config_mgr import AstrBotConfigManager
+from astrbot.core.pipeline.scheduler import PipelineScheduler
+
+from .platform import AstrMessageEvent
 
 
 class EventBus:
@@ -46,14 +47,15 @@ class EventBus:
 
         Args:
             event (AstrMessageEvent): 事件对象
+
         """
         # 如果有发送者名称: [平台名] 发送者名称/发送者ID: 消息概要
         if event.get_sender_name():
             logger.info(
-                f"[{conf_name}] [{event.get_platform_id()}({event.get_platform_name()})] {event.get_sender_name()}/{event.get_sender_id()}: {event.get_message_outline()}"
+                f"[{conf_name}] [{event.get_platform_id()}({event.get_platform_name()})] {event.get_sender_name()}/{event.get_sender_id()}: {event.get_message_outline()}",
             )
         # 没有发送者名称: [平台名] 发送者ID: 消息概要
         else:
             logger.info(
-                f"[{conf_name}] [{event.get_platform_id()}({event.get_platform_name()})] {event.get_sender_id()}: {event.get_message_outline()}"
+                f"[{conf_name}] [{event.get_platform_id()}({event.get_platform_name()})] {event.get_sender_id()}: {event.get_message_outline()}",
             )

astrbot/core/file_token_service.py

@@ -1,9 +1,9 @@
 import asyncio
 import os
-import uuid
-import time
-from urllib.parse import urlparse, unquote
 import platform
+import time
+import uuid
+from urllib.parse import unquote, urlparse
 
 
 class FileTokenService:
@@ -40,8 +40,8 @@ class FileTokenService:
 
         Raises:
             FileNotFoundError: 当路径不存在时抛出
-        """
 
+        """
         # 处理 file:///
         try:
             parsed_uri = urlparse(file_path)
@@ -61,7 +61,7 @@ class FileTokenService:
 
             if not os.path.exists(local_path):
                 raise FileNotFoundError(
-                    f"文件不存在: {local_path} (原始输入: {file_path})"
+                    f"文件不存在: {local_path} (原始输入: {file_path})",
                 )
 
             file_token = str(uuid.uuid4())
@@ -84,6 +84,7 @@ class FileTokenService:
         Raises:
             KeyError: 当令牌不存在或已过期时抛出
             FileNotFoundError: 当文件本身已被删除时抛出
+
         """
         async with self.lock:
             await self._cleanup_expired_tokens()

astrbot/core/initial_loader.py

@@ -1,5 +1,4 @@
-"""
-AstrBot 启动器，负责初始化和启动核心组件和仪表板服务器。
+"""AstrBot 启动器，负责初始化和启动核心组件和仪表板服务器。
 
 工作流程:
 1. 初始化核心生命周期, 传递数据库和日志代理实例到核心生命周期
@@ -8,10 +7,10 @@ AstrBot 启动器，负责初始化和启动核心组件和仪表板服务器。
 
 import asyncio
 import traceback
-from astrbot.core import logger
+
+from astrbot.core import LogBroker, logger
 from astrbot.core.core_lifecycle import AstrBotCoreLifecycle
 from astrbot.core.db import BaseDatabase
-from astrbot.core import LogBroker
 from astrbot.dashboard.server import AstrBotDashboard
 
 
@@ -39,7 +38,10 @@ class InitialLoader:
         webui_dir = self.webui_dir
 
         self.dashboard_server = AstrBotDashboard(
-            core_lifecycle, self.db, core_lifecycle.dashboard_shutdown_event, webui_dir
+            core_lifecycle,
+            self.db,
+            core_lifecycle.dashboard_shutdown_event,
+            webui_dir,
         )
 
         coro = self.dashboard_server.run()

astrbot/core/knowledge_base/chunking/__init__.py

@@ -1,6 +1,4 @@
-"""
-文档分块模块
-"""
+"""文档分块模块"""
 
 from .base import BaseChunker
 from .fixed_size import FixedSizeChunker

astrbot/core/knowledge_base/chunking/base.py

@@ -21,4 +21,5 @@ class BaseChunker(ABC):
 
         Returns:
             list[str]: 分块后的文本列表
+
         """