From 850b8c64d4aacc26f84db38b50ca64c9c07348b2 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Mon, 18 Aug 2025 12:02:13 +0000 Subject: [PATCH] Add OnToolStartEvent and OnToolEndEvent hook support Co-authored-by: Soulter <37870767+Soulter@users.noreply.github.com> --- astrbot/api/event/filter/__init__.py | 4 ++ .../agent/runners/tool_loop_agent_runner.py | 6 +-- .../process_stage/method/llm_request.py | 12 +++++ astrbot/core/star/register/__init__.py | 4 ++ astrbot/core/star/register/star_handler.py | 49 +++++++++++++++++++ astrbot/core/star/star_handler.py | 2 + astrbot/dashboard/routes/plugin.py | 2 + 7 files changed, 76 insertions(+), 3 deletions(-) diff --git a/astrbot/api/event/filter/__init__.py b/astrbot/api/event/filter/__init__.py index dd737e3f..ec8ac8d5 100644 --- a/astrbot/api/event/filter/__init__.py +++ b/astrbot/api/event/filter/__init__.py @@ -12,6 +12,8 @@ from astrbot.core.star.register import ( register_llm_tool as llm_tool, register_on_decorating_result as on_decorating_result, register_after_message_sent as after_message_sent, + register_on_tool_start as on_tool_start, + register_on_tool_end as on_tool_end, ) from astrbot.core.star.filter.event_message_type import ( @@ -46,4 +48,6 @@ __all__ = [ "on_decorating_result", "after_message_sent", "on_llm_response", + "on_tool_start", + "on_tool_end", ] diff --git a/astrbot/core/agent/runners/tool_loop_agent_runner.py b/astrbot/core/agent/runners/tool_loop_agent_runner.py index c38285f5..88d54234 100644 --- a/astrbot/core/agent/runners/tool_loop_agent_runner.py +++ b/astrbot/core/agent/runners/tool_loop_agent_runner.py @@ -272,7 +272,7 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]): try: await self.agent_hooks.on_tool_end( self.run_context, - func_tool_name, + func_tool, func_tool_args, resp, ) @@ -291,7 +291,7 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]): ) try: await self.agent_hooks.on_tool_end( - self.run_context, func_tool_name, func_tool_args, None + self.run_context, func_tool, func_tool_args, None ) except Exception as e: logger.error( @@ -304,7 +304,7 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]): try: await self.agent_hooks.on_tool_end( - self.run_context, func_tool_name, func_tool_args, None + self.run_context, func_tool, func_tool_args, None ) except Exception as e: logger.error( diff --git a/astrbot/core/pipeline/process_stage/method/llm_request.py b/astrbot/core/pipeline/process_stage/method/llm_request.py index 8f6e492f..e6ea1dd0 100644 --- a/astrbot/core/pipeline/process_stage/method/llm_request.py +++ b/astrbot/core/pipeline/process_stage/method/llm_request.py @@ -200,6 +200,18 @@ class FunctionToolExecutor(BaseFunctionToolExecutor[AstrAgentContext]): class MainAgentHooks(BaseAgentRunHooks[AgentContextWrapper]): + async def on_tool_start(self, run_context, tool, tool_args): + # 执行 Tool 开始事件钩子 + await call_event_hook( + run_context.event, EventType.OnToolStartEvent, tool, tool_args + ) + + async def on_tool_end(self, run_context, tool, tool_args, tool_result): + # 执行 Tool 完成事件钩子 + await call_event_hook( + run_context.event, EventType.OnToolEndEvent, tool, tool_args, tool_result + ) + async def on_agent_done(self, run_context, llm_response): # 执行事件钩子 await call_event_hook( diff --git a/astrbot/core/star/register/__init__.py b/astrbot/core/star/register/__init__.py index 55a4393d..f8049b68 100644 --- a/astrbot/core/star/register/__init__.py +++ b/astrbot/core/star/register/__init__.py @@ -14,6 +14,8 @@ from .star_handler import ( register_agent, register_on_decorating_result, register_after_message_sent, + register_on_tool_start, + register_on_tool_end, ) __all__ = [ @@ -32,4 +34,6 @@ __all__ = [ "register_agent", "register_on_decorating_result", "register_after_message_sent", + "register_on_tool_start", + "register_on_tool_end", ] diff --git a/astrbot/core/star/register/star_handler.py b/astrbot/core/star/register/star_handler.py index 8f82651a..a6d8f965 100644 --- a/astrbot/core/star/register/star_handler.py +++ b/astrbot/core/star/register/star_handler.py @@ -450,3 +450,52 @@ def register_after_message_sent(**kwargs): return awaitable return decorator + + +def register_on_tool_start(**kwargs): + """当 Tool 请求开始时的事件 + + Examples: + ```py + from astrbot.api.event import AstrMessageEvent + from astrbot.core.agent.tool import FunctionTool + + @register_on_tool_start() + async def test(event: AstrMessageEvent, tool: FunctionTool, tool_args: dict) -> None: + # 在工具调用开始时执行的逻辑 + logger.info(f"Tool {tool.name} started with args: {tool_args}") + ``` + + 请务必接收三个参数：event, tool, tool_args + """ + + def decorator(awaitable): + _ = get_handler_or_create(awaitable, EventType.OnToolStartEvent, **kwargs) + return awaitable + + return decorator + + +def register_on_tool_end(**kwargs): + """当 Tool 请求完成时的事件 + + Examples: + ```py + from astrbot.api.event import AstrMessageEvent + from astrbot.core.agent.tool import FunctionTool + import mcp.types + + @register_on_tool_end() + async def test(event: AstrMessageEvent, tool: FunctionTool, tool_args: dict, tool_result: mcp.types.CallToolResult | None) -> None: + # 在工具调用完成时执行的逻辑 + logger.info(f"Tool {tool.name} finished with result: {tool_result}") + ``` + + 请务必接收四个参数：event, tool, tool_args, tool_result + """ + + def decorator(awaitable): + _ = get_handler_or_create(awaitable, EventType.OnToolEndEvent, **kwargs) + return awaitable + + return decorator diff --git a/astrbot/core/star/star_handler.py b/astrbot/core/star/star_handler.py index 3a507708..945fefbd 100644 --- a/astrbot/core/star/star_handler.py +++ b/astrbot/core/star/star_handler.py @@ -86,6 +86,8 @@ class EventType(enum.Enum): OnLLMResponseEvent = enum.auto() # LLM 响应后 OnDecoratingResultEvent = enum.auto() # 发送消息前 OnCallingFuncToolEvent = enum.auto() # 调用函数工具 + OnToolStartEvent = enum.auto() # Tool 请求开始 + OnToolEndEvent = enum.auto() # Tool 请求完成 OnAfterMessageSentEvent = enum.auto() # 发送消息后 diff --git a/astrbot/dashboard/routes/plugin.py b/astrbot/dashboard/routes/plugin.py index c0276767..69281bb2 100644 --- a/astrbot/dashboard/routes/plugin.py +++ b/astrbot/dashboard/routes/plugin.py @@ -53,6 +53,8 @@ class PluginRoute(Route): EventType.OnLLMResponseEvent: "LLM 响应后", EventType.OnDecoratingResultEvent: "回复消息前", EventType.OnCallingFuncToolEvent: "函数工具", + EventType.OnToolStartEvent: "Tool 请求开始", + EventType.OnToolEndEvent: "Tool 请求完成", EventType.OnAfterMessageSentEvent: "发送消息后", }