rapids-bot · May 13, 2026 · Apr 8, 2026 · Apr 8, 2026 · Apr 8, 2026 · Apr 8, 2026
@@ -31,6 +31,7 @@
 from langchain_core.messages.utils import convert_to_openai_messages
 from langchain_core.runnables import Runnable
 from langchain_core.runnables import RunnableConfig
+from langchain_core.runnables.config import merge_configs
 from langchain_core.tools import BaseTool
 from langgraph.graph.state import CompiledStateGraph
 from langgraph.runtime import DEFAULT_RUNTIME
@@ -110,7 +111,10 @@ def __init__(self,
         self._runnable_config = RunnableConfig(callbacks=self.callbacks,
                                                configurable={"__pregel_runtime": DEFAULT_RUNTIME})
 
-    async def _stream_llm(self, runnable: Any, inputs: dict[str, Any]) -> AIMessage:
+    async def _stream_llm(self,
+                          runnable: Any,
+                          inputs: dict[str, Any],
+                          config: RunnableConfig | None = None) -> AIMessage:
         """
         Stream from LLM runnable. Retry logic is handled automatically by the underlying LLM client.
 
@@ -119,20 +123,27 @@ async def _stream_llm(self, runnable: Any, inputs: dict[str, Any]) -> AIMessage:
         ``AIMessage`` via ``_chunk_to_message``. This ensures that native tool calling
         (``use_native_tool_calling=True``) works correctly with the ReAct agent.
 
+        When a LangGraph runtime config is provided (for example, from an injected node config),
+        it is merged with the local runnable config so that both LangGraph's streaming callbacks
+        and the profiler callbacks fire together.
+
         Parameters
         ----------
         runnable : Any
             The LLM runnable (prompt | llm or similar)
         inputs : Dict[str, Any]
             The inputs to pass to the runnable
+        config : RunnableConfig | None
+            Optional LangGraph runtime config to merge with the local runnable config.
 
         Returns
         -------
         AIMessage
             The LLM response, including any tool_calls from native tool calling.
         """
+        effective_config = merge_configs(self._runnable_config, config) if config is not None else self._runnable_config
         chunks: list[AIMessageChunk] = []
-        async for chunk in runnable.astream(inputs, config=self._runnable_config):
+        async for chunk in runnable.astream(inputs, config=effective_config):
             chunks.append(chunk)
 
         if not chunks:

@@ -18,6 +18,7 @@
 
 from langchain_core.callbacks import AsyncCallbackHandler
 from langchain_core.language_models import BaseChatModel
+from langchain_core.runnables import RunnableConfig
 from langchain_core.tools import BaseTool
 from langgraph.graph import StateGraph
 from langgraph.graph.state import CompiledStateGraph
@@ -44,7 +45,7 @@ def __init__(self,
                          log_response_max_chars=log_response_max_chars)
 
     @abstractmethod
-    async def agent_node(self, state: BaseModel) -> BaseModel:
+    async def agent_node(self, state: BaseModel, config: RunnableConfig | None = None) -> BaseModel:
         pass
 
     @abstractmethod

@@ -32,6 +32,7 @@
 from langchain_core.prompts import ChatPromptTemplate
 from langchain_core.prompts import MessagesPlaceholder
 from langchain_core.runnables import Runnable
+from langchain_core.runnables import RunnableConfig
 from langchain_core.tools import BaseTool
 from pydantic import BaseModel
 from pydantic import Field
@@ -195,7 +196,7 @@ def _parse_tool_input(self, tool_input_str: str) -> tuple[typing.Any, bool]:
 
         return tool_input_str, False
 
-    async def agent_node(self, state: ReActGraphState):
+    async def agent_node(self, state: ReActGraphState, config: RunnableConfig | None = None):
         try:
             logger.debug("%s Starting the ReAct Agent Node", AGENT_LOG_PREFIX)
             # keeping a working state allows us to resolve parsing errors without polluting the agent scratchpad
@@ -217,9 +218,8 @@ async def agent_node(self, state: ReActGraphState):
                     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, "chat_history": chat_history
-                    })  # type: ignore
+                    inputs = {"question": question, "chat_history": chat_history}
+                    output_message = await self._stream_llm(self.agent, inputs, config=config)  # type: ignore
                     if isinstance(output_message.content, str):
                         raw_content = output_message.content
                         output_message.content = remove_r1_think_tags(raw_content)
