aiqtoolkit 1.2.0rc4__py3-none-any.whl → 1.2rc9__py3-none-any.whl

aiq/runtime/session.py

@@ -1,158 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import asyncio
-import contextvars
-import typing
-from collections.abc import Awaitable
-from collections.abc import Callable
-from contextlib import asynccontextmanager
-from contextlib import nullcontext
-
-from starlette.requests import HTTPConnection
-
-from aiq.builder.context import AIQContext
-from aiq.builder.context import AIQContextState
-from aiq.builder.workflow import Workflow
-from aiq.data_models.authentication import AuthenticatedContext
-from aiq.data_models.authentication import AuthFlowType
-from aiq.data_models.authentication import AuthProviderBaseConfig
-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,
-                      request: HTTPConnection | None = None,
-                      conversation_id: str | None = None,
-                      user_input_callback: Callable[[InteractionPrompt], Awaitable[HumanResponse]] = None,
-                      user_authentication_callback: Callable[[AuthProviderBaseConfig, AuthFlowType],
-                                                             Awaitable[AuthenticatedContext | None]] = 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)
-
-        token_user_authentication = None
-        if user_authentication_callback is not None:
-            token_user_authentication = self._context_state.user_auth_callback.set(user_authentication_callback)
-
-        if conversation_id is not None and request is None:
-            self._context_state.conversation_id.set(conversation_id)
-
-        self.set_metadata_from_http_request(request)
-
-        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)
-            if token_user_authentication is not None:
-                self._context_state.user_auth_callback.reset(token_user_authentication)
-
-    @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
-
-    def set_metadata_from_http_request(self, request: HTTPConnection | None) -> None:
-        """
-        Extracts and sets user metadata request attributes from a HTTP request.
-        If request is None, no attributes are set.
-        """
-        if request is None:
-            return
-
-        self._context.metadata._request.method = getattr(request, "method", None)
-        self._context.metadata._request.url_path = request.url.path
-        self._context.metadata._request.url_port = request.url.port
-        self._context.metadata._request.url_scheme = request.url.scheme
-        self._context.metadata._request.headers = request.headers
-        self._context.metadata._request.query_params = request.query_params
-        self._context.metadata._request.path_params = request.path_params
-        self._context.metadata._request.client_host = request.client.host
-        self._context.metadata._request.client_port = request.client.port
-        self._context.metadata._request.cookies = request.cookies
-
-        if request.headers.get("conversation-id"):
-            self._context_state.conversation_id.set(request.headers["conversation-id"])

aiq/runtime/user_metadata.py

@@ -1,130 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from starlette.datastructures import Headers
-from starlette.datastructures import QueryParams
-
-from aiq.data_models.api_server import Request
-
-
-class RequestAttributes:
-    """
-    The RequestAttributes class is responsible for managing user http and webscoket session
-    metadata. It provides a way to store and expose session attributes to workflow tools.
-    """
-
-    def __init__(self) -> None:
-        self._request: Request = Request()
-
-    @property
-    def method(self) -> str | None:
-        """
-        This property retrieves the HTTP method of the request.
-        It can be GET, POST, PUT, DELETE, etc.
-
-        Returns:
-            str | None
-        """
-        return self._request.method
-
-    @property
-    def url_path(self) -> str | None:
-        """
-        This property retrieves the path from the URL of the request.
-
-        Returns:
-            str | None
-        """
-        return self._request.url_path
-
-    @property
-    def url_port(self) -> int | None:
-        """
-        This property retrieves the port number from the URL of the request.
-
-        Returns:
-            int | None
-        """
-        return self._request.url_port
-
-    @property
-    def url_scheme(self) -> str | None:
-        """
-        This property retrieves the scheme from the URL of the request.
-
-        Returns:
-            str | None
-        """
-        return self._request.url_scheme
-
-    @property
-    def headers(self) -> Headers | None:
-        """
-        This property retrieves the headers from the request stored in a dictionary-like object.
-
-        Returns:
-            Headers | None
-        """
-        return self._request.headers
-
-    @property
-    def query_params(self) -> QueryParams | None:
-        """
-        This property retrieves the query parameters from the request stored in a dictionary-like object.
-
-        Returns:
-            QueryParams | None
-        """
-        return self._request.query_params
-
-    @property
-    def path_params(self) -> dict[str, str] | None:
-        """
-        This property retrieves the path parameters from the request stored in a dictionary-like object.
-
-        Returns:
-            dict[str, str] | None
-        """
-        return self._request.path_params
-
-    @property
-    def client_host(self) -> str | None:
-        """
-        This property retrieves the clients remote hostname or IP address.
-
-        Returns:
-            str | None
-        """
-        return self._request.client_host
-
-    @property
-    def client_port(self) -> int | None:
-        """
-        This property retrieves the clients remote port number from which the client is connecting to.
-
-        Returns:
-            int | None
-        """
-        return self._request.client_port
-
-    @property
-    def cookies(self) -> dict[str, str] | None:
-        """
-        This property retrieves the cookies from the request stored in a dictionary-like object.
-
-        Returns:
-            dict[str, str] | None
-        """
-        return self._request.cookies

aiq/settings/global_settings.py

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

@@ -1 +0,0 @@
-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/chat_completion.py

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