refine naming and some details

tests/metagpt/actions/test_execute_code.py

@@ -1,121 +0,0 @@
-import pytest
-
-from metagpt.actions.execute_code import ExecutePyCode, truncate
-
-
-@pytest.mark.asyncio
-async def test_code_running():
-    pi = ExecutePyCode()
-    output = await pi.run("print('hello world!')")
-    assert output[1] is True
-    output = await pi.run({"code": "print('hello world!')", "language": "python"})
-    assert output[1] is True
-
-
-@pytest.mark.asyncio
-async def test_split_code_running():
-    pi = ExecutePyCode()
-    output = await pi.run("x=1\ny=2")
-    output = await pi.run("z=x+y")
-    output = await pi.run("assert z==3")
-    assert output[1] is True
-
-
-@pytest.mark.asyncio
-async def test_execute_error():
-    pi = ExecutePyCode()
-    output = await pi.run("z=1/0")
-    assert output[1] is False
-
-
-@pytest.mark.asyncio
-async def test_plotting_code():
-    pi = ExecutePyCode()
-    code = """
-    import numpy as np
-    import matplotlib.pyplot as plt
-
-    # 生成随机数据
-    random_data = np.random.randn(1000)  # 生成1000个符合标准正态分布的随机数
-
-    # 绘制直方图
-    plt.hist(random_data, bins=30, density=True, alpha=0.7, color='blue', edgecolor='black')
-
-    # 添加标题和标签
-    plt.title('Histogram of Random Data')
-    plt.xlabel('Value')
-    plt.ylabel('Frequency')
-
-    # 显示图形
-    plt.show()
-    plt.close()
-    """
-    output = await pi.run(code)
-    assert output[1] is True
-
-
-def test_truncate():
-    # 代码执行成功
-    output, is_success = truncate("hello world", 5, True)
-    assert "Truncated to show only first 5 characters\nhello" in output
-    assert is_success
-    # 代码执行失败
-    output, is_success = truncate("hello world", 5, False)
-    assert "Truncated to show only last 5 characters\nworld" in output
-    assert not is_success
-    # 异步
-    output, is_success = truncate("<coroutine object", 5, True)
-    assert not is_success
-    assert "await" in output
-    # 重复的desc
-    result = "Executed code successfully. Truncated to show only first 5 characters\nhello"
-    output, is_success = truncate(result, 5, True)
-    assert is_success
-    assert output == result
-
-
-@pytest.mark.asyncio
-async def test_run_with_timeout():
-    pi = ExecutePyCode(timeout=1)
-    code = "import time; time.sleep(2)"
-    message, success = await pi.run(code)
-    assert not success
-    assert message.startswith("Cell execution timed out")
-
-
-@pytest.mark.asyncio
-async def test_run_code_text():
-    pi = ExecutePyCode()
-    message, success = await pi.run(code='print("This is a code!")', language="python")
-    assert success
-    assert message == "This is a code!\n"
-    message, success = await pi.run(code="# This is a code!", language="markdown")
-    assert success
-    assert message == "# This is a code!"
-    mix_text = "# Title!\n ```python\n print('This is a code!')```"
-    message, success = await pi.run(code=mix_text, language="markdown")
-    assert success
-    assert message == mix_text
-
-
-@pytest.mark.asyncio
-async def test_terminate():
-    pi = ExecutePyCode()
-    await pi.run(code='print("This is a code!")', language="python")
-    is_kernel_alive = await pi.nb_client.km.is_alive()
-    assert is_kernel_alive
-    await pi.terminate()
-    import time
-
-    time.sleep(2)
-    assert pi.nb_client.km is None
-
-
-@pytest.mark.asyncio
-async def test_reset():
-    pi = ExecutePyCode()
-    await pi.run(code='print("This is a code!")', language="python")
-    is_kernel_alive = await pi.nb_client.km.is_alive()
-    assert is_kernel_alive
-    await pi.reset()
-    assert pi.nb_client.km is None

tests/metagpt/actions/test_execute_nb_code.py

@@ -0,0 +1,123 @@
+import pytest
+
+from metagpt.actions.execute_nb_code import ExecuteNbCode, truncate
+
+
+@pytest.mark.asyncio
+async def test_code_running():
+    executor = ExecuteNbCode()
+    output, is_success = await executor.run("print('hello world!')")
+    assert is_success
+    output, is_success = await executor.run({"code": "print('hello world!')", "language": "python"})
+    assert is_success
+
+
+@pytest.mark.asyncio
+async def test_split_code_running():
+    executor = ExecuteNbCode()
+    _ = await executor.run("x=1\ny=2")
+    _ = await executor.run("z=x+y")
+    output, is_success = await executor.run("assert z==3")
+    assert is_success
+
+
+@pytest.mark.asyncio
+async def test_execute_error():
+    executor = ExecuteNbCode()
+    output, is_success = await executor.run("z=1/0")
+    assert not is_success
+
+
+PLOT_CODE = """
+import numpy as np
+import matplotlib.pyplot as plt
+
+# 生成随机数据
+random_data = np.random.randn(1000)  # 生成1000个符合标准正态分布的随机数
+
+# 绘制直方图
+plt.hist(random_data, bins=30, density=True, alpha=0.7, color='blue', edgecolor='black')
+
+# 添加标题和标签
+plt.title('Histogram of Random Data')
+plt.xlabel('Value')
+plt.ylabel('Frequency')
+
+# 显示图形
+plt.show()
+plt.close()
+"""
+
+
+@pytest.mark.asyncio
+async def test_plotting_code():
+    executor = ExecuteNbCode()
+    output, is_success = await executor.run(PLOT_CODE)
+    assert is_success
+
+
+def test_truncate():
+    # 代码执行成功
+    output, is_success = truncate("hello world", 5, True)
+    assert "Truncated to show only first 5 characters\nhello" in output
+    assert is_success
+    # 代码执行失败
+    output, is_success = truncate("hello world", 5, False)
+    assert "Truncated to show only last 5 characters\nworld" in output
+    assert not is_success
+    # 异步
+    output, is_success = truncate("<coroutine object", 5, True)
+    assert not is_success
+    assert "await" in output
+    # 重复的desc
+    result = "Executed code successfully. Truncated to show only first 5 characters\nhello"
+    output, is_success = truncate(result, 5, True)
+    assert is_success
+    assert output == result
+
+
+@pytest.mark.asyncio
+async def test_run_with_timeout():
+    executor = ExecuteNbCode(timeout=1)
+    code = "import time; time.sleep(2)"
+    message, success = await executor.run(code)
+    assert not success
+    assert message.startswith("Cell execution timed out")
+
+
+@pytest.mark.asyncio
+async def test_run_code_text():
+    executor = ExecuteNbCode()
+    message, success = await executor.run(code='print("This is a code!")', language="python")
+    assert success
+    assert message == "This is a code!\n"
+    message, success = await executor.run(code="# This is a code!", language="markdown")
+    assert success
+    assert message == "# This is a code!"
+    mix_text = "# Title!\n ```python\n print('This is a code!')```"
+    message, success = await executor.run(code=mix_text, language="markdown")
+    assert success
+    assert message == mix_text
+
+
+@pytest.mark.asyncio
+async def test_terminate():
+    executor = ExecuteNbCode()
+    await executor.run(code='print("This is a code!")', language="python")
+    is_kernel_alive = await executor.nb_client.km.is_alive()
+    assert is_kernel_alive
+    await executor.terminate()
+    import time
+
+    time.sleep(2)
+    assert executor.nb_client.km is None
+
+
+@pytest.mark.asyncio
+async def test_reset():
+    executor = ExecuteNbCode()
+    await executor.run(code='print("This is a code!")', language="python")
+    is_kernel_alive = await executor.nb_client.km.is_alive()
+    assert is_kernel_alive
+    await executor.reset()
+    assert executor.nb_client.km is None

tests/metagpt/actions/test_write_analysis_code.py

