nvidia-nat 1.3.0a20250910__py3-none-any.whl → 1.4.0a20251112__py3-none-any.whl

nat/utils/exception_handlers/automatic_retries.py

@@ -15,11 +15,13 @@
 import asyncio
 import copy
 import functools
+import gc
 import inspect
 import logging
 import re
 import time
 import types
+import weakref
 from collections.abc import Callable
 from collections.abc import Iterable
 from collections.abc import Sequence
@@ -31,9 +33,62 @@ Exc = tuple[type[BaseException], ...]  # exception classes
 CodePattern = int | str | range  # for retry_codes argument
 logger = logging.getLogger(__name__)
 
-# ──────────────────────────────────────────────────────────────────────────────
+# ─────────────────────────────────────────────────────────────
+#  Memory-optimized helpers
+# ─────────────────────────────────────────────────────────────
+
+
+def _shallow_copy_args(args: tuple, kwargs: dict) -> tuple[tuple, dict]:
+    """Create shallow copies of args and kwargs to avoid deep copy overhead."""
+    # For most use cases, shallow copy is sufficient and much faster
+    return tuple(args), dict(kwargs)
+
+
+def _deep_copy_args(args: tuple, kwargs: dict, skip_first: bool = False) -> tuple[tuple, dict]:
+    """Create deep copies of args and kwargs to prevent mutation issues.
+
+    Args:
+        args: Positional arguments to copy
+        kwargs: Keyword arguments to copy
+        skip_first: If True, skip copying the first arg (typically 'self')
+    """
+    if skip_first and args:
+        # Don't deep copy self, only the remaining arguments
+        return (args[0], ) + copy.deepcopy(args[1:]), copy.deepcopy(kwargs)
+    return copy.deepcopy(args), copy.deepcopy(kwargs)
+
+
+def _clear_exception_context(exc: BaseException) -> None:
+    """Clear exception traceback to free memory."""
+    if exc is None:
+        return
+
+    # Clear the exception's traceback to break reference cycles
+    # This is the main memory optimization
+    try:
+        exc.__traceback__ = None
+    except AttributeError:
+        pass
+
+    # Also try to clear any chained exceptions
+    try:
+        if hasattr(exc, '__cause__') and exc.__cause__ is not None:
+            _clear_exception_context(exc.__cause__)
+        if hasattr(exc, '__context__') and exc.__context__ is not None:
+            _clear_exception_context(exc.__context__)
+    except AttributeError:
+        pass
+
+
+def _run_gc_if_needed(attempt: int, gc_frequency: int = 3) -> None:
+    """Run garbage collection periodically to free memory."""
+    if attempt > 0 and attempt % gc_frequency == 0:
+        gc.collect()
+
+
+# ─────────────────────────────────────────────────────────────
 #  Helpers: status-code extraction & pattern matching
-# ──────────────────────────────────────────────────────────────────────────────
+# ─────────────────────────────────────────────────────────────
 _CODE_ATTRS = ("code", "status", "status_code", "http_status")
 
 
@@ -55,11 +110,12 @@ def _extract_status_code(exc: BaseException) -> int | 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.
+    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") + "$")
+    regex_pattern = escaped.replace(r"\*", r"\d").replace("x", r"\d")
+    return re.compile("^" + regex_pattern + "$")
 
 
 def _code_matches(code: int, pat: CodePattern) -> bool:
@@ -70,9 +126,9 @@ def _code_matches(code: int, pat: CodePattern) -> bool:
     return bool(_pattern_to_regex(pat).match(str(code)))
 
 
