aiqtoolkit 1.2.0.dev0__py3-none-any.whl → 1.2.0rc2__py3-none-any.whl

aiq/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/aiq/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"])

aiq/tool/mcp/exceptions.py

@@ -0,0 +1,142 @@
+# 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.
+
+from enum import Enum
+
+
+class MCPErrorCategory(str, Enum):
+    """Categories of MCP errors for structured handling."""
+    CONNECTION = "connection"
+    TIMEOUT = "timeout"
+    SSL = "ssl"
+    AUTHENTICATION = "authentication"
+    TOOL_NOT_FOUND = "tool_not_found"
+    PROTOCOL = "protocol"
+    UNKNOWN = "unknown"
+
+
+class MCPError(Exception):
+    """Base exception for MCP-related errors."""
+
+    def __init__(self,
+                 message: str,
+                 url: str,
+                 category: MCPErrorCategory = MCPErrorCategory.UNKNOWN,
+                 suggestions: list[str] | None = None,
+                 original_exception: Exception | None = None):
+        super().__init__(message)
+        self.url = url
+        self.category = category
+        self.suggestions = suggestions or []
+        self.original_exception = original_exception
+
+
+class MCPConnectionError(MCPError):
+    """Exception for MCP connection failures."""
+
+    def __init__(self, url: str, original_exception: Exception | None = None):
+        super().__init__(f"Unable to connect to MCP server at {url}",
+                         url=url,
+                         category=MCPErrorCategory.CONNECTION,
+                         suggestions=[
+                             "Please ensure the MCP server is running and accessible",
+                             "Check if the URL and port are correct"
+                         ],
+                         original_exception=original_exception)
+
+
+class MCPTimeoutError(MCPError):
+    """Exception for MCP timeout errors."""
+
+    def __init__(self, url: str, original_exception: Exception | None = None):
+        super().__init__(f"Connection timed out to MCP server at {url}",
+                         url=url,
+                         category=MCPErrorCategory.TIMEOUT,
+                         suggestions=[
+                             "The server may be overloaded or network is slow",
+                             "Try again in a moment or check network connectivity"
+                         ],
+                         original_exception=original_exception)
+
+
+class MCPSSLError(MCPError):
+    """Exception for MCP SSL/TLS errors."""
+
+    def __init__(self, url: str, original_exception: Exception | None = None):
+        super().__init__(f"SSL/TLS error connecting to {url}",
+                         url=url,
+                         category=MCPErrorCategory.SSL,
+                         suggestions=[
+                             "Check if the server requires HTTPS or has valid certificates",
+                             "Try using HTTP instead of HTTPS if appropriate"
+                         ],
+                         original_exception=original_exception)
+
+
+class MCPRequestError(MCPError):
+    """Exception for MCP request errors."""
+
+    def __init__(self, url: str, original_exception: Exception | None = None):
+        message = f"Request failed to MCP server at {url}"
+        if original_exception:
+            message += f": {original_exception}"
+
+        super().__init__(message,
+                         url=url,
+                         category=MCPErrorCategory.PROTOCOL,
+                         suggestions=["Check the server URL format and network settings"],
+                         original_exception=original_exception)
+
+
+class MCPToolNotFoundError(MCPError):
+    """Exception for when a specific MCP tool is not found."""
+
+    def __init__(self, tool_name: str, url: str, original_exception: Exception | None = None):
+        super().__init__(f"Tool '{tool_name}' not available at {url}",
+                         url=url,
+                         category=MCPErrorCategory.TOOL_NOT_FOUND,
+                         suggestions=[
+                             "Use 'aiq info mcp --detail' to see available tools",
+                             "Check that the tool name is spelled correctly"
+                         ],
+                         original_exception=original_exception)
+
+
+class MCPAuthenticationError(MCPError):
+    """Exception for MCP authentication failures."""
+
+    def __init__(self, url: str, original_exception: Exception | None = None):
+        super().__init__(f"Authentication failed when connecting to MCP server at {url}",
+                         url=url,
+                         category=MCPErrorCategory.AUTHENTICATION,
+                         suggestions=[
+                             "Check if the server requires authentication credentials",
+                             "Verify that your credentials are correct and not expired"
+                         ],
+                         original_exception=original_exception)
+
+
+class MCPProtocolError(MCPError):
+    """Exception for MCP protocol-related errors."""
+
+    def __init__(self, url: str, message: str = "Protocol error", original_exception: Exception | None = None):
+        super().__init__(f"{message} (MCP server at {url})",
+                         url=url,
+                         category=MCPErrorCategory.PROTOCOL,
+                         suggestions=[
+                             "Check that the MCP server is running and accessible at this URL",
+                             "Verify the server supports the expected MCP protocol version"
+                         ],
+                         original_exception=original_exception)

aiq/tool/mcp/mcp_client.py

@@ -27,6 +27,9 @@ from pydantic import BaseModel
 from pydantic import Field
 from pydantic import create_model
 
