nvidia-nat 1.3a20250819__py3-none-any.whl → 1.3.0a20250823__py3-none-any.whl

aiq/__init__.py

@@ -0,0 +1,66 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import importlib
+import importlib.abc
+import importlib.util
+import sys
+import warnings
+
+
+class CompatFinder(importlib.abc.MetaPathFinder):
+
+    def __init__(self, alias_prefix, target_prefix):
+        self.alias_prefix = alias_prefix
+        self.target_prefix = target_prefix
+
+    def find_spec(self, fullname, path, target=None):
+        if fullname == self.alias_prefix or fullname.startswith(self.alias_prefix + "."):
+            # Map aiq.something -> nat.something
+            target_name = self.target_prefix + fullname[len(self.alias_prefix):]
+            spec = importlib.util.find_spec(target_name)
+            if spec is None:
+                return None
+            # Wrap the loader so it loads under the alias name
+            return importlib.util.spec_from_loader(fullname, CompatLoader(fullname, target_name))
+        return None
+
+
+class CompatLoader(importlib.abc.Loader):
+
+    def __init__(self, alias_name, target_name):
+        self.alias_name = alias_name
+        self.target_name = target_name
+
+    def create_module(self, spec):
+        # Reuse the actual module so there's only one instance
+        target_module = importlib.import_module(self.target_name)
+        sys.modules[self.alias_name] = target_module
+        return target_module
+
+    def exec_module(self, module):
+        # Nothing to execute since the target is already loaded
+        pass
+
+
+# Register the compatibility finder
+sys.meta_path.insert(0, CompatFinder("aiq", "nat"))
+
+warnings.warn(
+    "!!! The 'aiq' namespace is deprecated and will be removed in a future release. "
+    "Please use the 'nat' namespace instead.",
+    DeprecationWarning,
+    stacklevel=2,
+)

nat/agent/base.py

@@ -234,6 +234,22 @@ class BaseAgent(ABC):
             logger.warning("%s Unexpected error during JSON parsing: %s", AGENT_LOG_PREFIX, str(e))
             return {"error": f"Unexpected parsing error: {str(e)}", "original_string": json_string}
 
+    def _get_chat_history(self, messages: list[BaseMessage]) -> str:
+        """
+        Get the chat history excluding the last message.
+
+        Parameters
+        ----------
+        messages : list[BaseMessage]
+            The messages to get the chat history from
+
+        Returns
+        -------
+        str
+            The chat history excluding the last message
+        """
+        return "\n".join([f"{message.type}: {message.content}" for message in messages[:-1]])
+
     @abstractmethod
     async def _build_graph(self, state_schema: type) -> CompiledGraph:
         pass

nat/agent/react_agent/agent.py

@@ -14,20 +14,23 @@
 # limitations under the License.
 
 import json
-# pylint: disable=R0917
 import logging
+import re
+import typing
 from json import JSONDecodeError
 
 from langchain_core.agents import AgentAction
 from langchain_core.agents import AgentFinish
 from langchain_core.callbacks.base import AsyncCallbackHandler
 from langchain_core.language_models import BaseChatModel
+from langchain_core.language_models import LanguageModelInput
 from langchain_core.messages.ai import AIMessage
 from langchain_core.messages.base import BaseMessage
 from langchain_core.messages.human import HumanMessage
 from langchain_core.messages.tool import ToolMessage
 from langchain_core.prompts import ChatPromptTemplate
 from langchain_core.prompts import MessagesPlaceholder
+from langchain_core.runnables import Runnable
 from langchain_core.runnables.config import RunnableConfig
 from langchain_core.tools import BaseTool
 from pydantic import BaseModel
@@ -44,7 +47,9 @@ from nat.agent.react_agent.output_parser import ReActOutputParser
 from nat.agent.react_agent.output_parser import ReActOutputParserException
 from nat.agent.react_agent.prompt import SYSTEM_PROMPT
 from nat.agent.react_agent.prompt import USER_PROMPT
-from nat.agent.react_agent.register import ReActAgentWorkflowConfig
+
+if typing.TYPE_CHECKING:
+    from nat.agent.react_agent.register import ReActAgentWorkflowConfig
 
 logger = logging.getLogger(__name__)
 
@@ -94,11 +99,27 @@ class ReActAgentGraph(DualNodeAgent):
                          f"{INPUT_SCHEMA_MESSAGE.format(schema=tools[-1].input_schema.model_fields)}")
         prompt = prompt.partial(tools=tool_names_and_descriptions, tool_names=tool_names)
         # construct the ReAct Agent
-        bound_llm = llm.bind(stop=["Observation:"])  # type: ignore
-        self.agent = prompt | bound_llm
+        self.agent = prompt | self._maybe_bind_llm_and_yield()
         self.tools_dict = {tool.name: tool for tool in tools}
         logger.debug("%s Initialized ReAct Agent Graph", AGENT_LOG_PREFIX)
 
+    def _maybe_bind_llm_and_yield(self) -> Runnable[LanguageModelInput, BaseMessage]:
+        """
+        Bind additional parameters to the LLM if needed
+        - if the LLM is a smart model, no need to bind any additional parameters
+        - if the LLM is a non-smart model, bind a stop sequence to the LLM
+
+        Returns:
+            Runnable[LanguageModelInput, BaseMessage]: The LLM with any additional parameters bound.
+        """
+        # models that don't need (or don't support)a stop sequence
+        smart_models = re.compile(r"gpt-?5", re.IGNORECASE)
+        if any(smart_models.search(getattr(self.llm, model, "")) for model in ["model", "model_name"]):
+            # no need to bind any additional parameters to the LLM
+            return self.llm
+        # add a stop sequence to the LLM
+        return self.llm.bind(stop=["Observation:"])
+
     def _get_tool(self, tool_name: str):
         try:
             return self.tools_dict.get(tool_name)
@@ -124,17 +145,19 @@ class ReActAgentGraph(DualNodeAgent):
                     if len(state.messages) == 0:
                         raise RuntimeError('No input received in state: "messages"')
                     # to check is any human input passed or not, if no input passed Agent will return the state
-                    content = str(state.messages[0].content)
+                    content = str(state.messages[-1].content)
                     if content.strip() == "":
                         logger.error("%s No human input passed to the agent.", AGENT_LOG_PREFIX)
                         state.messages += [AIMessage(content=NO_INPUT_ERROR_MESSAGE)]
                         return state
                     question = content
                     logger.debug("%s Querying agent, attempt: %s", AGENT_LOG_PREFIX, attempt)
-
+                    chat_history = self._get_chat_history(state.messages)
                     output_message = await self._stream_llm(
                         self.agent,
-                        {"question": question},
+                        {
+                            "question": question, "chat_history": chat_history
+                        },
                         RunnableConfig(callbacks=self.callbacks)  # type: ignore
                     )
 
@@ -152,13 +175,15 @@ class ReActAgentGraph(DualNodeAgent):
                         tool_response = HumanMessage(content=tool_response_content)
                         agent_scratchpad.append(tool_response)
                     agent_scratchpad += working_state
-                    question = str(state.messages[0].content)
+                    chat_history = self._get_chat_history(state.messages)
+                    question = str(state.messages[-1].content)
                     logger.debug("%s Querying agent, attempt: %s", AGENT_LOG_PREFIX, attempt)
 
-                    output_message = await self._stream_llm(self.agent, {
-                        "question": question, "agent_scratchpad": agent_scratchpad
-                    },
-                                                            RunnableConfig(callbacks=self.callbacks))
+                    output_message = await self._stream_llm(
+                        self.agent, {
+                            "question": question, "agent_scratchpad": agent_scratchpad, "chat_history": chat_history
+                        },
+                        RunnableConfig(callbacks=self.callbacks))
 
                     if self.detailed_logs:
                         logger.info(AGENT_CALL_LOG_MESSAGE, question, output_message.content)
@@ -326,7 +351,7 @@ class ReActAgentGraph(DualNodeAgent):
         return True
 
 
-def create_react_agent_prompt(config: ReActAgentWorkflowConfig) -> ChatPromptTemplate:
+def create_react_agent_prompt(config: "ReActAgentWorkflowConfig") -> ChatPromptTemplate:
     """
     Create a ReAct Agent prompt from the config.
 

nat/agent/react_agent/prompt.py

@@ -26,7 +26,7 @@ Use the following format exactly to ask the human to use a tool:
 Question: the input question you must answer
 Thought: you should always think about what to do
 Action: the action to take, should be one of [{tool_names}]
-Action Input: the input to the action (if there is no required input, include "Action Input: None")  
+Action Input: the input to the action (if there is no required input, include "Action Input: None")
 Observation: wait for the human to respond with the result from the tool, do not assume the response
 
 ... (this Thought/Action/Action Input/Observation can repeat N times. If you do not need to use a tool, or after asking the human to use any tools and waiting for the human to respond, you might know the final answer.)
@@ -37,5 +37,8 @@ Final Answer: the final answer to the original input question
 """
 
 USER_PROMPT = """
+Previous conversation history:
+{chat_history}
+
 Question: {question}
 """

