nvidia-nat 1.1.0a20251020__py3-none-any.whl

nat/tool/code_execution/test_code_execution_sandbox.py

@@ -0,0 +1,414 @@
+# 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.
+"""
+Test suite for Code Execution Sandbox using pytest.
+
+This module provides comprehensive testing for the code execution sandbox service,
+replacing the original bash script with a more maintainable Python implementation.
+"""
+
+import os
+from typing import Any
+
+import pytest
+import requests
+from requests.exceptions import ConnectionError
+from requests.exceptions import RequestException
+from requests.exceptions import Timeout
+
+
+class TestCodeExecutionSandbox:
+    """Test suite for the Code Execution Sandbox service."""
+
+    @pytest.fixture(scope="class")
+    def sandbox_config(self):
+        """Configuration for sandbox testing."""
+        return {
+            "url": os.environ.get("SANDBOX_URL", "http://127.0.0.1:6000/execute"),
+            "timeout": int(os.environ.get("SANDBOX_TIMEOUT", "30")),
+            "connection_timeout": 5
+        }
+
+    @pytest.fixture(scope="class", autouse=True)
+    def check_sandbox_running(self, sandbox_config):
+        """Check if sandbox server is running before running tests."""
+        try:
+            _ = requests.get(sandbox_config["url"], timeout=sandbox_config["connection_timeout"])
+            print(f"✓ Sandbox server is running at {sandbox_config['url']}")
+        except (ConnectionError, Timeout, RequestException):
+            pytest.skip(
+                f"Sandbox server is not running at {sandbox_config['url']}. "
+                "Please start it with: cd src/nat/tool/code_execution/local_sandbox && ./start_local_sandbox.sh")
+
+    def execute_code(self, sandbox_config: dict[str, Any], code: str, language: str = "python") -> dict[str, Any]:
+        """
+        Execute code in the sandbox and return the response.
+
+        Args:
+            sandbox_config: Configuration dictionary
+            code: Code to execute
+            language: Programming language (default: python)
+
+        Returns:
+            dictionary containing the response from the sandbox
+        """
+        payload = {"generated_code": code, "timeout": sandbox_config["timeout"], "language": language}
+
+        response = requests.post(
+            sandbox_config["url"],
+            json=payload,
+            timeout=sandbox_config["timeout"] + 5  # Add buffer to request timeout
+        )
+
+        # Ensure we got a response
+        response.raise_for_status()
+        return response.json()
+
+    def test_simple_print(self, sandbox_config):
+        """Test simple print statement execution."""
+        code = "print('Hello, World!')"
+        result = self.execute_code(sandbox_config, code)
+
+        assert result["process_status"] == "completed"
+        assert "Hello, World!" in result["stdout"]
+        assert result["stderr"] == ""
+
+    def test_basic_arithmetic(self, sandbox_config):
+        """Test basic arithmetic operations."""
+        code = """
+result = 2 + 3
+print(f'Result: {result}')
+"""
+        result = self.execute_code(sandbox_config, code)
+
+        assert result["process_status"] == "completed"
+        assert "Result: 5" in result["stdout"]
+        assert result["stderr"] == ""
+
+    def test_numpy_operations(self, sandbox_config):
+        """Test numpy dependency availability and operations."""
+        code = """
+import numpy as np
+arr = np.array([1, 2, 3, 4, 5])
+print(f'Array: {arr}')
+print(f'Mean: {np.mean(arr)}')
+"""
+        result = self.execute_code(sandbox_config, code)
+
+        assert result["process_status"] == "completed"
+        assert "Array: [1 2 3 4 5]" in result["stdout"]
+        assert "Mean: 3.0" in result["stdout"]
+        assert result["stderr"] == ""
+
+    def test_pandas_operations(self, sandbox_config):
+        """Test pandas dependency availability and operations."""
+        code = """
+import pandas as pd
+df = pd.DataFrame({'A': [1, 2, 3], 'B': [4, 5, 6]})
+print(df)
+print(f'Sum of column A: {df["A"].sum()}')
+"""
+        result = self.execute_code(sandbox_config, code)
+
+        assert result["process_status"] == "completed"
+        assert "Sum of column A: 6" in result["stdout"]
+        assert result["stderr"] == ""
+
+    def test_plotly_import(self, sandbox_config):
+        """Test plotly dependency availability."""
+        code = """
+import plotly.graph_objects as go
+print('Plotly imported successfully')
+fig = go.Figure()
+fig.add_trace(go.Scatter(x=[1, 2, 3], y=[4, 5, 6]))
+print('Plot created successfully')
+"""
+        result = self.execute_code(sandbox_config, code)
+
+        assert result["process_status"] == "completed"
+        assert "Plotly imported successfully" in result["stdout"]
+        assert "Plot created successfully" in result["stdout"]
+        assert result["stderr"] == ""
+
+    def test_syntax_error_handling(self, sandbox_config):
+        """Test handling of syntax errors."""
+        code = """
+print('Hello World'
+# Missing closing parenthesis
+"""
+        result = self.execute_code(sandbox_config, code)
+
+        assert result["process_status"] == "error"
+        assert "SyntaxError" in result["stderr"] or "SyntaxError" in result["stdout"]
+
+    def test_runtime_error_handling(self, sandbox_config):
+        """Test handling of runtime errors."""
+        code = """
+x = 1 / 0
+print('This should not print')
+"""
+        result = self.execute_code(sandbox_config, code)
+
+        assert result["process_status"] == "error"
+        assert "ZeroDivisionError" in result["stderr"] or "ZeroDivisionError" in result["stdout"]
+
+    def test_import_error_handling(self, sandbox_config):
+        """Test handling of import errors."""
+        code = """
+import nonexistent_module
+print('This should not print')
+"""
+        result = self.execute_code(sandbox_config, code)
+
+        assert result["process_status"] == "error"
+        assert "ModuleNotFoundError" in result["stderr"] or "ImportError" in result["stderr"]
+
+    def test_mixed_output(self, sandbox_config):
+        """Test code that produces both stdout and stderr output."""
+        code = """
+import sys
+print('This goes to stdout')
+print('This goes to stderr', file=sys.stderr)
+print('Back to stdout')
+"""
+        result = self.execute_code(sandbox_config, code)
+
+        assert result["process_status"] == "completed"
+        assert "This goes to stdout" in result["stdout"]
+        assert "Back to stdout" in result["stdout"]
+        assert "This goes to stderr" in result["stderr"]
+
+    def test_long_running_code(self, sandbox_config):
+        """Test code that takes some time to execute but completes within timeout."""
+        code = """
+import time
+for i in range(3):
+    print(f'Iteration {i}')
+    time.sleep(0.5)
+print('Completed')
+"""
+        result = self.execute_code(sandbox_config, code)
+
+        assert result["process_status"] == "completed"
+        assert "Iteration 0" in result["stdout"]
+        assert "Iteration 1" in result["stdout"]
+        assert "Iteration 2" in result["stdout"]
+        assert "Completed" in result["stdout"]
+        assert result["stderr"] == ""
+
+    def test_file_operations(self, sandbox_config):
+        """Test basic file operations in the sandbox."""
+        code = """
+import os
+print(f'Current directory: {os.getcwd()}')
+with open('test_file.txt', 'w') as f:
+    f.write('Hello, World!')
+with open('test_file.txt', 'r') as f:
+    content = f.read()
+print(f'File content: {content}')
+os.remove('test_file.txt')
+print('File operations completed')
+"""
+        result = self.execute_code(sandbox_config, code)
+
+        assert result["process_status"] == "completed"
+        assert "File content: Hello, World!" in result["stdout"]
+        assert "File operations completed" in result["stdout"]
+        assert result["stderr"] == ""
+
+    def test_file_persistence_create(self, sandbox_config):
+        """Test file persistence - create various file types."""
+        code = """
+import os
+import pandas as pd
+import numpy as np
+print('Current directory:', os.getcwd())
+print('Directory contents:', os.listdir('.'))
+
+# Create a test file
+with open('persistence_test.txt', 'w') as f:
+    f.write('Hello from sandbox persistence test!')
+
+# Create a CSV file
+df = pd.DataFrame({'A': [1, 2, 3], 'B': [4, 5, 6]})
+df.to_csv('persistence_test.csv', index=False)
+
+# Create a numpy array file
+arr = np.array([1, 2, 3, 4, 5])
+np.save('persistence_test.npy', arr)
+
+print('Files created:')
+for file in os.listdir('.'):
+    if 'persistence_test' in file:
+        print('  -', file)
+"""
+        result = self.execute_code(sandbox_config, code)
+
+        assert result["process_status"] == "completed"
+        assert "persistence_test.txt" in result["stdout"]
+        assert "persistence_test.csv" in result["stdout"]
+        assert "persistence_test.npy" in result["stdout"]
+        assert result["stderr"] == ""
+
+    def test_file_persistence_read(self, sandbox_config):
+        """Test file persistence - read back created files."""
+        code = """
+import pandas as pd
+import numpy as np
+
+# Read back the files we created
+print('=== Reading persistence_test.txt ===')
+with open('persistence_test.txt', 'r') as f:
+    content = f.read()
+    print(f'Content: {content}')
+
+print('\\n=== Reading persistence_test.csv ===')
+df = pd.read_csv('persistence_test.csv')
+print(df)
+print(f'DataFrame shape: {df.shape}')
+
+print('\\n=== Reading persistence_test.npy ===')
+arr = np.load('persistence_test.npy')
+print(f'Array: {arr}')
+print(f'Array sum: {np.sum(arr)}')
+
+print('\\n=== File persistence test PASSED! ===')
+"""
+        result = self.execute_code(sandbox_config, code)
+
+        assert result["process_status"] == "completed"
+        assert "Content: Hello from sandbox persistence test!" in result["stdout"]
+        assert "DataFrame shape: (3, 2)" in result["stdout"]
+        assert "Array: [1 2 3 4 5]" in result["stdout"]
+        assert "Array sum: 15" in result["stdout"]
+        assert "File persistence test PASSED!" in result["stdout"]
+        assert result["stderr"] == ""
+
+    def test_json_operations(self, sandbox_config):
+        """Test JSON file operations for persistence."""
+        code = """
+import json
+import os
+
+# Create a complex JSON file
+data = {
+    'test_name': 'sandbox_persistence',
+    'timestamp': '2024-07-03',
+    'results': {
+        'numpy_test': True,
+        'pandas_test': True,
+        'file_operations': True
+    },
+    'metrics': [1.5, 2.3, 3.7, 4.1],
+    'metadata': {
+        'working_dir': os.getcwd(),
+        'python_version': '3.x'
+    }
+}
+
+# Save JSON file
+with open('persistence_test.json', 'w') as f:
+    json.dump(data, f, indent=2)
+
+# Read it back
+with open('persistence_test.json', 'r') as f:
+    loaded_data = json.load(f)
+
+print('JSON file created and loaded successfully')
+print(f'Test name: {loaded_data["test_name"]}')
+print(f'Results count: {len(loaded_data["results"])}')
+print(f'Metrics: {loaded_data["metrics"]}')
+print('JSON persistence test completed!')
+"""
+        result = self.execute_code(sandbox_config, code)
+
+        assert result["process_status"] == "completed"
+        assert "JSON file created and loaded successfully" in result["stdout"]
+        assert "Test name: sandbox_persistence" in result["stdout"]
+        assert "Results count: 3" in result["stdout"]
+        assert "JSON persistence test completed!" in result["stdout"]
+        assert result["stderr"] == ""
+
+    def test_missing_generated_code_field(self, sandbox_config):
+        """Test request missing the generated_code field."""
+        payload = {"timeout": 10, "language": "python"}
+
+        response = requests.post(sandbox_config["url"], json=payload)
+
+        # Should return an error status code or error in response
+        assert response.status_code != 200 or "error" in response.json()
+
+    def test_missing_timeout_field(self, sandbox_config):
+        """Test request missing the timeout field."""
+        payload = {"generated_code": "print('test')", "language": "python"}
+
+        response = requests.post(sandbox_config["url"], json=payload)
+
+        # Should return error for missing timeout field
+        result = response.json()
+        assert response.status_code == 400 and result["process_status"] == "error"
+
+    def test_invalid_json(self, sandbox_config):
+        """Test request with invalid JSON."""
+        invalid_json = '{"generated_code": "print("test")", "timeout": 10}'
+
+        response = requests.post(sandbox_config["url"], data=invalid_json, headers={"Content-Type": "application/json"})
+
+        # Should return error for invalid JSON
+        assert response.status_code != 200
+
+    def test_non_json_request(self, sandbox_config):
+        """Test request with non-JSON content."""
+        response = requests.post(sandbox_config["url"], data="This is not JSON", headers={"Content-Type": "text/plain"})
+
+        # Should return error for non-JSON content
+        assert response.status_code != 200
+
+    def test_timeout_too_low(self, sandbox_config):
+        """Test request with timeout too low."""
+        code = """
+import time
+time.sleep(2.0)
+"""
+        payload = {"generated_code": code, "timeout": 1, "language": "python"}
+        response = requests.post(sandbox_config["url"], json=payload)
+        assert response.json()["process_status"] == "timeout"
+        assert response.status_code == 200
+
+
+# Pytest configuration and fixtures for command-line options
+def pytest_addoption(parser):
+    """Add custom command-line options for pytest."""
+    parser.addoption("--sandbox-url",
+                     action="store",
+                     default="http://127.0.0.1:6000/execute",
+                     help="Sandbox URL for testing")
+    parser.addoption("--sandbox-timeout",
+                     action="store",
+                     type=int,
+                     default=30,
+                     help="Timeout in seconds for sandbox operations")
+
+
+@pytest.fixture(scope="session", autouse=True)
+def setup_environment(request):
+    """Setup environment variables from command-line options."""
+    os.environ["SANDBOX_URL"] = request.config.getoption("--sandbox-url", "http://127.0.0.1:6000/execute")
+    os.environ["SANDBOX_TIMEOUT"] = str(request.config.getoption("--sandbox-timeout", 30))
+
+
+if __name__ == "__main__":
+    # Allow running as a script
+    pytest.main([__file__, "-v"])

