nvidia-nat 1.2.1rc1__py3-none-any.whl → 1.3.0__py3-none-any.whl

nat/agent/react_agent/register.py

@@ -22,26 +22,29 @@ from nat.builder.builder import Builder
 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.api_server import ChatRequestOrMessage
 from nat.data_models.api_server import ChatResponse
+from nat.data_models.api_server import Usage
+from nat.data_models.component_ref import FunctionGroupRef
 from nat.data_models.component_ref import FunctionRef
-from nat.data_models.component_ref import LLMRef
-from nat.data_models.function import FunctionBaseConfig
+from nat.data_models.optimizable import OptimizableField
+from nat.data_models.optimizable import OptimizableMixin
+from nat.data_models.optimizable import SearchSpace
 from nat.utils.type_converter import GlobalTypeConverter
 
 logger = logging.getLogger(__name__)
 
 
-class ReActAgentWorkflowConfig(FunctionBaseConfig, name="react_agent"):
+class ReActAgentWorkflowConfig(AgentBaseConfig, OptimizableMixin, name="react_agent"):
     """
     Defines a NAT function that uses a ReAct Agent performs reasoning inbetween tool calls, and utilizes the
     tool names and descriptions to select the optimal tool.
     """
-
-    tool_names: list[FunctionRef] = Field(default_factory=list,
-                                          description="The list of tools to provide to the react agent.")
-    llm_name: LLMRef = Field(description="The LLM model to use with the react agent.")
-    verbose: bool = Field(default=False, description="Set the verbosity of the react agent's logging.")
+    description: str = Field(default="ReAct Agent Workflow", description="The description of this functions use.")
+    tool_names: list[FunctionRef | FunctionGroupRef] = Field(
+        default_factory=list, description="The list of tools to provide to the react agent.")
     retry_agent_response_parsing_errors: bool = Field(
         default=True,
         validation_alias=AliasChoices("retry_agent_response_parsing_errors", "retry_parsing_errors"),
@@ -60,23 +63,29 @@ class ReActAgentWorkflowConfig(FunctionBaseConfig, name="react_agent"):
         description="Whether to pass tool call errors to agent. If False, failed tool calls will raise an exception.")
     include_tool_input_schema_in_tool_description: bool = Field(
         default=True, description="Specify inclusion of tool input schemas in the prompt.")
-    description: str = Field(default="ReAct Agent Workflow", description="The description of this functions use.")
+    normalize_tool_input_quotes: bool = Field(
+        default=True,
+        description="Whether to replace single quotes with double quotes in the tool input. "
+        "This is useful for tools that expect structured json input.")
     system_prompt: str | None = Field(
         default=None,
         description="Provides the SYSTEM_PROMPT to use with the agent")  # defaults to SYSTEM_PROMPT in prompt.py
     max_history: int = Field(default=15, description="Maximum number of messages to keep in the conversation history.")
-    use_openai_api: bool = Field(default=False,
-                                 description=("Use OpenAI API for the input/output types to the function. "
-                                              "If False, strings will be used."))
-    additional_instructions: str | None = Field(
-        default=None, description="Additional instructions to provide to the agent in addition to the base prompt.")
+    additional_instructions: str | None = OptimizableField(
+        default=None,
+        description="Additional instructions to provide to the agent in addition to the base prompt.",
+        space=SearchSpace(
+            is_prompt=True,
+            prompt="No additional instructions.",
+            prompt_purpose="Additional instructions to provide to the agent in addition to the base prompt.",
+        ))
 
 
 @register_function(config_type=ReActAgentWorkflowConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])
 async def react_agent_workflow(config: ReActAgentWorkflowConfig, builder: Builder):
     from langchain.schema import BaseMessage
     from langchain_core.messages import trim_messages
-    from langgraph.graph.graph import CompiledGraph
+    from langgraph.graph.state import CompiledStateGraph
 
     from nat.agent.base import AGENT_LOG_PREFIX
     from nat.agent.react_agent.agent import ReActAgentGraph
@@ -89,26 +98,41 @@ 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
     # construct the ReAct Agent Graph from the configured llm, prompt, and tools
-    graph: CompiledGraph = await ReActAgentGraph(
+    graph: CompiledStateGraph = await ReActAgentGraph(
         llm=llm,
         prompt=prompt,
         tools=tools,
         use_tool_schema=config.include_tool_input_schema_in_tool_description,
         detailed_logs=config.verbose,
+        log_response_max_chars=config.log_response_max_chars,
         retry_agent_response_parsing_errors=config.retry_agent_response_parsing_errors,
         parse_agent_response_max_retries=config.parse_agent_response_max_retries,
         tool_call_max_retries=config.tool_call_max_retries,
-        pass_tool_call_errors_to_agent=config.pass_tool_call_errors_to_agent).build_graph()
+        pass_tool_call_errors_to_agent=config.pass_tool_call_errors_to_agent,
+        normalize_tool_input_quotes=config.normalize_tool_input_quotes).build_graph()
+
+    async def _response_fn(chat_request_or_message: ChatRequestOrMessage) -> ChatResponse | str:
+        """
+        Main workflow entry function for the ReAct Agent.
+
+        This function invokes the ReAct Agent Graph and returns the response.
 
-    async def _response_fn(input_message: ChatRequest) -> ChatResponse:
+        Args:
+            chat_request_or_message (ChatRequestOrMessage): The input message to process
+
+        Returns:
+            ChatResponse | str: The response from the agent or error message
+        """
         try:
+            message = GlobalTypeConverter.get().convert(chat_request_or_message, to_type=ChatRequest)
+
             # initialize the starting state with the user query