nat/agent/react_agent/register.py

@@ -125,7 +125,7 @@ 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
+            output_message = state.messages[-1]
             return ChatResponse.from_string(str(output_message.content))
 
         except Exception as ex:

nat/agent/register.py

@@ -13,7 +13,6 @@
 # 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

nat/agent/rewoo_agent/agent.py

@@ -14,13 +14,13 @@
 # limitations under the License.
 
 import json
-# pylint: disable=R0917
 import logging
 from json import JSONDecodeError
 
 from langchain_core.callbacks.base import AsyncCallbackHandler
 from langchain_core.language_models import BaseChatModel
 from langchain_core.messages.ai import AIMessage
+from langchain_core.messages.base import BaseMessage
 from langchain_core.messages.human import HumanMessage
 from langchain_core.messages.tool import ToolMessage
 from langchain_core.prompts.chat import ChatPromptTemplate
@@ -43,6 +43,7 @@ logger = logging.getLogger(__name__)
 
 class ReWOOGraphState(BaseModel):
     """State schema for the ReWOO Agent Graph"""
+    messages: list[BaseMessage] = Field(default_factory=list)  # input and output of the ReWOO Agent
     task: HumanMessage = Field(default_factory=lambda: HumanMessage(content=""))  # the task provided by user
     plan: AIMessage = Field(
         default_factory=lambda: AIMessage(content=""))  # the plan generated by the planner to solve the task
@@ -183,10 +184,12 @@ class ReWOOAgentGraph(BaseAgent):
             if not task:
                 logger.error("%s No task provided to the ReWOO Agent. Please provide a valid task.", AGENT_LOG_PREFIX)
                 return {"result": NO_INPUT_ERROR_MESSAGE}
-
+            chat_history = self._get_chat_history(state.messages)
             plan = await self._stream_llm(
                 planner,
-                {"task": task},
+                {
+                    "task": task, "chat_history": chat_history
+                },
                 RunnableConfig(callbacks=self.callbacks)  # type: ignore
             )
 

nat/agent/rewoo_agent/prompt.py

@@ -87,6 +87,9 @@ Begin!
 """
 
 PLANNER_USER_PROMPT = """
+Previous conversation history:
+{chat_history}
+
 task: {task}
 """
 

nat/agent/rewoo_agent/register.py

@@ -124,15 +124,16 @@ async def rewoo_agent_workflow(config: ReWOOAgentWorkflowConfig, builder: Builde
                                                         token_counter=len,
                                                         start_on="human",
                                                         include_system=True)
-            task = HumanMessage(content=messages[0].content)
-            state = ReWOOGraphState(task=task)
+
+            task = HumanMessage(content=messages[-1].content)
+            state = ReWOOGraphState(messages=messages, task=task)
 
             # run the ReWOO Agent Graph
             state = await graph.ainvoke(state)
 
             # get and return the output from the state
             state = ReWOOGraphState(**state)
-            output_message = state.result.content  # pylint: disable=E1101
+            output_message = state.result.content
             return ChatResponse.from_string(output_message)
 
         except Exception as ex:

nat/agent/tool_calling_agent/agent.py

@@ -13,13 +13,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# pylint: disable=R0917
 import logging
+import typing
 
 from langchain_core.callbacks.base import AsyncCallbackHandler
 from langchain_core.language_models import BaseChatModel
+from langchain_core.messages import SystemMessage
 from langchain_core.messages.base import BaseMessage
-from langchain_core.runnables import RunnableConfig
+from langchain_core.runnables import RunnableLambda
+from langchain_core.runnables.config import RunnableConfig
 from langchain_core.tools import BaseTool
 from langgraph.prebuilt import ToolNode
 from pydantic import BaseModel
@@ -30,6 +32,9 @@ from nat.agent.base import AGENT_LOG_PREFIX
 from nat.agent.base import AgentDecision
 from nat.agent.dual_node import DualNodeAgent
 
+if typing.TYPE_CHECKING:
+    from nat.agent.tool_calling_agent.register import ToolCallAgentWorkflowConfig
+
 logger = logging.getLogger(__name__)
 
 
@@ -43,22 +48,51 @@ class ToolCallAgentGraph(DualNodeAgent):
     A tool Calling Agent utilizes the tool input parameters to select the optimal tool.  Supports handling tool errors.
     Argument "detailed_logs" toggles logging of inputs, outputs, and intermediate steps."""
 
