aiqtoolkit 1.2.0a20250707__py3-none-any.whl → 1.2.0rc1__py3-none-any.whl

aiq/tool/code_execution/local_sandbox/local_sandbox_server.py

@@ -12,16 +12,59 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
+import contextlib
 import logging
 import multiprocessing
+import os
 import resource
-import sys
+from enum import Enum
 from io import StringIO
 
 from flask import Flask
+from flask import Request
+from flask import Response
 from flask import request
+from pydantic import BaseModel
+from pydantic import Field
 
 app = Flask(__name__)
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.WARNING)
+
+
+class CodeExecutionStatus(str, Enum):
+    """
+    Status of code execution.
+    """
+    COMPLETED = "completed"
+    ERROR = "error"
+    TIMEOUT = "timeout"
+
+
+class CodeExecutionResult(BaseModel):
+    """
+    Result of code execution.
+    """
+    process_status: CodeExecutionStatus = Field(default=CodeExecutionStatus.COMPLETED,
+                                                description="Status of the process")
+    stdout: str = Field(description="Standard output of the process")
+    stderr: str = Field(description="Standard error of the process")
+
+
+class CodeExecutionResponse(Response):
+    """
+    Response class that returns a JSON response with the given status code and result.
+    """
+
+    def __init__(self, status_code: int, result: CodeExecutionResult):
+        super().__init__(status=status_code, mimetype="application/json", response=result.model_dump_json())
+
+    @classmethod
+    def with_error(cls, status_code: int, error_message: str) -> 'CodeExecutionResponse':
+        return cls(status_code,
+                   CodeExecutionResult(process_status=CodeExecutionStatus.ERROR, stdout="", stderr=error_message))
 
 
 @app.after_request
@@ -34,50 +77,122 @@ def add_hsts_header(response):
     return response
 
 
-def execute_python(generated_code, timeout):
+def execute_python(generated_code: str, timeout: float) -> CodeExecutionResult:
+    """
+    Execute Python code in a subprocess.
+
+    Args:
+        generated_code: The code to execute
+        timeout: The timeout for the execution
+
+    Returns:
+        CodeExecutionResult object containing the execution result
+    """
+
     # running in a separate process to ensure any kind of crashes are properly handled
     queue = multiprocessing.Queue()
     process = multiprocessing.Process(target=execute_code_subprocess, args=(generated_code, queue))
+
     process.start()
+    # wait until the process finishes or the timeout expires
     process.join(timeout=timeout)
-
-    if process.is_alive():  # didn't finish successfully
+    if process.exitcode is None:
         process.kill()
-        return {"process_status": "timeout", "stdout": "", "stderr": "Timed out\n"}
+        return CodeExecutionResult(process_status=CodeExecutionStatus.TIMEOUT, stdout="", stderr="Timed out\n")
 
     return queue.get()
 
 
 # need to memory-limit to avoid common errors of allocating too much
 # but this has to be done in a subprocess to not crush server itself
-def execute_code_subprocess(generated_code, queue):
-    limit = 1024 * 1024 * 1024 * 10  # 10gb - somehow with a smaller limit the server dies when numpy is used
-    resource.setrlimit(resource.RLIMIT_AS, (limit, limit))
-    resource.setrlimit(resource.RLIMIT_DATA, (limit, limit))
+def execute_code_subprocess(generated_code: str, queue):
+    """
+    Execute code in a subprocess.
+
+    Args:
+        generated_code: The code to execute
+        queue: The queue to put the result in
+    """
+
+    logger.debug("execute_code_subprocess started, PID: %s", os.getpid())
 
-    # this can be overriden inside generated code, so it's not a guaranteed protection
-    sys.stdout = StringIO()
     try:
-        exec(generated_code, {})  # pylint: disable=W0122
-        queue.put(sys.stdout.getvalue())
+        limit = 1024 * 1024 * 1024 * 10  # 10gb - somehow with a smaller limit the server dies when numpy is used
+        resource.setrlimit(resource.RLIMIT_AS, (limit, limit))
+        resource.setrlimit(resource.RLIMIT_DATA, (limit, limit))
     except Exception as e:
-        print(f"Error: {str(e)}")
-        queue.put({"process_status": "error", "stdout": "", "stderr": str(e) + "\n"})
+        logger.error("Failed to set resource limits, PID: %s, error: %s", os.getpid(), e)
+
+    stdout_capture = StringIO()
+    stderr_capture = StringIO()
+    try:
+        with contextlib.redirect_stdout(stdout_capture), contextlib.redirect_stderr(stderr_capture):
+            exec(generated_code, {})  # pylint: disable=W0122
+        logger.debug("execute_code_subprocess finished, PID: %s", os.getpid())
+        queue.put(CodeExecutionResult(stdout=stdout_capture.getvalue(), stderr=stderr_capture.getvalue()))
+    except Exception as e:
+        import traceback
+        with contextlib.redirect_stderr(stderr_capture):
+            traceback.print_exc()
+        logger.debug("execute_code_subprocess failed, PID: %s, error: %s", os.getpid(), e)
+        queue.put(
+            CodeExecutionResult(process_status=CodeExecutionStatus.ERROR,
+                                stdout=stdout_capture.getvalue(),
+                                stderr=stderr_capture.getvalue()))
+
+
+def do_execute(request: Request) -> CodeExecutionResponse:
+    """
+    Main function to handle execution requests.
+
+    Args:
+        request: Request object containing the execution request
+
+    Returns:
+        CodeExecutionResponse object containing the execution result
+    """
+    try:
+        # Check if request has JSON data
+        if not request.is_json:
+            return CodeExecutionResponse.with_error(400, "Request must be JSON")
+
+        # Get JSON data safely
+        json_data = request.get_json(silent=True)
+
+        if json_data is None:
+            return CodeExecutionResponse.with_error(400, "Invalid JSON data")
+
+        # Check for required fields
+        if 'generated_code' not in json_data:
+            return CodeExecutionResponse.with_error(400, "Missing required field: generated_code")
+
+        if 'timeout' not in json_data:
+            return CodeExecutionResponse.with_error(400, "Missing required field: timeout")
+
+        if 'language' not in json_data:
+            return CodeExecutionResponse.with_error(400, "Missing required field: language")
+
+        generated_code: str | None = json_data.get('generated_code', None)
+        assert generated_code is not None
+        timeout: float | None = json_data.get('timeout', None)
+        assert timeout is not None
+        language: str | None = json_data.get('language', None)
+        assert language is not None
+
+        if language != 'python':
+            return CodeExecutionResponse.with_error(400, "Only python execution is supported")
+
+        return CodeExecutionResponse(200, execute_python(generated_code, timeout))
+
+    except Exception as e:
+        return CodeExecutionResponse.with_error(500, f"Server error: {str(e)}")
 
 
 # Main Flask endpoint to handle execution requests
 @app.route("/execute", methods=["POST"])
 def execute():
