unable to implement memory with sql agent toolkit
Labels: name: area: memory, name: area: agent, name: auto:improvement

### System Info

I am using langchain 0.0.238, I want to use ConversationBufferMemory with sql-agent-toolkit. Please support in this regard.
I changed prompt :
# flake8: noqa

SQL_PREFIX = """You are an agent designed to interact with a SQL database.
Given an input question, create a syntactically correct {dialect} query to run, then look at the results of the query and return the answer.
Unless the user specifies a specific number of examples they wish to obtain, else fetch all results then fetch all results donot use 'LIMIT'.
When you are using 'WHERE' clause to fetch specific data, you have to apply 'WHERE' clause in several ways like exact matchusing '=' or substring match using 'LIKE'. 
You can order the results by a relevant column to return the most interesting examples in the database.
Never query for all the columns from a specific table, only ask for the relevant columns given the question. If you did notfind the column asked by the user, then return "data for the specified column doesnot exists." as the answer.
You have access to tools for interacting with the database.
Only use the below tools. Only use the information returned by the below tools to construct your final answer.
You MUST double check your query before executing it. If you get an error while executing a query, rewrite the query and try again.

DO NOT make any DML statements (INSERT, UPDATE, DELETE, DROP etc.) to the database.

If the question does not seem related to the database, just return "I don't know" as the answer.

This reminds you of your past memory and events:{chat_history}

Know that only the "Question" ("HumanMessage") and "Final answer" ("AIMessage") to each previous question are available inthe past memory section above.
Hence, always include enough detail in the Final answer so that you can answer subsequent questions.
If you believe any of the data computed to answer a question is likely to be useful in answering a possible follow-on question.

"""

SQL_SUFFIX = """Begin!

Question: {input}
Thought: I should look at the tables in the database to see what I can query.  Then I should query the schema of the most relevant tables.
{agent_scratchpad}"""

SQL_FUNCTIONS_SUFFIX = """I should look at the tables in the database to see what I can query.  Then I should query the schema of the most relevant tables."""

And changed base.py:
"""SQL agent."""
from typing import Any, Dict, List, Optional

from langchain.agents.agent import AgentExecutor, BaseSingleActionAgent
#from sql_agent_toolkits.sql.prompt import ( SQL_FUNCTIONS_SUFFIX,SQL_PREFIX,SQL_SUFFIX,)
from langchain.agents.agent_toolkits.sql.prompt import (
    SQL_FUNCTIONS_SUFFIX,
    SQL_PREFIX,
    SQL_SUFFIX,
)
from langchain.agents.agent_toolkits.sql.toolkit import SQLDatabaseToolkit
from langchain.agents.agent_types import AgentType
from langchain.agents.mrkl.base import ZeroShotAgent
from langchain.agents.mrkl.prompt import FORMAT_INSTRUCTIONS
from langchain.agents.openai_functions_agent.base import OpenAIFunctionsAgent
from langchain.callbacks.base import BaseCallbackManager
from langchain.chains.llm import LLMChain
from langchain.prompts.chat import (
    ChatPromptTemplate,
    HumanMessagePromptTemplate,
    MessagesPlaceholder,
)
from langchain.schema.language_model import BaseLanguageModel
from langchain.schema.messages import AIMessage, SystemMessage
from langchain.memory import ConversationBufferMemory

global memory
memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)


def create_sql_agent(
    llm: BaseLanguageModel,
    toolkit: SQLDatabaseToolkit,
    agent_type: AgentType = AgentType.ZERO_SHOT_REACT_DESCRIPTION,
    callback_manager: Optional[BaseCallbackManager] = None,
    prefix: str = SQL_PREFIX,
    suffix: Optional[str] = SQL_SUFFIX,
    format_instructions: str = FORMAT_INSTRUCTIONS,
    input_variables: Optional[List[str]] = None,
    top_k: int = 10,
    max_iterations: Optional[int] = 15,
    max_execution_time: Optional[float] = None,
    early_stopping_method: str = "force",
    verbose: bool = False,
    agent_executor_kwargs: Optional[Dict[str, Any]] = None,
    **kwargs: Dict[str, Any],
) -> AgentExecutor:
    """Construct an SQL agent from an LLM and tools."""
    tools = toolkit.get_tools()
    #prefix = prefix.format(dialect=toolkit.dialect, top_k=top_k)
    #print("prefix", prefix)
    if input_variables is None:
        input_variables = ["dialect", "input", "agent_scratchpad"]
    agent: BaseSingleActionAgent
    print("input_variables", input_variables)
    if agent_type == AgentType.ZERO_SHOT_REACT_DESCRIPTION:
        prompt = ZeroShotAgent.create_prompt(
            tools,
            prefix=prefix,
            suffix=suffix or SQL_SUFFIX,
            format_instructions=format_instructions,
            input_variables=input_variables,
        )
        llm_chain = LLMChain(
            llm=llm,
            prompt=prompt,
            callback_manager=callback_manager,
        )
        tool_names = [tool.name for tool in tools]
        agent = ZeroShotAgent(llm_chain=llm_chain, allowed_tools=tool_names, **kwargs)
