nvidia-nat 1.2.0rc5__py3-none-any.whl

aiq/front_ends/mcp/tool_converter.py

@@ -0,0 +1,242 @@
+# 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 json
+import logging
+from inspect import Parameter
+from inspect import Signature
+
+from mcp.server.fastmcp import FastMCP
+from pydantic import BaseModel
+
+from aiq.builder.function import Function
+from aiq.builder.function_base import FunctionBase
+from aiq.builder.workflow import Workflow
+
+logger = logging.getLogger(__name__)
+
+
+def create_function_wrapper(
+    function_name: str,
+    function: FunctionBase,
+    schema: type[BaseModel],
+    is_workflow: bool = False,
+):
+    """Create a wrapper function that exposes the actual parameters of an AIQ Function as an MCP tool.
+
+    Args:
+        function_name: The name of the function/tool
+        function: The AIQ Function object
+        schema: The input schema of the function
+        is_workflow: Whether the function is a Workflow
+
+    Returns:
+        A wrapper function suitable for registration with MCP
+    """
+    # Check if we're dealing with AIQChatRequest - special case
+    is_chat_request = False
+
+    # Check if the schema name is AIQChatRequest
+    if schema.__name__ == "AIQChatRequest" or (hasattr(schema, "__qualname__")
+                                               and "AIQChatRequest" in schema.__qualname__):
+        is_chat_request = True
+        logger.info("Function %s uses AIQChatRequest - creating simplified interface", function_name)
+
+        # For AIQChatRequest, we'll create a simple wrapper with just a query parameter
+        parameters = [Parameter(
+            name="query",
+            kind=Parameter.KEYWORD_ONLY,
+            default=Parameter.empty,
+            annotation=str,
+        )]
+    else:
+        # Regular case - extract parameter information from the input schema
+        # Extract parameter information from the input schema
+        param_fields = schema.model_fields
+
+        parameters = []
+        for name, field in param_fields.items():
+            # Get the field type and convert to appropriate Python type
+            field_type = field.annotation
+
+            # Add the parameter to our list
+            parameters.append(
+                Parameter(
+                    name=name,
+                    kind=Parameter.KEYWORD_ONLY,
+                    default=Parameter.empty if field.is_required else None,
+                    annotation=field_type,
+                ))
+
+    # Create the function signature WITHOUT the ctx parameter
+    # We'll handle this in the wrapper function internally
+    sig = Signature(parameters=parameters, return_annotation=str)
+
+    # Define the actual wrapper function that accepts ctx but doesn't expose it
+    def create_wrapper():
+
+        async def wrapper_with_ctx(**kwargs):
+            """Internal wrapper that will be called by MCP."""
+            # MCP will add a ctx parameter, extract it
+            ctx = kwargs.get("ctx")
+
+            # Remove ctx if present
+            if "ctx" in kwargs:
+                del kwargs["ctx"]
+
+            # Process the function call
+            if ctx:
+                ctx.info("Calling function %s with args: %s", function_name, json.dumps(kwargs, default=str))
+                await ctx.report_progress(0, 100)
+
+            try:
+                # Special handling for AIQChatRequest
+                if is_chat_request:
+                    from aiq.data_models.api_server import AIQChatRequest
+
+                    # Create a chat request from the query string
+                    query = kwargs.get("query", "")
+                    chat_request = AIQChatRequest.from_string(query)
+
+                    # Special handling for Workflow objects
+                    if is_workflow:
+                        # Workflows have a run method that is an async context manager
+                        # that returns an AIQRunner
+                        async with function.run(chat_request) as runner:
+                            # Get the result from the runner
+                            result = await runner.result(to_type=str)
+                    else:
+                        # Regular functions use ainvoke
+                        result = await function.ainvoke(chat_request, to_type=str)
+                else:
+                    # Regular handling
+                    # Handle complex input schema - if we extracted fields from a nested schema,
+                    # we need to reconstruct the input
+                    if len(schema.model_fields) == 1 and len(parameters) > 1:
+                        # Get the field name from the original schema
+                        field_name = next(iter(schema.model_fields.keys()))
+                        field_type = schema.model_fields[field_name].annotation
+
+                        # If it's a pydantic model, we need to create an instance
+                        if hasattr(field_type, "model_validate"):
+                            # Create the nested object
+                            nested_obj = field_type.model_validate(kwargs)
+                            # Call with the nested object
+                            kwargs = {field_name: nested_obj}
+
+                    # Call the AIQ function with the parameters - special handling for Workflow
+                    if is_workflow:
+                        # For workflow with regular input, we'll assume the first parameter is the input
+                        input_value = list(kwargs.values())[0] if kwargs else ""
+
+                        # Workflows have a run method that is an async context manager
+                        # that returns an AIQRunner
+                        async with function.run(input_value) as runner:
+                            # Get the result from the runner
+                            result = await runner.result(to_type=str)
+                    else:
+                        # Regular function call
+                        result = await function.acall_invoke(**kwargs)
+
+                # Report completion
+                if ctx:
+                    await ctx.report_progress(100, 100)
+
+                # Handle different result types for proper formatting
+                if isinstance(result, str):
+                    return result
+                if isinstance(result, (dict, list)):
+                    return json.dumps(result, default=str)
+                return str(result)
+            except Exception as e:
+                if ctx:
+                    ctx.error("Error calling function %s: %s", function_name, str(e))
+                raise
+
+        return wrapper_with_ctx
+
+    # Create the wrapper function
+    wrapper = create_wrapper()
+
+    # Set the signature on the wrapper function (WITHOUT ctx)
+    wrapper.__signature__ = sig
+    wrapper.__name__ = function_name
+
+    # Return the wrapper with proper signature
+    return wrapper
+
+
+def get_function_description(function: FunctionBase) -> str:
+    """
+    Retrieve a human-readable description for an AIQ function or workflow.
+
+    The description is determined using the following precedence:
+       1. If the function is a Workflow and has a 'description' attribute, use it.
+       2. If the Workflow's config has a 'topic', use it.
+       3. If the Workflow's config has a 'description', use it.
+       4. If the function is a regular Function, use its 'description' attribute.
+
+    Args:
+        function: The AIQ FunctionBase instance (Function or Workflow).
+
+    Returns:
+        The best available description string for the function.
+    """
+    function_description = ""
+
+    if isinstance(function, Workflow):
+        config = function.config
+
+        # Workflow doesn't have a description, but probably should
+        if hasattr(function, "description") and function.description:
+            function_description = function.description
+        # Try to get description from config
+        elif hasattr(config, "description") and config.description:
+            function_description = config.description
+        # Try to get anything that might be a description
+        elif hasattr(config, "topic") and config.topic:
+            function_description = config.topic
+
+    elif isinstance(function, Function):
+        function_description = function.description
+
+    return function_description
+
+
+def register_function_with_mcp(mcp: FastMCP, function_name: str, function: FunctionBase) -> None:
+    """Register an AIQ Function as an MCP tool.
+
+    Args:
+        mcp: The FastMCP instance
+        function_name: The name to register the function under
+        function: The AIQ Function to register
+    """
+    logger.info("Registering function %s with MCP", function_name)
+
+    # Get the input schema from the function
+    input_schema = function.input_schema
+    logger.info("Function %s has input schema: %s", function_name, input_schema)
+
+    # Check if we're dealing with a Workflow
+    is_workflow = isinstance(function, Workflow)
+    if is_workflow:
+        logger.info("Function %s is a Workflow", function_name)
+
+    # Get function description
+    function_description = get_function_description(function)
+
+    # Create and register the wrapper function with MCP
+    wrapper_func = create_function_wrapper(function_name, function, input_schema, is_workflow)
+    mcp.tool(name=function_name, description=function_description)(wrapper_func)

