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

nat/builder/context.py

@@ -19,6 +19,7 @@ from collections.abc import Awaitable
 from collections.abc import Callable
 from contextlib import contextmanager
 from contextvars import ContextVar
+from functools import cached_property
 
 from nat.builder.intermediate_step_manager import IntermediateStepManager
 from nat.builder.user_interaction_manager import UserInteractionManager
@@ -40,12 +41,12 @@ from nat.utils.reactive.subject import Subject
 class Singleton(type):
 
     def __init__(cls, name, bases, dict):
-        super(Singleton, cls).__init__(name, bases, dict)
+        super().__init__(name, bases, dict)
         cls.instance = None
 
     def __call__(cls, *args, **kw):
         if cls.instance is None:
-            cls.instance = super(Singleton, cls).__call__(*args, **kw)
+            cls.instance = super().__call__(*args, **kw)
         return cls.instance
 
 
@@ -67,14 +68,14 @@ class ContextState(metaclass=Singleton):
     def __init__(self):
         self.conversation_id: ContextVar[str | None] = ContextVar("conversation_id", default=None)
         self.user_message_id: ContextVar[str | None] = ContextVar("user_message_id", default=None)
+        self.workflow_run_id: ContextVar[str | None] = ContextVar("workflow_run_id", default=None)
+        self.workflow_trace_id: ContextVar[int | None] = ContextVar("workflow_trace_id", default=None)
         self.input_message: ContextVar[typing.Any] = ContextVar("input_message", default=None)
         self.user_manager: ContextVar[typing.Any] = ContextVar("user_manager", default=None)
-        self.metadata: ContextVar[RequestAttributes] = ContextVar("request_attributes", default=RequestAttributes())
-        self.event_stream: ContextVar[Subject[IntermediateStep] | None] = ContextVar("event_stream", default=Subject())
-        self.active_function: ContextVar[InvocationNode] = ContextVar("active_function",
-                                                                      default=InvocationNode(function_id="root",
-                                                                                             function_name="root"))
-        self.active_span_id_stack: ContextVar[list[str]] = ContextVar("active_span_id_stack", default=["root"])
+        self._metadata: ContextVar[RequestAttributes | None] = ContextVar("request_attributes", default=None)
+        self._event_stream: ContextVar[Subject[IntermediateStep] | None] = ContextVar("event_stream", default=None)
+        self._active_function: ContextVar[InvocationNode | None] = ContextVar("active_function", default=None)
+        self._active_span_id_stack: ContextVar[list[str] | None] = ContextVar("active_span_id_stack", default=None)
 
         # Default is a lambda no-op which returns NoneType
         self.user_input_callback: ContextVar[Callable[[InteractionPrompt], Awaitable[HumanResponse | None]]
@@ -85,6 +86,30 @@ class ContextState(metaclass=Singleton):
                                                      Awaitable[AuthenticatedContext]]
                                             | None] = ContextVar("user_auth_callback", default=None)
 
+    @property
+    def metadata(self) -> ContextVar[RequestAttributes]:
+        if self._metadata.get() is None:
+            self._metadata.set(RequestAttributes())
+        return typing.cast(ContextVar[RequestAttributes], self._metadata)
+
+    @property
+    def active_function(self) -> ContextVar[InvocationNode]:
+        if self._active_function.get() is None:
+            self._active_function.set(InvocationNode(function_id="root", function_name="root"))
+        return typing.cast(ContextVar[InvocationNode], self._active_function)
+
+    @property
+    def event_stream(self) -> ContextVar[Subject[IntermediateStep]]:
+        if self._event_stream.get() is None:
+            self._event_stream.set(Subject())
+        return typing.cast(ContextVar[Subject[IntermediateStep]], self._event_stream)
+
+    @property
+    def active_span_id_stack(self) -> ContextVar[list[str]]:
+        if self._active_span_id_stack.get() is None:
+            self._active_span_id_stack.set(["root"])
+        return typing.cast(ContextVar[list[str]], self._active_span_id_stack)
+
     @staticmethod
     def get() -> "ContextState":
         return ContextState()
