mirror of
https://github.com/deepset-ai/haystack.git
synced 2025-08-15 03:57:41 +00:00
8091ced8d5
* Extract ToolsManager, add it to Agent by the composition * PR feedback Massi --------- Co-authored-by: Massimiliano Pippi <mpippi@gmail.com> Co-authored-by: Darja Fokina <daria.f93@gmail.com>
64 lines
2.2 KiB
Python
64 lines
2.2 KiB
Python
from unittest import mock
|
|
|
|
import pytest
|
|
from haystack.agents.base import ToolsManager, Tool
|
|
|
|
|
|
@pytest.fixture
|
|
def tools_manager():
|
|
tools = [
|
|
Tool(name="ToolA", pipeline_or_node=mock.Mock(), description="Tool A Description"),
|
|
Tool(name="ToolB", pipeline_or_node=mock.Mock(), description="Tool B Description"),
|
|
]
|
|
return ToolsManager(tools=tools)
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_add_tool(tools_manager):
|
|
new_tool = Tool(name="ToolC", pipeline_or_node=mock.Mock(), description="Tool C Description")
|
|
tools_manager.add_tool(new_tool)
|
|
assert "ToolC" in tools_manager.tools
|
|
assert tools_manager.tools["ToolC"] == new_tool
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_get_tool_names(tools_manager):
|
|
assert tools_manager.get_tool_names() == "ToolA, ToolB"
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_get_tools(tools_manager):
|
|
tools = tools_manager.get_tools()
|
|
assert len(tools) == 2
|
|
assert tools[0].name == "ToolA"
|
|
assert tools[1].name == "ToolB"
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_get_tool_names_with_descriptions(tools_manager):
|
|
expected_output = "ToolA: Tool A Description\n" "ToolB: Tool B Description"
|
|
assert tools_manager.get_tool_names_with_descriptions() == expected_output
|
|
|
|
|
|
@pytest.mark.unit
|
|
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",
|
|
]
|
|
for example in examples:
|
|
tool_name, tool_input = tools_manager.extract_tool_name_and_tool_input(example)
|
|
assert tool_name == "Search" and tool_input == "Where was Jeremy McKinnon born"
|
|
|
|
negative_examples = [
|
|
"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 negative_examples:
|
|
tool_name, tool_input = tools_manager.extract_tool_name_and_tool_input(example)
|
|
assert tool_name is None and tool_input is None
|