aiq/front_ends/register.py

@@ -0,0 +1,22 @@
+# 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.
+
+# pylint: disable=unused-import
+# flake8: noqa
+# isort:skip_file
+
+from .console import register as console_register
+from .fastapi import register as fastapi_register
+from .mcp import register as mcp_register

aiq/front_ends/simple_base/__init__.py

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

aiq/front_ends/simple_base/simple_front_end_plugin_base.py

@@ -0,0 +1,54 @@
+# 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 logging
+from abc import ABC
+from abc import abstractmethod
+from io import StringIO
+
+import click
+
+from aiq.builder.front_end import FrontEndBase
+from aiq.builder.workflow_builder import WorkflowBuilder
+from aiq.data_models.front_end import FrontEndConfigT
+from aiq.runtime.session import AIQSessionManager
+
+logger = logging.getLogger(__name__)
+
+
+class SimpleFrontEndPluginBase(FrontEndBase[FrontEndConfigT], ABC):
+
+    async def pre_run(self):
+        pass
+
+    async def run(self):
+
+        # Must yield the workflow function otherwise it cleans up
+        async with WorkflowBuilder.from_config(config=self.full_config) as builder:
+
+            if logger.isEnabledFor(logging.INFO):
+                stream = StringIO()
+
+                self.full_config.print_summary(stream=stream)
+
+                click.echo(stream.getvalue())
+
+            workflow = builder.build()
+            session_manager = AIQSessionManager(workflow)
+            await self.run_workflow(session_manager)
+
+    @abstractmethod
+    async def run_workflow(self, session_manager: AIQSessionManager):
+        pass

