From 970c95cfcda01cee8279cd868b6fcefd46bcb7ef Mon Sep 17 00:00:00 2001 From: Mahdi Karami Date: Thu, 2 May 2024 16:12:34 +0330 Subject: [PATCH] feat(agent): Persistent execute_code session --- autogpts/autogpt/autogpt/agents/agent.py | 4 +-- autogpts/autogpt/autogpt/agents/base.py | 27 ++++++++++------- .../autogpt/autogpt/commands/execute_code.py | 27 +++++++++-------- .../tests/integration/test_execute_code.py | 29 ++++++++++++++----- 4 files changed, 54 insertions(+), 33 deletions(-) diff --git a/autogpts/autogpt/autogpt/agents/agent.py b/autogpts/autogpt/autogpt/agents/agent.py index 22ba063d986..83f2395a5f8 100644 --- a/autogpts/autogpt/autogpt/agents/agent.py +++ b/autogpts/autogpt/autogpt/agents/agent.py @@ -85,6 +85,7 @@ class AgentSettings(BaseAgentSettings): history: EpisodicActionHistory[OneShotAgentActionProposal] = Field( default_factory=EpisodicActionHistory[OneShotAgentActionProposal] ) + """(STATE) The action history of the agent.""" @@ -140,9 +141,6 @@ def __init__( """Timestamp the agent was created; only used for structured debug logging.""" self.log_cycle_handler = LogCycleHandler() - self.notebook = new_notebook() - kernel_name = self.notebook.metadata.get('kernelspec', {}).get('name', 'python') - _, self.python_kernel = start_new_kernel(kernel_name=kernel_name) """LogCycleHandler for structured debug logging.""" self.event_history = settings.history diff --git a/autogpts/autogpt/autogpt/agents/base.py b/autogpts/autogpt/autogpt/agents/base.py index cf8e3cac808..ec96ae93c7b 100644 --- a/autogpts/autogpt/autogpt/agents/base.py +++ b/autogpts/autogpt/autogpt/agents/base.py @@ -16,6 +16,8 @@ ) from colorama import Fore +from jupyter_client.manager import start_new_kernel +from nbformat.v4 import new_notebook from pydantic import BaseModel, Field, validator if TYPE_CHECKING: @@ -129,6 +131,16 @@ class BaseAgentSettings(SystemSettings): config: BaseAgentConfiguration = Field(default_factory=BaseAgentConfiguration) """The configuration for this BaseAgent subsystem instance.""" + notebook: Any = None + python_kernel: Any = None + + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.notebook = new_notebook() + self.python_kernel = start_new_kernel( + kernel_name=self.notebook.metadata.get("kernelspec", {}).get("name", "python") + )[1] + class AgentMeta(ABCMeta): def __call__(cls, *args, **kwargs): @@ -181,24 +193,21 @@ def send_token_limit(self) -> int: return self.config.send_token_limit or self.llm.max_tokens * 3 // 4 @abstractmethod - async def propose_action(self) -> BaseAgentActionProposal: - ... + async def propose_action(self) -> BaseAgentActionProposal: ... @abstractmethod async def execute( self, proposal: BaseAgentActionProposal, user_feedback: str = "", - ) -> ActionResult: - ... + ) -> ActionResult: ... @abstractmethod async def do_not_execute( self, denied_proposal: BaseAgentActionProposal, user_feedback: str, - ) -> ActionResult: - ... + ) -> ActionResult: ... def reset_trace(self): self._trace = [] @@ -206,14 +215,12 @@ def reset_trace(self): @overload async def run_pipeline( self, protocol_method: Callable[P, Iterator[T]], *args, retry_limit: int = 3 - ) -> list[T]: - ... + ) -> list[T]: ... @overload async def run_pipeline( self, protocol_method: Callable[P, None], *args, retry_limit: int = 3 - ) -> list[None]: - ... + ) -> list[None]: ... async def run_pipeline( self, diff --git a/autogpts/autogpt/autogpt/commands/execute_code.py b/autogpts/autogpt/autogpt/commands/execute_code.py index 372819029c7..a30e8e6e61c 100644 --- a/autogpts/autogpt/autogpt/commands/execute_code.py +++ b/autogpts/autogpt/autogpt/commands/execute_code.py @@ -1,11 +1,9 @@ """Commands to execute code""" -import logging import os -import queue import shlex +import queue import logging -import traceback import subprocess from pathlib import Path from tempfile import NamedTemporaryFile @@ -134,18 +132,23 @@ def execute_python_code(self, code: str) -> str: str: The STDOUT captured from the code when it ran. """ - tmp_code_file = NamedTemporaryFile( - "w", dir=self.workspace.root, suffix=".py", encoding="utf-8" - ) - tmp_code_file.write(code) - tmp_code_file.flush() - try: - return self.execute_python_file(tmp_code_file.name) + self.state.python_kernel.execute(code) + output_texts = [] + while True: + try: + io_msg = self.state.python_kernel.get_iopub_msg(timeout=1) + if 'name' in io_msg['content'].keys() and io_msg['content']['name'] == 'stdout': + output_texts.append(io_msg['content']['text']) + elif 'traceback' in io_msg['content'].keys(): + raise IOError(f"{io_msg['content']['evalue']}\n{io_msg['content']['traceback']}") + except queue.Empty: + break + + output = ''.join(output_texts) + return output except Exception as e: raise CommandExecutionError(*e.args) - finally: - tmp_code_file.close() @command( ["execute_python_file"], diff --git a/autogpts/autogpt/tests/integration/test_execute_code.py b/autogpts/autogpt/tests/integration/test_execute_code.py index a0425a4a169..f61a6afbe0c 100644 --- a/autogpts/autogpt/tests/integration/test_execute_code.py +++ b/autogpts/autogpt/tests/integration/test_execute_code.py @@ -12,7 +12,11 @@ is_docker_available, we_are_running_in_a_docker_container, ) -from autogpt.utils.exceptions import InvalidArgumentError, OperationNotAllowedError +from autogpt.utils.exceptions import ( + CommandExecutionError, + InvalidArgumentError, + OperationNotAllowedError, +) @pytest.fixture @@ -96,24 +100,33 @@ def test_execute_python_code( result: str = code_executor_component.execute_python_code(random_code) assert result.replace("\r", "") == f"Hello {random_string}!\n" -def test_several_execute_python_code(agent: Agent): - if not (sut.is_docker_available() or sut.we_are_running_in_a_docker_container()): + +def test_several_execute_python_code( + code_executor_component: CodeExecutorComponent, agent: Agent +): + if not (is_docker_available() or we_are_running_in_a_docker_container()): pytest.skip("Docker is not available") - result: str = sut.execute_python_code("a=3\nprint(a)", agent=agent) + result: str = code_executor_component.execute_python_code( + "a=3\nprint(a)" + ) assert result.replace("\r", "") == f"3\n" - result: str = sut.execute_python_code("a+=1\nprint(a)", agent=agent) + result: str = code_executor_component.execute_python_code( + "a+=1\nprint(a)" + ) assert result.replace("\r", "") == f"4\n" -def test_execute_python_code_raise_exception_scenario(agent: Agent): - if not (sut.is_docker_available() or sut.we_are_running_in_a_docker_container()): + +def test_execute_python_code_raise_exception_scenario(code_executor_component: CodeExecutorComponent ,agent: Agent): + if not (is_docker_available() or we_are_running_in_a_docker_container()): pytest.skip("Docker is not available") with pytest.raises( CommandExecutionError, match=r"name 'a' is not defined", ): - sut.execute_python_code("a+=1", agent=agent) + code_executor_component.execute_python_code("a+=1") + def test_execute_python_file_invalid( code_executor_component: CodeExecutorComponent, agent: Agent