-            messages: list[BaseMessage] = trim_messages(messages=[m.model_dump() for m in input_message.messages],
+            messages: list[BaseMessage] = trim_messages(messages=[m.model_dump() for m in message.messages],
                                                         max_tokens=config.max_history,
                                                         strategy="last",
                                                         token_counter=len,
@@ -125,25 +149,20 @@ async def react_agent_workflow(config: ReActAgentWorkflowConfig, builder: Builde
 
             # get and return the output from the state
             state = ReActGraphState(**state)
-            output_message = state.messages[-1]  # pylint: disable=E1136
-            return ChatResponse.from_string(str(output_message.content))
-
+            output_message = state.messages[-1]
+            content = str(output_message.content)
+
+            # Create usage statistics for the response
+            prompt_tokens = sum(len(str(msg.content).split()) for msg in message.messages)
+            completion_tokens = len(content.split()) if content else 0
+            total_tokens = prompt_tokens + completion_tokens
+            usage = Usage(prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, total_tokens=total_tokens)
+            response = ChatResponse.from_string(content, usage=usage)
+            if chat_request_or_message.is_string:
+                return GlobalTypeConverter.get().convert(response, to_type=str)
+            return response
         except Exception as ex:
-            logger.exception("%s ReAct Agent failed with exception: %s", AGENT_LOG_PREFIX, ex, exc_info=ex)
-            # here, we can implement custom error messages
-            if config.verbose:
-                return ChatResponse.from_string(str(ex))
-            return ChatResponse.from_string("I seem to be having a problem.")
-
-    if (config.use_openai_api):
-        yield FunctionInfo.from_fn(_response_fn, description=config.description)
-    else:
-
-        async def _str_api_fn(input_message: str) -> str:
-            oai_input = GlobalTypeConverter.get().try_convert(input_message, to_type=ChatRequest)
-
-            oai_output = await _response_fn(oai_input)
-
-            return GlobalTypeConverter.get().try_convert(oai_output, to_type=str)
+            logger.error("%s ReAct Agent failed with exception: %s", AGENT_LOG_PREFIX, str(ex))
+            raise
 
-        yield FunctionInfo.from_fn(_str_api_fn, description=config.description)
+    yield FunctionInfo.from_fn(_response_fn, description=config.description)

nat/agent/reasoning_agent/reasoning_agent.py

@@ -23,25 +23,22 @@ from nat.builder.builder import Builder
 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 FunctionRef
-from nat.data_models.component_ref import LLMRef
-from nat.data_models.function import FunctionBaseConfig
 
 logger = logging.getLogger(__name__)
 
 
-class ReasoningFunctionConfig(FunctionBaseConfig, name="reasoning_agent"):
+class ReasoningFunctionConfig(AgentBaseConfig, name="reasoning_agent"):
     """
     Defines a NAT function that performs reasoning on the input data.
     Output is passed to the next function in the workflow.
 
     Designed to be used with an InterceptingFunction.
     """
-
-    llm_name: LLMRef = Field(description="The name of the LLM to use for reasoning.")
+    description: str = Field(default="Reasoning Agent", description="The description of this function's use.")
     augmented_fn: FunctionRef = Field(description="The name of the function to reason on.")
-    verbose: bool = Field(default=False, description="Whether to log detailed information.")
     reasoning_prompt_template: str = Field(
         default=("You are an expert reasoning model task with creating a detailed execution plan"
                  " for a system that has the following description:\n\n"
@@ -102,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
 
@@ -113,11 +110,16 @@ async def build_reasoning_function(config: ReasoningFunctionConfig, builder: Bui
                          f"function without a description.")
 
     # Get the function dependencies of the augmented function
-    function_used_tools = builder.get_function_dependencies(config.augmented_fn).functions
+    function_dependencies = builder.get_function_dependencies(config.augmented_fn)
+    function_used_tools = set()
+    function_used_tools.update(function_dependencies.functions)
+    for function_group in function_dependencies.function_groups:
+        function_used_tools.update(builder.get_function_group_dependencies(function_group).functions)
+
     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
@@ -155,12 +157,12 @@ async def build_reasoning_function(config: ReasoningFunctionConfig, builder: Bui
             prompt = prompt.to_string()
 
             # Get the reasoning output from the LLM
-            reasoning_output = ""
+            reasoning_output = []
 
             async for chunk in llm.astream(prompt):
-                reasoning_output += chunk.content
+                reasoning_output.append(chunk.content)
 
-            reasoning_output = remove_r1_think_tags(reasoning_output)
+            reasoning_output = remove_r1_think_tags("".join(reasoning_output))
 
             output = await downstream_template.ainvoke(input={
                 "input_text": input_text, "reasoning_output": reasoning_output
@@ -198,12 +200,12 @@ async def build_reasoning_function(config: ReasoningFunctionConfig, builder: Bui
             prompt = prompt.to_string()
 
             # Get the reasoning output from the LLM
-            reasoning_output = ""
+            reasoning_output = []
 
             async for chunk in llm.astream(prompt):
-                reasoning_output += chunk.content
+                reasoning_output.append(chunk.content)
 
-            reasoning_output = remove_r1_think_tags(reasoning_output)
+            reasoning_output = remove_r1_think_tags("".join(reasoning_output))
 
             output = await downstream_template.ainvoke(input={
                 "input_text": input_text, "reasoning_output": reasoning_output

nat/agent/register.py

@@ -13,10 +13,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# pylint: disable=unused-import
 # flake8: noqa
 
 # Import any workflows which need to be automatically registered here
+from .prompt_optimizer import register as prompt_optimizer
 from .react_agent import register as react_agent
 from .reasoning_agent import reasoning_agent
 from .rewoo_agent import register as rewoo_agent