Merge branch 'fix_code_timeout' into 'code_intepreter'

fix timeout

See merge request agents/data_agents_opt!46

metagpt/actions/execute_code.py

@@ -4,6 +4,7 @@
 @Author  :   orange-crow
 @File    :   code_executor.py
 """
+import asyncio
 import re
 import traceback
 from abc import ABC, abstractmethod
@@ -81,6 +82,9 @@ class ExecutePyCode(ExecuteCode, Action):
     async def reset(self):
         """reset NotebookClient"""
         await self.terminate()
+
+        # sleep 1s to wait for the kernel to be cleaned up completely
+        await asyncio.sleep(1)
         await self.build()
         self.nb_client = NotebookClient(self.nb, timeout=self.timeout)
 
@@ -181,7 +185,11 @@ class ExecutePyCode(ExecuteCode, Action):
             await self.nb_client.async_execute_cell(cell, cell_index)
             return True, ""
         except CellTimeoutError:
-            return False, "TimeoutError"
+            assert self.nb_client.km is not None
+            await self.nb_client.km.interrupt_kernel()
+            await asyncio.sleep(1)
+            error_msg = "Cell execution timed out: Execution exceeded the time limit and was stopped; consider optimizing your code for better performance."
+            return False, error_msg
         except DeadKernelError:
             await self.reset()
             return False, "DeadKernelError"

tests/metagpt/actions/test_execute_code.py

@@ -96,4 +96,4 @@ async def test_run_with_timeout():
     code = "import time; time.sleep(2)"
     message, success = await pi.run(code)
     assert not success
-    assert message == "TimeoutError"
+    assert message.startswith("Cell execution timed out")