nvidia-nat 1.3.0a20250827__py3-none-any.whl → 1.3.0a20250829__py3-none-any.whl

nat/profiler/decorators/function_tracking.py

@@ -16,7 +16,9 @@
 import functools
 import inspect
 import uuid
+from collections.abc import Callable
 from typing import Any
+from typing import cast
 
 from pydantic import BaseModel
 
@@ -252,3 +254,126 @@ def track_function(func: Any = None, *, metadata: dict[str, Any] | None = None):
         return result
 
     return sync_wrapper
+
+
+def track_unregistered_function(func: Callable[..., Any] | None = None,
+                                *,
+                                name: str | None = None,
+                                metadata: dict[str, Any] | None = None) -> Callable[..., Any]:
+    """
+    Decorator that wraps any function with scope management and automatic tracking.
+
+    - Sets active function context using the function name
+    - Leverages Context.push_active_function for built-in tracking
+    - Avoids duplicate tracking entries by relying on the library's built-in systems
+    - Supports sync/async functions and generators
+
+    Args:
+        func: The function to wrap (auto-detected when used without parentheses)
+        name: Custom name to use for tracking instead of func.__name__
+        metadata: Additional metadata to include in tracking
+    """
+
+    # If called with parameters: @track_unregistered_function(name="...", metadata={...})
+    if func is None:
+
+        def decorator_wrapper(actual_func: Callable[..., Any]) -> Callable[..., Any]:
+            # Cast to ensure type checker understands this returns a callable
+            return cast(Callable[..., Any], track_unregistered_function(actual_func, name=name, metadata=metadata))
+
+        return decorator_wrapper
+
+    # Direct decoration: @track_unregistered_function or recursive call with actual function
+    function_name: str = name if name else func.__name__
+
+    # --- Validate metadata ---
+    if metadata is not None:
+        if not isinstance(metadata, dict):
+            raise TypeError("metadata must be a dict[str, Any].")
+        if any(not isinstance(k, str) for k in metadata.keys()):
+            raise TypeError("All metadata keys must be strings.")
+
+    trace_metadata = TraceMetadata(provided_metadata=metadata)
+
+    # --- Now detect the function type and wrap accordingly ---
+    if inspect.isasyncgenfunction(func):
+        # ---------------------
+        # ASYNC GENERATOR
+        # ---------------------
+
+        @functools.wraps(func)
+        async def async_gen_wrapper(*args, **kwargs):
+            context = Context.get()
+            input_data = (
+                *args,
+                kwargs,
+            )
+            # Only do context management - let push_active_function handle tracking
+            with context.push_active_function(function_name, input_data=input_data, metadata=trace_metadata) as manager:
+                final_outputs = []
+                async for item in func(*args, **kwargs):
+                    final_outputs.append(item)
+                    yield item
+
+                manager.set_output(final_outputs)
+
+        return async_gen_wrapper
+
+    if inspect.iscoroutinefunction(func):
+        # ---------------------
+        # ASYNC FUNCTION
+        # ---------------------
+        @functools.wraps(func)
+        async def async_wrapper(*args, **kwargs):
+            context = Context.get()
+            input_data = (
+                *args,
+                kwargs,
+            )
+
+            # Only do context management - let push_active_function handle tracking
+            with context.push_active_function(function_name, input_data=input_data, metadata=trace_metadata) as manager:
+                result = await func(*args, **kwargs)
+                manager.set_output(result)
+                return result
+
+        return async_wrapper
+
+    if inspect.isgeneratorfunction(func):
+        # ---------------------
+        # SYNC GENERATOR
+        # ---------------------
+        @functools.wraps(func)
+        def sync_gen_wrapper(*args, **kwargs):
+            context = Context.get()
+            input_data = (
+                *args,
+                kwargs,
+            )
+
+            # Only do context management - let push_active_function handle tracking
+            with context.push_active_function(function_name, input_data=input_data, metadata=trace_metadata) as manager:
+                final_outputs = []
+                for item in func(*args, **kwargs):
+                    final_outputs.append(item)
+                    yield item
+
+                manager.set_output(final_outputs)
+
+        return sync_gen_wrapper
+
+    @functools.wraps(func)
+    def sync_wrapper(*args, **kwargs):
+        context = Context.get()
+        input_data = (
+            *args,
+            kwargs,
+        )
+
+        # Only do context management - let push_active_function handle tracking
+        with context.push_active_function(function_name, input_data=input_data, metadata=trace_metadata) as manager:
+            result = func(*args, **kwargs)
+            manager.set_output(result)
+            return result
+
+    return sync_wrapper

