nvidia-nat 1.4.0a20251112__py3-none-any.whl → 1.4.0a20260113__py3-none-any.whl

aiq/__init__.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

nat/{front_ends/mcp → agent/auto_memory_wrapper}/__init__.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

nat/agent/auto_memory_wrapper/agent.py

@@ -0,0 +1,278 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026, 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 logging
+from collections.abc import Awaitable
+from collections.abc import Callable
+from typing import Any
+
+from langchain_core.messages import AIMessage
+from langchain_core.messages import BaseMessage
+from langchain_core.messages import HumanMessage
+from langchain_core.messages import SystemMessage
+from langgraph.graph import StateGraph
+from langgraph.graph.state import CompiledStateGraph
+
+from nat.builder.context import Context
+from nat.data_models.api_server import ChatRequest
+from nat.data_models.api_server import Message
+from nat.data_models.api_server import UserMessageContentRoleType
+from nat.memory.interfaces import MemoryEditor
+from nat.memory.models import MemoryItem
+
+from .state import AutoMemoryWrapperState
+
+logger = logging.getLogger(__name__)
+
+
+class AutoMemoryWrapperGraph:
+    """
+    Wraps any NAT agent to add automatic memory capture and retrieval.
+
+    The wrapper treats the inner agent as a black-box function that receives
+    a ChatRequest with multiple messages (including system messages with memory
+    context). The inner agent manages its own internal state (ReActGraphState,
+    ReWOOGraphState, etc.) - the wrapper never sees it.
+    """
+
+    def __init__(
+        self,
+        inner_agent_fn: Callable[
+            [ChatRequest], Awaitable[Any]],  # Inner agent as a Function (receives ChatRequest with multiple messages)
+        memory_editor: MemoryEditor,  # Zep/Mem0/Redis memory client
+        save_user_messages: bool = True,  # Auto-save user messages
+        retrieve_memory: bool = True,  # Auto-retrieve before agent
+        save_ai_responses: bool = True,  # Auto-save agent responses
+        search_params: dict[str, Any] | None = None,  # Backend-specific search parameters
+        add_params: dict[str, Any] | None = None  # Backend-specific add parameters
+    ):
+        self.inner_agent_fn = inner_agent_fn
+        self.memory_editor = memory_editor
+        self.save_user_messages = save_user_messages
+        self.retrieve_memory = retrieve_memory
+        self.save_ai_responses = save_ai_responses
+        self.search_params = search_params or {}
+        self.add_params = add_params or {}
+        self._context = Context.get()
+
+    def _get_user_id_from_context(self) -> str:
+        """
+        Extract user_id from runtime context.
+
+        Priority order:
+        1. user_manager.get_id() - For authenticated sessions (set via SessionManager.session())
+        2. X-User-ID HTTP header - For testing/simple auth without middleware
+        3. "default_user" - Fallback for development/testing without authentication
+
+        Returns:
+            str: The user ID for memory operations
+        """
+        # Priority 1: Get user_id from user_manager (for authenticated sessions)
+        user_manager = self._context.user_manager
+        if user_manager and hasattr(user_manager, 'get_id'):
+            try:
+                user_id = user_manager.get_id()
+                if user_id:
+                    logger.debug(f"Using user_id from user_manager: {user_id}")
+                    return user_id
+            except Exception as e:
+                logger.debug(f"Failed to get user_id from user_manager: {e}")
+
+        # Priority 2: Extract from X-User-ID HTTP header (temporary workaround for testing)
+        if self._context.metadata and self._context.metadata.headers:
+            user_id = self._context.metadata.headers.get("x-user-id")
+            if user_id:
+                logger.debug(f"Using user_id from X-User-ID header: {user_id}")
+                return user_id
+
+        # Fallback: default for development/testing
+        logger.debug("Using default user_id: default_user")
+        return "default_user"
+
+    def get_wrapper_node_count(self) -> int:
+        """
+        Calculate the number of wrapper nodes that will be added to the graph.
+
+        This is used to determine the recursion_limit overhead that the wrapper
+        adds on top of the inner agent's recursion needs.
+
+        Returns:
+            int: Number of nodes in the wrapper graph (1-4 depending on configuration)
+        """
+        count = 1  # inner_agent node is always present
+        if self.save_user_messages:
+            count += 1  # capture_user_message node
+        if self.retrieve_memory:
+            count += 1  # memory_retrieve node
+        if self.save_ai_responses:
+            count += 1  # capture_ai_response node
+        return count
+
+    @staticmethod
+    def _langchain_message_to_nat_message(lc_message: BaseMessage) -> Message:
+        """
+        Convert LangChain message to NAT Message format.
+
+        This is necessary to construct a proper ChatRequest with all messages
+        (including system messages with memory context) to pass to the inner agent.
+        """
+        if isinstance(lc_message, HumanMessage):
+            role = UserMessageContentRoleType.USER
+        elif isinstance(lc_message, AIMessage):
+            role = UserMessageContentRoleType.ASSISTANT
+        elif isinstance(lc_message, SystemMessage):
+            role = UserMessageContentRoleType.SYSTEM
+        else:
+            # Default to user for unknown message types
+            role = UserMessageContentRoleType.USER
+
+        return Message(role=role, content=str(lc_message.content))
+
+    async def capture_user_message_node(self, state: AutoMemoryWrapperState) -> AutoMemoryWrapperState:
+        """Captures user message to memory thread"""
+        if not self.save_user_messages or not state.messages:
+            return state
+
+        # Get the latest user message
+        user_message = state.messages[-1]
+        if isinstance(user_message, HumanMessage):
+            # Get user_id from runtime context
+            user_id = self._get_user_id_from_context()
+
+            # Add to memory
+            await self.memory_editor.add_items(
+                [MemoryItem(conversation=[{
+                    "role": "user", "content": str(user_message.content)
+                }], user_id=user_id)],
+                **self.add_params)
+        return state
+
+    async def memory_retrieve_node(self, state: AutoMemoryWrapperState) -> AutoMemoryWrapperState:
+        """Retrieves relevant memory from memory store"""
+        if not self.retrieve_memory or not state.messages:
+            return state
+
+        # Get the latest user message
+        user_message = state.messages[-1]
+
+        # Get user_id from runtime context
+        user_id = self._get_user_id_from_context()
+
+        # Retrieve memory from memory provider
+        memory_items = await self.memory_editor.search(
+            query=user_message.content,  # Reasonable default for memory retrieval
+            user_id=user_id,
+            **self.search_params  # User-configured params (e.g., top_k, mode)
+        )
+
+        # Extract memory strings and inject as system message if available
+        if memory_items:
+            # Extract memory field from each MemoryItem
+            memory_strings = [item.memory for item in memory_items if item.memory]
+            if memory_strings:
+                memory_text = "\n".join(memory_strings)
+                memory_message = SystemMessage(content=f"Relevant context from memory:\n{memory_text}")
+                # Insert before the last user message
+                state.messages.insert(-1, memory_message)
+
+        return state
+
+    async def inner_agent_node(self, state: AutoMemoryWrapperState) -> AutoMemoryWrapperState:
+        """
+        Calls the inner agent with a ChatRequest containing all messages.
+
+        The inner agent receives a ChatRequest with multiple messages (including
+        system messages with memory context), processes them using its own internal
+        state (ReActGraphState, ReWOOGraphState, etc.), and returns a ChatResponse.
+        """
+        # Convert all LangChain messages to NAT Message format
+        nat_messages = [self._langchain_message_to_nat_message(msg) for msg in state.messages]
+        chat_request = ChatRequest(messages=nat_messages)
+
+        # Call inner agent with ChatRequest - it manages its own state internally
+        response = await self.inner_agent_fn.ainvoke(chat_request)
+
+        # Extract content from response based on type
+        if hasattr(response, 'choices') and response.choices:
+            # ChatResponse object - extract from choices[0].message.content
+            response_text = response.choices[0].message.content or ""
+        elif hasattr(response, 'output'):
+            # GenerateResponse object - use output field
+            response_text = response.output
+        elif hasattr(response, 'value'):
+            # Some other response type with value field
+            response_text = str(response.value)
+        elif isinstance(response, str):
+            # Already a string
+            response_text = response
+        else:
+            # Last resort: convert to string
+            response_text = str(response)
+
+        # Add response to wrapper state
+        state.messages.append(AIMessage(content=response_text))
+        return state
+
+    async def capture_ai_response_node(self, state: AutoMemoryWrapperState) -> AutoMemoryWrapperState:
+        """Captures agent response to memory"""
+        if not self.save_ai_responses or not state.messages:
+            return state
+
+        # Get the latest AI message
+        ai_message = state.messages[-1]
+        if isinstance(ai_message, AIMessage):
+            # Get user_id from runtime context
+            user_id = self._get_user_id_from_context()
+
+            # Add to memory
+            await self.memory_editor.add_items(
+                [MemoryItem(conversation=[{
+                    "role": "assistant", "content": str(ai_message.content)
+                }], user_id=user_id)],
+                **self.add_params)
+        return state
+
+    def build_graph(self) -> CompiledStateGraph:
+        """Wraps inner agent with memory nodes"""
+        workflow = StateGraph(AutoMemoryWrapperState)
+
+        # Add nodes
+        if self.save_user_messages:
+            workflow.add_node("capture_user_message", self.capture_user_message_node)
+        if self.retrieve_memory:
+            workflow.add_node("memory_retrieve", self.memory_retrieve_node)
+        workflow.add_node("inner_agent", self.inner_agent_node)
+        if self.save_ai_responses:
+            workflow.add_node("capture_ai_response", self.capture_ai_response_node)
+
+        # Connect nodes based on enabled features
+        workflow.set_entry_point("capture_user_message" if self.save_user_messages else "memory_retrieve" if self.
+                                 retrieve_memory else "inner_agent")
+
+        if self.save_user_messages and self.retrieve_memory:
+            workflow.add_edge("capture_user_message", "memory_retrieve")
+            workflow.add_edge("memory_retrieve", "inner_agent")
+        elif self.save_user_messages:
+            workflow.add_edge("capture_user_message", "inner_agent")
+        elif self.retrieve_memory:
+            workflow.add_edge("memory_retrieve", "inner_agent")
+
+        if self.save_ai_responses:
+            workflow.add_edge("inner_agent", "capture_ai_response")
+            workflow.set_finish_point("capture_ai_response")
+        else:
+            workflow.set_finish_point("inner_agent")
+
+        return workflow.compile()

nat/agent/auto_memory_wrapper/register.py

@@ -0,0 +1,227 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026, 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 logging
+from collections.abc import AsyncGenerator
+from typing import Any
+
+from pydantic import Field
+
+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.component_ref import FunctionRef
+from nat.data_models.component_ref import MemoryRef
+
+logger = logging.getLogger(__name__)
+
+
+class AutoMemoryAgentConfig(AgentBaseConfig, name="auto_memory_agent"):
+    """
+    Wraps any NAT agent to provide automatic memory capture and retrieval.
+
+    This agent automatically captures user messages, retrieves relevant context,
+    and stores agent responses without requiring the LLM to invoke memory tools.
+
+    **Use this when:**
+    - You want guaranteed memory capture (not dependent on LLM tool calling)
+    - You need consistent memory operations across all interactions
+    - Your memory backend (Zep, Mem0) is designed for automatic memory management
+
+    **Use tool-based memory when:**
+    - You want the LLM to decide when to access memory
+    - Memory operations should be selective based on context
+
+    **Example:**
+
+    .. code-block:: yaml
+
+        functions:
+          my_react_agent:
+            _type: react_agent
+            llm_name: nim_llm
+            tool_names: [calculator, web_search]
+
+        memory:
+          zep_memory:
+            _type: nat.plugins.zep_cloud/zep_memory
+
+        workflow:
+          _type: auto_memory_agent
+          inner_agent_name: my_react_agent
+          memory_name: zep_memory
+          llm_name: nim_llm
+          verbose: true
+
+    **Multi-tenant User Isolation:**
+
+    User ID is automatically extracted from runtime context (user_manager.get_id()) for proper
+    multi-tenant memory isolation. Set user_manager via SessionManager.session() in production.
+    Defaults to "default_user" for testing/development. See README.md for deployment examples.
+    """
+
+    # Memory configuration
+    memory_name: MemoryRef = Field(..., description="Name of the memory backend (from memory section of config)")
+
+    # Reference to inner agent by NAME (not inline config)
+    inner_agent_name: FunctionRef = Field(..., description="Name of the agent workflow to wrap with automatic memory")
+
+    # Feature flags
+    save_user_messages_to_memory: bool = Field(
+        default=True, description="Automatically save user messages to memory before agent processing")
+    retrieve_memory_for_every_response: bool = Field(
+        default=True,
+        description=("Automatically retrieve memory context before agent processing. "
+                     "Set to false for save-only mode or when using tool-based retrieval."))
+    save_ai_messages_to_memory: bool = Field(
+        default=True, description="Automatically save AI agent responses to memory after generation")
+
+    # Memory retrieval configuration
+    search_params: dict[str, Any] = Field(
+        default_factory=dict,
+        description=(
+            "Backend-specific search parameters passed to memory_editor.search().\n"
+            "Common parameters:\n"
+            "  - top_k (int): Maximum results to return (default: 5)\n"
+            "  - mode (str): For Zep, 'basic' (fast) or 'summary' (comprehensive)\n\n"
+            "Additional parameters:\n"
+            "  - Any additional parameters that the chosen memory backend supports in its search function\n\n"))
+
+    # Memory addition configuration
+    add_params: dict[str, Any] = Field(
+        default_factory=dict,
+        description=(
+            "Backend-specific parameters passed to memory_editor.add_items().\n"
+            "For Zep:\n"
+            "  - ignore_roles (list[str]): Role types to exclude from graph memory (e.g., ['assistant'])\n"
+            "    Available roles: norole, system, assistant, user, function, tool\n\n"
+            "Additional parameters:\n"
+            "  - Any additional parameters that the chosen memory backend supports in its add_items function\n\n"))
+
+
+@register_function(config_type=AutoMemoryAgentConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])
+async def auto_memory_agent(config: AutoMemoryAgentConfig, builder: Builder) -> AsyncGenerator[FunctionInfo, None]:
+    """
+    Build the auto-memory agent that wraps another agent.
+
+    The inner agent is retrieved as a Function that receives a ChatRequest with
+    multiple messages (including system messages with memory context). It manages
+    its own internal state (ReActGraphState, etc.) and the wrapper never manipulates
+    that state.
+    """
+    from langchain_core.messages.human import HumanMessage
+    from langgraph.graph.state import CompiledStateGraph
+
+    from nat.agent.auto_memory_wrapper.agent import AutoMemoryWrapperGraph
+    from nat.agent.auto_memory_wrapper.state import AutoMemoryWrapperState
+    from nat.agent.base import AGENT_LOG_PREFIX
+
+    # Get memory editor from builder
+    memory_editor = await builder.get_memory_client(config.memory_name)
+
+    # Get inner agent as a Function (not a dict config)
+    # This gives us a function that accepts ChatRequest with multiple messages
+    inner_agent_fn = await builder.get_function(config.inner_agent_name)
+
+    # Get inner agent config to calculate recursion limits
+    inner_agent_config = builder.get_function_config(config.inner_agent_name)
+
+    # Calculate recursion_limit based on inner agent's configuration
+    # This ensures the wrapper is transparent - users only configure the inner agent's limits
+    # and the wrapper automatically accounts for its own overhead
+    inner_max_calls = None
+
+    if hasattr(inner_agent_config, 'max_tool_calls'):
+        # ReAct agent and similar agents use max_tool_calls
+        value = inner_agent_config.max_tool_calls
+        if value is not None and isinstance(value, int | float):
+            inner_max_calls = value
+
+    if inner_max_calls is None and hasattr(inner_agent_config, 'max_iterations'):
+        # Some agents use max_iterations as an alias
+        value = inner_agent_config.max_iterations
+        if value is not None and isinstance(value, int | float):
+            inner_max_calls = value
+
+    if inner_max_calls is None and hasattr(inner_agent_config, 'tool_call_max_retries'):
+        # ReWOO agent uses tool_call_max_retries - needs more steps per retry
+        value = inner_agent_config.tool_call_max_retries
+        if value is not None and isinstance(value, int | float):
+            inner_max_calls = value * 3
+
+    if inner_max_calls is None:
+        # Safe default for agents without explicit limits
+        inner_max_calls = 15
+
+    # Use same calculation formula as react_agent for consistency
+    # Formula: (max_tool_calls + 1) * 2 allows proper tool calling cycles with retries
+    # See src/nat/agent/react_agent/register.py:145 for reference
+    inner_agent_recursion = (int(inner_max_calls) + 1) * 2
+
+    # Create wrapper
+    wrapper_graph = AutoMemoryWrapperGraph(inner_agent_fn=inner_agent_fn,
+                                           memory_editor=memory_editor,
+                                           save_user_messages=config.save_user_messages_to_memory,
+                                           retrieve_memory=config.retrieve_memory_for_every_response,
+                                           save_ai_responses=config.save_ai_messages_to_memory,
+                                           search_params=config.search_params,
+                                           add_params=config.add_params)
+
+    # Calculate total recursion limit: wrapper overhead + inner agent needs
+    wrapper_node_count = wrapper_graph.get_wrapper_node_count()
+    total_recursion_limit = wrapper_node_count + inner_agent_recursion
+
+    logger.debug(f"{AGENT_LOG_PREFIX} Auto-memory wrapper calculated recursion_limit={total_recursion_limit} "
+                 f"(wrapper_overhead={wrapper_node_count} + inner_agent={inner_agent_recursion})")
+
+    # Build the graph
+    graph: CompiledStateGraph = wrapper_graph.build_graph()
+
+    async def _response_fn(input_message: str) -> str:
+        """
+        Main workflow entry function for the auto-memory agent.
+
+        Args:
+            input_message (str): The input message to process
+
+        Returns:
+            str: The response from the wrapped agent
+        """
+        try:
+            message = HumanMessage(content=input_message)
+            state = AutoMemoryWrapperState(messages=[message])
+
+            # Pass calculated recursion_limit to ensure wrapper + inner agent have enough steps
+            result_dict = await graph.ainvoke(state, config={'recursion_limit': total_recursion_limit})
+            result_state = AutoMemoryWrapperState(**result_dict)
+
+            output_message = result_state.messages[-1]
+            return str(output_message.content)
+
+        except Exception as ex:
+            logger.exception(f"{AGENT_LOG_PREFIX} Auto-memory agent failed with exception")
+            if config.verbose:
+                return str(ex)
+            return "Auto-memory agent failed"
+
+    try:
+        yield FunctionInfo.from_fn(_response_fn, description=config.description)
+    except GeneratorExit:
+        logger.debug("%s Workflow exited early!", AGENT_LOG_PREFIX)
+        raise
+    finally:
+        logger.debug("%s Cleaning up auto_memory_agent workflow.", AGENT_LOG_PREFIX)

