From 2ab13ab743b417b8edc6205a1372812f4816447b Mon Sep 17 00:00:00 2001
From: emarco177 <44670213+emarco177@users.noreply.github.com>
Date: Thu, 27 Jul 2023 23:46:06 +0300
Subject: [PATCH] added unit tests for mrkl output_parser.py (#8321)

- Description: added unit tests for mrkl output_parser.py,
  - Tag maintainer: @hinthornw
  - Twitter handle: EdenEmarco177
---
 .../langchain/agents/mrkl/output_parser.py    | 17 +++--
 .../agents/test_mrkl_output_parser.py         | 66 +++++++++++++++++++
 2 files changed, 78 insertions(+), 5 deletions(-)
 create mode 100644 libs/langchain/tests/unit_tests/agents/test_mrkl_output_parser.py

diff --git a/libs/langchain/langchain/agents/mrkl/output_parser.py b/libs/langchain/langchain/agents/mrkl/output_parser.py
index ca883fd131..37c95f46f8 100644
--- a/libs/langchain/langchain/agents/mrkl/output_parser.py
+++ b/libs/langchain/langchain/agents/mrkl/output_parser.py
@@ -6,6 +6,15 @@ from langchain.agents.mrkl.prompt import FORMAT_INSTRUCTIONS
 from langchain.schema import AgentAction, AgentFinish, OutputParserException
 
 FINAL_ANSWER_ACTION = "Final Answer:"
+MISSING_ACTION_AFTER_THOUGHT_ERROR_MESSAGE = (
+    "Invalid Format: Missing 'Action:' after 'Thought:"
+)
+MISSING_ACTION_INPUT_AFTER_ACTION_ERROR_MESSAGE = (
+    "Invalid Format: Missing 'Action Input:' after 'Action:'"
+)
+FINAL_ANSWER_AND_PARSABLE_ACTION_ERROR_MESSAGE = (
+    "Parsing LLM output produced both a final answer and a parse-able action:"
+)
 
 
 class MRKLOutputParser(AgentOutputParser):
@@ -23,8 +32,7 @@ class MRKLOutputParser(AgentOutputParser):
         if action_match:
             if includes_answer:
                 raise OutputParserException(
-                    "Parsing LLM output produced both a final answer "
-                    f"and a parse-able action: {text}"
+                    f"{FINAL_ANSWER_AND_PARSABLE_ACTION_ERROR_MESSAGE}: {text}"
                 )
             action = action_match.group(1).strip()
             action_input = action_match.group(2)
@@ -43,7 +51,7 @@ class MRKLOutputParser(AgentOutputParser):
         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:'",
+                observation=MISSING_ACTION_AFTER_THOUGHT_ERROR_MESSAGE,
                 llm_output=text,
                 send_to_llm=True,
             )
@@ -52,8 +60,7 @@ class MRKLOutputParser(AgentOutputParser):
         ):
             raise OutputParserException(
                 f"Could not parse LLM output: `{text}`",
-                observation="Invalid Format:"
-                " Missing 'Action Input:' after 'Action:'",
+                observation=MISSING_ACTION_INPUT_AFTER_ACTION_ERROR_MESSAGE,
                 llm_output=text,
                 send_to_llm=True,
             )
diff --git a/libs/langchain/tests/unit_tests/agents/test_mrkl_output_parser.py b/libs/langchain/tests/unit_tests/agents/test_mrkl_output_parser.py
new file mode 100644
index 0000000000..08a7f128a4
--- /dev/null
+++ b/libs/langchain/tests/unit_tests/agents/test_mrkl_output_parser.py
@@ -0,0 +1,66 @@
+import pytest
+
+from langchain.agents.mrkl.output_parser import (
+    MISSING_ACTION_AFTER_THOUGHT_ERROR_MESSAGE,
+    MISSING_ACTION_INPUT_AFTER_ACTION_ERROR_MESSAGE,
+    MRKLOutputParser,
+)
+from langchain.schema import AgentAction, AgentFinish, OutputParserException
+
+mrkl_output_parser = MRKLOutputParser()
+
+
+def test_valid_action_and_action_input_parse() -> None:
+    llm_output = """I can use the `foo` tool to achieve the goal.
+    Action: foo
+    Action Input: bar"""
+
+    agent_action: AgentAction = mrkl_output_parser.parse(llm_output)  # type: ignore
+    assert agent_action.tool == "foo"
+    assert agent_action.tool_input == "bar"
+
+
+def test_valid_final_answer_parse() -> None:
+    llm_output = """Final Answer: The best pizza to eat is margaritta """
+
+    agent_finish: AgentFinish = mrkl_output_parser.parse(llm_output)  # type: ignore
+    assert (
+        agent_finish.return_values.get("output")
+        == "The best pizza to eat is margaritta"
+    )
+
+
+def test_missing_action() -> None:
+    llm_output = """I can use the `foo` tool to achieve the goal."""
+
+    with pytest.raises(OutputParserException) as exception_info:
+        mrkl_output_parser.parse(llm_output)
+    assert (
+        exception_info.value.observation == MISSING_ACTION_AFTER_THOUGHT_ERROR_MESSAGE
+    )
+
+
+def test_missing_action_input() -> None:
+    llm_output = """I can use the `foo` tool to achieve the goal.
+    Action: foo"""
+
+    with pytest.raises(OutputParserException) as exception_info:
+        mrkl_output_parser.parse(llm_output)
+    assert (
+        exception_info.value.observation
+        == MISSING_ACTION_INPUT_AFTER_ACTION_ERROR_MESSAGE
+    )
+
+
+def test_final_answer_and_parsable_action() -> None:
+    llm_output = """Final Answer: The best pizza to eat is margaritta 
+        I can use the `foo` tool to achieve the goal.
+        Action: foo
+        Action Input: bar
+        """
+    with pytest.raises(OutputParserException) as exception_info:
+        mrkl_output_parser.parse(llm_output)
+    assert (
+        "Parsing LLM output produced both a final answer and a parse-able action"
+        in exception_info.value.args[0]
+    )