nvidia-nat 1.3.0.dev2__py3-none-any.whl → 1.3.0rc2__py3-none-any.whl

nat/builder/context.py

@@ -31,6 +31,7 @@ from nat.data_models.intermediate_step import IntermediateStep
 from nat.data_models.intermediate_step import IntermediateStepPayload
 from nat.data_models.intermediate_step import IntermediateStepType
 from nat.data_models.intermediate_step import StreamEventData
+from nat.data_models.intermediate_step import TraceMetadata
 from nat.data_models.invocation_node import InvocationNode
 from nat.runtime.user_metadata import RequestAttributes
 from nat.utils.reactive.subject import Subject
@@ -38,13 +39,13 @@ from nat.utils.reactive.subject import Subject
 
 class Singleton(type):
 
-    def __init__(cls, name, bases, dict):  # pylint: disable=W0622
-        super(Singleton, cls).__init__(name, bases, dict)
+    def __init__(cls, 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
 
 
@@ -65,14 +66,15 @@ 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]]
@@ -83,6 +85,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()
@@ -96,14 +122,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()
 
@@ -165,8 +191,32 @@ class Context:
         """
         return self._context_state.conversation_id.get()
 
+    @property
+    def user_message_id(self) -> str | None:
+        """
+        This property retrieves the user message ID which is the unique identifier for the current user message.
+        """
+        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, input_data: typing.Any | None):
+    def push_active_function(self,
+                             function_name: str,
+                             input_data: typing.Any | None,
+                             metadata: dict[str, typing.Any] | TraceMetadata | None = None):
         """
         Set the 'active_function' in context, push an invocation node,
         AND create an OTel child span for that function call.
@@ -187,7 +237,8 @@ class Context:
             IntermediateStepPayload(UUID=current_function_id,
                                     event_type=IntermediateStepType.FUNCTION_START,
                                     name=function_name,
-                                    data=StreamEventData(input=input_data)))
+                                    data=StreamEventData(input=input_data),
+                                    metadata=metadata))
 
         manager = ActiveFunctionContextManager()
 

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
@@ -61,7 +62,7 @@ class WorkflowEvalBuilder(WorkflowBuilder, EvalBuilder):
             # Store the evaluator
             self._evaluators[name] = ConfiguredEvaluator(config=config, instance=info_obj)
         except Exception as e:
-            logger.error("Error %s adding evaluator `%s` with config `%s`", e, name, config, exc_info=True)
+            logger.error("Error %s adding evaluator `%s` with config `%s`", e, name, config)
             raise
 
     @override
@@ -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, exc_info=True)
+                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]]
@@ -155,8 +159,9 @@ class Function(FunctionBase[InputT, StreamingOutputT, SingleOutputT], ABC):
 
                 return result
             except Exception as e:
-                logger.error("Error with ainvoke in function with input: %s.", value, exc_info=True)
-                raise e
+                err_msg = f"Error: {e}" if str(e).strip() else ""
+                logger.error("Error with ainvoke in function with input: %s. %s", value, err_msg)
+                raise
 
     @typing.final
     async def acall_invoke(self, *args, **kwargs):
@@ -186,14 +191,14 @@ class Function(FunctionBase[InputT, StreamingOutputT, SingleOutputT], ABC):
             input_obj = self.input_schema(*args, **kwargs)
 
             return await self.ainvoke(value=input_obj)
-        except Exception as e:
+        except Exception:
             logger.error(
                 "Error in acall_invoke() converting input to function schema. Both args and kwargs were "
                 "supplied which could not be converted to the input schema. args: %s\nkwargs: %s\nschema: %s",
                 args,
                 kwargs,
                 self.input_schema)
-            raise e
+            raise
 
     @abstractmethod
     async def _astream(self, value: InputT) -> AsyncGenerator[StreamingOutputT]:
@@ -252,8 +257,8 @@ class Function(FunctionBase[InputT, StreamingOutputT, SingleOutputT], ABC):
                 manager.set_output(final_output)
 
             except Exception as e:
-                logger.error("Error with astream in function with input: %s.", value, exc_info=True)
-                raise e
+                logger.error("Error with astream in function with input: %s. Error: %s", value, e)
+                raise
 
     @typing.final
     async def acall_stream(self, *args, **kwargs):
@@ -287,14 +292,14 @@ class Function(FunctionBase[InputT, StreamingOutputT, SingleOutputT], ABC):
 
                 async for x in self.astream(value=input_obj):
                     yield x
-            except Exception as e:
+            except Exception:
                 logger.error(
                     "Error in acall_stream() converting input to function schema. Both args and kwargs were "
                     "supplied which could not be converted to the input schema. args: %s\nkwargs: %s\nschema: %s",
                     args,
                     kwargs,
                     self.input_schema)
-                raise e
+                raise
 
 
 class LambdaFunction(Function[InputT, StreamingOutputT, SingleOutputT]):
@@ -342,3 +347,368 @@ 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, 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