aiq/llm/aws_bedrock_llm.py

@@ -0,0 +1,57 @@
+# 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 pydantic import AliasChoices
+from pydantic import ConfigDict
+from pydantic import Field
+
+from aiq.builder.builder import Builder
+from aiq.builder.llm import LLMProviderInfo
+from aiq.cli.register_workflow import register_llm_provider
+from aiq.data_models.llm import LLMBaseConfig
+from aiq.data_models.retry_mixin import RetryMixin
+
+
+class AWSBedrockModelConfig(LLMBaseConfig, RetryMixin, name="aws_bedrock"):
+    """An AWS Bedrock llm provider to be used with an LLM client."""
+
+    model_config = ConfigDict(protected_namespaces=())
+
+    # Completion parameters
+    model_name: str = Field(validation_alias=AliasChoices("model_name", "model"),
+                            serialization_alias="model",
+                            description="The model name for the hosted AWS Bedrock.")
+    temperature: float = Field(default=0.0, ge=0.0, le=1.0, description="Sampling temperature in [0, 1].")
+    max_tokens: int | None = Field(default=1024,
+                                   gt=0,
+                                   description="Maximum number of tokens to generate."
+                                   "This field is ONLY required when using AWS Bedrock with Langchain.")
+    context_size: int | None = Field(default=1024,
+                                     gt=0,
+                                     description="Maximum number of tokens to generate."
+                                     "This field is ONLY required when using AWS Bedrock with LlamaIndex.")
+
+    # Client parameters
+    region_name: str | None = Field(default="None", description="AWS region to use.")
+    base_url: str | None = Field(
+        default=None, description="Bedrock endpoint to use. Needed if you don't want to default to us-east-1 endpoint.")
+    credentials_profile_name: str | None = Field(
+        default=None, description="The name of the profile in the ~/.aws/credentials or ~/.aws/config files.")
+
+
+@register_llm_provider(config_type=AWSBedrockModelConfig)
+async def aws_bedrock_model(llm_config: AWSBedrockModelConfig, builder: Builder):
+
+    yield LLMProviderInfo(config=llm_config, description="A AWS Bedrock model for use with an LLM client.")

aiq/llm/nim_llm.py

@@ -0,0 +1,46 @@
+# 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 pydantic import AliasChoices
+from pydantic import ConfigDict
+from pydantic import Field
+from pydantic import PositiveInt
+
+from aiq.builder.builder import Builder
+from aiq.builder.llm import LLMProviderInfo
+from aiq.cli.register_workflow import register_llm_provider
+from aiq.data_models.llm import LLMBaseConfig
+from aiq.data_models.retry_mixin import RetryMixin
+
+
+class NIMModelConfig(LLMBaseConfig, RetryMixin, name="nim"):
+    """An NVIDIA Inference Microservice (NIM) llm provider to be used with an LLM client."""
+
+    model_config = ConfigDict(protected_namespaces=())
+
+    api_key: str | None = Field(default=None, description="NVIDIA API key to interact with hosted NIM.")
+    base_url: str | None = Field(default=None, description="Base url to the hosted NIM.")
+    model_name: str = Field(validation_alias=AliasChoices("model_name", "model"),
+                            serialization_alias="model",
+                            description="The model name for the hosted NIM.")
+    temperature: float = Field(default=0.0, description="Sampling temperature in [0, 1].")
+    top_p: float = Field(default=1.0, description="Top-p for distribution sampling.")
+    max_tokens: PositiveInt = Field(default=300, description="Maximum number of tokens to generate.")
+
+
+@register_llm_provider(config_type=NIMModelConfig)
+async def nim_model(llm_config: NIMModelConfig, builder: Builder):
+
+    yield LLMProviderInfo(config=llm_config, description="A NIM model for use with an LLM client.")

aiq/llm/openai_llm.py

@@ -0,0 +1,46 @@
+# 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 pydantic import AliasChoices
+from pydantic import ConfigDict
+from pydantic import Field
+
+from aiq.builder.builder import Builder
+from aiq.builder.llm import LLMProviderInfo
+from aiq.cli.register_workflow import register_llm_provider
+from aiq.data_models.llm import LLMBaseConfig
+from aiq.data_models.retry_mixin import RetryMixin
+
+
+class OpenAIModelConfig(LLMBaseConfig, RetryMixin, name="openai"):
+    """An OpenAI LLM provider to be used with an LLM client."""
+
+    model_config = ConfigDict(protected_namespaces=(), extra="allow")
+
+    api_key: str | None = Field(default=None, description="OpenAI API key to interact with hosted model.")
+    base_url: str | None = Field(default=None, description="Base url to the hosted model.")
+    model_name: str = Field(validation_alias=AliasChoices("model_name", "model"),
+                            serialization_alias="model",
+                            description="The OpenAI hosted model name.")
+    temperature: float = Field(default=0.0, description="Sampling temperature in [0, 1].")
+    top_p: float = Field(default=1.0, description="Top-p for distribution sampling.")
+    seed: int | None = Field(default=None, description="Random seed to set for generation.")
+    max_retries: int = Field(default=10, description="The max number of retries for the request.")
+
+
+@register_llm_provider(config_type=OpenAIModelConfig)
+async def openai_llm(config: OpenAIModelConfig, builder: Builder):
+
+    yield LLMProviderInfo(config=config, description="An OpenAI model for use with an LLM client.")

aiq/llm/register.py

@@ -0,0 +1,23 @@
+# 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.
+
+# pylint: disable=unused-import
+# flake8: noqa
+# isort:skip_file
+
+# Import any providers which need to be automatically registered here
+from . import aws_bedrock_llm
+from . import nim_llm
+from . import openai_llm

aiq/llm/utils/__init__.py

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

aiq/llm/utils/env_config_value.py

@@ -0,0 +1,94 @@
+# 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 os
+from abc import ABC
+from enum import Enum
+
+
+class EnvConfigValueSource(Enum):
+    ENV_DEFAULT = 1
+    CONSTRUCTOR = 2
+    ENV_OVERRIDE = 3
+
+
+class EnvConfigValue(ABC):
+    """
+    A wrapper for a string used as a configuration value which can be loaded from the system environment or injected via
+    the constructor. This class should be subclassed and the class fields `_ENV_KEY` and `_ENV_KEY_OVERRIDE` can be set
+    to enable environment-loading functionality. Convienience properties are available to check from where the value was
+    loaded.
+    """
+
+    _ENV_KEY: str | None = None
+    _ENV_KEY_OVERRIDE: str | None = None
+    _ALLOW_NONE: bool = False
+
+    def __init__(self, value: str | None = None, use_env: bool = True):
+        """
+        Parameters
+        ----------
+        value : str, optional
+            The value to be contained in the EnvConfigValue. If the value is `None`, an attempt will be made to load it
+            from the environment using `_ENV_KEY`. if the `_ENV_KEY_OVERRIDE` field is not `None`, an attempt will be
+            made to load that environment variable in place of the passed-in value.
+        use_env : bool
+            If False, all environment-loading logic will be bypassed and the passed-in value will be used as-is.
+            defaults to True.
+        """
+
+        self._source = EnvConfigValueSource.CONSTRUCTOR
+
+        if use_env:
+            if value is None and self.__class__._ENV_KEY is not None:
+                value = os.environ.get(self.__class__._ENV_KEY, None)
+                self._source = EnvConfigValueSource.ENV_DEFAULT
+
+            if self.__class__._ENV_KEY_OVERRIDE is not None and self.__class__._ENV_KEY_OVERRIDE in os.environ:
+                value = os.environ[self.__class__._ENV_KEY_OVERRIDE]
+                self._source = EnvConfigValueSource.ENV_OVERRIDE
+
+            if not self.__class__._ALLOW_NONE and value is None:
+
+                message = ("value must not be None, but provided value was None and no environment-based default or "
+                           "override was found.")
+
+                if self.__class__._ENV_KEY is None:
+                    raise ValueError(message)
+
+                raise ValueError(
+                    f"{message} Try passing a value to the constructor, or setting the `{self.__class__._ENV_KEY}` "
+                    "environment variable.")
+
+        else:
+            if not self.__class__._ALLOW_NONE and value is None:
+                raise ValueError("value must not be none")
+
+        assert isinstance(value, str) or value is None
+
+        self._value = value
+        self._use_env = use_env
+
+    @property
+    def source(self) -> EnvConfigValueSource:
+        return self._source
+
+    @property
+    def use_env(self) -> bool:
+        return self._use_env
+
+    @property
+    def value(self) -> str | None:
+        return self._value

aiq/llm/utils/error.py

@@ -0,0 +1,17 @@
+# 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_ERROR_MESSAGE = (
+    "{package} not found. Install it and other additional dependencies by running the following command:\n")

aiq/memory/__init__.py

@@ -0,0 +1,20 @@
+# 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.
+"""
+AIQ Toolkit Memory Module
+
+This package provides foundational classes and interfaces
+for managing text-based memory in AIQ Toolkit's LLM-based agents.
+"""