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

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

nat/builder/function_base.py

@@ -111,7 +111,7 @@ class FunctionBase(typing.Generic[InputT, StreamingOutputT, SingleOutputT], ABC)
         ValueError
             If the input type cannot be determined from the class definition
         """
-        for base_cls in self.__class__.__orig_bases__:  # pylint: disable=no-member # type: ignore
+        for base_cls in self.__class__.__orig_bases__:
 
             base_cls_args = typing.get_args(base_cls)
 
@@ -196,7 +196,7 @@ class FunctionBase(typing.Generic[InputT, StreamingOutputT, SingleOutputT], ABC)
         ValueError
             If the streaming output type cannot be determined from the class definition
         """
-        for base_cls in self.__class__.__orig_bases__:  # pylint: disable=no-member # type: ignore
+        for base_cls in self.__class__.__orig_bases__:
 
             base_cls_args = typing.get_args(base_cls)
 
@@ -269,7 +269,7 @@ class FunctionBase(typing.Generic[InputT, StreamingOutputT, SingleOutputT], ABC)
         ValueError
             If the single output type cannot be determined from the class definition
         """
-        for base_cls in self.__class__.__orig_bases__:  # pylint: disable=no-member # type: ignore
+        for base_cls in self.__class__.__orig_bases__:
 
             base_cls_args = typing.get_args(base_cls)
 

nat/builder/function_info.py

@@ -231,9 +231,9 @@ class FunctionDescriptor:
         else:
             annotations = [param.annotation for param in sig.parameters.values()]
 
-            is_input_typed = all([a != sig.empty for a in annotations])  # pylint: disable=use-a-generator
+            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()
@@ -372,18 +372,16 @@ class FunctionInfo:
 
             if (stream_to_single_fn is not None):
                 raise ValueError("Cannot provide both single_fn and stream_to_single_fn")
-        else:
-            if (stream_to_single_fn is not None and stream_fn is None):
-                raise ValueError("stream_fn must be provided if stream_to_single_fn is provided")
+        elif (stream_to_single_fn is not None and stream_fn is None):
+            raise ValueError("stream_fn must be provided if stream_to_single_fn is provided")
 
         if (stream_fn is not None):
             final_stream_fn = stream_fn
 
             if (single_to_stream_fn is not None):
                 raise ValueError("Cannot provide both stream_fn and single_to_stream_fn")
-        else:
-            if (single_to_stream_fn is not None and single_fn is None):
-                raise ValueError("single_fn must be provided if single_to_stream_fn is provided")
+        elif (single_to_stream_fn is not None and single_fn is None):
+            raise ValueError("single_fn must be provided if single_to_stream_fn is provided")
 
         if (single_fn is None and stream_fn is None):
             raise ValueError("At least one of single_fn or stream_fn must be provided")

nat/builder/user_interaction_manager.py

@@ -61,13 +61,13 @@ class UserInteractionManager:
 
         uuid_req = str(uuid.uuid4())
         status = InteractionStatus.IN_PROGRESS
-        timestamp = time.strftime("%Y-%m-%dT%H:%M:%SZ")
+        timestamp = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
         sys_human_interaction = InteractionPrompt(id=uuid_req, status=status, timestamp=timestamp, content=content)
 
         resp = await self._context_state.user_input_callback.get()(sys_human_interaction)
 
         # Rebuild a InteractionResponse object with the response
-        timestamp = time.strftime("%Y-%m-%dT%H:%M:%SZ")
+        timestamp = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
         status = InteractionStatus.COMPLETED
         sys_human_interaction = InteractionResponse(id=uuid_req, status=status, timestamp=timestamp, content=resp)
 

nat/builder/workflow.py

@@ -20,6 +20,7 @@ from typing import Any
 from nat.builder.context import ContextState
 from nat.builder.embedder import EmbedderProviderInfo
 from nat.builder.function import Function
+from nat.builder.function import FunctionGroup
 from nat.builder.function_base import FunctionBase
 from nat.builder.function_base import InputT
 from nat.builder.function_base import SingleOutputT
@@ -44,6 +45,7 @@ class Workflow(FunctionBase[InputT, StreamingOutputT, SingleOutputT]):
                  config: Config,
                  entry_fn: Function[InputT, StreamingOutputT, SingleOutputT],
                  functions: dict[str, Function] | None = None,
+                 function_groups: dict[str, FunctionGroup] | None = None,
                  llms: dict[str, LLMProviderInfo] | None = None,
                  embeddings: dict[str, EmbedderProviderInfo] | None = None,
                  memory: dict[str, MemoryEditor] | None = None,
@@ -59,6 +61,7 @@ class Workflow(FunctionBase[InputT, StreamingOutputT, SingleOutputT]):
 
         self.config = config
         self.functions = functions or {}
+        self.function_groups = function_groups or {}
         self.llms = llms or {}
         self.embeddings = embeddings or {}
         self.memory = memory or {}
@@ -83,6 +86,13 @@ class Workflow(FunctionBase[InputT, StreamingOutputT, SingleOutputT]):
 
         return self._entry_fn.has_single_output
 
+    async def get_all_exporters(self) -> dict[str, BaseExporter]:
+        return await self.exporter_manager.get_all_exporters()
+
+    @property
+    def exporter_manager(self) -> ExporterManager:
+        return self._exporter_manager.get()
+
     @asynccontextmanager
     async def run(self, message: InputT):
         """
@@ -93,7 +103,7 @@ class Workflow(FunctionBase[InputT, StreamingOutputT, SingleOutputT]):
         async with Runner(input_message=message,
                           entry_fn=self._entry_fn,
                           context_state=self._context_state,
-                          exporter_manager=self._exporter_manager.get()) as runner:
+                          exporter_manager=self.exporter_manager) as runner:
 
             # The caller can `yield runner` so they can do `runner.result()` or `runner.result_stream()`
             yield runner
@@ -119,6 +129,7 @@ class Workflow(FunctionBase[InputT, StreamingOutputT, SingleOutputT]):
                       config: Config,
                       entry_fn: Function[InputT, StreamingOutputT, SingleOutputT],
                       functions: dict[str, Function] | None = None,
+                      function_groups: dict[str, FunctionGroup] | None = None,
                       llms: dict[str, LLMProviderInfo] | None = None,
                       embeddings: dict[str, EmbedderProviderInfo] | None = None,
                       memory: dict[str, MemoryEditor] | None = None,
@@ -138,6 +149,7 @@ class Workflow(FunctionBase[InputT, StreamingOutputT, SingleOutputT]):
         return WorkflowImpl(config=config,
                             entry_fn=entry_fn,
                             functions=functions,
+                            function_groups=function_groups,
                             llms=llms,
                             embeddings=embeddings,
                             memory=memory,