⚡️ Speed up create_prompt() by 13% in libs/langchain/langchain/agents/mrkl/base.py

[codeflash-ai]

📄 create_prompt() in libs/langchain/langchain/agents/mrkl/base.py

📈 Performance went up by 13% (0.13x faster)

⏱️ Runtime went down from 240.60μs to 212.71μs

Explanation and details

(click to show)

In the original function, there are two list comprehensions where you're iterating over the same 'tools' object twice which can be costly if the 'tools' List is large.

Instead, you can combine these two list comprehensions into one loop on 'tools'. Also, we don't need to use join right away. We append the strings to lists and join only at the end, which should give performance improvement as string concatenation in Python creates a new string and is thus more expensive than list append.

Here is the optimized version.

Correctness verification

The new optimized code was tested for correctness. The results are listed below.

✅ 0 Passed − ⚙️ Existing Unit Tests

✅ 0 Passed − 🎨 Inspired Regression Tests

✅ 15 Passed − 🌀 Generated Regression Tests

(click to show generated tests)

# imports
import pytest
from typing import Sequence, Optional, List
from collections import namedtuple

# Assuming the existence of the following classes and constants
# as they are not defined within the provided code snippet.
Agent = object  # Placeholder for the actual Agent class
PromptTemplate = namedtuple('PromptTemplate', ['template', 'input_variables'])
BaseTool = namedtuple('BaseTool', ['name', 'description'])
PREFIX = "Tools available:"
SUFFIX = "End of tools."
FORMAT_INSTRUCTIONS = "Please use the following tools: {tool_names}."

# function to test
# ... (The ZeroShotAgent class definition remains the same as provided) ...

# Helper function to create a BaseTool object
def create_tool(name, description):
    return BaseTool(name=name, description=description)

# Test with different numbers of tools
def test_no_tools():
    prompt = ZeroShotAgent.create_prompt(tools=[])
    assert prompt.template == f"{PREFIX}\n\n{FORMAT_INSTRUCTIONS.format(tool_names='')}\n\n{SUFFIX}"
    assert prompt.input_variables == ["input", "agent_scratchpad"]

def test_single_tool():
    tool = create_tool("Tool1", "Description1")
    prompt = ZeroShotAgent.create_prompt(tools=[tool])
    assert tool.name in prompt.template
    assert tool.description in prompt.template
    assert prompt.input_variables == ["input", "agent_scratchpad"]

def test_multiple_tools():
    tools = [create_tool(f"Tool{i}", f"Description{i}") for i in range(3)]
    prompt = ZeroShotAgent.create_prompt(tools=tools)
    for tool in tools:
        assert tool.name in prompt.template
        assert tool.description in prompt.template
    assert prompt.input_variables == ["input", "agent_scratchpad"]

# Test with various tool attributes
@pytest.mark.parametrize("name, description", [
    ("", ""),
    ("NameWith\nNewline", "DescriptionWith\nNewline"),
    ("NameWith{Bracket", "DescriptionWith}Bracket"),
])
def test_tool_attributes(name, description):
    tool = create_tool(name, description)
    prompt = ZeroShotAgent.create_prompt(tools=[tool])
    assert name in prompt.template
    assert description in prompt.template

# Test with different prefix and suffix values
@pytest.mark.parametrize("prefix, suffix", [
    ("", ""),
    ("Custom Prefix", "Custom Suffix"),
    ("Prefix{With}Braces", "Suffix{With}Braces"),
])
def test_prefix_suffix(prefix, suffix):
    tool = create_tool("Tool1", "Description1")
    prompt = ZeroShotAgent.create_prompt(tools=[tool], prefix=prefix, suffix=suffix)
    assert prefix in prompt.template
    assert suffix in prompt.template

# Test with different input_variables
@pytest.mark.parametrize("input_variables", [
    None,
    [],
    ["custom_input"],
    ["input1", "input2"],
])
def test_input_variables(input_variables):
    tool = create_tool("Tool1", "Description1")
    prompt = ZeroShotAgent.create_prompt(tools=[tool], input_variables=input_variables)
    expected_variables = input_variables if input_variables is not None else ["input", "agent_scratchpad"]
    assert prompt.input_variables == expected_variables

# Test edge cases
def test_tools_with_non_string_attributes():
    tool = create_tool(None, 123)
    with pytest.raises(TypeError):
        ZeroShotAgent.create_prompt(tools=[tool])

def test_prefix_suffix_with_placeholders():
    tool = create_tool("Tool1", "Description1")
    prefix = "Prefix {unexpected}"
    suffix = "Suffix {unexpected}"
    with pytest.raises(KeyError):
        ZeroShotAgent.create_prompt(tools=[tool], prefix=prefix, suffix=suffix)

def test_format_instructions_without_placeholder():
    tool = create_tool("Tool1", "Description1")
    instructions = "No placeholder for tool names."
    prompt = ZeroShotAgent.create_prompt(tools=[tool], format_instructions=instructions)
    assert instructions in prompt.template

# Test error handling
def test_incorrect_types_for_tools():
    with pytest.raises(TypeError):
        ZeroShotAgent.create_prompt(tools="not a sequence")

def test_incorrect_types_for_prefix():
    with pytest.raises(TypeError):
        ZeroShotAgent.create_prompt(tools=[], prefix=123)

def test_incorrect_types_for_suffix():
    with pytest.raises(TypeError):
        ZeroShotAgent.create_prompt(tools=[], suffix=123)

def test_incorrect_types_for_format_instructions():
    with pytest.raises(TypeError):
        ZeroShotAgent.create_prompt(tools=[], format_instructions=123)