Merge branch 'add_test_for_ml_tools' into 'code_intepreter'

Add test for ml tools

See merge request agents/data_agents_opt!44

tests/metagpt/actions/test_write_analysis_code.py

@@ -3,8 +3,13 @@ import asyncio
 import pytest
 
 from metagpt.actions.execute_code import ExecutePyCode
-from metagpt.actions.write_analysis_code import WriteCodeByGenerate, WriteCodeWithTools
+from metagpt.actions.write_analysis_code import (
+    WriteCodeByGenerate,
+    WriteCodeWithTools,
+    WriteCodeWithToolsML,
+)
 from metagpt.logs import logger
+from metagpt.plan.planner import STRUCTURAL_CONTEXT
 from metagpt.schema import Message, Plan, Task
 
 
@@ -40,13 +45,15 @@ async def test_tool_recommendation():
     tools = await write_code._tool_recommendation(task, code_steps, available_tools)
 
     assert len(tools) == 1
-    assert tools[0] == ["fill_missing_value"]
+    assert tools[0] == "fill_missing_value"
 
 
 @pytest.mark.asyncio
 async def test_write_code_with_tools():
     write_code = WriteCodeWithTools()
-    messages = []
+    write_code_ml = WriteCodeWithToolsML()
+
+    requirement = "构造数据集并进行数据清洗"
     task_map = {
         "1": Task(
             task_id="1",
@@ -69,10 +76,6 @@ async def test_write_code_with_tools():
             instruction="对数据集进行数据清洗",
             task_type="data_preprocess",
             dependent_task_ids=["1"],
-            code_steps="""
-                {"Step 1": "对数据集进行去重",
-                "Step 2": "对数据集进行缺失值处理"}
-                """,
         ),
     }
     plan = Plan(
@@ -83,10 +86,22 @@ async def test_write_code_with_tools():
     )
     column_info = ""
 
-    code = await write_code.run(messages, plan, column_info)
+    context = STRUCTURAL_CONTEXT.format(
+        user_requirement=requirement,
+        context=plan.context,
+        tasks=list(task_map.values()),
+        current_task=plan.current_task.model_dump_json(),
+    )
+    context_msg = [Message(content=context, role="user")]
+
+    code = await write_code.run(context_msg, plan)
     assert len(code) > 0
     print(code)
 
+    code_with_ml = await write_code_ml.run([], plan, column_info)
+    assert len(code_with_ml) > 0
+    print(code_with_ml)
+
 
 @pytest.mark.asyncio
 async def test_write_code_to_correct_error():

tests/metagpt/tools/functions/libs/__init__.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
-# @Time    : 2023/11/17 10:24
+# @Time    : 2024/1/11 16:14
 # @Author  : lidanyang
 # @File    : __init__.py
 # @Desc    :

tests/metagpt/tools/functions/libs/test_data_preprocess.py

@@ -0,0 +1,111 @@
+from datetime import datetime
+
+import numpy as np
+import numpy.testing as npt
+import pandas as pd
+import pytest
+
+from metagpt.tools.functions.libs.data_preprocess import (
+    FillMissingValue,
+    LabelEncode,
+    MaxAbsScale,
+    MinMaxScale,
+    OneHotEncode,
+    OrdinalEncode,
+    RobustScale,
+    StandardScale,
+    get_column_info,
+)
+
+
+@pytest.fixture
+def mock_datasets():
+    return pd.DataFrame(
+        {
+            "num1": [1, 2, np.nan, 4, 5],
+            "cat1": ["A", "B", np.nan, "D", "A"],
+            "date1": [
+                datetime(2020, 1, 1),
+                datetime(2020, 1, 2),
+                datetime(2020, 1, 3),
+                datetime(2020, 1, 4),
+                datetime(2020, 1, 5),
+            ],
+        }
+    )
+
+
+def test_fill_missing_value(mock_datasets):
+    fm = FillMissingValue(features=["num1"], strategy="mean")
+    transformed = fm.fit_transform(mock_datasets.copy())
+
+    assert transformed["num1"].isnull().sum() == 0
+
+
+def test_min_max_scale(mock_datasets):
+    mms = MinMaxScale(features=["num1"])
+    transformed = mms.fit_transform(mock_datasets.copy())
+
+    npt.assert_allclose(transformed["num1"].min(), 0)
+    npt.assert_allclose(transformed["num1"].max(), 1)
+
+
+def test_standard_scale(mock_datasets):
+    ss = StandardScale(features=["num1"])
+    transformed = ss.fit_transform(mock_datasets.copy())
+
+    assert int(transformed["num1"].mean()) == 0
+    assert int(transformed["num1"].std()) == 1
+
+
+def test_max_abs_scale(mock_datasets):
+    mas = MaxAbsScale(features=["num1"])
+    transformed = mas.fit_transform(mock_datasets.copy())
+
+    npt.assert_allclose(transformed["num1"].abs().max(), 1)
+
+
+def test_robust_scale(mock_datasets):
+    rs = RobustScale(features=["num1"])
+    transformed = rs.fit_transform(mock_datasets.copy())
+
+    assert int(transformed["num1"].median()) == 0
+
+
+def test_ordinal_encode(mock_datasets):
+    oe = OrdinalEncode(features=["cat1"])
+    transformed = oe.fit_transform(mock_datasets.copy())
+
+    assert transformed["cat1"].max() == 2
+
+
+def test_one_hot_encode(mock_datasets):
+    ohe = OneHotEncode(features=["cat1"])
+    transformed = ohe.fit_transform(mock_datasets.copy())
+
+    assert transformed["cat1_A"].max() == 1
+
+
+def test_label_encode(mock_datasets):
+    le = LabelEncode(features=["cat1"])
+    transformed = le.fit_transform(mock_datasets.copy())
+
+    assert transformed["cat1"].max() == 3
+
+    # test transform with unseen data
+    test = mock_datasets.copy()
+    test["cat1"] = ["A", "B", "C", "D", "E"]
+    transformed = le.transform(test)
+    assert transformed["cat1"].max() == 4
+
+
+def test_get_column_info(mock_datasets):
+    df = mock_datasets
+    column_info = get_column_info(df)
+
+    assert column_info == {
+        "Category": ["cat1"],
+        "Numeric": ["num1"],
+        "Datetime": ["date1"],
+        "Others": [],
+    }