-# ──────────────────────────────────────────────────────────────────────────────
-#  Unified retry-decision helper
-# ──────────────────────────────────────────────────────────────────────────────
+# ─────────────────────────────────────────────────────────────
+#  Unified retry-decision helper (unchanged)
+# ─────────────────────────────────────────────────────────────
 def _want_retry(
     exc: BaseException,
     *,
@@ -106,9 +162,9 @@ def _want_retry(
     return False
 
 
-# ──────────────────────────────────────────────────────────────────────────────
-#  Core decorator factory (sync / async / (a)gen)
-# ──────────────────────────────────────────────────────────────────────────────
+# ─────────────────────────────────────────────────────────────
+#  Memory-optimized decorator factory
+# ─────────────────────────────────────────────────────────────
 def _retry_decorator(
     *,
     retries: int = 3,
@@ -117,9 +173,12 @@ def _retry_decorator(
     retry_on: Exc = (Exception, ),
     retry_codes: Sequence[CodePattern] | None = None,
     retry_on_messages: Sequence[str] | None = None,
-    deepcopy: bool = False,
+    shallow_copy: bool = True,  # Changed default to shallow copy
+    gc_frequency: int = 3,  # Run GC every N retries
+    clear_tracebacks: bool = True,  # Clear exception tracebacks
+    instance_context_aware: 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`
@@ -127,72 +186,213 @@ def _retry_decorator(
 
     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.
+    instance_context_aware:
+        If True, the decorator will check for a retry context flag on the first
+        argument (assumed to be 'self'). If the flag is set, retries are skipped
+        to prevent retry storms in nested method calls.
     """
 
     def decorate(fn: Callable[..., T]) -> Callable[..., T]:
-        use_deepcopy = deepcopy
+        use_shallow_copy = shallow_copy
+        use_context_aware = instance_context_aware
+        skip_self_in_deepcopy = instance_context_aware
+
+        class _RetryContext:
+            """Context manager for instance-level retry gating."""
+
+            __slots__ = ("_obj_ref", "_enabled", "_active")
+
+            def __init__(self, args: tuple[Any, ...]):
+                if use_context_aware and args:
+                    try:
+                        # Use weak reference to avoid keeping objects alive
+                        self._obj_ref = weakref.ref(args[0])
+                        self._enabled = True
+                    except TypeError:
+                        # Object doesn't support weak references
+                        self._obj_ref = None
+                        self._enabled = False
+                else:
+                    self._obj_ref = None
+                    self._enabled = False
+                self._active = False
+
+            def __enter__(self):
+                if not self._enabled or self._obj_ref is None:
+                    return False
+
+                obj = self._obj_ref()
+                if obj is None:
+                    return False
 
-        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
+                    # If already in retry context, skip retries
+                    if getattr(obj, "_in_retry_context", False):
+                        return True
+                    object.__setattr__(obj, "_in_retry_context", True)
+                    self._active = True
+                    return False
+                except Exception:
+                    # Cannot set attribute, disable context
+                    self._enabled = False
+                    return False
+
+            def __exit__(self, _exc_type, _exc, _tb):
+                if (self._enabled and self._active and self._obj_ref is not None):
+                    obj = self._obj_ref()
+                    if obj is not None:
+                        try:
+                            object.__setattr__(obj, "_in_retry_context", False)
+                        except Exception:
+                            pass
+
+        async def _call_with_retry_async(*args, **kw) -> T:
+            with _RetryContext(args) as already_in_context:
+                if already_in_context:
+                    return await fn(*args, **kw)
+
+                delay = base_delay
+                last_exception = None
+
+                for attempt in range(retries):
+                    # Copy args based on configuration
+                    if use_shallow_copy:
+                        call_args, call_kwargs = _shallow_copy_args(args, kw)
+                    else:
+                        call_args, call_kwargs = _deep_copy_args(args, kw, skip_first=skip_self_in_deepcopy)
+
+                    try:
+                        return await fn(*call_args, **call_kwargs)
+                    except retry_on as exc:
+                        last_exception = exc
+
+                        # Clear traceback to free memory
+                        if clear_tracebacks:
+                            _clear_exception_context(exc)
+
+                        # Run GC periodically
+                        _run_gc_if_needed(attempt, gc_frequency)
+
+                        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
+
+                if last_exception:
+                    raise last_exception
 
         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):
+            with _RetryContext(args) as already_in_context:
+                if already_in_context:
+                    async for item in fn(*args, **kw):
                         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
+
+                delay = base_delay
+                last_exception = None
+
+                for attempt in range(retries):
+                    if use_shallow_copy:
+                        call_args, call_kwargs = _shallow_copy_args(args, kw)
+                    else:
+                        call_args, call_kwargs = _deep_copy_args(args, kw, skip_first=skip_self_in_deepcopy)
+
+                    try:
+                        async for item in fn(*call_args, **call_kwargs):
+                            yield item
+                        return
+                    except retry_on as exc:
+                        last_exception = exc
+
+                        # Memory cleanup
+                        if clear_tracebacks:
+                            _clear_exception_context(exc)
+
+                        _run_gc_if_needed(attempt, gc_frequency)
+
+                        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
+
+                if last_exception:
+                    raise last_exception
 
         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)
+            with _RetryContext(args) as already_in_context:
+                if already_in_context:
+                    yield from fn(*args, **kw)
                     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
+
+                delay = base_delay
+                last_exception = None
+
+                for attempt in range(retries):
+                    if use_shallow_copy:
+                        call_args, call_kwargs = _shallow_copy_args(args, kw)
+                    else:
+                        call_args, call_kwargs = _deep_copy_args(args, kw, skip_first=skip_self_in_deepcopy)
+
+                    try:
+                        yield from fn(*call_args, **call_kwargs)
+                        return
+                    except retry_on as exc:
+                        last_exception = exc
+
+                        # Memory cleanup
+                        if clear_tracebacks:
+                            _clear_exception_context(exc)
+
+                        _run_gc_if_needed(attempt, gc_frequency)
+
+                        if not _want_retry(exc, code_patterns=retry_codes,
+                                           msg_substrings=retry_on_messages) or attempt == retries - 1:
+                            raise
+
+                        time.sleep(delay)
+                        delay *= backoff
+
+                if last_exception:
+                    raise last_exception
 
         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
+            with _RetryContext(args) as already_in_context:
+                if already_in_context:
+                    return fn(*args, **kw)
+
+                delay = base_delay
+                last_exception = None
+
+                for attempt in range(retries):
+                    if use_shallow_copy:
+                        call_args, call_kwargs = _shallow_copy_args(args, kw)
+                    else:
+                        call_args, call_kwargs = _deep_copy_args(args, kw, skip_first=skip_self_in_deepcopy)
+
+                    try:
+                        return fn(*call_args, **call_kwargs)
+                    except retry_on as exc:
+                        last_exception = exc
+
+                        # Memory cleanup
+                        if clear_tracebacks:
+                            _clear_exception_context(exc)
+
+                        _run_gc_if_needed(attempt, gc_frequency)
+
+                        if not _want_retry(exc, code_patterns=retry_codes,
+                                           msg_substrings=retry_on_messages) or attempt == retries - 1:
+                            raise
+
+                        time.sleep(delay)
+                        delay *= backoff
+
+                if last_exception:
+                    raise last_exception
 
         # Decide which wrapper to return
         if inspect.iscoroutinefunction(fn):
@@ -209,9 +409,6 @@ def _retry_decorator(
     return decorate
 
 
-# ──────────────────────────────────────────────────────────────────────────────
-#  Public helper : patch_with_retry
-# ──────────────────────────────────────────────────────────────────────────────
 def patch_with_retry(
     obj: Any,
     *,
@@ -221,7 +418,9 @@ def patch_with_retry(
     retry_on: Exc = (Exception, ),
     retry_codes: Sequence[CodePattern] | None = None,
     retry_on_messages: Sequence[str] | None = None,
-    deepcopy: bool = False,
+    deep_copy: bool = False,
+    gc_frequency: int = 3,
+    clear_tracebacks: bool = True,
 ) -> Any:
     """
     Patch *obj* instance-locally so **every public method** retries on failure.
@@ -237,6 +436,10 @@ def patch_with_retry(
         If True, each retry receives deep‑copied *args and **kwargs* to avoid
         mutating shared state between attempts.
     """
+
+    # Invert deep copy to keep function signature the same
+    shallow_copy = not deep_copy
+
     deco = _retry_decorator(
         retries=retries,
         base_delay=base_delay,
@@ -244,10 +447,13 @@ def patch_with_retry(
         retry_on=retry_on,
         retry_codes=retry_codes,
         retry_on_messages=retry_on_messages,
-        deepcopy=deepcopy,
+        shallow_copy=shallow_copy,
+        gc_frequency=gc_frequency,
+        clear_tracebacks=clear_tracebacks,
+        instance_context_aware=True,  # Prevent retry storms
     )
 
-    # Choose attribute source: the *class* to avoid triggering __getattr__
+    # Choose attribute source: the *class* to avoid __getattr__
     cls = obj if inspect.isclass(obj) else type(obj)
     cls_name = getattr(cls, "__name__", str(cls))
 
@@ -255,7 +461,7 @@ def patch_with_retry(
         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))):
+        if name.startswith("_") or isinstance(descriptor, property | staticmethod | classmethod):
             continue
 
         original = descriptor.__func__ if isinstance(descriptor, types.MethodType) else descriptor

nat/utils/io/yaml_tools.py

@@ -16,6 +16,7 @@
 import io
 import logging
 import typing
+from pathlib import Path
 
 import expandvars
 import yaml
@@ -44,23 +45,92 @@ def _interpolate_variables(value: str | int | float | bool | None) -> str | int
     return expandvars.expandvars(value)
 
 
-def yaml_load(config_path: StrPath) -> dict:
+def deep_merge(base: dict, override: dict) -> dict:
+    """
+    Recursively merge override dictionary into base dictionary.
+
+    Args:
+        base (dict): The base configuration dictionary.
+        override (dict): The override configuration dictionary.
+
+    Returns:
+        dict: The merged configuration dictionary.
+    """
+    result = base.copy()
+    for key, value in override.items():
+        if key in result and isinstance(result[key], dict) and isinstance(value, dict):
+            result[key] = deep_merge(result[key], value)
+        else:
+            result[key] = value
+    return result
+
+
+def yaml_load(config_path: StrPath, _visited: set[Path] | None = None) -> dict:
     """
     Load a YAML file and interpolate variables in the format
     ${VAR:-default_value}.
 
+    If the YAML file contains a "base" key, the file at that path will be
+    loaded first, and the current config will be merged on top of it. This enables
+    config inheritance to reduce duplication across similar configuration files.
+
     Args:
         config_path (StrPath): The path to the YAML file to load.
+        _visited (set[Path] | None): Internal parameter for circular dependency detection.
 
     Returns:
         dict: The processed configuration dictionary.
+
+    Raises:
+        TypeError: If the "base" key is not a string.
+        FileNotFoundError: If the base configuration file does not exist.
+        ValueError: If a circular dependency is detected in configuration inheritance.
     """
+    # Normalize the config path and detect circular dependencies
+    config_path_obj = Path(config_path).resolve()
+
+    if _visited is None:
+        _visited = set()
+
+    if config_path_obj in _visited:
+        raise ValueError(f"Circular dependency detected in configuration inheritance: {config_path_obj} "
+                         f"is already in the inheritance chain")
+
+    _visited.add(config_path_obj)
 
     # Read YAML file
-    with open(config_path, "r", encoding="utf-8") as stream:
+    with open(config_path_obj, encoding="utf-8") as stream:
         config_str = stream.read()
 
-    return yaml_loads(config_str)
+    config = yaml_loads(config_str)
+
+    # Check if config specifies a base for inheritance
+    if "base" in config:
+        base_path_str = config["base"]
+
+        # Validate that base is a string
+        if not isinstance(base_path_str, str):
+            raise TypeError(f"Configuration 'base' key must be a string, got {type(base_path_str).__name__}")
+
+        # Resolve base path relative to current config
+        if not Path(base_path_str).is_absolute():
+            base_path = config_path_obj.parent / base_path_str
+        else:
+            base_path = Path(base_path_str)
+
+        # Normalize and check if base file exists
+        base_path = base_path.resolve()
+        if not base_path.exists():
+            raise FileNotFoundError(f"Base configuration file not found: {base_path}")
+
+        # Load base config (recursively, so bases can have bases)
+        base_config = yaml_load(base_path, _visited=_visited)
+
+        # Perform deep merge and remove 'base' key from result
+        config = deep_merge(base_config, config)
+        config.pop("base", None)
+
+    return config
 
 
 def yaml_loads(config: str) -> dict:

nat/utils/log_levels.py

@@ -0,0 +1,25 @@
+# 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
+
+# Define log level choices
+LOG_LEVELS = {
+    'DEBUG': logging.DEBUG,
+    'INFO': logging.INFO,
+    'WARNING': logging.WARNING,
+    'ERROR': logging.ERROR,
+    'CRITICAL': logging.CRITICAL
+}

nat/utils/responses_api.py

@@ -0,0 +1,26 @@
+# 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.
+# pylint: disable=raising-format-tuple
+
+from nat.builder.framework_enum import LLMFrameworkEnum
+from nat.data_models.llm import APITypeEnum
+
+
+def validate_no_responses_api(llm_config, framework: LLMFrameworkEnum):
+    """Validate that the LLM config does not use the Responses API."""
+
+    if llm_config.api_type == APITypeEnum.RESPONSES:
+        raise ValueError(f"Responses API is not supported for config {str(type(llm_config))} in framework {framework}. "
+                         f"Please use a different API type.")

nat/utils/string_utils.py

@@ -36,3 +36,19 @@ def convert_to_str(value: Any) -> str:
         return str(value)
     else:
         raise ValueError(f"Unsupported type for conversion to string: {type(value)}")
+
+
+def truncate_string(text: str | None, max_length: int = 100) -> str | None:
+    """
+    Truncate a string to a maximum length, adding ellipsis if truncated.
+
+    Args:
+        text: The text to truncate (can be None)
+        max_length: Maximum allowed length (default: 100)
+
+    Returns:
+        The truncated text with ellipsis if needed, or None if input was None
+    """
+    if not text or len(text) <= max_length:
+        return text
+    return text[:max_length - 3] + "..."

nat/utils/type_converter.py

@@ -90,9 +90,17 @@ class TypeConverter:
         decomposed = DecomposedType(to_type)
 
         # 1) If data is already correct type, return it
-        if to_type is None or decomposed.is_instance((data, to_type)):
+        if to_type is None or decomposed.is_instance(data):
             return data
 
+        # 2) If data is a union type, try to convert to each type in the union
+        if decomposed.is_union:
+            for union_type in decomposed.args:
+                result = self._convert(data, union_type)
+                if result is not None:
+                    return result
+            return None
+
         root = decomposed.root
 
         # 2) Attempt direct in *this* converter
@@ -198,16 +206,17 @@ class TypeConverter:
         """
         visited = set()
         final = self._try_indirect_conversion(data, to_type, visited)
+        src_type = type(data)
         if final is not None:
             # Warn once if found a chain
-            self._maybe_warn_indirect(type(data), to_type)
+            self._maybe_warn_indirect(src_type, to_type)
             return final
 
         # If no success, try parent's indirect
         if self._parent is not None:
             parent_final = self._parent._try_indirect_convert(data, to_type)
             if parent_final is not None:
-                self._maybe_warn_indirect(type(data), to_type)
+                self._maybe_warn_indirect(src_type, to_type)
                 return parent_final
 
         return None