aiqtoolkit 1.1.0a20250429__py3-none-any.whl

aiq/runtime/session.py

@@ -0,0 +1,116 @@
+# 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 asyncio
+import contextvars
+import typing
+from collections.abc import Awaitable
+from collections.abc import Callable
+from contextlib import asynccontextmanager
+from contextlib import nullcontext
+
+from aiq.builder.context import AIQContext
+from aiq.builder.context import AIQContextState
+from aiq.builder.workflow import Workflow
+from aiq.data_models.config import AIQConfig
+from aiq.data_models.interactive import HumanResponse
+from aiq.data_models.interactive import InteractionPrompt
+
+_T = typing.TypeVar("_T")
+
+
+class UserManagerBase:
+    pass
+
+
+class AIQSessionManager:
+
+    def __init__(self, workflow: Workflow, max_concurrency: int = 8):
+        """
+        The AIQSessionManager class is used to run and manage a user workflow session. It runs and manages the context,
+        and configuration of a workflow with the specified concurrency.
+
+        Parameters
+        ----------
+        workflow : Workflow
+            The workflow to run
+        max_concurrency : int, optional
+            The maximum number of simultaneous workflow invocations, by default 8
+        """
+
+        if (workflow is None):
+            raise ValueError("Workflow cannot be None")
+
+        self._workflow: Workflow = workflow
+
+        self._max_concurrency = max_concurrency
+        self._context_state = AIQContextState.get()
+        self._context = AIQContext(self._context_state)
+
+        # We save the context because Uvicorn spawns a new process
+        # for each request, and we need to restore the context vars
+        self._saved_context = contextvars.copy_context()
+
+        if (max_concurrency > 0):
+            self._semaphore = asyncio.Semaphore(max_concurrency)
+        else:
+            # If max_concurrency is 0, then we don't need to limit the concurrency but we still need a context
+            self._semaphore = nullcontext()
+
+    @property
+    def config(self) -> AIQConfig:
+        return self._workflow.config
+
+    @property
+    def workflow(self) -> Workflow:
+        return self._workflow
+
+    @property
+    def context(self) -> AIQContext:
+        return self._context
+
+    @asynccontextmanager
+    async def session(self,
+                      user_manager=None,
+                      user_input_callback: Callable[[InteractionPrompt], Awaitable[HumanResponse]] = None):
+
+        token_user_input = None
+        if user_input_callback is not None:
+            token_user_input = self._context_state.user_input_callback.set(user_input_callback)
+
+        token_user_manager = None
+        if user_manager is not None:
+            token_user_manager = self._context_state.user_manager.set(user_manager)
+
+        try:
+            yield self
+        finally:
+            if token_user_manager is not None:
+                self._context_state.user_manager.reset(token_user_manager)
+            if token_user_input is not None:
+                self._context_state.user_input_callback.reset(token_user_input)
+
+    @asynccontextmanager
+    async def run(self, message):
+        """
+        Start a workflow run
+        """
+        async with self._semaphore:
+            # Apply the saved context
+            for k, v in self._saved_context.items():
+                k.set(v)
+
+            async with self._workflow.run(message) as runner:
+                yield runner

aiq/settings/global_settings.py

