diff --git a/camel/toolkits/code_execution.py b/camel/toolkits/code_execution.py index dbc845ecbc..bde875cd0a 100644 --- a/camel/toolkits/code_execution.py +++ b/camel/toolkits/code_execution.py @@ -22,6 +22,7 @@ ) from camel.toolkits import FunctionTool from camel.toolkits.base import BaseToolkit +from camel.utils import with_timeout class CodeExecutionToolkit(BaseToolkit): @@ -35,8 +36,8 @@ class CodeExecutionToolkit(BaseToolkit): by `eval()` without any security check. (default: :obj:`False`) import_white_list ( Optional[List[str]]): A list of allowed imports. (default: :obj:`None`) - require_confirm (bool): Whether to require confirmation before executing code. - (default: :obj:`False`) + require_confirm (bool): Whether to require confirmation before + executing code. (default: :obj:`False`) """ def __init__( @@ -102,12 +103,30 @@ def execute_code(self, code: str) -> str: str: The text output from the Code Interpreter tool call. """ output = self.interpreter.run(code, "python") - # ruff: noqa: E501 - content = f"Executed the code below:\n```py\n{code}\n```\n> Executed Results:\n{output}" + + # Format the output content + content = ( + f"Executed the code below:\n```py\n{code}\n```\n" + f"> Executed Results:\n{output}" + ) + if self.verbose: print(content) + return content + def execute_code_with_timeout(self, code: str, timeout: float = 10) -> str: + r"""Execute a given code snippet with a timeout. + + Args: + code (str): The input code to the Code Interpreter tool call. + timeout (float): The timeout (seconds) for the code execution. + + Returns: + str: The text output from the Code Interpreter tool call. + """ + return with_timeout(timeout)(self.execute_code)(code) + def get_tools(self) -> List[FunctionTool]: r"""Returns a list of FunctionTool objects representing the functions in the toolkit. diff --git a/camel/utils/__init__.py b/camel/utils/__init__.py index 2215d0d731..9525c029e3 100644 --- a/camel/utils/__init__.py +++ b/camel/utils/__init__.py @@ -36,6 +36,7 @@ text_extract_from_web, to_pascal, track_agent, + with_timeout, ) from .constants import Constants from .response_format import get_pydantic_model @@ -82,4 +83,5 @@ "get_pydantic_model", "download_github_subdirectory", "generate_prompt_for_structured_output", + "with_timeout", ] diff --git a/camel/utils/commons.py b/camel/utils/commons.py index a131f41770..215ca80d37 100644 --- a/camel/utils/commons.py +++ b/camel/utils/commons.py @@ -728,3 +728,32 @@ def generate_prompt_for_structured_output( {user_prompt} """ return final_prompt + + +def with_timeout(time: float): + r"""Decorator that adds a timeout to a function. + + Args: + time (float): The timeout in seconds. + + Returns: + The decorated function. + """ + import multiprocessing + + def decorator(func): + def wrapper(*args, **kwargs): + if time > 0: + with multiprocessing.Pool(processes=1) as pool: + future = pool.apply_async(func, args=args, kwds=kwargs) + try: + result = future.get(timeout=time) + except multiprocessing.TimeoutError: + result = "Timed out, execution terminated." + return result + else: + return func(*args, **kwargs) + + return wrapper + + return decorator diff --git a/examples/toolkits/code_execution_toolkit.py b/examples/toolkits/code_execution_toolkit.py index 696727dbb7..9f6a55f481 100644 --- a/examples/toolkits/code_execution_toolkit.py +++ b/examples/toolkits/code_execution_toolkit.py @@ -124,3 +124,24 @@ Weng earned $10.20 for 51 minutes of babysitting. =============================================================================== """ + + +time_consuming_code = ''' +import time +time.sleep(100) +''' + +code_execution_tool_kit = CodeExecutionToolkit( + sandbox='internal_python', import_white_list=['sympy'], unsafe_mode=True +) + +result = code_execution_tool_kit.execute_code_with_timeout( + time_consuming_code, timeout=5 +) + +print(result) +""" +=============================================================================== +Timed out, execution terminated. +=============================================================================== +"""