Understanding how langchain EntityMemory works

4 min readJun 29, 2023

Prompt: Person trying to recall from memory

Overview of EntityMemory Class

The EntityMemory class is an extension of the BaseChatMemory class. It is part of the memory module in the langchain.

Purpose

This class is designed to handle entities, a term often used to refer to named chunks of information in Natural Language Processing (NLP). The EntityMemory class uses language models and chat history to extract entities and store them.

Properties of EntityMemory Class

It has 3 main properties:

entityExtractionChain: an instance of LLMChain, which extracts entities from the conversation history.
entitySummarizationChain: another instance of LLMChain, which provides summaries for the identified entities.
entityStore: a store to manage the entities identified so far.

In addition to these, it maintains entityCache, a string array to hold the entities, and some other class members and methods such as llm, k, chatHistoryKey, entitiesKey, humanPrefix, aiPrefix, etc.

Key Methods in EntityMemory Class

The EntityMemory class has several methods:

constructor initialises those properties. More notably, it decides on the ‘entity extraction’ and ‘entity summarization’ processes.
loadMemoryVariables loads the memory variables. It gets messages from chatHistory, predicts with entityExtractionChain, and gets and stores summary of each entity from entityStore.
saveContext saves the context from the conversation into the buffer. Here it again predicts with entitySummarizationChain and sets the summary of an entity in entityStore.
clear, which clears the memory contents

Let’s see it in action

We’ll look at an example given in the langchain’s official documentation on the EntityMemory class. You can view the documentation here


const memory = new EntityMemory({
    llm: new OpenAI({ temperature: 0 }),
    chatHistoryKey: "history", // Default value
    entitiesKey: "entities", // Default value
  });
const model = new OpenAI({ temperature: 0.9 });
const chain = new LLMChain({
    llm: model,
    prompt: ENTITY_MEMORY_CONVERSATION_TEMPLATE, // Default prompt - must include the set chatHistoryKey and entitiesKey as input variables.
    memory,
  });
const res1 = await chain.call({ input: "Hi! I'm Jim." });

So what’s happening under the surface?

Calling the `.call` method

Loading entities from the EntityMemory object

The .call internally calls the loadMemoryVariables on the EntityMemory object. The path of execution is as follows:

The InputValues are loaded.
The chatHistory Messages are fetched and serialised.
The entityExtractionChain makes a prediction based on these serialised messages and the input.
The summarised entities and chat history are returned to the LLMChain

EntityMemory saving context

The .call then internally calls the saveContext method of the EntityMemory

First, it saves the chat history context by calling super.saveContext()
A prediction is made using entitySummarizationChain with the latest messages and input data.
If the prediction output is not “UNCHANGED”, it updates the entity store.

Visualising the entire flow

Here’s a sequence diagram to visualise the entire flow:

Mermaid sequence diagram generated by Genie

That’s it!

You now understand how EntityMemory works under the hood and is free to explore different ways that it can be improved upon.

In another article, i mentioned how we can create a continual learning cycle for autonomous agents using EntityMemory.

Although there are many deal breakers that come with the current solution, I believe that in time, the there will be innovations that come along to significantly make progress in our attempts to create a truly adaptive agent.

Side note

I actually built a tool, Genie, to aid me in writing a majority of this article. Genie is a tool that enables you to understand open-sourced projects easily. Here’s a walkthrough of how I came up with the materials using Genie.

More information on Genie

Free to use for developers
Our landing page: www.birdlabs.ai
Accessing the app directly: genie.birdlabs.ai

Using genie to write this article

Getting an overview understanding of the EntityMemory class

2. Genie generates a sequence diagram of the underlying processes that take place in the EntityMemory object when used in conjunction with a LLMChain