Seperate UserAgent into its Agent. New test for UserAgent

agents/buddy.py

@@ -4,35 +4,23 @@
 """
 
 
-import threading
-from autogen import AssistantAgent, UserProxyAgent
+from autogen import AssistantAgent
 from settings import (
     GPT3_5_TURBO_0613,
     CONFIG_LIST,
 )
-
-
-class UserAgent(UserProxyAgent):
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.current_message = None
-        self.message_event = threading.Event()
-
-    def set_current_message(self, message: str):
-        """Sets the current message from the user and signals that a message is available."""
-        self.current_message = message
-        self.message_event.set()
-
-    def get_human_input(self, prompt: str) -> str:
-        """Blocks until a message is available, then returns it."""
-        self.message_event.wait()  # Block until a message is available
-        message = self.current_message
-        self.message_event.clear()  # Reset the event for the next message
-        return message
+from .user_agent import UserAgent
 
 
 class Buddy:
-    """A class representing a chatbot for chatting with the assistant."""
+    """
+    A class representing a chatbot for chatting with the assistant.
+
+    Attributes:
+        conversation_context (list): A list of dictionaries representing the conversation context.
+        assistant (AssistantAgent): An instance of the AssistantAgent class.
+        user_proxy (UserAgent): An instance of the UserAgent class.
+    """
 
     def __init__(self):
         self.conversation_context = []
@@ -54,7 +42,15 @@ def __init__(self):
         )
 
     def start(self, message: str):
-        """Initiates a chat between the user and the assistant."""
+        """
+        Initiates a chat between the user and the assistant.
+
+        Args:
+            message (str): The message sent by the user.
+
+        Returns:
+            str: The reply sent by the assistant.
+        """
         # Store the user's message to the conversation context
         self.conversation_context.append({"role": "user", "content": message})
 

agents/user_agent.py

@@ -0,0 +1,54 @@
+"""
+This module defines the UserAgent class, which represents a user agent that can send messages
+ to the system.
+"""
+
+import threading
+from autogen import UserProxyAgent
+
+
+class UserAgent(UserProxyAgent):
+    """A class representing a user agent that can send messages to the system.
+
+    Attributes:
+      current_message (str): The current message from the user.
+      message_event (threading.Event): An event that signals when a message is available.
+
+    Methods:
+      set_current_message(message: str) -> None: Sets the current message from the user and
+                                                  signals that a message is available.
+      get_human_input(prompt: str) -> str: Blocks until a message is available, then returns it.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.current_message = None
+        self.message_event = threading.Event()
+
+    def set_current_message(self, message: str) -> None:
+        """
+        Sets the current message from the user and signals that a message is available.
+
+        Args:
+          message (str): The message to set as the current message.
+
+        Returns:
+          None
+        """
+        self.current_message = message
+        self.message_event.set()
+
+    def get_human_input(self, prompt: str) -> str:
+        """
+        Blocks until a message is available, then returns it.
+
+        Args:
+          prompt (str): The prompt to display to the user.
+
+        Returns:
+          str: The message entered by the user.
+        """
+        self.message_event.wait()  # Block until a message is available
+        message = self.current_message
+        self.message_event.clear()  # Reset the event for the next message
+        return message

requirements.txt

@@ -60,6 +60,7 @@ pyautogen==0.1.10
 pydeck==0.8.1b0
 Pygments==2.16.1
 pytest==7.4.2
+pytest-mock==3.11.1
 python-dateutil==2.8.2
 python-dotenv==1.0.0
 pytz==2023.3.post1

tests/agents/user_agent_test.py

@@ -0,0 +1,53 @@
+import sys
+import os
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")))
+
+from agents.user_agent import UserAgent
+
+import pytest
+
+
+@pytest.fixture
+def user_agent():
+    """Fixture to create an instance of UserAgent for testing."""
+    return UserAgent(name="test")
+
+
+def test_set_current_message(user_agent):
+    """Test the set_current_message method."""
+    message = "Test message"
+
+    # Set a message and check if it's correctly stored
+    user_agent.set_current_message(message)
+    assert user_agent.current_message == message
+
+    # Set another message and check if it's updated
+    new_message = "New test message"
+    user_agent.set_current_message(new_message)
+    assert user_agent.current_message == new_message
+
+
+def test_get_human_input(user_agent):
+    """Test the get_human_input method."""
+    prompt = "Enter a message: "
+
+    # Start a thread to simulate user input
+    import threading
+
+    def simulate_user_input():
+        user_agent.set_current_message("User input")
+
+    input_thread = threading.Thread(target=simulate_user_input)
+    input_thread.start()
+
+    # Check if the get_human_input method returns the user input
+    user_input = user_agent.get_human_input(prompt)
+    assert user_input == "User input"
+
+    # Wait for the input thread to finish
+    input_thread.join()
+
+
+if __name__ == "__main__":
+    pytest.main()