.ipynb
.pdf
How to create a custom Memory class
How to create a custom Memory class#
Although there are a few predefined types of memory in LangChain, it is highly possible you will want to add your own type of memory that is optimal for your application. This notebook covers how to do that.
For this notebook, we will add a custom memory type to ConversationChain. In order to add a custom memory class, we need to import the base memory class and subclass it.
from langchain import OpenAI, ConversationChain
from langchain.schema import BaseMemory
from pydantic import BaseModel
from typing import List, Dict, Any
In this example, we will write a custom memory class that uses spacy to extract entities and save information about them in a simple hash table. Then, during the conversation, we will look at the input text, extract any entities, and put any information about them into the context.
Please note that this implementation is pretty simple and brittle and probably not useful in a production setting. Its purpose is to showcase that you can add custom memory implementations.
For this, we will need spacy.
# !pip install spacy
# !python -m spacy download en_core_web_lg
import spacy
nlp = spacy.load('en_core_web_lg')
class SpacyEntityMemory(BaseMemory, BaseModel):
    """Memory class for storing information about entities."""
    # Define dictionary to store information about entities.
    entities: dict = {}
    # Define key to pass information about entities into prompt.
    memory_key: str = "entities"
        
    def clear(self):
        self.entities = {}
    @property
    def memory_variables(self) -> List[str]:
        """Define the variables we are providing to the prompt."""
        return [self.memory_key]
rtdocs_stable/api.python.langchain.com/en/stable/modules/memory/examples/custom_memory.html