+from aiq.tool.mcp.exceptions import MCPToolNotFoundError
+from aiq.utils.exception_handlers.mcp import mcp_exception_handler
+
 logger = logging.getLogger(__name__)
 
 
@@ -45,6 +48,7 @@ def model_from_mcp_schema(name: str, mcp_input_schema: dict) -> type[BaseModel]:
     }
 
     properties = mcp_input_schema.get("properties", {})
+    required_fields = set(mcp_input_schema.get("required", []))
     schema_dict = {}
 
     def _generate_valid_classname(class_name: str):
@@ -63,14 +67,34 @@ def model_from_mcp_schema(name: str, mcp_input_schema: dict) -> type[BaseModel]:
         elif json_type == "array" and "items" in field_properties:
             item_properties = field_properties.get("items", {})
             if item_properties.get("type") == "object":
-                item_type = model_from_mcp_schema(name=field_name, mcp_input_schema=field_properties)
+                item_type = model_from_mcp_schema(name=field_name, mcp_input_schema=item_properties)
             else:
-                item_type = _type_map.get(json_type, Any)
+                item_type = _type_map.get(item_properties.get("type", "string"), Any)
             field_type = list[item_type]
+        elif isinstance(json_type, list):
+            field_type = None
+            for t in json_type:
+                mapped = _type_map.get(t, Any)
+                field_type = mapped if field_type is None else field_type | mapped
+
+            return field_type, Field(
+                default=field_properties.get("default", None if "null" in json_type else ...),
+                description=field_properties.get("description", "")
+            )
         else:
             field_type = _type_map.get(json_type, Any)
 
-        default_value = field_properties.get("default", ...)
+        # Determine the default value based on whether the field is required
+        if field_name in required_fields:
+            # Field is required - use explicit default if provided, otherwise make it required
+            default_value = field_properties.get("default", ...)
+        else:
+            # Field is optional - use explicit default if provided, otherwise None
+            default_value = field_properties.get("default", None)
+            # Make the type optional if no default was provided
+            if "default" not in field_properties:
+                field_type = field_type | None
+
         nullable = field_properties.get("nullable", False)
         description = field_properties.get("description", "")
 
@@ -117,9 +141,16 @@ class MCPBuilder(MCPSSEClient):
         super().__init__(url)
         self._tools = None
 
+    @mcp_exception_handler
     async def get_tools(self):
         """
         Retrieve a dictionary of all tools served by the MCP server.
+
+        Returns:
+            Dict of tool name to MCPToolClient
+
+        Raises:
+            MCPError: If connection or tool retrieval fails
         """
         async with self.connect_to_sse_server() as session:
             response = await session.list_tools()
@@ -129,6 +160,7 @@ class MCPBuilder(MCPSSEClient):
             for tool in response.tools
         }
 
+    @mcp_exception_handler
     async def get_tool(self, tool_name: str) -> MCPToolClient:
         """
         Get an MCP Tool by name.
@@ -139,17 +171,19 @@ class MCPBuilder(MCPSSEClient):
         Returns:
             MCPToolClient for the configured tool.
 
-        Raise:
-            ValueError if no tool is available with that name.
+        Raises:
+            MCPToolNotFoundError: If no tool is available with that name
+            MCPError: If connection fails
         """
         if not self._tools:
             self._tools = await self.get_tools()
 
         tool = self._tools.get(tool_name)
         if not tool:
-            raise ValueError(f"Tool {tool_name} not available at {self.url}")
+            raise MCPToolNotFoundError(tool_name, self.url)
         return tool
 
+    @mcp_exception_handler
     async def call_tool(self, tool_name: str, tool_args: dict | None):
         async with self.connect_to_sse_server() as session:
             result = await session.call_tool(tool_name, tool_args)
@@ -200,6 +234,7 @@ class MCPToolClient(MCPSSEClient):
         """
         self._tool_description = description
 
+    @mcp_exception_handler
     async def acall(self, tool_args: dict) -> str:
         """
         Call the MCP tool with the provided arguments.

aiq/tool/mcp/mcp_tool.py

@@ -29,7 +29,7 @@ logger = logging.getLogger(__name__)
 
 class MCPToolConfig(FunctionBaseConfig, name="mcp_tool_wrapper"):
     """
-    Function which connects to a Model Context Protocol (MCP) server and wraps the selected tool as an AIQ Toolkit
+    Function which connects to a Model Context Protocol (MCP) server and wraps the selected tool as a NeMo Agent toolkit
     function.
     """
     # Add your custom configuration parameters here
@@ -75,7 +75,8 @@ async def mcp_tool(config: MCPToolConfig, builder: Builder):  # pylint: disable=
                 return await tool.acall(args)
 
             _ = tool.input_schema.model_validate(kwargs)
-            return await tool.acall(kwargs)
+            filtered_kwargs = {k: v for k, v in kwargs.items() if v is not None}
+            return await tool.acall(filtered_kwargs)
         except Exception as e:
             if config.return_exception:
                 if tool_input: