aiqtoolkit 1.1.0__py3-none-any.whl

aiq/tool/retriever.py

@@ -0,0 +1,89 @@
+# 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 pydantic import BaseModel
+from pydantic import Field
+
+from aiq.builder.builder import Builder
+from aiq.builder.function_info import FunctionInfo
+from aiq.cli.register_workflow import register_function
+from aiq.data_models.component_ref import RetrieverRef
+from aiq.data_models.function import FunctionBaseConfig
+from aiq.retriever.interface import AIQRetriever
+from aiq.retriever.models import RetrieverError
+from aiq.retriever.models import RetrieverOutput
+
+logger = logging.getLogger(__name__)
+
+
+class AIQRetrieverConfig(FunctionBaseConfig, name="aiq_retriever"):
+    """
+    AIQRetriever tool which provides a common interface for different vectorstores. Its
+    configuration uses clients, which are the vectorstore-specific implementaiton of the retriever interface.
+    """
+    retriever: RetrieverRef = Field(description="The retriever instance name from the workflow configuration object.")
+    raise_errors: bool = Field(
+        default=True,
+        description="If true the tool will raise exceptions, otherwise it will log them as warnings and return []",
+    )
+    topic: str | None = Field(default=None, description="Used to provide a more detailed tool description to the agent")
+    description: str | None = Field(default=None, description="If present it will be used as the tool description")
+
+
+def _get_description_from_config(config: AIQRetrieverConfig) -> str:
+    """
+    Generate a description of what the tool will do based on how it is configured.
+    """
+    description = "Retrieve document chunks{topic} which can be used to answer the provided question."
+
+    _topic = f" related to {config.topic}" if config.topic else ""
+
+    return description.format(topic=_topic) if not config.description else config.description
+
+
+@register_function(config_type=AIQRetrieverConfig)
+async def aiq_retriever_tool(config: AIQRetrieverConfig, builder: Builder):
+    """
+    Configure an AIQ Toolkit Retriever Tool which supports different clients such as Milvus and Nemo Retriever.
+
+    Args:
+        config: A config object with required parameters 'client' and 'client_config'
+        builder: A workflow builder object
+    """
+
+    class RetrieverInputSchema(BaseModel):
+        query: str = Field(description="The query to be searched in the configured data store")
+
+    client: AIQRetriever = await builder.get_retriever(config.retriever)
+
+    async def _retrieve(query: str) -> RetrieverOutput:
+        try:
+            retrieved_context = await client.search(query=query)
+            logger.info("Retrieved %s records for query %s.", len(retrieved_context), query)
+            return retrieved_context
+
+        except RetrieverError as e:
+            if config.raise_errors:
+                raise e
+            logger.warning("Retriever threw an error: %s. Returning an empty response.", e)
+            return RetrieverOutput(results=[])
+
+    yield FunctionInfo.from_fn(
+        fn=_retrieve,
+        input_schema=RetrieverInputSchema,
+        description=_get_description_from_config(config),
+    )

aiq/tool/server_tools.py

@@ -0,0 +1,63 @@
+# 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 aiq.builder.builder import Builder
+from aiq.builder.function_info import FunctionInfo
+from aiq.cli.register_workflow import register_function
+from aiq.data_models.function import FunctionBaseConfig
+
+
+class RequestAttributesTool(FunctionBaseConfig, name="current_request_attributes"):
+    """
+    A simple tool that demonstrates how to retrieve user-defined request attributes from HTTP requests
+    within workflow tools. Please refer to the 'general' section of the configuration file located in the
+    'examples/simple_calculator/configs/config-metadata.yml' directory to see how to define a custom route using a
+    YAML file and associate it with a corresponding function to acquire request attributes.
+    """
+    pass
+
+
+@register_function(config_type=RequestAttributesTool)
+async def current_request_attributes(config: RequestAttributesTool, builder: Builder):
+
+    from starlette.datastructures import Headers
+    from starlette.datastructures import QueryParams
+
+    async def _get_request_attributes(unused: str) -> str:
+
+        from aiq.builder.context import AIQContext
+        aiq_context = AIQContext.get()
+        method: str | None = aiq_context.metadata.method
+        url_path: str | None = aiq_context.metadata.url_path
+        url_scheme: str | None = aiq_context.metadata.url_scheme
+        headers: Headers | None = aiq_context.metadata.headers
+        query_params: QueryParams | None = aiq_context.metadata.query_params
+        path_params: dict[str, str] | None = aiq_context.metadata.path_params
+        client_host: str | None = aiq_context.metadata.client_host
+        client_port: int | None = aiq_context.metadata.client_port
+        cookies: dict[str, str] | None = aiq_context.metadata.cookies
+
+        return (f"Method: {method}, "
+                f"URL Path: {url_path}, "
+                f"URL Scheme: {url_scheme}, "
+                f"Headers: {dict(headers) if headers is not None else 'None'}, "
+                f"Query Params: {dict(query_params) if query_params is not None else 'None'}, "
+                f"Path Params: {path_params}, "
+                f"Client Host: {client_host}, "
+                f"Client Port: {client_port}, "
+                f"Cookies: {cookies}")
+
+    yield FunctionInfo.from_fn(_get_request_attributes,
+                               description="Returns the acquired user defined request attriubutes.")

