aiqtoolkit 1.2.0rc3__py3-none-any.whl → 1.2.0rc5__py3-none-any.whl

aiq/utils/exception_handlers/mcp.py

@@ -1,211 +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
-import ssl
-import sys
-from collections.abc import Callable
-from functools import wraps
-from typing import Any
-
-import httpx
-
-from aiq.tool.mcp.exceptions import MCPAuthenticationError
-from aiq.tool.mcp.exceptions import MCPConnectionError
-from aiq.tool.mcp.exceptions import MCPError
-from aiq.tool.mcp.exceptions import MCPProtocolError
-from aiq.tool.mcp.exceptions import MCPRequestError
-from aiq.tool.mcp.exceptions import MCPSSLError
-from aiq.tool.mcp.exceptions import MCPTimeoutError
-from aiq.tool.mcp.exceptions import MCPToolNotFoundError
-
-logger = logging.getLogger(__name__)
-
-
-def format_mcp_error(error: MCPError, include_traceback: bool = False) -> None:
-    """Format MCP errors for CLI display with structured logging and user guidance.
-
-    Logs structured error information for debugging and displays user-friendly
-    error messages with actionable suggestions to stderr.
-
-    Args:
-        error (MCPError): MCPError instance containing message, url, category, suggestions, and original_exception
-        include_traceback (bool, optional): Whether to include the traceback in the error message. Defaults to False.
-    """
-    # Log structured error information for debugging
-    logger.error("MCP operation failed: %s", error, exc_info=include_traceback)
-
-    # Display user-friendly suggestions
-    for suggestion in error.suggestions:
-        print(f"  → {suggestion}", file=sys.stderr)
-
-
-def _extract_url(args: tuple, kwargs: dict[str, Any], url_param: str, func_name: str) -> str:
-    """Extract URL from function arguments using clean fallback chain.
-
-    Args:
-        args: Function positional arguments
-        kwargs: Function keyword arguments
-        url_param (str): Parameter name containing the URL
-        func_name (str): Function name for logging
-
-    Returns:
-        str: URL string or "unknown" if extraction fails
-    """
-    # Try keyword arguments first
-    if url_param in kwargs:
-        return kwargs[url_param]
-
-    # Try self attribute (e.g., self.url)
-    if args and hasattr(args[0], url_param):
-        return getattr(args[0], url_param)
-
-    # Try common case: url as second parameter after self
-    if len(args) > 1 and url_param == "url":
-        return args[1]
-
-    # Fallback with warning
-    logger.warning("Could not extract URL for error handling in %s", func_name)
-    return "unknown"
-
-
-def extract_primary_exception(exceptions: list[Exception]) -> Exception:
-    """Extract the most relevant exception from a group.
-
-    Prioritizes connection errors over others for better user experience.
-
-    Args:
-        exceptions (list[Exception]): List of exceptions from ExceptionGroup
-
-    Returns:
-        Exception: Most relevant exception for user feedback
-    """
-    # Prioritize connection errors
-    for exc in exceptions:
-        if isinstance(exc, (httpx.ConnectError, ConnectionError)):
-            return exc
-
-    # Then timeout errors
-    for exc in exceptions:
-        if isinstance(exc, httpx.TimeoutException):
-            return exc
-
-    # Then SSL errors
-    for exc in exceptions:
-        if isinstance(exc, ssl.SSLError):
-            return exc
-
-    # Fall back to first exception
-    return exceptions[0]
-
-
-def convert_to_mcp_error(exception: Exception, url: str) -> MCPError:
-    """Convert single exception to appropriate MCPError.
-
-    Args:
-        exception (Exception): Single exception to convert
-        url (str): MCP server URL for context
-
-    Returns:
-        MCPError: Appropriate MCPError subclass
-    """
-    match exception:
-        case httpx.ConnectError() | ConnectionError():
-            return MCPConnectionError(url, exception)
-        case httpx.TimeoutException():
-            return MCPTimeoutError(url, exception)
-        case ssl.SSLError():
-            return MCPSSLError(url, exception)
-        case httpx.RequestError():
-            return MCPRequestError(url, exception)
-        case ValueError() if "Tool" in str(exception) and "not available" in str(exception):
-            # Extract tool name from error message if possible
-            tool_name = str(exception).split("Tool ")[1].split(" not available")[0] if "Tool " in str(
-                exception) else "unknown"
-            return MCPToolNotFoundError(tool_name, url, exception)
-        case _:
-            # Handle TaskGroup error message specifically
-            if "unhandled errors in a TaskGroup" in str(exception):
-                return MCPProtocolError(url, "Failed to connect to MCP server", exception)
-            if "unauthorized" in str(exception).lower() or "forbidden" in str(exception).lower():
-                return MCPAuthenticationError(url, exception)
-            return MCPError(f"Unexpected error: {exception}", url, original_exception=exception)
-
-
-def handle_mcp_exceptions(url_param: str = "url") -> Callable[..., Any]:
-    """Decorator that handles exceptions and converts them to MCPErrors.
-
-    This decorator wraps MCP client methods and converts low-level exceptions
-    to structured MCPError instances with helpful user guidance.
-
-    Args:
-        url_param (str): Name of the parameter or attribute containing the MCP server URL
-
-    Returns:
-        Callable[..., Any]: Decorated function
-
-    Example:
-        .. code-block:: python
-
-            @handle_mcp_exceptions("url")
-            async def get_tools(self, url: str):
-                # Method implementation
-                pass
-
-            @handle_mcp_exceptions("url")  # Uses self.url
-            async def get_tool(self):
-                # Method implementation
-                pass
-    """
-
-    def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
-
-        @wraps(func)
-        async def wrapper(*args, **kwargs):
-            try:
-                return await func(*args, **kwargs)
-            except MCPError:
-                # Re-raise MCPErrors as-is
-                raise
-            except Exception as e:
-                url = _extract_url(args, kwargs, url_param, func.__name__)
-
-                # Handle ExceptionGroup by extracting most relevant exception
-                if isinstance(e, ExceptionGroup):  # noqa: F821
-                    primary_exception = extract_primary_exception(list(e.exceptions))
-                    mcp_error = convert_to_mcp_error(primary_exception, url)
-                else:
-                    mcp_error = convert_to_mcp_error(e, url)
-
-                raise mcp_error from e
-
-        return wrapper
-
-    return decorator
-
-
-def mcp_exception_handler(func: Callable[..., Any]) -> Callable[..., Any]:
-    """Simplified decorator for methods that have self.url attribute.
-
-    This is a convenience decorator that assumes the URL is available as self.url.
-    Follows the same pattern as schema_exception_handler in this directory.
-
-    Args:
-        func (Callable[..., Any]): The function to decorate
-
-    Returns:
-        Callable[..., Any]: Decorated function
-    """
-    return handle_mcp_exceptions("url")(func)

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/model_processing.py

@@ -1,28 +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 re
-
-
-def remove_r1_think_tags(text: str):
-    pattern = r'(<think>)?.*?</think>\s*(.*)'
-
-    # Add re.DOTALL flag to make . match newlines
-    match = re.match(pattern, text, re.DOTALL)
-
-    if match:
-        return match.group(2)
-
-    return text

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/log_utils.py

@@ -1,37 +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
-
-
-class LogFilter(logging.Filter):
-    """
-    This class is used to filter log records based on a defined set of criteria.
-    """
-
-    def __init__(self, filter_criteria: list[str]):
-        self._filter_criteria = filter_criteria
-        super().__init__()
-
-    def filter(self, record: logging.LogRecord):
-        """
-        Evaluates whether a log record should be emitted based on the message content.
-
-        Returns:
-            False if the message content contains any of the filter criteria, True otherwise.
-        """
-        if any(match in record.getMessage() for match in self._filter_criteria):
-            return False
-        return True

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