LlamaIndex Tutorial (Python): adding memory to agents for beginners

This tutorial shows you how to give a LlamaIndex agent short-term memory in Python so it can remember prior turns in a conversation. You need this when your agent must answer follow-up questions, keep context across multiple messages, or avoid asking the user to repeat themselves.

What You'll Need

• •Python 3.10+
• •A virtual environment
• •llama-index
• •An OpenAI API key
• •python-dotenv for local env vars
• •Basic familiarity with LlamaIndex agents and tools

Install the packages first:

pip install llama-index python-dotenv

Set your API key in a .env file:

OPENAI_API_KEY=your_api_key_here

Step-by-Step

1. •Start by loading your environment and creating a simple tool the agent can call. For beginners, a single deterministic tool is enough to prove memory is working.

import os
from dotenv import load_dotenv

from llama_index.core.tools import FunctionTool

load_dotenv()

def get_policy_status(policy_id: str) -> str:
    policy_map = {
        "P1001": "Active",
        "P1002": "Pending renewal",
        "P1003": "Cancelled",
    }
    return f"Policy {policy_id} is {policy_map.get(policy_id, 'Unknown')}."

policy_tool = FunctionTool.from_defaults(fn=get_policy_status)

2. •Next, create the LLM and the memory object. ChatMemoryBuffer stores recent chat history and is the simplest way to add conversational state to an agent.

from llama_index.core import Settings
from llama_index.core.memory import ChatMemoryBuffer
from llama_index.llms.openai import OpenAI

Settings.llm = OpenAI(model="gpt-4o-mini")
memory = ChatMemoryBuffer.from_defaults(token_limit=2000)

3. •Now build the agent with your tool and memory attached. The key part is passing the same memory instance into the agent so it can reuse prior turns.

from llama_index.core.agent.workflow import FunctionAgent

agent = FunctionAgent(
    tools=[policy_tool],
    llm=Settings.llm,
    memory=memory,
    system_prompt=(
        "You are a policy support assistant. "
        "Use the tool when asked about policy status. "
        "Remember the user's policy ID if they mention it."
    ),
)

4. •Send a first message that establishes context, then ask a follow-up without repeating everything. If memory is wired correctly, the agent should keep track of the policy ID from the earlier turn.

response1 = await agent.run("My policy ID is P1002.")
print(response1)

response2 = await agent.run("What is its status?")
print(response2)

5. •If you want to inspect what got stored, read from the memory object directly. This is useful when debugging whether your conversation state is actually being retained.

chat_history = memory.get_all()
for message in chat_history:
    print(f"{message.role}: {message.content}")

Testing It

Run the script in an environment where OPENAI_API_KEY is set. The first response should acknowledge or store the policy ID, and the second response should answer using that remembered value instead of asking you to repeat it.

If it forgets the policy ID, check three things first: you reused the same memory object, your agent was created with that memory, and your model call actually completed successfully. Also make sure your token limit is not too small, because old messages can be trimmed once the buffer fills up.

A quick manual test is enough here: ask for a policy ID, then ask “What’s its status?” or “Can you repeat that?” If the second answer still contains the correct policy reference, memory is working.

Next Steps

• •Replace ChatMemoryBuffer with a persistent store when you need memory across sessions.
• •Add more tools and test whether the agent remembers tool outputs as well as user input.
• •Learn how to summarize long conversations so memory stays useful when chats get large.

Keep learning

• •The complete AI Agents Roadmap — my full 8-step breakdown
• •Free: The AI Agent Starter Kit — PDF checklist + starter code
• •Work with me — I build AI for banks and insurance companies

By Cyprian Aarons, AI Consultant at Topiax.