Improving Resilience of MRKL Agent

langchain/agents/agent.py

@@ -698,6 +698,9 @@ def _take_next_step(
                     color=color,
                     **tool_run_kwargs,
                 )
+            elif agent_action.tool == "Invalid LLM Output":
+                # Passing the error message regarding the invalid format in llm output
+                observation = str(agent_action.tool_input)
             else:
                 tool_run_kwargs = self.agent.tool_run_logging_kwargs()
                 observation = InvalidTool().run(

langchain/agents/mrkl/output_parser.py

@@ -21,7 +21,17 @@ def parse(self, text: str) -> Union[AgentAction, AgentFinish]:
         regex = r"Action\s*\d*\s*:(.*?)\nAction\s*\d*\s*Input\s*\d*\s*:[\s]*(.*)"
         match = re.search(regex, text, re.DOTALL)
         if not match:
-            raise OutputParserException(f"Could not parse LLM output: `{text}`")
+            if not re.search(r"Action\s*\d*\s*:(.*?)", text, re.DOTALL):
+                error_message = "Invalid Format: Missing 'Action:' after 'Thought:'"
+            elif not re.search(
+                r"\nAction\s*\d*\s*Input\s*\d*\s*:[\s]*(.*)", text, re.DOTALL
+            ):
+                error_message = (
+                    "Invalid Format: Missing 'Action Input:' after 'Action:'"
+                )
+            else:
+                raise OutputParserException(f"Could not parse LLM output: `{text}`")
+            return AgentAction("Invalid LLM Output", error_message, text)
         action = match.group(1).strip()
         action_input = match.group(2)
         return AgentAction(action, action_input.strip(" ").strip('"'), text)

tests/unit_tests/agents/test_mrkl.py

@@ -98,17 +98,19 @@ def test_get_final_answer_multiline() -> None:
 def test_bad_action_input_line() -> None:
     """Test handling when no action input found."""
     llm_output = "Thought: I need to search for NBA\n" "Action: Search\n" "Thought: NBA"
-    with pytest.raises(OutputParserException):
-        get_action_and_input(llm_output)
+    action, action_input = get_action_and_input(llm_output)
+    assert action == "Invalid LLM Output"
+    assert action_input == "Invalid Format: Missing 'Action Input:' after 'Action:'"
 
 
 def test_bad_action_line() -> None:
     """Test handling when no action input found."""
     llm_output = (
         "Thought: I need to search for NBA\n" "Thought: Search\n" "Action Input: NBA"
     )
-    with pytest.raises(OutputParserException):
-        get_action_and_input(llm_output)
+    action, action_input = get_action_and_input(llm_output)
+    assert action == "Invalid LLM Output"
+    assert action_input == "Invalid Format: Missing 'Action:' after 'Thought:'"
 
 
 def test_from_chains() -> None: