Refine agent background reply handling

This commit is contained in:
jxxghp
2026-04-30 00:25:23 +08:00
parent 11478faff3
commit 6532c60a3c
8 changed files with 128 additions and 50 deletions

View File

@@ -4,7 +4,7 @@ from unittest.mock import AsyncMock, patch
from langchain_core.messages import AIMessage
from app.agent import MoviePilotAgent
from app.agent import MoviePilotAgent, AgentManager, ReplyMode
from app.agent.memory import memory_manager
@@ -25,39 +25,11 @@ class _FakeAgent:
class AgentBackgroundOutputTest(unittest.IsolatedAsyncioTestCase):
async def test_background_non_streaming_still_sends_when_reply_not_suppressed(self):
async def test_background_non_streaming_does_not_send_by_default(self):
agent = MoviePilotAgent(session_id="bg-test", user_id="system")
agent.channel = None
agent.source = None
agent.suppress_user_reply = False
agent.persist_output_message = False
agent._tool_context = {"user_reply_sent": False}
agent._streamed_output = ""
agent.stream_handler = SimpleNamespace(
stop_streaming=AsyncMock(return_value=(False, ""))
)
agent._should_stream = lambda: False
agent._create_agent = lambda streaming=False: _FakeAgent(
[AIMessage(content="后台结果")]
)
agent.send_agent_message = AsyncMock()
agent._save_agent_message_to_db = AsyncMock()
with patch.object(memory_manager, "save_agent_messages") as save_messages:
await agent._execute_agent([])
agent.send_agent_message.assert_awaited_once_with(
"后台结果", title="MoviePilot助手"
)
agent._save_agent_message_to_db.assert_not_awaited()
save_messages.assert_called_once()
self.assertEqual("后台结果", agent._streamed_output)
async def test_background_non_streaming_persists_without_sending_when_reply_suppressed(self):
agent = MoviePilotAgent(session_id="bg-test", user_id="system")
agent.channel = None
agent.source = None
agent.suppress_user_reply = True
agent.reply_mode = ReplyMode.CAPTURE_ONLY
agent.persist_output_message = True
agent._tool_context = {"user_reply_sent": False}
agent._streamed_output = ""
@@ -81,6 +53,77 @@ class AgentBackgroundOutputTest(unittest.IsolatedAsyncioTestCase):
save_messages.assert_called_once()
self.assertEqual("后台结果", agent._streamed_output)
async def test_background_non_streaming_sends_when_reply_mode_dispatch(self):
agent = MoviePilotAgent(session_id="bg-test", user_id="system")
agent.channel = None
agent.source = None
agent.reply_mode = ReplyMode.DISPATCH
agent.persist_output_message = False
agent._tool_context = {"user_reply_sent": False}
agent._streamed_output = ""
agent.stream_handler = SimpleNamespace(
stop_streaming=AsyncMock(return_value=(False, ""))
)
agent._should_stream = lambda: False
agent._create_agent = lambda streaming=False: _FakeAgent(
[AIMessage(content="后台结果")]
)
agent.send_agent_message = AsyncMock()
agent._save_agent_message_to_db = AsyncMock()
with patch.object(memory_manager, "save_agent_messages") as save_messages:
await agent._execute_agent([])
agent.send_agent_message.assert_awaited_once_with(
"后台结果", title="MoviePilot助手"
)
agent._save_agent_message_to_db.assert_not_awaited()
save_messages.assert_called_once()
self.assertEqual("后台结果", agent._streamed_output)
async def test_background_non_streaming_persists_without_sending_when_capture_only(self):
agent = MoviePilotAgent(session_id="bg-test", user_id="system")
agent.channel = None
agent.source = None
agent.reply_mode = ReplyMode.CAPTURE_ONLY
agent.persist_output_message = True
agent._tool_context = {"user_reply_sent": False}
agent._streamed_output = ""
agent.stream_handler = SimpleNamespace(
stop_streaming=AsyncMock(return_value=(False, ""))
)
agent._should_stream = lambda: False
agent._create_agent = lambda streaming=False: _FakeAgent(
[AIMessage(content="后台结果")]
)
agent.send_agent_message = AsyncMock()
agent._save_agent_message_to_db = AsyncMock()
with patch.object(memory_manager, "save_agent_messages") as save_messages:
await agent._execute_agent([])
agent.send_agent_message.assert_not_awaited()
agent._save_agent_message_to_db.assert_awaited_once_with(
"后台结果", title="MoviePilot助手"
)
save_messages.assert_called_once()
self.assertEqual("后台结果", agent._streamed_output)
async def test_heartbeat_check_jobs_uses_dispatch_reply_mode(self):
manager = AgentManager()
with (
patch.object(manager, "_build_heartbeat_prompt", return_value="HEARTBEAT"),
patch.object(manager, "process_message", new=AsyncMock()) as process_message,
):
await manager.heartbeat_check_jobs()
process_message.assert_awaited_once()
self.assertEqual(
ReplyMode.DISPATCH,
process_message.await_args.kwargs["reply_mode"],
)
if __name__ == "__main__":
unittest.main()

View File

@@ -20,6 +20,7 @@ _stub_module("qbittorrentapi", TorrentFilesList=list)
_stub_module("transmission_rpc", File=object)
from app.agent.tools.factory import MoviePilotToolFactory
from app.agent import ReplyMode
from app.chain.search import SearchChain
from app.core.config import settings
@@ -125,7 +126,7 @@ class SearchChainAIRecommendTest(unittest.IsolatedAsyncioTestCase):
result = await chain._invoke_recommend_llm("Candidates")
self.assertEqual("[0, 2]", result)
self.assertTrue(captured["suppress_user_reply"])
self.assertEqual(ReplyMode.CAPTURE_ONLY, captured["reply_mode"])
self.assertFalse(captured["persist_output_message"])
self.assertFalse(captured["allow_message_tools"])