test_mrkl.py — langchain Source File

Source Code

"""Test MRKL functionality."""

import pytest
from langchain_core.agents import AgentAction
from langchain_core.exceptions import OutputParserException
from langchain_core.prompts import PromptTemplate
from langchain_core.tools import Tool

from langchain_classic.agents.mrkl.base import ZeroShotAgent
from langchain_classic.agents.mrkl.output_parser import MRKLOutputParser
from langchain_classic.agents.mrkl.prompt import FORMAT_INSTRUCTIONS, PREFIX, SUFFIX
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)
    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\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_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 when there is a new line before the action.

    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:
// ... (105 more lines)