tests/metagpt/tools/functions/libs/test_feature_engineering.py

@@ -0,0 +1,174 @@
+import numpy as np
+import pandas as pd
+import pytest
+from sklearn.datasets import fetch_california_housing, load_breast_cancer, load_iris
+
+from metagpt.tools.functions.libs.feature_engineering import (
+    CatCount,
+    CatCross,
+    ExtractTimeComps,
+    GeneralSelection,
+    GroupStat,
+    KFoldTargetMeanEncoder,
+    PolynomialExpansion,
+    SplitBins,
+    TargetMeanEncoder,
+    TreeBasedSelection,
+    VarianceBasedSelection,
+)
+
+
+@pytest.fixture
+def mock_dataset():
+    return pd.DataFrame(
+        {
+            "num1": [1, 2, np.nan, 4, 5, 6, 7, 3],
+            "num2": [1, 3, 2, 1, np.nan, 5, 6, 4],
+            "num3": [np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan],
+            "cat1": ["A", "B", np.nan, "D", "E", "C", "B", "A"],
+            "cat2": ["A", "A", "A", "A", "A", "A", "A", "A"],
+            "date1": [
+                "2020-01-01",
+                "2020-01-02",
+                "2020-01-03",
+                "2020-01-04",
+                "2020-01-05",
+                "2020-01-06",
+                "2020-01-07",
+                "2020-01-08",
+            ],
+            "label": [0, 1, 0, 1, 0, 1, 0, 1],
+        }
+    )
+
+
+def load_sklearn_data(data_name):
+    if data_name == "iris":
+        data = load_iris()
+    elif data_name == "breast_cancer":
+        data = load_breast_cancer()
+    elif data_name == "housing":
+        data = fetch_california_housing()
+    else:
+        raise ValueError("data_name not supported")
+
+    X, y, feature_names = data.data, data.target, data.feature_names
+    data = pd.DataFrame(X, columns=feature_names)
+    data["label"] = y
+    return data
+
+
+def test_polynomial_expansion(mock_dataset):
+    pe = PolynomialExpansion(cols=["num1", "num2", "label"], degree=2, label_col="label")
+    transformed = pe.fit_transform(mock_dataset)
+
+    assert len(transformed.columns) == len(mock_dataset.columns) + 3
+
+    # when too many columns
+    data = load_sklearn_data("breast_cancer")
+    cols = [c for c in data.columns if c != "label"]
+    pe = PolynomialExpansion(cols=cols, degree=2, label_col="label")
+    transformed = pe.fit_transform(data)
+
+    assert len(transformed.columns) == len(data.columns) + 55
+
+
+def test_cat_count(mock_dataset):
+    cc = CatCount(col="cat1")
+    transformed = cc.fit_transform(mock_dataset)
+
+    assert "cat1_cnt" in transformed.columns
+    assert transformed["cat1_cnt"][0] == 2
+
+
+def test_target_mean_encoder(mock_dataset):
+    tme = TargetMeanEncoder(col="cat1", label="label")
+    transformed = tme.fit_transform(mock_dataset)
+
+    assert "cat1_target_mean" in transformed.columns
+    assert transformed["cat1_target_mean"][0] == 0.5
+
+
+def test_kfold_target_mean_encoder(mock_dataset):
+    kfme = KFoldTargetMeanEncoder(col="cat1", label="label")
+    transformed = kfme.fit_transform(mock_dataset)
+
+    assert "cat1_kf_target_mean" in transformed.columns
+
+
+def test_cat_cross(mock_dataset):
+    cc = CatCross(cols=["cat1", "cat2"])
+    transformed = cc.fit_transform(mock_dataset)
+
+    assert "cat1_cat2" in transformed.columns
+
+    cc = CatCross(cols=["cat1", "cat2"], max_cat_num=3)
+    transformed = cc.fit_transform(mock_dataset)
+
+    assert "cat1_cat2" not in transformed.columns
+
+
+def test_group_stat(mock_dataset):
+    gs = GroupStat(group_col="cat1", agg_col="num1", agg_funcs=["mean", "sum"])
+    transformed = gs.fit_transform(mock_dataset)
+
+    assert "num1_mean_by_cat1" in transformed.columns
+    assert "num1_sum_by_cat1" in transformed.columns
+
+
+def test_split_bins(mock_dataset):
+    sb = SplitBins(cols=["num1"])
+    transformed = sb.fit_transform(mock_dataset)
+
+    assert transformed["num1"].nunique() <= 5
+    assert all(0 <= x < 5 for x in transformed["num1"])
+
+
+def test_extract_time_comps(mock_dataset):
+    time_comps = ["year", "month", "day", "hour", "dayofweek", "is_weekend"]
+    etc = ExtractTimeComps(time_col="date1", time_comps=time_comps)
+    transformed = etc.fit_transform(mock_dataset.copy())
+
+    for comp in time_comps:
+        assert comp in transformed.columns
+    assert transformed["year"][0] == 2020
+    assert transformed["month"][0] == 1
+    assert transformed["day"][0] == 1
+    assert transformed["hour"][0] == 0
+    assert transformed["dayofweek"][0] == 3
+    assert transformed["is_weekend"][0] == 0
+
+
+def test_general_selection(mock_dataset):
+    gs = GeneralSelection(label_col="label")
+    transformed = gs.fit_transform(mock_dataset.copy())
+
+    assert "num3" not in transformed.columns
+    assert "cat2" not in transformed.columns
+
+
+def test_tree_based_selection(mock_dataset):
+    # regression
+    data = load_sklearn_data("housing")
+    tbs = TreeBasedSelection(label_col="label", task_type="reg")
+    transformed = tbs.fit_transform(data)
+    assert len(transformed.columns) > 1
+
+    # classification
+    data = load_sklearn_data("breast_cancer")
+    tbs = TreeBasedSelection(label_col="label", task_type="cls")
+    transformed = tbs.fit_transform(data)
+    assert len(transformed.columns) > 1
+
+    # multi-classification
+    data = load_sklearn_data("iris")
+    tbs = TreeBasedSelection(label_col="label", task_type="mcls")
+    transformed = tbs.fit_transform(data)
+    assert len(transformed.columns) > 1
+
+
+def test_variance_based_selection(mock_dataset):
+    vbs = VarianceBasedSelection(label_col="label")
+    transformed = vbs.fit_transform(mock_dataset.copy())
+
+    assert "num3" not in transformed.columns

