langchain/tests/unit_tests/agents/test_agent.py

262 lines
7.8 KiB
Python
Raw Normal View History

"""Unit tests for agents."""
from typing import Any, List, Mapping, Optional
from pydantic import BaseModel
from langchain.agents import AgentExecutor, initialize_agent
from langchain.agents.tools import Tool
from langchain.callbacks.base import CallbackManager
from langchain.llms.base import LLM
from tests.unit_tests.callbacks.fake_callback_handler import FakeCallbackHandler
class FakeListLLM(LLM, BaseModel):
"""Fake LLM for testing that outputs elements of a list."""
responses: List[str]
i: int = -1
2022-12-15 15:53:32 +00:00
def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
"""Increment counter, and then return response in that index."""
self.i += 1
print(f"=== Mock Response #{self.i} ===")
print(self.responses[self.i])
return self.responses[self.i]
@property
def _identifying_params(self) -> Mapping[str, Any]:
return {}
@property
def _llm_type(self) -> str:
"""Return type of llm."""
return "fake_list"
def _get_agent(**kwargs: Any) -> AgentExecutor:
"""Get agent for testing."""
bad_action_name = "BadAction"
responses = [
f"I'm turning evil\nAction: {bad_action_name}\nAction Input: misalignment",
"Oh well\nAction: Final Answer\nAction Input: curses foiled again",
]
fake_llm = FakeListLLM(responses=responses)
tools = [
Tool(
name="Search",
func=lambda x: x,
description="Useful for searching",
),
Tool(
name="Lookup",
func=lambda x: x,
description="Useful for looking up things in a table",
),
]
agent = initialize_agent(
tools, fake_llm, agent="zero-shot-react-description", verbose=True, **kwargs
)
return agent
def test_agent_bad_action() -> None:
"""Test react chain when bad action given."""
agent = _get_agent()
output = agent.run("when was langchain made")
assert output == "curses foiled again"
2022-12-29 13:21:11 +00:00
def test_agent_stopped_early() -> None:
"""Test react chain when bad action given."""
agent = _get_agent(max_iterations=0)
output = agent.run("when was langchain made")
assert output == "Agent stopped due to max iterations."
def test_agent_with_callbacks_global() -> None:
"""Test react chain with callbacks by setting verbose globally."""
import langchain
langchain.verbose = True
handler = FakeCallbackHandler()
manager = CallbackManager(handlers=[handler])
tool = "Search"
2022-12-29 13:21:11 +00:00
responses = [
f"FooBarBaz\nAction: {tool}\nAction Input: misalignment",
2022-12-29 13:21:11 +00:00
"Oh well\nAction: Final Answer\nAction Input: curses foiled again",
]
fake_llm = FakeListLLM(responses=responses, callback_manager=manager, verbose=True)
2022-12-29 13:21:11 +00:00
tools = [
Tool(
name="Search",
func=lambda x: x,
description="Useful for searching",
callback_manager=manager,
),
2022-12-29 13:21:11 +00:00
]
agent = initialize_agent(
tools,
fake_llm,
agent="zero-shot-react-description",
verbose=True,
callback_manager=manager,
2022-12-29 13:21:11 +00:00
)
2022-12-29 13:21:11 +00:00
output = agent.run("when was langchain made")
assert output == "curses foiled again"
# 1 top level chain run runs, 2 LLMChain runs, 2 LLM runs, 1 tool run
assert handler.chain_starts == handler.chain_ends == 3
assert handler.llm_starts == handler.llm_ends == 2
assert handler.tool_starts == 2
assert handler.tool_ends == 1
# 1 extra agent action
assert handler.starts == 7
# 1 extra agent end
assert handler.ends == 7
assert handler.errors == 0
# during LLMChain
assert handler.text == 2
def test_agent_with_callbacks_local() -> None:
"""Test react chain with callbacks by setting verbose locally."""
import langchain
langchain.verbose = False
handler = FakeCallbackHandler()
manager = CallbackManager(handlers=[handler])
tool = "Search"
responses = [
f"FooBarBaz\nAction: {tool}\nAction Input: misalignment",
"Oh well\nAction: Final Answer\nAction Input: curses foiled again",
]
fake_llm = FakeListLLM(responses=responses, callback_manager=manager, verbose=True)
tools = [
Tool(
name="Search",
func=lambda x: x,
description="Useful for searching",
callback_manager=manager,
),
]
agent = initialize_agent(
tools,
fake_llm,
agent="zero-shot-react-description",
verbose=True,
callback_manager=manager,
)
agent.agent.llm_chain.verbose = True
output = agent.run("when was langchain made")
assert output == "curses foiled again"
# 1 top level chain run, 2 LLMChain starts, 2 LLM runs, 1 tool run
assert handler.chain_starts == handler.chain_ends == 3
assert handler.llm_starts == handler.llm_ends == 2
assert handler.tool_starts == 2
assert handler.tool_ends == 1
# 1 extra agent action
assert handler.starts == 7
# 1 extra agent end
assert handler.ends == 7
assert handler.errors == 0
# during LLMChain
assert handler.text == 2
def test_agent_with_callbacks_not_verbose() -> None:
"""Test react chain with callbacks but not verbose."""
import langchain
langchain.verbose = False
handler = FakeCallbackHandler()
manager = CallbackManager(handlers=[handler])
tool = "Search"
responses = [
f"FooBarBaz\nAction: {tool}\nAction Input: misalignment",
"Oh well\nAction: Final Answer\nAction Input: curses foiled again",
]
fake_llm = FakeListLLM(responses=responses, callback_manager=manager)
tools = [
Tool(
name="Search",
func=lambda x: x,
description="Useful for searching",
),
]
agent = initialize_agent(
tools,
fake_llm,
agent="zero-shot-react-description",
callback_manager=manager,
)
output = agent.run("when was langchain made")
assert output == "curses foiled again"
# 1 top level chain run, 2 LLMChain runs, 2 LLM runs, 1 tool run
assert handler.starts == 0
assert handler.ends == 0
assert handler.errors == 0
def test_agent_tool_return_direct() -> None:
"""Test agent using tools that return directly."""
tool = "Search"
responses = [
f"FooBarBaz\nAction: {tool}\nAction Input: misalignment",
"Oh well\nAction: Final Answer\nAction Input: curses foiled again",
]
fake_llm = FakeListLLM(responses=responses)
tools = [
Tool(
name="Search",
func=lambda x: x,
description="Useful for searching",
return_direct=True,
),
]
agent = initialize_agent(
tools,
fake_llm,
agent="zero-shot-react-description",
)
output = agent.run("when was langchain made")
assert output == "misalignment"
Pass kwargs from initialize_agent into agent classmethod (#799) # Problem I noticed that in order to change the prefix of the prompt in the `zero-shot-react-description` agent we had to dig around to subset strings deep into the agent's attributes. It requires the user to inspect a long chain of attributes and classes. `initialize_agent -> AgentExecutor -> Agent -> LLMChain -> Prompt from Agent.create_prompt` ``` python agent = initialize_agent( tools=tools, llm=fake_llm, agent="zero-shot-react-description" ) prompt_str = agent.agent.llm_chain.prompt.template new_prompt_str = change_prefix(prompt_str) agent.agent.llm_chain.prompt.template = new_prompt_str ``` # Implemented Solution `initialize_agent` accepts `**kwargs` but passes it to `AgentExecutor` but not `ZeroShotAgent`, by simply giving the kwargs to the agent class methods we can support changing the prefix and suffix for one agent while allowing future agents to take advantage of `initialize_agent`. ``` agent = initialize_agent( tools=tools, llm=fake_llm, agent="zero-shot-react-description", agent_kwargs={"prefix": prefix, "suffix": suffix} ) ``` To be fair, this was before finding docs around custom agents here: https://langchain.readthedocs.io/en/latest/modules/agents/examples/custom_agent.html?highlight=custom%20#custom-llmchain but i find that my use case just needed to change the prefix a little. # Changes * Pass kwargs to Agent class method * Added a test to check suffix and prefix --------- Co-authored-by: Jason Liu <jason@jxnl.coA>
2023-01-30 22:54:09 +00:00
def test_agent_with_new_prefix_suffix() -> None:
"""Test agent initilization kwargs with new prefix and suffix."""
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,
),
Pass kwargs from initialize_agent into agent classmethod (#799) # Problem I noticed that in order to change the prefix of the prompt in the `zero-shot-react-description` agent we had to dig around to subset strings deep into the agent's attributes. It requires the user to inspect a long chain of attributes and classes. `initialize_agent -> AgentExecutor -> Agent -> LLMChain -> Prompt from Agent.create_prompt` ``` python agent = initialize_agent( tools=tools, llm=fake_llm, agent="zero-shot-react-description" ) prompt_str = agent.agent.llm_chain.prompt.template new_prompt_str = change_prefix(prompt_str) agent.agent.llm_chain.prompt.template = new_prompt_str ``` # Implemented Solution `initialize_agent` accepts `**kwargs` but passes it to `AgentExecutor` but not `ZeroShotAgent`, by simply giving the kwargs to the agent class methods we can support changing the prefix and suffix for one agent while allowing future agents to take advantage of `initialize_agent`. ``` agent = initialize_agent( tools=tools, llm=fake_llm, agent="zero-shot-react-description", agent_kwargs={"prefix": prefix, "suffix": suffix} ) ``` To be fair, this was before finding docs around custom agents here: https://langchain.readthedocs.io/en/latest/modules/agents/examples/custom_agent.html?highlight=custom%20#custom-llmchain but i find that my use case just needed to change the prefix a little. # Changes * Pass kwargs to Agent class method * Added a test to check suffix and prefix --------- Co-authored-by: Jason Liu <jason@jxnl.coA>
2023-01-30 22:54:09 +00:00
]
prefix = "FooBarBaz"
suffix = "Begin now!\nInput: {input}\nThought: {agent_scratchpad}"
agent = initialize_agent(
tools=tools,
llm=fake_llm,
agent="zero-shot-react-description",
agent_kwargs={"prefix": prefix, "suffix": suffix},
)
# avoids "BasePromptTemplate" has no attribute "template" error
assert hasattr(agent.agent.llm_chain.prompt, "template")
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"