-    def __init__(self,
-                 llm: BaseChatModel,
-                 tools: list[BaseTool],
-                 callbacks: list[AsyncCallbackHandler] = None,
-                 detailed_logs: bool = False,
-                 handle_tool_errors: bool = True):
+    def __init__(
+        self,
+        llm: BaseChatModel,
+        tools: list[BaseTool],
+        prompt: str | None = None,
+        callbacks: list[AsyncCallbackHandler] = None,
+        detailed_logs: bool = False,
+        handle_tool_errors: bool = True,
+    ):
         super().__init__(llm=llm, tools=tools, callbacks=callbacks, detailed_logs=detailed_logs)
+        # 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
+            self.bound_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
+
+        if prompt is not None:
+            system_prompt = SystemMessage(content=prompt)
+            prompt_runnable = RunnableLambda(
+                lambda state: [system_prompt] + state.get("messages", []),
+                name="SystemPrompt",
+            )
+        else:
+            prompt_runnable = RunnableLambda(
+                lambda state: state.get("messages", []),
+                name="PromptPassthrough",
+            )
+
+        self.agent = prompt_runnable | self.bound_llm
+
         self.tool_caller = ToolNode(tools, handle_tool_errors=handle_tool_errors)
         logger.debug("%s Initialized Tool Calling Agent Graph", AGENT_LOG_PREFIX)
 
     async def agent_node(self, state: ToolCallAgentGraphState):
         try:
-            logger.debug('%s Starting the Tool Calling Agent Node', AGENT_LOG_PREFIX)
+            logger.debug("%s Starting the Tool Calling Agent Node", AGENT_LOG_PREFIX)
             if len(state.messages) == 0:
                 raise RuntimeError('No input received in state: "messages"')
-            response = await self.llm.ainvoke(state.messages, config=RunnableConfig(callbacks=self.callbacks))
+            response = await self.agent.ainvoke(
+                {"messages": state.messages},
+                config=RunnableConfig(callbacks=self.callbacks),
+            )
             if self.detailed_logs:
                 agent_input = "\n".join(str(message.content) for message in state.messages)
                 logger.info(AGENT_CALL_LOG_MESSAGE, agent_input, response)
@@ -75,16 +109,18 @@ class ToolCallAgentGraph(DualNodeAgent):
             last_message = state.messages[-1]
             if last_message.tool_calls:
                 # the agent wants to call a tool
-                logger.debug('%s Agent is calling a tool', AGENT_LOG_PREFIX)
+                logger.debug("%s Agent is calling a tool", AGENT_LOG_PREFIX)
                 return AgentDecision.TOOL
             if self.detailed_logs:
                 logger.debug("%s Final answer:\n%s", AGENT_LOG_PREFIX, state.messages[-1].content)
             return AgentDecision.END
         except Exception as ex:
-            logger.exception("%s Failed to determine whether agent is calling a tool: %s",
-                             AGENT_LOG_PREFIX,
-                             ex,
-                             exc_info=True)
+            logger.exception(
+                "%s Failed to determine whether agent is calling a tool: %s",
+                AGENT_LOG_PREFIX,
+                ex,
+                exc_info=True,
+            )
             logger.warning("%s Ending graph traversal", AGENT_LOG_PREFIX)
             return AgentDecision.END
 
@@ -92,14 +128,15 @@ class ToolCallAgentGraph(DualNodeAgent):
         try:
             logger.debug("%s Starting Tool Node", AGENT_LOG_PREFIX)
             tool_calls = state.messages[-1].tool_calls
-            tools = [tool.get('name') for tool in tool_calls]
+            tools = [tool.get("name") for tool in tool_calls]
             tool_input = state.messages[-1]
-            tool_response = await self.tool_caller.ainvoke(input={"messages": [tool_input]},
-                                                           config=RunnableConfig(callbacks=self.callbacks,
-                                                                                 configurable={}))
+            tool_response = await self.tool_caller.ainvoke(
+                input={"messages": [tool_input]},
+                config=RunnableConfig(callbacks=self.callbacks, configurable={}),
+            )
             # this configurable = {} argument is needed due to a bug in LangGraph PreBuilt ToolNode ^
 
-            for response in tool_response.get('messages'):
+            for response in tool_response.get("messages"):
                 if self.detailed_logs:
                     self._log_tool_response(str(tools), str(tool_input), response.content)
                 state.messages += [response]
@@ -112,8 +149,41 @@ class ToolCallAgentGraph(DualNodeAgent):
     async def build_graph(self):
         try:
             await super()._build_graph(state_schema=ToolCallAgentGraphState)
