diff --git a/tests/unit_tests/agent/__init__.py b/tests/unit_tests/agent/__init__.py new file mode 100644 index 000000000..aaafe5665 --- /dev/null +++ b/tests/unit_tests/agent/__init__.py @@ -0,0 +1 @@ +"""Unit tests for the agent module.""" diff --git a/tests/unit_tests/agent/test_agent_context.py b/tests/unit_tests/agent/test_agent_context.py new file mode 100644 index 000000000..a1ba65d77 --- /dev/null +++ b/tests/unit_tests/agent/test_agent_context.py @@ -0,0 +1,576 @@ +"""Unit tests for agent context classes. + +Tests for AgentContext, AgentGenerateContext, AgentMessage, and AgentReviewInfo +from dbgpt.agent.core.agent module. +""" + +from unittest.mock import MagicMock + +import pytest + +from dbgpt.agent.core.action.base import ActionOutput +from dbgpt.agent.core.agent import ( + Agent, + AgentContext, + AgentGenerateContext, + AgentMessage, + AgentReviewInfo, +) +from dbgpt.core.interface.message import ModelMessageRoleType + + +class TestAgentContext: + """Test suite for AgentContext class.""" + + def test_initialization_with_defaults(self): + """Test AgentContext initialization with default values.""" + context = AgentContext(conv_id="test_conv_123") + + assert context.conv_id == "test_conv_123" + assert context.gpts_app_code is None + assert context.gpts_app_name is None + assert context.language is None + assert context.max_chat_round == 100 + assert context.max_retry_round == 10 + assert context.max_new_tokens == 1024 + assert context.temperature == 0.5 + assert context.allow_format_str_template is False + assert context.verbose is False + assert context.app_link_start is False + assert context.enable_vis_message is True + + def test_initialization_with_custom_values(self): + """Test AgentContext initialization with custom values.""" + context = AgentContext( + conv_id="custom_conv", + gpts_app_code="app_001", + gpts_app_name="TestApp", + language="zh", + max_chat_round=50, + max_retry_round=5, + max_new_tokens=2048, + temperature=0.7, + allow_format_str_template=True, + verbose=True, + app_link_start=True, + enable_vis_message=False, + ) + + assert context.conv_id == "custom_conv" + assert context.gpts_app_code == "app_001" + assert context.gpts_app_name == "TestApp" + assert context.language == "zh" + assert context.max_chat_round == 50 + assert context.max_retry_round == 5 + assert context.max_new_tokens == 2048 + assert context.temperature == 0.7 + assert context.allow_format_str_template is True + assert context.verbose is True + assert context.app_link_start is True + assert context.enable_vis_message is False + + def test_to_dict(self): + """Test AgentContext.to_dict() method.""" + context = AgentContext( + conv_id="test_conv", + gpts_app_code="app_001", + language="en", + max_chat_round=50, + ) + + result = context.to_dict() + + assert isinstance(result, dict) + assert result["conv_id"] == "test_conv" + assert result["gpts_app_code"] == "app_001" + assert result["language"] == "en" + assert result["max_chat_round"] == 50 + assert "max_retry_round" in result + assert "temperature" in result + + +class TestAgentGenerateContext: + """Test suite for AgentGenerateContext class.""" + + def test_initialization_with_minimal_values(self): + """Test AgentGenerateContext with minimal required values.""" + mock_agent = MagicMock(spec=Agent) + mock_message = AgentMessage(content="test message") + + context = AgentGenerateContext(message=mock_message, sender=mock_agent) + + assert context.message == mock_message + assert context.sender == mock_agent + assert context.reviewer is None + assert context.silent is False + assert context.already_failed is False + assert context.last_speaker is None + assert context.already_started is False + assert context.begin_agent is None + assert context.rely_messages == [] + assert context.final is True + assert context.memory is None + assert context.agent_context is None + assert context.llm_client is None + assert context.round_index is None + + def test_initialization_with_full_values(self): + """Test AgentGenerateContext with all values.""" + mock_sender = MagicMock(spec=Agent) + mock_reviewer = MagicMock(spec=Agent) + mock_last_speaker = MagicMock(spec=Agent) + mock_message = AgentMessage(content="test message") + rely_msg1 = AgentMessage(content="rely message 1") + rely_msg2 = AgentMessage(content="rely message 2") + mock_agent_context = AgentContext(conv_id="test_conv") + + context = AgentGenerateContext( + message=mock_message, + sender=mock_sender, + reviewer=mock_reviewer, + silent=True, + already_failed=True, + last_speaker=mock_last_speaker, + already_started=True, + begin_agent="begin_agent_name", + rely_messages=[rely_msg1, rely_msg2], + final=False, + agent_context=mock_agent_context, + round_index=5, + ) + + assert context.message == mock_message + assert context.sender == mock_sender + assert context.reviewer == mock_reviewer + assert context.silent is True + assert context.already_failed is True + assert context.last_speaker == mock_last_speaker + assert context.already_started is True + assert context.begin_agent == "begin_agent_name" + assert len(context.rely_messages) == 2 + assert context.rely_messages[0] == rely_msg1 + assert context.rely_messages[1] == rely_msg2 + assert context.final is False + assert context.agent_context == mock_agent_context + assert context.round_index == 5 + + def test_rely_messages_handling(self): + """Test AgentGenerateContext handling of rely_messages.""" + mock_agent = MagicMock(spec=Agent) + mock_message = AgentMessage(content="main message") + rely_msg1 = AgentMessage(content="rely 1", role=ModelMessageRoleType.HUMAN) + rely_msg2 = AgentMessage(content="rely 2", role=ModelMessageRoleType.AI) + + context = AgentGenerateContext( + message=mock_message, + sender=mock_agent, + rely_messages=[rely_msg1, rely_msg2], + ) + + assert len(context.rely_messages) == 2 + assert context.rely_messages[0].content == "rely 1" + assert context.rely_messages[1].content == "rely 2" + + def test_already_failed_flag(self): + """Test AgentGenerateContext already_failed flag.""" + mock_agent = MagicMock(spec=Agent) + mock_message = AgentMessage(content="test") + + # Test with already_failed=False (default) + context1 = AgentGenerateContext(message=mock_message, sender=mock_agent) + assert context1.already_failed is False + + # Test with already_failed=True + context2 = AgentGenerateContext( + message=mock_message, sender=mock_agent, already_failed=True + ) + assert context2.already_failed is True + + def test_to_dict(self): + """Test AgentGenerateContext.to_dict() method.""" + mock_agent = MagicMock(spec=Agent) + mock_message = AgentMessage(content="test message") + mock_agent_context = AgentContext(conv_id="test_conv") + + context = AgentGenerateContext( + message=mock_message, + sender=mock_agent, + agent_context=mock_agent_context, + round_index=3, + ) + + result = context.to_dict() + + assert isinstance(result, dict) + assert "message" in result + assert "sender" in result + assert "agent_context" in result + assert "round_index" in result + assert result["round_index"] == 3 + + +class TestAgentMessage: + """Test suite for AgentMessage class.""" + + def test_initialization_with_defaults(self): + """Test AgentMessage initialization with default values.""" + message = AgentMessage() + + assert message.content is None + assert message.name is None + assert message.rounds == 0 + assert message.context is None + assert message.action_report is None + assert message.review_info is None + assert message.current_goal is None + assert message.model_name is None + assert message.role is None + assert message.success is True + assert message.resource_info is None + + def test_initialization_with_custom_values(self): + """Test AgentMessage initialization with custom values.""" + action_report = ActionOutput(content="Action result", is_exe_success=True) + review_info = AgentReviewInfo(approve=True, comments="Looks good") + + message = AgentMessage( + content="Test content", + name="TestAgent", + rounds=3, + context={"key": "value"}, + action_report=action_report, + review_info=review_info, + current_goal="Test goal", + model_name="gpt-4", + role=ModelMessageRoleType.AI, + success=False, + resource_info={"resource": "data"}, + ) + + assert message.content == "Test content" + assert message.name == "TestAgent" + assert message.rounds == 3 + assert message.context == {"key": "value"} + assert message.action_report == action_report + assert message.review_info == review_info + assert message.current_goal == "Test goal" + assert message.model_name == "gpt-4" + assert message.role == ModelMessageRoleType.AI + assert message.success is False + assert message.resource_info == {"resource": "data"} + + def test_copy(self): + """Test AgentMessage.copy() method.""" + action_report = ActionOutput(content="Action result") + review_info = AgentReviewInfo(approve=True, comments="Good") + original = AgentMessage( + content="Original content", + name="Agent1", + rounds=2, + context={"key": "value"}, + action_report=action_report, + review_info=review_info, + current_goal="Goal", + model_name="gpt-3.5", + role=ModelMessageRoleType.HUMAN, + success=True, + resource_info={"res": "info"}, + ) + + copied = original.copy() + + # Check that values are copied + assert copied.content == original.content + assert copied.name == original.name + assert copied.rounds == original.rounds + assert copied.current_goal == original.current_goal + assert copied.model_name == original.model_name + assert copied.role == original.role + assert copied.success == original.success + assert copied.action_report == original.action_report + assert copied.resource_info == original.resource_info + + # Check that context is deep copied (for dict) + assert copied.context == original.context + assert copied.context is not original.context # Different objects + + # Check review_info is copied + assert copied.review_info.approve == original.review_info.approve + assert copied.review_info.comments == original.review_info.comments + + def test_copy_with_string_context(self): + """Test AgentMessage.copy() with string context.""" + original = AgentMessage(content="Test", context="string context") + copied = original.copy() + + assert copied.context == original.context + # String context is immutable, so it's the same object + assert copied.context is original.context + + def test_to_dict(self): + """Test AgentMessage.to_dict() method.""" + action_report = ActionOutput( + content="Action result", is_exe_success=True, view="chart" + ) + message = AgentMessage( + content="Test content", + name="Agent1", + rounds=2, + action_report=action_report, + role=ModelMessageRoleType.AI, + ) + + result = message.to_dict() + + assert isinstance(result, dict) + assert result["content"] == "Test content" + assert result["name"] == "Agent1" + assert result["rounds"] == 2 + assert result["role"] == ModelMessageRoleType.AI + assert "action_report" in result + assert isinstance(result["action_report"], dict) + assert result["action_report"]["content"] == "Action result" + assert result["action_report"]["is_exe_success"] is True + + def test_to_dict_without_action_report(self): + """Test AgentMessage.to_dict() without action_report.""" + message = AgentMessage(content="Simple message", rounds=1) + result = message.to_dict() + + assert result["content"] == "Simple message" + assert result["rounds"] == 1 + assert result["action_report"] is None + + def test_to_llm_message(self): + """Test AgentMessage.to_llm_message() method.""" + message = AgentMessage( + content="Test content", + context={"key": "value"}, + role=ModelMessageRoleType.HUMAN, + name="Agent1", + rounds=5, + ) + + result = message.to_llm_message() + + assert isinstance(result, dict) + assert result["content"] == "Test content" + assert result["context"] == {"key": "value"} + assert result["role"] == ModelMessageRoleType.HUMAN + # to_llm_message only returns content, context, and role + assert "name" not in result + assert "rounds" not in result + + def test_from_llm_message(self): + """Test AgentMessage.from_llm_message() factory method.""" + llm_msg = { + "content": "Test content", + "context": {"key": "value"}, + "role": ModelMessageRoleType.AI, + "rounds": 3, + } + + message = AgentMessage.from_llm_message(llm_msg) + + assert message.content == "Test content" + assert message.context == {"key": "value"} + assert message.role == ModelMessageRoleType.AI + assert message.rounds == 3 + + def test_from_llm_message_with_missing_fields(self): + """Test AgentMessage.from_llm_message() with missing fields.""" + llm_msg = {"content": "Just content"} + + message = AgentMessage.from_llm_message(llm_msg) + + assert message.content == "Just content" + assert message.context is None + assert message.role is None + assert message.rounds == 0 # Default value + + def test_from_messages(self): + """Test AgentMessage.from_messages() factory method.""" + messages_data = [ + { + "content": "Message 1", + "role": ModelMessageRoleType.HUMAN, + "rounds": 1, + "name": "User", + }, + { + "content": "Message 2", + "role": ModelMessageRoleType.AI, + "rounds": 2, + "name": "Agent", + "success": False, + }, + {"content": "Message 3"}, + ] + + messages = AgentMessage.from_messages(messages_data) + + assert len(messages) == 3 + assert messages[0].content == "Message 1" + assert messages[0].role == ModelMessageRoleType.HUMAN + assert messages[0].rounds == 1 + assert messages[0].name == "User" + assert messages[1].content == "Message 2" + assert messages[1].role == ModelMessageRoleType.AI + assert messages[1].success is False + assert messages[2].content == "Message 3" + + def test_from_messages_filters_unknown_fields(self): + """Test that from_messages filters out unknown fields.""" + messages_data = [ + { + "content": "Message 1", + "unknown_field": "should be ignored", + "another_unknown": 123, + } + ] + + messages = AgentMessage.from_messages(messages_data) + + assert len(messages) == 1 + assert messages[0].content == "Message 1" + # Unknown fields should not cause errors + + def test_action_report_serialization(self): + """Test action_report is properly serialized in to_dict().""" + action_report = ActionOutput( + content="Action completed", + is_exe_success=True, + thoughts="I thought about it", + observations="I observed this", + next_speakers=["Agent2", "Agent3"], + ) + message = AgentMessage(content="Test", action_report=action_report) + + result = message.to_dict() + + assert isinstance(result["action_report"], dict) + assert result["action_report"]["content"] == "Action completed" + assert result["action_report"]["is_exe_success"] is True + assert result["action_report"]["thoughts"] == "I thought about it" + assert result["action_report"]["observations"] == "I observed this" + assert result["action_report"]["next_speakers"] == ["Agent2", "Agent3"] + + def test_context_deep_copy_dict(self): + """Test that dict context is deep copied.""" + original_context = {"key": "value", "nested": {"inner": "data"}} + original = AgentMessage(content="Test", context=original_context) + + copied = original.copy() + + # Modify the copied context + copied.context["key"] = "modified" + copied.context["nested"]["inner"] = "modified_data" + + # Original should not be affected + assert original.context["key"] == "value" + assert original.context["nested"]["inner"] == "data" + + def test_get_dict_context(self): + """Test AgentMessage.get_dict_context() method.""" + # Test with dict context + message1 = AgentMessage(content="Test", context={"key": "value"}) + result1 = message1.get_dict_context() + assert result1 == {"key": "value"} + + # Test with string context + message2 = AgentMessage(content="Test", context="string context") + result2 = message2.get_dict_context() + assert result2 == {} + + # Test with None context + message3 = AgentMessage(content="Test", context=None) + result3 = message3.get_dict_context() + assert result3 == {} + + +class TestAgentReviewInfo: + """Test suite for AgentReviewInfo class.""" + + def test_initialization_with_defaults(self): + """Test AgentReviewInfo initialization with default values.""" + review = AgentReviewInfo() + + assert review.approve is False + assert review.comments is None + + def test_initialization_with_custom_values(self): + """Test AgentReviewInfo initialization with custom values.""" + review = AgentReviewInfo(approve=True, comments="Excellent work!") + + assert review.approve is True + assert review.comments == "Excellent work!" + + def test_copy(self): + """Test AgentReviewInfo.copy() method.""" + original = AgentReviewInfo(approve=True, comments="Good job") + copied = original.copy() + + assert copied.approve == original.approve + assert copied.comments == original.comments + + # Modify copied + copied.approve = False + copied.comments = "Needs improvement" + + # Original should not be affected + assert original.approve is True + assert original.comments == "Good job" + + def test_to_dict(self): + """Test AgentReviewInfo.to_dict() method.""" + review = AgentReviewInfo(approve=True, comments="Looks good") + result = review.to_dict() + + assert isinstance(result, dict) + assert result["approve"] is True + assert result["comments"] == "Looks good" + + def test_to_dict_with_defaults(self): + """Test AgentReviewInfo.to_dict() with default values.""" + review = AgentReviewInfo() + result = review.to_dict() + + assert isinstance(result, dict) + assert result["approve"] is False + assert result["comments"] is None + + +@pytest.mark.parametrize( + "conv_id,max_chat_round,temperature", + [ + ("conv_1", 50, 0.5), + ("conv_2", 100, 0.7), + ("conv_3", 200, 0.3), + ], +) +def test_agent_context_parametrized(conv_id, max_chat_round, temperature): + """Parametrized test for AgentContext with different values.""" + context = AgentContext( + conv_id=conv_id, max_chat_round=max_chat_round, temperature=temperature + ) + + assert context.conv_id == conv_id + assert context.max_chat_round == max_chat_round + assert context.temperature == temperature + + +@pytest.mark.parametrize( + "content,rounds,success", + [ + ("Message 1", 1, True), + ("Message 2", 5, False), + ("Message 3", 10, True), + ], +) +def test_agent_message_parametrized(content, rounds, success): + """Parametrized test for AgentMessage with different values.""" + message = AgentMessage(content=content, rounds=rounds, success=success) + + assert message.content == content + assert message.rounds == rounds + assert message.success == success diff --git a/tests/unit_tests/agent/test_awel_agent_operator.py b/tests/unit_tests/agent/test_awel_agent_operator.py new file mode 100644 index 000000000..18ca9d5b3 --- /dev/null +++ b/tests/unit_tests/agent/test_awel_agent_operator.py @@ -0,0 +1,746 @@ +"""Unit tests for AWEL agent operators. + +Tests for AWELAgentOperator, WrappedAgentOperator, and AgentBranchOperator +from dbgpt.agent.core.plan.awel.agent_operator module. +""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from dbgpt.agent.core.action.base import ActionOutput +from dbgpt.agent.core.agent import ( + Agent, + AgentContext, + AgentGenerateContext, + AgentMessage, +) +from dbgpt.agent.core.base_agent import ConversableAgent +from dbgpt.agent.core.memory.agent_memory import AgentMemory +from dbgpt.agent.core.plan.awel.agent_operator import ( + AgentBranchOperator, + AWELAgentOperator, + WrappedAgentOperator, +) +from dbgpt.agent.core.plan.awel.agent_operator_resource import AWELAgent +from dbgpt.core.interface.message import ModelMessageRoleType + + +@pytest.fixture +def mock_agent(): + """Create a mock Agent instance.""" + agent = MagicMock(spec=Agent) + agent.name = "TestAgent" + agent.role = "test_role" + agent.send = AsyncMock() + agent.generate_reply = AsyncMock() + return agent + + +@pytest.fixture +def mock_conversable_agent(): + """Create a mock ConversableAgent instance.""" + agent = MagicMock(spec=ConversableAgent) + agent.name = "TestConversableAgent" + agent.role = "conversable_role" + agent.fixed_subgoal = None + agent.send = AsyncMock() + agent.generate_reply = AsyncMock() + return agent + + +@pytest.fixture +def mock_agent_memory(): + """Create a mock AgentMemory instance.""" + memory = MagicMock(spec=AgentMemory) + memory.structure_clone = MagicMock(return_value=memory) + return memory + + +@pytest.fixture +def sample_agent_context(): + """Create a sample AgentContext.""" + return AgentContext(conv_id="test_conv", max_chat_round=10, temperature=0.5) + + +@pytest.fixture +def sample_message(): + """Create a sample AgentMessage.""" + return AgentMessage(content="Test message content", rounds=0, success=True) + + +@pytest.fixture +def sample_generate_context(mock_agent, sample_message, sample_agent_context): + """Create a sample AgentGenerateContext.""" + return AgentGenerateContext( + message=sample_message, + sender=mock_agent, + reviewer=None, + silent=False, + already_failed=False, + rely_messages=[], + agent_context=sample_agent_context, + llm_client=None, + ) + + +@pytest.fixture +def mock_awel_agent(): + """Create a mock AWELAgent instance.""" + awel_agent = MagicMock(spec=AWELAgent) + awel_agent.agent_profile = "TestAgentProfile" + awel_agent.role_name = "TestRole" + awel_agent.llm_config = None + awel_agent.fixed_subgoal = None + awel_agent.agent_prompt = None + awel_agent.resources = [] + return awel_agent + + +class TestWrappedAgentOperator: + """Test suite for WrappedAgentOperator class.""" + + @pytest.mark.asyncio + async def test_map_basic_execution(self, mock_agent, sample_message): + """Test WrappedAgentOperator.map() with basic execution flow.""" + # Setup + operator = WrappedAgentOperator(agent=mock_agent) + rely_msg = AgentMessage( + content="Previous message", role=ModelMessageRoleType.HUMAN + ) + input_context = AgentGenerateContext( + message=sample_message, + sender=mock_agent, + rely_messages=[rely_msg], + ) + + # Mock the agent's generate_reply to return a successful message + reply_message = AgentMessage( + content="Agent reply", success=True, role=ModelMessageRoleType.AI + ) + mock_agent.generate_reply.return_value = reply_message + + # Execute + result = await operator.map(input_context) + + # Verify + assert isinstance(result, AgentGenerateContext) + assert result.sender == mock_agent # Sender is now the agent + assert len(result.rely_messages) == 2 # Human message + AI reply + assert result.rely_messages[0].role == ModelMessageRoleType.HUMAN + assert result.rely_messages[1].role == ModelMessageRoleType.AI + assert result.already_started is True + assert result.last_speaker == mock_agent + + # Verify agent.send was called + mock_agent.send.assert_called_once() + + # Verify agent.generate_reply was called + mock_agent.generate_reply.assert_called_once() + + @pytest.mark.asyncio + async def test_map_empty_message_raises_exception(self, mock_agent): + """Test WrappedAgentOperator.map() raises ValueError when message is empty.""" + operator = WrappedAgentOperator(agent=mock_agent) + input_context = AgentGenerateContext(message=None, sender=mock_agent) + + with pytest.raises(ValueError, match="The message is empty"): + await operator.map(input_context) + + @pytest.mark.asyncio + async def test_map_empty_sender_raises_exception(self, mock_agent, sample_message): + """Test WrappedAgentOperator.map() raises ValueError when sender is empty.""" + operator = WrappedAgentOperator(agent=mock_agent) + input_context = AgentGenerateContext(message=sample_message, sender=None) + + with pytest.raises(ValueError, match="The sender is empty"): + await operator.map(input_context) + + @pytest.mark.asyncio + async def test_map_unsuccessful_reply_raises_exception( + self, mock_agent, sample_message + ): + """Test WrappedAgentOperator.map() raises exception on unsuccessful reply.""" + operator = WrappedAgentOperator(agent=mock_agent) + input_context = AgentGenerateContext(message=sample_message, sender=mock_agent) + + # Mock the agent's generate_reply to return an unsuccessful message + reply_message = AgentMessage( + content="Failed to complete task", + success=False, + role=ModelMessageRoleType.AI, + ) + mock_agent.generate_reply.return_value = reply_message + + with pytest.raises( + ValueError, + match="The task failed at step test_role and the attempt to repair it", + ): + await operator.map(input_context) + + @pytest.mark.asyncio + async def test_map_with_memory(self, mock_agent, sample_message, mock_agent_memory): + """Test WrappedAgentOperator.map() with memory handling.""" + operator = WrappedAgentOperator(agent=mock_agent) + input_context = AgentGenerateContext( + message=sample_message, sender=mock_agent, memory=mock_agent_memory + ) + + reply_message = AgentMessage(content="Reply", success=True) + mock_agent.generate_reply.return_value = reply_message + + result = await operator.map(input_context) + + # Verify memory is cloned + mock_agent_memory.structure_clone.assert_called_once() + assert result.memory is not None + + @pytest.mark.asyncio + async def test_map_current_goal_construction(self, mock_agent, sample_message): + """Test WrappedAgentOperator.map() constructs current_goal correctly.""" + operator = WrappedAgentOperator(agent=mock_agent) + input_context = AgentGenerateContext(message=sample_message, sender=mock_agent) + + reply_message = AgentMessage(content="Reply", success=True) + mock_agent.generate_reply.return_value = reply_message + + await operator.map(input_context) + + # Verify the message passed to generate_reply has current_goal set + call_args = mock_agent.generate_reply.call_args + received_message = call_args.kwargs["received_message"] + assert received_message.current_goal is not None + assert received_message.current_goal.startswith("[TestAgent]:") + + +class TestAWELAgentOperator: + """Test suite for AWELAgentOperator class.""" + + @pytest.mark.asyncio + async def test_map_returns_early_if_already_failed( + self, mock_awel_agent, sample_generate_context + ): + """Test AWELAgentOperator.map() returns early when already_failed is True.""" + operator = AWELAgentOperator(awel_agent=mock_awel_agent) + sample_generate_context.already_failed = True + + result = await operator.map(sample_generate_context) + + assert result == sample_generate_context + assert result.already_failed is True + + @pytest.mark.asyncio + async def test_map_empty_message_raises_exception( + self, mock_awel_agent, mock_agent, sample_agent_context + ): + """Test AWELAgentOperator.map() raises ValueError when message is empty.""" + operator = AWELAgentOperator(awel_agent=mock_awel_agent) + input_context = AgentGenerateContext( + message=None, sender=mock_agent, agent_context=sample_agent_context + ) + + with pytest.raises(ValueError, match="The message is empty"): + await operator.map(input_context) + + @pytest.mark.asyncio + async def test_map_empty_sender_raises_exception( + self, mock_awel_agent, sample_message, sample_agent_context + ): + """Test AWELAgentOperator.map() raises ValueError when sender is empty.""" + operator = AWELAgentOperator(awel_agent=mock_awel_agent) + input_context = AgentGenerateContext( + message=sample_message, sender=None, agent_context=sample_agent_context + ) + + # Need to mock get_agent to avoid issues + with patch.object( + operator, "get_agent", new_callable=AsyncMock + ) as mock_get_agent: + mock_get_agent.return_value = MagicMock(spec=ConversableAgent) + + with pytest.raises(ValueError, match="The sender is empty"): + await operator.map(input_context) + + @pytest.mark.asyncio + async def test_map_normal_execution( + self, mock_awel_agent, mock_conversable_agent, sample_generate_context + ): + """Test AWELAgentOperator.map() normal execution flow.""" + operator = AWELAgentOperator(awel_agent=mock_awel_agent) + + # Mock get_agent to return our conversable agent + with patch.object( + operator, "get_agent", new_callable=AsyncMock + ) as mock_get_agent: + mock_get_agent.return_value = mock_conversable_agent + + # Mock agent's generate_reply + reply_message = AgentMessage( + content="Agent response", success=True, role=ModelMessageRoleType.AI + ) + mock_conversable_agent.generate_reply.return_value = reply_message + + result = await operator.map(sample_generate_context) + + # Verify + assert isinstance(result, AgentGenerateContext) + assert result.sender == mock_conversable_agent + assert len(result.rely_messages) == 2 # HUMAN + AI messages + assert result.rely_messages[0].role == ModelMessageRoleType.HUMAN + assert result.rely_messages[1].role == ModelMessageRoleType.AI + assert result.already_started is True + assert result.already_failed is False + assert result.last_speaker == mock_conversable_agent + + @pytest.mark.asyncio + async def test_map_with_fixed_subgoal( + self, mock_awel_agent, mock_conversable_agent, sample_generate_context + ): + """Test AWELAgentOperator.map() with fixed_subgoal.""" + mock_conversable_agent.fixed_subgoal = "Fixed goal for testing" + operator = AWELAgentOperator(awel_agent=mock_awel_agent) + + with patch.object( + operator, "get_agent", new_callable=AsyncMock + ) as mock_get_agent: + mock_get_agent.return_value = mock_conversable_agent + + reply_message = AgentMessage(content="Response", success=True) + mock_conversable_agent.generate_reply.return_value = reply_message + + await operator.map(sample_generate_context) + + # Verify that fixed_subgoal overrides message content + call_args = mock_conversable_agent.generate_reply.call_args + received_message = call_args.kwargs["received_message"] + assert received_message.content == "Fixed goal for testing" + assert "Fixed goal for testing" in received_message.current_goal + + @pytest.mark.asyncio + async def test_map_without_fixed_subgoal( + self, mock_awel_agent, mock_conversable_agent, sample_generate_context + ): + """Test AWELAgentOperator.map() without fixed_subgoal.""" + mock_conversable_agent.fixed_subgoal = None + operator = AWELAgentOperator(awel_agent=mock_awel_agent) + + with patch.object( + operator, "get_agent", new_callable=AsyncMock + ) as mock_get_agent: + mock_get_agent.return_value = mock_conversable_agent + + reply_message = AgentMessage(content="Response", success=True) + mock_conversable_agent.generate_reply.return_value = reply_message + + await operator.map(sample_generate_context) + + # Verify that message content is used + call_args = mock_conversable_agent.generate_reply.call_args + received_message = call_args.kwargs["received_message"] + assert received_message.content == "Test message content" + + @pytest.mark.asyncio + async def test_map_failure_sets_already_failed( + self, mock_awel_agent, mock_conversable_agent, sample_generate_context + ): + """Test AWELAgentOperator.map() sets already_failed on unsuccessful reply.""" + operator = AWELAgentOperator(awel_agent=mock_awel_agent) + + with patch.object( + operator, "get_agent", new_callable=AsyncMock + ) as mock_get_agent: + mock_get_agent.return_value = mock_conversable_agent + + # Mock agent's generate_reply to return unsuccessful message + reply_message = AgentMessage( + content="Failed", success=False, role=ModelMessageRoleType.AI + ) + mock_conversable_agent.generate_reply.return_value = reply_message + + result = await operator.map(sample_generate_context) + + # Verify already_failed is set to True + assert result.already_failed is True + + @pytest.mark.asyncio + async def test_map_increments_rounds( + self, mock_awel_agent, mock_conversable_agent, sample_generate_context + ): + """Test AWELAgentOperator.map() increments message rounds.""" + sample_generate_context.message.rounds = 5 + operator = AWELAgentOperator(awel_agent=mock_awel_agent) + + with patch.object( + operator, "get_agent", new_callable=AsyncMock + ) as mock_get_agent: + mock_get_agent.return_value = mock_conversable_agent + + reply_message = AgentMessage(content="Response", success=True) + mock_conversable_agent.generate_reply.return_value = reply_message + + await operator.map(sample_generate_context) + + # Verify rounds are incremented + call_args = mock_conversable_agent.generate_reply.call_args + received_message = call_args.kwargs["received_message"] + assert received_message.rounds == 6 + + @pytest.mark.asyncio + async def test_map_with_begin_agent_not_started( + self, mock_awel_agent, mock_conversable_agent, sample_generate_context + ): + """Test AWELAgentOperator.map() when begin_agent is set but not started.""" + sample_generate_context.begin_agent = "other_agent" + sample_generate_context.already_started = False + mock_conversable_agent.role = "test_role" + operator = AWELAgentOperator(awel_agent=mock_awel_agent) + + with patch.object( + operator, "get_agent", new_callable=AsyncMock + ) as mock_get_agent: + mock_get_agent.return_value = mock_conversable_agent + + result = await operator.map(sample_generate_context) + + # Should return early without executing + assert result == sample_generate_context + mock_conversable_agent.generate_reply.assert_not_called() + + @pytest.mark.asyncio + async def test_map_with_begin_agent_matching( + self, mock_awel_agent, mock_conversable_agent, sample_generate_context + ): + """Test AWELAgentOperator.map() when begin_agent matches current agent.""" + sample_generate_context.begin_agent = "conversable_role" + sample_generate_context.already_started = False + mock_conversable_agent.role = "conversable_role" + operator = AWELAgentOperator(awel_agent=mock_awel_agent) + + with patch.object( + operator, "get_agent", new_callable=AsyncMock + ) as mock_get_agent: + mock_get_agent.return_value = mock_conversable_agent + + reply_message = AgentMessage(content="Response", success=True) + mock_conversable_agent.generate_reply.return_value = reply_message + + result = await operator.map(sample_generate_context) + + # Should execute and set is_retry_chat=True + assert result.sender == mock_conversable_agent + mock_conversable_agent.generate_reply.assert_called_once() + call_args = mock_conversable_agent.generate_reply.call_args + assert call_args.kwargs.get("is_retry_chat") is True + + @pytest.mark.asyncio + async def test_map_with_memory_clone( + self, + mock_awel_agent, + mock_conversable_agent, + sample_generate_context, + mock_agent_memory, + ): + """Test AWELAgentOperator.map() clones memory.""" + sample_generate_context.memory = mock_agent_memory + operator = AWELAgentOperator(awel_agent=mock_awel_agent) + + with patch.object( + operator, "get_agent", new_callable=AsyncMock + ) as mock_get_agent: + mock_get_agent.return_value = mock_conversable_agent + + reply_message = AgentMessage(content="Response", success=True) + mock_conversable_agent.generate_reply.return_value = reply_message + + result = await operator.map(sample_generate_context) + + # Verify memory is cloned + mock_agent_memory.structure_clone.assert_called_once() + assert result.memory is not None + + @pytest.mark.asyncio + async def test_map_message_role_conversion( + self, mock_awel_agent, mock_conversable_agent, sample_generate_context + ): + """Test AWELAgentOperator.map() converts message roles correctly.""" + operator = AWELAgentOperator(awel_agent=mock_awel_agent) + + with patch.object( + operator, "get_agent", new_callable=AsyncMock + ) as mock_get_agent: + mock_get_agent.return_value = mock_conversable_agent + + reply_message = AgentMessage(content="Response", success=True) + mock_conversable_agent.generate_reply.return_value = reply_message + + result = await operator.map(sample_generate_context) + + # Check that roles are set correctly + assert result.rely_messages[0].role == ModelMessageRoleType.HUMAN + assert result.rely_messages[1].role == ModelMessageRoleType.AI + + @pytest.mark.asyncio + async def test_get_agent_builds_agent( + self, mock_awel_agent, sample_generate_context + ): + """Test AWELAgentOperator.get_agent() builds agent correctly.""" + operator = AWELAgentOperator(awel_agent=mock_awel_agent) + operator.llm_client = MagicMock() + + # Mock the agent class and its builder pattern + mock_agent_class = MagicMock() + mock_agent_instance = MagicMock(spec=ConversableAgent) + mock_agent_instance.role = "test_role" + + # Setup the builder chain + mock_builder = MagicMock() + mock_builder.bind.return_value = mock_builder + mock_builder.build = AsyncMock(return_value=mock_agent_instance) + mock_agent_class.return_value = mock_builder + + with patch( + "dbgpt.agent.core.plan.awel.agent_operator.get_agent_manager" + ) as mock_get_manager: + mock_manager = MagicMock() + mock_manager.get_by_name.return_value = mock_agent_class + mock_get_manager.return_value = mock_manager + + with patch( + "dbgpt.agent.core.plan.awel.agent_operator.get_resource_manager" + ) as mock_get_resource_manager: + mock_resource_manager = MagicMock() + mock_resource_manager.build_resource.return_value = None + mock_get_resource_manager.return_value = mock_resource_manager + + result = await operator.get_agent(sample_generate_context) + + assert result == mock_agent_instance + mock_manager.get_by_name.assert_called_once_with("TestAgentProfile") + + +class TestAgentBranchOperator: + """Test suite for AgentBranchOperator class.""" + + @pytest.mark.asyncio + async def test_branches_returns_branch_map(self): + """Test AgentBranchOperator.branches() returns correct branch map.""" + operator = AgentBranchOperator() + + # Create mock downstream AWELAgentOperator nodes + mock_awel_agent1 = MagicMock(spec=AWELAgent) + mock_awel_agent1.agent_profile = "Agent1" + + mock_awel_agent2 = MagicMock(spec=AWELAgent) + mock_awel_agent2.agent_profile = "Agent2" + + mock_operator1 = MagicMock(spec=AWELAgentOperator) + mock_operator1.awel_agent = mock_awel_agent1 + mock_operator1.node_name = "node1" + + mock_operator2 = MagicMock(spec=AWELAgentOperator) + mock_operator2.awel_agent = mock_awel_agent2 + mock_operator2.node_name = "node2" + + # Mock downstream attribute + operator.downstream = [mock_operator1, mock_operator2] + + result = await operator.branches() + + assert isinstance(result, dict) + assert len(result) == 2 + + # The keys should be functions and values should be node names + list(result.keys()) + branch_targets = list(result.values()) + assert branch_targets == ["node1", "node2"] + + @pytest.mark.asyncio + async def test_branch_function_checks_next_speakers(self): + """Test that branch functions check action_report.next_speakers.""" + operator = AgentBranchOperator() + + mock_awel_agent = MagicMock(spec=AWELAgent) + mock_awel_agent.agent_profile = "TargetAgent" + + mock_operator = MagicMock(spec=AWELAgentOperator) + mock_operator.awel_agent = mock_awel_agent + mock_operator.node_name = "target_node" + + operator.downstream = [mock_operator] + + branch_map = await operator.branches() + branch_func = list(branch_map.keys())[0] + + # Test with matching next_speakers + action_output = ActionOutput( + content="Test", next_speakers=["TargetAgent", "OtherAgent"] + ) + last_message = AgentMessage(content="Test", action_report=action_output) + context = AgentGenerateContext( + message=AgentMessage(content="Main"), + sender=MagicMock(spec=Agent), + rely_messages=[last_message], + ) + + result = branch_func(context) + assert result is True + + # Test with non-matching next_speakers + action_output2 = ActionOutput(content="Test", next_speakers=["OtherAgent"]) + last_message2 = AgentMessage(content="Test", action_report=action_output2) + context2 = AgentGenerateContext( + message=AgentMessage(content="Main"), + sender=MagicMock(spec=Agent), + rely_messages=[last_message2], + ) + + result2 = branch_func(context2) + assert result2 is False + + @pytest.mark.asyncio + async def test_branch_function_returns_false_without_action_report(self): + """Test branch function returns False when action_report is None.""" + operator = AgentBranchOperator() + + mock_awel_agent = MagicMock(spec=AWELAgent) + mock_awel_agent.agent_profile = "TargetAgent" + + mock_operator = MagicMock(spec=AWELAgentOperator) + mock_operator.awel_agent = mock_awel_agent + mock_operator.node_name = "target_node" + + operator.downstream = [mock_operator] + + branch_map = await operator.branches() + branch_func = list(branch_map.keys())[0] + + # Test without action_report + last_message = AgentMessage(content="Test", action_report=None) + context = AgentGenerateContext( + message=AgentMessage(content="Main"), + sender=MagicMock(spec=Agent), + rely_messages=[last_message], + ) + + result = branch_func(context) + assert result is False + + @pytest.mark.asyncio + async def test_branch_function_returns_false_without_next_speakers(self): + """Test branch function returns False when next_speakers is None.""" + operator = AgentBranchOperator() + + mock_awel_agent = MagicMock(spec=AWELAgent) + mock_awel_agent.agent_profile = "TargetAgent" + + mock_operator = MagicMock(spec=AWELAgentOperator) + mock_operator.awel_agent = mock_awel_agent + mock_operator.node_name = "target_node" + + operator.downstream = [mock_operator] + + branch_map = await operator.branches() + branch_func = list(branch_map.keys())[0] + + # Test with action_report but no next_speakers + action_output = ActionOutput(content="Test", next_speakers=None) + last_message = AgentMessage(content="Test", action_report=action_output) + context = AgentGenerateContext( + message=AgentMessage(content="Main"), + sender=MagicMock(spec=Agent), + rely_messages=[last_message], + ) + + result = branch_func(context) + assert result is False + + @pytest.mark.asyncio + async def test_branches_filters_non_awel_operators(self): + """Test branches() filters out non-AWELAgentOperator nodes.""" + operator = AgentBranchOperator() + + mock_awel_agent = MagicMock(spec=AWELAgent) + mock_awel_agent.agent_profile = "Agent1" + + mock_operator1 = MagicMock(spec=AWELAgentOperator) + mock_operator1.awel_agent = mock_awel_agent + mock_operator1.node_name = "node1" + + # Add a non-AWELAgentOperator node + mock_other_operator = MagicMock() + + operator.downstream = [mock_operator1, mock_other_operator] + + result = await operator.branches() + + # Should only include the AWELAgentOperator + assert len(result) == 1 + + +@pytest.mark.parametrize( + "already_failed,should_execute", + [ + (True, False), # Should return early + (False, True), # Should execute + ], +) +@pytest.mark.asyncio +async def test_awel_operator_already_failed_parametrized( + already_failed, should_execute, mock_awel_agent, mock_conversable_agent +): + """Parametrized test for AWELAgentOperator with already_failed flag.""" + operator = AWELAgentOperator(awel_agent=mock_awel_agent) + message = AgentMessage(content="Test", rounds=0) + mock_sender = MagicMock(spec=Agent) + + input_context = AgentGenerateContext( + message=message, + sender=mock_sender, + already_failed=already_failed, + agent_context=AgentContext(conv_id="test"), + ) + + if should_execute: + with patch.object( + operator, "get_agent", new_callable=AsyncMock + ) as mock_get_agent: + mock_get_agent.return_value = mock_conversable_agent + reply_message = AgentMessage(content="Response", success=True) + mock_conversable_agent.generate_reply.return_value = reply_message + + result = await operator.map(input_context) + assert result.sender == mock_conversable_agent + else: + result = await operator.map(input_context) + assert result == input_context + + +@pytest.mark.parametrize( + "success,expected_already_failed", + [ + (True, False), + (False, True), + ], +) +@pytest.mark.asyncio +async def test_awel_operator_success_flag_parametrized( + success, expected_already_failed, mock_awel_agent, mock_conversable_agent +): + """Parametrized test for AWELAgentOperator success handling.""" + operator = AWELAgentOperator(awel_agent=mock_awel_agent) + message = AgentMessage(content="Test", rounds=0) + mock_sender = MagicMock(spec=Agent) + + input_context = AgentGenerateContext( + message=message, + sender=mock_sender, + agent_context=AgentContext(conv_id="test"), + ) + + with patch.object(operator, "get_agent", new_callable=AsyncMock) as mock_get_agent: + mock_get_agent.return_value = mock_conversable_agent + reply_message = AgentMessage(content="Response", success=success) + mock_conversable_agent.generate_reply.return_value = reply_message + + result = await operator.map(input_context) + assert result.already_failed == expected_already_failed