aiqtoolkit 1.1.0a20250515__py3-none-any.whl → 1.1.0a20251020__py3-none-any.whl

aiq/tool/code_execution/code_sandbox.py

@@ -1,188 +0,0 @@
-# Copyright (c) 2024-2025, NVIDIA CORPORATION.  All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import abc
-import json
-import logging
-from urllib.parse import urljoin
-
-import requests
-from pydantic import HttpUrl
-
-logger = logging.getLogger(__file__)
-
-
-class Sandbox(abc.ABC):
-    """Code execution sandbox.
-
-    Args:
-        host: Optional[str] = '127.0.0.1' - Host of the sandbox server.
-            Can also be specified through NEMO_SKILLS_SANDBOX_HOST env var.
-        port: Optional[str] = '5000' - Port of the sandbox server.
-            Can also be specified through NEMO_SKILLS_SANDBOX_PORT env var.
-        ssh_server: Optional[str] = None - SSH server for tunneling requests.
-            Useful if server is running on slurm cluster to which there is an ssh access.
-            Can also be specified through NEMO_SKILLS_SSH_SERVER env var.
-        ssh_key_path: Optional[str] = None - Path to the ssh key for tunneling.
-            Can also be specified through NEMO_SKILLS_SSH_KEY_PATH env var.
-    """
-
-    def __init__(
-        self,
-        *,
-        uri: HttpUrl,
-    ):
-        self.url = self._get_execute_url(uri)
-        session = requests.Session()
-        adapter = requests.adapters.HTTPAdapter(pool_maxsize=1500, pool_connections=1500, max_retries=3)
-        session.mount('http://', adapter)
-        session.mount('https://', adapter)
-        self.http_session = session
-
-    def _send_request(self, request, timeout):
-        output = self.http_session.post(
-            url=self.url,
-            data=json.dumps(request),
-            timeout=timeout,
-            headers={"Content-Type": "application/json"},
-        )
-        # retrying 502 errors
-        if output.status_code == 502:
-            raise requests.exceptions.Timeout
-
-        return self._parse_request_output(output)
-
-    @abc.abstractmethod
-    def _parse_request_output(self, output):
-        pass
-
-    @abc.abstractmethod
-    def _get_execute_url(self, uri):
-        pass
-
-    @abc.abstractmethod
-    def _prepare_request(self, generated_code, timeout):
-        pass
-
-    async def execute_code(
-        self,
-        generated_code: str,
-        timeout: float = 10.0,
-        language: str = "python",
-        max_output_characters: int = 1000,
-    ) -> tuple[dict, str]:
-
-        generated_code = generated_code.lstrip().rstrip().lstrip("`").rstrip("`")
-        code_to_execute = """
-import traceback
-import json
-import os
-import warnings
-import contextlib
-import io
-warnings.filterwarnings('ignore')
-os.environ['OPENBLAS_NUM_THREADS'] = '16'
-"""
-
-        code_to_execute += f"""
-\ngenerated_code = {repr(generated_code)}\n
-stdout = io.StringIO()
-stderr = io.StringIO()
-
-with contextlib.redirect_stdout(stdout), contextlib.redirect_stderr(stderr):
-    try:
-        exec(generated_code)
-        status = "completed"
-    except Exception:
-        status = "error"
-        stderr.write(traceback.format_exc())
-stdout = stdout.getvalue()
-stderr = stderr.getvalue()
-if len(stdout) > {max_output_characters}:
-    stdout = stdout[:{max_output_characters}] + "<output cut>"
-if len(stderr) > {max_output_characters}:
-    stderr = stderr[:{max_output_characters}] + "<output cut>"
-if stdout:
-    stdout += "\\n"
-if stderr:
-    stderr += "\\n"
-output = {{"process_status": status, "stdout": stdout, "stderr": stderr}}
-print(json.dumps(output))
-"""
-        request = self._prepare_request(code_to_execute, timeout)
-        try:
-            output = self._send_request(request, timeout)
-        except requests.exceptions.Timeout:
-            output = {"process_status": "timeout", "stdout": "", "stderr": "Timed out\n"}
-        return output
-
-
-class LocalSandbox(Sandbox):
-    """Locally hosted sandbox."""
-
-    def _get_execute_url(self, uri):
-        return urljoin(str(uri), "execute")
-
-    def _parse_request_output(self, output):
-        try:
-            return output.json()
-        except json.JSONDecodeError as e:
-            logger.exception("Error  parsing output: %s. %s", output.text, e)
-            return {'process_status': 'error', 'stdout': '', 'stderr': 'Unknown error'}
-
-    def _prepare_request(self, generated_code, timeout, language='python', **kwargs):
-        return {
-            "generated_code": generated_code,
-            "timeout": timeout,
-            "language": language,
-        }
-
-
-class PistonSandbox(Sandbox):
-    """Piston sandbox (https://github.com/engineer-man/piston)"""
-
-    def _get_execute_url(self, uri):
-        return urljoin(str(uri), "execute")
-
-    def _parse_request_output(self, output):
-        output = output.json()
-        if output['run']['signal'] == "SIGKILL":
-            return {'result': None, 'error_message': 'Unknown error: SIGKILL'}
-        return json.loads(output['run']['output'])
-
-    def _prepare_request(self, generated_code: str, timeout, **kwargs):
-        return {
-            "language": "py",
-            "version": "3.10.0",
-            "files": [{
-                "content": generated_code,
-            }],
-            "stdin": "",
-            "args": [],
-            "run_timeout": timeout * 1000.0,  # milliseconds
-            "compile_memory_limit": -1,
-            "run_memory_limit": -1,
-        }
-
-
-sandboxes = {
-    'local': LocalSandbox,
-    'piston': PistonSandbox,
-}
-
-
-def get_sandbox(sandbox_type: str = "local", **kwargs):
-    """A helper function to make it easier to set sandbox through cmd."""
-    sandbox_class = sandboxes[sandbox_type.lower()]
-    return sandbox_class(**kwargs)

