Spaces:
Sleeping
Sleeping
File size: 5,164 Bytes
cfd3735 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 |
"""Test MRKL functionality."""
from typing import Tuple
import pytest
from langchain.agents.mrkl.base import ZeroShotAgent
from langchain.agents.mrkl.output_parser import MRKLOutputParser
from langchain.agents.mrkl.prompt import FORMAT_INSTRUCTIONS, PREFIX, SUFFIX
from langchain.agents.tools import Tool
from langchain.prompts import PromptTemplate
from langchain.schema import AgentAction, OutputParserException
from tests.unit_tests.llms.fake_llm import FakeLLM
def get_action_and_input(text: str) -> Tuple[str, str]:
output = MRKLOutputParser().parse(text)
if isinstance(output, AgentAction):
return output.tool, str(output.tool_input)
else:
return "Final Answer", output.return_values["output"]
def test_get_action_and_input() -> None:
"""Test getting an action from text."""
llm_output = (
"Thought: I need to search for NBA\n" "Action: Search\n" "Action Input: NBA"
)
action, action_input = get_action_and_input(llm_output)
assert action == "Search"
assert action_input == "NBA"
def test_get_action_and_input_whitespace() -> None:
"""Test getting an action from text."""
llm_output = "Thought: I need to search for NBA\nAction: Search \nAction Input: NBA"
action, action_input = get_action_and_input(llm_output)
assert action == "Search"
assert action_input == "NBA"
def test_get_action_and_input_newline() -> None:
"""Test getting an action from text where Action Input is a code snippet."""
llm_output = (
"Now I need to write a unittest for the function.\n\n"
"Action: Python\nAction Input:\n```\nimport unittest\n\nunittest.main()\n```"
)
action, action_input = get_action_and_input(llm_output)
assert action == "Python"
assert action_input == "```\nimport unittest\n\nunittest.main()\n```"
def test_get_action_and_input_newline_after_keyword() -> None:
"""Test getting an action and action input from the text
when there is a new line before the action
(after the keywords "Action:" and "Action Input:")
"""
llm_output = """
I can use the `ls` command to list the contents of the directory \
and `grep` to search for the specific file.
Action:
Terminal
Action Input:
ls -l ~/.bashrc.d/
"""
action, action_input = get_action_and_input(llm_output)
assert action == "Terminal"
assert action_input == "ls -l ~/.bashrc.d/\n"
def test_get_final_answer() -> None:
"""Test getting final answer."""
llm_output = (
"Thought: I need to search for NBA\n"
"Action: Search\n"
"Action Input: NBA\n"
"Observation: founded in 1994\n"
"Thought: I can now answer the question\n"
"Final Answer: 1994"
)
action, action_input = get_action_and_input(llm_output)
assert action == "Final Answer"
assert action_input == "1994"
def test_get_final_answer_new_line() -> None:
"""Test getting final answer."""
llm_output = (
"Thought: I need to search for NBA\n"
"Action: Search\n"
"Action Input: NBA\n"
"Observation: founded in 1994\n"
"Thought: I can now answer the question\n"
"Final Answer:\n1994"
)
action, action_input = get_action_and_input(llm_output)
assert action == "Final Answer"
assert action_input == "1994"
def test_get_final_answer_multiline() -> None:
"""Test getting final answer that is multiline."""
llm_output = (
"Thought: I need to search for NBA\n"
"Action: Search\n"
"Action Input: NBA\n"
"Observation: founded in 1994 and 1993\n"
"Thought: I can now answer the question\n"
"Final Answer: 1994\n1993"
)
action, action_input = get_action_and_input(llm_output)
assert action == "Final Answer"
assert action_input == "1994\n1993"
def test_bad_action_input_line() -> None:
"""Test handling when no action input found."""
llm_output = "Thought: I need to search for NBA\n" "Action: Search\n" "Thought: NBA"
with pytest.raises(OutputParserException):
get_action_and_input(llm_output)
def test_bad_action_line() -> None:
"""Test handling when no action input found."""
llm_output = (
"Thought: I need to search for NBA\n" "Thought: Search\n" "Action Input: NBA"
)
with pytest.raises(OutputParserException):
get_action_and_input(llm_output)
def test_from_chains() -> None:
"""Test initializing from chains."""
chain_configs = [
Tool(name="foo", func=lambda x: "foo", description="foobar1"),
Tool(name="bar", func=lambda x: "bar", description="foobar2"),
]
agent = ZeroShotAgent.from_llm_and_tools(FakeLLM(), chain_configs)
expected_tools_prompt = "foo: foobar1\nbar: foobar2"
expected_tool_names = "foo, bar"
expected_template = "\n\n".join(
[
PREFIX,
expected_tools_prompt,
FORMAT_INSTRUCTIONS.format(tool_names=expected_tool_names),
SUFFIX,
]
)
prompt = agent.llm_chain.prompt
assert isinstance(prompt, PromptTemplate)
assert prompt.template == expected_template
|