@@ -98,14 +123,14 @@ class Context:
     @property
     def input_message(self):
         """
-            Retrieves the input message from the context state.
+        Retrieves the input message from the context state.
 
-            The input_message property is used to access the message stored in the
-            context state. This property returns the message as it is currently
-            maintained in the context.
+        The input_message property is used to access the message stored in the
+        context state. This property returns the message as it is currently
+        maintained in the context.
 
-            Returns:
-                str: The input message retrieved from the context state.
+        Returns:
+            str: The input message retrieved from the context state.
         """
         return self._context_state.input_message.get()
 
@@ -143,7 +168,7 @@ class Context:
         """
         return UserInteractionManager(self._context_state)
 
-    @property
+    @cached_property
     def intermediate_step_manager(self) -> IntermediateStepManager:
         """
         Retrieves the intermediate step manager instance from the current context state.
@@ -174,6 +199,20 @@ class Context:
         """
         return self._context_state.user_message_id.get()
 
+    @property
+    def workflow_run_id(self) -> str | None:
+        """
+        Returns a stable identifier for the current workflow/agent invocation (UUID string).
+        """
+        return self._context_state.workflow_run_id.get()
+
+    @property
+    def workflow_trace_id(self) -> int | None:
+        """
+        Returns the 128-bit trace identifier for the current run, used as the OpenTelemetry trace_id.
+        """
+        return self._context_state.workflow_trace_id.get()
+
     @contextmanager
     def push_active_function(self,
                              function_name: str,

nat/builder/eval_builder.py

@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import asyncio
 import dataclasses
 import logging
 from contextlib import asynccontextmanager
@@ -90,17 +91,20 @@ class WorkflowEvalBuilder(WorkflowBuilder, EvalBuilder):
         return self.eval_general_config.output_dir
 
     @override
-    def get_all_tools(self, wrapper_type: LLMFrameworkEnum | str):
-        tools = []
+    async def get_all_tools(self, wrapper_type: LLMFrameworkEnum | str):
         tool_wrapper_reg = self._registry.get_tool_wrapper(llm_framework=wrapper_type)
-        for fn_name in self._functions:
-            fn = self.get_function(fn_name)
+
+        async def get_tool(fn_name: str):
+            fn = await self.get_function(fn_name)
             try:
-                tools.append(tool_wrapper_reg.build_fn(fn_name, fn, self))
+                return tool_wrapper_reg.build_fn(fn_name, fn, self)
             except Exception:
                 logger.exception("Error fetching tool `%s`", fn_name)
+                return None
 
-        return tools
+        tasks = [get_tool(fn_name) for fn_name in self._functions]
+        tools = await asyncio.gather(*tasks, return_exceptions=False)
+        return [tool for tool in tools if tool is not None]
 
     def _log_build_failure_evaluator(self,
                                      failing_evaluator_name: str,
@@ -127,11 +131,12 @@ class WorkflowEvalBuilder(WorkflowBuilder, EvalBuilder):
                                 remaining_components,
                                 original_error)
 
-    async def populate_builder(self, config: Config):
+    @override
+    async def populate_builder(self, config: Config, skip_workflow: bool = False):
         # Skip setting workflow if workflow config is EmptyFunctionConfig
-        skip_workflow = isinstance(config.workflow, EmptyFunctionConfig)
+        skip_workflow = skip_workflow or isinstance(config.workflow, EmptyFunctionConfig)
 
-        await super().populate_builder(config, skip_workflow)
+        await super().populate_builder(config, skip_workflow=skip_workflow)
 
         # Initialize progress tracking for evaluators
         completed_evaluators = []

nat/builder/framework_enum.py

@@ -22,3 +22,4 @@ class LLMFrameworkEnum(str, Enum):
     CREWAI = "crewai"
     SEMANTIC_KERNEL = "semantic_kernel"
     AGNO = "agno"
+    ADK = "adk"

nat/builder/front_end.py

@@ -37,7 +37,7 @@ class FrontEndBase(typing.Generic[FrontEndConfigT], ABC):
 
         super().__init__()
 
-        self._full_config: "Config" = full_config
+        self._full_config: Config = full_config
         self._front_end_config: FrontEndConfigT = typing.cast(FrontEndConfigT, full_config.general.front_end)
 
     @property

nat/builder/function.py

@@ -14,12 +14,14 @@
 # limitations under the License.
 
 import logging
+import re
 import typing
 from abc import ABC
 from abc import abstractmethod
 from collections.abc import AsyncGenerator
 from collections.abc import Awaitable
 from collections.abc import Callable
+from collections.abc import Sequence
 
 from pydantic import BaseModel
 
@@ -29,7 +31,9 @@ from nat.builder.function_base import InputT
 from nat.builder.function_base import SingleOutputT
 from nat.builder.function_base import StreamingOutputT
 from nat.builder.function_info import FunctionInfo
+from nat.data_models.function import EmptyFunctionConfig
 from nat.data_models.function import FunctionBaseConfig
+from nat.data_models.function import FunctionGroupBaseConfig
 
 _InvokeFnT = Callable[[InputT], Awaitable[SingleOutputT]]
 _StreamFnT = Callable[[InputT], AsyncGenerator[StreamingOutputT]]
@@ -342,3 +346,369 @@ class LambdaFunction(Function[InputT, StreamingOutputT, SingleOutputT]):
             pass
 
         return FunctionImpl(config=config, info=info, instance_name=instance_name)
+
+
+class FunctionGroup:
+    """
+    A group of functions that can be used together, sharing the same configuration, context, and resources.
+    """
+
+    def __init__(self,
+                 *,
+                 config: FunctionGroupBaseConfig,
+                 instance_name: str | None = None,
+                 filter_fn: Callable[[Sequence[str]], Awaitable[Sequence[str]]] | None = None):
+        """
+        Creates a new function group.
+
+        Parameters
+        ----------
+        config : FunctionGroupBaseConfig
+            The configuration for the function group.
+        instance_name : str | None, optional
+            The name of the function group. If not provided, the type of the function group will be used.
+        filter_fn : Callable[[Sequence[str]], Awaitable[Sequence[str]]] | None, optional
+            A callback function to additionally filter the functions in the function group dynamically when
+            the functions are accessed via any accessor method.
+        """
+        self._config = config
+        self._instance_name = instance_name or config.type
+        self._functions: dict[str, Function] = dict()
+        self._filter_fn = filter_fn
+        self._per_function_filter_fn: dict[str, Callable[[str], Awaitable[bool]]] = dict()
+
+    def add_function(self,
+                     name: str,
+                     fn: Callable,
+                     *,
+                     input_schema: type[BaseModel] | None = None,
+                     description: str | None = None,
+                     converters: list[Callable] | None = None,
+                     filter_fn: Callable[[str], Awaitable[bool]] | None = None):
+        """
+        Adds a function to the function group.
+
+        Parameters
+        ----------
+        name : str
+            The name of the function.
+        fn : Callable
+            The function to add to the function group.
+        input_schema : type[BaseModel] | None, optional
+            The input schema for the function.
+        description : str | None, optional
+            The description of the function.
+        converters : list[Callable] | None, optional
+            The converters to use for the function.
+        filter_fn : Callable[[str], Awaitable[bool]] | None, optional
+            A callback to determine if the function should be included in the function group. The
+            callback will be called with the function name. The callback is invoked dynamically when
+            the functions are accessed via any accessor method such as `get_accessible_functions`,
+            `get_included_functions`, `get_excluded_functions`, `get_all_functions`.
+
+        Raises
+        ------
+        ValueError
+            When the function name is empty or blank.
+            When the function name contains invalid characters.
+            When the function already exists in the function group.
+        """
+        if not name.strip():
+            raise ValueError("Function name cannot be empty or blank")
+        if not re.match(r"^[a-zA-Z0-9_.-]+$", name):
+            raise ValueError(
+                f"Function name can only contain letters, numbers, underscores, periods, and hyphens: {name}")
+        if name in self._functions:
+            raise ValueError(f"Function {name} already exists in function group {self._instance_name}")
+
+        info = FunctionInfo.from_fn(fn, input_schema=input_schema, description=description, converters=converters)
+        full_name = self._get_fn_name(name)
+        lambda_fn = LambdaFunction.from_info(config=EmptyFunctionConfig(), info=info, instance_name=full_name)
+        self._functions[name] = lambda_fn
+        if filter_fn:
+            self._per_function_filter_fn[name] = filter_fn
+
+    def get_config(self) -> FunctionGroupBaseConfig:
+        """
+        Returns the configuration for the function group.
+
+        Returns
+        -------
+        FunctionGroupBaseConfig
+            The configuration for the function group.
+        """
+        return self._config
+
+    def _get_fn_name(self, name: str) -> str:
+        return f"{self._instance_name}.{name}"
+
+    async def _fn_should_be_included(self, name: str) -> bool:
+        if name not in self._per_function_filter_fn:
+            return True
+        return await self._per_function_filter_fn[name](name)
+
+    async def _get_all_but_excluded_functions(
+        self,
+        filter_fn: Callable[[Sequence[str]], Awaitable[Sequence[str]]] | None = None,
+    ) -> dict[str, Function]:
+        """
+        Returns a dictionary of all functions in the function group except the excluded functions.
+        """
+        missing = set(self._config.exclude) - set(self._functions.keys())
+        if missing:
+            raise ValueError(f"Unknown excluded functions: {sorted(missing)}")
+
+        if filter_fn is None:
+            if self._filter_fn is None:
+
+                async def identity_filter(x: Sequence[str]) -> Sequence[str]:
+                    return x
+
+                filter_fn = identity_filter
+            else:
+                filter_fn = self._filter_fn
+
+        excluded = set(self._config.exclude)
+        included = set(await filter_fn(list(self._functions.keys())))
+
+        result = {}
+        for name in self._functions:
+            if name in excluded:
+                continue
+            if not await self._fn_should_be_included(name):
+                continue
+            if name not in included:
+                continue
+            result[self._get_fn_name(name)] = self._functions[name]
+
+        return result
+
+    async def get_accessible_functions(
+        self,
+        filter_fn: Callable[[Sequence[str]], Awaitable[Sequence[str]]] | None = None,
+    ) -> dict[str, Function]:
+        """
+        Returns a dictionary of all accessible functions in the function group.
+
+        First, the functions are filtered by the function group's configuration.
+        If the function group is configured to:
+        - include some functions, this will return only the included functions.
+        - not include or exclude any function, this will return all functions in the group.
+        - exclude some functions, this will return all functions in the group except the excluded functions.
+
+        Then, the functions are filtered by filter function and per-function filter functions.
+
+        Parameters
+        ----------
+        filter_fn : Callable[[Sequence[str]], Awaitable[Sequence[str]]] | None, optional
+            A callback function to additionally filter the functions in the function group dynamically. If not provided
+            then fall back to the function group's filter function. If no filter function is set for the function group
+            all functions will be returned.
+
+        Returns
+        -------
+        dict[str, Function]
+            A dictionary of all accessible functions in the function group.
+
+        Raises
+        ------
+        ValueError
+            When the function group is configured to include functions that are not found in the group.
+        """
+        if self._config.include:
+            return await self.get_included_functions(filter_fn=filter_fn)
+        if self._config.exclude:
+            return await self._get_all_but_excluded_functions(filter_fn=filter_fn)
+        return await self.get_all_functions(filter_fn=filter_fn)
+
+    async def get_excluded_functions(
+        self,
+        filter_fn: Callable[[Sequence[str]], Awaitable[Sequence[str]]] | None = None,
+    ) -> dict[str, Function]:
+        """
+        Returns a dictionary of all functions in the function group which are configured to be excluded or filtered
+        out by a filter function or per-function filter function.
+
+        Parameters
+        ----------
+        filter_fn : Callable[[Sequence[str]], Awaitable[Sequence[str]]] | None, optional
+            A callback function to additionally filter the functions in the function group dynamically. If not provided
+            then fall back to the function group's filter function. If no filter function is set for the function group
+            then no functions will be added to the returned dictionary.
+
+        Returns
+        -------
+        dict[str, Function]
+            A dictionary of all excluded functions in the function group.
+
+        Raises
+        ------
+        ValueError
+            When the function group is configured to exclude functions that are not found in the group.
+        """
+        missing = set(self._config.exclude) - set(self._functions.keys())
+        if missing:
+            raise ValueError(f"Unknown excluded functions: {sorted(missing)}")
+
+        if filter_fn is None:
+            if self._filter_fn is None:
+
+                async def identity_filter(x: Sequence[str]) -> Sequence[str]:
+                    return x
+
+                filter_fn = identity_filter
+            else:
+                filter_fn = self._filter_fn
+
+        excluded = set(self._config.exclude)
+        included = set(await filter_fn(list(self._functions.keys())))
+
+        result = {}
+        for name in self._functions:
+            is_excluded = False
+            if name in excluded:
+                is_excluded = True
+            elif not await self._fn_should_be_included(name):
+                is_excluded = True
+            elif name not in included:
+                is_excluded = True
+
+            if is_excluded:
+                result[self._get_fn_name(name)] = self._functions[name]
+
+        return result
+
+    async def get_included_functions(
+        self,
+        filter_fn: Callable[[Sequence[str]], Awaitable[Sequence[str]]] | None = None,
+    ) -> dict[str, Function]:
+        """
+        Returns a dictionary of all functions in the function group which are:
+        - configured to be included and added to the global function registry
+        - not configured to be excluded.
+        - not filtered out by a filter function.
+
+        Parameters
+        ----------
+        filter_fn : Callable[[Sequence[str]], Awaitable[Sequence[str]]] | None, optional
+            A callback function to additionally filter the functions in the function group dynamically. If not provided
+            then fall back to the function group's filter function. If no filter function is set for the function group
+            all functions will be returned.
+
+        Returns
+        -------
+        dict[str, Function]
+            A dictionary of all included functions in the function group.
+
+        Raises
+        ------
+        ValueError
+            When the function group is configured to include functions that are not found in the group.
+        """
+        missing = set(self._config.include) - set(self._functions.keys())
+        if missing:
+            raise ValueError(f"Unknown included functions: {sorted(missing)}")
+
+        if filter_fn is None:
+            if self._filter_fn is None:
+
+                async def identity_filter(x: Sequence[str]) -> Sequence[str]:
+                    return x
+
+                filter_fn = identity_filter
+            else:
+                filter_fn = self._filter_fn
+
+        included = set(await filter_fn(list(self._config.include)))
+        result = {}
+        for name in included:
+            if await self._fn_should_be_included(name):
+                result[self._get_fn_name(name)] = self._functions[name]
+        return result
+
+    async def get_all_functions(
+        self,
+        filter_fn: Callable[[Sequence[str]], Awaitable[Sequence[str]]] | None = None,
+    ) -> dict[str, Function]:
+        """
+        Returns a dictionary of all functions in the function group, regardless if they are included or excluded.
+
+        If a filter function has been set, the returned functions will additionally be filtered by the callback.
+
+        Parameters
+        ----------
+        filter_fn : Callable[[Sequence[str]], Awaitable[Sequence[str]]] | None, optional
+            A callback function to additionally filter the functions in the function group dynamically. If not provided
+            then fall back to the function group's filter function. If no filter function is set for the function group
+            all functions will be returned.
+
+        Returns
+        -------
+        dict[str, Function]
+            A dictionary of all functions in the function group.
+        """
+        if filter_fn is None:
+            if self._filter_fn is None:
+
+                async def identity_filter(x: Sequence[str]) -> Sequence[str]:
+                    return x
+
+                filter_fn = identity_filter
+            else:
+                filter_fn = self._filter_fn
+
+        included = set(await filter_fn(list(self._functions.keys())))
+        result = {}
+        for name in included:
+            if await self._fn_should_be_included(name):
+                result[self._get_fn_name(name)] = self._functions[name]
+        return result
+
+    def set_filter_fn(self, filter_fn: Callable[[Sequence[str]], Awaitable[Sequence[str]]]):
+        """
+        Sets the filter function for the function group.
+
+        Parameters
+        ----------
+        filter_fn : Callable[[Sequence[str]], Awaitable[Sequence[str]]]
+            The filter function to set for the function group.
+        """
+        self._filter_fn = filter_fn
+
+    def set_per_function_filter_fn(self, name: str, filter_fn: Callable[[str], Awaitable[bool]]):
+        """
+        Sets the a per-function filter function for the a function within the function group.
+
+        Parameters
+        ----------
+        name : str
+            The name of the function.
+        filter_fn : Callable[[str], Awaitable[bool]]
+            The per-function filter function to set for the function group.
+
+        Raises
+        ------
+        ValueError
+            When the function is not found in the function group.
+        """
+        if name not in self._functions:
+            raise ValueError(f"Function {name} not found in function group {self._instance_name}")
+        self._per_function_filter_fn[name] = filter_fn
+
+    def set_instance_name(self, instance_name: str):
+        """
+        Sets the instance name for the function group.
+
+        Parameters
+        ----------
+        instance_name : str
+            The instance name to set for the function group.
+        """
+        self._instance_name = instance_name
+
+    @property
+    def instance_name(self) -> str:
+        """
+        Returns the instance name for the function group.
+        """
+        return self._instance_name

nat/builder/function_info.py

@@ -233,7 +233,7 @@ class FunctionDescriptor:
 
             is_input_typed = all([a != sig.empty for a in annotations])
 
-            input_type = tuple[*annotations] if is_input_typed else None  # noqa: syntax-error
+            input_type = tuple[*annotations] if is_input_typed else None
 
         # Get the base type here removing all annotations and async generators
         output_annotation_decomp = DecomposedType(sig.return_annotation).get_base_type()

nat/builder/intermediate_step_manager.py

@@ -16,6 +16,8 @@
 import dataclasses
 import logging
 import typing
+import weakref
+from typing import ClassVar
 
 from nat.data_models.intermediate_step import IntermediateStep
 from nat.data_models.intermediate_step import IntermediateStepPayload
@@ -46,11 +48,19 @@ class IntermediateStepManager:
     Manages updates to the NAT Event Stream for intermediate steps
     """
 