-    generated_code = request.json['generated_code']
-    timeout = request.json['timeout']
-    language = request.json.get('language', 'python')
-
-    if language == 'python':
-        return execute_python(generated_code, timeout)
-    return {"process_status": "error", "stdout": "", "stderr": "Only python execution is supported"}
+    return do_execute(request)
 
 
 if __name__ == '__main__':
-    log = logging.getLogger('werkzeug')
-    log.setLevel(logging.WARNING)
     app.run(port=6000)

aiq/tool/code_execution/local_sandbox/sandbox.requirements.txt

@@ -1,4 +1,6 @@
 numpy
 pandas
 scipy
-ipython
+ipython
+plotly
+pydantic

aiq/tool/code_execution/local_sandbox/start_local_sandbox.sh

@@ -14,12 +14,37 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Usage: ./start_local_sandbox.sh [SANDBOX_NAME] [OUTPUT_DATA_PATH]
 # NOTE: needs to run from the root of the repo!
 
+DOCKER_COMMAND=${DOCKER_COMMAND:-"docker"}
 SANDBOX_NAME=${1:-'local-sandbox'}
 NUM_THREADS=10
 
+# Get the output_data directory path for mounting
+# Priority: command line argument > environment variable > default path (current directory)
+OUTPUT_DATA_PATH=${2:-${OUTPUT_DATA_PATH:-$(pwd)}}
 
-docker build --tag=${SANDBOX_NAME} --build-arg="UWSGI_PROCESSES=$((${NUM_THREADS} * 10))" --build-arg="UWSGI_CHEAPER=${NUM_THREADS}" -f Dockerfile.sandbox .
+echo "Starting sandbox with container name: ${SANDBOX_NAME}"
+echo "Mounting output_data directory: ${OUTPUT_DATA_PATH}"
 
-docker run --network=host --rm --name=local-sandbox ${SANDBOX_NAME}
+# Verify the path exists before mounting, create if it doesn't
+if [ ! -d "${OUTPUT_DATA_PATH}" ]; then
+    echo "Output data directory does not exist, creating: ${OUTPUT_DATA_PATH}"
+    mkdir -p "${OUTPUT_DATA_PATH}"
+fi
+
+# Check if the Docker image already exists
+if ! ${DOCKER_COMMAND} images ${SANDBOX_NAME} | grep -q "${SANDBOX_NAME}"; then
+    echo "Docker image not found locally. Building ${SANDBOX_NAME}..."
+    ${DOCKER_COMMAND} build --tag=${SANDBOX_NAME} --build-arg="UWSGI_PROCESSES=$((${NUM_THREADS} * 10))" --build-arg="UWSGI_CHEAPER=${NUM_THREADS}" -f Dockerfile.sandbox .
+else
+    echo "Using existing Docker image: ${SANDBOX_NAME}"
+fi
+
+# Mount the output_data directory directly so files created in container appear in the local directory
+${DOCKER_COMMAND} run --rm --name=local-sandbox \
+  --network=host \
+  -v "${OUTPUT_DATA_PATH}:/workspace" \
+  -w /workspace \
+  ${SANDBOX_NAME}

aiq/tool/code_execution/register.py

@@ -46,7 +46,11 @@ async def code_execution_tool(config: CodeExecutionToolConfig, builder: Builder)
     class CodeExecutionInputSchema(BaseModel):
         generated_code: str = Field(description="String containing the code to be executed")
 
-    sandbox = get_sandbox(sandbox_type=config.sandbox_type, uri=config.uri)
+    # Create sandbox without working_directory
+    sandbox_kwargs = {"uri": config.uri}
+
+    sandbox = get_sandbox(sandbox_type=config.sandbox_type, **sandbox_kwargs)
+    logger.info(f"[DEBUG] Created sandbox of type: {config.sandbox_type}")
 
     async def _execute_code(generated_code: str) -> dict:
         logger.info("Executing code in the sandbox at %s", config.uri)
@@ -54,12 +58,12 @@ async def code_execution_tool(config: CodeExecutionToolConfig, builder: Builder)
             output = await sandbox.execute_code(
                 generated_code=generated_code,
                 language="python",
-                timeout=config.timeout,
+                timeout_seconds=config.timeout,
                 max_output_characters=config.max_output_characters,
             )
         except Exception as e:
             logger.exception("Error when executing code in the sandbox, %s", e)
-            return {"process_status": "error", "stdout": "", "stderr": e}
+            return {"process_status": "error", "stdout": "", "stderr": str(e)}
         return output
 
     yield FunctionInfo.from_fn(