Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 20 additions & 0 deletions astrbot/core/agent/runners/tool_loop_agent_runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -970,6 +970,26 @@ async def step(self):
)

await self._complete_with_assistant_response(llm_resp)
# Re-query uses text_chat(), so its reply has no stream chunks.
# Supply them after hooks without changing llm_result ordering.
if self.streaming:
if llm_resp.reasoning_content:
yield AgentResponse(
type="streaming_delta",
data=AgentResponseData(
chain=MessageChain(type="reasoning").message(
llm_resp.reasoning_content,
),
),
)
chain = llm_resp.result_chain
if not chain and llm_resp.completion_text:
chain = MessageChain().message(llm_resp.completion_text)
if chain:
yield AgentResponse(
type="streaming_delta",
data=AgentResponseData(chain=chain),
)
return
else:
llm_resp.tools_call_name = requery_resp.tools_call_name
Expand Down
139 changes: 137 additions & 2 deletions tests/test_tool_loop_agent_runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
from pathlib import Path
from types import SimpleNamespace
from typing import Any, cast
from unittest.mock import AsyncMock
from unittest.mock import AsyncMock, MagicMock

import pytest

Expand All @@ -18,8 +18,10 @@
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_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

Expand Down Expand Up @@ -1653,7 +1655,8 @@ async def test_follow_up_ticket_not_consumed_when_no_next_tool_call(


@pytest.mark.asyncio
async def test_skills_like_requery_passes_extra_user_content_parts():
@pytest.mark.parametrize("streaming", [False, True])
async def test_skills_like_requery_passes_extra_user_content_parts(streaming):
"""skills-like 模式 re-query 时应传递 extra_user_content_parts(如 image_caption)"""
from astrbot.core.agent.message import TextPart

Expand Down Expand Up @@ -1719,6 +1722,7 @@ async def text_chat(self, **kwargs) -> LLMResponse:
tool_executor=cast(Any, MockToolExecutor()),
agent_hooks=MockHooks(),
tool_schema_mode="skills_like",
streaming=streaming,
)

async for _ in runner.step():
Expand All @@ -1733,6 +1737,137 @@ async def text_chat(self, **kwargs) -> LLMResponse:
assert parts[0].text == "<image_caption>一张猫的照片</image_caption>"


@pytest.mark.asyncio
@pytest.mark.parametrize(
("streaming", "stream_to_general", "show_reasoning"),
[
(True, False, False),
(True, False, True),
(True, True, True),
(False, False, True),
],
)
@pytest.mark.parametrize("use_result_chain", [False, True])
async def test_skills_like_requery_reply_reaches_stream_bridge_once(
runner,
provider_request,
mock_tool_executor,
mock_hooks,
streaming,
stream_to_general,
show_reasoning,
use_result_chain,
):
"""Deliver a non-streaming re-query reply through the real runner and bridge."""
final_text = "The search is complete: two pushes."
reasoning = "The existing tool result is sufficient."

class RequeryReplyProvider(MockProvider):
async def text_chat(self, **kwargs) -> LLMResponse:
self.call_count += 1
if self.call_count == 1:
return LLMResponse(
role="assistant",
completion_text="Let me check.",
tools_call_name=["test_tool"],
tools_call_args=[{}],
tools_call_ids=["select_tool"],
)
assert self.call_count == 2
return LLMResponse(
role="assistant",
completion_text=None if use_result_chain else final_text,
result_chain=MessageChain().message(final_text)
if use_result_chain
else None,
reasoning_content=reasoning,
)

provider = RequeryReplyProvider()
event = MagicMock()
event.is_stopped.return_value = False
event.get_extra.return_value = None
event.get_platform_name.return_value = "lark"
await runner.reset(
provider=provider,
request=provider_request,
run_context=ContextWrapper(context=MockAgentContext(event)),
tool_executor=mock_tool_executor,
agent_hooks=mock_hooks,
streaming=streaming,
tool_schema_mode="skills_like",
)
original_step = runner.step
final_events = []
hooks_at_emission = []

async def recorded_step():
async for response in original_step():
chain = response.data["chain"]
if chain.get_plain_text() in (final_text, reasoning):
final_events.append((response.type, chain.type))
hooks_at_emission.append(mock_hooks.agent_done_called)
yield response

runner.step = recorded_step
chains = [
chain
async for chain in run_agent(
runner,
stream_to_general=stream_to_general,
show_reasoning=show_reasoning,
)
]
assert sum(chain.get_plain_text() == final_text for chain in chains) == 1
assert sum(chain.get_plain_text() == "Let me check." for chain in chains) == 1
assert sum(chain.get_plain_text() == reasoning for chain in chains) == int(
streaming and not stream_to_general and show_reasoning
)
expected_types = ["llm_result", "streaming_delta"] if streaming else ["llm_result"]
assert final_events == [
(response_type, chain_type)
for response_type in expected_types
for chain_type in ("reasoning", None)
]
# Preserve existing llm_result ordering; only new deltas follow the hooks.
assert hooks_at_emission == [False, False] + ([True, True] if streaming else [])
assert runner.done()
assert runner.get_final_llm_resp().completion_text == final_text
assert runner.run_context.messages[-1].content[-1].text == final_text
assert not mock_hooks.tool_start_called
assert provider.call_count == 2


@pytest.mark.asyncio
async def test_normal_streaming_reply_is_not_duplicated_by_stream_bridge(
runner,
mock_provider,
provider_request,
mock_tool_executor,
mock_hooks,
):
mock_provider.should_call_tools = False
event = MagicMock()
event.is_stopped.return_value = False
event.get_extra.return_value = None
event.get_platform_name.return_value = "lark"
await runner.reset(
provider=mock_provider,
request=provider_request,
run_context=ContextWrapper(context=MockAgentContext(event)),
tool_executor=mock_tool_executor,
agent_hooks=mock_hooks,
streaming=True,
tool_schema_mode="skills_like",
)

chains = [chain async for chain in run_agent(runner)]

assert [chain.get_plain_text() for chain in chains] == ["这是我的最终回答"]
assert mock_provider.call_count == 1
assert runner.done()


def test_skills_like_requery_preserves_existing_context_prefix():
messages = [
Message(role="system", content="stable system prompt"),
Expand Down
Loading