Files
langchain/docs/examples/memory/agent_with_memory.ipynb
T

10 KiB

Adding Memory to an Agent

This notebook goes over adding memory to an Agent. Before going through this notebook, please walkthrough the following notebooks, as this will build on top of both of them:

In order to add a memory to an agent we are going to the the following steps:

  1. We are going to create an LLMChain with memory.
  2. We are going to use that LLMChain to create a custom Agent.

For the purposes of this exercise, we are going to create a simple custom Agent that has access to a search tool and utilizes the ConversationBufferMemory class.

In [1]:
from langchain.agents import ZeroShotAgent, Tool
from langchain.chains.conversation.memory import ConversationBufferMemory
from langchain import OpenAI, SerpAPIWrapper, LLMChain
In [2]:
search = SerpAPIWrapper()
tools = [
    Tool(
        name = "Search",
        func=search.run,
        description="useful for when you need to answer questions about current events"
    )
]

Notice the usage of the chat_history variable in the PromptTemplate, which matches up with the dynamic key name in the ConversationBufferMemory.

In [3]:
prefix = """Have a conversation with a human, answering the following questions as best you can. You have access to the following tools:"""
suffix = """Begin!"

{chat_history}
Question: {input}"""

prompt = ZeroShotAgent.create_prompt(
    tools, 
    prefix=prefix, 
    suffix=suffix, 
    input_variables=["input", "chat_history"]
)
memory = ConversationBufferMemory(memory_key="chat_history")

We can now construct the LLMChain, with the Memory object, and then create the agent.

In [4]:
llm_chain = LLMChain(llm=OpenAI(temperature=0), prompt=prompt, memory=memory)
agent = ZeroShotAgent(llm_chain=llm_chain, tools=tools, verbose=True)
In [5]:
agent.run("How many people live in canada?")
Out [5]:

> Entering new chain...
How many people live in canada?
Thought: I should look up how many people live in canada
Action: Search
Action Input: "How many people live in canada?"
Observation: The current population of Canada is 38,533,678 as of Friday, November 25, 2022, based on Worldometer elaboration of the latest United Nations data. · Canada 2020 ...
Thought: I now know the final answer
Final Answer: The current population of Canada is 38,533,678 as of Friday, November 25, 2022, based on Worldometer elaboration of the latest United Nations data.
> Finished chain.
'The current population of Canada is 38,533,678 as of Friday, November 25, 2022, based on Worldometer elaboration of the latest United Nations data.'

To test the memory of this agent, we can ask a followup question that relies on information in the previous exchange to be answered correctly.

In [6]:
agent.run("what is their national anthem called?")
Out [6]:

> Entering new chain...
what is their national anthem called?
Thought:
AI:  I should look up the name of Canada's national anthem
Action: Search
Action Input: "What is the name of Canada's national anthem?"
Observation: After 100 years of tradition, O Canada was proclaimed Canada's national anthem in 1980. The music for O Canada was composed in 1880 by Calixa ...
Thought:
AI:  I now know the final answer
Final Answer: After 100 years of tradition, O Canada was proclaimed Canada's national anthem in 1980. The music for O Canada was composed in 1880 by Calixa Lavallée.
> Finished chain.
"After 100 years of tradition, O Canada was proclaimed Canada's national anthem in 1980. The music for O Canada was composed in 1880 by Calixa Lavallée."

We can see that the agent remembered that the previous question was about Canada, and properly asked Google Search what the name of Canada's national anthem was.

For fun, let's compare this to an agent that does NOT have memory.

In [7]:
prefix = """Have a conversation with a human, answering the following questions as best you can. You have access to the following tools:"""
suffix = """Begin!"

Question: {input}"""

prompt = ZeroShotAgent.create_prompt(
    tools, 
    prefix=prefix, 
    suffix=suffix, 
    input_variables=["input"]
)
llm_chain = LLMChain(llm=OpenAI(temperature=0), prompt=prompt)
agent_without_memory = ZeroShotAgent(llm_chain=llm_chain, tools=tools, verbose=True)
In [10]:
agent_without_memory.run("How many people live in canada?")
Out [10]:

> Entering new chain...
How many people live in canada?
Thought: I should look up how many people live in canada
Action: Search
Action Input: "How many people live in canada?"
Observation: The current population of Canada is 38,533,678 as of Friday, November 25, 2022, based on Worldometer elaboration of the latest United Nations data. · Canada 2020 ...
Thought: I now know the final answer
Final Answer: The current population of Canada is 38,533,678
> Finished chain.
'The current population of Canada is 38,533,678'
In [11]:
agent_without_memory.run("what is their national anthem called?")
Out [11]:

> Entering new chain...
what is their national anthem called?
Thought: I should probably look this up
Action: Search
Action Input: "What is the national anthem of [country]"
Observation: Most nation states have an anthem, defined as "a song, as of praise, devotion, or patriotism"; most anthems are either marches or hymns in style.
Thought: I now know the final answer
Final Answer: The national anthem is called "the national anthem."
> Finished chain.
'The national anthem is called "the national anthem."'
In [ ]: