From eb9d14faebf9927deb1c9ffdf5f8d5499b8993bf Mon Sep 17 00:00:00 2001 From: Vladimir Blagojevic Date: Thu, 18 May 2023 16:39:31 +0200 Subject: [PATCH] fix: Adjust tool pattern to support multi-line inputs (#4801) * Add support for multi line tool input * Fix failing agent test, additional test_tools_manager.py tests * Allow empty tool input, add more tests * More unit tests * String formatting * Small str fix --- haystack/agents/base.py | 4 +- test/agents/test_tools_manager.py | 64 ++++++++++++++++++++++++++++++- 2 files changed, 65 insertions(+), 3 deletions(-) diff --git a/haystack/agents/base.py b/haystack/agents/base.py index 6cd11ff7b..538caa02b 100644 --- a/haystack/agents/base.py +++ b/haystack/agents/base.py @@ -120,7 +120,7 @@ class ToolsManager: def __init__( self, tools: Optional[List[Tool]] = None, - tool_pattern: str = r'Tool:\s*(\w+)\s*Tool Input:\s*("?)([^"\n]+)\2\s*', + tool_pattern: str = r"Tool:\s*(\w+)\s*Tool Input:\s*(?:\"([\s\S]*?)\"|((?:.|\n)*))\s*", ): """ :param tools: A list of tools to add to the ToolManager. Each tool must have a unique name. @@ -198,7 +198,7 @@ class ToolsManager: tool_match = re.search(self.tool_pattern, llm_response) if tool_match: tool_name = tool_match.group(1) - tool_input = tool_match.group(3) + tool_input = tool_match.group(2) or tool_match.group(3) return tool_name.strip('" []\n').strip(), tool_input.strip('" \n') return None, None diff --git a/test/agents/test_tools_manager.py b/test/agents/test_tools_manager.py index 48ddbe125..f726a064c 100644 --- a/test/agents/test_tools_manager.py +++ b/test/agents/test_tools_manager.py @@ -48,7 +48,7 @@ def test_extract_tool_name_and_tool_input(tools_manager): examples = [ "need to find out what city he was born.\nTool: Search\nTool Input: Where was Jeremy McKinnon born", "need to find out what city he was born.\n\nTool: Search\n\nTool Input: Where was Jeremy McKinnon born", - "need to find out what city he was born. Tool: Search Tool Input: Where was Jeremy McKinnon born", + 'need to find out what city he was born. Tool: Search Tool Input: "Where was Jeremy McKinnon born"', ] for example in examples: tool_name, tool_input = tools_manager.extract_tool_name_and_tool_input(example) @@ -98,3 +98,65 @@ def test_tool_invocation(): # same but for the document with unittest.mock.patch("haystack.pipelines.Pipeline.run", return_value=[Document("mocked_document")]): assert tool.run("input") == "mocked_document" + + +@pytest.mark.unit +def test_extract_tool_name_and_tool_multi_line_input(tools_manager): + # new pattern being supported but with backward compatibility for the old + text = ( + "We need to find out the following information:\n" + "1. What city was Jeremy McKinnon born in?\n" + "2. What's the capital of Germany?\n" + "Tool: Search\n" + "Tool Input: Where was Jeremy\n McKinnon born\n and where did he grow up?" + ) + + tool_name, tool_input = tools_manager.extract_tool_name_and_tool_input(text) + assert tool_name == "Search" and tool_input == "Where was Jeremy\n McKinnon born\n and where did he grow up?" + + # tool input is empty + text2 = ( + "We need to find out the following information:\n" + "1. What city was Jeremy McKinnon born in?\n" + "2. What's the capital of Germany?\n" + "Tool: Search\n" + "Tool Input:" + ) + tool_name, tool_input = tools_manager.extract_tool_name_and_tool_input(text2) + assert tool_name == "Search" and tool_input == "" + + # Case where the tool name and tool input are provided with extra whitespaces + text3 = " Tool: Search \n Tool Input: What is the tallest building in the world? " + tool_name, tool_input = tools_manager.extract_tool_name_and_tool_input(text3) + assert tool_name.strip() == "Search" and tool_input.strip() == "What is the tallest building in the world?" + + # Case where the tool name is provided but the tool input line is not provided at all + # Tool input is not optional, so this should return None for both tool name and tool input + text4 = ( + "We need to find out the following information:\n" + "1. Who is the current president of the United States?\n" + "Tool: Search\n" + ) + tool_name, tool_input = tools_manager.extract_tool_name_and_tool_input(text4) + assert tool_name is None and tool_input is None + + # Case where neither the tool name nor the tool input is provided + text5 = "We need to find out the following information:\n 1. What is the population of India?" + tool_name, tool_input = tools_manager.extract_tool_name_and_tool_input(text5) + assert tool_name is None and tool_input is None + + # Case where the tool name and tool input are provided with extra whitespaces and new lines + text6 = " Tool: Search \n Tool Input: \nWhat is the tallest \nbuilding in the world? " + tool_name, tool_input = tools_manager.extract_tool_name_and_tool_input(text6) + assert tool_name.strip() == "Search" and tool_input.strip() == "What is the tallest \nbuilding in the world?" + + +@pytest.mark.unit +def test_extract_tool_name_and_empty_tool_input(tools_manager): + examples = [ + "need to find out what city he was born.\nTool: Search\nTool Input:", + "need to find out what city he was born.\nTool: Search\nTool Input: ", + ] + for example in examples: + tool_name, tool_input = tools_manager.extract_tool_name_and_tool_input(example) + assert tool_name == "Search" and tool_input == ""