feat(agent): Persistent execute_code session

autogpts/autogpt/autogpt/agents/agent.py

@@ -5,9 +5,6 @@
 from datetime import datetime
 from typing import TYPE_CHECKING, Optional
 
-from nbformat.v4 import new_notebook
-from jupyter_client.manager import start_new_kernel
-
 import sentry_sdk
 from pydantic import Field
 
@@ -140,9 +137,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

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,39 +193,34 @@ 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 = []
 
     @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,

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"],

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