aiqtoolkit 1.2.0a20250706__py3-none-any.whl → 1.2.0a20250730__py3-none-any.whl

aiq/tool/chat_completion.py

@@ -0,0 +1,74 @@
+# 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.
+"""
+Simple Completion Function for AIQ Toolkit
+
+This module provides a simple completion function that can handle
+natural language queries and perform basic text completion tasks.
+"""
+
+from pydantic import Field
+
+from aiq.builder.builder import Builder
+from aiq.builder.framework_enum import LLMFrameworkEnum
+from aiq.cli.register_workflow import register_function
+from aiq.data_models.component_ref import LLMRef
+from aiq.data_models.function import FunctionBaseConfig
+
+
+class ChatCompletionConfig(FunctionBaseConfig, name="chat_completion"):
+    """Configuration for the Chat Completion Function."""
+
+    system_prompt: str = Field(("You are a helpful AI assistant. Provide clear, accurate, and helpful "
+                                "responses to user queries. You can give general advice, recommendations, "
+                                "tips, and engage in conversation. Be helpful and informative."),
+                               description="The system prompt to use for chat completion.")
+
+    llm_name: LLMRef = Field(description="The LLM to use for generating responses.")
+
+
+@register_function(config_type=ChatCompletionConfig)
+async def register_chat_completion(config: ChatCompletionConfig, builder: Builder):
+    """Registers a chat completion function that can handle natural language queries."""
+
+    # Get the LLM from the builder context using the configured LLM reference
+    # Use LangChain framework wrapper since we're using LangChain-based LLM
+    llm = await builder.get_llm(config.llm_name, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
+
+    async def _chat_completion(query: str) -> str:
+        """A simple chat completion function that responds to natural language queries.
+
+        Args:
+            query: The user's natural language query
+
+        Returns:
+            A helpful response to the query
+        """
+        try:
+            # Create a simple prompt with the system message and user query
+            prompt = f"{config.system_prompt}\n\nUser: {query}\n\nAssistant:"
+
+            # Generate response using the LLM
+            response = await llm.ainvoke(prompt)
+
+            return response
+
+        except Exception as e:
+            # Fallback response if LLM call fails
+            return (f"I apologize, but I encountered an error while processing your "
+                    f"query: '{query}'. Please try rephrasing your question or try "
+                    f"again later. Error: {str(e)}")
+
+    yield _chat_completion

aiq/tool/code_execution/README.md

@@ -0,0 +1,152 @@
+<!--
+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.
+-->
+
+# Code Execution Sandbox
+
+A secure, containerized Python code execution environment that allows safe execution of Python code with comprehensive error handling and debugging capabilities.
+
+## Overview
+
+The Code Execution Sandbox provides:
+- **Secure code execution** in isolated Docker containers
+- **Comprehensive error handling** with detailed stdout/stderr capture
+- **Multiple input formats** including raw code, dictionary format, and markdown
+- **Dependency management** with pre-installed libraries
+- **Flexible configuration** with customizable timeouts and output limits
+- **Robust debugging** with extensive logging and error reporting
+
+## Quick Start
+
+### Step 1: Start the Sandbox Server
+
+Navigate to the local sandbox directory and start the server:
+
+```bash
+cd src/aiq/tool/code_execution/local_sandbox
+./start_local_sandbox.sh
+```
+
+The script will:
+- Build the Docker image if it doesn't exist
+- Start the sandbox server on port 6000
+- Mount your working directory for file operations
+
+#### Advanced Usage:
+```bash
+# Custom container name
+./start_local_sandbox.sh my-sandbox
+
+# Custom output directory
+./start_local_sandbox.sh my-sandbox /path/to/output
+
+# Using environment variable
+export OUTPUT_DATA_PATH=/path/to/output
+./start_local_sandbox.sh
+```
+
+### Step 2: Test the Installation
+
+Run the comprehensive test suite to verify everything is working:
+
+```bash
+cd src/aiq/tool/code_execution
+pytest test_code_execution_sandbox.py
+```
+
+Note: a running instance of a local sandbox is required.
+
+## Using the Code Execution Tool
+
+### Basic Usage
+
+The sandbox accepts HTTP POST requests to `http://localhost:6000/execute` with JSON payloads:
+
+```bash
+curl -X POST \
+  -H "Content-Type: application/json" \
+  -d '{
+    "generated_code": "print(\"Hello, World!\")",
+    "timeout": 30,
+    "language": "python"
+  }' \
+  http://localhost:6000/execute
+```
+
+### Supported Input Formats
+
+#### 1. Raw Python Code
+```json
+{
+  "generated_code": "import numpy as np\nprint(np.array([1, 2, 3]))",
+  "timeout": 30,
+  "language": "python"
+}
+```
+
+#### 2. Dictionary Format
+```json
+{
+  "generated_code": "{'generated_code': 'print(\"Hello from dict format\")'}",
+  "timeout": 30,
+  "language": "python"
+}
+```
+
+#### 3. Markdown Code Blocks
+```json
+{
+  "generated_code": "```python\nprint('Hello from markdown')\n```",
+  "timeout": 30,
+  "language": "python"
+}
+```
+
+### Response Format
+
+The sandbox returns JSON responses with the following structure:
+
+```json
+{
+  "process_status": "completed|error|timeout",
+  "stdout": "Standard output content",
+  "stderr": "Standard error content"
+}
+```
+
+## Configuration Options
+
+### Sandbox Configuration
+
+- **URI**: Default `http://127.0.0.1:6000`
+- **Timeout**: Default 10 seconds (configurable)
+- **Max Output Characters**: Default 1000 characters
+- **Memory Limit**: 10GB (configurable in Docker)
+- **Working Directory**: Mounted volume for file operations
+
+### Environment Variables
+
+- `OUTPUT_DATA_PATH`: Custom path for file operations
+- `SANDBOX_HOST`: Custom sandbox host
+- `SANDBOX_PORT`: Custom sandbox port
+
+## Security Considerations
+
+- **Isolated execution**: All code runs in Docker containers
+- **Resource limits**: Memory and CPU limits prevent resource exhaustion
+- **Network isolation**: Containers have limited network access
+- **File system isolation**: Mounted volumes provide controlled file access
+- **Process isolation**: Each execution runs in a separate process

aiq/tool/code_execution/code_sandbox.py

@@ -15,11 +15,16 @@
 import abc
 import json
 import logging
+import textwrap
+from typing import Any
 from urllib.parse import urljoin
 
 import requests
+import requests.adapters
 from pydantic import HttpUrl
 
+from aiq.utils.type_utils import override
+
 logger = logging.getLogger(__file__)
 
 
@@ -43,18 +48,18 @@ class Sandbox(abc.ABC):
         *,
         uri: HttpUrl,
     ):
