aiqtoolkit 1.1.0a20250515__py3-none-any.whl → 1.1.0a20251020__py3-none-any.whl

aiq/tool/retriever.py

@@ -1,89 +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 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

@@ -1,63 +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 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

@@ -1,58 +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 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

@@ -1,43 +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.
-
-
-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

@@ -1,114 +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 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

@@ -1,119 +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 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

@@ -1,74 +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.
-
-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