remove ToolTypesEnum

metagpt/tools/libs/data_preprocess.py

@@ -13,10 +13,10 @@ from sklearn.preprocessing import (
     StandardScaler,
 )
 
-from metagpt.tools.tool_data_type import ToolTypeEnum
 from metagpt.tools.tool_registry import register_tool
+from metagpt.tools.tool_types import ToolTypes
 
-TOOL_TYPE = ToolTypeEnum.DATA_PREPROCESS.value
+TOOL_TYPE = ToolTypes.DATA_PREPROCESS.type_name
 
 
 class MLProcess(object):

metagpt/tools/libs/feature_engineering.py

@@ -16,10 +16,10 @@ from sklearn.model_selection import KFold
 from sklearn.preprocessing import KBinsDiscretizer, PolynomialFeatures
 
 from metagpt.tools.libs.data_preprocess import MLProcess
-from metagpt.tools.tool_data_type import ToolTypeEnum
 from metagpt.tools.tool_registry import register_tool
+from metagpt.tools.tool_types import ToolTypes
 
-TOOL_TYPE = ToolTypeEnum.FEATURE_ENGINEERING.value
+TOOL_TYPE = ToolTypes.FEATURE_ENGINEERING.type_name
 
 
 @register_tool(tool_type=TOOL_TYPE)

metagpt/tools/libs/gpt_v_generator.py

@@ -12,8 +12,8 @@ from pathlib import Path
 import requests
 
 from metagpt.const import DEFAULT_WORKSPACE_ROOT
-from metagpt.tools.tool_data_type import ToolTypeEnum
 from metagpt.tools.tool_registry import register_tool
+from metagpt.tools.tool_types import ToolTypes
 
 ANALYZE_LAYOUT_PROMPT = """You are now a UI/UX, please generate layout information for this image:
 
@@ -30,7 +30,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(tool_type=ToolTypeEnum.IMAGE2WEBPAGE.value)
+@register_tool(tool_type=ToolTypes.IMAGE2WEBPAGE.type_name)
 class GPTvGenerator:
     def __init__(self):
         from metagpt.config2 import config

metagpt/tools/libs/sd_engine.py

@@ -16,8 +16,8 @@ from PIL import Image, PngImagePlugin
 #
 from metagpt.const import SD_OUTPUT_FILE_REPO, SOURCE_ROOT
 from metagpt.logs import logger
-from metagpt.tools.tool_data_type import ToolTypeEnum
 from metagpt.tools.tool_registry import register_tool
+from metagpt.tools.tool_types import ToolTypes
 
 payload = {
     "prompt": "",
@@ -53,7 +53,7 @@ payload = {
 default_negative_prompt = "(easynegative:0.8),black, dark,Low resolution"
 
 
-@register_tool(tool_type=ToolTypeEnum.STABLE_DIFFUSION.value)
+@register_tool(tool_type=ToolTypes.STABLE_DIFFUSION.type_name)
 class SDEngine:
     def __init__(self, sd_url=""):
         # Initialize the SDEngine with configuration

metagpt/tools/libs/web_scraping.py

@@ -1,9 +1,9 @@
-from metagpt.tools.tool_data_type import ToolTypeEnum
 from metagpt.tools.tool_registry import register_tool
+from metagpt.tools.tool_types import ToolTypes
 from metagpt.tools.web_browser_engine_playwright import PlaywrightWrapper
 
 
-@register_tool(tool_type=ToolTypeEnum.WEBSCRAPING.value)
+@register_tool(tool_type=ToolTypes.WEBSCRAPING.type_name)
 async def scrape_web_playwright(url, *urls):
     """
     Scrape and save the HTML structure and inner text content of a web page using Playwright.

metagpt/tools/tool_data_type.py

@@ -1,26 +1,9 @@
-from enum import Enum
-
 from pydantic import BaseModel
 
 