nat/profiler/profile_runner.py

@@ -282,7 +282,7 @@ class ProfilerRunner:
                 fitted_model = model_trainer.train(all_steps)
                 logger.info("Fitted model for forecasting.")
             except Exception as e:
-                logger.exception("Fitting model failed. %s", e, exc_info=True)
+                logger.exception("Fitting model failed. %s", e)
                 return ProfilerResults()
 
             if self.write_output:

nat/profiler/utils.py

@@ -175,7 +175,7 @@ def create_standardized_dataframe(requests_data: list[list[IntermediateStep]]) -
                                  event_type=step.event_type).model_dump(), )
 
     except Exception as e:
-        logger.exception("Error creating standardized DataFrame: %s", e, exc_info=True)
+        logger.exception("Error creating standardized DataFrame: %s", e)
         return pd.DataFrame()
 
     if not all_rows:

nat/registry_handlers/local/local_handler.py

@@ -133,7 +133,7 @@ class LocalRegistryHandler(AbstractRegistryHandler):
                                                            "message": msg,
                                                            "action": ActionEnum.SEARCH
                                                        })
-            logger.exception(validated_search_response.status.message, exc_info=True)
+            logger.exception(validated_search_response.status.message)
 
             yield validated_search_response
 
@@ -168,7 +168,7 @@ class LocalRegistryHandler(AbstractRegistryHandler):
             validated_remove_response = RemoveResponse(status={
                 "status": StatusEnum.ERROR, "message": msg, "action": ActionEnum.REMOVE
             })  # type: ignore
-            logger.exception(validated_remove_response.status.message, exc_info=True)
+            logger.exception(validated_remove_response.status.message)
 
             yield validated_remove_response
 

nat/registry_handlers/package_utils.py

