lionagi 0.0.305__py3-none-any.whl → 0.0.307__py3-none-any.whl

lionagi/tests/test_core/test_base_branch.py

@@ -1,427 +1,427 @@
-from lionagi.core.branch.base_branch import BaseBranch
-from lionagi.core.branch.util import MessageUtil
-from lionagi.core.messages.schema import System
-from lionagi.core.schema import DataLogger
-
-
-import unittest
-from unittest.mock import patch, MagicMock
-import pandas as pd
-from datetime import datetime
-import json
-
-
-class TestBaseBranch(unittest.TestCase):
-
-    def setUp(self):
-        # Patching DataLogger to avoid filesystem interactions during tests
-        self.patcher = patch(
-            "lionagi.core.branch.base_branch.DataLogger", autospec=True
-        )
-        self.MockDataLogger = self.patcher.start()
-        self.addCleanup(self.patcher.stop)
-
-        self.branch = BaseBranch()
-        # Example messages to populate the DataFrame for testing
-        self.test_messages = [
-            {
-                "node_id": "1",
-                "timestamp": "2021-01-01 00:00:00",
-                "role": "system",
-                "sender": "system",
-                "content": json.dumps({"system_info": "System message"}),
-            },
-            {
-                "node_id": "2",
-                "timestamp": "2021-01-01 00:01:00",
-                "role": "user",
-                "sender": "user1",
-                "content": json.dumps({"instruction": "User message"}),
-            },
-            {
-                "node_id": "3",
-                "timestamp": "2021-01-01 00:02:00",
-                "role": "assistant",
-                "sender": "assistant",
-                "content": json.dumps({"response": "Assistant response"}),
-            },
-            {
-                "node_id": "4",
-                "timestamp": "2021-01-01 00:03:00",
-                "role": "assistant",
-                "sender": "action_request",
-                "content": json.dumps({"action_request": "Action request"}),
-            },
-            {
-                "node_id": "5",
-                "timestamp": "2021-01-01 00:04:00",
-                "role": "assistant",
-                "sender": "action_response",
-                "content": json.dumps({"action_response": "Action response"}),
-            },
-        ]
-        self.branch = BaseBranch(messages=pd.DataFrame(self.test_messages))
-
-    def test_init_with_empty_messages(self):
-        """Test __init__ method with no messages provided."""
-        branch = BaseBranch()
-        self.assertTrue(branch.messages.empty)
-
-    def test_init_with_given_messages(self):
-        """Test __init__ method with provided messages DataFrame."""
-        messages = pd.DataFrame(
-            [
-                [
-                    "0",
-                    datetime(2024, 1, 1),
-                    "system",
-                    "system",
-                    json.dumps({"system_info": "Hi"}),
-                ]
-            ],
-            columns=["node_id", "timestamp", "role", "sender", "content"],
-        )
-        branch = BaseBranch(messages=messages)
-        self.assertFalse(branch.messages.empty)
-
-    def test_add_message(self):
-        """Test adding a message."""
-        message_info = {"info": "Test Message"}
-        with patch.object(
-            MessageUtil, "create_message", return_value=System(system=message_info)
-        ) as mocked_create_message:
-            self.branch.add_message(system=message_info)
-            mocked_create_message.assert_called_once()
-            self.assertEqual(len(self.branch.messages), 6)
-
-    def test_chat_messages_without_sender(self):
-        """Test chat_messages property without including sender information."""
-        chat_messages = self.branch.chat_messages
-        self.assertEqual(len(chat_messages), 5)
-        self.assertNotIn("Sender", chat_messages[0]["content"])
-
-    def test_chat_messages_with_sender(self):
-        """Test retrieving chat messages with sender information included."""
-        chat_messages_with_sender = self.branch.chat_messages_with_sender
-        expected_prefixes = [
-            "Sender system: ",
-            "Sender user1: ",
-            "Sender assistant: ",
-            "Sender action_request: ",
-            "Sender action_response: ",
-        ]
-
-        # Verify the number of messages returned matches the number added
-        self.assertEqual(len(chat_messages_with_sender), len(self.test_messages))
-
-        # Check each message for correct sender prefix in content
-        for i, message_dict in enumerate(chat_messages_with_sender):
-            self.assertTrue(
-                message_dict["content"].startswith(expected_prefixes[i]),
-                msg=f"Message content does not start with expected sender prefix. Found: {message_dict['content']}",
-            )
-
-        # Optionally, verify the content matches after removing the prefix
-        for i, message_dict in enumerate(chat_messages_with_sender):
-            prefix, content = message_dict["content"].split(": ", 1)
-            self.assertEqual(
-                content,
-                self.test_messages[i]["content"],
-                msg=f"Message content does not match after removing sender prefix. Found: {content}, Expected: {self.test_messages[i]['content']}",
-            )
-
-    def test_last_message(self):
-        """Test retrieving the last message."""
-        last_message = self.branch.last_message
-        self.assertEqual(
-            json.loads(last_message.iloc[0]["content"]),
-            {"action_response": "Action response"},
-        )
-
-    def test_last_message_content(self):
-        """Test retrieving the content of the last message."""
-        content = self.branch.last_message_content
-        self.assertEqual(content, {"action_response": "Action response"})
-
-    def test_first_system_message(self):
-        """Test retrieving the first 'system' message."""
-        first_system_message = self.branch.first_system
-        self.assertEqual(
-            json.loads(first_system_message.iloc[0]["content"]),
-            {"system_info": "System message"},
-        )
-
-    def test_last_response(self):
-        """Test retrieving the last 'assistant' message."""
-        last_response = self.branch.last_response
-        self.assertEqual(last_response.iloc[0]["sender"], "action_response")
-
-    def test_last_response_content(self):
-        """Test extracting content of the last 'assistant' message."""
-        content = self.branch.last_response_content
-        self.assertEqual(content, {"action_response": "Action response"})
-
-    def test_action_request_messages(self):
-        """Test filtering messages sent by 'action_request'."""
-        action_requests = self.branch.action_request
-        self.assertTrue(all(action_requests["sender"] == "action_request"))
-
-    def test_action_response_messages(self):
-        """Test filtering messages sent by 'action_response'."""
-        action_responses = self.branch.action_response
-        self.assertTrue(all(action_responses["sender"] == "action_response"))
-
-    def test_responses(self):
-        """Test filtering of 'assistant' role messages."""
-        responses = self.branch.responses
-        # Verify that all returned messages have the 'assistant' role
-        self.assertTrue(all(responses["role"] == "assistant"))
-        # Optionally, check the count matches the expected number of 'assistant' messages
-        expected_count = sum(
-            1 for msg in self.test_messages if msg["role"] == "assistant"
-        )
-        self.assertEqual(len(responses), expected_count)
-
-    def test_assistant_responses(self):
-        """Test filtering of 'assistant' messages excluding action requests/responses."""
-        assistant_responses = self.branch.assistant_responses
-        # Verify that no returned messages are from 'action_request' or 'action_response' senders
-        self.assertTrue(all(assistant_responses["sender"] != "action_request"))
-        self.assertTrue(all(assistant_responses["sender"] != "action_response"))
-        # Verify that all returned messages have the 'assistant' role
-        self.assertTrue(all(assistant_responses["role"] == "assistant"))
-
-    def test_info(self):
-        """Test summarization of message counts by role."""
-        info = self.branch.info
-        # Verify that the dictionary contains keys for each role and 'total'
-        self.assertIn("assistant", info)
-        self.assertIn("user", info)
-        self.assertIn("system", info)
-        self.assertIn("total", info)
-        # Optionally, verify the counts match expected values
-        for role in ["assistant", "user", "system"]:
-            expected_count = sum(1 for msg in self.test_messages if msg["role"] == role)
-            self.assertEqual(info[role], expected_count)
-        self.assertEqual(info["total"], len(self.test_messages))
-
-    def test_sender_info(self):
-        """Test summarization of message counts by sender."""
-        sender_info = self.branch.sender_info
-        # Verify that the dictionary contains keys for each sender and the counts match
-        for sender in set(msg["sender"] for msg in self.test_messages):
-            expected_count = sum(
-                1 for msg in self.test_messages if msg["sender"] == sender
-            )
-            self.assertEqual(sender_info.get(sender, 0), expected_count)
-
-    def test_describe(self):
-        """Test detailed description of the branch."""
-        description = self.branch.describe
-        # Verify that the description contains expected keys
-        self.assertIn("total_messages", description)
-        self.assertIn("summary_by_role", description)
-        self.assertIn("messages", description)
-        # Optionally, verify the accuracy of the values
-        self.assertEqual(description["total_messages"], len(self.test_messages))
-        self.assertEqual(len(description["messages"]), min(5, len(self.test_messages)))
-
-    @patch("lionagi.libs.ln_dataframe.read_csv")
-    def test_from_csv(cls, mock_read_csv):
-        # Define a mock return value for read_csv
-        mock_messages_df = pd.DataFrame(
-            {
-                "node_id": ["1", "2"],
-                "timestamp": [datetime(2021, 1, 1), datetime(2021, 1, 1)],
-                "role": ["system", "user"],
-                "sender": ["system", "user1"],
-                "content": [
-                    json.dumps({"system_info": "System message"}),
-                    json.dumps({"instruction": "User message"}),
-                ],
-            }
-        )
-        mock_read_csv.return_value = mock_messages_df
-
-        # Call the from_csv class method
-        branch = BaseBranch.from_csv(filename="dummy.csv")
-
-        # Verify that read_csv was called correctly
-        mock_read_csv.assert_called_once_with("dummy.csv")
-
-        # Verify that the branch instance contains the correct messages
-        pd.testing.assert_frame_equal(branch.messages, mock_messages_df)
-
-    @patch("lionagi.libs.ln_dataframe.read_json")
-    def test_from_json(cls, mock_read_csv):
-        # Define a mock return value for read_csv
-        mock_messages_df = pd.DataFrame(
-            {
-                "node_id": ["1", "2"],
-                "timestamp": [datetime(2021, 1, 1), datetime(2021, 1, 1)],
-                "role": ["system", "user"],
-                "sender": ["system", "user1"],
-                "content": [
-                    json.dumps({"system_info": "System message"}),
-                    json.dumps({"instruction": "User message"}),
-                ],
-            }
-        )
-        mock_read_csv.return_value = mock_messages_df
-
-        # Call the from_csv class method
-        branch = BaseBranch.from_json_string(filename="dummy.json")
-
-        # Verify that read_csv was called correctly
-        mock_read_csv.assert_called_once_with("dummy.json")
-
-        # Verify that the branch instance contains the correct messages
-        pd.testing.assert_frame_equal(branch.messages, mock_messages_df)
-
-    @patch(
-        "lionagi.libs.sys_util.SysUtil.create_path",
-        return_value="path/to/messages.csv",
-    )
-    @patch.object(pd.DataFrame, "to_csv")
-    def test_to_csv(self, mock_to_csv, mock_create_path):
-        self.branch.datalogger = MagicMock()
-        self.branch.datalogger.persist_path = "data/logs/"
-
-        self.branch.to_csv_file("messages.csv", verbose=False, clear=False)
-
-        mock_create_path.assert_called_once_with(
-            self.branch.datalogger.persist_path,
-            "messages.csv",
-            timestamp=True,
-            dir_exist_ok=True,
-            time_prefix=False,
-        )
-        mock_to_csv.assert_called_once_with("path/to/messages.csv")
-
-        # Verify that messages are not cleared after exporting
-        assert not self.branch.messages.empty
-
-    @patch(
-        "lionagi.libs.sys_util.SysUtil.create_path",
-        return_value="path/to/messages.json",
-    )
-    @patch.object(pd.DataFrame, "to_json")
-    def test_to_json(self, mock_to_json, mock_create_path):
-        self.branch.datalogger = MagicMock()
-        self.branch.datalogger.persist_path = "data/logs/"
-
-        self.branch.to_json_file("messages.json", verbose=False, clear=False)
-
-        mock_create_path.assert_called_once_with(
-            self.branch.datalogger.persist_path,
-            "messages.json",
-            timestamp=True,
-            dir_exist_ok=True,
-            time_prefix=False,
-        )
-        mock_to_json.assert_called_once_with(
-            "path/to/messages.json", orient="records", lines=True, date_format="iso"
-        )
-
-        # Verify that messages are not cleared after exporting
-        assert not self.branch.messages.empty
-
-    # def test_log_to_csv(self):
-    #     self.branch.log_to_csv('log.csv', verbose=False, clear=False)
-
-    #     self.branch.datalogger.to_csv_file.assert_called_once_with(filename='log.csv', dir_exist_ok=True, timestamp=True,
-    #                                             time_prefix=False, verbose=False, clear=False)
-
-    # def test_log_to_json(self):
-    #     branch = BaseBranch()
-    #     branch.log_to_json('log.json', verbose=False, clear=False)
-
-    #     self.branch.datalogger.to_json_file.assert_called_once_with(filename='log.json', dir_exist_ok=True, timestamp=True,
-    #                                              time_prefix=False, verbose=False, clear=False)
-
-    def test_remove_message(self):
-        """Test removing a message from the branch based on its node ID."""
-        initial_length = len(self.branch.messages)
-        node_id_to_remove = "2"
-        self.branch.remove_message(node_id_to_remove)
-        final_length = len(self.branch.messages)
-        self.assertNotIn(node_id_to_remove, self.branch.messages["node_id"].tolist())
-        self.assertEqual(final_length, initial_length - 1)
-
-    def test_update_message(self):
-        """Test updating a specific message's column identified by node_id with a new value."""
-        node_id_to_update = "2"
-        new_value = "Updated content"
-        self.branch.update_message(node_id_to_update, "content", new_value)
-        updated_value = self.branch.messages.loc[
-            self.branch.messages["node_id"] == node_id_to_update, "content"
-        ].values[0]
-        self.assertEqual(updated_value, new_value)
-
-    # def test_change_first_system_message(self):
-    #     """Test updating the first system message with new content and/or sender."""
-    #     new_system_content = {"system_info": "Updated system message"}
-    #     self.branch.change_first_system_message(new_system_content['system_info'])
-    #     first_system_message_content = self.branch.messages.loc[self.branch.messages['role'] == 'system', 'content'].iloc[0]
-    #     self.assertIn(json.dumps({"system_info": "Updated system message"}), first_system_message_content)
-
-    def test_rollback(self):
-        """Test removing the last 'n' messages from the branch."""
-        steps_to_remove = 2
-        initial_length = len(self.branch.messages)
-        self.branch.rollback(steps_to_remove)
-        final_length = len(self.branch.messages)
-        self.assertEqual(final_length, initial_length - steps_to_remove)
-
-    def test_clear_messages(self):
-        """Test clearing all messages from the branch."""
-        self.branch.clear_messages()
-        self.assertTrue(self.branch.messages.empty)
-
-    def test_replace_keyword(self):
-        """Test replacing occurrences of a specified keyword with a replacement string."""
-        keyword = "Assistant response"
-        replacement = "Helper feedback"
-        self.branch.replace_keyword(keyword, replacement)
-        self.assertTrue(
-            any(replacement in message for message in self.branch.messages["content"])
-        )
-
-    def test_search_keywords(self):
-        """Test filtering messages by a specified keyword or list of keywords."""
-        keyword_to_search = "Assistant"
-        filtered_messages = self.branch.search_keywords(keyword_to_search)
-        self.assertTrue(
-            all(
-                keyword_to_search in content for content in filtered_messages["content"]
-            )
-        )
-
-    def test_extend(self):
-        """Test extending branch messages with additional messages."""
-        additional_messages = pd.DataFrame(
-            [
-                {
-                    "node_id": "6",
-                    "timestamp": datetime(2021, 1, 1),
-                    "role": "user",
-                    "sender": "user2",
-                    "content": json.dumps({"instruction": "Another user message"}),
-                }
-            ]
-        )
-        initial_length = len(self.branch.messages)
-        self.branch.extend(additional_messages)
-        self.assertEqual(len(self.branch.messages), initial_length + 1)
-
-    def test_filter_by(self):
-        """Test filtering messages by various criteria."""
-        # Filtering by role
-        filtered_messages = self.branch.filter_by(role="user")
-        self.assertTrue(
-            all(msg["role"] == "user" for _, msg in filtered_messages.iterrows())
-        )
-
-
-if __name__ == "__main__":
-    unittest.main()
+# from lionagi.core.branch.base_branch import BaseBranch
+# from lionagi.core.branch.util import MessageUtil
+# from lionagi.core.messages.schema import System
+# from lionagi.core.schema import DataLogger
+
+
+# import unittest
+# from unittest.mock import patch, MagicMock
+# import pandas as pd
+# from datetime import datetime
+# import json
+
+
+# class TestBaseBranch(unittest.TestCase):
+
+#     def setUp(self):
+#         # Patching DataLogger to avoid filesystem interactions during tests
+#         self.patcher = patch(
+#             "lionagi.core.branch.base_branch.DataLogger", autospec=True
+#         )
+#         self.MockDataLogger = self.patcher.start()
+#         self.addCleanup(self.patcher.stop)
+
+#         self.branch = BaseBranch()
+#         # Example messages to populate the DataFrame for testing
+#         self.test_messages = [
+#             {
+#                 "node_id": "1",
+#                 "timestamp": "2021-01-01 00:00:00",
+#                 "role": "system",
+#                 "sender": "system",
+#                 "content": json.dumps({"system_info": "System message"}),
+#             },
+#             {
+#                 "node_id": "2",
+#                 "timestamp": "2021-01-01 00:01:00",
+#                 "role": "user",
+#                 "sender": "user1",
+#                 "content": json.dumps({"instruction": "User message"}),
+#             },
+#             {
+#                 "node_id": "3",
+#                 "timestamp": "2021-01-01 00:02:00",
+#                 "role": "assistant",
+#                 "sender": "assistant",
+#                 "content": json.dumps({"response": "Assistant response"}),
+#             },
+#             {
+#                 "node_id": "4",
+#                 "timestamp": "2021-01-01 00:03:00",
+#                 "role": "assistant",
+#                 "sender": "action_request",
+#                 "content": json.dumps({"action_request": "Action request"}),
+#             },
+#             {
+#                 "node_id": "5",
+#                 "timestamp": "2021-01-01 00:04:00",
+#                 "role": "assistant",
+#                 "sender": "action_response",
+#                 "content": json.dumps({"action_response": "Action response"}),
+#             },
+#         ]
+#         self.branch = BaseBranch(messages=pd.DataFrame(self.test_messages))
+
+#     def test_init_with_empty_messages(self):
+#         """Test __init__ method with no messages provided."""
+#         branch = BaseBranch()
+#         self.assertTrue(branch.messages.empty)
+
+#     def test_init_with_given_messages(self):
+#         """Test __init__ method with provided messages DataFrame."""
+#         messages = pd.DataFrame(
+#             [
+#                 [
+#                     "0",
+#                     datetime(2024, 1, 1),
+#                     "system",
+#                     "system",
+#                     json.dumps({"system_info": "Hi"}),
+#                 ]
+#             ],
+#             columns=["node_id", "timestamp", "role", "sender", "content"],
+#         )
+#         branch = BaseBranch(messages=messages)
+#         self.assertFalse(branch.messages.empty)
+
+#     def test_add_message(self):
+#         """Test adding a message."""
+#         message_info = {"info": "Test Message"}
+#         with patch.object(
+#             MessageUtil, "create_message", return_value=System(system=message_info)
+#         ) as mocked_create_message:
+#             self.branch.add_message(system=message_info)
+#             mocked_create_message.assert_called_once()
+#             self.assertEqual(len(self.branch.messages), 6)
+
+#     def test_chat_messages_without_sender(self):
+#         """Test chat_messages property without including sender information."""
+#         chat_messages = self.branch.chat_messages
+#         self.assertEqual(len(chat_messages), 5)
+#         self.assertNotIn("Sender", chat_messages[0]["content"])
+
+#     def test_chat_messages_with_sender(self):
+#         """Test retrieving chat messages with sender information included."""
+#         chat_messages_with_sender = self.branch.chat_messages_with_sender
+#         expected_prefixes = [
+#             "Sender system: ",
+#             "Sender user1: ",
+#             "Sender assistant: ",
+#             "Sender action_request: ",
+#             "Sender action_response: ",
+#         ]
+
+#         # Verify the number of messages returned matches the number added
+#         self.assertEqual(len(chat_messages_with_sender), len(self.test_messages))
+
+#         # Check each message for correct sender prefix in content
+#         for i, message_dict in enumerate(chat_messages_with_sender):
+#             self.assertTrue(
+#                 message_dict["content"].startswith(expected_prefixes[i]),
+#                 msg=f"Message content does not start with expected sender prefix. Found: {message_dict['content']}",
+#             )
+
+#         # Optionally, verify the content matches after removing the prefix
+#         for i, message_dict in enumerate(chat_messages_with_sender):
+#             prefix, content = message_dict["content"].split(": ", 1)
+#             self.assertEqual(
+#                 content,
+#                 self.test_messages[i]["content"],
+#                 msg=f"Message content does not match after removing sender prefix. Found: {content}, Expected: {self.test_messages[i]['content']}",
+#             )
+
+#     def test_last_message(self):
+#         """Test retrieving the last message."""
+#         last_message = self.branch.last_message
+#         self.assertEqual(
+#             json.loads(last_message.iloc[0]["content"]),
+#             {"action_response": "Action response"},
+#         )
+
+#     def test_last_message_content(self):
+#         """Test retrieving the content of the last message."""
+#         content = self.branch.last_message_content
+#         self.assertEqual(content, {"action_response": "Action response"})
+
+#     def test_first_system_message(self):
+#         """Test retrieving the first 'system' message."""
+#         first_system_message = self.branch.first_system
+#         self.assertEqual(
+#             json.loads(first_system_message.iloc[0]["content"]),
+#             {"system_info": "System message"},
+#         )
+
+#     def test_last_response(self):
+#         """Test retrieving the last 'assistant' message."""
+#         last_response = self.branch.last_response
+#         self.assertEqual(last_response.iloc[0]["sender"], "action_response")
+
+#     def test_last_response_content(self):
+#         """Test extracting content of the last 'assistant' message."""
+#         content = self.branch.last_response_content
+#         self.assertEqual(content, {"action_response": "Action response"})
+
+#     def test_action_request_messages(self):
+#         """Test filtering messages sent by 'action_request'."""
+#         action_requests = self.branch.action_request
+#         self.assertTrue(all(action_requests["sender"] == "action_request"))
+
+#     def test_action_response_messages(self):
+#         """Test filtering messages sent by 'action_response'."""
+#         action_responses = self.branch.action_response
+#         self.assertTrue(all(action_responses["sender"] == "action_response"))
+
+#     def test_responses(self):
+#         """Test filtering of 'assistant' role messages."""
+#         responses = self.branch.responses
+#         # Verify that all returned messages have the 'assistant' role
+#         self.assertTrue(all(responses["role"] == "assistant"))
+#         # Optionally, check the count matches the expected number of 'assistant' messages
+#         expected_count = sum(
+#             1 for msg in self.test_messages if msg["role"] == "assistant"
+#         )
+#         self.assertEqual(len(responses), expected_count)
+
+#     def test_assistant_responses(self):
+#         """Test filtering of 'assistant' messages excluding action requests/responses."""
+#         assistant_responses = self.branch.assistant_responses
+#         # Verify that no returned messages are from 'action_request' or 'action_response' senders
+#         self.assertTrue(all(assistant_responses["sender"] != "action_request"))
+#         self.assertTrue(all(assistant_responses["sender"] != "action_response"))
+#         # Verify that all returned messages have the 'assistant' role
+#         self.assertTrue(all(assistant_responses["role"] == "assistant"))
+
+#     def test_info(self):
+#         """Test summarization of message counts by role."""
+#         info = self.branch.info
+#         # Verify that the dictionary contains keys for each role and 'total'
+#         self.assertIn("assistant", info)
+#         self.assertIn("user", info)
+#         self.assertIn("system", info)
+#         self.assertIn("total", info)
+#         # Optionally, verify the counts match expected values
+#         for role in ["assistant", "user", "system"]:
+#             expected_count = sum(1 for msg in self.test_messages if msg["role"] == role)
+#             self.assertEqual(info[role], expected_count)
+#         self.assertEqual(info["total"], len(self.test_messages))
+
+#     def test_sender_info(self):
+#         """Test summarization of message counts by sender."""
+#         sender_info = self.branch.sender_info
+#         # Verify that the dictionary contains keys for each sender and the counts match
+#         for sender in set(msg["sender"] for msg in self.test_messages):
+#             expected_count = sum(
+#                 1 for msg in self.test_messages if msg["sender"] == sender
+#             )
+#             self.assertEqual(sender_info.get(sender, 0), expected_count)
+
+#     def test_describe(self):
+#         """Test detailed description of the branch."""
+#         description = self.branch.describe
+#         # Verify that the description contains expected keys
+#         self.assertIn("total_messages", description)
+#         self.assertIn("summary_by_role", description)
+#         self.assertIn("messages", description)
+#         # Optionally, verify the accuracy of the values
+#         self.assertEqual(description["total_messages"], len(self.test_messages))
+#         self.assertEqual(len(description["messages"]), min(5, len(self.test_messages)))
+
+#     @patch("lionagi.libs.ln_dataframe.read_csv")
+#     def test_from_csv(cls, mock_read_csv):
+#         # Define a mock return value for read_csv
+#         mock_messages_df = pd.DataFrame(
+#             {
+#                 "node_id": ["1", "2"],
+#                 "timestamp": [datetime(2021, 1, 1), datetime(2021, 1, 1)],
+#                 "role": ["system", "user"],
+#                 "sender": ["system", "user1"],
+#                 "content": [
+#                     json.dumps({"system_info": "System message"}),
+#                     json.dumps({"instruction": "User message"}),
+#                 ],
+#             }
+#         )
+#         mock_read_csv.return_value = mock_messages_df
+
+#         # Call the from_csv class method
+#         branch = BaseBranch.from_csv(filename="dummy.csv")
+
+#         # Verify that read_csv was called correctly
+#         mock_read_csv.assert_called_once_with("dummy.csv")
+
+#         # Verify that the branch instance contains the correct messages
+#         pd.testing.assert_frame_equal(branch.messages, mock_messages_df)
+
+#     @patch("lionagi.libs.ln_dataframe.read_json")
+#     def test_from_json(cls, mock_read_csv):
+#         # Define a mock return value for read_csv
+#         mock_messages_df = pd.DataFrame(
+#             {
+#                 "node_id": ["1", "2"],
+#                 "timestamp": [datetime(2021, 1, 1), datetime(2021, 1, 1)],
+#                 "role": ["system", "user"],
+#                 "sender": ["system", "user1"],
+#                 "content": [
+#                     json.dumps({"system_info": "System message"}),
+#                     json.dumps({"instruction": "User message"}),
+#                 ],
+#             }
+#         )
+#         mock_read_csv.return_value = mock_messages_df
+
+#         # Call the from_csv class method
+#         branch = BaseBranch.from_json_string(filename="dummy.json")
+
+#         # Verify that read_csv was called correctly
+#         mock_read_csv.assert_called_once_with("dummy.json")
+
+#         # Verify that the branch instance contains the correct messages
+#         pd.testing.assert_frame_equal(branch.messages, mock_messages_df)
+
+#     @patch(
+#         "lionagi.libs.sys_util.SysUtil.create_path",
+#         return_value="path/to/messages.csv",
+#     )
+#     @patch.object(pd.DataFrame, "to_csv")
+#     def test_to_csv(self, mock_to_csv, mock_create_path):
+#         self.branch.datalogger = MagicMock()
+#         self.branch.datalogger.persist_path = "data/logs/"
+
+#         self.branch.to_csv_file("messages.csv", verbose=False, clear=False)
+
+#         mock_create_path.assert_called_once_with(
+#             self.branch.datalogger.persist_path,
+#             "messages.csv",
+#             timestamp=True,
+#             dir_exist_ok=True,
+#             time_prefix=False,
+#         )
+#         mock_to_csv.assert_called_once_with("path/to/messages.csv")
+
+#         # Verify that messages are not cleared after exporting
+#         assert not self.branch.messages.empty
+
+#     @patch(
+#         "lionagi.libs.sys_util.SysUtil.create_path",
+#         return_value="path/to/messages.json",
+#     )
+#     @patch.object(pd.DataFrame, "to_json")
+#     def test_to_json(self, mock_to_json, mock_create_path):
+#         self.branch.datalogger = MagicMock()
+#         self.branch.datalogger.persist_path = "data/logs/"
+
+#         self.branch.to_json_file("messages.json", verbose=False, clear=False)
+
+#         mock_create_path.assert_called_once_with(
+#             self.branch.datalogger.persist_path,
+#             "messages.json",
+#             timestamp=True,
+#             dir_exist_ok=True,
+#             time_prefix=False,
+#         )
+#         mock_to_json.assert_called_once_with(
+#             "path/to/messages.json", orient="records", lines=True, date_format="iso"
+#         )
+
+#         # Verify that messages are not cleared after exporting
+#         assert not self.branch.messages.empty
+
+#     # def test_log_to_csv(self):
+#     #     self.branch.log_to_csv('log.csv', verbose=False, clear=False)
+
+#     #     self.branch.datalogger.to_csv_file.assert_called_once_with(filename='log.csv', dir_exist_ok=True, timestamp=True,
+#     #                                             time_prefix=False, verbose=False, clear=False)
+
+#     # def test_log_to_json(self):
+#     #     branch = BaseBranch()
+#     #     branch.log_to_json('log.json', verbose=False, clear=False)
+
+#     #     self.branch.datalogger.to_json_file.assert_called_once_with(filename='log.json', dir_exist_ok=True, timestamp=True,
+#     #                                              time_prefix=False, verbose=False, clear=False)
+
+#     def test_remove_message(self):
+#         """Test removing a message from the branch based on its node ID."""
+#         initial_length = len(self.branch.messages)
+#         node_id_to_remove = "2"
+#         self.branch.remove_message(node_id_to_remove)
+#         final_length = len(self.branch.messages)
+#         self.assertNotIn(node_id_to_remove, self.branch.messages["node_id"].tolist())
+#         self.assertEqual(final_length, initial_length - 1)
+
+#     def test_update_message(self):
+#         """Test updating a specific message's column identified by node_id with a new value."""
+#         node_id_to_update = "2"
+#         new_value = "Updated content"
+#         self.branch.update_message(node_id_to_update, "content", new_value)
+#         updated_value = self.branch.messages.loc[
+#             self.branch.messages["node_id"] == node_id_to_update, "content"
+#         ].values[0]
+#         self.assertEqual(updated_value, new_value)
+
+#     # def test_change_first_system_message(self):
+#     #     """Test updating the first system message with new content and/or sender."""
+#     #     new_system_content = {"system_info": "Updated system message"}
+#     #     self.branch.change_first_system_message(new_system_content['system_info'])
+#     #     first_system_message_content = self.branch.messages.loc[self.branch.messages['role'] == 'system', 'content'].iloc[0]
+#     #     self.assertIn(json.dumps({"system_info": "Updated system message"}), first_system_message_content)
+
+#     def test_rollback(self):
+#         """Test removing the last 'n' messages from the branch."""
+#         steps_to_remove = 2
+#         initial_length = len(self.branch.messages)
+#         self.branch.rollback(steps_to_remove)
+#         final_length = len(self.branch.messages)
+#         self.assertEqual(final_length, initial_length - steps_to_remove)
+
+#     def test_clear_messages(self):
+#         """Test clearing all messages from the branch."""
+#         self.branch.clear_messages()
+#         self.assertTrue(self.branch.messages.empty)
+
+#     def test_replace_keyword(self):
+#         """Test replacing occurrences of a specified keyword with a replacement string."""
+#         keyword = "Assistant response"
+#         replacement = "Helper feedback"
+#         self.branch.replace_keyword(keyword, replacement)
+#         self.assertTrue(
+#             any(replacement in message for message in self.branch.messages["content"])
+#         )
+
+#     def test_search_keywords(self):
+#         """Test filtering messages by a specified keyword or list of keywords."""
+#         keyword_to_search = "Assistant"
+#         filtered_messages = self.branch.search_keywords(keyword_to_search)
+#         self.assertTrue(
+#             all(
+#                 keyword_to_search in content for content in filtered_messages["content"]
+#             )
+#         )
+
+#     def test_extend(self):
+#         """Test extending branch messages with additional messages."""
+#         additional_messages = pd.DataFrame(
+#             [
+#                 {
+#                     "node_id": "6",
+#                     "timestamp": datetime(2021, 1, 1),
+#                     "role": "user",
+#                     "sender": "user2",
+#                     "content": json.dumps({"instruction": "Another user message"}),
+#                 }
+#             ]
+#         )
+#         initial_length = len(self.branch.messages)
+#         self.branch.extend(additional_messages)
+#         self.assertEqual(len(self.branch.messages), initial_length + 1)
+
+#     def test_filter_by(self):
+#         """Test filtering messages by various criteria."""
+#         # Filtering by role
+#         filtered_messages = self.branch.filter_by(role="user")
+#         self.assertTrue(
+#             all(msg["role"] == "user" for _, msg in filtered_messages.iterrows())
+#         )
+
+
+# if __name__ == "__main__":
+#     unittest.main()