@@ -250,10 +250,8 @@ async def agent_node(self, state: ReActGraphState):
                     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, "chat_history": chat_history
-                        })  # type: ignore
+                    inputs = {"question": question, "agent_scratchpad": agent_scratchpad, "chat_history": chat_history}
+                    output_message = await self._stream_llm(self.agent, inputs, config=config)  # type: ignore
                     if isinstance(output_message.content, str):
                         raw_content = output_message.content
                         output_message.content = remove_r1_think_tags(raw_content)

@@ -31,7 +31,7 @@
 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.)
-Use the following format once you have the final answer:
+Use the following format once you have the final answer. Important: write "Final Answer:" exactly as shown, with no surrounding asterisks or markdown formatting:
 
 Thought: I now know the final answer
 Final Answer: the final answer to the original input question

@@ -14,6 +14,8 @@
 # limitations under the License.
 
 import logging
+import uuid
+from collections.abc import AsyncGenerator
 
 from pydantic import AliasChoices
 from pydantic import Field
@@ -26,12 +28,14 @@
 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 ChatResponseChunk
 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.optimizable import OptimizableField
 from nat.data_models.optimizable import OptimizableMixin
 from nat.data_models.optimizable import SearchSpace
+from nat.utils.io.model_processing import remove_r1_think_tags
 from nat.utils.type_converter import GlobalTypeConverter
 
 logger = logging.getLogger(__name__)
@@ -92,14 +96,17 @@ class ReActAgentWorkflowConfig(AgentBaseConfig, OptimizableMixin, name="react_ag
 
 @register_function(config_type=ReActAgentWorkflowConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])
 async def react_agent_workflow(config: ReActAgentWorkflowConfig, builder: Builder):
+    from langchain_core.messages import AIMessageChunk
     from langchain_core.messages import BaseMessage
     from langchain_core.messages import trim_messages
+    from langgraph.errors import GraphRecursionError
     from langgraph.graph.state import CompiledStateGraph
 
     from nat.plugins.langchain.agent.base import AGENT_LOG_PREFIX
     from nat.plugins.langchain.agent.react_agent.agent import ReActAgentGraph
     from nat.plugins.langchain.agent.react_agent.agent import ReActGraphState
     from nat.plugins.langchain.agent.react_agent.agent import create_react_agent_prompt
+    from nat.plugins.langchain.agent.react_agent.output_parser import FINAL_ANSWER_PATTERN
 
     prompt = create_react_agent_prompt(config)
 
@@ -176,4 +183,73 @@ async def _response_fn(chat_request_or_message: ChatRequestOrMessage) -> ChatRes
             logger.error("%s ReAct Agent failed with exception: %s", AGENT_LOG_PREFIX, str(ex))
             raise
 