-        self.url = self._get_execute_url(uri)
+        self.url: str = 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
+        self.http_session: requests.Session = session
 
-    def _send_request(self, request, timeout):
+    def _send_request(self, request: dict[str, Any], timeout_seconds: float) -> dict[str, str]:
         output = self.http_session.post(
             url=self.url,
             data=json.dumps(request),
-            timeout=timeout,
+            timeout=timeout_seconds,
             headers={"Content-Type": "application/json"},
         )
         # retrying 502 errors
@@ -64,104 +69,180 @@ class Sandbox(abc.ABC):
         return self._parse_request_output(output)
 
     @abc.abstractmethod
-    def _parse_request_output(self, output):
+    def _parse_request_output(self, output: requests.Response) -> dict[str, str]:
         pass
 
     @abc.abstractmethod
-    def _get_execute_url(self, uri):
+    def _get_execute_url(self, uri: HttpUrl) -> str:
         pass
 
     @abc.abstractmethod
-    def _prepare_request(self, generated_code, timeout):
+    def _prepare_request(self, generated_code: str, timeout_seconds: float) -> dict[str, Any]:
         pass
 
     async def execute_code(
         self,
         generated_code: str,
-        timeout: float = 10.0,
+        timeout_seconds: float = 10.0,
         language: str = "python",
         max_output_characters: int = 1000,
-    ) -> tuple[dict, str]:
+    ) -> dict[str, 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)
+        if language != "python":
+            raise ValueError(f"Language {language} not supported")
+
+        generated_code = generated_code.strip().strip("`")
+        code_to_execute = textwrap.dedent("""
+            import traceback
+            import json
+            import os
+            import warnings
+            import contextlib
+            import io
+            warnings.filterwarnings('ignore')
+            os.environ['OPENBLAS_NUM_THREADS'] = '16'
+        """).strip()
+
+        # Use json.dumps to properly escape the generated_code instead of repr()
+        escaped_code = json.dumps(generated_code)
+        code_to_execute += textwrap.dedent(f"""
+
+            generated_code = {escaped_code}
+
+            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))
+        """).strip()
+        request = self._prepare_request(code_to_execute, timeout_seconds)
         try:
-            output = self._send_request(request, timeout)
+            return self._send_request(request, timeout_seconds)
         except requests.exceptions.Timeout:
-            output = {"process_status": "timeout", "stdout": "", "stderr": "Timed out\n"}
-        return output
+            return {"process_status": "timeout", "stdout": "", "stderr": "Timed out\n"}
 
 
 class LocalSandbox(Sandbox):
     """Locally hosted sandbox."""
 
-    def _get_execute_url(self, uri):
+    def __init__(self, *, uri: HttpUrl):
+        super().__init__(uri=uri)
+
+    @override
+    def _get_execute_url(self, uri: HttpUrl) -> str:
         return urljoin(str(uri), "execute")
 
