AstrBot 3.5.6__py3-none-any.whl → 4.7.0__py3-none-any.whl

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

@@ -0,0 +1,204 @@
+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
+
+
+class FaissVecDB(BaseVecDB):
+    """A class to represent a vector database."""
+
+    def __init__(
+        self,
+        doc_store_path: str,
+        index_store_path: str,
+        embedding_provider: EmbeddingProvider,
+        rerank_provider: RerankProvider | None = None,
+    ):
+        self.doc_store_path = doc_store_path
+        self.index_store_path = index_store_path
+        self.embedding_provider = embedding_provider
+        self.document_storage = DocumentStorage(doc_store_path)
+        self.embedding_storage = EmbeddingStorage(
+            embedding_provider.get_dim(),
+            index_store_path,
+        )
+        self.embedding_provider = embedding_provider
+        self.rerank_provider = rerank_provider
+
+    async def initialize(self):
+        await self.document_storage.initialize()
+
+    async def insert(
+        self,
+        content: str,
+        metadata: dict | None = None,
+        id: str | None = None,
+    ) -> int:
+        """插入一条文本和其对应向量，自动生成 ID 并保持一致性。"""
+        metadata = metadata or {}
+        str_id = id or str(uuid.uuid4())  # 使用 UUID 作为原始 ID
+
+        vector = await self.embedding_provider.get_embedding(content)
+        vector = np.array(vector, dtype=np.float32)
+
+        # 使用 DocumentStorage 的方法插入文档
+        int_id = await self.document_storage.insert_document(str_id, content, metadata)
+
+        # 插入向量到 FAISS
+        await self.embedding_storage.insert(vector, int_id)
+        return int_id
+
+    async def insert_batch(
+        self,
+        contents: list[str],
+        metadatas: list[dict] | None = None,
+        ids: list[str] | None = None,
+        batch_size: int = 32,
+        tasks_limit: int = 3,
+        max_retries: int = 3,
+        progress_callback=None,
+    ) -> list[int]:
+        """批量插入文本和其对应向量，自动生成 ID 并保持一致性。
+
+        Args:
+            progress_callback: 进度回调函数，接收参数 (current, total)
+
+        """
+        metadatas = metadatas or [{} for _ in contents]
+        ids = ids or [str(uuid.uuid4()) for _ in contents]
+
+        start = time.time()
+        logger.debug(f"Generating embeddings for {len(contents)} contents...")
+        vectors = await self.embedding_provider.get_embeddings_batch(
+            contents,
+            batch_size=batch_size,
+            tasks_limit=tasks_limit,
+            max_retries=max_retries,
+            progress_callback=progress_callback,
+        )
+        end = time.time()
+        logger.debug(
+            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,
+        )
+
+        # 批量插入向量到 FAISS
+        vectors_array = np.array(vectors).astype("float32")
+        await self.embedding_storage.insert_batch(vectors_array, int_ids)
+        return int_ids
+
+    async def retrieve(
+        self,
+        query: str,
+        k: int = 5,
+        fetch_k: int = 20,
+        rerank: bool = False,
+        metadata_filters: dict | None = None,
+    ) -> list[Result]:
+        """搜索最相似的文档。
+
+        Args:
+            query (str): 查询文本
+            k (int): 返回的最相似文档的数量
+            fetch_k (int): 在根据 metadata 过滤前从 FAISS 中获取的数量
+            rerank (bool): 是否使用重排序。这需要在实例化时提供 rerank_provider, 如果未提供并且 rerank 为 True, 不会抛出异常。
+            metadata_filters (dict): 元数据过滤器
+
+        Returns:
+            List[Result]: 查询结果
+
+        """
+        embedding = await self.embedding_provider.get_embedding(query)
+        scores, indices = await self.embedding_storage.search(
+            vector=np.array([embedding]).astype("float32"),
+            k=fetch_k if metadata_filters else k,
+        )
+        if len(indices[0]) == 0 or indices[0][0] == -1:
+            return []
+        # normalize scores
+        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],
+        )
+        if not fetched_docs:
+            return []
+        result_docs: list[Result] = []
+
+        idx_pos = {fetch_doc["id"]: idx for idx, fetch_doc in enumerate(fetched_docs)}
+        for i, indice_idx in enumerate(indices[0]):
+            pos = idx_pos.get(indice_idx)
+            if pos is None:
+                continue
+            fetch_doc = fetched_docs[pos]
+            score = scores[0][i]
+            result_docs.append(Result(similarity=float(score), data=fetch_doc))
+
+        top_k_results = result_docs[:k]
+
+        if rerank and self.rerank_provider:
+            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,
+            )
+            top_k_results = [
+                top_k_results[reranked_result.index]
+                for reranked_result in reranked_results
+            ]
+
+        return top_k_results
+
+    async def delete(self, doc_id: str):
+        """删除一条文档块（chunk）"""
+        # 获得对应的 int id
+        result = await self.document_storage.get_document_by_doc_id(doc_id)
+        int_id = result["id"] if result else None
+        if int_id is None:
+            return
+
+        # 使用 DocumentStorage 的删除方法
+        await self.document_storage.delete_document_by_doc_id(doc_id)
+        await self.embedding_storage.delete([int_id])
+
+    async def close(self):
+        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 {},
+        )
+        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,
+        )
+        doc_ids: list[int] = [doc["id"] for doc in docs]
+        await self.embedding_storage.delete(doc_ids)
+        await self.document_storage.delete_documents(metadata_filters=metadata_filters)