aiq/utils/data_models/schema_validator.py

@@ -0,0 +1,58 @@
+# 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 yaml
+from pydantic import ValidationError
+
+from ..exception_handlers.schemas import schema_exception_handler
+from ..exception_handlers.schemas import yaml_exception_handler
+
+
+@schema_exception_handler
+def validate_schema(metadata, Schema):  # pylint: disable=invalid-name
+
+    try:
+        return Schema(**metadata)
+    except ValidationError as e:
+
+        raise e
+
+
+@yaml_exception_handler
+def validate_yaml(ctx, param, value):  # pylint: disable=unused-argument
+    """
+    Validate that the file is a valid YAML file
+
+    Parameters
+    ----------
+    ctx: Click context
+    param: Click parameter
+    value: Path to YAML file
+
+    Returns
+    -------
+    str: Path to valid YAML file
+
+    Raises
+    ------
+    ValueError: If file is invalid or unreadable
+    """
+    if value is None:
+        return None
+
+    with open(value, 'r', encoding="utf-8") as f:
+        yaml.safe_load(f)
+
+    return value

aiq/utils/debugging_utils.py

@@ -0,0 +1,43 @@
+# 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.
+
+
+def is_debugger_attached() -> bool:
+    """
+    Check if a debugger is attached to the current process.
+
+    Returns
+    -------
+    bool
+        True if a debugger is attached, False otherwise
+    """
+    import sys
+
+    if "debugpy" in sys.modules:
+
+        import debugpy
+
+        return debugpy.is_client_connected()
+
+    trace_func = sys.gettrace()
+
+    # The presence of a trace function and pydevd means a debugger is attached
+    if (trace_func is not None):
+        trace_module = getattr(trace_func, "__module__", None)
+
+        if (trace_module is not None and trace_module.find("pydevd") != -1):
+            return True
+
+    return False

aiq/utils/exception_handlers/schemas.py

@@ -0,0 +1,114 @@
+# 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 logging
+
+import yaml
+from pydantic import ValidationError
+
+logger = logging.getLogger(__name__)
+
+
+def schema_exception_handler(func, **kwargs):  # pylint: disable=unused-argument
+    """
+    A decorator that handles `ValidationError` exceptions for schema validation functions.
+
+    This decorator wraps a function that performs schema validation using Pydantic.
+    If a `ValidationError` is raised, it logs detailed error messages and raises a `ValueError` with the combined error
+    messages.
+
+    Parameters
+    ----------
+    func : callable
+        The function to be decorated. This function is expected to perform schema validation.
+
+    kwargs : dict
+        Additional keyword arguments to be passed to the function.
+
+    Returns
+    -------
+    callable
+        The wrapped function that executes `func` with exception handling.
+
+    Raises
+    ------
+    ValueError
+        If a `ValidationError` is caught, this decorator logs the error details and raises a `ValueError` with the
+        combined error messages.
+
+    Notes
+    -----
+    This decorator is particularly useful for functions that validate configurations or data models,
+    ensuring that any validation errors are logged and communicated clearly.
+
+    Examples
+    --------
+    >>> @schema_exception_handler
+    ... def validate_config(config_data):
+    ...     schema = MySchema(**config_data)
+    ...     return schema
+    ...
+    >>> try:
+    ...     validate_config(invalid_config)
+    ... except ValueError as e:
+    ...     logger.error("Caught error: %s", e)
+    Caught error: Invalid configuration: field1: value is not a valid integer; field2: field required
+    """
+
+    def inner_function(*args, **kwargs):
+        try:
+            return func(*args, **kwargs)
+        except ValidationError as e:
+            error_messages = "; ".join([f"{error['loc'][0]}: {error['msg']}" for error in e.errors()])
+            log_error_message = f"Invalid configuration: {error_messages}"
+            logger.error(log_error_message)
+            raise ValueError(log_error_message) from e
+
+    return inner_function
+
+
+def yaml_exception_handler(func):
+    """
+    A decorator that handles YAML parsing exceptions.
+
+    This decorator wraps a function that performs YAML file operations.
+    If a YAML-related error occurs, it logs the error and raises a ValueError
+    with a clear error message.
+
+    Returns
+    -------
+    callable
+        The wrapped function that executes `func` with YAML exception handling.
+
+    Raises
+    ------
+    ValueError
+        If a YAML error is caught, with details about the parsing failure.
+    """
+
+    def inner_function(*args, **kwargs):
+        try:
+            return func(*args, **kwargs)
+        except yaml.YAMLError as e:
+            log_error_message = f"Invalid YAML configuration: {str(e)}"
+            logger.error(log_error_message)
+            raise ValueError(log_error_message) from e
+
+        except Exception as e:
+            log_error_message = f"Error reading YAML file: {str(e)}"
+            logger.error(log_error_message)
+            raise ValueError(log_error_message) from e
+
+    return inner_function

