Updates

Author: Josh-XT

README.md

@@ -18,15 +18,13 @@ pip install safeexecute
 
 ## Usage
 
-Set your working directory which is where any potential outputs of the code will be stored. This is also where the code will be executed.
+You can pass an entire message from a langauge model into the `code` field and it will parse out any Python code blocks and execute them.  If anywhere in the `code` says `pip install <package>`, it will install the package in the container before executing the code.
 
 ```python
-import os
 from safeexecute import execute_python_code
 
-working_directory = os.path.join(os.getcwd(), "WORKSPACE")
 code = "print('Hello, World!')"
 
-result = execute_python_code(code=code, working_directory=working_directory)
+result = execute_python_code(code=code)
 print(result)
 ```

safeexecute/__init__.py

@@ -18,15 +18,16 @@ def install_docker_image():
     return client
 
 
-def execute_python_code(code: str, working_directory: str) -> str:
-    # Create working directory if it doesn't exist
+def execute_python_code(code: str, working_directory: str = None) -> str:
+    if working_directory is None:
+        working_directory = os.path.join(os.getcwd(), "WORKSPACE")
     if not os.path.exists(working_directory):
         os.makedirs(working_directory)
     # Check if there are any package requirements in the code to install
     package_requirements = re.findall(r"pip install (.*)", code)
+    # Strip out python code blocks if they exist in the code
     if "```python" in code:
         code = code.split("```python")[1].split("```")[0]
-    # Create a temporary Python file in the WORKSPACE directory
     temp_file = os.path.join(working_directory, "temp.py")
     with open(temp_file, "w") as f:
         f.write(code)

setup.py

@@ -8,7 +8,7 @@
 
 setup(
     name="safeexecute",
-    version="0.0.5",
+    version="0.0.6",
     description="Safe way to execute Python code with containerization.",
     long_description=long_description,
     long_description_content_type="text/markdown",