astrbot/core/event_bus.py

@@ -1,5 +1,4 @@
-"""
-事件总线, 用于处理事件的分发和处理
+"""事件总线, 用于处理事件的分发和处理
 事件总线是一个异步队列, 用于接收各种消息事件, 并将其发送到Scheduler调度器进行处理
 其中包含了一个无限循环的调度函数, 用于从事件队列中获取新的事件, 并创建一个新的异步任务来执行管道调度器的处理逻辑
 
@@ -13,45 +12,50 @@ class:
 
 import asyncio
 from asyncio import Queue
-from astrbot.core.pipeline.scheduler import PipelineScheduler
+
 from astrbot.core import logger
+from astrbot.core.astrbot_config_mgr import AstrBotConfigManager
+from astrbot.core.pipeline.scheduler import PipelineScheduler
+
 from .platform import AstrMessageEvent
 
 
 class EventBus:
-    """事件总线: 用于处理事件的分发和处理
-
-    维护一个异步队列, 来接受各种消息事件
-    """
-
-    def __init__(self, event_queue: Queue, pipeline_scheduler: PipelineScheduler):
+    """用于处理事件的分发和处理"""
+
+    def __init__(
+        self,
+        event_queue: Queue,
+        pipeline_scheduler_mapping: dict[str, PipelineScheduler],
+        astrbot_config_mgr: AstrBotConfigManager = None,
+    ):
         self.event_queue = event_queue  # 事件队列
-        self.pipeline_scheduler = pipeline_scheduler  # 管道调度器
+        # abconf uuid -> scheduler
+        self.pipeline_scheduler_mapping = pipeline_scheduler_mapping
+        self.astrbot_config_mgr = astrbot_config_mgr
 
     async def dispatch(self):
-        """无限循环的调度函数, 从事件队列中获取新的事件, 打印日志并创建一个新的异步任务来执行管道调度器的处理逻辑"""
         while True:
-            event: AstrMessageEvent = (
-                await self.event_queue.get()
-            )  # 从事件队列中获取新的事件
-            self._print_event(event)  # 打印日志
-            asyncio.create_task(
-                self.pipeline_scheduler.execute(event)
-            )  # 创建新的异步任务来执行管道调度器的处理逻辑
-
-    def _print_event(self, event: AstrMessageEvent):
+            event: AstrMessageEvent = await self.event_queue.get()
+            conf_info = self.astrbot_config_mgr.get_conf_info(event.unified_msg_origin)
+            self._print_event(event, conf_info["name"])
+            scheduler = self.pipeline_scheduler_mapping.get(conf_info["id"])
+            asyncio.create_task(scheduler.execute(event))
+
+    def _print_event(self, event: AstrMessageEvent, conf_name: str):
         """用于记录事件信息
 
         Args:
             event (AstrMessageEvent): 事件对象
+
         """
         # 如果有发送者名称: [平台名] 发送者名称/发送者ID: 消息概要
         if event.get_sender_name():
             logger.info(
-                f"[{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"[{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/exceptions.py

@@ -0,0 +1,9 @@
+from __future__ import annotations
+
+
+class AstrBotError(Exception):
+    """Base exception for all AstrBot errors."""
+
+
+class ProviderNotFoundError(AstrBotError):
+    """Raised when a specified provider is not found."""

astrbot/core/file_token_service.py

@@ -0,0 +1,98 @@
+import asyncio
+import os
+import platform
+import time
+import uuid
+from urllib.parse import unquote, urlparse
+
+
+class FileTokenService:
+    """维护一个简单的基于令牌的文件下载服务，支持超时和懒清除。"""
+
+    def __init__(self, default_timeout: float = 300):
+        self.lock = asyncio.Lock()
+        self.staged_files = {}  # token: (file_path, expire_time)
+        self.default_timeout = default_timeout
+
+    async def _cleanup_expired_tokens(self):
+        """清理过期的令牌"""
+        now = time.time()
+        expired_tokens = [
+            token for token, (_, expire) in self.staged_files.items() if expire < now
+        ]
+        for token in expired_tokens:
+            self.staged_files.pop(token, None)
+
+    async def check_token_expired(self, file_token: str) -> bool:
+        async with self.lock:
+            await self._cleanup_expired_tokens()
+            return file_token not in self.staged_files
+
+    async def register_file(self, file_path: str, timeout: float | None = None) -> str:
+        """向令牌服务注册一个文件。
+
+        Args:
+            file_path(str): 文件路径
+            timeout(float): 超时时间，单位秒（可选）
+
+        Returns:
+            str: 一个单次令牌
+
+        Raises:
+            FileNotFoundError: 当路径不存在时抛出
+
+        """
+        # 处理 file:///
+        try:
+            parsed_uri = urlparse(file_path)
+            if parsed_uri.scheme == "file":
+                local_path = unquote(parsed_uri.path)
+                if platform.system() == "Windows" and local_path.startswith("/"):
+                    local_path = local_path[1:]
+            else:
+                # 如果没有 file:/// 前缀，则认为是普通路径
+                local_path = file_path
+        except Exception:
+            # 解析失败时，按原路径处理
+            local_path = file_path
+
+        async with self.lock:
+            await self._cleanup_expired_tokens()
+
+            if not os.path.exists(local_path):
+                raise FileNotFoundError(
+                    f"文件不存在: {local_path} (原始输入: {file_path})",
+                )
+
+            file_token = str(uuid.uuid4())
+            expire_time = time.time() + (
+                timeout if timeout is not None else self.default_timeout
+            )
+            # 存储转换后的真实路径
+            self.staged_files[file_token] = (local_path, expire_time)
+            return file_token
+
+    async def handle_file(self, file_token: str) -> str:
+        """根据令牌获取文件路径，使用后令牌失效。
+
+        Args:
+            file_token(str): 注册时返回的令牌
+
+        Returns:
+            str: 文件路径
+
+        Raises:
+            KeyError: 当令牌不存在或已过期时抛出
+            FileNotFoundError: 当文件本身已被删除时抛出
+
+        """
+        async with self.lock:
+            await self._cleanup_expired_tokens()
+
+            if file_token not in self.staged_files:
+                raise KeyError(f"无效或过期的文件 token: {file_token}")
+
+            file_path, _ = self.staged_files.pop(file_token)
+            if not os.path.exists(file_path):
+                raise FileNotFoundError(f"文件不存在: {file_path}")
+            return file_path

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
 
 
@@ -22,25 +21,35 @@ class InitialLoader:
         self.db = db
         self.logger = logger
         self.log_broker = log_broker
+        self.webui_dir: str | None = None
 
     async def start(self):
         core_lifecycle = AstrBotCoreLifecycle(self.log_broker, self.db)
 
-        core_task = []
         try:
             await core_lifecycle.initialize()
-            core_task = core_lifecycle.start()
         except Exception as e:
             logger.critical(traceback.format_exc())
             logger.critical(f"😭 初始化 AstrBot 失败：{e} !!!")
+            return
+
+        core_task = core_lifecycle.start()
+
+        webui_dir = self.webui_dir
 
         self.dashboard_server = AstrBotDashboard(
-            core_lifecycle, self.db, core_lifecycle.dashboard_shutdown_event
+            core_lifecycle,
+            self.db,
+            core_lifecycle.dashboard_shutdown_event,
+            webui_dir,
         )
-        task = asyncio.gather(
-            core_task, self.dashboard_server.run()
-        )  # 启动核心任务和仪表板服务器
 
+        coro = self.dashboard_server.run()
+        if coro:
+            # 启动核心任务和仪表板服务器
+            task = asyncio.gather(core_task, coro)
+        else:
+            task = core_task
         try:
             await task  # 整个AstrBot在这里运行
         except asyncio.CancelledError:

astrbot/core/knowledge_base/chunking/__init__.py

@@ -0,0 +1,9 @@
+"""文档分块模块"""
+
+from .base import BaseChunker
+from .fixed_size import FixedSizeChunker
+
+__all__ = [
+    "BaseChunker",
+    "FixedSizeChunker",
+]

astrbot/core/knowledge_base/chunking/base.py

@@ -0,0 +1,25 @@
+"""文档分块器基类
+
+定义了文档分块处理的抽象接口。
+"""
+
+from abc import ABC, abstractmethod
+
+
+class BaseChunker(ABC):
+    """分块器基类
+
+    所有分块器都应该继承此类并实现 chunk 方法。
+    """
+
+    @abstractmethod
+    async def chunk(self, text: str, **kwargs) -> list[str]:
+        """将文本分块
+
+        Args:
+            text: 输入文本
+
+        Returns:
+            list[str]: 分块后的文本列表
+
+        """

astrbot/core/knowledge_base/chunking/fixed_size.py

@@ -0,0 +1,59 @@
+"""固定大小分块器
+
+按照固定的字符数将文本分块,支持重叠区域。
+"""
+
+from .base import BaseChunker
+
+
+class FixedSizeChunker(BaseChunker):
+    """固定大小分块器
+
+    按照固定的字符数分块,并支持块之间的重叠。
+    """
+
+    def __init__(self, chunk_size: int = 512, chunk_overlap: int = 50):
+        """初始化分块器
+
+        Args:
+            chunk_size: 块的大小(字符数)
+            chunk_overlap: 块之间的重叠字符数
+
+        """
+        self.chunk_size = chunk_size
+        self.chunk_overlap = chunk_overlap
+
+    async def chunk(self, text: str, **kwargs) -> list[str]:
+        """固定大小分块
+
+        Args:
+            text: 输入文本
+            chunk_size: 每个文本块的最大大小
+            chunk_overlap: 每个文本块之间的重叠部分大小
+
+        Returns:
+            list[str]: 分块后的文本列表
+
+        """
+        chunk_size = kwargs.get("chunk_size", self.chunk_size)
+        chunk_overlap = kwargs.get("chunk_overlap", self.chunk_overlap)
+
+        chunks = []
+        start = 0
+        text_len = len(text)
+
+        while start < text_len:
+            end = start + chunk_size
+            chunk = text[start:end]
+
+            if chunk:
+                chunks.append(chunk)
+
+            # 移动窗口,保留重叠部分
+            start = end - chunk_overlap
+
+            # 防止无限循环: 如果重叠过大,直接移到end
+            if start >= end or chunk_overlap >= chunk_size:
+                start = end
+
+        return chunks