-class ToolTypeEnum(Enum):
-    EDA = "eda"
-    DATA_PREPROCESS = "data_preprocess"
-    FEATURE_ENGINEERING = "feature_engineering"
-    MODEL_TRAIN = "model_train"
-    MODEL_EVALUATE = "model_evaluate"
-    STABLE_DIFFUSION = "stable_diffusion"
-    IMAGE2WEBPAGE = "image2webpage"
-    WEBSCRAPING = "web_scraping"
-    OTHER = "other"
-
-    def __missing__(self, key):
-        return self.OTHER
-
-
 class ToolType(BaseModel):
     name: str
-    desc: str
+    desc: str = ""
     usage_prompt: str = ""
 
 

metagpt/tools/tool_registry.py

@@ -11,12 +11,13 @@ import re
 from collections import defaultdict
 
 import yaml
-from pydantic import BaseModel
+from pydantic import BaseModel, field_validator
 
 from metagpt.const import TOOL_SCHEMA_PATH
 from metagpt.logs import logger
 from metagpt.tools.tool_convert import convert_code_to_tool_schema
 from metagpt.tools.tool_data_type import Tool, ToolSchema, ToolType
+from metagpt.tools.tool_types import ToolTypes
 
 
 class ToolRegistry(BaseModel):
@@ -24,16 +25,16 @@ class ToolRegistry(BaseModel):
     tool_types: dict = {}
     tools_by_types: dict = defaultdict(dict)  # two-layer k-v, {tool_type: {tool_name: {...}, ...}, ...}
 
-    def register_tool_type(self, tool_type: ToolType, verbose: bool = False):
-        self.tool_types[tool_type.name] = tool_type
-        if verbose:
-            logger.info(f"tool type {tool_type.name} registered")
+    @field_validator("tool_types", mode="before")
+    @classmethod
+    def init_tool_types(cls, tool_types: ToolTypes):
+        return {tool_type.type_name: tool_type.value for tool_type in tool_types}
 
     def register_tool(
         self,
         tool_name,
         tool_path,
-        schema_path=None,
+        schema_path="",
         tool_code="",
         tool_type="other",
         tool_source_object=None,
@@ -44,6 +45,16 @@ class ToolRegistry(BaseModel):
         if self.has_tool(tool_name):
             return
 
+        if tool_type not in self.tool_types:
+            # register new tool type on the fly
+            logger.warning(
+                f"{tool_type} not previously defined, will create a temporary ToolType with just a name. This ToolType is only effective during this runtime. You may consider add this ToolType with more configs permanently at metagpt.tools.tool_types"
+            )
+            temp_tool_type_obj = ToolType(name=tool_type)
+            self.tool_types[tool_type] = temp_tool_type_obj
+            if verbose:
+                logger.info(f"tool type {tool_type} registered")
+
         schema_path = schema_path or TOOL_SCHEMA_PATH / tool_type / f"{tool_name}.yml"
 
         if not os.path.exists(schema_path):
@@ -93,16 +104,10 @@ class ToolRegistry(BaseModel):
 
 
 # Registry instance
-TOOL_REGISTRY = ToolRegistry()
+TOOL_REGISTRY = ToolRegistry(tool_types=ToolTypes)
 
 
-def register_tool_type(cls):
-    """register a tool type to registry"""
-    TOOL_REGISTRY.register_tool_type(tool_type=cls())
-    return cls
-
-
-def register_tool(tool_name="", tool_type="other", schema_path=None, **kwargs):
+def register_tool(tool_name: str = "", tool_type: str = "other", schema_path: str = "", **kwargs):
     """register a tool to registry"""
 
     def decorator(cls, tool_name=tool_name):

metagpt/tools/tool_types.py

@@ -1,3 +1,5 @@
+from enum import Enum
+
 from metagpt.prompts.tool_types import (
     DATA_PREPROCESS_PROMPT,
     FEATURE_ENGINEERING_PROMPT,
@@ -5,64 +7,74 @@ from metagpt.prompts.tool_types import (
     MODEL_EVALUATE_PROMPT,
     MODEL_TRAIN_PROMPT,
 )
-from metagpt.tools.tool_data_type import ToolType, ToolTypeEnum
-from metagpt.tools.tool_registry import register_tool_type
+from metagpt.tools.tool_data_type import ToolType
+
+Eda = ToolType(name="eda", desc="For performing exploratory data analysis")
+
+DataPreprocess = ToolType(
+    name="data_preprocess",
+    desc="Only for changing value inplace.",
+    usage_prompt=DATA_PREPROCESS_PROMPT,
+)
 
 
-@register_tool_type
-class EDA(ToolType):
-    name: str = ToolTypeEnum.EDA.value
-    desc: str = "For performing exploratory data analysis"
+FeatureEngineering = ToolType(
+    name="feature_engineering",
+    desc="Only for creating new columns for input data.",
+    usage_prompt=FEATURE_ENGINEERING_PROMPT,
+)
 
 
-@register_tool_type
-class DataPreprocess(ToolType):
-    name: str = ToolTypeEnum.DATA_PREPROCESS.value
-    desc: str = "Only for changing value inplace."
-    usage_prompt: str = DATA_PREPROCESS_PROMPT
+ModelTrain = ToolType(
+    name="model_train",
+    desc="Only for training model.",
+    usage_prompt=MODEL_TRAIN_PROMPT,
+)
 
 
-@register_tool_type
-class FeatureEngineer(ToolType):
-    name: str = ToolTypeEnum.FEATURE_ENGINEERING.value
-    desc: str = "Only for creating new columns for input data."
-    usage_prompt: str = FEATURE_ENGINEERING_PROMPT
+ModelEvaluate = ToolType(
+    name="model_evaluate",
+    desc="Only for evaluating model.",
+    usage_prompt=MODEL_EVALUATE_PROMPT,
+)
 
 
-@register_tool_type
-class ModelTrain(ToolType):
-    name: str = ToolTypeEnum.MODEL_TRAIN.value
-    desc: str = "Only for training model."
-    usage_prompt: str = MODEL_TRAIN_PROMPT
+StableDiffusion = ToolType(
+    name="stable_diffusion",
+    desc="Related to text2image, image2image using stable diffusion model.",
+)
 
 
-@register_tool_type
-class ModelEvaluate(ToolType):
-    name: str = ToolTypeEnum.MODEL_EVALUATE.value
-    desc: str = "Only for evaluating model."
-    usage_prompt: str = MODEL_EVALUATE_PROMPT
+Image2Webpage = ToolType(
+    name="image2webpage",
+    desc="For converting image into webpage code.",
+    usage_prompt=IMAGE2WEBPAGE_PROMPT,
+)
 
 
-@register_tool_type
-class StableDiffusion(ToolType):
-    name: str = ToolTypeEnum.STABLE_DIFFUSION.value
-    desc: str = "Related to text2image, image2image using stable diffusion model."
+WebScraping = ToolType(
+    name="web_scraping",
+    desc="For scraping data from web pages.",
+)
 
 
-@register_tool_type
-class Image2Webpage(ToolType):
-    name: str = ToolTypeEnum.IMAGE2WEBPAGE.value
-    desc: str = "For converting image into webpage code."
-    usage_prompt: str = IMAGE2WEBPAGE_PROMPT
+Other = ToolType(name="other", desc="Any tools not in the defined categories")
 
 
-@register_tool_type
-class WebScraping(ToolType):
-    name: str = ToolTypeEnum.WEBSCRAPING.value
-    desc: str = "For scraping data from web pages."
+class ToolTypes(Enum):
+    EDA = Eda
+    DATA_PREPROCESS = DataPreprocess
+    FEATURE_ENGINEERING = FeatureEngineering
+    MODEL_TRAIN = ModelTrain
+    MODEL_EVALUATE = ModelEvaluate
+    STABLE_DIFFUSION = StableDiffusion
+    IMAGE2WEBPAGE = Image2Webpage
+    WEBSCRAPING = WebScraping
+    OTHER = Other
 
+    def __missing__(self, key):
+        return self.OTHER
 
-@register_tool_type
-class Other(ToolType):
-    name: str = ToolTypeEnum.OTHER.value
-    desc: str = "Any tools not in the defined categories"
+    @property
+    def type_name(self):
+        return self.value.name