add memory to create_pandas_dataframe_agent in Langchain

Question

I am trying to add memory to create_pandas_dataframe_agent to perform post processing on a model that I trained using Langchain. I am using the following code at the moment.

from langchain.llms import OpenAI
import pandas as pd

df = pd.read_csv('titanic.csv')
agent = create_pandas_dataframe_agent(OpenAI(temperature=0), [df], verbose=True)

I tried adding memory = ConversationBufferMemory(memory_key="chat_history") but that didnt help

score 1 · Answer 1 · answered Jun 27 '23 at 13:41

Answer from github issue https://github.com/hwchase17/langchain/issues/3106

"""Agent for working with pandas objects."""
from typing import Any, List, Optional

from langchain.agents.agent import AgentExecutor
# from langchain.agents.agent_toolkits.pandas.prompt import PREFIX, SUFFIX
from langchain.agents import ZeroShotAgent
from langchain.callbacks.base import BaseCallbackManager
from langchain.chains.llm import LLMChain
from langchain.llms.base import BaseLLM
from langchain.tools.python.tool import PythonAstREPLTool
from langchain.memory import ConversationBufferMemory

memory = ConversationBufferMemory(memory_key="chat_history")
def create_pandas_dataframe_agent(
    llm: BaseLLM,
    df: Any,
    callback_manager: Optional[BaseCallbackManager] = None,
    prefix: str = PREFIX,
    suffix: str = SUFFIX,
    input_variables: Optional[List[str]] = None,
    verbose: bool = False,
    return_intermediate_steps: bool = False,
    max_iterations: Optional[int] = 15,
    max_execution_time: Optional[float] = None,
    early_stopping_method: str = "force",
    **kwargs: Any,
) -> AgentExecutor:
    """Construct a pandas agent from an LLM and dataframe."""
    import pandas as pd

    if not isinstance(df, pd.DataFrame):
        raise ValueError(f"Expected pandas object, got {type(df)}")
    if input_variables is None:
        input_variables = ["df", "input", "agent_scratchpad"]
    tools = [PythonAstREPLTool(locals={"df": df})]
    
    PREFIX = """
            You are working with a pandas dataframe in Python. The name of the dataframe is `df`.
            You should use the tools below to answer the question posed of you:"""
    
    SUFFIX = """
            This is the result of `print(df.head())`:
            {df}
            Begin!
            {chat_history}
            Question: {input}
            {agent_scratchpad}"""

    prompt = ZeroShotAgent.create_prompt(
        tools, 
        prefix=PREFIX,
        suffix=SUFFIX, 
        input_variables=["df", "input", "chat_history", "agent_scratchpad"]
    )
    
    print(prompt)
    
    partial_prompt = prompt.partial(df=str(df.head()))
    
    llm_chain = LLMChain(
        llm=llm,
        prompt=partial_prompt,
        callback_manager=callback_manager,
    )
    
    tool_names = [tool.name for tool in tools]
    
    agent = ZeroShotAgent(
        llm_chain=llm_chain,
        allowed_tools=tool_names,
        callback_manager=callback_manager,
        **kwargs,
    )
    return AgentExecutor.from_agent_and_tools(
        agent=agent,
        tools=tools,
        verbose=verbose,
        return_intermediate_steps=return_intermediate_steps,
        max_iterations=max_iterations,
        max_execution_time=max_execution_time,
        early_stopping_method=early_stopping_method,
        callback_manager=callback_manager,
        memory = memory
    )
```*emphasized text*

score 0 · Answer 2 · answered Aug 31 '23 at 05:09

You can fix it by modifying the create_pandas_dataframe_agent to accept the memory object and then passing that along to the AgentCreator like below:

return AgentExecutor.from_agent_and_tools(
    agent=agent,
    tools=tools,
    verbose=verbose,
    return_intermediate_steps=return_intermediate_steps,
    max_iterations=max_iterations,
    max_execution_time=max_execution_time,
    early_stopping_method=early_stopping_method,
    memory=memory,
)

add memory to create_pandas_dataframe_agent in Langchain

2 Answers2