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

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

@@ -6,7 +6,8 @@ import asyncio
 import copy
 import json
 import traceback
-from typing import AsyncGenerator, Union
+from datetime import timedelta
+from collections.abc import AsyncGenerator
 from astrbot.core.conversation_mgr import Conversation
 from astrbot.core import logger
 from astrbot.core.message.components import Image
@@ -32,6 +33,7 @@ from astrbot.core.star.star_handler import EventType
 from astrbot.core.utils.metrics import Metric
 from ...context import PipelineContext, call_event_hook, call_handler
 from ..stage import Stage
+from ..utils import inject_kb_context
 from astrbot.core.provider.register import llm_tools
 from astrbot.core.star.star_handler import star_map
 from astrbot.core.astr_agent_context import AstrAgentContext
@@ -43,7 +45,7 @@ except (ModuleNotFoundError, ImportError):
 
 
 AgentContextWrapper = ContextWrapper[AstrAgentContext]
-AgentRunner = ToolLoopAgentRunner[AgentContextWrapper]
+AgentRunner = ToolLoopAgentRunner[AstrAgentContext]
 
 
 class FunctionToolExecutor(BaseFunctionToolExecutor[AstrAgentContext]):
@@ -101,7 +103,7 @@ class FunctionToolExecutor(BaseFunctionToolExecutor[AstrAgentContext]):
 
         request = ProviderRequest(
             prompt=input_,
-            system_prompt=tool.description,
+            system_prompt=tool.description or "",
             image_urls=[],  # 暂时不传递原始 agent 的上下文
             contexts=[],  # 暂时不传递原始 agent 的上下文
             func_tool=toolset,
@@ -185,21 +187,33 @@ class FunctionToolExecutor(BaseFunctionToolExecutor[AstrAgentContext]):
             handler=awaitable,
             **tool_args,
         )
-        async for resp in wrapper:
-            if resp is not None:
-                if isinstance(resp, mcp.types.CallToolResult):
-                    yield resp
+        # async for resp in wrapper:
+        while True:
+            try:
+                resp = await asyncio.wait_for(
+                    anext(wrapper),
+                    timeout=run_context.context.tool_call_timeout,
+                )
+                if resp is not None:
+                    if isinstance(resp, mcp.types.CallToolResult):
+                        yield resp
+                    else:
+                        text_content = mcp.types.TextContent(
+                            type="text",
+                            text=str(resp),
+                        )
+                        yield mcp.types.CallToolResult(content=[text_content])
                 else:
-                    text_content = mcp.types.TextContent(
-                        type="text",
-                        text=str(resp),
-                    )
-                    yield mcp.types.CallToolResult(content=[text_content])
-            else:
-                # NOTE: Tool 在这里直接请求发送消息给用户
-                # TODO: 是否需要判断 event.get_result() 是否为空?
-                # 如果为空,则说明没有发送消息给用户,并且返回值为空,将返回一个特殊的 TextContent,其内容如"工具没有返回内容"
-                yield None
+                    # NOTE: Tool 在这里直接请求发送消息给用户
+                    # TODO: 是否需要判断 event.get_result() 是否为空?
+                    # 如果为空,则说明没有发送消息给用户,并且返回值为空,将返回一个特殊的 TextContent,其内容如"工具没有返回内容"
+                    yield None
+            except asyncio.TimeoutError:
+                raise Exception(
+                    f"tool {tool.name} execution timeout after {run_context.context.tool_call_timeout} seconds."
+                )
+            except StopAsyncIteration:
+                break
 
     @classmethod
     async def _execute_mcp(
@@ -217,13 +231,16 @@ class FunctionToolExecutor(BaseFunctionToolExecutor[AstrAgentContext]):
         res = await session.call_tool(
             name=tool.name,
             arguments=tool_args,
+            read_timeout_seconds=timedelta(
+                seconds=run_context.context.tool_call_timeout
+            ),
         )
         if not res:
             return
         yield res
 
 
-class MainAgentHooks(BaseAgentRunHooks[AgentContextWrapper]):
+class MainAgentHooks(BaseAgentRunHooks[AstrAgentContext]):
     async def on_agent_done(self, run_context, llm_response):
         # 执行事件钩子
         await call_event_hook(
@@ -307,6 +324,7 @@ class LLMRequestSubStage(Stage):
         )
         self.streaming_response: bool = settings["streaming_response"]
         self.max_step: int = settings.get("max_agent_step", 30)
+        self.tool_call_timeout: int = settings.get("tool_call_timeout", 60)
         if isinstance(self.max_step, bool):  # workaround: #2622
             self.max_step = 30
         self.show_tool_use: bool = settings.get("show_tool_use_status", True)
@@ -320,7 +338,7 @@ class LLMRequestSubStage(Stage):
 
         self.conv_manager = ctx.plugin_manager.context.conversation_manager
 
-    def _select_provider(self, event: AstrMessageEvent) -> Provider | None:
+    def _select_provider(self, event: AstrMessageEvent):
         """选择使用的 LLM 提供商"""
         sel_provider = event.get_extra("selected_provider")
         _ctx = self.ctx.plugin_manager.context
@@ -350,7 +368,7 @@ class LLMRequestSubStage(Stage):
 
     async def process(
         self, event: AstrMessageEvent, _nested: bool = False
-    ) -> Union[None, AsyncGenerator[None, None]]:
+    ) -> None | AsyncGenerator[None, None]:
         req: ProviderRequest | None = None
 
         if not self.ctx.astrbot_config["provider_settings"]["enable"]:
@@ -365,6 +383,9 @@ class LLMRequestSubStage(Stage):
         provider = self._select_provider(event)
         if provider is None:
             return
+        if not isinstance(provider, Provider):
+            logger.error(f"选择的提供商类型无效({type(provider)})，跳过 LLM 请求处理。")
+            return
 
         if event.get_extra("provider_request"):
             req = event.get_extra("provider_request")
@@ -399,6 +420,14 @@ class LLMRequestSubStage(Stage):
         if not req.prompt and not req.image_urls:
             return
 
+        # 应用知识库
+        try:
+            await inject_kb_context(
+                umo=event.unified_msg_origin, p_ctx=self.ctx, req=req
+            )
+        except Exception as e:
+            logger.error(f"调用知识库时遇到问题: {e}")
+
         # 执行请求 LLM 前事件钩子。
         if await call_event_hook(event, EventType.OnLLMRequestEvent, req):
             return
@@ -463,6 +492,9 @@ class LLMRequestSubStage(Stage):
                     new_tool_set.add_tool(tool)
             req.func_tool = new_tool_set
 
+        # 备份 req.contexts
+        backup_contexts = copy.deepcopy(req.contexts)
+
         # run agent
         agent_runner = AgentRunner()
         logger.debug(
@@ -473,6 +505,7 @@ class LLMRequestSubStage(Stage):
             first_provider_request=req,
             curr_provider_request=req,
             streaming=self.streaming_response,
+            tool_call_timeout=self.tool_call_timeout,
         )
         await agent_runner.reset(
             provider=provider,
@@ -499,8 +532,10 @@ class LLMRequestSubStage(Stage):
                         chain = (
                             MessageChain().message(final_llm_resp.completion_text).chain
                         )
-                    else:
+                    elif final_llm_resp.result_chain:
                         chain = final_llm_resp.result_chain.chain
+                    else:
+                        chain = MessageChain().chain
                     event.set_result(
                         MessageEventResult(
                             chain=chain,
@@ -511,6 +546,9 @@ class LLMRequestSubStage(Stage):
             async for _ in run_agent(agent_runner, self.max_step, self.show_tool_use):
                 yield
 
+        # 恢复备份的 contexts
+        req.contexts = backup_contexts
+
         await self._save_to_history(event, req, agent_runner.get_final_llm_resp())
 
         # 异步处理 WebChat 特殊情况
@@ -529,6 +567,8 @@ class LLMRequestSubStage(Stage):
         self, event: AstrMessageEvent, req: ProviderRequest, prov: Provider
     ):
         """处理 WebChat 平台的特殊情况，包括第一次 LLM 对话时总结对话内容生成 title"""
+        if not req.conversation:
+            return
         conversation = await self.conv_manager.get_conversation(
             event.unified_msg_origin, req.conversation.cid
         )

astrbot/core/pipeline/process_stage/utils.py

@@ -0,0 +1,80 @@
+from ..context import PipelineContext
+from astrbot.core.provider.entities import ProviderRequest
+from astrbot.api import logger, sp
+
+
+async def inject_kb_context(
+    umo: str,
+    p_ctx: PipelineContext,
+    req: ProviderRequest,
+) -> None:
+    """inject knowledge base context into the provider request
+
+    Args:
+        umo: Unique message object (session ID)
+        p_ctx: Pipeline context
+        req: Provider request
+    """
+
+    kb_mgr = p_ctx.plugin_manager.context.kb_manager
+
+    # 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 = p_ctx.astrbot_config.get("kb_names", [])
+        top_k = p_ctx.astrbot_config.get("kb_final_top_k", 5)
+        logger.debug(f"[知识库] 使用全局配置，知识库数量: {len(kb_names)}")
+
+    top_k_fusion = p_ctx.astrbot_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=req.prompt,
+        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)} 条相关知识块")
+        req.system_prompt = f"{formatted}\n\n{req.system_prompt or ''}"

astrbot/core/pipeline/scheduler.py

@@ -74,7 +74,7 @@ class PipelineScheduler:
         await self._process_stages(event)
 
         # 如果没有发送操作, 则发送一个空消息, 以便于后续的处理
-        if event.get_platform_name() == "webchat":
+        if event.get_platform_name() in ["webchat", "wecom_ai_bot"]:
             await event.send(None)
 
         logger.debug("pipeline 执行完毕。")

astrbot/core/platform/astr_message_event.py

@@ -4,7 +4,7 @@ import re
 import hashlib
 import uuid
 
-from typing import List, Union, Optional, AsyncGenerator, TypeVar, Any
+from typing import List, Union, Optional, AsyncGenerator, Any
 
 from astrbot import logger
 from astrbot.core.db.po import Conversation
@@ -26,8 +26,6 @@ from .astrbot_message import AstrBotMessage, Group
 from .platform_metadata import PlatformMetadata
 from .message_session import MessageSession, MessageSesion  # noqa
 
-_VT = TypeVar("_VT")
-
 
 class AstrMessageEvent(abc.ABC):
     def __init__(
@@ -92,8 +90,10 @@ class AstrMessageEvent(abc.ABC):
         """
         return self.message_str
 
-    def _outline_chain(self, chain: List[BaseMessageComponent]) -> str:
+    def _outline_chain(self, chain: Optional[List[BaseMessageComponent]]) -> str:
         outline = ""
+        if not chain:
+            return outline
         for i in chain:
             if isinstance(i, Plain):
                 outline += i.text
@@ -175,9 +175,7 @@ class AstrMessageEvent(abc.ABC):
         """
         self._extras[key] = value
 
-    def get_extra(
-        self, key: str | None = None, default: _VT = None
-    ) -> dict[str, Any] | _VT:
+    def get_extra(self, key: str | None = None, default=None) -> Any:
         """
         获取额外的信息。
         """
@@ -265,6 +263,9 @@ class AstrMessageEvent(abc.ABC):
         """
         if isinstance(result, str):
             result = MessageEventResult().message(result)
+        # 兼容外部插件或调用方传入的 chain=None 的情况，确保为可迭代列表
+        if isinstance(result, MessageEventResult) and result.chain is None:
+            result.chain = []
         self._result = result
 
     def stop_event(self):

astrbot/core/platform/manager.py

@@ -82,6 +82,10 @@ class PlatformManager:
                     from .sources.wecom.wecom_adapter import (
                         WecomPlatformAdapter,  # noqa: F401
                     )
+                case "wecom_ai_bot":
+                    from .sources.wecom_ai_bot.wecomai_adapter import (
+                        WecomAIBotAdapter,  # noqa: F401
+                    )
                 case "weixin_official_account":
                     from .sources.weixin_official_account.weixin_offacc_adapter import (
                         WeixinOfficialAccountPlatformAdapter,  # noqa: F401