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

astrbot/core/pipeline/process_stage/method/star_request.py

@@ -1,16 +1,17 @@
-"""
-本地 Agent 模式的 AstrBot 插件调用 Stage
-"""
+"""本地 Agent 模式的 AstrBot 插件调用 Stage"""
+
+import traceback
+from collections.abc import AsyncGenerator
+from typing import Any
 
-from ...context import PipelineContext
-from ..stage import Stage
-from typing import Dict, Any, List, AsyncGenerator, Union
-from astrbot.core.platform.astr_message_event import AstrMessageEvent
-from astrbot.core.message.message_event_result import MessageEventResult
 from astrbot.core import logger
-from astrbot.core.star.star_handler import StarHandlerMetadata
+from astrbot.core.message.message_event_result import MessageEventResult
+from astrbot.core.platform.astr_message_event import AstrMessageEvent
 from astrbot.core.star.star import star_map
-import traceback
+from astrbot.core.star.star_handler import StarHandlerMetadata
+
+from ...context import PipelineContext, call_handler
+from ..stage import Stage
 
 
 class StarRequestSubStage(Stage):
@@ -21,36 +22,29 @@ class StarRequestSubStage(Stage):
         self.ctx = ctx
 
     async def process(
-        self, event: AstrMessageEvent
-    ) -> Union[None, AsyncGenerator[None, None]]:
-        activated_handlers: List[StarHandlerMetadata] = event.get_extra(
-            "activated_handlers"
+        self,
+        event: AstrMessageEvent,
+    ) -> AsyncGenerator[None, None]:
+        activated_handlers: list[StarHandlerMetadata] = event.get_extra(
+            "activated_handlers",
         )
