lionagi 0.0.114__py3-none-any.whl → 0.0.116__py3-none-any.whl

lionagi/services/base_api_service.py

@@ -5,7 +5,9 @@ import tiktoken
 import logging
 import aiohttp
 from typing import Generator, NoReturn, Dict, Any, Optional
-from .service_objs import BaseService, RateLimiter, StatusTracker, AsyncQueue
+from ..objs.status_tracker import StatusTracker
+from ..objs.abc_objs import  BaseService, RateLimiter
+from ..objs.async_queue import AsyncQueue
 
 class BaseAPIRateLimiter(RateLimiter):
 
@@ -20,7 +22,7 @@ class BaseAPIRateLimiter(RateLimiter):
     ) -> None:
         self = cls(max_requests_per_minute, max_tokens_per_minute)
         if not os.getenv("env_readthedocs"):
-            self.rate_limit_replenisher_task = await asyncio.create_task(
+            self.rate_limit_replenisher_task = asyncio.create_task(
                 self.rate_limit_replenisher()
             )
         return self
@@ -43,6 +45,7 @@ class BaseAPIRateLimiter(RateLimiter):
             self.available_request_capacity = self.max_requests_per_minute
             self.available_token_capacity = self.max_tokens_per_minute
 
+# credit to OpenAI for the following function
     def calculate_num_token(
         self,
         payload: Dict[str, Any] = None,
@@ -130,13 +133,16 @@ class BaseAPIService(BaseService):
     def __init__(self, api_key: str = None, 
                  status_tracker = None,
                  queue = None, endpoint=None, schema=None, 
-                 ratelimiter=None, max_requests_per_minute=None, max_tokens_per_minute=None) -> None:
+                 ratelimiter=BaseAPIRateLimiter, max_requests_per_minute=None, max_tokens_per_minute=None) -> None:
         self.api_key = api_key
         self.status_tracker = status_tracker or StatusTracker()
         self.queue = queue or AsyncQueue()
         self.endpoint=endpoint
         self.schema = schema
-        self.rate_limiter = ratelimiter(max_requests_per_minute, max_tokens_per_minute)
+        self.max_requests_per_minute = max_requests_per_minute
+        self.max_tokens_per_minute = max_tokens_per_minute
+        self.rate_limiter_class = ratelimiter
+        self.rate_limiter = None
     
     @staticmethod                    
     def api_methods(http_session, method="post"):
@@ -168,6 +174,10 @@ class BaseAPIService(BaseService):
             yield task_id
             task_id += 1
 
+    async def _init(self):
+        if self.rate_limiter is None:
+            self.rate_limiter = await self.rate_limiter_class.create(self.max_requests_per_minute, self.max_tokens_per_minute)
+
     async def _call_api(self, http_session, endpoint_, method="post", payload: Dict[str, any] =None) -> Optional[Dict[str, any]]:
         endpoint_ = self.api_endpoint_from_url("https://api.openai.com/v1/"+endpoint_)
         
@@ -224,4 +234,4 @@ class BaseAPIService(BaseService):
         except Exception as e:
                 self.status_tracker.num_tasks_failed += 1
                 raise e
-            
+            

lionagi/services/oai.py

@@ -15,7 +15,7 @@ class OpenAIService(BaseAPIService):
         max_attempts: int = 3,
         max_requests_per_minute: int = 500,
         max_tokens_per_minute: int = 150_000,
-        ratelimiter = BaseAPIRateLimiter ,
+        ratelimiter = BaseAPIRateLimiter,
         status_tracker = None,
         queue = None,
     ):
@@ -31,4 +31,4 @@ class OpenAIService(BaseAPIService):
 
     async def serve(self, payload, endpoint_="chat/completions", method="post"):
         return await self._serve(payload=payload, endpoint_=endpoint_, method=method)
-    
+    

lionagi/services/openrouter.py

@@ -2,8 +2,7 @@ from os import getenv
 from .base_api_service import BaseAPIService, BaseAPIRateLimiter
 
 class OpenRouterService(BaseAPIService):
-    _key_scheme = "OPENROUTER_API_KEY"
-
+    key_scheme = "OPENROUTER_API_KEY"
     base_url = "https://openrouter.ai/api/v1/"
 
     def __init__(
@@ -18,7 +17,7 @@ class OpenRouterService(BaseAPIService):
         queue = None,
     ):
         super().__init__(
-            api_key = api_key or getenv(self._key_scheme),
+            api_key = api_key or getenv(self.key_scheme),
             status_tracker = status_tracker,
             queue = queue,
             ratelimiter=ratelimiter,

lionagi/structures/graph.py

@@ -0,0 +1,96 @@
+from pydantic import Field
+
+from lionagi.schema.base_node import BaseNode
+from .relationship import Relationship
+from lionagi.utils.call_util import lcall
+
+
+class Graph(BaseNode):
+    nodes: dict = Field(default={})
+    relationships: dict = Field(default={})
+    node_relationships: dict = Field(default={})
+
+    def add_node(self, node: BaseNode):
+        self.nodes[node.id_] = node
+        self.node_relationships[node.id_] = {'in': {}, 'out': {}}
+
+    def add_relationship(self, relationships: Relationship):
+        if relationships.source_node_id not in self.node_relationships.keys():
+            raise KeyError(f'node {relationships.source_node_id} is not found.')
+        if relationships.target_node_id not in self.node_relationships.keys():
+            raise KeyError(f'node {relationships.target_node_id} is not found.')
+
+        self.relationships[relationships.id_] = relationships
+        self.node_relationships[relationships.source_node_id]['out'][relationships.id_] = relationships.target_node_id
+        self.node_relationships[relationships.target_node_id]['in'][relationships.id_] = relationships.source_node_id
+
+    def get_node_relationships(self, node: BaseNode = None, out_edge=True):
+        if node is None:
+            return list(self.relationships.values())
+
+        if node.id_ not in self.nodes.keys():
+            raise KeyError(f'node {node.id_} is not found')
+
+        if out_edge:
+            relationship_ids = list(self.node_relationships[node.id_]['out'].keys())
+            relationships = lcall(relationship_ids, lambda i: self.relationships[i])
+            return relationships
+        else:
+            relationship_ids = list(self.node_relationships[node.id_]['in'].keys())
+            relationships = lcall(relationship_ids, lambda i: self.relationships[i])
+            return relationships
+
+    def remove_node(self, node: BaseNode):
+        if node.id_ not in self.nodes.keys():
+            raise KeyError(f'node {node.id_} is not found')
+
+        out_relationship = self.node_relationships[node.id_]['out']
+        for relationship_id, node_id in out_relationship.items():
+            self.node_relationships[node_id]['in'].pop(relationship_id)
+            self.relationships.pop(relationship_id)
+
+        in_relationship = self.node_relationships[node.id_]['in']
+        for relationship_id, node_id in in_relationship.items():
+            self.node_relationships[node_id]['out'].pop(relationship_id)
+            self.relationships.pop(relationship_id)
+
+        self.node_relationships.pop(node.id_)
+        return self.nodes.pop(node.id_)
+
+    def remove_relationship(self, relationship: Relationship):
+        if relationship.id_ not in self.relationships.keys():
+            raise KeyError(f'relationship {relationship.id_} is not found')
+
+        self.node_relationships[relationship.source_node_id]['out'].pop(relationship.id_)
+        self.node_relationships[relationship.target_node_id]['in'].pop(relationship.id_)
+
+        return self.relationships.pop(relationship.id_)
+
+    def node_exists(self, node: BaseNode):
+        if node.id_ in self.nodes.keys():
+            return True
+        else:
+            return False
+
+    def relationship_exists(self, relationship: Relationship):
+        if relationship.id_ in self.relationships.keys():
+            return True
+        else:
+            return False
+
+    def to_networkx(self, **kwargs):
+        import networkx as nx
+        g = nx.DiGraph(**kwargs)
+        for node_id, node in self.nodes.items():
+            node_info = node.to_dict()
+            node_info.pop('node_id')
+            g.add_node(node_id, **node_info)
+
+        for relationship_id, relationship in self.relationships.items():
+            relationship_info = relationship.to_dict()
+            relationship_info.pop('node_id')
+            source_node_id = relationship_info.pop('source_node_id')
+            target_node_id = relationship_info.pop('target_node_id')
+            g.add_edge(source_node_id, target_node_id, **relationship_info)
+
+        return g

lionagi/{structure → structures}/relationship.py

@@ -1,6 +1,6 @@
 from pydantic import Field
 from typing import Dict, Optional, Any
-from ..schema.base_schema import BaseNode
+from ..schema.base_node import BaseNode
 
 
 class Relationship(BaseNode):
@@ -125,4 +125,12 @@ class Relationship(BaseNode):
             raise ValueError(f"Target node {self.source_node_id} does not exist")
         else :
             raise ValueError(f"Source node {self.target_node_id} does not exist")
-        
+
+    def __str__(self) -> str:
+        """Returns a simple string representation of the Relationship."""
+        return f"Relationship (id_={self.id_}, from={self.source_node_id}, to={self.target_node_id}, label={self.label})"
+
+    def __repr__(self) -> str:
+        """Returns a detailed string representation of the Relationship."""
+        return f"Relationship(id_={self.id_}, from={self.source_node_id}, to={self.target_node_id}, content={self.content}, " \
+               f"metadata={self.metadata}, label={self.label})"

lionagi/structures/structure.py

@@ -0,0 +1,102 @@
+from typing import TypeVar
+from .graph import Graph
+from ..schema import BaseNode
+from .relationship import Relationship
+
+T = TypeVar('T', bound='BaseNode')
+R = TypeVar('R', bound='Relationship')
+
+
+class Structure(BaseNode):
+    """
+    Represents the structure of a graph consisting of nodes and relationships.
+    """
+    graph: Graph
+
+    def add_node(self, node: T) -> None:
+        """
+        Adds a node to the structure.
+        
+        Args:
+            node (T): The node instance to be added.
+        """
+        self.graph.add_node(node)
+
+    def add_relationship(self, relationship: R) -> None:
+        """
+        Adds a relationship to the structure.
+        
+        Args:
+            relationship (R): The relationship instance to be added.
+        """
+        self.graph.add_relationship(relationship)
+
+    # type can be dict or list
+    # @staticmethod
+    # def _typed_return(type: Type[Union[Dict, List]],
+    #                   obj: Optional[Dict[str, Any]] = None
+    #                   ) -> Union[Dict[str, Any], List[Any]]:
+    #     """
+    #     Returns the object in the specified type format.
+    #
+    #     Args:
+    #         type (Type[Union[Dict, List]]): The type to return the object as (dict or list).
+    #
+    #         obj (Optional[Dict[str, Any]]): The object to be converted.
+    #
+    #     Returns:
+    #         Union[Dict[str, Any], List[Any]]: The object in the specified type format.
+    #     """
+    #     if type is list:
+    #         return list(obj.values())
+    #     return obj
+    
+    def get_relationships(self) -> list[R]:
+        return self.graph.get_node_relationships()
+
+    def get_node_relationships(self, node: T, out_edge=True) -> R:
+        return self.graph.get_node_relationships(node, out_edge)
+
+    def node_exist(self, node: T) -> bool:
+        """
+        Checks if a node exists in the structure.
+        
+        Args:
+            node (T): The node instance or node ID to check for existence.
+            
+        Returns:
+            bool: True if the node exists, False otherwise.
+        """
+
+        return self.graph.node_exist(node)
+    
+    def relationship_exist(self, relationship: R) -> bool:
+        """
+        Checks if a relationship exists in the structure.
+        
+        Args:
+            relationship (R): The relationship instance to check for existence.
+            
+        Returns:
+            bool: True if the relationship exists, False otherwise.
+        """
+        return self.graph.relationship_exists(relationship)
+
+    def remove_node(self, node: T) -> T:
+        """
+        Removes a node and its associated relationships from the structure.
+        
+        Args:
+            node (T): The node instance or node ID to be removed.
+        """
+        return self.graph.remove_node(node)
+
+    def remove_relationship(self, relationship: R) -> R:
+        """
+        Removes a relationship from the structure.
+        
+        Args:
+            relationship (R): The relationship instance to be removed.
+        """
+        return self.graph.remove_relationship(relationship)
+        

lionagi/tests/test_api_util.py

@@ -0,0 +1,46 @@
+import unittest
+from unittest.mock import MagicMock
+
+# Assuming the Python module with the above functions is named 'api_utils'
+from lionagi.utils.api_util import *
+
+class TestApiUtils(unittest.TestCase):
+
+    def test_api_method_valid(self):
+        session = MagicMock()
+        methods = ['post', 'delete', 'head', 'options', 'patch']
+        for method in methods:
+            with self.subTest(method=method):
+                func = api_method(session, method)
+                self.assertTrue(callable(func))
+
+    def test_api_method_invalid(self):
+        session = MagicMock()
+        with self.assertRaises(ValueError):
+            api_method(session, 'get')
+
+    def test_api_error(self):
+        response_with_error = {'error': 'Something went wrong'}
+        response_without_error = {'result': 'Success'}
+
+        self.assertTrue(api_error(response_with_error))
+        self.assertFalse(api_error(response_without_error))
+
+    def test_api_rate_limit_error(self):
+        response_with_rate_limit_error = {'error': {'message': 'Rate limit exceeded'}}
+        response_without_rate_limit_error = {'error': {'message': 'Another error'}}
+
+        self.assertTrue(api_rate_limit_error(response_with_rate_limit_error))
+        self.assertFalse(api_rate_limit_error(response_without_rate_limit_error))
+
+    def test_api_endpoint_from_url(self):
+        url_with_endpoint = "https://api.example.com/v1/users"
+        url_without_endpoint = "https://api.example.com/users"
+        url_invalid = "Just a string"
+
+        self.assertEqual(api_endpoint_from_url(url_with_endpoint), 'users')
+        self.assertEqual(api_endpoint_from_url(url_without_endpoint), '')
+        self.assertEqual(api_endpoint_from_url(url_invalid), '')
+
+if __name__ == '__main__':
+    unittest.main()

lionagi/tests/test_call_util.py

@@ -0,0 +1,115 @@
+# import asyncio
+# import unittest
+# from unittest.mock import patch, MagicMock
+# from lionagi.utils.sys_util import create_copy
+# from lionagi.utils.flat_util import to_list
+# from typing import Callable
+
+# # Assuming the Python module with the above functions is named 'call_utils'
+# from lionagi.utils.call_util import *
+
+# class TestCallUtils(unittest.TestCase):
+
+#     def setUp(self):
+#         self.sample_input = [1, 2, 3]
+#         self.sample_func = MagicMock(return_value=42)
+        
+#     def test_hcall(self):
+#         with patch('time.sleep', return_value=None):
+#             result = hcall(self.sample_input, self.sample_func)
+#             self.sample_func.assert_called_once_with(self.sample_input)
+#             self.assertEqual(result, 42)
+
+#     def test_ahcall(self):
+#         async def async_test():
+#             self.sample_func.reset_mock()
+#             with patch('asyncio.sleep', return_value=None):
+#                 result = await ahcall(self.sample_input, self.sample_func)
+#                 self.sample_func.assert_called_once_with(self.sample_input)
+#                 self.assertEqual(result, 42)
+
+#         asyncio.run(async_test())
+
+#     def test_lcall(self):
+#         expected_result = [42, 42, 42]
+#         result = lcall(self.sample_input, self.sample_func)
+#         self.assertEqual(result, expected_result)
+#         calls = [unittest.mock.call(item) for item in self.sample_input]
+#         self.sample_func.assert_has_calls(calls, any_order=True)
+
+#     def test_alcall(self):
+#         async def async_test():
+#             expected_result = [42, 42, 42]
+#             self.sample_func.reset_mock()
+#             result = await alcall(self.sample_input, self.sample_func)
+#             self.assertEqual(result, expected_result)
+#             calls = [unittest.mock.call(item) for item in self.sample_input]
+#             self.sample_func.assert_has_calls(calls, any_order=True)
+
+#         asyncio.run(async_test())
+
+
+# class TestCallUtils2(unittest.TestCase):
+
+#     def setUp(self):
+#         self.sample_input = [1, 2, 3]
+#         self.sample_func = MagicMock(side_effect=lambda x: x + 1)
+#         self.sample_async_func = MagicMock(side_effect=lambda x: x + 1)
+        
+#     def test_mcall_single_function(self):
+#         result = mcall(self.sample_input, self.sample_func)
+#         self.sample_func.assert_has_calls([unittest.mock.call(i) for i in self.sample_input], any_order=True)
+#         self.assertEqual(result, [2, 3, 4])
+
+#     def test_mcall_multiple_functions(self):
+#         # Define multiple functions
+#         funcs = [
+#             MagicMock(side_effect=lambda x: x + 1),
+#             MagicMock(side_effect=lambda x: x * 2),
+#             MagicMock(side_effect=lambda x: x - 1)
+#         ]
+#         result = mcall(self.sample_input, funcs)
+#         for i, func in enumerate(funcs):
+#             func.assert_called_once_with(self.sample_input[i])
+#         self.assertEqual(result, [2, 4, 2])
+
+#     def test_amcall_single_function(self):
+#         async def async_test():
+#             result = await amcall(self.sample_input, self.sample_async_func)
+#             self.sample_async_func.assert_has_calls([unittest.mock.call(i) for i in self.sample_input], any_order=True)
+#             self.assertEqual(result, [2, 3, 4])
+
+#         asyncio.run(async_test())
+
+#     def test_amcall_multiple_functions(self):
+#         # Define multiple asynchronous functions
+#         async_funcs = [
+#             MagicMock(side_effect=lambda x: x + 1),
+#             MagicMock(side_effect=lambda x: x * 2),
+#             MagicMock(side_effect=lambda x: x - 1)
+#         ]
+
+#         async def async_test():
+#             result = await amcall(self.sample_input, async_funcs)
+#             for i, func in enumerate(async_funcs):
+#                 func.assert_called_once_with(self.sample_input[i])
+#             self.assertEqual(result, [2, 4, 2])
+
+#         asyncio.run(async_test())
+
+#     def test_ecall(self):
+#         funcs = [MagicMock(side_effect=lambda x: x * x)]
+#         result = ecall(self.sample_input, funcs)
+#         self.assertEqual(result, [[1, 4, 9]])
+
+#     def test_aecall(self):
+#         async_funcs = [MagicMock(side_effect=lambda x: x * x)]
+
+#         async def async_test():
+#             result = await aecall(self.sample_input, async_funcs)
+#             self.assertEqual(result, [[1, 4, 9]])
+
+#         asyncio.run(async_test())
+
+# if __name__ == '__main__':
+#     unittest.main()

lionagi/tests/test_convert_util.py

@@ -0,0 +1,202 @@
+import unittest
+from lionagi.utils.convert_util import *
+
+# Test cases for the function
+class TestStrToNum(unittest.TestCase):
+    def test_valid_int(self):
+        self.assertEqual(str_to_num("123"), 123)
+        self.assertEqual(str_to_num("-123"), -123)
+    
+    def test_valid_float(self):
+        self.assertEqual(str_to_num("123.45", num_type=float), 123.45)
+        self.assertEqual(str_to_num("-123.45", num_type=float), -123.45)
+    
+    def test_precision(self):
+        self.assertEqual(str_to_num("123.456", num_type=float, precision=1), 123.5)
+        self.assertEqual(str_to_num("123.444", num_type=float, precision=2), 123.44)
+    
+    def test_bounds(self):
+        self.assertEqual(str_to_num("10", lower_bound=5, upper_bound=15), 10)
+        with self.assertRaises(ValueError):
+            str_to_num("20", upper_bound=15)
+        with self.assertRaises(ValueError):
+            str_to_num("2", lower_bound=5)
+    
+    def test_invalid_input(self):
+        with self.assertRaises(ValueError):
+            str_to_num("abc")
+        with self.assertRaises(ValueError):
+            str_to_num("123abc", num_type=str)
+    
+    def test_no_numeric_value(self):
+        with self.assertRaises(ValueError):
+            str_to_num("No numbers here")
+
+
+
+# Functions to be tested
+def dict_to_xml(data: Dict[str, Any], root_tag: str = 'node') -> str:
+    root = ET.Element(root_tag)
+    _build_xml(root, data)
+    return ET.tostring(root, encoding='unicode')
+
+def _build_xml(element: ET.Element, data: Any):
+    if isinstance(data, dict):
+        for key, value in data.items():
+            sub_element = ET.SubElement(element, key)
+            _build_xml(sub_element, value)
+    elif isinstance(data, list):
+        for item in data:
+            item_element = ET.SubElement(element, 'item')
+            _build_xml(item_element, item)
+    else:
+        element.text = str(data)
+
+def xml_to_dict(element: ET.Element) -> Dict[str, Any]:
+    dict_data = {}
+    for child in element:
+        if list(child):
+            dict_data[child.tag] = xml_to_dict(child)
+        else:
+            dict_data[child.tag] = child.text
+    return dict_data
+
+# Test cases for the functions
+class TestDictXMLConversion(unittest.TestCase):
+    def setUp(self):
+        self.data = {
+            'name': 'John',
+            'age': 30,
+            'children': [
+                {'name': 'Alice', 'age': 5},
+                {'name': 'Bob', 'age': 7}
+            ]
+        }
+        self.root_tag = 'person'
+        self.xml = dict_to_xml(self.data, self.root_tag)
+        self.xml_element = ET.fromstring(self.xml)
+
+    def test_dict_to_xml(self):
+        self.assertIn('<name>John</name>', self.xml)
+        self.assertIn('<age>30</age>', self.xml)
+        self.assertIn('<children>', self.xml)
+        self.assertIn('<item>', self.xml)
+
+    # def test_xml_to_dict(self):
+    #     data_from_xml = xml_to_dict(self.xml_element)
+    #     self.assertEqual(data_from_xml, self.data)
+
+    # def test_xml_to_dict_to_xml(self):
+    #     data_from_xml = xml_to_dict(self.xml_element)
+    #     xml_from_dict = dict_to_xml(data_from_xml, self.root_tag)
+    #     self.assertEqual(xml_from_dict, self.xml)
+
+    # def test_invalid_input(self):
+    #     with self.assertRaises(TypeError):
+    #         dict_to_xml("not a dict", self.root_tag)
+
+
+
+class TestDocstringExtraction(unittest.TestCase):
+    def test_google_style_extraction(self):
+        def sample_func(arg1, arg2):
+            """
+            This is a sample function.
+
+            Args:
+                arg1 (int): The first argument.
+                arg2 (str): The second argument.
+
+            Returns:
+                bool: The truth value.
+            """
+            return True
+
+        description, params = extract_docstring_details_google(sample_func)
+        self.assertEqual(description, "This is a sample function.")
+        self.assertEqual(params, {
+            'arg1': 'The first argument.',
+            'arg2': 'The second argument.'
+        })
+
+    def test_rest_style_extraction(self):
+        def sample_func(arg1, arg2):
+            """
+            This is a sample function.
+
+            :param int arg1: The first argument.
+            :param str arg2: The second argument.
+            :return: The truth value.
+            :rtype: bool
+            """
+            return True
+
+        description, params = extract_docstring_details_rest(sample_func)
+        self.assertEqual(description, "This is a sample function.")
+        self.assertEqual(params, {
+            'arg1': 'The first argument.',
+            'arg2': 'The second argument.'
+        })
+
+    def test_extract_docstring_details_with_invalid_style(self):
+        def sample_func(arg1, arg2):
+            return True
+
+        with self.assertRaises(ValueError):
+            extract_docstring_details(sample_func, style='unsupported')
+
+class TestPythonToJsonTypeConversion(unittest.TestCase):
+    def test_python_to_json_type_conversion(self):
+        self.assertEqual(python_to_json_type('str'), 'string')
+        self.assertEqual(python_to_json_type('int'), 'number')
+        self.assertEqual(python_to_json_type('float'), 'number')
+        self.assertEqual(python_to_json_type('list'), 'array')
+        self.assertEqual(python_to_json_type('tuple'), 'array')
+        self.assertEqual(python_to_json_type('bool'), 'boolean')
+        self.assertEqual(python_to_json_type('dict'), 'object')
+        self.assertEqual(python_to_json_type('nonexistent'), 'object')
+
+class TestFunctionToSchema(unittest.TestCase):
+    def test_func_to_schema(self):
+        def sample_func(arg1: int, arg2: str) -> bool:
+            """
+            This is a sample function.
+
+            Args:
+                arg1 (int): The first argument.
+                arg2 (str): The second argument.
+
+            Returns:
+                bool: The truth value.
+            """
+            return True
+
+        schema = func_to_schema(sample_func, style='google')
+        expected_schema = {
+            "type": "function",
+            "function": {
+                "name": "sample_func",
+                "description": "This is a sample function.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "arg1": {
+                            "type": "number",
+                            "description": "The first argument."
+                        },
+                        "arg2": {
+                            "type": "string",
+                            "description": "The second argument."
+                        }
+                    },
+                    "required": ["arg1", "arg2"]
+                }
+            }
+        }
+        self.assertEqual(schema, expected_schema)
+
+if __name__ == '__main__':
+    unittest.main()
+    
+    
+# --------------------------------------------------