nat/agent/auto_memory_wrapper/state.py

@@ -0,0 +1,30 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026, 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.
+
+from langchain_core.messages import BaseMessage
+from pydantic import BaseModel
+from pydantic import Field
+
+
+class AutoMemoryWrapperState(BaseModel):
+    """
+    Simple wrapper state - only needs to track messages.
+
+    The inner agent manages its own complex state internally
+    (ReActGraphState, ReWOOGraphState, etc.). The wrapper
+    never sees or manipulates the inner agent's state.
+    """
+    messages: list[BaseMessage] = Field(default_factory=list,
+                                        description="Conversation messages with context injection")

nat/agent/base.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

nat/agent/dual_node.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

nat/agent/prompt_optimizer/prompt.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

nat/agent/prompt_optimizer/register.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2021-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

nat/agent/react_agent/agent.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -33,6 +33,7 @@ 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 langgraph.runtime import DEFAULT_RUNTIME
 from pydantic import BaseModel
 from pydantic import Field
 
@@ -122,7 +123,7 @@ class ReActAgentGraph(DualNodeAgent):
         """
         # 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"]):
+        if smart_models.search(str(getattr(self.llm, "model", ""))):
             # no need to bind any additional parameters to the LLM
             return self.llm
         # add a stop sequence to the LLM
@@ -162,7 +163,8 @@ class ReActAgentGraph(DualNodeAgent):
                         {
                             "question": question, "chat_history": chat_history
                         },
-                        RunnableConfig(callbacks=self.callbacks)  # type: ignore
+                        RunnableConfig(callbacks=self.callbacks,
+                                       configurable={"__pregel_runtime": DEFAULT_RUNTIME})  # type: ignore
                     )
 
                     if self.detailed_logs:
@@ -184,10 +186,13 @@ class ReActAgentGraph(DualNodeAgent):
                     logger.debug("%s Querying agent, attempt: %s", AGENT_LOG_PREFIX, attempt)
 
                     output_message = await self._stream_llm(
-                        self.agent, {
+                        self.agent,
+                        {
                             "question": question, "agent_scratchpad": agent_scratchpad, "chat_history": chat_history
                         },
-                        RunnableConfig(callbacks=self.callbacks))
+                        RunnableConfig(callbacks=self.callbacks,
+                                       configurable={"__pregel_runtime": DEFAULT_RUNTIME})  # type: ignore
+                    )
 
                     if self.detailed_logs:
                         logger.info(AGENT_CALL_LOG_MESSAGE, question, output_message.content)
@@ -321,10 +326,12 @@ class ReActAgentGraph(DualNodeAgent):
                 tool_input = tool_input_str
 
         # Call tool once with the determined input (either parsed dict or raw string)
-        tool_response = await self._call_tool(requested_tool,
-                                              tool_input,
-                                              RunnableConfig(callbacks=self.callbacks),
-                                              max_retries=self.tool_call_max_retries)
+        tool_response = await self._call_tool(
+            requested_tool,
+            tool_input,
+            RunnableConfig(callbacks=self.callbacks,
+                           configurable={"__pregel_runtime": DEFAULT_RUNTIME}),  # type: ignore
+            max_retries=self.tool_call_max_retries)
 
         if self.detailed_logs:
             self._log_tool_response(requested_tool.name, tool_input, str(tool_response.content))

nat/agent/react_agent/output_parser.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -15,7 +15,7 @@
 
 import re
 
-from langchain.agents.agent import AgentOutputParser
+from langchain_classic.agents.agent import AgentOutputParser
 from langchain_core.agents import AgentAction
 from langchain_core.agents import AgentFinish
 from langchain_core.exceptions import LangChainException