-    def _parse_request_output(self, output):
+    @override
+    def _parse_request_output(self, output: requests.Response) -> dict[str, str]:
         try:
-            return output.json()
+            output_json = output.json()
+            assert isinstance(output_json, dict)
+            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'}
+            logger.exception("Error parsing output: %s. %s", output.text, e)
+            return {'process_status': 'error', 'stdout': '', 'stderr': f'Unknown error: {e} \"{output.text}\"'}
 
-    def _prepare_request(self, generated_code, timeout, language='python', **kwargs):
-        return {
+    @override
+    def _prepare_request(self,
+                         generated_code: str,
+                         timeout_seconds: float,
+                         language: str = "python",
+                         **kwargs) -> dict[str, Any]:
+        request = {
             "generated_code": generated_code,
-            "timeout": timeout,
+            "timeout": timeout_seconds,
             "language": language,
         }
+        return request
+
+    @override
+    async def execute_code(
+        self,
+        generated_code: str,
+        timeout_seconds: float = 10.0,
+        language: str = "python",
+        max_output_characters: int = 1000,
+    ) -> dict[str, str]:
+        """Override execute_code to bypass the wrapper logic and send user code directly to our server."""
+
+        logger.debug("Raw input generated_code: %s", generated_code)
+
+        # The input appears to be a string representation of a dictionary
+        # We need to parse it and extract the actual code
+        try:
+            # Try to evaluate the string as a Python literal (dictionary)
+            import ast
+            parsed_dict = ast.literal_eval(generated_code)
+            if isinstance(parsed_dict, dict) and 'generated_code' in parsed_dict:
+                actual_code = parsed_dict['generated_code']
+                assert isinstance(actual_code, str)
+                logger.debug("Extracted code from dict: %s...", actual_code[:100])
+            else:
+                # If it's not a dict or doesn't have the expected key, use as-is
+                actual_code = generated_code
+                logger.debug("Using code as-is: %s...", actual_code[:100])
+        except (ValueError, SyntaxError):
+            # If parsing fails, use the input as-is
+            actual_code = generated_code
+            logger.debug("Failed to parse, using as-is: %s...", actual_code[:100])
+
+        # Clean the actual code more carefully to avoid removing backticks that are part of Python code
+        # remove all leading/trailing whitespace -- strip()
+        # remove all leading/trailing backticks -- strip("`")
+        # may potentially start with python, so just trim from the front.
+        POTENTIAL_PREFIXES = ["python"]
+        actual_code = actual_code.strip().strip("`")
+        for prefix in POTENTIAL_PREFIXES:
+            if actual_code.startswith(prefix):
+                actual_code = actual_code[len(prefix):]
+                break
+
+        # Send the user's code directly to our server without any wrapper logic
+        # Our server already handles stdout/stderr capture and error handling
+        request = self._prepare_request(actual_code, timeout_seconds, language)
+        try:
+            return self._send_request(request, timeout_seconds)
+        except requests.exceptions.Timeout:
+            return {"process_status": "timeout", "stdout": "", "stderr": "Timed out\n"}
 
 
 class PistonSandbox(Sandbox):
     """Piston sandbox (https://github.com/engineer-man/piston)"""
 
-    def _get_execute_url(self, uri):
+    @override
+    def _get_execute_url(self, uri: HttpUrl) -> str:
         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'])
+    @override
+    def _parse_request_output(self, output: requests.Response) -> dict[str, str]:
+        output_json = output.json()
+        assert isinstance(output_json, dict)
+        assert 'run' in output_json
+        run_json = output_json['run']
+        assert isinstance(run_json, dict)
+        if run_json["code"] != 0:
+            return {'process_status': "error", 'stdout': run_json['stdout'], 'stderr': run_json['stderr']}
+        return {'process_status': "completed", 'stdout': run_json['stdout'], 'stderr': run_json['stderr']}
 
-    def _prepare_request(self, generated_code: str, timeout, **kwargs):
+    @override
+    def _prepare_request(self, generated_code: str, timeout_seconds: float, **kwargs) -> dict[str, Any]:
         return {
             "language": "py",
             "version": "3.10.0",
@@ -170,19 +251,17 @@ class PistonSandbox(Sandbox):
             }],
             "stdin": "",
             "args": [],
-            "run_timeout": timeout * 1000.0,  # milliseconds
+            "run_timeout": timeout_seconds * 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."""
+    sandboxes = {
+        'local': LocalSandbox,
+        'piston': PistonSandbox,
+    }
     sandbox_class = sandboxes[sandbox_type.lower()]
     return sandbox_class(**kwargs)

aiq/tool/code_execution/local_sandbox/.gitignore

@@ -0,0 +1 @@
+persistence_test.*