nat/tool/code_execution/utils.py

@@ -0,0 +1,100 @@
+# 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]

nat/tool/datetime_tools.py

@@ -0,0 +1,82 @@
+# 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.
+
+import datetime
+import zoneinfo
+
+from starlette.datastructures import Headers
+
+from nat.builder.builder import Builder
+from nat.builder.function_info import FunctionInfo
+from nat.cli.register_workflow import register_function
+from nat.data_models.function import FunctionBaseConfig
+from nat.settings.global_settings import GlobalSettings
+
+
+class CurrentTimeToolConfig(FunctionBaseConfig, name="current_datetime"):
+    """
+    Simple tool which returns the current date and time in human readable format with timezone information. By default,
+    the timezone is in Etc/UTC. If the user provides a timezone in the header, we will use it. Timezone will be
+    provided in IANA zone name format. For example, "America/New_York" or "Etc/UTC".
+    """
+    pass
+
+
+def _get_timezone_obj(headers: Headers | None) -> zoneinfo.ZoneInfo | datetime.tzinfo:
+    # Default to UTC
+    timezone_obj = zoneinfo.ZoneInfo("Etc/UTC")
+
+    if headers:
+        # If user has provided a timezone in the header, we will prioritize on using it
+        timezone_header = headers.get("x-timezone")
+        if timezone_header:
+            try:
+                timezone_obj = zoneinfo.ZoneInfo(timezone_header)
+            except Exception:
+                pass
+    else:
+        # Only if a timezone is not in the header, we will determine default timezone based on global settings
+        fallback_tz = GlobalSettings.get().fallback_timezone
+
+        if fallback_tz == "system":
+            # Use the system's local timezone. Avoid requiring external deps.
+            timezone_obj = datetime.datetime.now().astimezone().tzinfo or zoneinfo.ZoneInfo("Etc/UTC")
+
+    return timezone_obj
+
+
+@register_function(config_type=CurrentTimeToolConfig)
+async def current_datetime(_config: CurrentTimeToolConfig, _builder: Builder):
+
+    async def _get_current_time(unused: str) -> str:
+
+        del unused  # Unused parameter to avoid linting error
+
+        from nat.builder.context import Context
+        nat_context = Context.get()
+
+        headers: Headers | None = nat_context.metadata.headers
+
+        timezone_obj = _get_timezone_obj(headers)
+
+        now = datetime.datetime.now(timezone_obj)
+        now_machine_readable = now.strftime("%Y-%m-%d %H:%M:%S %z")
+
+        # Returns the current time in machine readable format with timezone offset.
+        return f"The current time of day is {now_machine_readable}"
+
+    yield FunctionInfo.from_fn(
+        _get_current_time,
+        description="Returns the current date and time in human readable format with timezone information.")