From 8b371eaad3b7cc6505dd51aedb9198836bd6fea7 Mon Sep 17 00:00:00 2001
From: yzlin <yzlin@fuzhi.ai>
Date: Wed, 27 Mar 2024 11:41:04 +0800
Subject: [PATCH] add test cases for tool convert

---
 metagpt/tools/tool_convert.py            |  2 +-
 tests/metagpt/tools/test_tool_convert.py | 93 +++++++++++++++++++++++-
 2 files changed, 93 insertions(+), 2 deletions(-)

diff --git a/metagpt/tools/tool_convert.py b/metagpt/tools/tool_convert.py
index e6894762a..7d7374fde 100644
--- a/metagpt/tools/tool_convert.py
+++ b/metagpt/tools/tool_convert.py
@@ -141,7 +141,7 @@ class CodeVisitor(ast.NodeVisitor):
         if node.returns:
             return_annotation = f" -> {ast.unparse(node.returns)}"
 
-        return f"({' ,'.join(args)}){return_annotation}"
+        return f"({', '.join(args)}){return_annotation}"
 
     def get_tool_schemas(self):
         return self.tool_schemas
diff --git a/tests/metagpt/tools/test_tool_convert.py b/tests/metagpt/tools/test_tool_convert.py
index 061a619ce..4798d32b0 100644
--- a/tests/metagpt/tools/test_tool_convert.py
+++ b/tests/metagpt/tools/test_tool_convert.py
@@ -2,7 +2,10 @@ from typing import Literal, Union
 
 import pandas as pd
 
-from metagpt.tools.tool_convert import convert_code_to_tool_schema
+from metagpt.tools.tool_convert import (
+    convert_code_to_tool_schema,
+    convert_code_to_tool_schema_ast,
+)
 
 
 class DummyClass:
@@ -128,3 +131,91 @@ def test_convert_code_to_tool_schema_function():
 def test_convert_code_to_tool_schema_async_function():
     schema = convert_code_to_tool_schema(dummy_async_fn)
     assert schema.get("type") == "async_function"
+
+
+TEST_CODE_FILE_TEXT = '''
+import pandas as pd  # imported obj should not be parsed
+from some_module1 import some_imported_function, SomeImportedClass  # imported obj should not be parsed
+from ..some_module2 import some_imported_function2  # relative import should not result in an error
+
+class MyClass:
+    """This is a MyClass docstring."""
+    def __init__(self, arg1):
+        """This is the constructor docstring."""
+        self.arg1 = arg1
+
+    def my_method(self, arg2: Union[list[str], str], arg3: pd.DataFrame, arg4: int = 1, arg5: Literal["a","b","c"] = "a") -> Tuple[int, str]:
+        """
+        This is a method docstring.
+        
+        Args:
+            arg2 (Union[list[str], str]): A union of a list of strings and a string.
+            ...
+        
+        Returns:
+            Tuple[int, str]: A tuple of an integer and a string.
+        """
+        return self.arg4 + arg5
+    
+    async def my_async_method(self, some_arg) -> str:
+        return "hi"
+    
+    def _private_method(self):  # private should not be parsed
+        return "private"
+
+def my_function(arg1, arg2) -> dict:
+    """This is a function docstring."""
+    return arg1 + arg2
+
+def my_async_function(arg1, arg2) -> dict:
+    return arg1 + arg2
+
+def _private_function():  # private should not be parsed
+    return "private"
+'''
+
+
+def test_convert_code_to_tool_schema_ast():
+    expected = {
+        "MyClass": {
+            "type": "class",
+            "description": "This is a MyClass docstring.",
+            "methods": {
+                "__init__": {
+                    "type": "function",
+                    "description": "This is the constructor docstring.",
+                    "signature": "(self, arg1)",
+                    "parameters": "",
+                },
+                "my_method": {
+                    "type": "function",
+                    "description": "This is a method docstring. ",
+                    "signature": "(self, arg2: Union[list[str], str], arg3: pd.DataFrame, arg4: int = 1, arg5: Literal['a', 'b', 'c'] = 'a') -> Tuple[int, str]",
+                    "parameters": "Args: arg2 (Union[list[str], str]): A union of a list of strings and a string. ... Returns: Tuple[int, str]: A tuple of an integer and a string.",
+                },
+                "my_async_method": {
+                    "type": "async_function",
+                    "description": "",
+                    "signature": "(self, some_arg) -> str",
+                    "parameters": "",
+                },
+            },
+            "code": 'class MyClass:\n    """This is a MyClass docstring."""\n    def __init__(self, arg1):\n        """This is the constructor docstring."""\n        self.arg1 = arg1\n\n    def my_method(self, arg2: Union[list[str], str], arg3: pd.DataFrame, arg4: int = 1, arg5: Literal["a","b","c"] = "a") -> Tuple[int, str]:\n        """\n        This is a method docstring.\n        \n        Args:\n            arg2 (Union[list[str], str]): A union of a list of strings and a string.\n            ...\n        \n        Returns:\n            Tuple[int, str]: A tuple of an integer and a string.\n        """\n        return self.arg4 + arg5\n    \n    async def my_async_method(self, some_arg) -> str:\n        return "hi"\n    \n    def _private_method(self):  # private should not be parsed\n        return "private"',
+        },
+        "my_function": {
+            "type": "function",
+            "description": "This is a function docstring.",
+            "signature": "(arg1, arg2) -> dict",
+            "parameters": "",
+            "code": 'def my_function(arg1, arg2) -> dict:\n    """This is a function docstring."""\n    return arg1 + arg2',
+        },
+        "my_async_function": {
+            "type": "function",
+            "description": "",
+            "signature": "(arg1, arg2) -> dict",
+            "parameters": "",
+            "code": "def my_async_function(arg1, arg2) -> dict:\n    return arg1 + arg2",
+        },
+    }
+    schemas = convert_code_to_tool_schema_ast(TEST_CODE_FILE_TEXT)
+    assert schemas == expected