feat: add pre-commit as pylint

.github/workflows/pylint.yaml

@@ -0,0 +1,11 @@
+name: pylint
+
+on: [pull_request, push]
+
+jobs:
+  pre-commit:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v3
+    - uses: actions/setup-python@v3
+    - uses: pre-commit/[email protected]

.pre-commit-config.yaml

@@ -0,0 +1,26 @@
+repos:
+-   repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v2.3.0
+    hooks:
+    -   id: end-of-file-fixer
+    -   id: trailing-whitespace
+    -   id: check-added-large-files
+-   repo: https://github.com/psf/black
+    rev: 23.10.0
+    hooks:
+    -   id: black
+-   repo: https://github.com/pre-commit/mirrors-mypy
+    rev: v1.7.1
+    hooks:
+    -   id: mypy
+        additional_dependencies: [pydantic]
+        exclude: tests
+-  repo: https://github.com/astral-sh/ruff-pre-commit
+   rev: v0.1.7
+   hooks:
+   -    id: ruff
+        args: [--fix]
+-  repo: https://github.com/gitleaks/gitleaks
+   rev: v8.17.0
+   hooks:
+   - id: gitleaks

Makefile

@@ -0,0 +1,25 @@
+PHONY: hello env lint
+
+RED=\033[0;31m
+GREEN=\033[0;32m
+YELLOW=\033[0;33m
+BLUE=\033[0;34m
+MAGENTA=\033[0;35m
+CYAN=\033[0;36m
+RESET=\033[0m
+
+hello:
+	@echo "${MAGENTA}Hello, $$(whoami)!${RESET}"
+	@echo "${GREEN}Current Time:${RESET}\t\t${YELLOW}$$(date)${RESET}"
+	@echo "${GREEN}Working Directory:${RESET}\t${YELLOW}$$(pwd)${RESET}"
+	@echo "${GREEN}Shell:${RESET}\t\t\t${YELLOW}$$(echo $$SHELL)${RESET}"
+	@echo "${GREEN}Terminal:${RESET}\t\t${YELLOW}$$(echo $$TERM)${RESET}"
+
+env:
+	@echo "To activate the Poetry environment, run:"
+	@echo "source $$(poetry env info --path)/bin/activate"
+
+lint:
+	@echo "Running linter..."
+	@source $$(poetry env info --path)/bin/activate && pre-commit run --all-files
+	@echo "Done."

llm_sandbox/session.py

@@ -1,25 +1,86 @@
-from typing import List
+import os
+import docker
+from typing import List, Optional, Union
 
-from .const import SupportedLanguage
+from docker.models.images import Image
+from llm_sandbox.utils import image_exists
+from llm_sandbox.const import SupportedLanguage
 
 
 class SandboxSession:
-    def __init__(self, image: str = None, dockerfile: str = None, lang: str = SupportedLanguage.PYTHON):
-        # raise error if both the image and dockerfile exist since we only need one
+    def __init__(
+        self,
+        image: Optional[str] = None,
+        dockerfile: Optional[str] = None,
+        lang: str = SupportedLanguage.PYTHON,
+        keep_template: bool = False,
+    ):
+        """
+        Create a new sandbox session
+        :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
+        """
         if image and dockerfile:
             raise ValueError("Only one of image or dockerfile should be provided")
 
-        # raise error if neither the image nor dockerfile exist since we need one
         if not image and not dockerfile:
             raise ValueError("Either image or dockerfile should be provided")
 
-        self.lang = lang
+        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.path = None
+        self.keep_template = keep_template
+        self.is_create_template: bool = False
 
     def open(self):
-        raise NotImplementedError
+        warning_str = (
+            "Since the `keep_image` flag is set to True the image and container will not be removed after the session "
+            "ends and remains for future use."
+        )
+        if self.dockerfile:
+            self.path = os.path.dirname(self.dockerfile)
+            self.image, _ = self.client.images.build(
+                path=self.path,
+                dockerfile=os.path.basename(self.dockerfile),
+                tag="sandbox",
+            )
+            self.is_create_template = True
+            f_str = f"Built image {self.image.tags[-1]} from {self.dockerfile}"
+            f_str = f"{f_str}. {warning_str}" if self.keep_template else f_str
+            print(f_str)
+
+        if isinstance(self.image, str):
+            if not image_exists(self.client, self.image):
+                self.image = self.client.images.pull(self.image)
+                self.is_create_template = True
+                f_str = f"Pulled image {self.image.tags[-1]}"
+                f_str = f"{f_str}. {warning_str}" if self.keep_template else f_str
+                print(f_str)
+            else:
+                self.image = self.client.images.get(self.image)
+                print(f"Using image {self.image.tags[-1]}")
+
+        self.container = self.client.containers.create(
+            self.image, name="sandbox", detach=True
+        )
 
     def close(self):
-        raise NotImplementedError
+        if self.container:
+            self.container.remove(force=True)
+            self.container = None
+
+        if self.is_create_template and not self.keep_template:
+            if isinstance(self.image, str):
+                self.client.images.remove(self.image)
+            elif isinstance(self.image, Image):
+                self.image.remove(force=True)
+            else:
+                raise ValueError("Invalid image type")
 
     def run(self, code: str, libraries: List = []):
         raise NotImplementedError
@@ -38,4 +99,9 @@ def __enter__(self):
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        self.close()
+        self.close()
+
+
+if __name__ == "__main__":
+    with SandboxSession(dockerfile="tests/busybox.Dockerfile") as session:
+        session.run("print('Hello, World!')")

llm_sandbox/utils.py

@@ -1,9 +1,7 @@
 import docker
 import docker.errors
 
-from typing import List
 from docker import DockerClient
-from docker.models.images import Image
 
 
 def image_exists(client: DockerClient, image: str) -> bool:

tests/busybox.Dockerfile

@@ -1 +1 @@
-FROM busybox:latest
+FROM busybox:latest