From ee3b8e89b387d83b2b2a9d93d1f723e207322632 Mon Sep 17 00:00:00 2001
From: Harrison Chase <hw.chase.17@gmail.com>
Date: Sun, 25 Dec 2022 09:53:36 -0500
Subject: [PATCH] better parsing of agent output (#418)

---
 langchain/agents/mrkl/base.py        | 25 +++++++++----------------
 tests/unit_tests/agents/test_mrkl.py | 15 +++++++++++++++
 2 files changed, 24 insertions(+), 16 deletions(-)

diff --git a/langchain/agents/mrkl/base.py b/langchain/agents/mrkl/base.py
index 5a502f46527..222dfeb749e 100644
--- a/langchain/agents/mrkl/base.py
+++ b/langchain/agents/mrkl/base.py
@@ -1,6 +1,7 @@
 """Attempt to implement MRKL systems as described in arxiv.org/pdf/2205.00445.pdf."""
 from __future__ import annotations
 
+import re
 from typing import Any, Callable, List, NamedTuple, Optional, Tuple
 
 from langchain.agents.agent import Agent, AgentExecutor
@@ -28,22 +29,14 @@ class ChainConfig(NamedTuple):
 
 def get_action_and_input(llm_output: str) -> Tuple[str, str]:
     """Parse out the action and input from the LLM output."""
-    ps = [p for p in llm_output.split("\n") if p]
-    if ps[-1].startswith("Final Answer"):
-        directive = ps[-1][len(FINAL_ANSWER_ACTION) :]
-        return "Final Answer", directive
-    if not ps[-1].startswith("Action Input: "):
-        raise ValueError(
-            "The last line does not have an action input, "
-            "something has gone terribly wrong."
-        )
-    if not ps[-2].startswith("Action: "):
-        raise ValueError(
-            "The second to last line does not have an action, "
-            "something has gone terribly wrong."
-        )
-    action = ps[-2][len("Action: ") :]
-    action_input = ps[-1][len("Action Input: ") :]
+    if FINAL_ANSWER_ACTION in llm_output:
+        return "Final Answer", llm_output.split(FINAL_ANSWER_ACTION)[-1]
+    regex = r"Action: (.*?)\nAction Input: (.*)"
+    match = re.search(regex, llm_output)
+    if not match:
+        raise ValueError(f"Could not parse LLM output: `{llm_output}`")
+    action = match.group(1)
+    action_input = match.group(2)
     return action, action_input.strip(" ").strip('"')
 
 
diff --git a/tests/unit_tests/agents/test_mrkl.py b/tests/unit_tests/agents/test_mrkl.py
index b1841c8d64c..d001ba8c810 100644
--- a/tests/unit_tests/agents/test_mrkl.py
+++ b/tests/unit_tests/agents/test_mrkl.py
@@ -34,6 +34,21 @@ def test_get_final_answer() -> None:
     assert action_input == "1994"
 
 
+def test_get_final_answer_multiline() -> None:
+    """Test getting final answer that is multiline."""
+    llm_output = (
+        "Thought: I need to search for NBA\n"
+        "Action: Search\n"
+        "Action Input: NBA\n"
+        "Observation: founded in 1994 and 1993\n"
+        "Thought: I can now answer the question\n"
+        "Final Answer: 1994\n1993"
+    )
+    action, action_input = get_action_and_input(llm_output)
+    assert action == "Final Answer"
+    assert action_input == "1994\n1993"
+
+
 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"