aiqtoolkit 1.2.0.dev0__py3-none-any.whl → 1.2.0rc1__py3-none-any.whl

aiq/tool/register.py

@@ -17,6 +17,7 @@
 # flake8: noqa
 
 # Import any tools which need to be automatically registered here
+from . import chat_completion
 from . import datetime_tools
 from . import document_search
 from . import github_tools

aiq/tool/server_tools.py

@@ -23,8 +23,8 @@ 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.
+    'examples/getting_started/simple_web_query/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
 
@@ -39,6 +39,7 @@ async def current_request_attributes(config: RequestAttributesTool, builder: Bui
 
         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
@@ -48,6 +49,7 @@ async def current_request_attributes(config: RequestAttributesTool, builder: Bui
         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
+        conversation_id: str | None = aiq_context.conversation_id
 
         return (f"Method: {method}, "
                 f"URL Path: {url_path}, "
@@ -57,7 +59,8 @@ async def current_request_attributes(config: RequestAttributesTool, builder: Bui
                 f"Path Params: {path_params}, "
                 f"Client Host: {client_host}, "
                 f"Client Port: {client_port}, "
-                f"Cookies: {cookies}")
+                f"Cookies: {cookies}, "
+                f"Conversation Id: {conversation_id}")
 
     yield FunctionInfo.from_fn(_get_request_attributes,
                                description="Returns the acquired user defined request attriubutes.")

aiq/utils/exception_handlers/automatic_retries.py

@@ -0,0 +1,289 @@
+# 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 copy
+import functools
+import inspect
+import logging
+import re
+import time
+import types
+from collections.abc import Callable
+from collections.abc import Iterable
+from collections.abc import Sequence
+from typing import Any
+from typing import TypeVar
+
+# pylint: disable=inconsistent-return-statements
+
+T = TypeVar("T")
+Exc = tuple[type[BaseException], ...]  # exception classes
+CodePattern = int | str | range  # for retry_codes argument
+logger = logging.getLogger(__name__)
+
+# ──────────────────────────────────────────────────────────────────────────────
+#  Helpers: status-code extraction & pattern matching
+# ──────────────────────────────────────────────────────────────────────────────
+_CODE_ATTRS = ("code", "status", "status_code", "http_status")
+
+
+def _extract_status_code(exc: BaseException) -> int | None:
+    """Return a numeric status code found inside *exc*, else None."""
+    for attr in _CODE_ATTRS:
+        if hasattr(exc, attr):
+            try:
+                return int(getattr(exc, attr))
+            except (TypeError, ValueError):
+                pass
+    if exc.args:
+        try:
+            return int(exc.args[0])
+        except (TypeError, ValueError):
+            pass
+    return None
+
+
+def _pattern_to_regex(pat: str) -> re.Pattern[str]:
+    """
+    Convert simple wildcard pattern (“4xx”, “5*”, “40x”) to a ^regex$.
+    Rule:  ‘x’ or ‘*’ ⇒ any digit.
+    """
+    escaped = re.escape(pat)
+    return re.compile("^" + escaped.replace(r"\*", r"\d").replace("x", r"\d") + "$")
+
+
+def _code_matches(code: int, pat: CodePattern) -> bool:
+    if isinstance(pat, int):
+        return code == pat
+    if isinstance(pat, range):
+        return code in pat
+    return bool(_pattern_to_regex(pat).match(str(code)))
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+#  Unified retry-decision helper
+# ──────────────────────────────────────────────────────────────────────────────
+def _want_retry(
+    exc: BaseException,
+    *,
+    code_patterns: Sequence[CodePattern] | None,
+    msg_substrings: Sequence[str] | None,
+) -> bool:
+    """
+    Return True if the exception satisfies *either* (when provided):
+       • code_patterns  – matches status-code pattern(s)
+       • msg_substrings – contains any of the substrings (case-insensitive)
+    """
+
+    if not code_patterns and not msg_substrings:
+        logger.info("Retrying on exception %s without extra filters", exc)
+        return True
+
+    # -------- status-code filter --------
+    if code_patterns is not None:
+        code = _extract_status_code(exc)
+        if any(_code_matches(code, p) for p in code_patterns):
+            logger.info("Retrying on exception %s with matched code %s", exc, code)
+            return True
+
+    # -------- message filter -----------
+    if msg_substrings is not None:
+        msg = str(exc).lower()
+        if any(s.lower() in msg for s in msg_substrings):
+            logger.info("Retrying on exception %s with matched message %s", exc, msg)
+            return True
+
+    return False
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+#  Core decorator factory (sync / async / (a)gen)
+# ──────────────────────────────────────────────────────────────────────────────
+def _retry_decorator(
+    *,
+    retries: int = 3,
+    base_delay: float = 0.25,
+    backoff: float = 2.0,
+    retry_on: Exc = (Exception, ),
+    retry_codes: Sequence[CodePattern] | None = None,
+    retry_on_messages: Sequence[str] | None = None,
+    deepcopy: bool = False,
+) -> Callable[[Callable[..., T]], Callable[..., T]]:
+    """
+    Build a decorator that retries with exponential back-off *iff*:
+
+      • the raised exception is an instance of one of `retry_on`
+      • AND `_want_retry()` returns True (i.e. matches codes/messages filters)
+
+    If both `retry_codes` and `retry_on_messages` are None, all exceptions are retried.
+
+    deepcopy:
+        If True, each retry receives deep‑copied *args and **kwargs* to avoid
+        mutating shared state between attempts.
+    """
+
+    def decorate(fn: Callable[..., T]) -> Callable[..., T]:
+        use_deepcopy = deepcopy
+
+        async def _call_with_retry_async(*args, **kw) -> T:
+            delay = base_delay
+            for attempt in range(retries):
+                call_args = copy.deepcopy(args) if use_deepcopy else args
+                call_kwargs = copy.deepcopy(kw) if use_deepcopy else kw
+                try:
+                    return await fn(*call_args, **call_kwargs)
+                except retry_on as exc:
+                    if (not _want_retry(exc, code_patterns=retry_codes, msg_substrings=retry_on_messages)
+                            or attempt == retries - 1):
+                        raise
+                    await asyncio.sleep(delay)
+                    delay *= backoff
+
+        async def _agen_with_retry(*args, **kw):
+            delay = base_delay
+            for attempt in range(retries):
+                call_args = copy.deepcopy(args) if use_deepcopy else args
+                call_kwargs = copy.deepcopy(kw) if use_deepcopy else kw
+                try:
+                    async for item in fn(*call_args, **call_kwargs):
+                        yield item
+                    return
+                except retry_on as exc:
+                    if (not _want_retry(exc, code_patterns=retry_codes, msg_substrings=retry_on_messages)
+                            or attempt == retries - 1):
+                        raise
+                    await asyncio.sleep(delay)
+                    delay *= backoff
+
+        def _gen_with_retry(*args, **kw) -> Iterable[Any]:
+            delay = base_delay
+            for attempt in range(retries):
+                call_args = copy.deepcopy(args) if use_deepcopy else args
+                call_kwargs = copy.deepcopy(kw) if use_deepcopy else kw
+                try:
+                    yield from fn(*call_args, **call_kwargs)
+                    return
+                except retry_on as exc:
+                    if (not _want_retry(exc, code_patterns=retry_codes, msg_substrings=retry_on_messages)
+                            or attempt == retries - 1):
+                        raise
+                    time.sleep(delay)
+                    delay *= backoff
+
+        def _sync_with_retry(*args, **kw) -> T:
+            delay = base_delay
+            for attempt in range(retries):
+                call_args = copy.deepcopy(args) if use_deepcopy else args
+                call_kwargs = copy.deepcopy(kw) if use_deepcopy else kw
+                try:
+                    return fn(*call_args, **call_kwargs)
+                except retry_on as exc:
+                    if (not _want_retry(exc, code_patterns=retry_codes, msg_substrings=retry_on_messages)
+                            or attempt == retries - 1):
+                        raise
+                    time.sleep(delay)
+                    delay *= backoff
+
+        # Decide which wrapper to return
+        if inspect.iscoroutinefunction(fn):
+            wrapper = _call_with_retry_async
+        elif inspect.isasyncgenfunction(fn):
+            wrapper = _agen_with_retry
+        elif inspect.isgeneratorfunction(fn):
+            wrapper = _gen_with_retry
+        else:
+            wrapper = _sync_with_retry
+
+        return functools.wraps(fn)(wrapper)  # type: ignore[return-value]
+
+    return decorate
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+#  Public helper : patch_with_retry
+# ──────────────────────────────────────────────────────────────────────────────
+def patch_with_retry(
+    obj: Any,
+    *,
+    retries: int = 3,
+    base_delay: float = 0.25,
+    backoff: float = 2.0,
+    retry_on: Exc = (Exception, ),
+    retry_codes: Sequence[CodePattern] | None = None,
+    retry_on_messages: Sequence[str] | None = None,
+    deepcopy: bool = False,
+) -> Any:
+    """
+    Patch *obj* instance-locally so **every public method** retries on failure.
+
+    Extra filters
+    -------------
+    retry_codes
+        Same as before – ints, ranges, or wildcard strings (“4xx”, “5*”…).
+    retry_on_messages
+        List of *substring* patterns.  We retry only if **any** pattern
+        appears (case-insensitive) in `str(exc)`.
+    deepcopy:
+        If True, each retry receives deep‑copied *args and **kwargs* to avoid
+        mutating shared state between attempts.
+    """
+    deco = _retry_decorator(
+        retries=retries,
+        base_delay=base_delay,
+        backoff=backoff,
+        retry_on=retry_on,
+        retry_codes=retry_codes,
+        retry_on_messages=retry_on_messages,
+        deepcopy=deepcopy,
+    )
+
+    # Choose attribute source: the *class* to avoid triggering __getattr__
+    cls = obj if inspect.isclass(obj) else type(obj)
+    cls_name = getattr(cls, "__name__", str(cls))
+
+    for name, _ in inspect.getmembers(cls, callable):
+        descriptor = inspect.getattr_static(cls, name)
+
+        # Skip dunders, privates and all descriptors we must not wrap
+        if (name.startswith("_") or isinstance(descriptor, (property, staticmethod, classmethod))):
+            continue
+
+        original = descriptor.__func__ if isinstance(descriptor, types.MethodType) else descriptor
+        wrapped = deco(original)
+
+        try:  # instance‑level first
+            if not inspect.isclass(obj):
+                object.__setattr__(obj, name, types.MethodType(wrapped, obj))
+                continue
+        except Exception as exc:
+            logger.info(
+                "Instance‑level patch failed for %s.%s (%s); "
+                "falling back to class‑level patch.",
+                cls_name,
+                name,
+                exc,
+            )
+
+        try:  # class‑level fallback
+            setattr(cls, name, wrapped)
+        except Exception as exc:
+            logger.info(
+                "Cannot patch method %s.%s with automatic retries: %s",
+                cls_name,
+                name,
+                exc,
+            )
+
+    return obj

aiq/utils/exception_handlers/mcp.py

@@ -0,0 +1,211 @@
+# 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/io/model_processing.py

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

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

@@ -0,0 +1,38 @@
+# 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 typing import Any
+
+from pydantic import BaseModel
+
+
+def convert_to_str(value: Any) -> str:
+    """
+    Convert a value to a string representation.
+    Handles various types including lists, dictionaries, and other objects.
+    """
+    if isinstance(value, str):
+        return value
+
+    if isinstance(value, list):
+        return ", ".join(map(str, value))
+    elif isinstance(value, BaseModel):
+        return value.model_dump_json(exclude_none=True, exclude_unset=True)
+    elif isinstance(value, dict):
+        return ", ".join(f"{k}: {v}" for k, v in value.items())
+    elif hasattr(value, '__str__'):
+        return str(value)
+    else:
+        raise ValueError(f"Unsupported type for conversion to string: {type(value)}")