@@ -0,0 +1,318 @@
+# 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 json
+import logging
+import os
+import typing
+from collections.abc import Callable
+from contextlib import contextmanager
+from copy import deepcopy
+
+from platformdirs import user_config_dir
+from pydantic import ConfigDict
+from pydantic import Discriminator
+from pydantic import Tag
+from pydantic import ValidationError
+from pydantic import ValidationInfo
+from pydantic import ValidatorFunctionWrapHandler
+from pydantic import field_validator
+
+from aiq.cli.type_registry import GlobalTypeRegistry
+from aiq.cli.type_registry import RegisteredInfo
+from aiq.data_models.common import HashableBaseModel
+from aiq.data_models.common import TypedBaseModel
+from aiq.data_models.common import TypedBaseModelT
+from aiq.data_models.registry_handler import RegistryHandlerBaseConfig
+
+logger = logging.getLogger(__name__)
+
+
+class Settings(HashableBaseModel):
+
+    model_config = ConfigDict(extra="forbid")
+
+    # Registry Handeler Configuration
+    channels: dict[str, RegistryHandlerBaseConfig] = {}
+
+    _configuration_directory: typing.ClassVar[str]
+    _settings_changed_hooks: typing.ClassVar[list[Callable[[], None]]] = []
+    _settings_changed_hooks_active: bool = True
+
+    @field_validator("channels", mode="wrap")
+    @classmethod
+    def validate_components(cls, value: typing.Any, handler: ValidatorFunctionWrapHandler, info: ValidationInfo):
+
+        try:
+            return handler(value)
+        except ValidationError as err:
+
+            for e in err.errors():
+                if e['type'] == 'union_tag_invalid' and len(e['loc']) > 0:
+                    requested_type = e['loc'][0]
+
+                    if (info.field_name == "channels"):
+                        registered_keys = GlobalTypeRegistry.get().get_registered_registry_handlers()
+                    else:
+                        assert False, f"Unknown field name {info.field_name} in validator"
+
+                    # Check and see if the there are multiple full types which match this short type
+                    matching_keys = [k for k in registered_keys if k.local_name == requested_type]
+
+                    assert len(matching_keys) != 1, "Exact match should have been found. Contact developers"
+
+                    matching_key_names = [x.full_type for x in matching_keys]
+                    registered_key_names = [x.full_type for x in registered_keys]
+
+                    if (len(matching_keys) == 0):
+                        # This is a case where the requested type is not found. Show a helpful message about what is
+                        # available
+                        raise ValueError(
+                            f"Requested {info.field_name} type `{requested_type}` not found. "
+                            "Have you ensured the necessary package has been installed with `uv pip install`?"
+                            "\nAvailable {} names:\n - {}".format(info.field_name,
+                                                                  '\n - '.join(registered_key_names))) from err
+
+                    # This is a case where the requested type is ambiguous.
+                    raise ValueError(f"Requested {info.field_name} type `{requested_type}` is ambiguous. " +
+                                     f"Matched multiple {info.field_name} by their local name: {matching_key_names}. " +
+                                     f"Please use the fully qualified {info.field_name} name." +
+                                     "\nAvailable {} names:\n - {}".format(info.field_name,
+                                                                           '\n - '.join(registered_key_names))) from err
+
+            raise
+
+    @classmethod
+    def rebuild_annotations(cls):
+
+        def compute_annotation(cls: type[TypedBaseModelT], registrations: list[RegisteredInfo[TypedBaseModelT]]):
+
+            while (len(registrations) < 2):
+                registrations.append(RegisteredInfo[TypedBaseModelT](full_type=f"_ignore/{len(registrations)}",
+                                                                     config_type=cls))
+
+            short_names: dict[str, int] = {}
+            type_list: list[tuple[str, type[TypedBaseModelT]]] = []
+
+            # For all keys in the list, split the key by / and increment the count of the last element
+            for key in registrations:
+                short_names[key.local_name] = short_names.get(key.local_name, 0) + 1
+
+                type_list.append((key.full_type, key.config_type))
+
+            # Now loop again and if the short name is unique, then create two entries, for the short and full name
+            for key in registrations:
+
+                if (short_names[key.local_name] == 1):
+                    type_list.append((key.local_name, key.config_type))
+
+            # pylint: disable=consider-alternative-union-syntax
+            return typing.Union[tuple(typing.Annotated[x_type, Tag(x_id)] for x_id, x_type in type_list)]
+
+        RegistryHandlerAnnotation = dict[
+            str,
+            typing.Annotated[compute_annotation(RegistryHandlerBaseConfig,
+                                                GlobalTypeRegistry.get().get_registered_registry_handlers()),
+                             Discriminator(TypedBaseModel.discriminator)]]
+
+        should_rebuild = False
+
+        channels_field = cls.model_fields.get("channels")
+        if channels_field is not None and channels_field.annotation != RegistryHandlerAnnotation:
+            channels_field.annotation = RegistryHandlerAnnotation
+            should_rebuild = True
+
+        if (should_rebuild):
+            cls.model_rebuild(force=True)
+
+    @property
+    def channel_names(self) -> list:
+        return list(self.channels.keys())
+
+    @property
+    def configuration_directory(self) -> str:
+        return self._configuration_directory
+
+    @property
+    def configuration_file(self) -> str:
+        return os.path.join(self.configuration_directory, "config.json")
+
+    @staticmethod
+    def from_file():
+
+        configuration_directory = os.getenv("AIQ_CONFIG_DIR", user_config_dir(appname="aiq"))
+
+        if not os.path.exists(configuration_directory):
+            os.makedirs(configuration_directory, exist_ok=True)
+
+        configuration_file = os.path.join(configuration_directory, "config.json")
+
+        file_path = os.path.join(configuration_directory, "config.json")
+
+        if (not os.path.exists(configuration_file)):
+            loaded_config = {}
+        else:
+            with open(file_path, mode="r", encoding="utf-8") as f:
+                loaded_config = json.load(f)
+
+        settings = Settings(**loaded_config)
+        settings.set_configuration_directory(configuration_directory)
+        return settings
+
+    def set_configuration_directory(self, directory: str, remove: bool = False) -> None:
+        if (remove):
+            if os.path.exists(self.configuration_directory):
+                os.rmdir(self.configuration_directory)
+        self.__class__._configuration_directory = directory
+
+    def reset_configuration_directory(self, remove: bool = False) -> None:
+        if (remove):
+            if os.path.exists(self.configuration_directory):
+                os.rmdir(self.configuration_directory)
+        self._configuration_directory = os.getenv("AIQ_CONFIG_DIR", user_config_dir(appname="aiq"))
+
+    def _save_settings(self) -> None:
+
+        if not os.path.exists(self.configuration_directory):
+            os.mkdir(self.configuration_directory)
+
+        with open(self.configuration_file, mode="w", encoding="utf-8") as f:
+            f.write(self.model_dump_json(indent=4, by_alias=True, serialize_as_any=True))
+
+        self._settings_changed()
+
+    def update_settings(self, config_obj: "dict | Settings"):
+        self._update_settings(config_obj)
+
+    def _update_settings(self, config_obj: "dict | Settings"):
+
+        if isinstance(config_obj, Settings):
+            config_obj = config_obj.model_dump(serialize_as_any=True, by_alias=True)
+
+        self._revalidate(config_dict=config_obj)
+
+        self._save_settings()
+
+    def _revalidate(self, config_dict) -> bool:
+
+        try:
+            validated_data = self.__class__(**config_dict)
+
+            for field in validated_data.model_fields:
+                match field:
+                    case "channels":
+                        self.channels = validated_data.channels
+                    case _:
+                        raise ValueError(f"Encountered invalid model field: {field}")
+
+            return True
+
+        except Exception as e:
+            logger.exception("Unable to validate user settings configuration: %s", e, exc_info=True)
+            return False
+
+    def print_channel_settings(self, channel_type: str | None = None) -> None:
+
+        import yaml
+
+        remote_channels = self.model_dump(serialize_as_any=True, by_alias=True)
+
+        if (not remote_channels or not remote_channels.get("channels")):
+            logger.warning("No configured channels to list.")
+            return
+
+        if (channel_type is not None):
+            filter_channels = []
+            for channel, settings in remote_channels.items():
+                if (settings["type"] != channel_type):
+                    filter_channels.append(channel)
+            for channel in filter_channels:
+                del remote_channels[channel]
+
+        if (remote_channels):
+            logger.info(yaml.dump(remote_channels, allow_unicode=True, default_flow_style=False))
+
+    def override_settings(self, config_file: str) -> "Settings":
+
+        from aiq.utils.io.yaml_tools import yaml_load
+
+        override_settings_dict = yaml_load(config_file)
+
+        settings_dict = self.model_dump()
+        updated_settings = {**override_settings_dict, **settings_dict}
+        self._update_settings(config_obj=updated_settings)
+
+        return self
+
+    def _settings_changed(self):
+
+        if (not self._settings_changed_hooks_active):
+            return
+
+        for hook in self._settings_changed_hooks:
+            hook()
+
+    @contextmanager
+    def pause_settings_changed_hooks(self):
+
+        self._settings_changed_hooks_active = False
+
+        try:
+            yield
+        finally:
+            self._settings_changed_hooks_active = True
+
+            # Ensure that the registration changed hooks are called
+            self._settings_changed()
+
+    def add_settings_changed_hook(self, cb: Callable[[], None]) -> None:
+
+        self._settings_changed_hooks.append(cb)
+
+
+GlobalTypeRegistry.get().add_registration_changed_hook(lambda: Settings.rebuild_annotations())
+
+
+class GlobalSettings:
+
+    _global_settings: Settings | None = None
+
+    @staticmethod
+    def get() -> Settings:
+
+        if (GlobalSettings._global_settings is None):
+            from aiq.runtime.loader import PluginTypes
+            from aiq.runtime.loader import discover_and_register_plugins
+
+            discover_and_register_plugins(PluginTypes.REGISTRY_HANDLER)
+
+            GlobalSettings._global_settings = Settings.from_file()
+
+        return GlobalSettings._global_settings
+
+    @staticmethod
+    @contextmanager
+    def push():
+
+        saved = GlobalSettings.get()
+        settings = deepcopy(saved)
+
+        try:
+            GlobalSettings._global_settings = settings
+
+            yield settings
+        finally:
+            GlobalSettings._global_settings = saved
+            GlobalSettings._global_settings._settings_changed()

aiq/test/.namespace

@@ -0,0 +1 @@
+Note: This is a python namespace package and this directory should remain empty. Do NOT add a `__init__.py` file or any other files to this directory. This file is also needed to ensure the directory exists in git.

aiq/tool/code_execution/code_sandbox.py

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

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