lionagi 0.1.2__py3-none-any.whl → 0.2.1__py3-none-any.whl

lionagi/tests/test_core/test_branch.py

@@ -1,294 +1,115 @@
-# from lionagi.core.branch.branch import Branch
-# from lionagi.core.tool.tool_manager import ToolManager, func_to_tool
-# from lionagi.core.schema import DataLogger
-# from lionagi.core.branch.util import MessageUtil
-
-
-# import unittest
-# from unittest.mock import patch, MagicMock
-# import pandas as pd
-# import json
-# from collections import deque
-
-
-# class TestBranch(unittest.TestCase):
-#     def setUp(self):
-#         # Assuming no need for actual files or external services for initialization
-#         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 = Branch(
-#             branch_name="TestBranch", messages=pd.DataFrame(self.test_messages)
-#         )
-
-#         def sample_func(param1: int) -> bool:
-#             """Sample function.
-
-#             Args:
-#                 param1 (int): Description of param1.
-
-#             Returns:
-#                 bool: Description of return value.
-#             """
-#             return True
-
-#         self.tool = func_to_tool(sample_func)
-
-#     def test_initialization(self):
-#         """Test the initialization of the Branch class."""
-#         self.assertEqual(self.branch.branch_name, "TestBranch")
-#         self.assertIsInstance(self.branch.tool_manager, ToolManager)
-#         self.assertIsInstance(self.branch.datalogger, DataLogger)
-#         self.assertEqual(self.branch.sender, "system")
-
-#     def test_has_tools_property(self):
-#         """Test the has_tools property."""
-#         # Initially, no tools are registered
-#         self.assertFalse(self.branch.has_tools)
-
-#         # Mock tool registration
-#         self.branch.register_tools(self.tool)
-#         self.assertTrue(self.branch.has_tools)
-
-
-#     # @patch("lionagi.core.branch.BaseBranch._from_csv")
-#     # def test_from_csv(self, mock_from_csv):
-#     #     """Test creating a Branch instance from a CSV file."""
-#     #     filepath = "path/to/your/csvfile.csv"
-#     #     Branch.from_csv(filepath=filepath, branch_name="TestBranchFromCSV")
-#     #     mock_from_csv.assert_called_once_with(
-#     #         filepath=filepath,
-#     #         read_kwargs=None,
-#     #         branch_name="TestBranchFromCSV",
-#     #         service=None,
-#     #         llmconfig=None,
-#     #         tools=None,
-#     #         datalogger=None,
-#     #         persist_path=None,
-#     #         tool_manager=None,
-#     #     )
-
-#     # @patch("lionagi.core.branch.BaseBranch._from_json")
-#     # def test_from_json(self, mock_from_json):
-#     #     """Test creating a Branch instance from a JSON file."""
-#     #     filepath = "path/to/your/jsonfile.json"
-#     #     Branch.from_json_string(filepath=filepath, branch_name="TestBranchFromJSON")
-#     #     mock_from_json.assert_called_once_with(
-#     #         filepath=filepath,
-#     #         read_kwargs=None,
-#     #         branch_name="TestBranchFromJSON",
-#     #         service=None,
-#     #         llmconfig=None,
-#     #         tools=None,
-#     #         datalogger=None,
-#     #         persist_path=None,
-#     #         tool_manager=None,
-#     #     )
-
-
-#     def test_messages_describe(self):
-#         """Test the messages_describe method for accuracy."""
-#         # Assuming self.branch has been set up with some messages
-#         description = self.branch.messages_describe()
-#         self.assertIn("total_messages", description)
-#         self.assertIn("summary_by_role", description)
-#         self.assertIn("summary_by_sender", description)
-#         self.assertIn("registered_tools", description)
-
-#     def test_merge_branch(self):
-#         """Test merging another Branch instance into the current one."""
-#         mes = [
-#             {
-#                 "node_id": "6",
-#                 "timestamp": "2021-01-01 00:01:00",
-#                 "role": "user",
-#                 "sender": "user1",
-#                 "content": json.dumps({"instruction": "User message"}),
-#             }
-#         ]
-#         other_branch = Branch(branch_name="OtherBranch", messages=pd.DataFrame(mes))
-
-#         original_message_count = len(self.branch.messages)
-#         self.branch.merge_branch(other_branch)
-#         merged_message_count = len(self.branch.messages)
-#         self.assertTrue(merged_message_count > original_message_count)
-
-#     def test_register_and_delete_tools(self):
-#         """Test tool registration and deletion."""
-#         self.branch.register_tools(self.tool)
-#         self.assertIn("sample_func", self.branch.tool_manager.registry)
-#         self.branch.delete_tools(self.tool, verbose=False)
-#         self.assertNotIn("sample_func", self.branch.tool_manager.registry)
-
-#     def test_send(self):
-#         """Test sending a mail package."""
-#         package = {"data": "example"}
-#         self.branch.send(recipient="BranchB", category="messages", package=package)
-#         self.assertEqual(len(self.branch.pending_outs), 1)
-#         mail = self.branch.pending_outs[0]
-#         self.assertEqual(mail.sender, "system")
-#         self.assertEqual(mail.recipient, "BranchB")
-#         self.assertEqual(mail.category, "messages")
-#         self.assertEqual(mail.package, package)
-
-#     # def test_is_invoked_true(self):
-#     #     branch = Branch()
-
-#     #     mock_messages = [
-#     #         MessageUtil.to_json_content({"action_response": {"function": "func_name", "arguments": {}, "output": "result"}})
-#     #     ]
-#     #     branch.messages = pd.DataFrame(mock_messages, columns=['content'])
-#     #     self.assertTrue(branch._is_invoked())
-
-#     def test_is_invoked_false(self):
-#         """Test that _is_invoked returns False when the last message is not a valid action response."""
-#         self.assertFalse(self.branch._is_invoked())
-
-
-# class TestBranchReceive(unittest.TestCase):
-#     def setUp(self):
-#         self.branch = Branch(branch_name="TestBranch")
-#         # Set up a mock sender and initial pending_ins structure
-#         self.sender = "MockSender"
-#         self.branch.pending_ins[self.sender] = deque()
-
-#     # @patch("lionagi.core.mail.BaseMail")
-#     # @patch("lionagi.core.branch.util.MessageUtil.validate_messages")
-#     # def test_receive_messages(self, mock_validate_messages, mock_base_mail):
-#     #     # Prepare a mock mail package with messages
-#     #     messages_df = pd.DataFrame(
-#     #         [
-#     #             {
-#     #                 "node_id": "1",
-#     #                 "timestamp": "2021-01-01 00:00:00",
-#     #                 "role": "system",
-#     #                 "sender": "system",
-#     #                 "content": json.dumps({"system_info": "System message"}),
-#     #             }
-#     #         ]
-#     #     )
-#     #     mail_package_messages = MagicMock(category="messages", package=messages_df)
-#     #     self.branch.pending_ins[self.sender].append(mail_package_messages)
-
-#     #     # Test receiving messages
-#     #     self.branch.receive(self.sender)
-#     #     mock_validate_messages.assert_called_once_with(messages_df)
-#     #     self.assertTrue(len(self.branch.messages) > 0)
-#     #     self.assertEqual(self.branch.pending_ins, {})
-
-#     # def test_receive_tools(self):
-#     #     def sample_func(param1: int) -> bool:
-#     #         """Sample function.
-
-#     #         Args:
-#     #             param1 (int): Description of param1.
-
-#     #         Returns:
-#     #             bool: Description of return value.
-#     #         """
-#     #         return True
-
-#     #     tool = func_to_tool(sample_func)
-#     #     mail_package_tools = MagicMock(category="tools", package=tool)
-#     #     self.branch.pending_ins[self.sender].append(mail_package_tools)
-
-#     #     # Test receiving tools
-#     #     self.branch.receive(self.sender)
-#     #     self.assertIn(tool, self.branch.tool_manager.registry.values())
-
-#     def test_receive_service(self):
-#         # Prepare a mock mail package with a service
-#         from lionagi.libs.ln_api import BaseService
-
-#         service = BaseService()
-#         mail_package_service = MagicMock(category="provider", package=service)
-#         self.branch.pending_ins[self.sender].append(mail_package_service)
-
-#         # Test receiving service
-#         self.branch.receive(self.sender)
-#         self.assertEqual(self.branch.service, service)
-
-#     def test_receive_llmconfig(self):
-#         # Prepare a mock mail package with llmconfig
-#         llmconfig = self.branch.llmconfig.copy()
-#         mail_package_llmconfig = MagicMock(category="llmconfig", package=llmconfig)
-#         self.branch.pending_ins[self.sender].append(mail_package_llmconfig)
-
-#         # Test receiving llmconfig
-#         self.branch.receive(self.sender)
-#         self.assertEqual(llmconfig, self.branch.llmconfig)
-
-#     def test_invalid_format(self):
-#         # Test handling of invalid package format
-#         invalid_package = MagicMock(category="messages", package="Not a DataFrame")
-#         self.branch.pending_ins[self.sender].append(invalid_package)
-
-#         with self.assertRaises(ValueError) as context:
-#             self.branch.receive(self.sender)
-#         self.assertTrue("Invalid messages format" in str(context.exception))
-
-#     def test_receive_all(self):
-#         messages_df = pd.DataFrame(
-#             [
-#                 {
-#                     "node_id": "1",
-#                     "timestamp": "2021-01-01 00:00:00",
-#                     "role": "system",
-#                     "sender": "system",
-#                     "content": json.dumps({"system_info": "System message"}),
-#                 }
-#             ]
-#         )
-#         mail_package_messages = MagicMock(category="messages", package=messages_df)
-#         self.branch.pending_ins[self.sender].append(mail_package_messages)
-
-#         llmconfig = self.branch.llmconfig.copy()
-#         mail_package_llmconfig = MagicMock(category="llmconfig", package=llmconfig)
-#         self.branch.pending_ins[self.sender].append(mail_package_llmconfig)
-
-#         self.branch.receive_all()
-#         self.assertTrue(
-#             not self.branch.pending_ins,
-#             "pending_ins should be empty or contain only skipped requests",
-#         )
-#         self.assertTrue(..., "Additional assertions based on your implementation")
-
-
-# # Chatflow: call_chatcompletion, chat, ReAct, auto_followup
-
-# if __name__ == "__main__":
-#     unittest.main()
+import unittest
+from unittest.mock import MagicMock, patch
+import lionagi as li
+from lionagi.core.message import System, Instruction, AssistantResponse, ActionResponse
+from lionagi.core.collections import Pile, Progression, Exchange
+from lionagi.core.action.tool_manager import ToolManager
+
+
+class TestBranch(unittest.TestCase):
+
+    def setUp(self):
+        self.branch = li.Branch()
+
+    def test_initialize_branch(self):
+        self.assertIsInstance(self.branch, li.Branch)
+        self.assertIsInstance(self.branch.messages, Pile)
+        self.assertIsInstance(self.branch.progress, Progression)
+        self.assertIsInstance(self.branch.tool_manager, ToolManager)
+        self.assertIsInstance(self.branch.mailbox, Exchange)
+        self.assertIsInstance(self.branch.imodel, li.iModel)
+
+    def test_add_message_system(self):
+        self.branch.add_message(
+            system="You are a helpful assistant, let's think step by step"
+        )
+        self.assertEqual(len(self.branch.messages), 1)
+        self.assertEqual(
+            self.branch.messages[0].content,
+            {"system_info": "You are a helpful assistant, let's think step by step"},
+        )
+
+    def test_to_df(self):
+        self.branch.add_message(
+            system="You are a helpful assistant, let's think step by step"
+        )
+        df = self.branch.to_df()
+        self.assertEqual(df.iloc[0]["message_type"], "System")
+        self.assertEqual(df.iloc[0]["role"], "system")
+
+    def test_to_chat_messages(self):
+        self.branch.add_message(
+            system="You are a helpful assistant, let's think step by step"
+        )
+        chat_msgs = self.branch.to_chat_messages()
+        self.assertEqual(chat_msgs[0]["role"], "system")
+        self.assertEqual(
+            chat_msgs[0]["content"],
+            "You are a helpful assistant, let's think step by step",
+        )
+
+    @patch("lionagi.Branch.chat")
+    async def test_chat(self, mock_chat):
+        mock_chat.return_value = (
+            "Rain poured, but their love shone brighter than any storm."
+        )
+        response = await self.branch.chat("tell me a 10 word story", logprobs=True)
+        self.assertEqual(
+            response, "Rain poured, but their love shone brighter than any storm."
+        )
+        mock_chat.assert_called_once_with("tell me a 10 word story", logprobs=True)
+
+    def test_metadata(self):
+        self.branch.add_message(
+            system="You are a helpful assistant, let's think step by step"
+        )
+        self.assertIn("last_updated", self.branch.messages[0].metadata)
+
+    # def test_register_tools(self):
+    #     tool = MagicMock()
+    #     self.branch.register_tools([tool])
+    #     self.assertIn(tool, self.branch.tool_manager.registry.values())
+
+    # def test_delete_tools(self):
+    #     tool = MagicMock()
+    #     tool.schema_ = {"function": {"name": "test_tool"}}
+    #     self.branch.register_tools([tool])
+    #     self.branch.delete_tools([tool])
+    #     self.assertNotIn("test_tool", self.branch.tool_manager.registry)
+
+    def test_send_receive_mail(self):
+        self.branch.send = MagicMock()
+        self.branch.receive = MagicMock()
+        package = MagicMock()
+        self.branch.send(recipient="recipient_id", category="message", package=package)
+        self.branch.receive(sender="recipient_id")
+        self.branch.send.assert_called_once_with(
+            recipient="recipient_id", category="message", package=package
+        )
+        self.branch.receive.assert_called_once_with(sender="recipient_id")
+
+    async def test_chat_with_tool(self):
+        async def mock_multiply(number1, number2, number3=1):
+            return number1 * number2 * number3
+
+        instruction = """
+        solve the following problem
+        """
+        context = """
+        I have 730_000 trees, with average 123 apples per tree, each weigh 0.4 lbs. 
+        20 percent are bad and sold for 0.1 dollar per lbs, 30 percent are sold to 
+        brewery for 0.3 dollar per apple, what is my revenue?
+        """
+
+        self.branch = li.Branch(
+            "act like a calculator, invoke tool uses", tools=[mock_multiply]
+        )
+        response = await self.branch.chat(
+            instruction=instruction, context=context, tools=True
+        )
+        self.assertIsNotNone(response)
+        self.assertIn("revenue", response.lower())
+
+
+if __name__ == "__main__":
+    unittest.main()