tests/metagpt/tools/functions/register/test_register.py

@@ -1,55 +0,0 @@
-#!/usr/bin/env python
-# -*- coding: utf-8 -*-
-# @Time    : 2023/11/17 10:24
-# @Author  : lidanyang
-# @File    : test_register.py
-# @Desc    :
-import pytest
-
-from metagpt.tools.functions.register.register import FunctionRegistry
-from metagpt.tools.functions.schemas.base import ToolSchema, tool_field
-
-
-@pytest.fixture
-def registry():
-    return FunctionRegistry()
-
-
-class AddNumbers(ToolSchema):
-    """Add two numbers"""
-
-    num1: int = tool_field(description="First number")
-    num2: int = tool_field(description="Second number")
-
-
-def test_register(registry):
-    @registry.register("module1", AddNumbers)
-    def add_numbers(num1, num2):
-        return num1 + num2
-
-    assert len(registry.functions["module1"]) == 1
-    assert "add_numbers" in registry.functions["module1"]
-
-    with pytest.raises(ValueError):
-
-        @registry.register("module1", AddNumbers)
-        def add_numbers(num1, num2):
-            return num1 + num2
-
-    func = registry.get("module1", "add_numbers")
-    assert func["func"](1, 2) == 3
-    assert func["schema"] == {
-        "name": "add_numbers",
-        "description": "Add two numbers",
-        "parameters": {
-            "type": "object",
-            "properties": {
-                "num1": {"description": "First number", "type": "int"},
-                "num2": {"description": "Second number", "type": "int"},
-            },
-            "required": ["num1", "num2"],
-        },
-    }
-
-    module1_funcs = registry.get_all_by_module("module1")
-    assert len(module1_funcs) == 1