LangMem SDK: Personalizing AI Agents with Semantic Memory

While interacting with AI agents, we often find ourselves repeatedly sharing the same preferences, facts, and information. This lack of long-term memory means the agent cannot learn from past conversations or adapt its responses. Imagine if these AI agents could remember your preferences, learn from previous interactions, and optimize its behavior accordingly, retaining the knowledge of your actions, facts, and events over time. Eventually, this makes the agents behave more smartly in the conversations over time. In this article, we’ll uncover how the LangMem SDK equips your agents to learn and adapt using long-term memory. Also, we can segregate memory based on different users and maintain it persistently.

Learning Objectives

• Understanding what is LangMem SDK that is recently introduced by LangChain.
• Learn the memory storage concept it employs to function as a long-term memory.
• Explore the LangMem SDK tool and how to use it.
• Delve into the applications and benefits of integrating LangMem SDK with AI agents.

This article was published as a part of the Data Science Blogathon.

Introduction to LangMem SDK

Recently, Langchain introduced a Software Development Kit (SDK) called LangMem for long-term memory storage that can be integrated with AI agents. Its tools provide functionality to extract information from the conversations. These tools help the agents remember user preferences and provide facts, which eventually fine-tune the prompt and refine the agent’s behavior. In every interaction with the agent, stored memory gets updated. Hence, it helps in developing a better context-aware, coherent, and optimized AI agent by adapting based on memory and providing enhanced performance.

Memory Storage Concept

LangMem has a core memory API that can store memory in any storage, such as backend database or in-memory vector store. Its function is independent of the database with which we are integrating LangMem. Its tools provide direct access to the agents to perform these memory management operations. It extracts new memory based on the conversation, understands the context, and updates the existing memory. This is known as the concept of semantic memory, which is explained in detail below.

Semantic Memory

Semantic memory stores facts that we cannot get from LLM or in the knowledge base. In this, memory is stored in the key-value pair as facts, which we get from the user’s conversations. This semantic memory can be used by the agent later on to retrieve contextually similar memory and adapt the agent accordingly. This type of memory is useful when we cannot get data from the pre-trained models and any integrated knowledge base and we need customization and personalization of the agent on the go. Also, it focuses on the importance of the information, as in the most frequently used information to store. It helps in maintaining a balance between memory creation and consolidation.

As we can see in the below diagram, when the user is specifying client locations during conversation, parallelly it is validating the data and updating the memory vectors as shown below to keep the updated data.

How to Setup and Install LangMem

Let us now look into the setup and install LangMem below:

Step 1: Install Packages

To integrate LangMem with AI agent we need to install langmem package firstly

!pip install -U langmem #This package for Integrating LangMem
!pip install langchain
!pip install langgraph
!pip install -qU "langchain[groq]"

Step 2: Configuring the API keys

Configure the API keys for the provider in the environment variables whose model you plan to use. We will be using open source model from Groq here, so exporting the Groq API key as an environment variable.

Export GROQ_API_KEY = "<your groq api key>"

Step 3: Import Necessary Packages

from langgraph.prebuilt import create_react_agent
from langchain.chat_models import init_chat_model
from langgraph.store.memory import InMemoryStore
from langgraph.store.memory import InMemorySaver
from langmem import create_manage_memory_tool, create_search_memory_tool

In the above imports, create_react_agent is used to create the AI agent with which we will be integrating the langmem. It is based on the Reasoning and Acting (ReAct) concept. init_chat_model is where we will be initializing the chat model and providing a model name that needs to be used for the agent. InMemoryStore is a memory store that we will be using to store vectors. create_manage_memory_tool and create_search_memory_tool are the langmem tools that will be used to extract, manage and optimize the agent by storing long-term memory.

Step 4: Define the Langmem Tools

In the below code snippet, namespace helps in identifying and segmenting the information stored, such as here we are passing “agent_memory” as the namespace. You can give any name you wish to provide.

create_manage_memory_tool is used for storing new information and create_search_memory_tool is used for retrieving past information using semantic search.

tools = [
        create_manage_memory_tool(namespace=("agent_memory",)),
        create_search_memory_tool(namespace=("agent_memory",)),]

Step 5: Setting up a Memory Store

Here we are using InMemory store and defining openai embeddings that we want to use to create embeddings. We are passing “dims” as 1536, which will create an embedding vector of 1536 dimensions, and these vectors will be stored in memory.

store = InMemoryStore(
    index={
        "dims": 1536,
        "embed": "openai:text-embedding-3-small",
    }
)

Step 6: Initialize the Chat Model

Next step is to initialize the chat model that we want to use. we are using open source model llama3 to initialize the agent

model = init_chat_model("llama3-8b-8192", model_provider="groq")

Step 7: Adding Checkpointer

We are adding short-term memory called a checkpointer 

checkpointer = InMemorySaver()

Step 8: Activating the Agents

Finally, we are activating the agent as below by passing all the above-defined parameters.

agent_executor = create_react_agent(llm=model, tools=tools, checkpointer=checkpointer, store=store)

Now, we can execute the agent and test it by doing general interaction with the agent. Langmem tools work in the backend, and memory retrieval happens automatically. We don’t have to pass anything explicitly to the agent.

text = "Hi, Please create two weeks itinerary in short for my Europe trip from India in mid-budget in bullet points"
session_id = 1
result = agent_executor.invoke({"messages": [{"role": "user", "content": text}]}, config={"configurable":{"session_id": session_id}})
print(result["messages"][-1].content)

Output

"Day 1: Arrive in London, UK\n- Arrive in London.\n- Evening: Explore London Eye, Big Ben, Trafalgar Square.\n
Day 2: Explore London\n- Visit The British Museum, Buckingham Palace, Westminster Abbey, Tower Bridge.\n- Evening: Dinner at Borough Market or a pub.\n
Day 3: Day Trip to Stonehenge & Bath\n- Visit Stonehenge and Roman Baths in Bath.\n- Return to London.\n
Day 4: Travel to Paris\n- Take Eurostar train (2.5 hours) to Paris.\n- Evening: Eiffel Tower, Champs-Élysées.\n
Day 5: Explore Paris\n- Visit Louvre Museum, Notre Dame, Montmartre, Sacré-Cœur.\n
Day 6: Day Trip to Versailles\n- Visit Versailles Palace and gardens.\n- Return to Paris.\n
Day 7: Travel to Amsterdam\n- Take Thalys train (3.5 hours) to Amsterdam.\n- Evening: Explore Canals, Anne Frank House, Dam Square.\n
Day 8: Explore Amsterdam\n- Visit Van Gogh Museum, Rijksmuseum, Vondelpark.\n- Canal tour.\n
Day 9: Travel to Berlin\n- Take train (6-7 hours) to Berlin.\n- Evening: Brandenburg Gate, Berlin Wall.\n
Day 10: Explore Berlin\n- Visit Berlin Wall Memorial, Pergamon Museum, East Side Gallery.\n- Evening: Explore Kreuzberg district.\n
Day 11: Travel to Prague\n- Take train (4.5 hours) to Prague.\n- Evening: Explore Old Town Square, Astronomical Clock, Charles Bridge.\n
Day 12: Explore Prague\n- Visit Prague Castle, St. Vitus Cathedral, Golden Lane.\n- Evening: Visit Petrin Hill.\n
Day 13: Travel to Vienna\n- Take train (4 hours) to Vienna.\n- Evening: Explore St. Stephen's Cathedral, Hofburg Palace.\n
Day 14: Explore Vienna & Departure\n- Visit Schönbrunn Palace, Belvedere Palace, Kunsthistorisches Museum.\n- Departure.\n
General Tips:\n- Use city passes for discounts on transport and attractions.\n- Book trains in advance for better prices.\n- Opt for mid-range accommodations like boutique hotels or Airbnb.\n"

Now, we are querying the agent to provide us with the travel options as a follow-up question in the same session_id

text = "Also, pls suggest travel options for above plan in short"
session_id = 1
result = agent_executor.invoke({"messages": [{"role": "user", "content": text}]}, config={"configurable":{"session_id": session_id}})
print(result["messages"][-1].content)

Output

"Day 1: Arrive in London\nTravel Option: Flight from India to London.\n
Day 2-3: Explore London & Stonehenge\nTravel Option: Use London Underground, rent a car for Stonehenge.\n
Day 4-5: Travel to Paris & Explore\nTravel Option: Eurostar train to Paris, use Paris Métro.\n
Day 6: Versailles Day Trip\nTravel Option: Take RER C train.\n
Day 7-8: Travel to Amsterdam & Explore\nTravel Option: Thalys train, bike or canal cruise.\n
Day 9-10: Travel to Berlin & Explore\nTravel Option: Train to Berlin, use U-Bahn/S-Bahn.\n
Day 11-12: Travel to Prague & Explore\nTravel Option: Train to Prague"

Now, we will change the session_id and query a follow-up question and check whether it’s retrieving data based on stored memory or not:

text = "Also, pls suggest food options for my Euope trip"
session_id = 2
result = agent_executor.invoke({"messages": [{"role": "user", "content": text}]}, config={"configurable":{"session_id": session_id}})
print(result["messages"][-1].content)

Output