lionagi/tests/test_core/test_form.py

@@ -0,0 +1,46 @@
+import unittest
+from lionagi import Form
+
+
+class TestForm(unittest.TestCase):
+
+    def setUp(self):
+        self.form = Form(assignment="input1, input2 -> output")
+
+    def test_initial_state(self):
+        self.assertEqual(self.form.input_fields, ["input1", "input2"])
+        self.assertEqual(self.form.requested_fields, ["output"])
+        self.assertFalse(self.form.workable)
+        self.assertFalse(self.form.filled)
+        self.assertEqual(
+            self.form.work_fields, {"input1": None, "input2": None, "output": None}
+        )
+
+    def test_fill_input_fields(self):
+        self.form.fill(input1=1, input2=2)
+        self.assertTrue(self.form.workable)
+        self.assertFalse(self.form.filled)
+        self.assertEqual(
+            self.form.work_fields, {"input1": 1, "input2": 2, "output": None}
+        )
+
+    def test_fill_all_fields(self):
+        self.form.fill(input1=1, input2=2)
+        self.form.fill(output=3)
+        self.assertTrue(self.form.filled)
+        self.assertFalse(self.form.workable)
+        self.assertEqual(self.form.work_fields, {"input1": 1, "input2": 2, "output": 3})
+
+    def test_fill_once(self):
+        self.form.fill(input1=1, input2=2)
+        self.form.fill(output=3)
+        with self.assertRaises(ValueError) as context:
+            self.form.fill(input1=2, input2=3)
+        self.assertTrue(
+            "Form is filled, cannot be worked on again" in str(context.exception)
+        )
+        self.assertEqual(self.form.work_fields, {"input1": 1, "input2": 2, "output": 3})
+
+
+if __name__ == "__main__":
+    unittest.main()