+    # Class-level tracking for debugging and monitoring
+    _instance_count: ClassVar[int] = 0
+    _active_instances: ClassVar[set[weakref.ref]] = set()
+
     def __init__(self, context_state: "ContextState"):  # noqa: F821
         self._context_state = context_state
 
         self._outstanding_start_steps: dict[str, OpenStep] = {}
 
+        # Track instance creation
+        IntermediateStepManager._instance_count += 1
+        IntermediateStepManager._active_instances.add(weakref.ref(self, self._cleanup_instance_tracking))
+
     def push_intermediate_step(self, payload: IntermediateStepPayload) -> None:
         """
         Pushes an intermediate step to the NAT Event Stream
@@ -91,7 +101,10 @@ class IntermediateStepManager:
             open_step = self._outstanding_start_steps.pop(payload.UUID, None)
 
             if (open_step is None):
-                logger.warning("Step id %s not found in outstanding start steps", payload.UUID)
+                logger.warning(
+                    "Step id %s not found in outstanding start steps. "
+                    "This may occur if the step was started in a different context or already completed.",
+                    payload.UUID)
                 return
 
             parent_step_id = open_step.step_parent_id
@@ -147,7 +160,8 @@ class IntermediateStepManager:
             if (open_step is None):
                 logger.warning(
                     "Created a chunk for step %s, but no matching start step was found. "
-                    "Chunks must be created with the same ID as the start step.",
+                    "Chunks must be created with the same ID as the start step. "
+                    "This may occur if the step was started in a different context.",
                     payload.UUID)
                 return
 
@@ -172,3 +186,25 @@ class IntermediateStepManager:
         """
 
         return self._context_state.event_stream.get().subscribe(on_next, on_error, on_complete)
+
+    @classmethod
+    def _cleanup_instance_tracking(cls, ref: weakref.ref) -> None:
+        """Cleanup callback for weakref when instance is garbage collected."""
+        cls._active_instances.discard(ref)
+
+    @classmethod
+    def get_active_instance_count(cls) -> int:
+        """Get the number of active IntermediateStepManager instances.
+
+        Returns:
+            int: Number of active instances (cleaned up automatically via weakref)
+        """
+        return len(cls._active_instances)
+
+    def get_outstanding_step_count(self) -> int:
+        """Get the number of outstanding (started but not ended) steps.
+
+        Returns:
+            int: Number of steps that have been started but not yet ended
+        """
+        return len(self._outstanding_start_steps)