-    yield FunctionInfo.from_fn(_response_fn, description=config.description)
+    async def _stream_fn(chat_request_or_message: ChatRequestOrMessage) -> AsyncGenerator[ChatResponseChunk]:
+        """
+        Streaming workflow entry function for the ReAct Agent.
+
+        Uses graph.astream with stream_mode="messages" to yield token-level content chunks from the LLM,
+        enabling real-time SSE streaming over the OpenAI-compatible /v1/chat/completions endpoint.
+
+        Args:
+            chat_request_or_message (ChatRequestOrMessage): The input message to process
+
+        Yields:
+            ChatResponseChunk: Streaming chunks containing content deltas
+        """
+        chunk_id = str(uuid.uuid4())
+        try:
+            message = GlobalTypeConverter.get().convert(chat_request_or_message, to_type=ChatRequest)
+            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 = ReActGraphState(messages=messages)
+
+            # buffer tokens until "Final Answer:" is found, then yield only the answer content
+            buffer = ""
+            found_final_answer = False
+
+            async for msg, metadata in graph.astream(
+                    state,
+                    config={'recursion_limit': (config.max_tool_calls + 1) * 2},
+                    stream_mode="messages"):
+                if not isinstance(msg, AIMessageChunk):
+                    continue
+                if not isinstance(metadata, dict) or metadata.get("langgraph_node") != "agent":
+                    continue
+                if isinstance(msg.content, str) and msg.content and not msg.tool_call_chunks:
+                    if found_final_answer:
+                        yield ChatResponseChunk.create_streaming_chunk(msg.content, id_=chunk_id)
+                    else:
+                        buffer += msg.content
+                        cleaned_buffer = remove_r1_think_tags(buffer)
+                        match = FINAL_ANSWER_PATTERN.search(cleaned_buffer)
+                        if match:
+                            found_final_answer = True
+                            after_marker = cleaned_buffer[match.end():]
+                            if after_marker:
+                                yield ChatResponseChunk.create_streaming_chunk(after_marker, id_=chunk_id)
+                            buffer = ""
+
+            # fallback: if the LLM answered directly without ReAct format, yield the stripped buffer
+            if not found_final_answer and buffer:
+                yield ChatResponseChunk.create_streaming_chunk(remove_r1_think_tags(buffer), id_=chunk_id)
+
+        except GraphRecursionError:
+            logger.warning(
+                "%s ReAct Agent reached its maximum iteration limit (%d) without producing a final answer. "
+                "This typically means the LLM kept calling tools instead of returning a response.",
+                AGENT_LOG_PREFIX,
+                config.max_tool_calls)
+            yield ChatResponseChunk.create_streaming_chunk(
+                f"The react agent could not produce a final answer within {config.max_tool_calls} "
+                "iterations. The agent repeatedly called tools without converging on a response.",
+                id_=chunk_id,
+            )
+        except Exception as ex:
+            logger.error("%s ReAct Agent streaming failed with exception: %s", AGENT_LOG_PREFIX, ex)
+            raise
+
+    yield FunctionInfo.create(single_fn=_response_fn, stream_fn=_stream_fn, description=config.description)
@@ -24,6 +24,7 @@
 from langchain_core.messages import ToolMessage
 from langchain_core.messages.ai import UsageMetadata
 from langchain_core.messages.base import BaseMessage
+from langchain_core.runnables import RunnableConfig
 from langchain_core.runnables import RunnableLambda
 from langchain_core.tools import BaseTool
 from langgraph.graph import StateGraph
@@ -132,7 +133,7 @@ async def _invoke_llm(self, state: ToolCallAgentGraphState):
             response = response + c
         return response
 
-    async def agent_node(self, state: ToolCallAgentGraphState):
+    async def agent_node(self, state: ToolCallAgentGraphState, config: RunnableConfig | None = None):
         try:
             logger.debug("%s Starting the Tool Calling Agent Node", AGENT_LOG_PREFIX)
             if len(state.messages) == 0:

@@ -153,6 +153,42 @@ async def mock_astream(inputs, **kwargs):
         assert isinstance(result, AIMessage)
         assert result.content == ""
 
+    async def test_stream_llm_uses_runnable_config_by_default(self, base_agent):
+        """Test that _stream_llm passes self._runnable_config to astream when config=None."""
+        captured: dict = {}
+        base_agent._runnable_config = RunnableConfig(tags=["internal"])
+
+        async def mock_astream(inputs, **kwargs):
+            captured["config"] = kwargs.get("config")
+            yield AIMessageChunk(content="hello")
+
+        mock_runnable = Mock()
+        mock_runnable.astream = mock_astream
+
+        await base_agent._stream_llm(mock_runnable, {})
+
+        assert captured["config"] is base_agent._runnable_config
+
+    async def test_stream_llm_merges_external_config(self, base_agent):
+        """Test that an external config is merged with _runnable_config before passing to astream."""
+        captured: dict = {}
+        base_agent._runnable_config = RunnableConfig(tags=["internal"])
+        external_config = RunnableConfig(tags=["external"])
+
+        async def mock_astream(inputs, **kwargs):
+            captured["config"] = kwargs.get("config")
+            yield AIMessageChunk(content="hello")
+
+        mock_runnable = Mock()
+        mock_runnable.astream = mock_astream
+
+        await base_agent._stream_llm(mock_runnable, {}, config=external_config)
+
+        effective = captured["config"]
+        assert effective is not base_agent._runnable_config, "merged config should be a new object"
+        assert "internal" in effective.get("tags", [])
+        assert "external" in effective.get("tags", [])
+
 
 class TestCallLLM:
     """Test the _call_llm method."""