lionagi/tests/test_core/test_report.py

@@ -0,0 +1,105 @@
+import unittest
+from lionagi.core.report.form import Form
+from lionagi.core.report.report import Report
+
+
+class TestForm(unittest.TestCase):
+
+    def setUp(self):
+        self.form = Form(assignment="input1, input2 -> output")
+
+    def test_initial_fields(self):
+        self.assertEqual(self.form.input_fields, ["input1", "input2"])
+        self.assertEqual(self.form.requested_fields, ["output"])
+        self.assertFalse(self.form.workable)
+        self.assertFalse(self.form.filled)
+        self.assertEqual(
+            self.form.work_fields, {"input1": None, "input2": None, "output": None}
+        )
+
+    def test_fill_input_fields(self):
+        self.form.fill(input1=1, input2=2)
+        self.assertTrue(self.form.workable)
+        self.assertFalse(self.form.filled)
+        self.assertEqual(
+            self.form.work_fields, {"input1": 1, "input2": 2, "output": None}
+        )
+
+    def test_fill_output_field(self):
+        self.form.fill(input1=1, input2=2)
+        self.form.fill(output=3)
+        self.assertTrue(self.form.filled)
+        self.assertEqual(self.form.work_fields, {"input1": 1, "input2": 2, "output": 3})
+        self.assertFalse(self.form.workable)
+
+    def test_fill_again_raises_error(self):
+        self.form.fill(input1=1, input2=2, output=3)
+        with self.assertRaises(ValueError):
+            self.form.fill(input1=2, input2=3)
+
+
+class TestReport(unittest.TestCase):
+
+    def setUp(self):
+        self.report = Report(assignment="a, b -> c")
+
+    def test_initial_fields(self):
+        self.assertEqual(self.report.input_fields, ["a", "b"])
+        self.assertEqual(self.report.requested_fields, ["c"])
+        self.assertEqual(self.report.work_fields, {"a": None, "b": None, "c": None})
+        self.assertFalse(self.report.filled)
+        self.assertFalse(self.report.workable)
+
+    def test_fill_input_fields(self):
+        self.report.fill(a=3, b=4)
+        self.assertEqual(self.report.work_fields, {"a": 3, "b": 4, "c": None})
+        self.assertTrue(self.report.workable)
+
+    def test_fill_output_field(self):
+        self.report.fill(a=3, b=4)
+        self.report.fill(c=4)
+        self.assertEqual(self.report.work_fields, {"a": 3, "b": 4, "c": 4})
+        self.assertTrue(self.report.filled)
+        self.assertFalse(self.report.workable)
+
+    def test_next_forms_none(self):
+        self.assertIsNone(self.report.next_forms())
+
+    def test_complex_assignment(self):
+        self.report = Report(
+            assignment="a, b -> h",
+            form_assignments=[
+                "a, b -> c",
+                "a -> e",
+                "b -> f",
+                "c -> g",
+                "e, f, g -> h",
+            ],
+            a=3,
+            b=4,
+        )
+        self.assertEqual(
+            self.report.work_fields,
+            {
+                "a": None,
+                "b": None,
+                "c": None,
+                "e": None,
+                "f": None,
+                "g": None,
+                "h": None,
+            },
+        )
+        self.report.fill(c=5, e=6, f=7, g=8, h=10)
+        self.assertEqual(
+            self.report.work_fields,
+            {"a": None, "b": None, "c": 5, "e": 6, "f": 7, "g": 8, "h": 10},
+        )
+        self.assertFalse(self.report.filled)
+        self.assertFalse(self.report.workable)
+        with self.assertRaises(ValueError):
+            self.report.fill(c="xx", e="yy")
+
+
+if __name__ == "__main__":
+    unittest.main()