improve input normalization

yczhang-nv · yczhang-nv · commit d824332649fa · 2025-08-28T21:15:53.000-07:00
Signed-off-by: Yuchen Zhang &lt;yuchenz@nvidia.com&gt;
diff --git a/docs/source/workflows/about/react-agent.md b/docs/source/workflows/about/react-agent.md
@@ -91,6 +91,8 @@ functions:
 
 * `pass_tool_call_errors_to_agent`: Defaults to `True`.  If set to `True`, the agent will pass tool call errors to the agent.  If set to `False`, the agent will raise an exception.
 
+* `replace_single_quotes_with_double_quotes_in_tool_input`: Defaults to `True`.  If set to `True`, the agent will replace single quotes with double quotes in the tool input.  This is useful for tools that expect structured json input.
+
 * `description`:  Defaults to `"ReAct Agent Workflow"`.  When the ReAct agent is configured as a function, this config option allows us to control the tool description (for example, when used as a tool within another agent).
 
 * `system_prompt`:  Optional.  Allows us to override the system prompt for the ReAct agent.
diff --git a/src/nat/agent/react_agent/agent.py b/src/nat/agent/react_agent/agent.py
@@ -77,7 +77,8 @@ def __init__(self,
                  retry_agent_response_parsing_errors: bool = True,
                  parse_agent_response_max_retries: int = 1,
                  tool_call_max_retries: int = 1,
-                 pass_tool_call_errors_to_agent: bool = True):
+                 pass_tool_call_errors_to_agent: bool = True,
+                 replace_single_quotes_with_double_quotes_in_tool_input: bool = True):
         super().__init__(llm=llm,
                          tools=tools,
                          callbacks=callbacks,
@@ -87,6 +88,8 @@ def __init__(self,
                                                  if retry_agent_response_parsing_errors else 1)
         self.tool_call_max_retries = tool_call_max_retries
         self.pass_tool_call_errors_to_agent = pass_tool_call_errors_to_agent
+        self.replace_single_quotes_with_double_quotes_in_tool_input = \
+            replace_single_quotes_with_double_quotes_in_tool_input
         logger.debug(
             "%s Filling the prompt variables 'tools' and 'tool_names', using the tools provided in the config.",
             AGENT_LOG_PREFIX)
@@ -286,35 +289,45 @@ async def tool_node(self, state: ReActGraphState):
                      agent_thoughts.tool_input)
 
         # Run the tool. Try to use structured input, if possible.
+        tool_input_str = agent_thoughts.tool_input.strip()
+
         try:
-            tool_input_str = str(agent_thoughts.tool_input).strip().replace("'", '"')
-            tool_input_dict = json.loads(tool_input_str) if tool_input_str != 'None' else tool_input_str
+            tool_input = json.loads(tool_input_str) if tool_input_str != 'None' else tool_input_str
             logger.debug("%s Successfully parsed structured tool input from Action Input", AGENT_LOG_PREFIX)
 
-            tool_response = await self._call_tool(requested_tool,
-                                                  tool_input_dict,
-                                                  RunnableConfig(callbacks=self.callbacks),
-                                                  max_retries=self.tool_call_max_retries)
-
-            if self.detailed_logs:
-                self._log_tool_response(requested_tool.name, tool_input_dict, str(tool_response.content))
-
-        except JSONDecodeError as ex:
-            logger.debug(
-                "%s Unable to parse structured tool input from Action Input. Using Action Input as is."
-                "\nParsing error: %s",
-                AGENT_LOG_PREFIX,
-                ex,
-                exc_info=True)
-            tool_input_str = str(agent_thoughts.tool_input)
-
-            tool_response = await self._call_tool(requested_tool,
-                                                  tool_input_str,
-                                                  RunnableConfig(callbacks=self.callbacks),
-                                                  max_retries=self.tool_call_max_retries)
+        except JSONDecodeError as original_ex:
+            if self.replace_single_quotes_with_double_quotes_in_tool_input:
+                # If initial JSON parsing fails, try with quote normalization as a fallback
+                normalized_str = tool_input_str.replace("'", '"')
+                try:
+                    tool_input = json.loads(normalized_str)
+                    logger.debug("%s Successfully parsed structured tool input after quote normalization",
+                                 AGENT_LOG_PREFIX)
+                except JSONDecodeError:
+                    # the quote normalization failed, use raw string input
+                    logger.debug(
+                        "%s Unable to parse structured tool input after quote normalization. Using Action Input as is."
+                        "\nParsing error: %s",
+                        AGENT_LOG_PREFIX,
+                        original_ex)
+                    tool_input = tool_input_str
+            else:
+                # use raw string input
+                logger.debug(
+                    "%s Unable to parse structured tool input from Action Input. Using Action Input as is."
+                    "\nParsing error: %s",
+                    AGENT_LOG_PREFIX,
+                    original_ex)
+                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)
 
         if self.detailed_logs:
-            self._log_tool_response(requested_tool.name, tool_input_str, str(tool_response.content))
+            self._log_tool_response(requested_tool.name, tool_input, str(tool_response.content))
 
         if not self.pass_tool_call_errors_to_agent:
             if tool_response.status == "error":
diff --git a/src/nat/agent/react_agent/register.py b/src/nat/agent/react_agent/register.py
@@ -62,6 +62,10 @@ class ReActAgentWorkflowConfig(FunctionBaseConfig, name="react_agent"):
     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.")
+    replace_single_quotes_with_double_quotes_in_tool_input: 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
@@ -107,7 +111,9 @@ async def react_agent_workflow(config: ReActAgentWorkflowConfig, builder: Builde
         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,
+        replace_single_quotes_with_double_quotes_in_tool_input=config.
+        replace_single_quotes_with_double_quotes_in_tool_input).build_graph()
 
     async def _response_fn(input_message: ChatRequest) -> ChatResponse:
         try:
diff --git a/tests/nat/agent/test_react.py b/tests/nat/agent/test_react.py
@@ -590,3 +590,252 @@ def test_config_mixed_alias_usage():
     assert config.parse_agent_response_max_retries == 12
     assert config.max_tool_calls == 28
     assert config.tool_call_max_retries == 1  # default value
+
+
+# Tests for quote normalization in tool input parsing
+async def test_tool_node_json_input_with_double_quotes(mock_react_agent):
+    """Test that valid JSON with double quotes is parsed correctly."""
+    tool_input = '{"query": "search term", "limit": 5}'
+    mock_state = ReActGraphState(agent_scratchpad=[AgentAction(tool='Tool A', tool_input=tool_input, log='test')])
+
+    response = await mock_react_agent.tool_node(mock_state)
+    response = response.tool_responses[-1]
+
+    assert isinstance(response, ToolMessage)
+    assert response.name == "Tool A"
+    # When JSON is successfully parsed, the mock tool receives a dict and LangChain extracts the "query" value
+    assert response.content == "search term"  # The mock tool extracts the query field value
+
+
+async def test_tool_node_json_input_with_single_quotes_normalization_enabled(mock_react_agent):
+    """Test that JSON with single quotes is normalized to double quotes when normalization is enabled."""
+    # Agent should have normalization enabled by default
+    assert mock_react_agent.replace_single_quotes_with_double_quotes_in_tool_input is True
+
+    tool_input_single_quotes = "{'query': 'search term', 'limit': 5}"
+    mock_state = ReActGraphState(
+        agent_scratchpad=[AgentAction(tool='Tool A', tool_input=tool_input_single_quotes, log='test')])
+
+    response = await mock_react_agent.tool_node(mock_state)
+    response = response.tool_responses[-1]
+
+    assert isinstance(response, ToolMessage)
+    assert response.name == "Tool A"
+    # With quote normalization enabled, single quotes get normalized and JSON is parsed successfully
+    # The mock tool then receives a dict and LangChain extracts the "query" value
+    assert response.content == "search term"
+
+
+async def test_tool_node_json_input_with_single_quotes_normalization_disabled(mock_config_react_agent,
+                                                                              mock_llm,
+                                                                              mock_tool):
+    """Test that JSON with single quotes is NOT normalized when normalization is disabled."""
+    tools = [mock_tool('Tool A'), mock_tool('Tool B')]
+    prompt = create_react_agent_prompt(mock_config_react_agent)
+
+    # Create agent with quote normalization disabled
+    agent = ReActAgentGraph(llm=mock_llm,
+                            prompt=prompt,
+                            tools=tools,
+                            detailed_logs=mock_config_react_agent.verbose,
+                            replace_single_quotes_with_double_quotes_in_tool_input=False)
+
+    assert agent.replace_single_quotes_with_double_quotes_in_tool_input is False
+
+    tool_input_single_quotes = "{'query': 'search term', 'limit': 5}"
+    mock_state = ReActGraphState(
+        agent_scratchpad=[AgentAction(tool='Tool A', tool_input=tool_input_single_quotes, log='test')])
+
+    response = await agent.tool_node(mock_state)
+    response = response.tool_responses[-1]
+
+    assert isinstance(response, ToolMessage)
+    assert response.name == "Tool A"
+    # Should use the raw string input since JSON parsing fails and normalization is disabled
+    assert response.content == tool_input_single_quotes
+
+
+async def test_tool_node_invalid_json_fallback_to_string(mock_react_agent):
+    """Test that invalid JSON falls back to using the raw string input."""
+    # Invalid JSON that cannot be fixed by quote normalization
+    tool_input_invalid = "{'query': 'search term', 'limit': }"
+    mock_state = ReActGraphState(
+        agent_scratchpad=[AgentAction(tool='Tool A', tool_input=tool_input_invalid, log='test')])
+
+    response = await mock_react_agent.tool_node(mock_state)
+    response = response.tool_responses[-1]
+
+    assert isinstance(response, ToolMessage)
+    assert response.name == "Tool A"
+    # Should fall back to using the raw string
+    assert response.content == tool_input_invalid
+
+
+async def test_tool_node_string_input_no_json_parsing(mock_react_agent):
+    """Test that plain string input is used as-is without attempting JSON parsing."""
+    tool_input_string = "simple string input"
+    mock_state = ReActGraphState(
+        agent_scratchpad=[AgentAction(tool='Tool A', tool_input=tool_input_string, log='test')])
+
+    response = await mock_react_agent.tool_node(mock_state)
+    response = response.tool_responses[-1]
+
+    assert isinstance(response, ToolMessage)
+    assert response.name == "Tool A"
+    assert response.content == tool_input_string
+
+
+async def test_tool_node_none_input(mock_react_agent):
+    """Test that 'None' input is handled correctly."""
+    tool_input_none = "None"
+    mock_state = ReActGraphState(agent_scratchpad=[AgentAction(tool='Tool A', tool_input=tool_input_none, log='test')])
+
+    response = await mock_react_agent.tool_node(mock_state)
+    response = response.tool_responses[-1]
+
+    assert isinstance(response, ToolMessage)
+    assert response.name == "Tool A"
+    assert response.content == tool_input_none
+
+
+async def test_tool_node_nested_json_with_single_quotes(mock_react_agent):
+    """Test that complex nested JSON with single quotes is normalized correctly."""
+    # Complex nested JSON with single quotes - doesn't have a "query" field so would return the full dict
+    tool_input_nested = \
+        "{'user': {'name': 'John', 'preferences': {'theme': 'dark', 'notifications': True}}, 'action': 'update'}"
+    mock_state = ReActGraphState(
+        agent_scratchpad=[AgentAction(tool='Tool A', tool_input=tool_input_nested, log='test')])
+
+    response = await mock_react_agent.tool_node(mock_state)
+    response = response.tool_responses[-1]
+
+    assert isinstance(response, ToolMessage)
+    assert response.name == "Tool A"
+    # Since this JSON doesn't have a "query" field, the mock tool receives the full dict
+    # and LangChain can't extract a "query" parameter, so it falls back to default behavior
+    assert "John" in str(response.content) or isinstance(response.content, dict)
+
+
+async def test_tool_node_mixed_quotes_in_json(mock_config_react_agent, mock_llm, mock_tool):
+    """Test that JSON with mixed quotes is handled appropriately."""
+    # This creates a scenario with mixed quotes that might be challenging to normalize
+    tools = [mock_tool('Tool A')]
+    prompt = create_react_agent_prompt(mock_config_react_agent)
+
+    agent = ReActAgentGraph(llm=mock_llm, prompt=prompt, tools=tools, detailed_logs=False)
+
+    # Mixed quotes - this is challenging JSON to normalize
+    tool_input_mixed = '''{'outer': "inner string with 'nested quotes'", 'number': 42}'''
+    mock_state = ReActGraphState(agent_scratchpad=[AgentAction(tool='Tool A', tool_input=tool_input_mixed, log='test')])
+
+    response = await agent.tool_node(mock_state)
+    response = response.tool_responses[-1]
+
+    assert isinstance(response, ToolMessage)
+    assert response.name == "Tool A"
+    # Mixed quotes are complex to normalize, so it likely falls back to raw string input
+    assert response.content == tool_input_mixed
+
+
+async def test_tool_node_whitespace_handling(mock_react_agent):
+    """Test that whitespace in tool input is handled correctly."""
+    # Tool input with leading/trailing whitespace
+    tool_input_whitespace = "  {'query': 'search term'}  "
+    mock_state = ReActGraphState(
+        agent_scratchpad=[AgentAction(tool='Tool A', tool_input=tool_input_whitespace, log='test')])
+
+    response = await mock_react_agent.tool_node(mock_state)
+    response = response.tool_responses[-1]
+
+    assert isinstance(response, ToolMessage)
+    assert response.name == "Tool A"
+    # With whitespace trimmed and quote normalization, JSON is parsed and "query" value is extracted
+    assert response.content == "search term"
+
+
+def test_config_replace_single_quotes_default():
+    """Test that replace_single_quotes_with_double_quotes_in_tool_input defaults to True."""
+    config = ReActAgentWorkflowConfig(tool_names=['test'], llm_name='test')
+    assert config.replace_single_quotes_with_double_quotes_in_tool_input is True
+
+
+def test_config_replace_single_quotes_explicit_false():
+    """Test that replace_single_quotes_with_double_quotes_in_tool_input can be set to False."""
+    config = ReActAgentWorkflowConfig(tool_names=['test'],
+                                      llm_name='test',
+                                      replace_single_quotes_with_double_quotes_in_tool_input=False)
+    assert config.replace_single_quotes_with_double_quotes_in_tool_input is False
+
+
+def test_react_agent_init_with_quote_normalization_param(mock_config_react_agent, mock_llm, mock_tool):
+    """Test that ReActAgentGraph initialization respects the quote normalization parameter."""
+    tools = [mock_tool('Tool A'), mock_tool('Tool B')]
+    prompt = create_react_agent_prompt(mock_config_react_agent)
+
+    # Test with normalization enabled
+    agent_enabled = ReActAgentGraph(llm=mock_llm,
+                                    prompt=prompt,
+                                    tools=tools,
+                                    detailed_logs=False,
+                                    replace_single_quotes_with_double_quotes_in_tool_input=True)
+    assert agent_enabled.replace_single_quotes_with_double_quotes_in_tool_input is True
+
+    # Test with normalization disabled
+    agent_disabled = ReActAgentGraph(llm=mock_llm,
+                                     prompt=prompt,
+                                     tools=tools,
+                                     detailed_logs=False,
+                                     replace_single_quotes_with_double_quotes_in_tool_input=False)
+    assert agent_disabled.replace_single_quotes_with_double_quotes_in_tool_input is False
+
+
+# Additional test to specifically verify the JSON parsing logic with quote normalization
+async def test_quote_normalization_json_parsing_logic(mock_config_react_agent, mock_llm):
+    """Test the specific quote normalization logic in JSON parsing."""
+    from langchain_core.tools import BaseTool
+
+    # Create a custom tool that returns the exact input it receives
+    class ExactInputTool(BaseTool):
+        name: str = "ExactInputTool"
+        description: str = "Returns exactly what it receives"
+
+        async def _arun(self, query, **kwargs):
+            return f"Received: {query} (type: {type(query).__name__})"
+
+        def _run(self, query, **kwargs):
+            return f"Received: {query} (type: {type(query).__name__})"
+
+    tools = [ExactInputTool()]
+    prompt = create_react_agent_prompt(mock_config_react_agent)
+
+    # Test with quote normalization enabled
+    agent_enabled = ReActAgentGraph(llm=mock_llm,
+                                    prompt=prompt,
+                                    tools=tools,
+                                    detailed_logs=False,
+                                    replace_single_quotes_with_double_quotes_in_tool_input=True)
+
+    # Test with single quotes - should be normalized and parsed as JSON
+    tool_input_single = "{'query': 'test', 'count': 42}"
+    mock_state = ReActGraphState(
+        agent_scratchpad=[AgentAction(tool='ExactInputTool', tool_input=tool_input_single, log='test')])
+    response = await agent_enabled.tool_node(mock_state)
+    response_content = response.tool_responses[-1].content
+
+    # Should receive the "query" field value from the parsed JSON dict
+    # This proves that quote normalization worked and JSON was successfully parsed
+    assert "Received: test (type: str)" in response_content
+
+    # Test with quote normalization disabled
+    agent_disabled = ReActAgentGraph(llm=mock_llm,
+                                     prompt=prompt,
+                                     tools=tools,
+                                     detailed_logs=False,
+                                     replace_single_quotes_with_double_quotes_in_tool_input=False)
+
+    response = await agent_disabled.tool_node(mock_state)
+    response_content = response.tool_responses[-1].content
+
+    # Should receive the raw string (JSON parsing failed due to no normalization)
+    # The full JSON string should be passed as the query parameter
+    assert tool_input_single in response_content and "type: str" in response_content
