nvidia-nat 1.3.0a20250923__py3-none-any.whl → 1.3.0a20250925__py3-none-any.whl

nat/agent/react_agent/agent.py

@@ -59,6 +59,7 @@ class ReActGraphState(BaseModel):
     messages: list[BaseMessage] = Field(default_factory=list)  # input and output of the ReAct Agent
     agent_scratchpad: list[AgentAction] = Field(default_factory=list)  # agent thoughts / intermediate steps
     tool_responses: list[BaseMessage] = Field(default_factory=list)  # the responses from any tool calls
+    final_answer: str | None = Field(default=None)  # the final answer from the ReAct Agent
 
 
 class ReActAgentGraph(DualNodeAgent):
@@ -204,6 +205,7 @@ class ReActAgentGraph(DualNodeAgent):
                         # this is where we handle the final output of the Agent, we can clean-up/format/postprocess here
                         # the final answer goes in the "messages" state channel
                         state.messages += [AIMessage(content=final_answer)]
+                        state.final_answer = final_answer
                     else:
                         # the agent wants to call a tool, ensure the thoughts are preserved for the next agentic cycle
                         agent_output.log = output_message.content
@@ -242,10 +244,9 @@ class ReActAgentGraph(DualNodeAgent):
     async def conditional_edge(self, state: ReActGraphState):
         try:
             logger.debug("%s Starting the ReAct Conditional Edge", AGENT_LOG_PREFIX)
-            if len(state.messages) > 1:
-                # the ReAct Agent has finished executing, the last agent output was AgentFinish
-                last_message_content = str(state.messages[-1].content)
-                logger.debug("%s Final answer:\n%s", AGENT_LOG_PREFIX, last_message_content)
+            if state.final_answer:
+                # the ReAct Agent has finished executing
+                logger.debug("%s Final answer:\n%s", AGENT_LOG_PREFIX, state.final_answer)
                 return AgentDecision.END
             # else the agent wants to call a tool
             agent_output = state.agent_scratchpad[-1]

nat/agent/react_agent/register.py

@@ -99,7 +99,7 @@ async def react_agent_workflow(config: ReActAgentWorkflowConfig, builder: Builde
     llm = await builder.get_llm(config.llm_name, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
     # the agent can run any installed tool, simply install the tool and add it to the config file
     # the sample tool provided can easily be copied or changed
-    tools = builder.get_tools(tool_names=config.tool_names, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
+    tools = await builder.get_tools(tool_names=config.tool_names, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
     if not tools:
         raise ValueError(f"No tools specified for ReAct Agent '{config.llm_name}'")
     # configure callbacks, for sending intermediate steps
@@ -118,6 +118,17 @@ async def react_agent_workflow(config: ReActAgentWorkflowConfig, builder: Builde
         normalize_tool_input_quotes=config.normalize_tool_input_quotes).build_graph()
 
     async def _response_fn(input_message: ChatRequest) -> ChatResponse:
+        """
+        Main workflow entry function for the ReAct Agent.
+
+        This function invokes the ReAct Agent Graph and returns the response.
+
+        Args:
+            input_message (ChatRequest): The input message to process
+
+        Returns:
+            ChatResponse: The response from the agent or error message
+        """
         try:
             # initialize the starting state with the user query
             messages: list[BaseMessage] = trim_messages(messages=[m.model_dump() for m in input_message.messages],

nat/agent/reasoning_agent/reasoning_agent.py

@@ -99,7 +99,7 @@ async def build_reasoning_function(config: ReasoningFunctionConfig, builder: Bui
     llm: BaseChatModel = await builder.get_llm(config.llm_name, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
 
     # Get the augmented function's description
-    augmented_function = builder.get_function(config.augmented_fn)
+    augmented_function = await builder.get_function(config.augmented_fn)
 
     # For now, we rely on runtime checking for type conversion
 
@@ -119,7 +119,7 @@ async def build_reasoning_function(config: ReasoningFunctionConfig, builder: Bui
     tool_names_with_desc: list[tuple[str, str]] = []
 
     for tool in function_used_tools:
-        tool_impl = builder.get_function(tool)
+        tool_impl = await builder.get_function(tool)
         tool_names_with_desc.append((tool, tool_impl.description if hasattr(tool_impl, "description") else ""))
 
     # Draft the reasoning prompt for the augmented function

nat/agent/rewoo_agent/register.py

@@ -108,7 +108,7 @@ async def rewoo_agent_workflow(config: ReWOOAgentWorkflowConfig, builder: Builde
 
     # the agent can run any installed tool, simply install the tool and add it to the config file
     # the sample tool provided can easily be copied or changed
-    tools = builder.get_tools(tool_names=config.tool_names, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
+    tools = await builder.get_tools(tool_names=config.tool_names, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
     if not tools:
         raise ValueError(f"No tools specified for ReWOO Agent '{config.llm_name}'")
 
@@ -125,6 +125,17 @@ async def rewoo_agent_workflow(config: ReWOOAgentWorkflowConfig, builder: Builde
         raise_tool_call_error=config.raise_tool_call_error).build_graph()
 
     async def _response_fn(input_message: ChatRequest) -> ChatResponse:
+        """
+        Main workflow entry function for the ReWOO Agent.
+
+        This function invokes the ReWOO Agent Graph and returns the response.
+
+        Args:
+            input_message (ChatRequest): The input message to process
+
+        Returns:
+            ChatResponse: The response from the agent or error message
+        """
         try:
             # initialize the starting state with the user query
             messages: list[BaseMessage] = trim_messages(messages=[m.model_dump() for m in input_message.messages],

nat/agent/tool_calling_agent/register.py

@@ -22,6 +22,7 @@ from nat.builder.framework_enum import LLMFrameworkEnum
 from nat.builder.function_info import FunctionInfo
 from nat.cli.register_workflow import register_function
 from nat.data_models.agent import AgentBaseConfig
+from nat.data_models.api_server import ChatRequest
 from nat.data_models.component_ref import FunctionGroupRef
 from nat.data_models.component_ref import FunctionRef
 
@@ -38,6 +39,8 @@ class ToolCallAgentWorkflowConfig(AgentBaseConfig, name="tool_calling_agent"):
         default_factory=list, description="The list of tools to provide to the tool calling agent.")
     handle_tool_errors: bool = Field(default=True, description="Specify ability to handle tool calling errors.")
     max_iterations: int = Field(default=15, description="Number of tool calls before stoping the tool calling agent.")
+    max_history: int = Field(default=15, description="Maximum number of messages to keep in the conversation history.")
+
     system_prompt: str | None = Field(default=None, description="Provides the system prompt to use with the agent.")
     additional_instructions: str | None = Field(default=None,
                                                 description="Additional instructions appended to the system prompt.")
@@ -47,7 +50,8 @@ class ToolCallAgentWorkflowConfig(AgentBaseConfig, name="tool_calling_agent"):
 
 @register_function(config_type=ToolCallAgentWorkflowConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])
 async def tool_calling_agent_workflow(config: ToolCallAgentWorkflowConfig, builder: Builder):
-    from langchain_core.messages.human import HumanMessage
+    from langchain_core.messages import trim_messages
+    from langchain_core.messages.base import BaseMessage
     from langgraph.graph.state import CompiledStateGraph
 
     from nat.agent.base import AGENT_LOG_PREFIX
@@ -60,13 +64,13 @@ async def tool_calling_agent_workflow(config: ToolCallAgentWorkflowConfig, build
     llm = await builder.get_llm(config.llm_name, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
     # the agent can run any installed tool, simply install the tool and add it to the config file
     # the sample tools provided can easily be copied or changed
-    tools = builder.get_tools(tool_names=config.tool_names, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
+    tools = await builder.get_tools(tool_names=config.tool_names, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
     if not tools:
         raise ValueError(f"No tools specified for Tool Calling Agent '{config.llm_name}'")
 
     # convert return_direct FunctionRef objects to BaseTool objects
-    return_direct_tools = builder.get_tools(tool_names=config.return_direct,
-                                            wrapper_type=LLMFrameworkEnum.LANGCHAIN) if config.return_direct else None
+    return_direct_tools = await builder.get_tools(
+        tool_names=config.return_direct, wrapper_type=LLMFrameworkEnum.LANGCHAIN) if config.return_direct else None
 
     # construct the Tool Calling Agent Graph from the configured llm, and tools
     graph: CompiledStateGraph = await ToolCallAgentGraph(llm=llm,
@@ -77,11 +81,27 @@ async def tool_calling_agent_workflow(config: ToolCallAgentWorkflowConfig, build
                                                          handle_tool_errors=config.handle_tool_errors,
                                                          return_direct=return_direct_tools).build_graph()
 
-    async def _response_fn(input_message: str) -> str:
+    async def _response_fn(input_message: ChatRequest) -> str:
+        """
+        Main workflow entry function for the Tool Calling Agent.
+
+        This function invokes the Tool Calling Agent Graph and returns the response.
+
+        Args:
+            input_message (ChatRequest): The input message to process
+
+        Returns:
+            str: The response from the agent or error message
+        """
         try:
             # initialize the starting state with the user query
-            input_message = HumanMessage(content=input_message)
-            state = ToolCallAgentGraphState(messages=[input_message])
+            messages: list[BaseMessage] = trim_messages(messages=[m.model_dump() for m in input_message.messages],
+                                                        max_tokens=config.max_history,
+                                                        strategy="last",
+                                                        token_counter=len,
+                                                        start_on="human",
+                                                        include_system=True)
+            state = ToolCallAgentGraphState(messages=messages)
 
             # run the Tool Calling Agent Graph
             state = await graph.ainvoke(state, config={'recursion_limit': (config.max_iterations + 1) * 2})
@@ -92,7 +112,7 @@ async def tool_calling_agent_workflow(config: ToolCallAgentWorkflowConfig, build
             # get and return the output from the state
             state = ToolCallAgentGraphState(**state)
             output_message = state.messages[-1]
-            return output_message.content
+            return str(output_message.content)
         except Exception as ex:
             logger.exception("%s Tool Calling Agent failed with exception: %s", AGENT_LOG_PREFIX, ex)
             if config.verbose:

nat/builder/builder.py

@@ -45,12 +45,16 @@ from nat.data_models.memory import MemoryBaseConfig
 from nat.data_models.object_store import ObjectStoreBaseConfig
 from nat.data_models.retriever import RetrieverBaseConfig
 from nat.data_models.ttc_strategy import TTCStrategyBaseConfig
+from nat.experimental.decorators.experimental_warning_decorator import experimental
 from nat.experimental.test_time_compute.models.stage_enums import PipelineTypeEnum
 from nat.experimental.test_time_compute.models.stage_enums import StageTypeEnum
 from nat.memory.interfaces import MemoryEditor
 from nat.object_store.interfaces import ObjectStore
 from nat.retriever.interface import Retriever
 
+if typing.TYPE_CHECKING:
+    from nat.experimental.test_time_compute.models.strategy_base import StrategyBase
+
 
 class UserManagerHolder():
 
@@ -72,19 +76,20 @@ class Builder(ABC):
         pass
 
     @abstractmethod
-    def get_function(self, name: str | FunctionRef) -> Function:
+    async def get_function(self, name: str | FunctionRef) -> Function:
         pass
 
     @abstractmethod
-    def get_function_group(self, name: str | FunctionGroupRef) -> FunctionGroup:
+    async def get_function_group(self, name: str | FunctionGroupRef) -> FunctionGroup:
         pass
 
-    def get_functions(self, function_names: Sequence[str | FunctionRef]) -> list[Function]:
-
-        return [self.get_function(name) for name in function_names]
+    async def get_functions(self, function_names: Sequence[str | FunctionRef]) -> list[Function]:
+        tasks = [self.get_function(name) for name in function_names]
+        return list(await asyncio.gather(*tasks, return_exceptions=False))
 
-    def get_function_groups(self, function_group_names: Sequence[str | FunctionGroupRef]) -> list[FunctionGroup]:
-        return [self.get_function_group(name) for name in function_group_names]
+    async def get_function_groups(self, function_group_names: Sequence[str | FunctionGroupRef]) -> list[FunctionGroup]:
+        tasks = [self.get_function_group(name) for name in function_group_names]
+        return list(await asyncio.gather(*tasks, return_exceptions=False))
 
     @abstractmethod
     def get_function_config(self, name: str | FunctionRef) -> FunctionBaseConfig:
@@ -107,17 +112,17 @@ class Builder(ABC):
         pass
 
     @abstractmethod
-    def get_tools(self,
-                  tool_names: Sequence[str | FunctionRef | FunctionGroupRef],
-                  wrapper_type: LLMFrameworkEnum | str) -> list[typing.Any]:
+    async def get_tools(self,
+                        tool_names: Sequence[str | FunctionRef | FunctionGroupRef],
+                        wrapper_type: LLMFrameworkEnum | str) -> list[typing.Any]:
         pass
 
     @abstractmethod
-    def get_tool(self, fn_name: str | FunctionRef, wrapper_type: LLMFrameworkEnum | str) -> typing.Any:
+    async def get_tool(self, fn_name: str | FunctionRef, wrapper_type: LLMFrameworkEnum | str) -> typing.Any:
         pass
 
     @abstractmethod
-    async def add_llm(self, name: str | LLMRef, config: LLMBaseConfig):
+    async def add_llm(self, name: str | LLMRef, config: LLMBaseConfig) -> typing.Any:
         pass
 
     @abstractmethod
@@ -138,7 +143,9 @@ class Builder(ABC):
         pass
 
     @abstractmethod
-    async def add_auth_provider(self, name: str | AuthenticationRef, config: AuthProviderBaseConfig):
+    @experimental(feature_name="Authentication")
+    async def add_auth_provider(self, name: str | AuthenticationRef,
+                                config: AuthProviderBaseConfig) -> AuthProviderBase:
         pass
 
     @abstractmethod
@@ -154,7 +161,7 @@ class Builder(ABC):
         return list(auth_providers)
 
     @abstractmethod
-    async def add_object_store(self, name: str | ObjectStoreRef, config: ObjectStoreBaseConfig):
+    async def add_object_store(self, name: str | ObjectStoreRef, config: ObjectStoreBaseConfig) -> ObjectStore:
         pass
 
     async def get_object_store_clients(self, object_store_names: Sequence[str | ObjectStoreRef]) -> list[ObjectStore]:
@@ -172,7 +179,7 @@ class Builder(ABC):
         pass
 
     @abstractmethod
-    async def add_embedder(self, name: str | EmbedderRef, config: EmbedderBaseConfig):
+    async def add_embedder(self, name: str | EmbedderRef, config: EmbedderBaseConfig) -> None:
         pass
 
     async def get_embedders(self, embedder_names: Sequence[str | EmbedderRef],
@@ -193,17 +200,18 @@ class Builder(ABC):
         pass
 
     @abstractmethod
-    async def add_memory_client(self, name: str | MemoryRef, config: MemoryBaseConfig):
+    async def add_memory_client(self, name: str | MemoryRef, config: MemoryBaseConfig) -> MemoryEditor:
         pass
 
-    def get_memory_clients(self, memory_names: Sequence[str | MemoryRef]) -> list[MemoryEditor]:
+    async def get_memory_clients(self, memory_names: Sequence[str | MemoryRef]) -> list[MemoryEditor]:
         """
         Return a list of memory clients for the specified names.
         """
-        return [self.get_memory_client(n) for n in memory_names]
+        tasks = [self.get_memory_client(n) for n in memory_names]
+        return list(await asyncio.gather(*tasks, return_exceptions=False))
 
     @abstractmethod
-    def get_memory_client(self, memory_name: str | MemoryRef) -> MemoryEditor:
+    async def get_memory_client(self, memory_name: str | MemoryRef) -> MemoryEditor:
         """
         Return the instantiated memory client for the given name.
         """
@@ -214,12 +222,12 @@ class Builder(ABC):
         pass
 
     @abstractmethod
-    async def add_retriever(self, name: str | RetrieverRef, config: RetrieverBaseConfig):
+    async def add_retriever(self, name: str | RetrieverRef, config: RetrieverBaseConfig) -> None:
         pass
 
     async def get_retrievers(self,
                              retriever_names: Sequence[str | RetrieverRef],
-                             wrapper_type: LLMFrameworkEnum | str | None = None):
+                             wrapper_type: LLMFrameworkEnum | str | None = None) -> list[Retriever]:
 
         tasks = [self.get_retriever(n, wrapper_type=wrapper_type) for n in retriever_names]
 
@@ -251,14 +259,15 @@ class Builder(ABC):
         pass
 
     @abstractmethod
-    async def add_ttc_strategy(self, name: str | str, config: TTCStrategyBaseConfig):
+    @experimental(feature_name="TTC")
+    async def add_ttc_strategy(self, name: str | TTCStrategyRef, config: TTCStrategyBaseConfig):
         pass
 
     @abstractmethod
     async def get_ttc_strategy(self,
                                strategy_name: str | TTCStrategyRef,
                                pipeline_type: PipelineTypeEnum,
-                               stage_type: StageTypeEnum):
+                               stage_type: StageTypeEnum) -> "StrategyBase":
         pass
 
     @abstractmethod
@@ -304,5 +313,5 @@ class EvalBuilder(ABC):
         pass
 
     @abstractmethod
-    def get_all_tools(self, wrapper_type: LLMFrameworkEnum | str) -> list[typing.Any]:
+    async def get_all_tools(self, wrapper_type: LLMFrameworkEnum | str) -> list[typing.Any]:
         pass

nat/builder/component_utils.py

@@ -158,7 +158,7 @@ def recursive_componentref_discovery(cls: TypedBaseModel, value: typing.Any,
             yield from recursive_componentref_discovery(cls, field_data, field_info.annotation)
     if (decomposed_type.is_union):
         for arg in decomposed_type.args:
-            if arg is typing.Any or (isinstance(value, DecomposedType(arg).root)):
+            if arg is typing.Any or DecomposedType(arg).is_instance(value):
                 yield from recursive_componentref_discovery(cls, value, arg)
     else:
         for arg in decomposed_type.args:

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"