From 8293d96f24484ccb22556a2fc6f61e804f754820 Mon Sep 17 00:00:00 2001 From: Omri Grossman Date: Sat, 15 Apr 2023 23:15:28 +0300 Subject: [PATCH] Added tests for `create_chat_message` and `generate_context` methods of the chat module --- tests/test_chat.py | 116 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 116 insertions(+) create mode 100644 tests/test_chat.py diff --git a/tests/test_chat.py b/tests/test_chat.py new file mode 100644 index 00000000..5bf652e6 --- /dev/null +++ b/tests/test_chat.py @@ -0,0 +1,116 @@ + +# Generated by CodiumAI +import unittest +import time +from unittest.mock import patch + +from autogpt.chat import create_chat_message, generate_context + + + +""" +Code Analysis + +Objective: +The objective of the function is to create a chat message with the given role and content and return it as a dictionary. + +Inputs: +The function takes two inputs: +- role (str): The role of the message sender, e.g., "system", "user", or "assistant". +- content (str): The content of the message. + +Flow: +The function takes the role and content as input and creates a dictionary containing the role and content of the message. It then returns the dictionary as output. + +Outputs: +The function returns a dictionary containing the role and content of the message. + +""" + + +class TestChat(unittest.TestCase): + + # Tests that the function returns a dictionary with the correct keys and values when valid strings are provided for role and content. + def test_happy_path_role_content(self): + result = create_chat_message("system", "Hello, world!") + self.assertEqual(result, {"role": "system", "content": "Hello, world!"}) + + # Tests that the function returns a dictionary with the correct keys and values when empty strings are provided for role and content. + def test_empty_role_content(self): + result = create_chat_message("", "") + self.assertEqual(result, {"role": "", "content": ""}) + + # Tests the behavior of the generate_context function when all input parameters are empty. + @patch("time.strftime") + def test_generate_context_empty_inputs(self, mock_strftime): + # Mock the time.strftime function to return a fixed value + mock_strftime.return_value = "Sat Apr 15 00:00:00 2023" + # Arrange + prompt = "" + relevant_memory = "" + full_message_history = [] + model = "gpt-3.5-turbo-0301" + + # Act + result = generate_context(prompt, relevant_memory, full_message_history, model) + + # Assert + expected_result = (-1, 47, 3, [ + {"role": "system", "content": ""}, + {"role": "system", "content": f"The current time and date is {time.strftime('%c')}"}, + {"role": "system", "content": f"This reminds you of these events from your past:\n\n\n"}, + ]) + self.assertEqual(result, expected_result) + + # Tests that the function successfully generates a current_context given valid inputs. + def test_generate_context_valid_inputs(self): + # Given + prompt = "What is your favorite color?" + relevant_memory = "You once painted your room blue." + full_message_history = [ + create_chat_message("user", "Hi there!"), + create_chat_message("assistant", "Hello! How can I assist you today?"), + create_chat_message("user", "Can you tell me a joke?"), + create_chat_message("assistant", "Why did the tomato turn red? Because it saw the salad dressing!"), + create_chat_message("user", "Haha, that's funny."), + ] + model = "gpt-3.5-turbo-0301" + + # When + result = generate_context(prompt, relevant_memory, full_message_history, model) + + # Then + self.assertIsInstance(result[0], int) + self.assertIsInstance(result[1], int) + self.assertIsInstance(result[2], int) + self.assertIsInstance(result[3], list) + self.assertGreaterEqual(result[0], 0) + self.assertGreaterEqual(result[1], 0) + self.assertGreaterEqual(result[2], 0) + self.assertGreaterEqual(len(result[3]), 3) # current_context should have at least 3 messages + self.assertLessEqual(result[1], 2048) # token limit for GPT-3.5-turbo-0301 is 2048 tokens + + # Tests that the function works correctly with valid inputs. + def test_generate_context_valid_inputs(self): + # Arrange + prompt = "Hello, how can I assist you today?" + relevant_memory = "You previously mentioned needing help with a software issue." + full_message_history = [ + create_chat_message("user", "Can you help me with a software issue?"), + create_chat_message("assistant", "Of course, what seems to be the problem?"), + create_chat_message("user", "I keep getting an error message."), + create_chat_message("assistant", "Let's try some troubleshooting steps."), + ] + model = "gpt-3.5-turbo-0301" + + # Act + next_message_index, tokens_used, insertion_index, context = generate_context( + prompt, relevant_memory, full_message_history, model + ) + + # Assert + self.assertEqual(next_message_index, 3) + self.assertGreater(tokens_used, 0) + self.assertLessEqual(tokens_used, 2048) + self.assertEqual(insertion_index, 3) + self.assertEqual(len(context), 3)