#         memory = ConversationBufferMemory(memory_key="chat_history")
    elif agent_type == AgentType.OPENAI_FUNCTIONS:
        messages = [
            SystemMessage(content=prefix),
            HumanMessagePromptTemplate.from_template("{input}"),
            AIMessage(content=suffix or SQL_FUNCTIONS_SUFFIX),
            MessagesPlaceholder(variable_name="agent_scratchpad"),
        ]
        input_variables = ["input", "agent_scratchpad"]
        _prompt = ChatPromptTemplate(input_variables=input_variables, messages=messages)

        agent = OpenAIFunctionsAgent(
            llm=llm,
            prompt=_prompt,
            tools=tools,
            callback_manager=callback_manager,
            **kwargs,
        )
    else:
        raise ValueError(f"Agent type {agent_type} not supported at the moment.")

    return AgentExecutor.from_agent_and_tools(
        agent=agent,
        tools=tools,
        callback_manager=callback_manager,
        verbose=verbose,
        max_iterations=max_iterations,
        max_execution_time=max_execution_time,
        early_stopping_method=early_stopping_method,
        memory=memory,
        **(agent_executor_kwargs or {}),
    )


### Who can help?

_No response_

### Information

- [ ] The official example notebooks/scripts
- [ ] My own modified scripts

### Related Components

- [ ] LLMs/Chat Models
- [ ] Embedding Models
- [ ] Prompts / Prompt Templates / Prompt Selectors
- [ ] Output Parsers
- [ ] Document Loaders
- [ ] Vector Stores / Retrievers
- [X] Memory
- [X] Agents / Agent Executors
- [X] Tools / Toolkits
- [ ] Chains
- [ ] Callbacks/Tracing
- [ ] Async

### Reproduction

I have changed Prompt.py and base.py to provide chat_history as input but I am running into errors from create_prompt method using ZeroShotAgent or from LLMChain(expecting inputs in dictionary)
Error:

input_variables ['dialect', 'input', 'agent_scratchpad']
input_variables from /opt/miniconda/lib/python3.10/site-packages/langchain/agents/mrkl/ ['dialect', 'input', 'agent_scratchpad']
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[1], line 2
      1 from tools.db_llm_with_memory import *
----> 2 db_llm_with_memory("which all enzymes have urate as substrate")

File ~/notebooks/atossa-usecase/atossa_usecase/tools/db_llm_with_memory.py:42, in db_llm_with_memory(query)
     35     agent = create_sql_agent(
     36         llm=ChatOpenAI(temperature=0),
     37         toolkit=toolkit,
     38         verbose=True,
     39         agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,memory=memory
     40     )
     41 #     agent_chain = AgentExecutor.from_agent_and_tools(agent=agent, tools=toolkit, verbose=True, memory=memory)
---> 42     agent.run(query)

File /opt/miniconda/lib/python3.10/site-packages/langchain/chains/base.py:440, in Chain.run(self, callbacks, tags, metadata, *args, **kwargs)
    438     if len(args) != 1:
    439         raise ValueError("`run` supports only one positional argument.")
--> 440     return self(args[0], callbacks=callbacks, tags=tags, metadata=metadata)[
    441         _output_key
    442     ]
    444 if kwargs and not args:
    445     return self(kwargs, callbacks=callbacks, tags=tags, metadata=metadata)[
    446         _output_key
    447     ]

File /opt/miniconda/lib/python3.10/site-packages/langchain/chains/base.py:220, in Chain.__call__(self, inputs, return_only_outputs, callbacks, tags, metadata, include_run_info)
    185 def __call__(
    186     self,
    187     inputs: Union[Dict[str, Any], Any],
   (...)
    193     include_run_info: bool = False,
    194 ) -> Dict[str, Any]:
    195     """Execute the chain.
    196 
    197     Args:
   (...)
    218             `Chain.output_keys`.
    219     """
--> 220     inputs = self.prep_inputs(inputs)
    221     callback_manager = CallbackManager.configure(
    222         callbacks,
    223         self.callbacks,
   (...)
    228         self.metadata,
    229     )
    230     new_arg_supported = inspect.signature(self._call).parameters.get("run_manager")

File /opt/miniconda/lib/python3.10/site-packages/langchain/chains/base.py:364, in Chain.prep_inputs(self, inputs)
    362         _input_keys = _input_keys.difference(self.memory.memory_variables)
    363     if len(_input_keys) != 1:
--> 364         raise ValueError(
    365             f"A single string input was passed in, but this chain expects "
    366             f"multiple inputs ({_input_keys}). When a chain expects "
    367             f"multiple inputs, please call it by passing in a dictionary, "
    368             "eg `chain({'foo': 1, 'bar': 2})`"
    369         )
    370     inputs = {list(_input_keys)[0]: inputs}
    371 if self.memory is not None:

ValueError: A single string input was passed in, but this chain expects multiple inputs ({'dialect', 'input'}). When a chain expects multiple inputs, please call it by passing in a dictionary, eg `chain({'foo': 1, 'bar': 2})`


### Expected behavior

Please provide some simple addon behaviour to use memory with sql agent