Merge pull request #1000 from garylin2099/di_fixes

unit tests & legacy codes

metagpt/strategy/task_type.py

@@ -54,6 +54,20 @@ class TaskType(Enum):
     )
     OTHER = TaskTypeDef(name="other", desc="Any tasks not in the defined categories")
 
+    # Legacy TaskType to support tool recommendation using type match. You don't need to define task types if you have no human priors to inject.
+    TEXT2IMAGE = TaskTypeDef(
+        name="text2image",
+        desc="Related to text2image, image2image using stable diffusion model.",
+    )
+    WEBSCRAPING = TaskTypeDef(
+        name="web scraping",
+        desc="For scraping data from web pages.",
+    )
+    EMAIL_LOGIN = TaskTypeDef(
+        name="email login",
+        desc="For logging to an email.",
+    )
+
     @property
     def type_name(self):
         return self.value.name

metagpt/tools/libs/email_login.py

@@ -23,7 +23,7 @@ IMAP_SERVERS = {
 }
 
 
-@register_tool()
+@register_tool(tags=["email login"])
 def email_login_imap(email_address, email_password):
     """
     Use imap_tools package to log in to your email (the email that supports IMAP protocol) to verify and return the account object.

metagpt/tools/libs/gpt_v_generator.py

@@ -28,7 +28,7 @@ As the design pays tribute to large companies, sometimes it is normal for some c
 Now, please generate the corresponding webpage code including HTML, CSS and JavaScript:"""
 
 
-@register_tool(include_functions=["__init__", "generate_webpages", "save_webpages"])
+@register_tool(tags=["image2webpage"], include_functions=["__init__", "generate_webpages", "save_webpages"])
 class GPTvGenerator:
     """Class for generating webpage code from a given webpage screenshot.
 

metagpt/tools/tool_convert.py

@@ -12,6 +12,8 @@ def convert_code_to_tool_schema(obj, include: list[str] = None):
     if inspect.isclass(obj):
         schema = {"type": "class", "description": remove_spaces(docstring), "methods": {}}
         for name, method in inspect.getmembers(obj, inspect.isfunction):
+            if name.startswith("_") and name != "__init__":  # skip private methodss
+                continue
             if include and name not in include:
                 continue
             # method_doc = inspect.getdoc(method)

tests/metagpt/actions/di/test_write_analysis_code.py

@@ -5,7 +5,7 @@ from metagpt.schema import Message
 
 
 @pytest.mark.asyncio
-async def test_write_code():
+async def test_write_code_with_plan():
     write_code = WriteAnalysisCode()
 
     user_requirement = "Run data analysis on sklearn Iris dataset, include a plot"
@@ -16,9 +16,29 @@ async def test_write_code():
     assert "sklearn" in code
 
 
+@pytest.mark.asyncio
+async def test_write_code_with_tools():
+    write_code = WriteAnalysisCode()
+
+    user_requirement = "Preprocess sklearn Wine recognition dataset and train a model to predict wine class (20% as validation), and show validation accuracy."
+    tool_info = """
+    ## Capabilities
+    - You can utilize pre-defined tools in any code lines from 'Available Tools' in the form of Python class or function.
+    - You can freely combine the use of any other public packages, like sklearn, numpy, pandas, etc..
+
+    ## Available Tools:
+    Each tool is described in JSON format. When you call a tool, import the tool from its path first.
+    {'FillMissingValue': {'type': 'class', 'description': 'Completing missing values with simple strategies.', 'methods': {'__init__': {'type': 'function', 'description': 'Initialize self. ', 'signature': '(self, features: \'list\', strategy: "Literal[\'mean\', \'median\', \'most_frequent\', \'constant\']" = \'mean\', fill_value=None)', 'parameters': 'Args: features (list): Columns to be processed. strategy (Literal["mean", "median", "most_frequent", "constant"], optional): The imputation strategy, notice \'mean\' and \'median\' can only be used for numeric features. Defaults to \'mean\'. fill_value (int, optional): Fill_value is used to replace all occurrences of missing_values. Defaults to None.'}, 'fit': {'type': 'function', 'description': 'Fit a model to be used in subsequent transform. ', 'signature': "(self, df: 'pd.DataFrame')", 'parameters': 'Args: df (pd.DataFrame): The input DataFrame.'}, 'fit_transform': {'type': 'function', 'description': 'Fit and transform the input DataFrame. ', 'signature': "(self, df: 'pd.DataFrame') -> 'pd.DataFrame'", 'parameters': 'Args: df (pd.DataFrame): The input DataFrame. Returns: pd.DataFrame: The transformed DataFrame.'}, 'transform': {'type': 'function', 'description': 'Transform the input DataFrame with the fitted model. ', 'signature': "(self, df: 'pd.DataFrame') -> 'pd.DataFrame'", 'parameters': 'Args: df (pd.DataFrame): The input DataFrame. Returns: pd.DataFrame: The transformed DataFrame.'}}, 'tool_path': 'metagpt/tools/libs/data_preprocess.py'}
+    """
+
+    code = await write_code.run(user_requirement=user_requirement, tool_info=tool_info)
+    assert len(code) > 0
+    assert "metagpt.tools.libs" in code
+
+
 @pytest.mark.asyncio
 async def test_debug_with_reflection():
-    user_requirement = "Run data analysis on sklearn Iris dataset, include a plot"
+    user_requirement = "read a dataset test.csv and print its head"
 
     plan_status = """
     ## Finished Tasks