From b7f392fdd6edbeecf212c006ee27da5380cb4271 Mon Sep 17 00:00:00 2001 From: blob42 Date: Tue, 28 Mar 2023 06:10:34 +0000 Subject: [PATCH] [agent_executor] convenience func: lookup tool by name (#2001) A quick convenience function to lookup a tool by name Co-authored-by: blob42 --- langchain/agents/agent.py | 4 ++++ tests/unit_tests/agents/test_agent.py | 22 ++++++++++++++++++++++ 2 files changed, 26 insertions(+) diff --git a/langchain/agents/agent.py b/langchain/agents/agent.py index fa52fe16..0587228f 100644 --- a/langchain/agents/agent.py +++ b/langchain/agents/agent.py @@ -363,6 +363,10 @@ class AgentExecutor(Chain, BaseModel): else: return self.agent.return_values + def lookup_tool(self, name: str) -> BaseTool: + """Lookup tool by name.""" + return {tool.name: tool for tool in self.tools}[name] + def _should_continue(self, iterations: int) -> bool: if self.max_iterations is None: return True diff --git a/tests/unit_tests/agents/test_agent.py b/tests/unit_tests/agents/test_agent.py index e9d08317..7db3894d 100644 --- a/tests/unit_tests/agents/test_agent.py +++ b/tests/unit_tests/agents/test_agent.py @@ -289,3 +289,25 @@ def test_agent_with_new_prefix_suffix() -> None: prompt_str = agent.agent.llm_chain.prompt.template assert prompt_str.startswith(prefix), "Prompt does not start with prefix" assert prompt_str.endswith(suffix), "Prompt does not end with suffix" + + +def test_agent_lookup_tool() -> None: + """Test agent lookup tool.""" + fake_llm = FakeListLLM( + responses=["FooBarBaz\nAction: Search\nAction Input: misalignment"] + ) + tools = [ + Tool( + name="Search", + func=lambda x: x, + description="Useful for searching", + return_direct=True, + ), + ] + agent = initialize_agent( + tools=tools, + llm=fake_llm, + agent="zero-shot-react-description", + ) + + assert agent.lookup_tool("Search") == tools[0]