chore: setup workspace for exchange (#105)

2024-10-02 11:05:43 -07:00
parent aa70408ff7
commit 92fe8e7008
63 changed files with 5895 additions and 120 deletions
@@ -0,0 +1,161 @@
+import attrs
+from exchange.tool import Tool
+
+
+def get_current_weather(location: str) -> None:
+    """Get the current weather in a given location
+
+    Args:
+        location (str): The city and state, e.g. San Francisco, CA
+    """
+    pass
+
+
+def test_load():
+    tool = Tool.from_function(get_current_weather)
+
+    expected = {
+        "name": "get_current_weather",
+        "description": "Get the current weather in a given location",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "location": {
+                    "type": "string",
+                    "description": "The city and state, e.g. San Francisco, CA",
+                },
+            },
+            "required": ["location"],
+        },
+        "function": get_current_weather,
+    }
+
+    assert attrs.asdict(tool) == expected
+
+
+def another_function(
+    param1: int,
+    param2: str,
+    param3: bool,
+    param4: float,
+    param5: list[int],
+    param6: dict[str, float],
+) -> None:
+    """
+    This is another example function with various types
+
+    Args:
+        param1 (int): Description for param1
+        param2 (str): Description for param2
+        param3 (bool): Description for param3
+        param4 (float): Description for param4
+        param5 (list[int]): Description for param5
+        param6 (dict[str, float]): Description for param6
+    """
+    pass
+
+
+def test_load_types():
+    tool = Tool.from_function(another_function)
+    expected_schema = {
+        "type": "object",
+        "properties": {
+            "param1": {"type": "integer", "description": "Description for param1"},
+            "param2": {"type": "string", "description": "Description for param2"},
+            "param3": {"type": "boolean", "description": "Description for param3"},
+            "param4": {"type": "number", "description": "Description for param4"},
+            "param5": {
+                "type": "array",
+                "items": {"type": "integer"},
+                "description": "Description for param5",
+            },
+            "param6": {
+                "type": "object",
+                "additionalProperties": {"type": "number"},
+                "description": "Description for param6",
+            },
+        },
+        "required": ["param1", "param2", "param3", "param4", "param5", "param6"],
+    }
+    assert tool.parameters == expected_schema
+
+
+def numpy_function(param1: int, param2: str) -> None:
+    """
+    This function uses numpy style docstrings
+
+    Parameters
+    ----------
+    param1 : int
+        Description for param1
+    param2 : str
+        Description for param2
+    """
+    pass
+
+
+def test_load_numpy_style():
+    tool = Tool.from_function(numpy_function)
+    expected_schema = {
+        "type": "object",
+        "properties": {
+            "param1": {"type": "integer", "description": "Description for param1"},
+            "param2": {"type": "string", "description": "Description for param2"},
+        },
+        "required": ["param1", "param2"],
+    }
+    assert tool.parameters == expected_schema
+
+
+def sphinx_function(param1: int, param2: str, param3: bool) -> None:
+    """
+    This function uses sphinx style docstrings
+
+    :param param1: Description for param1
+    :type param1: int
+    :param param2: Description for param2
+    :type param2: str
+    :param param3: Description for param3
+    :type param3: bool
+    """
+    pass
+
+
+def test_load_sphinx_style():
+    tool = Tool.from_function(sphinx_function)
+    expected_schema = {
+        "type": "object",
+        "properties": {
+            "param1": {"type": "integer", "description": "Description for param1"},
+            "param2": {"type": "string", "description": "Description for param2"},
+            "param3": {"type": "boolean", "description": "Description for param3"},
+        },
+        "required": ["param1", "param2", "param3"],
+    }
+    assert tool.parameters == expected_schema
+
+
+class FunctionLike:
+    def __init__(self, state: int) -> None:
+        self.state = state
+
+    def __call__(self, param1: int) -> int:
+        """Example
+
+        Args:
+            param1 (int): Description for param1
+        """
+        return self.state + param1
+
+
+def test_load_stateful_class():
+    tool = Tool.from_function(FunctionLike(1))
+    expected_schema = {
+        "type": "object",
+        "properties": {
+            "param1": {"type": "integer", "description": "Description for param1"},
+        },
+        "required": ["param1"],
+    }
+    assert tool.parameters == expected_schema
+    assert tool.function(2) == 3