aiq/tool/code_execution/local_sandbox/Dockerfile.sandbox

@@ -1,60 +0,0 @@
-# Copyright (c) 2024-2025, NVIDIA CORPORATION.  All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-# Use the base image with Python 3.10 and Flask
-FROM tiangolo/uwsgi-nginx-flask:python3.10
-
-# Install dependencies required for Lean 4 and other tools
-RUN apt-get update && \
-    apt-get install -y curl git && \
-    curl https://raw.githubusercontent.com/leanprover/elan/master/elan-init.sh -sSf | sh -s -- -y && \
-    /root/.elan/bin/elan toolchain install leanprover/lean4:v4.12.0 && \
-    /root/.elan/bin/elan default leanprover/lean4:v4.12.0 && \
-    /root/.elan/bin/elan self update
-
-# Set environment variables to include Lean and elan/lake in the PATH
-ENV PATH="/root/.elan/bin:$PATH"
-
-# Create Lean project directory and initialize a new Lean project with Mathlib4
-RUN mkdir -p /lean4 && cd /lean4 && \
-    /root/.elan/bin/lake new my_project && \
-    cd my_project && \
-    echo 'leanprover/lean4:v4.12.0' > lean-toolchain && \
-    echo 'require mathlib from git "https://github.com/leanprover-community/mathlib4" @ "v4.12.0"' >> lakefile.lean
-
-# Download and cache Mathlib4 to avoid recompiling, then build the project
-RUN cd /lean4/my_project && \
-    /root/.elan/bin/lake exe cache get && \
-    /root/.elan/bin/lake build
-
-# Set environment variables to include Lean project path
-ENV LEAN_PATH="/lean4/my_project"
-ENV PATH="/lean4/my_project:$PATH"
-
-# Set up application code and install Python dependencies
-COPY sandbox.requirements.txt /app/requirements.txt
-RUN pip install --no-cache-dir -r /app/requirements.txt
-COPY local_sandbox_server.py /app/main.py
-
-# Set the working directory to /app
-WORKDIR /app
-
-# Set Flask app environment variables and ports
-ARG UWSGI_CHEAPER
-ENV UWSGI_CHEAPER=$UWSGI_CHEAPER
-
-ARG UWSGI_PROCESSES
-ENV UWSGI_PROCESSES=$UWSGI_PROCESSES
-
-ENV LISTEN_PORT=6000

aiq/tool/code_execution/local_sandbox/__init__.py

@@ -1,13 +0,0 @@
-# Copyright (c) 2024-2025, NVIDIA CORPORATION.  All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.

aiq/tool/code_execution/local_sandbox/local_sandbox_server.py

@@ -1,83 +0,0 @@
-# Copyright (c) 2024-2025, NVIDIA CORPORATION.  All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import logging
-import multiprocessing
-import resource
-import sys
-from io import StringIO
-
-from flask import Flask
-from flask import request
-
-app = Flask(__name__)
-
-
-@app.after_request
-def add_hsts_header(response):
-    response.headers['Strict-Transport-Security'] = 'max-age=31536000; includeSubDomains'
-    response.headers['X-Content-Type-Options'] = 'nosniff'
-    response.headers['X-Frame-Options'] = 'SAMEORIGIN'
-    response.headers['X-XSS-Protection'] = '1; mode=block'
-
-    return response
-
-
-def execute_python(generated_code, timeout):
-    # 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()
-    process.join(timeout=timeout)
-
-    if process.is_alive():  # didn't finish successfully
-        process.kill()
-        return {"process_status": "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))
-
-    # 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())
-    except Exception as e:
-        print(f"Error: {str(e)}")
-        queue.put({"process_status": "error", "stdout": "", "stderr": str(e) + "\n"})
-
-
-# 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"}
-
-
-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 +0,0 @@
-numpy
-pandas
-scipy
-ipython

aiq/tool/code_execution/local_sandbox/start_local_sandbox.sh

@@ -1,25 +0,0 @@
-#!/bin/bash
-
-# Copyright (c) 2024-2025, NVIDIA CORPORATION.  All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-# NOTE: needs to run from the root of the repo!
-
-SANDBOX_NAME=${1:-'local-sandbox'}
-NUM_THREADS=10
-
-
-docker build --tag=${SANDBOX_NAME} --build-arg="UWSGI_PROCESSES=$((${NUM_THREADS} * 10))" --build-arg="UWSGI_CHEAPER=${NUM_THREADS}" -f Dockerfile.sandbox .
-
-docker run --network=host --rm --name=local-sandbox ${SANDBOX_NAME}

aiq/tool/code_execution/register.py

@@ -1,70 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import logging
-from typing import Literal
-
-from pydantic import BaseModel
-from pydantic import Field
-from pydantic import HttpUrl
-
-from aiq.builder.builder import Builder
-from aiq.builder.function_info import FunctionInfo
-from aiq.cli.register_workflow import register_function
-from aiq.data_models.function import FunctionBaseConfig
-
-logger = logging.getLogger(__name__)
-
-
-class CodeExecutionToolConfig(FunctionBaseConfig, name="code_execution"):
-    """
-    Tool for executing python code in a remotely hosted sandbox environment.
-    """
-    uri: HttpUrl = Field(default=HttpUrl("http://127.0.0.1:6000"),
-                         description="URI for the code execution sandbox server")
-    sandbox_type: Literal["local", "piston"] = Field(default="local", description="The type of code execution sandbox")
-    timeout: float = Field(default=10.0, description="Number of seconds to wait for a code execution request")
-    max_output_characters: int = Field(default=1000, description="Maximum number of characters that can be returned")
-
-
-@register_function(config_type=CodeExecutionToolConfig)
-async def code_execution_tool(config: CodeExecutionToolConfig, builder: Builder):
-    from aiq.tool.code_execution.code_sandbox import get_sandbox
-
-    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)
-
-    async def _execute_code(generated_code: str) -> dict:
-        logger.info("Executing code in the sandbox at %s", config.uri)
-        try:
-            output = await sandbox.execute_code(
-                generated_code=generated_code,
-                language="python",
-                timeout=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 output
-
-    yield FunctionInfo.from_fn(
-        fn=_execute_code,
-        input_schema=CodeExecutionInputSchema,
-        description="""Executes the provied 'generated_code' in a python sandbox environment and returns
-        a dictionary containing stdout, stderr, and the execution status, as well as a session_id. The
-        session_id can be used to append to code that was previously executed.""")

aiq/tool/code_execution/utils.py

@@ -1,100 +0,0 @@
-# Copyright (c) 2024-2025, NVIDIA CORPORATION.  All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import glob
-import logging
-import re
-
-LOG = logging.getLogger(__name__)
-
-
-def format_code_output(execution_dict: dict[str, str],
-                       code_output_begin: str,
-                       code_output_end: str,
-                       code_output_format: str = 'llama'):
-    """Formatting code output to be displayed as an llm expects it."""
-    if code_output_format == 'llama':
-        output = execution_dict["process_status"]
-        if execution_dict['stdout']:
-            output += f"\n[stdout]\n{execution_dict['stdout']}[/stdout]"
-        if execution_dict['stderr']:
-            output += f"\n[stderr]\n{execution_dict['stderr']}[/stderr]"
-        output = f"{code_output_begin}\n\n{output}{code_output_end}\n\n"
-    elif code_output_format == 'qwen':
-        output = ""
-        if execution_dict['stdout']:
-            output += f"{execution_dict['stdout']}"
-        if execution_dict['stderr']:
-            output += f"{execution_dict['stderr']}"
-        if execution_dict['stderr'] and execution_dict['stdout']:
-            LOG.warning("Both stdout and stderr are not empty. This shouldn't normally happen! %s", execution_dict)
-        output = f"{code_output_begin}{output}{code_output_end}"
-    else:
-        raise ValueError(f"Unknown code_output_format: {code_output_format}")
-
-    # wrapping with code output separators
-    return output
-
-
-def _extract_between_separators(generation: str, separators: tuple[str, str], extract_all: bool = False):
-    """Extracting all text between last occurrence of separators[0] and [1].
-
-    If extract_all is True, returning a list with all occurrences of text between separators.
-    """
-    if extract_all:
-        separators = [re.escape(sp) for sp in separators]
-        pattern = f'{separators[0]}(.*?){separators[1]}'
-        return re.findall(pattern, generation, re.DOTALL)
-    return generation.split(separators[0])[-1].split(separators[1])[0]
-
-
-def extract_code_to_execute(generation: str, code_begin: str, code_end: str, extract_all: bool = False):
-    return _extract_between_separators(generation, [code_begin, code_end], extract_all)
-
-
-def extract_code_output(generation: str, code_output_begin: str, code_output_end: str, extract_all: bool = False):
-    return _extract_between_separators(generation, [code_output_begin, code_output_end], extract_all)
-
-
-def unroll_files(input_files):
-    if len(input_files) == 0:
-        raise ValueError("No files found with the given pattern.")
-    total_files = 0
-    for file_pattern in input_files:
-        for file in sorted(glob.glob(file_pattern, recursive=True)):
-            total_files += 1
-            yield file
-    if total_files == 0:
-        raise ValueError("No files found with the given pattern.")
-
-
-def python_doc_to_cmd_help(doc_class, docs_prefix="", arg_prefix=""):
-    """Converts python doc to cmd help format.
-
-    Will color the args and change the format to match what we use in cmd help.
-    """
-    all_args = docs_prefix
-    all_args += doc_class.__doc__.split("Args:")[1].rstrip()
-    # \033[92m ... \033[0m - green in terminal
-    colored_args = ""
-    for line in all_args.split("\n"):
-        if "        " in line and " - " in line:
-            # add colors
-            line = line.replace("        ", "        \033[92m").replace(" - ", "\033[0m - ")
-            # fixing arg format
-            line = line.replace('        \033[92m', f'        \033[92m{arg_prefix}')
-        # fixing indent
-        line = line.replace("        ", "    ").replace("    ", "  ")
-        colored_args += line + '\n'
-    return colored_args[:-1]

aiq/tool/datetime_tools.py

@@ -1,42 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from aiq.builder.builder import Builder
-from aiq.builder.function_info import FunctionInfo
-from aiq.cli.register_workflow import register_function
-from aiq.data_models.function import FunctionBaseConfig
-
-
-class CurrentTimeToolConfig(FunctionBaseConfig, name="current_datetime"):
-    """
-    Simple tool which returns the current date and time in human readable format.
-    """
-    pass
-
-
-@register_function(config_type=CurrentTimeToolConfig)
-async def current_datetime(config: CurrentTimeToolConfig, builder: Builder):
-
-    import datetime
-
-    async def _get_current_time(unused: str) -> str:
-
-        now = datetime.datetime.now()  # Get current time
-        now_human_readable = now.strftime(("%Y-%m-%d %H:%M:%S"))
-
-        return f"The current time of day is {now_human_readable}"  # Format time in H:MM AM/PM format
-
-    yield FunctionInfo.from_fn(_get_current_time,
-                               description="Returns the current date and time in human readable format.")