feat: commit docker container before remove

Makefile

@@ -1,4 +1,4 @@
-PHONY: hello env lint
+PHONY: hello env lint unittest
 
 RED=\033[0;31m
 GREEN=\033[0;32m
@@ -23,3 +23,8 @@ lint:
 	@echo "Running linter..."
 	@source $$(poetry env info --path)/bin/activate && pre-commit run --all-files
 	@echo "Done."
+
+unittest:
+	@echo "Running tests..."
+	@source $$(poetry env info --path)/bin/activate && poetry run python -m unittest discover -s tests -v
+	@echo "Done."

example/code_runner.py

@@ -0,0 +1,20 @@
+from llm_sandbox import SandboxSession
+
+
+def run_python_code():
+    with SandboxSession(lang="python", keep_template=True, verbose=True) as session:
+        output = session.run("print('Hello, World!')")
+        print(output)
+
+        output = session.run(
+            "import numpy as np\nprint(np.random.rand())", libraries=["numpy"]
+        )
+        print(output)
+
+        session.execute_command("pip install pandas")
+        output = session.run("import pandas as pd\nprint(pd.__version__)")
+        print(output)
+
+
+if __name__ == "__main__":
+    run_python_code()

llm_sandbox/const.py

@@ -11,6 +11,16 @@ class SupportedLanguage:
     RUBY = "ruby"
 
 
+@dataclass
+class DefaultImage:
+    PYTHON = "python:3.9.19-bullseye"
+    JAVA = "openjdk:11.0.12-jdk-bullseye"
+    JAVASCRIPT = "node:16.6.1-bullseye"
+    CPP = "gcc:11.2.0-bullseye"
+    GO = "golang:1.17.0-bullseye"
+    RUBY = "ruby:3.0.2-bullseye"
+
+
 SupportedLanguageValues = [
     v for k, v in SupportedLanguage.__dict__.items() if not k.startswith("__")
 ]

llm_sandbox/session.py

@@ -5,13 +5,14 @@
 from typing import List, Optional, Union
 
 from docker.models.images import Image
+from docker.models.containers import Container
 from llm_sandbox.utils import (
     image_exists,
     get_libraries_installation_command,
     get_code_file_extension,
     get_code_execution_command,
 )
-from llm_sandbox.const import SupportedLanguage, SupportedLanguageValues
+from llm_sandbox.const import SupportedLanguage, SupportedLanguageValues, DefaultImage
 
 
 class SandboxSession:
@@ -34,19 +35,19 @@ def __init__(
         if image and dockerfile:
             raise ValueError("Only one of image or dockerfile should be provided")
 
-        if not image and not dockerfile:
-            raise ValueError("Either image or dockerfile should be provided")
-
         if lang not in SupportedLanguageValues:
             raise ValueError(
                 f"Language {lang} is not supported. Must be one of {SupportedLanguageValues}"
             )
 
+        if not image and not dockerfile:
+            image = DefaultImage.__dict__[lang.upper()]
+
         self.lang: str = lang
         self.client: docker.DockerClient = docker.from_env()
         self.image: Union[Image, str] = image
         self.dockerfile: Optional[str] = dockerfile
-        self.container = None
+        self.container: Optional[Container] = None
         self.path = None
         self.keep_template = keep_template
         self.is_create_template: bool = False
@@ -89,6 +90,9 @@ def open(self):
 
     def close(self):
         if self.container:
+            if isinstance(self.image, Image):
+                self.container.commit(self.image.tags[-1])
+
             self.container.remove(force=True)
             self.container = None
 
@@ -168,7 +172,10 @@ def copy_to_runtime(self, src: str, dest: str):
         tarstream.seek(0)
         self.container.put_archive(os.path.dirname(dest), tarstream)
 
-    def execute_command(self, command: str):
+    def execute_command(self, command: Optional[str]):
+        if not command:
+            raise ValueError("Command cannot be empty")
+
         if not self.container:
             raise RuntimeError(
                 "Session is not open. Please call open() method before executing commands."
@@ -177,12 +184,19 @@ def execute_command(self, command: str):
         if self.verbose:
             print(f"Executing command: {command}")
 
-        exit_code, output = self.container.exec_run(command)
+        _, exec_log = self.container.exec_run(command, stream=True)
+        output = ""
+
         if self.verbose:
-            print(f"Output: {output.decode()}")
-            print(f"Exit code: {exit_code}")
+            print("Output:", end=" ")
+
+        for chunk in exec_log:
+            chunk_str = chunk.decode("utf-8")
+            output += chunk_str
+            if self.verbose:
+                print(chunk_str, end="")
 
-        return exit_code, output.decode()
+        return output
 
     def __enter__(self):
         self.open()

tests/test_session.py

@@ -47,6 +47,7 @@ def test_open_with_image(self):
     def test_close(self):
         mock_container = MagicMock()
         self.session.container = mock_container
+        mock_container.commit.return_values = MagicMock(tags=["python:3.9.19-bullseye"])
 
         self.session.close()
         mock_container.remove.assert_called_once()
@@ -113,13 +114,16 @@ def test_execute_command(self):
         self.session.container = mock_container
 
         command = "echo 'Hello'"
-        mock_container.exec_run.return_value = (0, b"Hello\n")
+        mock_container.exec_run.return_value = (0, iter([b"Hello\n"]))
 
-        exit_code, output = self.session.execute_command(command)
-        mock_container.exec_run.assert_called_with(command)
-        self.assertEqual(exit_code, 0)
+        output = self.session.execute_command(command)
+        mock_container.exec_run.assert_called_with(command, stream=True)
         self.assertEqual(output, "Hello\n")
 
+    def test_execute_empty_command(self):
+        with self.assertRaises(ValueError):
+            self.session.execute_command("")
+
 
 if __name__ == "__main__":
     unittest.main()