fixed tests and role value

awslabs · Oct 9, 2024 · cf446d9 · cf446d9
1 parent 80f6162
commit cf446d9
Show file tree

Hide file tree

Showing 7 changed files with 14 additions and 14 deletions.
diff --git a/python/src/multi_agent_orchestrator/agents/amazon_bedrock_agent.py b/python/src/multi_agent_orchestrator/agents/amazon_bedrock_agent.py
@@ -77,13 +77,13 @@ async def process_request(
                     Logger.warn("Received a chunk event with no chunk data")
 
             return ConversationMessage(
-                role=ParticipantRole.ASSISTANT,
+                role=ParticipantRole.ASSISTANT.value,
                 content=[{"text": completion}]
             )
 
         except (BotoCoreError, ClientError) as error:
             Logger.error(f"Error processing request: {str(error)}")
             return ConversationMessage(
-                role=ParticipantRole.ASSISTANT,
+                role=ParticipantRole.ASSISTANT.value,
                 content=[{"text": "Sorry, I encountered an error while processing your request."}]
             )
diff --git a/python/src/multi_agent_orchestrator/agents/bedrock_translator_agent.py b/python/src/multi_agent_orchestrator/agents/bedrock_translator_agent.py
@@ -62,13 +62,13 @@ async def process_request(self,
         # Check if input is a number and return it as-is if true
         if input_text.isdigit():
             return ConversationMessage(
-                role=ParticipantRole.ASSISTANT,
+                role=ParticipantRole.ASSISTANT.value,
                 content=[{"text": input_text}]
             )
 
         # Prepare user message
         user_message = ConversationMessage(
-            role=ParticipantRole.USER,
+            role=ParticipantRole.USER.value,
             content=[{"text": f"<userinput>{input_text}</userinput>"}]
         )
 
@@ -121,7 +121,7 @@ async def process_request(self,
 
                         # Return the translated text
                         return ConversationMessage(
-                            role=ParticipantRole.ASSISTANT,
+                            role=ParticipantRole.ASSISTANT.value,
                             content=[{"text": translation}]
                         )
 

diff --git a/python/src/multi_agent_orchestrator/agents/chain_agent.py b/python/src/multi_agent_orchestrator/agents/chain_agent.py
@@ -81,6 +81,6 @@ def is_conversation_message(response: any) -> bool:
 
     def create_default_response(self) -> ConversationMessage:
         return ConversationMessage(
-            role=ParticipantRole.ASSISTANT,
+            role=ParticipantRole.ASSISTANT.value,
             content=[{"text": self.default_output}]
         )
diff --git a/python/src/multi_agent_orchestrator/agents/comprehend_filter_agent.py b/python/src/multi_agent_orchestrator/agents/comprehend_filter_agent.py
@@ -90,7 +90,7 @@ async def process_request(self,
 
             # If no issues, return the original input as a ConversationMessage
             return ConversationMessage(
-                role=ParticipantRole.ASSISTANT,
+                role=ParticipantRole.ASSISTANT.value,
                 content=[{"text": input_text}]
             )
 

diff --git a/python/src/multi_agent_orchestrator/agents/lex_bot_agent.py b/python/src/multi_agent_orchestrator/agents/lex_bot_agent.py
@@ -49,7 +49,7 @@ async def process_request(self, input_text: str, user_id: str, session_id: str,
             )
 
             return ConversationMessage(
-                role=ParticipantRole.ASSISTANT,
+                role=ParticipantRole.ASSISTANT.value,
                 content=[{"text": concatenated_content or "No response from Lex bot."}]
             )
 

diff --git a/python/src/tests/agents/test_amazon_bedrock_agent.py b/python/src/tests/agents/test_amazon_bedrock_agent.py
@@ -43,7 +43,7 @@ async def test_process_request_success(bedrock_agent):
     )
 
     assert isinstance(result, ConversationMessage)
-    assert result.role == ParticipantRole.ASSISTANT
+    assert result.role == ParticipantRole.ASSISTANT.value
     assert result.content == [{"text": "Hello, world!"}]
 
     bedrock_agent.client.invoke_agent.assert_called_once_with(
@@ -68,7 +68,7 @@ async def test_process_request_error(bedrock_agent):
     )
 
     assert isinstance(result, ConversationMessage)
-    assert result.role == ParticipantRole.ASSISTANT
+    assert result.role == ParticipantRole.ASSISTANT.value
     assert result.content == [{"text": "Sorry, I encountered an error while processing your request."}]
 
 @pytest.mark.asyncio
@@ -89,7 +89,7 @@ async def test_process_request_empty_chunk(bedrock_agent):
     )
 
     assert isinstance(result, ConversationMessage)
-    assert result.role == ParticipantRole.ASSISTANT
+    assert result.role == ParticipantRole.ASSISTANT.value
     assert result.content == [{"text": "Hello"}]
 
 @pytest.mark.asyncio
@@ -110,5 +110,5 @@ async def test_process_request_with_additional_params(bedrock_agent):
     )
 
     assert isinstance(result, ConversationMessage)
-    assert result.role == ParticipantRole.ASSISTANT
+    assert result.role == ParticipantRole.ASSISTANT.value
     assert result.content == [{"text": "Response with additional params"}]
diff --git a/python/src/tests/agents/test_lex_bot_agent.py b/python/src/tests/agents/test_lex_bot_agent.py
@@ -50,7 +50,7 @@ async def test_process_request_success(lex_bot_agent, mock_lex_client):
     )
 
     assert isinstance(result, ConversationMessage)
-    assert result.role == ParticipantRole.ASSISTANT
+    assert result.role == ParticipantRole.ASSISTANT.value
     assert result.content == [{"text": "Hello How can I help?"}]
 
     mock_lex_client.recognize_text.assert_called_once_with(
@@ -71,7 +71,7 @@ async def test_process_request_no_response(lex_bot_agent, mock_lex_client):
     )
 
     assert isinstance(result, ConversationMessage)
-    assert result.role == ParticipantRole.ASSISTANT
+    assert result.role == ParticipantRole.ASSISTANT.value
     assert result.content == [{"text": "No response from Lex bot."}]
 
 @pytest.mark.asyncio