Agent not running tools when instantiating agent from class.

[JasonS09]

Hello. I'm developing a custom agent that can generate images connected to a local instance of Stable Diffusion. However I wanted to know why if I put the logic for building the agent inside a class, the agent works but it's silently unable to call the tools functions. No error or nothing, it just simply ignores them even though the output is correctly formatted. I'm not putting this as an issue because I think this is something more basic and I'm overlooking it.

Here's my code:

from langchain.chains import LLMChain
from langchain.prompts import StringPromptTemplate
from langchain.agents import Tool, AgentExecutor, LLMSingleActionAgent, AgentOutputParser
from langchain.memory import ConversationBufferWindowMemory
from langchain.schema import AgentAction, AgentFinish, OutputParserException
from langchain.tools import DuckDuckGoSearchResults
from langchain.utilities import DuckDuckGoSearchAPIWrapper

from schemas import ChatResponse

from typing import List, Union

import re

from comfy import StableDiffusionGenerator

# Set up a prompt template
class CustomPromptTemplate(StringPromptTemplate):
    # The template to use
    template: str
    # The list of tools available
    tools: List[Tool]

    def format(self, **kwargs) -> str:
        # Get the intermediate steps (AgentAction, Observation tuples)
        # Format them in a particular way
        intermediate_steps = kwargs.pop("intermediate_steps")
        thoughts = ""
        for action, observation in intermediate_steps:
            thoughts += action.log
            thoughts += f"\nObservation: {observation}\nThought: "
        # Set the agent_scratchpad variable to that value
        kwargs["agent_scratchpad"] = thoughts
        # Create a tools variable from the list of tools provided
        kwargs["tools"] = "\n".join([f"{tool.name}: {tool.description}" for tool in self.tools])
        # Create a list of tool names for the tools provided
        kwargs["tool_names"] = ", ".join([tool.name for tool in self.tools])
        return self.template.format(**kwargs)

class CustomOutputParser(AgentOutputParser):

    def parse(self, llm_output: str) -> Union[AgentAction, AgentFinish]:
        if "Final Answer:" in llm_output:
            return AgentFinish(
                # Return values is generally always a dictionary with a single `output` key
                # It is not recommended to try anything else at the moment :)
                return_values={"output": llm_output.split("Final Answer:")[-1].strip()},
                log=llm_output,
            )
        # Parse out the action and action input
        regex = r"Action\s*\d*\s*:(.*?)\n.*?Action Input\s*\d*\s*:[\s]*(.*)"
        match = re.search(regex, llm_output, re.DOTALL)
        if not match:
            raise OutputParserException(f"Could not parse LLM output: `{llm_output}`")
        action = match.group(1).strip()
        action_input = match.group(2)
        # Return the action and action input
        return AgentAction(tool=action, tool_input=action_input.strip(" ").strip('"'), log=llm_output)
    
class CustomAgent:
    def __init__(self):
        #Filled with images from Stable Diffusion tool output
        self.response = ChatResponse(sender="bot", message="", type="stream", images=[])
        llm = KoboldApiLLM(endpoint="http://localhost:5001", max_length=512)

        wrapper = DuckDuckGoSearchAPIWrapper()
        search = DuckDuckGoSearchResults(wrapper=wrapper)
        sd = StableDiffusionGenerator()

        def run_stable_diffusion(prompt):
            try:
                msg, images = sd.get_images(prompt)
                self.response.images = images
                return msg
            except:
                return "Error when generating images."   

        tools = [
            Tool(
                func=search.run,
                name="Search",
                description="Useful for when you need to answer questions about current events or concepts.",
                handle_tool_error=True
            ),
            Tool(
                func=run_stable_diffusion,
                name="Generate Images",
                description="Useful for when you need to generate images or visually illustrate a concept.",
                handle_tool_error=True
            )
        ]

        # Set up the base template
        template_with_history = """Complete the objective as best you can. You have access to the following tools:

        {tools}

        Use the following format:

        Question: the input question you must answer
        Thought: Do you need to take an action? If not, skip to Final Answer.
        Action: the name of the action to take, should be one of the following: {tool_names}
        Action Input: the input to the action
        Observation: the result of the action
        ... (this Thought/Action/Action Input/Observation can repeat N times)
        Thought: I now know the final answer
        Final Answer: the final answer to the original input question. Make sure to give long, detailed final answers if needed.
        Begin!

        Previous conversation history:
        {history}

        Question: {input}
        {agent_scratchpad}"""

        prompt_with_history = CustomPromptTemplate(
            template=template_with_history,
            tools=tools,
            # This omits the `agent_scratchpad`, `tools`, and `tool_names` variables because those are generated dynamically
            # This includes the `intermediate_steps` variable because that is needed
            input_variables=["input", "intermediate_steps", "history"]
        )

        output_parser = CustomOutputParser()
        llm_chain = LLMChain(llm=llm, prompt=prompt_with_history)

        tool_names = [tool.name for tool in tools]
        agent = LLMSingleActionAgent(
            llm_chain=llm_chain,
            output_parser=output_parser,
            stop=["Begin!", "Question:"],
            allowed_tools=tool_names
        )

        memory=ConversationBufferWindowMemory(k=2)

        self.agent_executor = AgentExecutor(
            agent=agent, tools=tools, verbose=True, 
            memory=memory, handle_parsing_errors=True)
    
    def chat(self, input):
        self.response = ChatResponse(sender="bot", message="", type="stream", images=[])
        self.response.message = self.agent_executor.run(input)
        return self.response
        

agent = CustomAgent()
agent.chat("Can you generate a pic of a cute plant?")

If I put all the logic outside the class, and run it as a script, it works as expected and the agent calls the function for image generation. What is wrong with my code then? It has to do with the way I designed the class.

[JasonS09]

I finally understood what is the issue. Seems I was missing one stop sequence for LLMSingleActionAgent (Observation).

[fabiodevn1]

Could you explain more?