-            logger.debug("%s Tool Calling Agent Graph built and compiled successfully", AGENT_LOG_PREFIX)
+            logger.debug(
+                "%s Tool Calling Agent Graph built and compiled successfully",
+                AGENT_LOG_PREFIX,
+            )
             return self.graph
         except Exception as ex:
-            logger.exception("%s Failed to build Tool Calling Agent Graph: %s", AGENT_LOG_PREFIX, ex, exc_info=ex)
+            logger.exception(
+                "%s Failed to build Tool Calling Agent Graph: %s",
+                AGENT_LOG_PREFIX,
+                ex,
+                exc_info=ex,
+            )
             raise ex
+
+
+def create_tool_calling_agent_prompt(config: "ToolCallAgentWorkflowConfig") -> str | None:
+    """
+    Create a Tool Calling Agent prompt from the config.
+
+    Args:
+        config (ToolCallAgentWorkflowConfig): The config to use for the prompt.
+
+    Returns:
+        ChatPromptTemplate: The Tool Calling Agent prompt.
+    """
+    # the Tool Calling Agent prompt can be customized via config option system_prompt and additional_instructions.
+
+    if config.system_prompt:
+        prompt_str = config.system_prompt
+    else:
+        prompt_str = ""
+
+    if config.additional_instructions:
+        prompt_str += f" {config.additional_instructions}"
+
+    if len(prompt_str) > 0:
+        return prompt_str
+    return None

nat/agent/tool_calling_agent/register.py

@@ -41,6 +41,9 @@ class ToolCallAgentWorkflowConfig(FunctionBaseConfig, name="tool_calling_agent")
     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.")
     max_iterations: int = Field(default=15, description="Number of tool calls before stoping the tool calling agent.")
+    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.")
 
 
 @register_function(config_type=ToolCallAgentWorkflowConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])
@@ -49,10 +52,11 @@ async def tool_calling_agent_workflow(config: ToolCallAgentWorkflowConfig, build
     from langgraph.graph.graph import CompiledGraph
 
     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
@@ -61,18 +65,10 @@ async def tool_calling_agent_workflow(config: ToolCallAgentWorkflowConfig, build
     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
-
     # construct the Tool Calling Agent Graph from the configured llm, and tools
     graph: CompiledGraph = await ToolCallAgentGraph(llm=llm,
                                                     tools=tools,
+                                                    prompt=prompt,
                                                     detailed_logs=config.verbose,
                                                     handle_tool_errors=config.handle_tool_errors).build_graph()
 
@@ -90,7 +86,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]  # pylint: disable=E1136
+            output_message = state.messages[-1]
             return output_message.content
         except Exception as ex:
             logger.exception("%s Tool Calling Agent failed with exception: %s", AGENT_LOG_PREFIX, ex, exc_info=ex)

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

nat/authentication/register.py

@@ -13,7 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# pylint: disable=unused-import
 # flake8: noqa
 
 from nat.authentication.api_key import register as register_api_key

nat/builder/builder.py

@@ -58,7 +58,7 @@ class UserManagerHolder():
         return self._context.user_manager.get_id()
 
 
-class Builder(ABC):  # pylint: disable=too-many-public-methods
+class Builder(ABC):
 
     @abstractmethod
     async def add_function(self, name: str | FunctionRef, config: FunctionBaseConfig) -> Function:

nat/builder/context.py

@@ -38,7 +38,7 @@ from nat.utils.reactive.subject import Subject
 
 class Singleton(type):
 
-    def __init__(cls, name, bases, dict):  # pylint: disable=W0622
+    def __init__(cls, name, bases, dict):
         super(Singleton, cls).__init__(name, bases, dict)
         cls.instance = None
 
@@ -65,6 +65,7 @@ class ContextState(metaclass=Singleton):
 
     def __init__(self):
         self.conversation_id: ContextVar[str | None] = ContextVar("conversation_id", default=None)
+        self.user_message_id: ContextVar[str | None] = ContextVar("user_message_id", default=None)
         self.input_message: ContextVar[typing.Any] = ContextVar("input_message", default=None)
         self.user_manager: ContextVar[typing.Any] = ContextVar("user_manager", default=None)
         self.metadata: ContextVar[RequestAttributes] = ContextVar("request_attributes", default=RequestAttributes())
@@ -165,6 +166,13 @@ class Context:
         """
         return self._context_state.conversation_id.get()
 
+    @property
+    def user_message_id(self) -> str | None:
+        """
+        This property retrieves the user message ID which is the unique identifier for the current user message.
+        """
+        return self._context_state.user_message_id.get()
+
     @contextmanager
     def push_active_function(self, function_name: str, input_data: typing.Any | None):
         """

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)