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

nat/agent/tool_calling_agent/register.py

@@ -21,66 +21,91 @@ 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.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.utils.type_converter import GlobalTypeConverter
 
 logger = logging.getLogger(__name__)
 
 
-class ToolCallAgentWorkflowConfig(FunctionBaseConfig, name="tool_calling_agent"):
+class ToolCallAgentWorkflowConfig(AgentBaseConfig, name="tool_calling_agent"):
     """
     A Tool Calling Agent requires an LLM which supports tool calling. A tool Calling Agent utilizes the tool
     input parameters to select the optimal tool.  Supports handling tool errors.
     """
-
-    tool_names: list[FunctionRef] = Field(default_factory=list,
-                                          description="The list of tools to provide to the tool calling agent.")
-    llm_name: LLMRef = Field(description="The LLM model to use with the tool calling agent.")
-    verbose: bool = Field(default=False, description="Set the verbosity of the tool calling agent's logging.")
-    handle_tool_errors: bool = Field(default=True, description="Specify ability to handle tool calling errors.")
     description: str = Field(default="Tool Calling Agent Workflow", description="Description of this functions use.")
+    tool_names: list[FunctionRef | FunctionGroupRef] = Field(
+        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.")
+    return_direct: list[FunctionRef] | None = Field(
+        default=None, description="List of tool names that should return responses directly without LLM processing.")
 
 
 @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 langgraph.graph.graph import CompiledGraph
+    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
+    from nat.agent.tool_calling_agent.agent import ToolCallAgentGraph
+    from nat.agent.tool_calling_agent.agent import ToolCallAgentGraphState
+    from nat.agent.tool_calling_agent.agent import create_tool_calling_agent_prompt
 
-    from .agent import ToolCallAgentGraph
-    from .agent import ToolCallAgentGraphState
-
+    prompt = create_tool_calling_agent_prompt(config)
     # we can choose an LLM for the ReAct agent in the config file
     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}'")
 
-    # some LLMs support tool calling
-    # these models accept the tool's input schema and decide when to use a tool based on the input's relevance
-    try:
-        # in tool calling agents, we bind the tools to the LLM, to pass the tools' input schemas at runtime
-        llm = llm.bind_tools(tools)
-    except NotImplementedError as ex:
-        logger.error("%s Failed to bind tools: %s", AGENT_LOG_PREFIX, ex, exc_info=True)
-        raise ex
+    # convert return_direct FunctionRef objects to BaseTool objects
+    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: CompiledGraph = await ToolCallAgentGraph(llm=llm,
-                                                    tools=tools,
-                                                    detailed_logs=config.verbose,
-                                                    handle_tool_errors=config.handle_tool_errors).build_graph()
-
-    async def _response_fn(input_message: str) -> str:
+    graph: CompiledStateGraph = await ToolCallAgentGraph(llm=llm,
+                                                         tools=tools,
+                                                         prompt=prompt,
+                                                         detailed_logs=config.verbose,
+                                                         log_response_max_chars=config.log_response_max_chars,
+                                                         handle_tool_errors=config.handle_tool_errors,
+                                                         return_direct=return_direct_tools).build_graph()
+
+    async def _response_fn(chat_request_or_message: ChatRequestOrMessage) -> str:
+        """
+        Main workflow entry function for the Tool Calling Agent.
+
+        This function invokes the Tool Calling Agent Graph and returns the response.
+
+        Args:
+            chat_request_or_message (ChatRequestOrMessage): The input message to process
+
+        Returns:
+            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
-            input_message = HumanMessage(content=input_message)
-            state = ToolCallAgentGraphState(messages=[input_message])
+            messages: list[BaseMessage] = trim_messages(messages=[m.model_dump() for m in 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})
@@ -90,17 +115,15 @@ 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]  # pylint: disable=E1136
-            return output_message.content
+            output_message = state.messages[-1]
+            return str(output_message.content)
         except Exception as ex:
-            logger.exception("%s Tool Calling Agent failed with exception: %s", AGENT_LOG_PREFIX, ex, exc_info=ex)
-            if config.verbose:
-                return str(ex)
-            return "I seem to be having a problem."
+            logger.error("%s Tool Calling Agent failed with exception: %s", AGENT_LOG_PREFIX, ex)
+            raise
 
     try:
         yield FunctionInfo.from_fn(_response_fn, description=config.description)
     except GeneratorExit:
-        logger.exception("%s Workflow exited early!", AGENT_LOG_PREFIX, exc_info=True)
+        logger.exception("%s Workflow exited early!", AGENT_LOG_PREFIX)
     finally:
         logger.debug("%s Cleaning up react_agent workflow.", AGENT_LOG_PREFIX)

nat/authentication/api_key/api_key_auth_provider.py

@@ -31,7 +31,7 @@ class APIKeyAuthProvider(AuthProviderBase[APIKeyAuthProviderConfig]):
     # fmt: off
     def __init__(self,
                  config: APIKeyAuthProviderConfig,
-                 config_name: str | None = None) -> None:  # pylint: disable=unused-argument
+                 config_name: str | None = None) -> None:
         assert isinstance(config, APIKeyAuthProviderConfig), ("Config is not APIKeyAuthProviderConfig")
         super().__init__(config)
     # fmt: on
@@ -80,7 +80,7 @@ class APIKeyAuthProvider(AuthProviderBase[APIKeyAuthProviderConfig]):
 
         raise ValueError(f"Unsupported header auth scheme: {header_auth_scheme}")
 
-    async def authenticate(self, user_id: str | None = None) -> AuthResult | None:
+    async def authenticate(self, user_id: str | None = None, **kwargs) -> AuthResult | None:
         """
         Authenticate the user using the API key credentials.