mirror of
https://github.com/jxxghp/MoviePilot.git
synced 2026-05-23 07:26:46 +00:00
224 lines
8.6 KiB
Python
224 lines
8.6 KiB
Python
import asyncio
|
||
import unittest
|
||
from unittest.mock import AsyncMock, Mock, patch
|
||
|
||
from app.agent.prompt import prompt_manager
|
||
from app.agent.tools.factory import MoviePilotToolFactory
|
||
from app.agent.tools.impl.ask_user_choice import (
|
||
AskUserChoiceTool,
|
||
UserChoiceOptionInput,
|
||
)
|
||
from app.helper.interaction import (
|
||
AgentInteractionOption,
|
||
agent_interaction_manager,
|
||
)
|
||
from app.chain.message import MessageChain
|
||
from app.core.config import settings
|
||
from app.schemas.types import MessageChannel
|
||
|
||
|
||
class TestAgentInteraction(unittest.TestCase):
|
||
def tearDown(self):
|
||
agent_interaction_manager.clear()
|
||
|
||
def test_prompt_injects_choice_tool_hint_only_for_button_channels(self):
|
||
telegram_prompt = prompt_manager.get_agent_prompt(
|
||
channel=MessageChannel.Telegram.value
|
||
)
|
||
wechat_prompt = prompt_manager.get_agent_prompt(
|
||
channel=MessageChannel.Wechat.value
|
||
)
|
||
|
||
self.assertIn("ask_user_choice", telegram_prompt)
|
||
self.assertNotIn("ask_user_choice", wechat_prompt)
|
||
|
||
def test_factory_injects_choice_tool_only_for_button_channels(self):
|
||
with patch(
|
||
"app.agent.tools.factory.PluginManager.get_plugin_agent_tools",
|
||
return_value=[],
|
||
):
|
||
telegram_tools = MoviePilotToolFactory.create_tools(
|
||
session_id="session-1",
|
||
user_id="10001",
|
||
channel=MessageChannel.Telegram.value,
|
||
source="telegram-test",
|
||
username="tester",
|
||
)
|
||
wechat_tools = MoviePilotToolFactory.create_tools(
|
||
session_id="session-2",
|
||
user_id="10001",
|
||
channel=MessageChannel.Wechat.value,
|
||
source="wechat-test",
|
||
username="tester",
|
||
)
|
||
|
||
self.assertIn("ask_user_choice", [tool.name for tool in telegram_tools])
|
||
self.assertNotIn("ask_user_choice", [tool.name for tool in wechat_tools])
|
||
|
||
def test_choice_tool_returns_direct_after_sending_interaction(self):
|
||
"""发送按钮后应结束当前 Agent 轮次，等待用户选择作为新消息进入。"""
|
||
tool = AskUserChoiceTool(session_id="session-1", user_id="10001")
|
||
|
||
self.assertTrue(tool.return_direct)
|
||
|
||
def test_choice_tool_sends_buttons_and_registers_pending_request(self):
|
||
tool = AskUserChoiceTool(session_id="session-1", user_id="10001")
|
||
tool.set_message_attr(
|
||
channel=MessageChannel.Telegram.value,
|
||
source="telegram-test",
|
||
username="tester",
|
||
)
|
||
tool.set_agent_context(agent_context={})
|
||
|
||
with patch(
|
||
"app.agent.tools.impl.ask_user_choice.ToolChain.async_post_message",
|
||
new=AsyncMock(),
|
||
) as async_post_message:
|
||
result = asyncio.run(
|
||
tool.run(
|
||
message="请选择要执行的操作",
|
||
options=[
|
||
UserChoiceOptionInput(label="继续下载", value="继续下载"),
|
||
UserChoiceOptionInput(label="先看详情", value="先看详情"),
|
||
],
|
||
title="需要你的选择",
|
||
)
|
||
)
|
||
|
||
self.assertIn("等待用户选择", result)
|
||
self.assertTrue(tool._agent_context.get("user_reply_sent"))
|
||
notification = async_post_message.await_args.args[0]
|
||
self.assertEqual(notification.text, "请选择要执行的操作")
|
||
self.assertEqual(sum(len(row) for row in notification.buttons), 2)
|
||
|
||
callback_data = notification.buttons[0][0]["callback_data"]
|
||
_, _, request_id, option_index = callback_data.split(":")
|
||
resolved = agent_interaction_manager.resolve(
|
||
request_id, int(option_index), "10001"
|
||
)
|
||
self.assertIsNotNone(resolved)
|
||
_, option = resolved
|
||
self.assertEqual(option.value, "继续下载")
|
||
|
||
def test_choice_tool_blocks_after_feedback_quality_rejection(self):
|
||
tool = AskUserChoiceTool(session_id="session-feedback", user_id="10001")
|
||
tool.set_message_attr(
|
||
channel=MessageChannel.Telegram.value,
|
||
source="telegram-test",
|
||
username="tester",
|
||
)
|
||
tool.set_agent_context(
|
||
agent_context={"feedback_issue_rejected_quality": True}
|
||
)
|
||
|
||
with patch(
|
||
"app.agent.tools.impl.ask_user_choice.ToolChain.async_post_message",
|
||
new=AsyncMock(),
|
||
) as async_post_message:
|
||
result = asyncio.run(
|
||
tool.run(
|
||
message="测试ISSUE提交被系统质量校验拦截，请选择：",
|
||
options=[
|
||
UserChoiceOptionInput(
|
||
label="提供真实问题描述重新提交",
|
||
value="提供真实问题描述重新提交",
|
||
),
|
||
UserChoiceOptionInput(
|
||
label="取消测试，了解原因",
|
||
value="取消测试，了解原因",
|
||
),
|
||
],
|
||
)
|
||
)
|
||
|
||
self.assertIn("质量门槛拒绝", result)
|
||
async_post_message.assert_not_awaited()
|
||
|
||
def test_agent_interaction_callback_routes_selected_value_back_to_agent(self):
|
||
chain = MessageChain()
|
||
request = agent_interaction_manager.create_request(
|
||
session_id="session-choice",
|
||
user_id="10001",
|
||
channel=MessageChannel.Telegram.value,
|
||
source="telegram-test",
|
||
username="tester",
|
||
title="需要你的选择",
|
||
prompt="请选择",
|
||
options=[
|
||
AgentInteractionOption(label="电影", value="我选择电影"),
|
||
AgentInteractionOption(label="电视剧", value="我选择电视剧"),
|
||
],
|
||
)
|
||
|
||
with patch.object(settings, "AI_AGENT_ENABLE", True), patch.object(
|
||
chain.messagehelper, "put"
|
||
) as message_put, patch.object(
|
||
chain.messageoper, "add"
|
||
) as message_add, patch.object(
|
||
chain, "edit_message", return_value=True
|
||
) as edit_message, patch(
|
||
"app.chain.message.agent_manager.process_message",
|
||
new_callable=AsyncMock,
|
||
) as process_message, patch(
|
||
"app.chain.message.asyncio.run_coroutine_threadsafe",
|
||
side_effect=lambda coro, _loop: (coro.close(), Mock())[1],
|
||
):
|
||
handled = chain._handle_callback(
|
||
text=f"CALLBACK:agent_interaction:choice:{request.request_id}:1",
|
||
channel=MessageChannel.Telegram,
|
||
source="telegram-test",
|
||
userid="10001",
|
||
username="tester",
|
||
original_message_id=123,
|
||
original_chat_id="456",
|
||
)
|
||
|
||
self.assertTrue(handled)
|
||
edit_message.assert_called_once_with(
|
||
channel=MessageChannel.Telegram,
|
||
source="telegram-test",
|
||
message_id=123,
|
||
chat_id="456",
|
||
title="需要你的选择",
|
||
text="请选择\n\n已选择：电影",
|
||
)
|
||
process_message.assert_called_once()
|
||
kwargs = process_message.call_args.kwargs
|
||
self.assertEqual(kwargs["message"], "我选择电影")
|
||
self.assertEqual(kwargs["session_id"], "session-choice")
|
||
self.assertEqual(kwargs["channel"], MessageChannel.Telegram.value)
|
||
self.assertEqual(kwargs["source"], "telegram-test")
|
||
self.assertNotIn("processing_status", kwargs)
|
||
message_put.assert_called_once()
|
||
message_add.assert_called_once()
|
||
|
||
def test_legacy_agent_choice_callback_still_supported(self):
|
||
chain = MessageChain()
|
||
request = agent_interaction_manager.create_request(
|
||
session_id="session-choice",
|
||
user_id="10001",
|
||
channel=MessageChannel.Telegram.value,
|
||
source="telegram-test",
|
||
username="tester",
|
||
title=None,
|
||
prompt="请选择",
|
||
options=[AgentInteractionOption(label="电影", value="我选择电影")],
|
||
)
|
||
|
||
with patch.object(chain, "_handle_ai_message") as handle_ai_message, patch.object(
|
||
chain.messagehelper, "put"
|
||
), patch.object(chain.messageoper, "add"):
|
||
chain._handle_callback(
|
||
text=f"CALLBACK:agent_choice:{request.request_id}:1",
|
||
channel=MessageChannel.Telegram,
|
||
source="telegram-test",
|
||
userid="10001",
|
||
username="tester",
|
||
)
|
||
|
||
handle_ai_message.assert_called_once()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|