-        handlers_parsed_params: Dict[str, Dict[str, Any]] = event.get_extra(
-            "handlers_parsed_params"
+        handlers_parsed_params: dict[str, dict[str, Any]] = event.get_extra(
+            "handlers_parsed_params",
         )
         if not handlers_parsed_params:
             handlers_parsed_params = {}
 
         for handler in activated_handlers:
-            # 检查处理器是否在当前平台兼容
-            if (
-                hasattr(handler, "platform_compatible")
-                and handler.platform_compatible is False
-            ):
-                logger.debug(
-                    f"处理器 {handler.handler_name} 在当前平台不兼容，跳过执行"
+            params = handlers_parsed_params.get(handler.handler_full_name, {})
+            md = star_map.get(handler.handler_module_path)
+            if not md:
+                logger.warning(
+                    f"Cannot find plugin for given handler module path: {handler.handler_module_path}",
                 )
                 continue
-
-            params = handlers_parsed_params.get(handler.handler_full_name, {})
+            logger.debug(f"plugin -> {md.name} - {handler.handler_name}")
             try:
-                if handler.handler_module_path not in star_map:
-                    continue
-                logger.debug(
-                    f"plugin -> {star_map.get(handler.handler_module_path).name} - {handler.handler_name}"
-                )
-                wrapper = self._call_handler(self.ctx, event, handler.handler, **params)
+                wrapper = call_handler(event, handler.handler, **params)
                 async for ret in wrapper:
                     yield ret
                 event.clear_result()  # 清除上一个 handler 的结果
@@ -59,7 +53,7 @@ class StarRequestSubStage(Stage):
                 logger.error(f"Star {handler.handler_full_name} handle error: {e}")
 
                 if event.is_at_or_wake_command:
-                    ret = f":(\n\n在调用插件 {star_map.get(handler.handler_module_path).name} 的处理函数 {handler.handler_name} 时出现异常：{e}"
+                    ret = f":(\n\n在调用插件 {md.name} 的处理函数 {handler.handler_name} 时出现异常：{e}"
                     event.set_result(MessageEventResult().message(ret))
                     yield
                     event.clear_result()

astrbot/core/pipeline/process_stage/stage.py

@@ -1,12 +1,14 @@
-from typing import List, Union, AsyncGenerator
-from ..stage import Stage, register_stage
-from ..context import PipelineContext
-from .method.llm_request import LLMRequestSubStage
-from .method.star_request import StarRequestSubStage
+from collections.abc import AsyncGenerator
+
+from astrbot.core import logger
 from astrbot.core.platform.astr_message_event import AstrMessageEvent
-from astrbot.core.star.star_handler import StarHandlerMetadata
 from astrbot.core.provider.entities import ProviderRequest
-from astrbot.core import logger
+from astrbot.core.star.star_handler import StarHandlerMetadata
+
+from ..context import PipelineContext
+from ..stage import Stage, register_stage
+from .method.agent_request import AgentRequestSubStage
+from .method.star_request import StarRequestSubStage
 
 
 @register_stage
@@ -15,18 +17,22 @@ class ProcessStage(Stage):
         self.ctx = ctx
         self.config = ctx.astrbot_config
         self.plugin_manager = ctx.plugin_manager
-        self.llm_request_sub_stage = LLMRequestSubStage()
-        await self.llm_request_sub_stage.initialize(ctx)
 
+        # initialize agent sub stage
+        self.agent_sub_stage = AgentRequestSubStage()
+        await self.agent_sub_stage.initialize(ctx)
+
+        # initialize star request sub stage
         self.star_request_sub_stage = StarRequestSubStage()
         await self.star_request_sub_stage.initialize(ctx)
 
     async def process(
-        self, event: AstrMessageEvent
-    ) -> Union[None, AsyncGenerator[None, None]]:
+        self,
+        event: AstrMessageEvent,
+    ) -> None | AsyncGenerator[None, None]:
         """处理事件"""
-        activated_handlers: List[StarHandlerMetadata] = event.get_extra(
-            "activated_handlers"
+        activated_handlers: list[StarHandlerMetadata] = event.get_extra(
+            "activated_handlers",
         )
         # 有插件 Handler 被激活
         if activated_handlers:
@@ -36,7 +42,7 @@ class ProcessStage(Stage):
                     # Handler 的 LLM 请求
                     event.set_extra("provider_request", resp)
                     _t = False
-                    async for _ in self.llm_request_sub_stage.process(event):
+                    async for _ in self.agent_sub_stage.process(event):
                         _t = True
                         yield
                     if not _t:
@@ -64,5 +70,5 @@ class ProcessStage(Stage):
                     logger.info("未找到可用的 LLM 提供商，请先前往配置服务提供商。")
                     return
 
-                async for _ in self.llm_request_sub_stage.process(event):
+                async for _ in self.agent_sub_stage.process(event):
                     yield

astrbot/core/pipeline/process_stage/utils.py

@@ -0,0 +1,125 @@
+from pydantic import Field
+from pydantic.dataclasses import dataclass
+
+from astrbot.api import logger, sp
+from astrbot.core.agent.run_context import ContextWrapper
+from astrbot.core.agent.tool import FunctionTool, ToolExecResult
+from astrbot.core.astr_agent_context import AstrAgentContext
+from astrbot.core.star.context import Context
+
+
+@dataclass
+class KnowledgeBaseQueryTool(FunctionTool[AstrAgentContext]):
+    name: str = "astr_kb_search"
+    description: str = (
+        "Query the knowledge base for facts or relevant context. "
+        "Use this tool when the user's question requires factual information, "
+        "definitions, background knowledge, or previously indexed content. "
+        "Only send short keywords or a concise question as the query."
+    )
+    parameters: dict = Field(
+        default_factory=lambda: {
+            "type": "object",
+            "properties": {
+                "query": {
+                    "type": "string",
+                    "description": "A concise keyword query for the knowledge base.",
+                },
+            },
+            "required": ["query"],
+        }
+    )
+
+    async def call(
+        self, context: ContextWrapper[AstrAgentContext], **kwargs
+    ) -> ToolExecResult:
+        query = kwargs.get("query", "")
+        if not query:
+            return "error: Query parameter is empty."
+        result = await retrieve_knowledge_base(
+            query=kwargs.get("query", ""),
+            umo=context.context.event.unified_msg_origin,
+            context=context.context.context,
+        )
+        if not result:
+            return "No relevant knowledge found."
+        return result
+
+
+async def retrieve_knowledge_base(
+    query: str,
+    umo: str,
+    context: Context,
+) -> str | None:
+    """Inject knowledge base context into the provider request
+
+    Args:
+        umo: Unique message object (session ID)
+        p_ctx: Pipeline context
+    """
+    kb_mgr = context.kb_manager
+    config = context.get_config(umo=umo)
+
+    # 1. 优先读取会话级配置
+    session_config = await sp.session_get(umo, "kb_config", default={})
+
+    if session_config and "kb_ids" in session_config:
+        # 会话级配置
+        kb_ids = session_config.get("kb_ids", [])
+
+        # 如果配置为空列表，明确表示不使用知识库
+        if not kb_ids:
+            logger.info(f"[知识库] 会话 {umo} 已被配置为不使用知识库")
+            return
+
+        top_k = session_config.get("top_k", 5)
+
+        # 将 kb_ids 转换为 kb_names
+        kb_names = []
+        invalid_kb_ids = []
+        for kb_id in kb_ids:
+            kb_helper = await kb_mgr.get_kb(kb_id)
+            if kb_helper:
+                kb_names.append(kb_helper.kb.kb_name)
+            else:
+                logger.warning(f"[知识库] 知识库不存在或未加载: {kb_id}")
+                invalid_kb_ids.append(kb_id)
+
+        if invalid_kb_ids:
+            logger.warning(
+                f"[知识库] 会话 {umo} 配置的以下知识库无效: {invalid_kb_ids}",
+            )
+
+        if not kb_names:
+            return
+
+        logger.debug(f"[知识库] 使用会话级配置，知识库数量: {len(kb_names)}")
+    else:
+        kb_names = config.get("kb_names", [])
+        top_k = config.get("kb_final_top_k", 5)
+        logger.debug(f"[知识库] 使用全局配置，知识库数量: {len(kb_names)}")
+
+    top_k_fusion = config.get("kb_fusion_top_k", 20)
+
+    if not kb_names:
+        return
+
+    logger.debug(f"[知识库] 开始检索知识库，数量: {len(kb_names)}, top_k={top_k}")
+    kb_context = await kb_mgr.retrieve(
+        query=query,
+        kb_names=kb_names,
+        top_k_fusion=top_k_fusion,
+        top_m_final=top_k,
+    )
+
+    if not kb_context:
+        return
+
+    formatted = kb_context.get("context_text", "")
+    if formatted:
+        results = kb_context.get("results", [])
+        logger.debug(f"[知识库] 为会话 {umo} 注入了 {len(results)} 条相关知识块")
+        return formatted
+
+
+KNOWLEDGE_BASE_QUERY_TOOL = KnowledgeBaseQueryTool()

astrbot/core/pipeline/rate_limit_check/stage.py

@@ -1,18 +1,19 @@
 import asyncio
-from datetime import datetime, timedelta
 from collections import defaultdict, deque
-from typing import DefaultDict, Deque, Union, AsyncGenerator
-from ..stage import Stage, register_stage
-from ..context import PipelineContext
-from astrbot.core.platform.astr_message_event import AstrMessageEvent
+from collections.abc import AsyncGenerator
+from datetime import datetime, timedelta
+
 from astrbot.core import logger
 from astrbot.core.config.astrbot_config import RateLimitStrategy
+from astrbot.core.platform.astr_message_event import AstrMessageEvent
+
+from ..context import PipelineContext
+from ..stage import Stage, register_stage
 
 
 @register_stage
 class RateLimitStage(Stage):
-    """
-    检查是否需要限制消息发送的限流器。
+    """检查是否需要限制消息发送的限流器。
 
     使用 Fixed Window 算法。
     如果触发限流，将 stall 流水线，直到下一个时间窗口来临时自动唤醒。
@@ -20,32 +21,30 @@ class RateLimitStage(Stage):
 
     def __init__(self):
         # 存储每个会话的请求时间队列
-        self.event_timestamps: DefaultDict[str, Deque[datetime]] = defaultdict(deque)
+        self.event_timestamps: defaultdict[str, deque[datetime]] = defaultdict(deque)
         # 为每个会话设置一个锁，避免并发冲突
-        self.locks: DefaultDict[str, asyncio.Lock] = defaultdict(asyncio.Lock)
+        self.locks: defaultdict[str, asyncio.Lock] = defaultdict(asyncio.Lock)
         # 限流参数
         self.rate_limit_count: int = 0
         self.rate_limit_time: timedelta = timedelta(0)
 
     async def initialize(self, ctx: PipelineContext) -> None:
-        """
-        初始化限流器，根据配置设置限流参数。
-        """
+        """初始化限流器，根据配置设置限流参数。"""
         self.rate_limit_count = ctx.astrbot_config["platform_settings"]["rate_limit"][
             "count"
         ]
         self.rate_limit_time = timedelta(
-            seconds=ctx.astrbot_config["platform_settings"]["rate_limit"]["time"]
+            seconds=ctx.astrbot_config["platform_settings"]["rate_limit"]["time"],
         )
         self.rl_strategy = ctx.astrbot_config["platform_settings"]["rate_limit"][
             "strategy"
         ]  # stall or discard
 
     async def process(
-        self, event: AstrMessageEvent
-    ) -> Union[None, AsyncGenerator[None, None]]:
-        """
-        检查并处理限流逻辑。如果触发限流，流水线会 stall 并在窗口期后自动恢复。
+        self,
+        event: AstrMessageEvent,
+    ) -> None | AsyncGenerator[None, None]:
+        """检查并处理限流逻辑。如果触发限流，流水线会 stall 并在窗口期后自动恢复。
 
         Args:
             event (AstrMessageEvent): 当前消息事件。
@@ -53,48 +52,47 @@ class RateLimitStage(Stage):
 
         Returns:
             MessageEventResult: 继续或停止事件处理的结果。
+
         """
         session_id = event.session_id
         now = datetime.now()
 
         async with self.locks[session_id]:  # 确保同一会话不会并发修改队列
-            timestamps = self.event_timestamps[session_id]
-
-            self._remove_expired_timestamps(timestamps, now)
-
-            if len(timestamps) >= self.rate_limit_count:
-                # 达到限流阈值，计算下一个窗口的时间
+            # 检查并处理限流，可能需要多次检查直到满足条件
+            while True:
+                timestamps = self.event_timestamps[session_id]
+                self._remove_expired_timestamps(timestamps, now)
+
+                if len(timestamps) < self.rate_limit_count:
+                    timestamps.append(now)
+                    break
                 next_window_time = timestamps[0] + self.rate_limit_time
-                stall_duration = (next_window_time - now).total_seconds()
+                stall_duration = (next_window_time - now).total_seconds() + 0.3
 
                 match self.rl_strategy:
                     case RateLimitStrategy.STALL.value:
                         logger.info(
-                            f"会话 {session_id} 被限流。根据限流策略，此会话处理将被暂停 {stall_duration:.2f} 秒。"
+                            f"会话 {session_id} 被限流。根据限流策略，此会话处理将被暂停 {stall_duration:.2f} 秒。",
                         )
                         await asyncio.sleep(stall_duration)
+                        now = datetime.now()
                     case RateLimitStrategy.DISCARD.value:
-                        # event.set_result(MessageEventResult().message(f"会话 {session_id} 被限流。根据限流策略，此请求已被丢弃，直到您的限额于 {stall_duration:.2f} 秒后重置。"))
                         logger.info(
-                            f"会话 {session_id} 被限流。根据限流策略，此请求已被丢弃，直到限额于 {stall_duration:.2f} 秒后重置。"
+                            f"会话 {session_id} 被限流。根据限流策略，此请求已被丢弃，直到限额于 {stall_duration:.2f} 秒后重置。",
                         )
                         return event.stop_event()
 
-                self._remove_expired_timestamps(
-                    timestamps, now + timedelta(seconds=stall_duration)
-                )
-
-            timestamps.append(now)
-
     def _remove_expired_timestamps(
-        self, timestamps: Deque[datetime], now: datetime
+        self,
+        timestamps: deque[datetime],
+        now: datetime,
     ) -> None:
-        """
-        移除时间窗口外的时间戳。
+        """移除时间窗口外的时间戳。
 
         Args:
             timestamps (Deque[datetime]): 当前会话的时间戳队列。
             now (datetime): 当前时间，用于计算过期时间。
+
         """
         expiry_threshold: datetime = now - self.rate_limit_time
         while timestamps and timestamps[0] < expiry_threshold: