2022-12-07 05:52:48 +00:00
|
|
|
"""Unit tests for agents."""
|
|
|
|
|
|
|
|
from typing import Any, List, Mapping, Optional
|
|
|
|
|
2022-12-13 14:46:01 +00:00
|
|
|
from pydantic import BaseModel
|
|
|
|
|
2023-02-18 21:40:43 +00:00
|
|
|
from langchain.agents import AgentExecutor, initialize_agent
|
|
|
|
from langchain.agents.tools import Tool
|
2023-01-04 15:54:25 +00:00
|
|
|
from langchain.callbacks.base import CallbackManager
|
2022-12-07 05:52:48 +00:00
|
|
|
from langchain.llms.base import LLM
|
2023-01-04 15:54:25 +00:00
|
|
|
from tests.unit_tests.callbacks.fake_callback_handler import FakeCallbackHandler
|
2022-12-07 05:52:48 +00:00
|
|
|
|
|
|
|
|
2022-12-13 14:46:01 +00:00
|
|
|
class FakeListLLM(LLM, BaseModel):
|
2022-12-07 05:52:48 +00:00
|
|
|
"""Fake LLM for testing that outputs elements of a list."""
|
|
|
|
|
2022-12-13 14:46:01 +00:00
|
|
|
responses: List[str]
|
|
|
|
i: int = -1
|
2022-12-07 05:52:48 +00:00
|
|
|
|
2022-12-15 15:53:32 +00:00
|
|
|
def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
|
2022-12-07 05:52:48 +00:00
|
|
|
"""Increment counter, and then return response in that index."""
|
|
|
|
self.i += 1
|
2023-01-28 16:05:20 +00:00
|
|
|
print(f"=== Mock Response #{self.i} ===")
|
|
|
|
print(self.responses[self.i])
|
2022-12-07 05:52:48 +00:00
|
|
|
return self.responses[self.i]
|
|
|
|
|
|
|
|
@property
|
|
|
|
def _identifying_params(self) -> Mapping[str, Any]:
|
|
|
|
return {}
|
|
|
|
|
2022-12-13 14:46:01 +00:00
|
|
|
@property
|
|
|
|
def _llm_type(self) -> str:
|
|
|
|
"""Return type of llm."""
|
|
|
|
return "fake_list"
|
|
|
|
|
2022-12-07 05:52:48 +00:00
|
|
|
|
2023-01-04 15:54:25 +00:00
|
|
|
def _get_agent(**kwargs: Any) -> AgentExecutor:
|
|
|
|
"""Get agent for testing."""
|
2022-12-07 05:52:48 +00:00
|
|
|
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",
|
|
|
|
]
|
2022-12-13 14:46:01 +00:00
|
|
|
fake_llm = FakeListLLM(responses=responses)
|
2022-12-07 05:52:48 +00:00
|
|
|
tools = [
|
2023-02-18 21:40:43 +00:00
|
|
|
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",
|
|
|
|
),
|
2022-12-07 05:52:48 +00:00
|
|
|
]
|
|
|
|
agent = initialize_agent(
|
2023-01-04 15:54:25 +00:00
|
|
|
tools, fake_llm, agent="zero-shot-react-description", verbose=True, **kwargs
|
2022-12-07 05:52:48 +00:00
|
|
|
)
|
2023-01-04 15:54:25 +00:00
|
|
|
return agent
|
|
|
|
|
|
|
|
|
|
|
|
def test_agent_bad_action() -> None:
|
|
|
|
"""Test react chain when bad action given."""
|
|
|
|
agent = _get_agent()
|
2022-12-07 05:52:48 +00:00
|
|
|
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."""
|
2023-01-04 15:54:25 +00:00
|
|
|
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 = [
|
2023-01-04 15:54:25 +00:00
|
|
|
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",
|
|
|
|
]
|
2023-01-04 15:54:25 +00:00
|
|
|
fake_llm = FakeListLLM(responses=responses, callback_manager=manager, verbose=True)
|
2022-12-29 13:21:11 +00:00
|
|
|
tools = [
|
2023-02-18 21:40:43 +00:00
|
|
|
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,
|
2023-01-04 15:54:25 +00:00
|
|
|
callback_manager=manager,
|
2022-12-29 13:21:11 +00:00
|
|
|
)
|
2023-01-04 15:54:25 +00:00
|
|
|
|
2022-12-29 13:21:11 +00:00
|
|
|
output = agent.run("when was langchain made")
|
2023-01-04 15:54:25 +00:00
|
|
|
assert output == "curses foiled again"
|
|
|
|
|
2023-01-28 16:05:20 +00:00
|
|
|
# 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
|
2023-02-21 06:54:15 +00:00
|
|
|
assert handler.tool_starts == 2
|
|
|
|
assert handler.tool_ends == 1
|
|
|
|
# 1 extra agent action
|
|
|
|
assert handler.starts == 7
|
2023-01-04 15:54:25 +00:00
|
|
|
# 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 = [
|
2023-02-18 21:40:43 +00:00
|
|
|
Tool(
|
|
|
|
name="Search",
|
|
|
|
func=lambda x: x,
|
|
|
|
description="Useful for searching",
|
|
|
|
callback_manager=manager,
|
|
|
|
),
|
2023-01-04 15:54:25 +00:00
|
|
|
]
|
|
|
|
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"
|
|
|
|
|
2023-01-28 16:05:20 +00:00
|
|
|
# 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
|
2023-02-21 06:54:15 +00:00
|
|
|
assert handler.tool_starts == 2
|
|
|
|
assert handler.tool_ends == 1
|
|
|
|
# 1 extra agent action
|
|
|
|
assert handler.starts == 7
|
2023-01-04 15:54:25 +00:00
|
|
|
# 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 = [
|
2023-02-18 21:40:43 +00:00
|
|
|
Tool(
|
|
|
|
name="Search",
|
|
|
|
func=lambda x: x,
|
|
|
|
description="Useful for searching",
|
|
|
|
),
|
2023-01-04 15:54:25 +00:00
|
|
|
]
|
|
|
|
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
|
2023-01-06 14:40:32 +00:00
|
|
|
|
|
|
|
|
|
|
|
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 = [
|
2023-02-18 21:40:43 +00:00
|
|
|
Tool(
|
|
|
|
name="Search",
|
|
|
|
func=lambda x: x,
|
|
|
|
description="Useful for searching",
|
|
|
|
return_direct=True,
|
|
|
|
),
|
2023-01-06 14:40:32 +00:00
|
|
|
]
|
|
|
|
agent = initialize_agent(
|
|
|
|
tools,
|
|
|
|
fake_llm,
|
|
|
|
agent="zero-shot-react-description",
|
|
|
|
)
|
|
|
|
|
|
|
|
output = agent.run("when was langchain made")
|
|
|
|
assert output == "misalignment"
|
2023-01-30 22:54:09 +00:00
|
|
|
|
|
|
|
|
2023-03-13 14:54:29 +00:00
|
|
|
def test_agent_tool_return_direct_in_intermediate_steps() -> 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",
|
|
|
|
return_intermediate_steps=True,
|
|
|
|
)
|
|
|
|
|
|
|
|
resp = agent("when was langchain made")
|
|
|
|
assert resp["output"] == "misalignment"
|
|
|
|
assert len(resp["intermediate_steps"]) == 1
|
|
|
|
action, _action_intput = resp["intermediate_steps"][0]
|
|
|
|
assert action.tool == "Search"
|
|
|
|
|
|
|
|
|
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 = [
|
2023-02-18 21:40:43 +00:00
|
|
|
Tool(
|
|
|
|
name="Search",
|
|
|
|
func=lambda x: x,
|
|
|
|
description="Useful for searching",
|
|
|
|
return_direct=True,
|
|
|
|
),
|
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"
|
2023-03-26 03:03:18 +00:00
|
|
|
|
|
|
|
|
|
|
|
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]
|