@@ -2,7 +2,7 @@ import asyncio
 
 import pytest
 
-from metagpt.actions.execute_code import ExecutePyCode
+from metagpt.actions.execute_nb_code import ExecuteNbCode
 from metagpt.actions.write_analysis_code import WriteCodeByGenerate, WriteCodeWithTools
 from metagpt.logs import logger
 from metagpt.plan.planner import STRUCTURAL_CONTEXT
@@ -13,7 +13,7 @@ from metagpt.schema import Message, Plan, Task
 @pytest.mark.asyncio
 async def test_write_code_by_list_plan():
     write_code = WriteCodeByGenerate()
-    execute_code = ExecutePyCode()
+    execute_code = ExecuteNbCode()
     messages = []
     plan = ["随机生成一个pandas DataFrame时间序列", "绘制这个时间序列的直方图", "求均值"]
     for task in plan:

tests/metagpt/roles/run_code_interpreter.py

@@ -1,6 +1,6 @@
 import fire
 
-from metagpt.actions.execute_code import ExecutePyCode
+from metagpt.actions.execute_nb_code import ExecuteNbCode
 from metagpt.const import DATA_PATH
 from metagpt.logs import logger
 from metagpt.roles.code_interpreter import CodeInterpreter
@@ -35,7 +35,7 @@ async def run_code_interpreter(role_class, requirement, auto_run, use_tools, sav
         logger.info("Resuming from history trajectory")
         plan, nb = load_history(save_dir)
         role.planner.plan = Plan(**plan)
-        role.execute_code = ExecutePyCode(nb)
+        role.execute_code = ExecuteNbCode(nb)
 
     else:
         logger.info("Run from scratch")

tests/metagpt/roles/test_code_interpreter.py

@@ -7,7 +7,7 @@ from metagpt.roles.code_interpreter import CodeInterpreter
 @pytest.mark.asyncio
 @pytest.mark.parametrize("auto_run", [(True), (False)])
 async def test_code_interpreter(mocker, auto_run):
-    mocker.patch("metagpt.actions.execute_code.ExecutePyCode.run", return_value=("a successful run", True))
+    mocker.patch("metagpt.actions.execute_nb_code.ExecuteNbCode.run", return_value=("a successful run", True))
     mocker.patch("builtins.input", return_value="confirm")
 
     requirement = "Run data analysis on sklearn Iris dataset, include a plot"

tests/metagpt/roles/test_ml_engineer.py

@@ -1,6 +1,6 @@
 import pytest
 
-from metagpt.actions.execute_code import ExecutePyCode
+from metagpt.actions.execute_nb_code import ExecuteNbCode
 from metagpt.logs import logger
 from metagpt.roles.ml_engineer import MLEngineer
 from metagpt.schema import Message, Plan, Task
@@ -72,7 +72,7 @@ async def test_mle_update_data_columns(mocker):
 @pytest.mark.asyncio
 async def test_mle_debug_code(mocker):
     mle = MLEngineer(auto_run=True, use_tools=True)
-    mle.working_memory.add(Message(content=ErrorStr, cause_by=ExecutePyCode))
+    mle.working_memory.add(Message(content=ErrorStr, cause_by=ExecuteNbCode))
     mle.latest_code = CODE
     mle.debug_context = DebugContext
     code, _ = await mle._write_code()

tests/metagpt/utils/test_save_code.py

@@ -8,7 +8,7 @@ import os
 import nbformat
 import pytest
 
-from metagpt.actions.execute_code import ExecutePyCode
+from metagpt.actions.execute_nb_code import ExecuteNbCode
 from metagpt.utils.save_code import DATA_PATH, save_code_file
 
 
@@ -33,7 +33,7 @@ def test_save_code_file_json():
 @pytest.mark.asyncio
 async def test_save_code_file_notebook():
     code = "print('Hello, World!')"
-    executor = ExecutePyCode()
+    executor = ExecuteNbCode()
     await executor.run(code)
     # Save as a Notebook file
     save_code_file("example_nb", executor.nb, file_format="ipynb")