feat: test python interactive

examples/python_interactive.py

@@ -0,0 +1,5 @@
+from llm_sandbox import PythonInteractiveSandboxSession
+
+
+with PythonInteractiveSandboxSession(verbose=True, keep_template=True) as session:
+    out = session.run_cell("print('Hello, World!')")

llm_sandbox/__init__.py

@@ -1 +1 @@
-from .session import SandboxSession  # noqa: F401
+from .session import SandboxSession, PythonInteractiveSandboxSession  # noqa: F401

llm_sandbox/docker.py

@@ -1,7 +1,9 @@
 import io
 import os
+import time
 import docker
 import tarfile
+import threading
 from typing import List, Optional, Union
 
 from docker.models.images import Image
@@ -30,6 +32,10 @@ def __init__(
         lang: str = SupportedLanguage.PYTHON,
         keep_template: bool = False,
         verbose: bool = False,
+        network_disabled: bool = False,
+        network_mode: Optional[str] = "bridge",
+        remove: bool = True,
+        read_only: bool = False,
     ):
         """
         Create a new sandbox session
@@ -70,6 +76,10 @@ def __init__(
         self.keep_template = keep_template
         self.is_create_template: bool = False
         self.verbose = verbose
+        self.network_disabled = network_disabled
+        self.network_mode = network_mode
+        self.remove = remove
+        self.read_only = read_only
 
     def open(self):
         warning_str = (
@@ -104,7 +114,15 @@ def open(self):
                 if self.verbose:
                     print(f"Using image {self.image.tags[-1]}")
 
-        self.container = self.client.containers.run(self.image, detach=True, tty=True)
+        self.container = self.client.containers.run(
+            self.image,
+            detach=True,
+            tty=True,
+            remove=self.remove,
+            network_disabled=self.network_disabled,
+            network_mode=self.network_mode,
+            read_only=self.read_only
+        )
 
     def close(self):
         if self.container:
@@ -260,3 +278,92 @@ def execute_command(
                 print(chunk_str, end="")
 
         return ConsoleOutput(output)
+
+
+import threading
+import time
+
+
+class PythonInteractiveSandboxDockerSession(SandboxDockerSession):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.python_interpreter = None
+        self.interpreter_thread = None
+        self.interpreter_output = ""
+        self.input_pipe = None
+
+    def open(self):
+        """Open a persistent Docker session with a running Python interpreter."""
+        if not self.container:
+            super().open()
+            print("Checking if Python interpreter is already running..")
+
+            # Start the Python interpreter
+            exec_result = self.container.exec_run(
+                cmd="python -i", tty=True, stdin=True, stdout=True, stderr=True, stream=True
+            )
+            self.input_pipe = exec_result.input
+            self.python_interpreter = exec_result.output
+
+            # Start the interpreter in a separate thread to handle output
+            self.interpreter_thread = threading.Thread(target=self._read_interpreter_output)
+            self.interpreter_thread.start()
+
+            if self.verbose:
+                print(f"Python interpreter started in the container {self.container.short_id}")
+        else:
+            if self.verbose:
+                print("Session is already open. Skipping..")
+
+    def _read_interpreter_output(self):
+        """Helper function to read the Python interpreter output."""
+        for chunk in self.python_interpreter:
+            self.interpreter_output += chunk.decode("utf-8")
+            time.sleep(0.1)  # Simulate waiting for more output
+
+    def close(self):
+        """Close the Docker session and stop the Python interpreter."""
+        if self.input_pipe:
+            # Send an exit command to the interpreter
+            self.input_pipe.write(b"exit()\n")
+            self.input_pipe.flush()
+
+            if self.verbose:
+                print("Python interpreter stopped")
+
+        # Wait for the interpreter thread to stop
+        if self.interpreter_thread and self.interpreter_thread.is_alive():
+            self.interpreter_thread.join()
+
+        super().close()
+
+    def run_cell(self, code: str) -> ConsoleOutput:
+        """
+        Run a cell in the Python interpreter.
+        :param code: Python code to execute
+        :return: Output of the executed code
+        """
+        if not self.container:
+            raise RuntimeError(
+                "Session is not open. Please call open() method before running code."
+            )
+
+        if self.verbose:
+            print(f"Running cell: {code}")
+
+        # Write the code to the interpreter
+        if self.input_pipe:
+            self.input_pipe.write(code.encode("utf-8") + b"\n")
+            self.input_pipe.flush()
+
+        # Allow some time for execution
+        time.sleep(1)
+
+        # Read the output from the interpreter
+        output = self.interpreter_output
+        self.interpreter_output = ""  # Reset the output buffer
+
+        if self.verbose:
+            print(f"Output: {output}")
+
+        return ConsoleOutput(output)

llm_sandbox/session.py

@@ -2,7 +2,7 @@
 from typing import Optional, Union
 from kubernetes import client as k8s_client
 from llm_sandbox.const import SupportedLanguage
-from llm_sandbox.docker import SandboxDockerSession
+from llm_sandbox.docker import SandboxDockerSession, PythonInteractiveSandboxDockerSession
 from llm_sandbox.kubernetes import SandboxKubernetesSession
 
 
@@ -17,6 +17,7 @@ def __new__(
         verbose: bool = False,
         use_kubernetes: bool = False,
         kube_namespace: Optional[str] = "default",
+        **kwargs,
     ):
         """
         Create a new sandbox session
@@ -48,3 +49,47 @@ def __new__(
             keep_template=keep_template,
             verbose=verbose,
         )
+
+
+class PythonInteractiveSandboxSession:
+    def __new__(
+        cls,
+        client: Union[docker.DockerClient, k8s_client.CoreV1Api] = None,
+        image: Optional[str] = None,
+        dockerfile: Optional[str] = None,
+        lang: str = SupportedLanguage.PYTHON,
+        keep_template: bool = False,
+        verbose: bool = False,
+        use_kubernetes: bool = False,
+        kube_namespace: Optional[str] = "default",
+    ):
+        """
+        Create a new sandbox session
+        :param client: Either Docker or Kubernetes client, if not provided, a new client will be created based on local context
+        :param image: Docker image to use
+        :param dockerfile: Path to the Dockerfile, if image is not provided
+        :param lang: Language of the code
+        :param keep_template: if True, the image and container will not be removed after the session ends
+        :param verbose: if True, print messages (default is True)
+        :param use_kubernetes: if True, use Kubernetes instead of Docker (default is False)
+        :param kube_namespace: Kubernetes namespace to use (only if 'use_kubernetes' is True), default is 'default'
+        """
+        if use_kubernetes:
+            return SandboxKubernetesSession(
+                client=client,
+                image=image,
+                dockerfile=dockerfile,
+                lang=lang,
+                keep_template=keep_template,
+                verbose=verbose,
+                kube_namespace=kube_namespace,
+            )
+
+        return PythonInteractiveSandboxDockerSession(
+            client=client,
+            image=image,
+            dockerfile=dockerfile,
+            lang=lang,
+            keep_template=keep_template,
+            verbose=verbose,
+        )