From 1ac3319e45ff90c988b0b9f12d8456bd711e084d Mon Sep 17 00:00:00 2001 From: Harrison Chase Date: Sun, 15 Jan 2023 16:39:27 -0800 Subject: [PATCH] simplify parsing of the final answer (#621) --- langchain/agents/mrkl/base.py | 4 ++-- tests/unit_tests/agents/test_mrkl.py | 15 +++++++++++++++ 2 files changed, 17 insertions(+), 2 deletions(-) diff --git a/langchain/agents/mrkl/base.py b/langchain/agents/mrkl/base.py index 222dfeb7..2a7e264c 100644 --- a/langchain/agents/mrkl/base.py +++ b/langchain/agents/mrkl/base.py @@ -10,7 +10,7 @@ from langchain.agents.tools import Tool from langchain.llms.base import BaseLLM from langchain.prompts import PromptTemplate -FINAL_ANSWER_ACTION = "Final Answer: " +FINAL_ANSWER_ACTION = "Final Answer:" class ChainConfig(NamedTuple): @@ -30,7 +30,7 @@ class ChainConfig(NamedTuple): def get_action_and_input(llm_output: str) -> Tuple[str, str]: """Parse out the action and input from the LLM output.""" if FINAL_ANSWER_ACTION in llm_output: - return "Final Answer", llm_output.split(FINAL_ANSWER_ACTION)[-1] + return "Final Answer", llm_output.split(FINAL_ANSWER_ACTION)[-1].strip() regex = r"Action: (.*?)\nAction Input: (.*)" match = re.search(regex, llm_output) if not match: diff --git a/tests/unit_tests/agents/test_mrkl.py b/tests/unit_tests/agents/test_mrkl.py index d001ba8c..032bafa4 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_new_line() -> None: + """Test getting final answer.""" + llm_output = ( + "Thought: I need to search for NBA\n" + "Action: Search\n" + "Action Input: NBA\n" + "Observation: founded in 1994\n" + "Thought: I can now answer the question\n" + "Final Answer:\n1994" + ) + action, action_input = get_action_and_input(llm_output) + assert action == "Final Answer" + assert action_input == "1994" + + def test_get_final_answer_multiline() -> None: """Test getting final answer that is multiline.""" llm_output = (