diff --git a/astrbot/core/agent/hooks.py b/astrbot/core/agent/hooks.py index 74ca6335b31..76c67df4bfe 100644 --- a/astrbot/core/agent/hooks.py +++ b/astrbot/core/agent/hooks.py @@ -9,7 +9,18 @@ class BaseAgentRunHooks(Generic[TContext]): + def should_buffer_streaming_response( + self, run_context: ContextWrapper[TContext] + ) -> bool: + """Whether response hooks may modify streamed response content.""" + return False + async def on_agent_begin(self, run_context: ContextWrapper[TContext]) -> None: ... + async def on_agent_response( + self, + run_context: ContextWrapper[TContext], + llm_response: LLMResponse, + ) -> None: ... async def on_tool_start( self, run_context: ContextWrapper[TContext], diff --git a/astrbot/core/agent/runners/tool_loop_agent_runner.py b/astrbot/core/agent/runners/tool_loop_agent_runner.py index f8418fe4205..bd8f993bd39 100644 --- a/astrbot/core/agent/runners/tool_loop_agent_runner.py +++ b/astrbot/core/agent/runners/tool_loop_agent_runner.py @@ -838,14 +838,17 @@ async def step(self): return self.run_context.messages = processed_messages self._simple_print_message_role("[AftCompact]", self.run_context.messages) + buffer_streaming_response = self.streaming and bool( + self.agent_hooks.should_buffer_streaming_response(self.run_context) + ) + buffered_streaming_responses: list[AgentResponse] = [] async for llm_response in self._iter_llm_responses_with_fallback(): if llm_response.is_chunk: if self.stats.time_to_first_token == 0: self.stats.time_to_first_token = time.time() - self.stats.start_time - if llm_response.reasoning_content: - yield AgentResponse( + response = AgentResponse( type="streaming_delta", data=AgentResponseData( chain=MessageChain(type="reasoning").message( @@ -853,18 +856,30 @@ async def step(self): ), ), ) + if buffer_streaming_response: + buffered_streaming_responses.append(response) + else: + yield response if llm_response.result_chain: - yield AgentResponse( + response = AgentResponse( type="streaming_delta", data=AgentResponseData(chain=llm_response.result_chain), ) + if buffer_streaming_response: + buffered_streaming_responses.append(response) + else: + yield response elif llm_response.completion_text: - yield AgentResponse( + response = AgentResponse( type="streaming_delta", data=AgentResponseData( chain=MessageChain().message(llm_response.completion_text), ), ) + if buffer_streaming_response: + buffered_streaming_responses.append(response) + else: + yield response if self._is_stop_requested(): break continue @@ -894,16 +909,22 @@ async def step(self): break # got final response if self._is_stop_requested(): + for response in buffered_streaming_responses: + yield response yield await self._finalize_aborted_step() return if not llm_resp_result: + for response in buffered_streaming_responses: + yield response return # 处理 LLM 响应 llm_resp = llm_resp_result if llm_resp.role == "err": + for response in buffered_streaming_responses: + yield response # 如果 LLM 响应错误,转换到错误状态 self.final_llm_resp = llm_resp self.stats.end_time = time.time() @@ -923,6 +944,38 @@ async def step(self): if not llm_resp.tools_call_name: await self._complete_with_assistant_response(llm_resp) + else: + try: + await self.agent_hooks.on_agent_response(self.run_context, llm_resp) + except Exception as e: + logger.error(f"Error in on_agent_response hook: {e}", exc_info=True) + if not llm_resp.tools_call_name: + await self._complete_with_assistant_response(llm_resp) + + # Response hooks can modify the complete response. Emit streaming deltas + # only after they run so clients never receive content the hook changed. + if buffer_streaming_response: + if llm_resp.reasoning_content: + yield AgentResponse( + type="streaming_delta", + data=AgentResponseData( + chain=MessageChain(type="reasoning").message( + llm_resp.reasoning_content, + ), + ), + ) + if llm_resp.result_chain: + yield AgentResponse( + type="streaming_delta", + data=AgentResponseData(chain=llm_resp.result_chain), + ) + elif llm_resp.completion_text: + yield AgentResponse( + type="streaming_delta", + data=AgentResponseData( + chain=MessageChain().message(llm_resp.completion_text), + ), + ) # 返回 LLM 结果 if llm_resp.reasoning_content: diff --git a/astrbot/core/astr_agent_hooks.py b/astrbot/core/astr_agent_hooks.py index 3155257d712..d53dcef869e 100644 --- a/astrbot/core/astr_agent_hooks.py +++ b/astrbot/core/astr_agent_hooks.py @@ -7,10 +7,24 @@ from astrbot.core.agent.tool import FunctionTool from astrbot.core.astr_agent_context import AstrAgentContext from astrbot.core.pipeline.context_utils import call_event_hook -from astrbot.core.star.star_handler import EventType +from astrbot.core.star.star import star_map +from astrbot.core.star.star_handler import EventType, star_handlers_registry class MainAgentHooks(BaseAgentRunHooks[AstrAgentContext]): + def should_buffer_streaming_response( + self, run_context: ContextWrapper[AstrAgentContext] + ) -> bool: + event = run_context.context.event + for handler in star_handlers_registry.get_handlers_by_event_type( + EventType.OnLLMResponseEvent, + plugins_name=event.plugins_name, + ): + plugin = star_map.get(handler.handler_module_path) + if plugin and not plugin.reserved: + return True + return False + async def on_agent_begin( self, run_context: ContextWrapper[AstrAgentContext] ) -> None: @@ -40,6 +54,13 @@ async def on_agent_done(self, run_context, llm_response) -> None: llm_response, ) + async def on_agent_response(self, run_context, llm_response) -> None: + await call_event_hook( + run_context.context.event, + EventType.OnLLMResponseEvent, + llm_response, + ) + async def on_tool_start( self, run_context: ContextWrapper[AstrAgentContext], diff --git a/tests/test_tool_loop_agent_runner.py b/tests/test_tool_loop_agent_runner.py index ef8e818bb5c..f4adb441f35 100644 --- a/tests/test_tool_loop_agent_runner.py +++ b/tests/test_tool_loop_agent_runner.py @@ -18,12 +18,14 @@ from astrbot.core.agent.run_context import ContextWrapper from astrbot.core.agent.runners.tool_loop_agent_runner import ToolLoopAgentRunner from astrbot.core.agent.tool import FunctionTool, ToolSet +from astrbot.core.astr_agent_hooks import MainAgentHooks from astrbot.core.astr_agent_run_util import run_agent from astrbot.core.astr_agent_tool_exec import FunctionToolExecutor from astrbot.core.exceptions import EmptyModelOutputError from astrbot.core.message.message_event_result import MessageChain from astrbot.core.provider.entities import LLMResponse, ProviderRequest, TokenUsage from astrbot.core.provider.provider import Provider +from astrbot.core.star.star_handler import EventType class MockProvider(Provider): @@ -398,9 +400,18 @@ class MockHooks(BaseAgentRunHooks): def __init__(self): self.agent_begin_called = False self.agent_done_called = False + self.agent_responses = [] + self.event_order = [] + self.events = [] + self.buffer_streaming_response = False + self.rewrite_response = False + self.clear_tool_calls = False self.tool_start_called = False self.tool_end_called = False + def should_buffer_streaming_response(self, run_context): + return self.buffer_streaming_response + async def on_agent_begin(self, run_context): self.agent_begin_called = True @@ -413,6 +424,15 @@ async def on_tool_end(self, run_context, tool, tool_args, tool_result): async def on_agent_done(self, run_context, llm_response): self.agent_done_called = True + async def on_agent_response(self, run_context, llm_response): + self.events.append("response_hook") + if self.clear_tool_calls: + llm_response.tools_call_name = [] + if self.rewrite_response: + llm_response.completion_text = f"hooked: {llm_response.completion_text}" + self.agent_responses.append(llm_response) + self.event_order.append("hook") + class ClearingAgentBeginHooks(MockHooks): async def on_agent_begin(self, run_context): @@ -912,6 +932,211 @@ async def test_hooks_called_with_max_step( assert mock_hooks.tool_end_called, "on_tool_end应该被调用" +@pytest.mark.asyncio +async def test_tool_call_response_invokes_response_hook( + runner, mock_provider, provider_request, mock_tool_executor, mock_hooks +): + mock_provider.should_call_tools = True + mock_provider.max_calls_before_normal_response = 100 + mock_hooks.buffer_streaming_response = True + mock_hooks.rewrite_response = True + + await runner.reset( + provider=mock_provider, + request=provider_request, + run_context=ContextWrapper(context=None), + tool_executor=mock_tool_executor, + agent_hooks=mock_hooks, + streaming=True, + ) + + responses = [] + async for response in runner.step(): + responses.append(response) + if ( + response.type == "llm_result" + and response.data["chain"].get_plain_text() + == "hooked: 我需要使用工具来帮助您" + ): + mock_hooks.event_order.append("text") + if response.type == "streaming_delta": + mock_hooks.events.append("streaming_delta") + + assert len(mock_hooks.agent_responses) == 1 + assert ( + mock_hooks.agent_responses[0].completion_text + == "hooked: 我需要使用工具来帮助您" + ) + assert mock_hooks.agent_responses[0].tools_call_name == ["test_tool"] + assert not mock_hooks.agent_done_called + assert mock_hooks.event_order == ["hook", "text"] + assert any( + response.type == "llm_result" + and response.data["chain"].get_plain_text() == "hooked: 我需要使用工具来帮助您" + for response in responses + ) + assert any( + response.type == "streaming_delta" + and response.data["chain"].get_plain_text() == "hooked: 我需要使用工具来帮助您" + for response in responses + ) + assert mock_hooks.events.index("response_hook") < mock_hooks.events.index( + "streaming_delta" + ) + + +@pytest.mark.asyncio +async def test_cleared_tool_calls_finalize_as_assistant_response( + runner, mock_provider, provider_request, mock_tool_executor, mock_hooks +): + mock_provider.should_call_tools = True + mock_provider.max_calls_before_normal_response = 100 + mock_hooks.clear_tool_calls = True + + await runner.reset( + provider=mock_provider, + request=provider_request, + run_context=ContextWrapper(context=None), + tool_executor=mock_tool_executor, + agent_hooks=mock_hooks, + streaming=False, + ) + + responses = [response async for response in runner.step()] + + assert runner.done() + assert runner.get_final_llm_resp().tools_call_name == [] + assert mock_hooks.agent_done_called + assert not mock_hooks.tool_start_called + assert any(response.type == "llm_result" for response in responses) + + +@pytest.mark.asyncio +async def test_buffered_stream_chunks_are_flushed_when_provider_fails( + runner, provider_request, mock_tool_executor, mock_hooks +): + class FailingStreamingProvider(MockProvider): + async def text_chat_stream(self, **kwargs): + yield LLMResponse( + role="assistant", + completion_text="partial response", + is_chunk=True, + ) + raise RuntimeError("stream failed") + + mock_hooks.buffer_streaming_response = True + + await runner.reset( + provider=FailingStreamingProvider(), + request=provider_request, + run_context=ContextWrapper(context=None), + tool_executor=mock_tool_executor, + agent_hooks=mock_hooks, + streaming=True, + ) + + responses = [response async for response in runner.step()] + + assert any( + response.type == "streaming_delta" + and response.data["chain"].get_plain_text() == "partial response" + for response in responses + ) + assert any(response.type == "err" for response in responses) + + +@pytest.mark.asyncio +async def test_buffered_partial_stream_is_flushed_when_stopped( + runner, provider_request, mock_tool_executor, mock_hooks +): + class PausedStreamingProvider(MockProvider): + def __init__(self): + super().__init__() + self.chunk_yielded = asyncio.Event() + + async def text_chat_stream(self, **kwargs): + yield LLMResponse( + role="assistant", + completion_text="partial response", + is_chunk=True, + ) + self.chunk_yielded.set() + await asyncio.Future() + + provider = PausedStreamingProvider() + mock_hooks.buffer_streaming_response = True + + await runner.reset( + provider=provider, + request=provider_request, + run_context=ContextWrapper(context=None), + tool_executor=mock_tool_executor, + agent_hooks=mock_hooks, + streaming=True, + ) + + async def collect_responses(): + return [response async for response in runner.step()] + + step_task = asyncio.create_task(collect_responses()) + await asyncio.wait_for(provider.chunk_yielded.wait(), timeout=5) + runner.request_stop() + responses = await asyncio.wait_for(step_task, timeout=5) + + assert any( + response.type == "streaming_delta" + and response.data["chain"].get_plain_text() == "partial response" + for response in responses + ) + assert any(response.type == "aborted" for response in responses) + + +@pytest.mark.asyncio +async def test_main_agent_response_hook_dispatches_llm_response_event(monkeypatch): + event = MagicMock() + llm_response = LLMResponse( + role="assistant", + completion_text="intermediate response", + tools_call_name=["test_tool"], + tools_call_args=[{"query": "test"}], + tools_call_ids=["call_123"], + ) + call_event_hook = AsyncMock() + monkeypatch.setattr( + "astrbot.core.astr_agent_hooks.call_event_hook", call_event_hook + ) + + await MainAgentHooks().on_agent_response( + ContextWrapper(context=MockAgentContext(event)), llm_response + ) + + call_event_hook.assert_awaited_once_with( + event, EventType.OnLLMResponseEvent, llm_response + ) + + +@pytest.mark.parametrize( + ("reserved", "expected"), + [(True, False), (False, True)], +) +def test_main_agent_buffers_only_for_external_response_handlers( + monkeypatch, reserved, expected +): + from astrbot.core.star.star import star_map + + module_path = "tests.response_handler" + handler = SimpleNamespace(handler_module_path=module_path) + event = SimpleNamespace(plugins_name=None) + run_context = ContextWrapper(context=SimpleNamespace(event=event)) + monkeypatch.setattr( + "astrbot.core.astr_agent_hooks.star_handlers_registry.get_handlers_by_event_type", + lambda *args, **kwargs: [handler], + ) + monkeypatch.setitem(star_map, module_path, SimpleNamespace(reserved=reserved)) + + assert MainAgentHooks().should_buffer_streaming_response(run_context) is expected + + @pytest.mark.asyncio async def test_tool_result_includes_all_calltoolresult_content( runner, mock_provider, provider_request, mock_hooks, monkeypatch