langchain/docs/modules/memory/examples/custom_memory.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "94e33ebe",
   "metadata": {},
   "source": [
    "# How to create a custom Memory class\n",
    "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."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bdfd0305",
   "metadata": {},
   "source": [
    "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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "6d787ef2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain import OpenAI, ConversationChain\n",
    "from langchain.schema import BaseMemory\n",
    "from pydantic import BaseModel\n",
    "from typing import List, Dict, Any"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9489e5e1",
   "metadata": {},
   "source": [
    "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.\n",
    "\n",
    "* 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.\n",
    "\n",
    "For this, we will need spacy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48a5dd13",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install spacy\n",
    "# !python -m spacy download en_core_web_lg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "ff065f58",
   "metadata": {},
   "outputs": [],
   "source": [
    "import spacy\n",
    "nlp = spacy.load('en_core_web_lg')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "1d45d429",
   "metadata": {},
   "outputs": [],
   "source": [
    "class SpacyEntityMemory(BaseMemory, BaseModel):\n",
    "    \"\"\"Memory class for storing information about entities.\"\"\"\n",
    "\n",
    "    # Define dictionary to store information about entities.\n",
    "    entities: dict = {}\n",
    "    # Define key to pass information about entities into prompt.\n",
    "    memory_key: str = \"entities\"\n",
    "        \n",
    "    def clear(self):\n",
    "        self.entities = {}\n",
    "\n",
    "    @property\n",
    "    def memory_variables(self) -> List[str]:\n",
    "        \"\"\"Define the variables we are providing to the prompt.\"\"\"\n",
    "        return [self.memory_key]\n",
    "\n",
    "    def load_memory_variables(self, inputs: Dict[str, Any]) -> Dict[str, str]:\n",
    "        \"\"\"Load the memory variables, in this case the entity key.\"\"\"\n",
    "        # Get the input text and run through spacy\n",
    "        doc = nlp(inputs[list(inputs.keys())[0]])\n",
    "        # Extract known information about entities, if they exist.\n",
    "        entities = [self.entities[str(ent)] for ent in doc.ents if str(ent) in self.entities]\n",
    "        # Return combined information about entities to put into context.\n",
    "        return {self.memory_key: \"\\n\".join(entities)}\n",
    "\n",
    "    def save_context(self, inputs: Dict[str, Any], outputs: Dict[str, str]) -> None:\n",
    "        \"\"\"Save context from this conversation to buffer.\"\"\"\n",
    "        # Get the input text and run through spacy\n",
    "        text = inputs[list(inputs.keys())[0]]\n",
    "        doc = nlp(text)\n",
    "        # For each entity that was mentioned, save this information to the dictionary.\n",
    "        for ent in doc.ents:\n",
    "            ent_str = str(ent)\n",
    "            if ent_str in self.entities:\n",
    "                self.entities[ent_str] += f\"\\n{text}\"\n",
    "            else:\n",
    "                self.entities[ent_str] = text"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "429ba264",
   "metadata": {},
   "source": [
    "We now define a prompt that takes in information about entities as well as user input"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "c05159b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.prompts.prompt import PromptTemplate\n",
    "\n",
    "template = \"\"\"The following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know. You are provided with information about entities the Human mentions, if relevant.\n",
    "\n",
    "Relevant entity information:\n",
    "{entities}\n",
    "\n",
    "Conversation:\n",
    "Human: {input}\n",
    "AI:\"\"\"\n",
    "prompt = PromptTemplate(\n",
    "    input_variables=[\"entities\", \"input\"], template=template\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db611041",
   "metadata": {},
   "source": [
    "And now we put it all together!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "f08dc8ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = OpenAI(temperature=0)\n",
    "conversation = ConversationChain(llm=llm, prompt=prompt, verbose=True, memory=SpacyEntityMemory())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "92a5f685",
   "metadata": {},
   "source": [
    "In the first example, with no prior knowledge about Harrison, the \"Relevant entity information\" section is empty."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "5b96e836",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new ConversationChain chain...\u001b[0m\n",
      "Prompt after formatting:\n",
      "\u001b[32;1m\u001b[1;3mThe following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know. You are provided with information about entities the Human mentions, if relevant.\n",
      "\n",
      "Relevant entity information:\n",
      "\n",
      "\n",
      "Conversation:\n",
      "Human: Harrison likes machine learning\n",
      "AI:\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished ConversationChain chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "\" That's great to hear! Machine learning is a fascinating field of study. It involves using algorithms to analyze data and make predictions. Have you ever studied machine learning, Harrison?\""
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conversation.predict(input=\"Harrison likes machine learning\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b1faa743",
   "metadata": {},
   "source": [
    "Now in the second example, we can see that it pulls in information about Harrison."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "4bca7070",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new ConversationChain chain...\u001b[0m\n",
      "Prompt after formatting:\n",
      "\u001b[32;1m\u001b[1;3mThe following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know. You are provided with information about entities the Human mentions, if relevant.\n",
      "\n",
      "Relevant entity information:\n",
      "Harrison likes machine learning\n",
      "\n",
      "Conversation:\n",
      "Human: What do you think Harrison's favorite subject in college was?\n",
      "AI:\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished ConversationChain chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "' From what I know about Harrison, I believe his favorite subject in college was machine learning. He has expressed a strong interest in the subject and has mentioned it often.'"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conversation.predict(input=\"What do you think Harrison's favorite subject in college was?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "58b856e3",
   "metadata": {},
   "source": [
    "Again, 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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1994600",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}