Building an AI Agent with Memory and Adaptability

Full Code Tutorial

Hey folks! 👋

In one of my previous blog posts, I took you through the basics of AI Agents - explaining what they are and walking through a simple implementation using LangGraph. We saw how AI Agents revolutionize the landscape by orchestrating multiple capabilities while maintaining a holistic understanding of tasks.

This time, we're taking things to the next level by implementing three powerful types of memory using LangGraph and LangMem. We'll build these capabilities into an email assistant as our practical example, creating an agent that gets better at understanding your communication needs over time - like a trusted executive assistant rather than a temporary contractor who needs instructions every day.

This blog post is a tutorial based on, and a simplified version of, the course “Long-Term Agentic Memory With LangGraph” by Harrison Chase and Andrew Ng on DeepLearning.AI.

Link to my GitHub repo that contains a code tutorial on memory in agents: Repo Link

Why Memory Matters: From Forgetful Bots to Intelligent Assistants

Imagine having a new assistant who, every single morning, walks into your office with absolutely no recollection of anything that happened the day before. You'd need to re-explain who your key clients are, remind them of ongoing projects, and constantly repeat your preferences. By lunchtime, you'd be exhausted and questioning your hiring decision!

The same principle applies to AI agents. Without memory, they're stuck in a perpetual "present," unable to learn from past interactions or build a deeper understanding of your needs.

Think of it like this:

Without Memory: Your AI is like a goldfish, swimming in circles, treating each email as if it's the first time it's ever seen anything from that sender or on that topic. It's like having an assistant with a clipboard but no filing cabinet – they can handle what's right in front of them, but there's no accumulated wisdom.

With Memory: Your AI becomes like a seasoned executive assistant who remembers that Bob from Accounting always needs extra context in responses, that emails from your biggest client should be prioritized even when they don't say "urgent," and that meeting requests on Fridays should be gently redirected to Monday mornings when you're freshest. The agent builds a mental model of your world and operates within it intelligently.

The email agent that we will implement in this tutorial will leverage three key types of memory, much like the human brain does:

Semantic Memory (Facts): This is your agent's "encyclopedia" – facts it learns about your world. Just as you know that Paris is the capital of France without remembering exactly when or how you learned it, your agent will remember that "Alice is the point of contact for API documentation" or "John prefers morning meetings." It's knowledge that exists independent of specific experiences.

Episodic Memory (Examples): This is your agent's "photo album" – specific memories of past events and interactions. Like how you might remember exactly where you were when you heard important news, your agent will remember, "Last time this client emailed about deadline extensions, my response was too rigid and created friction" or "When emails contain the phrase 'quick question,' they usually require detailed technical explanations that aren't quick at all."

Procedural Memory (Instructions): This is your agent's "muscle memory" – learned behaviors that become automatic. Just as you don't consciously think about each keystroke when typing, your agent will internalize processes like "Always prioritize emails about API documentation" or "Use a more helpful tone in responses to technical questions."

By weaving these types of memory together, our agent becomes truly intelligent and personalized. Imagine an assistant who:

• Notices that emails from a particular client often require follow-up if not answered within 24 hours

• Learns your writing style and tone, adapting formal responses for external communications and casual ones for team members

• Remembers complex project contexts without you having to explain them repeatedly

• Gets better at predicting which emails you'll want to see versus handle automatically

Let's translate this vision into reality by implementing it! :)

Don’t worry, you don’t need to copy-paste the code to run it, you can find everything here: 👉 Full Code Tutorial on GitHub.

Building Our Email Agent: Step-by-Step

Let's get our hands dirty and start building!

1. Imports and Setup

import os
from dotenv import load_dotenv
from typing import TypedDict, Literal, Annotated, List
from langgraph.graph import StateGraph, START, END, add_messages
from langchain.chat_models import init_chat_model
from langchain_core.tools import tool
from langchain.prompts import PromptTemplate
from langchain.schema import HumanMessage
from pydantic import BaseModel, Field
from langgraph.store.memory import InMemoryStore  # For storing memories
from langmem import create_manage_memory_tool, create_search_memory_tool 

Now, let's initialize our environment and tools:

# Load environment variables (your OpenAI API key for example)
load_dotenv()

# Initialize the LLM
llm = init_chat_model("openai:gpt-4o-mini")

# Initialize the memory store (we'll use an in-memory store for simplicity)
store = InMemoryStore(index={"embed": "openai:text-embedding-3-small"})

The memory store is particularly important here - it's like giving our agent a brain where it can store and retrieve information. We're using an in-memory store for simplicity, but in a production environment, you might want to use a persistent database.

2. Defining Our Agent's "Brain": The State 🧠

Now we need to design our agent's working memory - the mental scratchpad where it keeps track of what it's currently processing. This is different from the long-term memory we'll implement later - it's more like the things you actively hold in mind while working on a task.

class State(TypedDict):
    email_input: dict  # The incoming email
    messages: Annotated[list, add_messages]  # The conversation history
    triage_result: str # The result of the triage (ignore, notify, respond)

This State object holds three crucial pieces of information:

• The current email being processed

• The ongoing conversation (if any)

• The decision about how to handle the email

Think of it like a doctor's clipboard during patient rounds - it contains the immediate information needed to make decisions, while the patient's full medical history remains in the chart room (our memory store).

3. The Triage Center: Deciding What to Do (with Episodic Memory) 🤔

First, we'll create a structure for our agent to explain its reasoning and classification:

class Router(BaseModel):
    reasoning: str = Field(description="Step-by-step reasoning behind the classification.")
    classification: Literal["ignore", "respond", "notify"] = Field(
        description="The classification of an email: 'ignore', 'notify', or 'respond'."
    )

llm_router = llm.with_structured_output(Router)

To leverage episodic memory, we need a way to format examples from past interactions:

def format_few_shot_examples(examples):
    formatted_examples = []
    for eg in examples:
        email = eg.value['email']
        label = eg.value['label']
        formatted_examples.append(
            f"From: {email['author']}\nSubject: {email['subject']}\nBody: {email['email_thread'][:300]}...\n\nClassification: {label}"
        )
    return "\n\n".join(formatted_examples)

This function transforms our stored examples into a format that helps the model learn from them - like showing a new employee a training manual with annotated examples of how to handle different situations.

Now, let's create our email triage function that uses episodic memory:

def triage_email(state: State, config: dict, store: InMemoryStore) -> dict:
    email = state["email_input"]
    user_id = config["configurable"]["langgraph_user_id"]
    namespace = ("email_assistant", user_id, "examples")  # Namespace for episodic memory

    # Retrieve relevant examples from memory
    examples = store.search(namespace, query=str(email))
    formatted_examples = format_few_shot_examples(examples)

    prompt_template = PromptTemplate.from_template("""You are an email triage assistant.  Classify the following email:
    From: {author}
    To: {to}
    Subject: {subject}
    Body: {email_thread}

    Classify as 'ignore', 'notify', or 'respond'.

    Here are some examples of previous classifications:
    {examples}
    """)

    prompt = prompt_template.format(examples=formatted_examples, **email)
    messages = [HumanMessage(content=prompt)]
    result = llm_router.invoke(messages)
    return {"triage_result": result.classification}

This function is the core of episodic memory at work. When a new email arrives, it doesn't analyze it in isolation - it searches for similar emails from the past and sees how those were handled. It's like a doctor who remembers, "The last three patients with these symptoms responded well to this treatment."

4. Defining Tools with Semantic Memory 🧠📚🔤

Now let's give our agent some tools to work with. First, basic abilities to write emails and check calendars:

@tool
def write_email(to: str, subject: str, content: str) -> str:
    """Write and send an email."""
    print(f"Sending email to {to} with subject '{subject}'\nContent:\n{content}\n")
    return f"Email sent to {to} with subject '{subject}'"

@tool
def check_calendar_availability(day: str) -> str:
    """Check calendar availability for a given day."""
    return f"Available times on {day}: 9:00 AM, 2:00 PM, 4:00 PM"

Note that these are simplified implementations for demonstration purposes. In a production environment, you would connect these functions to actual email and calendar APIs (like Gmail API or Microsoft Graph API). For example, the write_email function would interact with a real email service to send messages, and check_calendar_availability would query your actual calendar data.

Now, let's add the semantic memory tools - our agent's ability to store and retrieve facts about the world:

# Create LangMem memory tools 
manage_memory_tool = create_manage_memory_tool(namespace=("email_assistant", "{langgraph_user_id}", "collection"))
search_memory_tool = create_search_memory_tool(namespace=("email_assistant", "{langgraph_user_id}", "collection"))

tools = [write_email, check_calendar_availability, manage_memory_tool, search_memory_tool]

The manage_memory_tool is like giving our agent a notebook where it can jot down important facts, while the search_memory_tool lets it flip through that notebook to find relevant information when needed.

