From a2eeaf3d43b7738320713bc9bd49d7a5f22df615 Mon Sep 17 00:00:00 2001 From: Harrison Chase Date: Sat, 21 Jan 2023 16:03:48 -0800 Subject: [PATCH] strip whitespace (#680) --- langchain/agents/mrkl/base.py | 2 +- tests/unit_tests/agents/test_mrkl.py | 8 ++++++++ 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/langchain/agents/mrkl/base.py b/langchain/agents/mrkl/base.py index 76af6dc7..a885f54d 100644 --- a/langchain/agents/mrkl/base.py +++ b/langchain/agents/mrkl/base.py @@ -41,7 +41,7 @@ def get_action_and_input(llm_output: str) -> Tuple[str, str]: match = re.search(regex, llm_output) if not match: raise ValueError(f"Could not parse LLM output: `{llm_output}`") - action = match.group(1) + action = match.group(1).strip() 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 032bafa4..6c87f914 100644 --- a/tests/unit_tests/agents/test_mrkl.py +++ b/tests/unit_tests/agents/test_mrkl.py @@ -19,6 +19,14 @@ def test_get_action_and_input() -> None: assert action_input == "NBA" +def test_get_action_and_input_whitespace() -> None: + """Test getting an action from text.""" + llm_output = "Thought: I need to search for NBA\nAction: Search \nAction Input: NBA" + action, action_input = get_action_and_input(llm_output) + assert action == "Search" + assert action_input == "NBA" + + def test_get_final_answer() -> None: """Test getting final answer.""" llm_output = (