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

nat/agent/react_agent/register.py

@@ -22,26 +22,28 @@ 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 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,7 +62,10 @@ 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
@@ -68,15 +73,21 @@ class ReActAgentWorkflowConfig(FunctionBaseConfig, name="react_agent"):
     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,23 +100,36 @@ 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(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],
@@ -125,15 +149,19 @@ 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 input_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)
+            return ChatResponse.from_string(content, usage=usage)
 
         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.")
+            logger.exception("%s ReAct Agent failed with exception: %s", AGENT_LOG_PREFIX, str(ex))
+            raise RuntimeError
 
     if (config.use_openai_api):
         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

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