"Day 1: Arrive in London\nTravel Option: Flight from India to London.\nFood Option: Try Fish and Chips, English Breakfast.\n
Day 2-3: Explore London & Stonehenge\nTravel Option: Use London Underground, rent a car for Stonehenge.\nFood Option: Try Sunday Roast, Beef Wellington.\n
Day 4-5: Travel to Paris & Explore\nTravel Option: Eurostar train to Paris, use Paris Métro.\nFood Option: Try Croissants, Escargot, Coq au Vin.\n
Day 6: Versailles Day Trip\nTravel Option: Take RER C train.\nFood Option: Try French Pastries and Tarte Tatin.\n
Day 7-8: Travel to Amsterdam & Explore\nTravel Option: Thalys train, bike or canal cruise.\nFood Option: Try Stroopwafels, Dutch Pancakes.\n
Day 9-10: Travel to Berlin & Explore\nTravel Option: Train to Berlin, use U-Bahn/S-Bahn.\nFood Option: Try Currywurst, Pretzel, and Sauerkraut.\n
Day 11-12: Travel to Prague & Explore\nTravel Option: Train to Prague, use public trams/metro.\nFood Option: Try Svíčková, Trdelník, Pilsner Beer.\n
Day 13-14: Travel to Vienna & Departure\nTravel Option: Train to Vienna, use Vienna public transport.\nFood Option: Try Wiener Schnitzel, Sachertorte.\n
General Tips: Book trains early, use city passes, opt for mid-range accommodations.\n
"

From the above output, we can see that it remembers the Europe trip plan and suggests food options on top of the same plan. If we do not add LangMem here, then if we change the session_id, the agent won’t be able to remember the context for that user as it had only short-term memory. Here, it tries to find semantic emddings from the memory store and enhance the response for that user. In this way, we can utilize the long-term memory utility of the LangMem to make agents smart.

Leveraging LangMem and Segregating Memory for Multiple Users

If we have multiple users for the AI agent, then we would want separate long-term memory for each user to provide a better conversation experience. We can segregate memories by creating different namespaces and specifying the user_id during runtime to get or update memory for that particular user and maintain privacy. Below is example of how we can achieve it.

namespace = {"agent_memory", "{user_id}"}
text = "travel options for my Europe trip"
session_id = 2
user_id = "ab"
result = agent_executor.invoke({"messages": [{"role": "user", "content": text}]}, config={"configurable":{"session_id": session_id, "user_id": user_id}})

Now, we can see by searching the memory store directly to verify whether it is storing the data separately based on user_id or not. Also to validate memory segregation.

items = store.search(("agent_memory",))
for item in items:
    print(item.namespace, item.value)

Output:

As seen in the above output, information will be stored based on namespace and user_id. In this way, LangMem smartly manages long-term memory for each user. It helps in data privacy and avoids data leakage.

Benefits of Integrating LangMem with AI Agents

• Consistency: The customer will experience continuity across multiple sessions, where the assistant “remembers” past issues and preferences.
• Efficiency: The AI assistant can answer questions or resolve issues faster because it doesn’t need to ask the customer for the same information repeatedly.
• Personalization: The system can offer more personalized solutions, recommendations, and services, which boosts customer satisfaction.

Usecases

• Customer Interaction: A customer reaches out to the AI assistant about a technical issue with a product they purchased six months ago. The assistant can immediately retrieve the product’s warranty details, troubleshoot past support tickets, and offer personalized solutions based on previous interactions.
• Long-Term Memory in Action: Over time, the AI assistant learns more about the customer’s preferences, such as preferred communication tone (formal vs. informal), product usage patterns, or service upgrade preferences. This results in more effective and efficient interactions, saving both the customer and company time.
• Continuous Improvement: As the AI assistant interacts with the customer more, LangMem enables it to retain insights and refine its responses based on past interactions, improving the overall user experience.

Conclusion

In essence, we can see that if we use LangMem’s capabilities to retain memory efficiently, then it can be a valuable integration with the AI agent. It will help companies achieve the enhanced performance from the AI agents. Proper namespace segregation should be done for different contexts, such as user-specific and general. Efficient persistent data storage usage for memory management. Following these practices and incorporating LangMem tools, we can enhance the agents over time by utilizing smart memory capabilities.

Key Takeaways

• LangMem SDK enables AI agents to retain long-term memory, improving conversational coherence and adaptability.
• It employs semantic memory storage to dynamically store and retrieve user-specific facts and preferences.
• LangMem supports integration with various databases, enhancing flexibility in memory management.
• The SDK allows multiple users to maintain separate memory spaces, ensuring personalized AI interactions.
• With LangMem, AI agents evolve over time, refining responses and optimizing behavior based on past interactions.

Frequently Asked Questions

The media shown in this article is not owned by Analytics Vidhya and is used at the Author’s discretion.