aiq/utils/io/yaml_tools.py

@@ -0,0 +1,119 @@
+# 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 io
+import logging
+import typing
+
+import expandvars
+import yaml
+
+from aiq.utils.type_utils import StrPath
+
+logger = logging.getLogger(__name__)
+
+
+def _interpolate_variables(value: str | int | float | bool | None) -> str | int | float | bool | None:
+    """
+    Interpolate variables in a string with the format ${VAR:-default_value}.
+    If the variable is not set, the default value will be used.
+    If no default value is provided, an empty string will be used.
+
+    Args:
+        value (str | int | float | bool | None): The value to interpolate variables in.
+
+    Returns:
+        str | int | float | bool | None: The value with variables interpolated.
+    """
+
+    if not isinstance(value, str):
+        return value
+
+    return expandvars.expandvars(value)
+
+
+def yaml_load(config_path: StrPath) -> dict:
+    """
+    Load a YAML file and interpolate variables in the format
+    ${VAR:-default_value}.
+
+    Args:
+        config_path (StrPath): The path to the YAML file to load.
+
+    Returns:
+        dict: The processed configuration dictionary.
+    """
+
+    # Read YAML file
+    with open(config_path, "r", encoding="utf-8") as stream:
+        config_str = stream.read()
+
+    return yaml_loads(config_str)
+
+
+def yaml_loads(config: str) -> dict:
+    """
+    Load a YAML string and interpolate variables in the format
+    ${VAR:-default_value}.
+
+    Args:
+        config (str): The YAML string to load.
+
+    Returns:
+        dict: The processed configuration dictionary.
+    """
+
+    interpolated_config_str = _interpolate_variables(config)
+    assert isinstance(interpolated_config_str, str), "Config must be a string"
+
+    stream = io.StringIO(interpolated_config_str)
+    stream.seek(0)
+
+    # Load the YAML data
+    try:
+        config_data = yaml.safe_load(stream)
+    except yaml.YAMLError as e:
+        logger.error("Error loading YAML: %s", interpolated_config_str, exc_info=True)
+        raise ValueError(f"Error loading YAML: {e}") from e
+
+    assert isinstance(config_data, dict)
+
+    return config_data
+
+
+def yaml_dump(config: dict, fp: typing.TextIO) -> None:
+    """
+    Dump a configuration dictionary to a YAML file.
+
+    Args:
+        config (dict): The configuration dictionary to dump.
+        fp (typing.TextIO): The file pointer to write the YAML to.
+    """
+    yaml.dump(config, stream=fp, indent=2, sort_keys=False)
+    fp.flush()
+
+
+def yaml_dumps(config: dict) -> str:
+    """
+    Dump a configuration dictionary to a YAML string.
+
+    Args:
+        config (dict): The configuration dictionary to dump.
+
+    Returns:
+        str: The YAML string.
+    """
+
+    return yaml.dump(config, indent=2)

aiq/utils/metadata_utils.py

@@ -0,0 +1,74 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from pydantic_core import PydanticUndefined
+
+from aiq.data_models.common import TypedBaseModelT
+from aiq.utils.type_utils import DecomposedType
+
+
+def generate_config_type_docs(config_type: TypedBaseModelT) -> str:
+    """Generates a docstring from configuration object to facilitate discovery.
+
+    Args:
+        config_type (TypedBaseModelT): A component configuration object.
+
+    Returns:
+        str: An enriched docstring, including model attributes and default values.
+    """
+
+    # Get the docstring
+    description_formatting = []
+    # Ensure uniform formatting of docstring
+    docstring = (config_type.__doc__ or "").strip().strip(".")
+    docstring = docstring + "." if docstring != "" else "Description unavailable."
+    description_formatting.append(docstring)
+    description_formatting.append("")
+    description_formatting.append("  Args:")
+
+    # Iterate over fields to get their documentation
+    for field_name, field_info in config_type.model_fields.items():
+
+        if (field_name == "type"):
+            field_name = "_type"
+
+        decomponsed_type = DecomposedType(field_info.annotation)
+
+        if not (decomponsed_type.is_union):
+            annotation = field_info.annotation.__name__
+        else:
+            annotation = field_info.annotation
+
+        default_string = ""
+        if ((field_info.get_default() is not PydanticUndefined) and (field_name != "_type")):
+            if issubclass(type(field_info.get_default()), str):
+                default_value = f'"{field_info.get_default()}"'
+            else:
+                default_value = field_info.get_default()
+            default_string += f" Defaults to {default_value}."
+
+        # Ensure uniform formatting of field info
+        field_info_description = (field_info.description or "").strip(".")
+        if field_info_description != "":
+            field_info_description = field_info_description + "."
+        else:
+            field_info_description = "Description unavailable."
+
+        parameter_string = f"    {field_name} ({annotation}): {field_info_description}{default_string}"
+        description_formatting.append(parameter_string)
+
+    description = "\n".join(description_formatting)
+
+    return description