@@ -396,7 +396,7 @@ def get_transitive_dependencies(distribution_names: list[str]) -> dict[str, set[
                     except importlib.metadata.PackageNotFoundError:
                         pass
 
-                logger.error("Distribution %s not found (tried common variations)", dist_name)
+                logger.error("Distribution %s not found (tried common variations)", dist_name, exc_info=True)
                 result[dist_name] = set()
 
     return result

nat/registry_handlers/pypi/pypi_handler.py

@@ -85,7 +85,7 @@ class PypiRegistryHandler(AbstractRegistryHandler):
             validated_publish_response = PublishResponse(status={
                 "status": StatusEnum.ERROR, "message": msg, "action": ActionEnum.PUBLISH
             })
-            logger.exception(validated_publish_response.status.message, exc_info=True)
+            logger.exception(validated_publish_response.status.message)
 
             yield validated_publish_response
 
@@ -149,7 +149,7 @@ class PypiRegistryHandler(AbstractRegistryHandler):
             validated_pull_response = PullResponse(status={
                 "status": StatusEnum.ERROR, "message": msg, "action": ActionEnum.PULL
             })
-            logger.exception(validated_pull_response.status.message, exc_info=True)
+            logger.exception(validated_pull_response.status.message)
 
             yield validated_pull_response
 
@@ -212,7 +212,7 @@ class PypiRegistryHandler(AbstractRegistryHandler):
 
         except Exception as e:
             msg = f"Error searching for artifacts: {e}"
-            logger.exception(msg, exc_info=True)
+            logger.exception(msg)
             validated_search_response = SearchResponse(params=query,
                                                        status={
                                                            "status": StatusEnum.ERROR,

nat/registry_handlers/rest/rest_handler.py

@@ -88,7 +88,7 @@ class RestRegistryHandler(AbstractRegistryHandler):
             validated_publish_response = PublishResponse(status={
                 "status": StatusEnum.ERROR, "message": msg, "action": ActionEnum.PUBLISH
             })
-            logger.exception(validated_publish_response.status.message, exc_info=True)
+            logger.exception(validated_publish_response.status.message)
 
             yield validated_publish_response
 
@@ -155,7 +155,7 @@ class RestRegistryHandler(AbstractRegistryHandler):
             validated_pull_response = PullResponse(status={
                 "status": StatusEnum.ERROR, "message": msg, "action": ActionEnum.PULL
             })
-            logger.exception(validated_pull_response.status.message, exc_info=True)
+            logger.exception(validated_pull_response.status.message)
 
             yield validated_pull_response
 
@@ -193,7 +193,7 @@ class RestRegistryHandler(AbstractRegistryHandler):
                                                            "message": msg,
                                                            "action": ActionEnum.SEARCH
                                                        })
-            logger.exception(validated_search_response.status.message, exc_info=True)
+            logger.exception(validated_search_response.status.message)
 
             yield validated_search_response
 
@@ -228,7 +228,7 @@ class RestRegistryHandler(AbstractRegistryHandler):
             validated_remove_response = RemoveResponse(status={
                 "status": StatusEnum.ERROR, "message": msg, "action": ActionEnum.REMOVE
             })
-            logger.exception(validated_remove_response.status.message, exc_info=True)
+            logger.exception(validated_remove_response.status.message)
 
             yield validated_remove_response
 

nat/retriever/milvus/retriever.py

@@ -154,7 +154,7 @@ class MilvusRetriever(Retriever):
                 return _wrap_milvus_results(results, content_field=self.content_field)
 
         except Exception as e:
-            logger.exception("Exception when retrieving results from milvus for query %s: %s", query, e)
+            logger.error("Exception when retrieving results from milvus for query %s: %s", query, e)
             raise RetrieverError(f"Error when retrieving documents from {collection_name} for query '{query}'") from e
 
     async def _search(self,

nat/retriever/nemo_retriever/retriever.py

@@ -143,7 +143,7 @@ class NemoRetriever(Retriever):
                 return _wrap_nemo_results(output=output, content_field="content")
 
         except Exception as e:
-            logger.exception("Encountered an error when retrieving results from Nemo Retriever: %s", e)
+            logger.error("Encountered an error when retrieving results from Nemo Retriever: %s", e)
             raise CollectionUnavailableError(
                 f"Error when retrieving documents from {collection_name} for query '{query}'") from e
 

nat/runtime/loader.py

@@ -210,7 +210,7 @@ def discover_and_register_plugins(plugin_type: PluginTypes):
                 # Optionally, you can mark the plugin as unavailable or take other actions
 
             except Exception:
-                logger.exception("An error occurred while loading plugin '%s': {e}", entry_point.name, exc_info=True)
+                logger.exception("An error occurred while loading plugin '%s'", entry_point.name)
 
             finally:
                 count += 1

nat/runtime/runner.py

@@ -149,7 +149,7 @@ class Runner:
 
             return result
         except Exception as e:
-            logger.exception("Error running workflow: %s", e)
+            logger.error("Error running workflow: %s", e)
             event_stream = self._context_state.event_stream.get()
             if event_stream:
                 event_stream.on_complete()
@@ -181,7 +181,7 @@ class Runner:
                     event_stream.on_complete()
 
         except Exception as e:
-            logger.exception("Error running workflow: %s", e)
+            logger.error("Error running workflow: %s", e)
             event_stream = self._context_state.event_stream.get()
             if event_stream:
                 event_stream.on_complete()

nat/settings/global_settings.py

@@ -231,7 +231,7 @@ class Settings(HashableBaseModel):
             return True
 
         except Exception as e:
-            logger.exception("Unable to validate user settings configuration: %s", e, exc_info=True)
+            logger.exception("Unable to validate user settings configuration: %s", e)
             return False
 
     def print_channel_settings(self, channel_type: str | None = None) -> None:

nat/tool/code_execution/local_sandbox/local_sandbox_server.py

@@ -121,7 +121,7 @@ def execute_code_subprocess(generated_code: str, queue):
         resource.setrlimit(resource.RLIMIT_AS, (limit, limit))
         resource.setrlimit(resource.RLIMIT_DATA, (limit, limit))
     except Exception as e:
-        logger.error("Failed to set resource limits, PID: %s, error: %s", os.getpid(), e)
+        logger.exception("Failed to set resource limits, PID: %s, error: %s", os.getpid(), e)
 
     stdout_capture = StringIO()
     stderr_capture = StringIO()

nat/tool/nvidia_rag.py

@@ -86,7 +86,7 @@ async def nvidia_rag_tool(config: NVIDIARAGToolConfig, builder: Builder):
                     [await aformat_document(doc, document_prompt) for doc in docs])
                 return parsed_output
             except Exception as e:
-                logger.exception("Error while running the tool", exc_info=True)
+                logger.exception("Error while running the tool")
                 return f"Error while running the tool: {e}"
 
         yield FunctionInfo.from_fn(

nat/tool/retriever.py

@@ -78,8 +78,9 @@ async def retriever_tool(config: RetrieverConfig, builder: Builder):
 
         except RetrieverError as e:
             if config.raise_errors:
-                raise e
-            logger.warning("Retriever threw an error: %s. Returning an empty response.", e)
+                logger.error("Retriever threw an error: %s.", e)
+                raise
+            logger.exception("Retriever threw an error: %s. Returning an empty response.", e)
             return RetrieverOutput(results=[])
 
     yield FunctionInfo.from_fn(

nat/utils/io/yaml_tools.py

@@ -85,7 +85,7 @@ def yaml_loads(config: str) -> dict:
     try:
         config_data = yaml.safe_load(stream)
     except yaml.YAMLError as e:
-        logger.error("Error loading YAML: %s", interpolated_config_str, exc_info=True)
+        logger.error("Error loading YAML: %s", interpolated_config_str)
         raise ValueError(f"Error loading YAML: {e}") from e
 
     assert isinstance(config_data, dict)

nat/utils/reactive/observer.py

@@ -64,7 +64,7 @@ class Observer(ObserverBase[_T_in_contra]):
                 try:
                     self._on_error(exc)
                 except Exception as e:
-                    logger.exception("Error in on_error callback: %s", e, exc_info=True)
+                    logger.exception("Error in on_error callback: %s", e)
 
     def on_complete(self) -> None:
         if not self._stopped:
@@ -73,4 +73,4 @@ class Observer(ObserverBase[_T_in_contra]):
                 try:
                     self._on_complete()
                 except Exception as e:
-                    logger.exception("Error in on_complete callback: %s", e, exc_info=True)
+                    logger.exception("Error in on_complete callback: %s", e)

nat/utils/settings/global_settings.py

@@ -55,7 +55,7 @@ def configure_registry_channel(config_type: RegistryHandlerBaseConfig, channel_n
                 channel_registry_pre[field] = getattr(validated_field_model, field)
                 break
             except Exception as e:
-                logger.exception(e, exc_info=True)
+                logger.exception(e)
                 logger.warning("Invalid '%s' input, input must be of type %s.", field, info.annotation)
 
     validated_model = config_type(**channel_registry_pre)
@@ -78,7 +78,7 @@ def add_channel_interative(channel_type: str) -> None:
     try:
         ChannelConfigType = registry.get_registered_channel_info_by_channel_type(channel_type=channel_type).config_type
     except Exception as e:
-        logger.exception("Invalid channel type: %s", e, exc_info=True)
+        logger.exception("Invalid channel type: %s", e)
         return
 
     while (True):

{nvidia_nat-1.3.0a20250827.dist-info → nvidia_nat-1.3.0a20250829.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nvidia-nat
-Version: 1.3.0a20250827
+Version: 1.3.0a20250829
 Summary: NVIDIA NeMo Agent toolkit
 Author: NVIDIA Corporation
 Maintainer: NVIDIA Corporation
@@ -243,6 +243,8 @@ Provides-Extra: agno
 Requires-Dist: nvidia-nat-agno; extra == "agno"
 Provides-Extra: crewai
 Requires-Dist: nvidia-nat-crewai; extra == "crewai"
+Provides-Extra: data-flywheel
+Requires-Dist: nvidia-nat-data-flywheel; extra == "data-flywheel"
 Provides-Extra: ingestion
 Requires-Dist: nvidia-nat-ingestion; extra == "ingestion"
 Provides-Extra: langchain