5. The Response Agent: Creating Our Core Assistant (with Semantic Memory ) 🤖💡💬

We'll now create the core agent that handles responses using all its memory systems:

from langgraph.prebuilt import create_react_agent

def create_agent_prompt(state, config, store):
    messages = state['messages']
    user_id = config["configurable"]["langgraph_user_id"]
    
    # Get the current response prompt from procedural memory
    system_prompt = store.get(("email_assistant", user_id, "prompts"), "response_prompt").value
    
    return [{"role": "system", "content": system_prompt}] + messages

response_agent = create_react_agent(
    tools=tools,
    prompt=create_agent_prompt,
    store=store,
    model=llm  # Using 'model' instead of 'llm'
)

If you want to learn how to control your prompts correctly, I’ve written a dedicated book for this (newsletter subscribers are eligible for a 33% discount)

This function creates a prompt that pulls instructions from procedural memory and passes the current conversation along with it. It's like a manager who checks the company handbook before responding to a complex situation, ensuring they follow the latest protocols.

Note that we've now set up two key memory systems:

• Episodic memory (in the triage function)

• Semantic memory (in these agent tools)

But we still need to add procedural memory to complete our agent's cognitive abilities. This will come in the following sections where we'll enable our agent to refine its own behavior over time based on feedback.

6. Building the Graph: Connecting the Pieces

Now, let's bring everything together into a cohesive workflow:

workflow = StateGraph(State)

# Update this line to pass the store to the node
workflow.add_node("triage", lambda state, config: triage_email(state, config, store))
workflow.add_node("response_agent", response_agent)

def route_based_on_triage(state):
  if state["triage_result"] == "respond":
    return "response_agent"
  else:
    return END

workflow.add_edge(START, "triage")
workflow.add_conditional_edges("triage", route_based_on_triage,
                              {
                                  "response_agent": "response_agent",
                                  END: END
                              })

# Compile the graph
agent = workflow.compile(store=store)

This is how our agent looks:

This workflow defines the logical path for our agent:

1. First, triage the incoming email using episodic memory

2. If it needs a response, activate the response agent with semantic memory

3. Otherwise, end the process (for "ignore" or "notify" emails)

It's like setting up assembly line stations in a factory - each piece has its own job, but they work together to create the final product.

7. Let's Run It! (and Store Some Memories) 🧠💭📸

Time to put our agent to the test:

email_input = {
    "author": "Alice Smith <alice.smith@company.com>",
    "to": "John Doe <john.doe@company.com>",
    "subject": "Quick question about API documentation",
    "email_thread": """Hi John,

I was reviewing the API documentation and noticed a few endpoints are missing. Could you help?

Thanks,
Alice""",
}

config = {"configurable": {"langgraph_user_id": "test_user"}} # Set the user ID!
inputs = {"email_input": email_input, "messages": []}

for output in agent.stream(inputs, config=config): # Pass the config
    for key, value in output.items():
        print(f"-----\n{key}:")
        print(value)
    print("-----")

response_agent: "I'm here to help you with any questions you have about API documentation. If you're facing issues like missing endpoints or need clarification on specific sections, just let me know! I can provide you with guidance and useful resources. Feel free to share the details of your inquiry, and I'll get right on it!"

Let's also add a training example to our episodic memory, to help the agent recognize spam in the future:

#add few shot examples to memory
example1 = {
    "email": {
        "author": "Spammy Marketer <spam@example.com>",
        "to": "John Doe <john.doe@company.com>",
        "subject": "BIG SALE!!!",
        "email_thread": "Buy our product now and get 50% off!",
    },
    "label": "ignore",
}
store.put(("email_assistant", "test_user", "examples"), "spam_example", example1)

This is like training a new assistant: "See this kind of email? You can safely ignore these." The more examples we provide, the more nuanced the agent's understanding becomes.

8. Adding Procedural Memory 🧠⚙️🏃 (Updating Instructions) - The Final Touch!

Now for the most sophisticated memory system - procedural memory that allows our agent to improve its own instructions based on feedback:

initial_triage_prompt = """You are an email triage assistant. Classify the following email:
From: {author}
To: {to}
Subject: {subject}
Body: {email_thread}

Classify as 'ignore', 'notify', or 'respond'.

Here are some examples of previous classifications:
{examples}
"""

initial_response_prompt = """You are a helpful assistant. Use the tools available, including memory tools, to assist the user."""

# Store these prompts in the memory store
store.put(("email_assistant", "test_user", "prompts"), "triage_prompt", initial_triage_prompt)

store.put(("email_assistant", "test_user", "prompts"), "response_prompt", initial_response_prompt)

These are our starting points - the basic instructions our agent begins with. But unlike most systems, ours can evolve these instructions over time.

Let's create a version of our triage function that pulls its instructions from memory:

def triage_email_with_procedural_memory(state: State, config: dict, store: InMemoryStore) -> dict:
    email = state["email_input"]
    user_id = config["configurable"]["langgraph_user_id"]
    
    # Retrieve the current triage prompt (procedural memory)
    current_prompt_template = store.get(("email_assistant", user_id, "prompts"), "triage_prompt").value
    
    # Retrieve relevant examples from memory (episodic memory)
    namespace = ("email_assistant", user_id, "examples")
    examples = store.search(namespace, query=str(email))
    formatted_examples = format_few_shot_examples(examples)
    
    # Format the prompt
    prompt = PromptTemplate.from_template(current_prompt_template).format(examples=formatted_examples, **email)
    messages = [HumanMessage(content=prompt)]
    result = llm_router.invoke(messages)
    return {"triage_result": result.classification}

This function integrates procedural memory (the current prompt template) with episodic memory (relevant examples) to make triage decisions.

Now, let's create a function that can improve our prompts based on feedback:

from langmem import create_multi_prompt_optimizer

def optimize_prompts(feedback: str, config: dict, store: InMemoryStore):
    """Improve our prompts based on feedback."""
    user_id = config["configurable"]["langgraph_user_id"]
    
    # Get current prompts
    triage_prompt = store.get(("email_assistant", user_id, "prompts"), "triage_prompt").value
    response_prompt = store.get(("email_assistant", user_id, "prompts"), "response_prompt").value
    
    # Create a more relevant test example based on our actual email
    sample_email = {
        "author": "Alice Smith <alice.smith@company.com>",
        "to": "John Doe <john.doe@company.com>",
        "subject": "Quick question about API documentation",
        "email_thread": "Hi John, I was reviewing the API documentation and noticed a few endpoints are missing. Could you help? Thanks, Alice",
    }
    
    # Create the optimizer
    optimizer = create_multi_prompt_optimizer(llm)
    
    # Create a more relevant conversation trajectory with feedback
    conversation = [
        {"role": "system", "content": response_prompt},
        {"role": "user", "content": f"I received this email: {sample_email}"},
        {"role": "assistant", "content": "How can I assist you today?"}
    ]
    
    # Format prompts
    prompts = [
        {"name": "triage", "prompt": triage_prompt},
        {"name": "response", "prompt": response_prompt}
    ]
    
    
    # More relevant trajectories 
    trajectories = [(conversation, {"feedback": feedback})]
    result = optimizer.invoke({"trajectories": trajectories, "prompts": prompts})
        
    # Extract the improved prompts
    improved_triage_prompt = next(p["prompt"] for p in result if p["name"] == "triage")
    improved_response_prompt = next(p["prompt"] for p in result if p["name"] == "response")
        
     documentation or missing endpoints are high priority and should ALWAYS be classified as 'respond'."
        improved_response_prompt = response_prompt + "\n\nWhen responding to emails about documentation or API issues, acknowledge the specific issue mentioned and offer specific assistance rather than generic responses."
    
    # Store the improved prompts
    store.put(("email_assistant", user_id, "prompts"), "triage_prompt", improved_triage_prompt)
    store.put(("email_assistant", user_id, "prompts"), "response_prompt", improved_response_prompt)
    
    print(f"Triage prompt improved: {improved_triage_prompt[:100]}...")
    print(f"Response prompt improved: {improved_response_prompt[:100]}...")
    
    return "Prompts improved based on feedback!"

This function is the essence of procedural memory. It takes feedback like 'You're not prioritizing API documentation emails correctly' and uses it to rewrite the agent's core instructions. The optimizer works like a coach watching game footage, studying what went wrong and updating the playbook accordingly. It analyzes conversation examples alongside feedback, then refines the prompts that guide the agent's behavior. Instead of just memorizing specific corrections, it absorbs the underlying lessons into its overall approach, similar to how a chef improves recipes based on customer feedback rather than simply following different instructions each time.

9. Let's Run Our Complete Memory-Enhanced Agent!

Now let's bring everything together into a complete system that can evolve over time:

def create_email_agent(store):
    # Define the workflow
    workflow = StateGraph(State)
    workflow.add_node("triage", lambda state, config: triage_email_with_procedural_memory(state, config, store))
    
    # Create a fresh response agent that will use the latest prompts
    response_agent = create_react_agent(
        tools=tools,
        prompt=create_agent_prompt,
        store=store,
        model=llm
    )
    
    workflow.add_node("response_agent", response_agent)
    
    # The routing logic remains the same
    workflow.add_edge(START, "triage")
    workflow.add_conditional_edges("triage", route_based_on_triage,
                                {
                                    "response_agent": "response_agent",
                                    END: END
                                })
    
    # Compile and return the graph
    return workflow.compile(store=store)

This function creates a fresh agent that uses the most current version of our prompts - ensuring it always reflects our latest learnings and feedback.

Now this is the final version of our agent (including everything):

Let's run it twice - once with the original settings, and once after we've provided feedback to improve it:

# First process the original email to capture "before" behavior
print("\n\nProcessing original email BEFORE optimization...\n\n")
agent = create_email_agent(store)  # Create a fresh agent
for output in agent.stream(inputs, config=config):
    for key, value in output.items():
        print(f"-----\n{key}:")
        print(value)
    print("-----")

# Add a specific example to episodic memory
api_doc_example = {
    "email": {
        "author": "Developer <dev@company.com>",
        "to": "John Doe <john.doe@company.com>",
        "subject": "API Documentation Issue", 
        "email_thread": "Found missing endpoints in the API docs. Need urgent update.",
    },
    "label": "respond",
}
store.put(("email_assistant", "test_user", "examples"), "api_doc_example", api_doc_example)
print("Added API documentation example to episodic memory")

# Provide feedback
feedback = """The agent didn't properly recognize that emails about API documentation issues 
are high priority and require immediate attention. When an email mentions 
'API documentation', it should always be classified as 'respond' with a helpful tone.
Also, instead of just responding with 'How can I assist you today?', the agent should 
acknowledge the specific documentation issue mentioned and offer assistance."""

# Optimize prompts
optimize_prompts(feedback, config, store)

# Process the SAME email after optimization with a FRESH agent
print("\n\nProcessing the SAME email AFTER optimization with a fresh agent...\n\n")
new_agent = create_email_agent(store)  # Create a fresh agent with updated prompts
for output in new_agent.stream(inputs, config=config):
    for key, value in output.items():
        print(f"-----\n{key}:")
        print(value)
    print("-----")

Let’s look at the output:

- Processing original email BEFORE optimization...

response_agent: "It seems like you're inquiring about API documentation. If you have specific questions or need certain information, please let me know! I can assist you with that or help you find particular resources related to API documentation."

- Added API documentation example to episodic memory Triage prompt improved: You are an email triage assistant. Classify the following email: From: {author} To: {to} Subject:... Response prompt improved: You are a helpful assistant. Prioritize inquiries about critical topics like API documentation by ac...

- Processing the SAME email AFTER optimization with a fresh agent...

response_agent: "It sounds like you're inquiring about API documentation, and I noticed you mentioned a specific issue regarding missing endpoints. Can you please provide more details about which endpoints you're looking for or the specific API you're referring to? This will help me guide you more effectively in resolving the issue."

Look at the difference in responses! After our feedback, the agent should:

1. More consistently recognize API documentation issues as high-priority

2. Provide more specific, helpful responses that acknowledge the actual issue

3. Offer concrete assistance rather than generic platitudes

Conclusion

We've now built an email agent that's far more than a simple script. Like a skilled human assistant who grows more valuable over time, our agent builds a multi-faceted memory system:

1. Semantic Memory: A knowledge base of facts about your work context, contacts, and preferences

2. Episodic Memory: A collection of specific examples that guide decision-making through pattern recognition

3. Procedural Memory: The ability to improve its own processes based on feedback and experience

This agent demonstrates how combining different types of memory creates an assistant that actually learns from interactions and gets better over time.

Imagine coming back from a two-week vacation to find that your AI assistant has not only kept your inbox under control but has done so in a way that reflects your priorities and communication style. The spam is gone, the urgent matters were flagged appropriately, and routine responses were handled so well that recipients didn't even realize they were talking to an AI. That's the power of memory-enhanced agents.

This is just a starting point! You can extend this agent with more sophisticated tools, persistent storage for long-term memory, fine-grained feedback mechanisms, and even collaborative capabilities that let multiple agents share knowledge while maintaining privacy boundaries.