Improving Resilience of MRKL Agent

langchain/agents/agent.py

@@ -773,7 +773,11 @@ def _take_next_step(
                 raise e
             text = str(e)
             if isinstance(self.handle_parsing_errors, bool):
-                observation = "Invalid or incomplete response"
+                if e.send_to_llm:
+                    observation = str(e.observation)
+                    text = str(e.llm_output)
+                else:
+                    observation = "Invalid or incomplete response"
             elif isinstance(self.handle_parsing_errors, str):
                 observation = self.handle_parsing_errors
             elif callable(self.handle_parsing_errors):

langchain/agents/mrkl/output_parser.py

@@ -23,7 +23,25 @@ def parse(self, text: str) -> Union[AgentAction, AgentFinish]:
         )
         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*:[\s]*(.*?)", text, re.DOTALL):
+                raise OutputParserException(
+                    f"Could not parse LLM output: `{text}`",
+                    observation="Invalid Format: Missing 'Action:' after 'Thought:'",
+                    llm_output=text,
+                    send_to_llm=True,
+                )
+            elif not re.search(
+                r"[\s]*Action\s*\d*\s*Input\s*\d*\s*:[\s]*(.*)", text, re.DOTALL
+            ):
+                raise OutputParserException(
+                    f"Could not parse LLM output: `{text}`",
+                    observation="Invalid Format:"
+                    " Missing 'Action Input:' after 'Action:'",
+                    llm_output=text,
+                    send_to_llm=True,
+                )
+            else:
+                raise OutputParserException(f"Could not parse LLM output: `{text}`")
         action = match.group(1).strip()
         action_input = match.group(2)
         return AgentAction(action, action_input.strip(" ").strip('"'), text)

langchain/schema.py

@@ -369,7 +369,23 @@ class OutputParserException(ValueError):
     errors will be raised.
     """
 
-    pass
+    def __init__(
+        self,
+        error: Any,
+        observation: str | None = None,
+        llm_output: str | None = None,
+        send_to_llm: bool = False,
+    ):
+        super(OutputParserException, self).__init__(error)
+        if send_to_llm:
+            if observation is None or llm_output is None:
+                raise ValueError(
+                    "Arguments 'observation' & 'llm_output'"
+                    " are required if 'send_to_llm' is True"
+                )
+        self.observation = observation
+        self.llm_output = llm_output
+        self.send_to_llm = send_to_llm
 
 
 class BaseDocumentTransformer(ABC):

tests/unit_tests/agents/test_mrkl.py

@@ -119,17 +119,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):
+    with pytest.raises(OutputParserException) as e_info:
         get_action_and_input(llm_output)
+    assert e_info.value.observation is not None
 
 
 def test_bad_action_line() -> None:
-    """Test handling when no action input found."""
+    """Test handling when no action found."""
     llm_output = (
         "Thought: I need to search for NBA\n" "Thought: Search\n" "Action Input: NBA"
     )
-    with pytest.raises(OutputParserException):
+    with pytest.raises(OutputParserException) as e_info:
         get_action_and_input(llm_output)
+    assert e_info.value.observation is not None
 
 
 def test_from_chains() -> None: