aiqtoolkit 1.2.0rc4__py3-none-any.whl → 1.2rc9__py3-none-any.whl

aiq/tool/register.py

@@ -1,38 +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.
-
-# pylint: disable=unused-import
-# 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
-from . import nvidia_rag
-from . import retriever
-from . import server_tools
-from .code_execution import register
-from .github_tools import create_github_commit
-from .github_tools import create_github_issue
-from .github_tools import create_github_pr
-from .github_tools import get_github_file
-from .github_tools import get_github_issue
-from .github_tools import get_github_pr
-from .github_tools import update_github_issue
-from .mcp import mcp_tool
-from .memory_tools import add_memory_tool
-from .memory_tools import delete_memory_tool
-from .memory_tools import get_memory_tool

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,66 +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/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
-
-
-@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
-        conversation_id: str | None = aiq_context.conversation_id
-
-        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}, "
-                f"Conversation Id: {conversation_id}")
-
-    yield FunctionInfo.from_fn(_get_request_attributes,
-                               description="Returns the acquired user defined request attributes.")

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

@@ -1,32 +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 argparse
-import json
-
-from aiq.runtime.loader import get_all_aiq_entrypoints_distro_mapping
-
-
-def dump_distro_mapping(path: str):
-    mapping = get_all_aiq_entrypoints_distro_mapping()
-    with open(path, "w", encoding="utf-8") as f:
-        json.dump(mapping, f, indent=4)
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--path", type=str, required=True)
-    args = parser.parse_args()
-    dump_distro_mapping(args.path)

aiq/utils/exception_handlers/automatic_retries.py

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