Update app.py

app.py

@@ -1,66 +1,65 @@
-from fastapi import FastAPI, Form
 import subprocess
 import sys
-import tempfile
 import os
-import importlib.util
-
-CACHE_DIR = "/tmp/.cache_packages"
-
-# Ensure the cache dir exists
-os.makedirs(CACHE_DIR, exist_ok=True)
+import tempfile
+from fastapi import FastAPI
+from pydantic import BaseModel, Field
 
-# Add cache dir to Python path
-if CACHE_DIR not in sys.path:
-    sys.path.insert(0, CACHE_DIR)
+# Defines the structure of the request body
+class ExecutionRequest(BaseModel):
+    package_name: str = Field(..., example="numpy", description="The name of the pip package to install.")
+    script_content: str = Field(..., example="import numpy as np; print(np.array([1, 2, 3]))", description="The Python script content to execute.")
 
 app = FastAPI()
 
-def install_package_if_needed(package_name: str):
-    """Install a package to CACHE_DIR only if not already installed."""
-    if importlib.util.find_spec(package_name) is None:
-        subprocess.check_call([
-            sys.executable, "-m", "pip", "install",
-            "--target", CACHE_DIR, package_name
-        ])
+@app.post("/execute/")
+def install_and_execute(request: ExecutionRequest):
+    """
+    Installs a specified pip package and then executes a provided Python script.
 
-@app.post("/run")
-def run_code(code: str = Form(...)):
-    try:
-        lines = code.strip().split("\n")
-        
-        # Separate pip installs from Python code
-        pip_commands = [line.strip() for line in lines if line.strip().startswith("pip install")]
-        python_code_lines = [line for line in lines if not line.strip().startswith("pip install")]
-
-        # Handle pip installs
-        for cmd in pip_commands:
-            parts = cmd.split()
-            for pkg in parts[2:]:  # Skip 'pip' and 'install'
-                try:
-                    install_package_if_needed(pkg)
-                except subprocess.CalledProcessError as e:
-                    return {"stdout": "", "stderr": f"Error installing: {pkg}\n{str(e)}"}
+    This endpoint first uses pip to install the package named in the `package_name`
+    field. It then writes the `script_content` to a temporary file and
+    executes it using the same Python interpreter that is running the FastAPI
+    application.
 
-        # Save Python code to temp file
-        python_code = "\n".join(python_code_lines)
-        with tempfile.NamedTemporaryFile(delete=False, suffix=".py") as temp_file:
-            temp_file.write(python_code.encode())
-            temp_path = temp_file.name
+    The output from both the pip installation and the script execution is
+    captured and returned in the JSON response. Any errors during the process
+    are also captured and returned.
+    """
+    pip_output = ""
+    script_output = ""
+    error = ""
 
-        # Run Python code
-        result = subprocess.run(
-            [sys.executable, temp_path],
-            capture_output=True,
+    try:
+        # Execute the pip install command
+        pip_output = subprocess.check_output(
+            [sys.executable, "-m", "pip", "install", request.package_name],
+            stderr=subprocess.STDOUT,
             text=True
         )
 
-        os.unlink(temp_path)  # Clean up
+        # Create a temporary file with a .py extension to save the script content
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".py", mode='w') as temp_script:
+            temp_script.write(request.script_content)
+            temp_script_path = temp_script.name
+
+        try:
+            # Execute the Python script
+            script_output = subprocess.check_output(
+                [sys.executable, temp_script_path],
+                stderr=subprocess.STDOUT,
+                text=True
+            )
+        finally:
+            # Clean up the temporary file
+            os.remove(temp_script_path)
 
-        return {
-            "stdout": result.stdout,
-            "stderr": result.stderr
-        }
+    except subprocess.CalledProcessError as e:
+        # If any command fails, capture the output as an error
+        error = e.output
 
-    except Exception as e:
-        return {"error": str(e)}
+    return {
+        "pip_install_output": pip_output,
+        "script_execution_output": script_output,
+        "error_output": error
+    }