SQLite-backed Entity Memory (langchain-ai#5129)

# SQLite-backed Entity Memory

Following the initiative of
langchain-ai#2397 I think it would be
helpful to be able to persist Entity Memory on disk by default

Co-authored-by: Dev 2049 <dev.dev2049@gmail.com>

docs/modules/memory/examples/entity_memory_with_sqlite.ipynb

@@ -0,0 +1,191 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "eg0Hwptz9g5q"
+   },
+   "source": [
+    "# Entity Memory with SQLite storage\n",
+    "\n",
+    "In this walkthrough we'll create a simple conversation chain which uses ConversationEntityMemory backed by a SqliteEntityStore."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {
+    "id": "2wUMSUoF8ffn"
+   },
+   "outputs": [],
+   "source": [
+    "from langchain.chains import ConversationChain\n",
+    "from langchain.llms import OpenAI\n",
+    "from langchain.memory import ConversationEntityMemory\n",
+    "from langchain.memory.entity import SQLiteEntityStore\n",
+    "from langchain.memory.prompt import ENTITY_MEMORY_CONVERSATION_TEMPLATE"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {
+    "id": "8TpJZti99gxV"
+   },
+   "outputs": [],
+   "source": [
+    "entity_store=SQLiteEntityStore()\n",
+    "llm = OpenAI(temperature=0)\n",
+    "memory = ConversationEntityMemory(llm=llm, entity_store=entity_store)\n",
+    "conversation = ConversationChain(\n",
+    "    llm=llm, \n",
+    "    prompt=ENTITY_MEMORY_CONVERSATION_TEMPLATE,\n",
+    "    memory=memory,\n",
+    "    verbose=True,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "HEAHG1L79ca1"
+   },
+   "source": [
+    "Notice the usage of `EntitySqliteStore` as parameter to `entity_store` on the `memory` property."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {
+    "colab": {
+     "base_uri": "https://localhost:8080/",
+     "height": 437
+    },
+    "id": "BzXphJWf_TAZ",
+    "outputId": "de7fc966-e0fd-4daf-a9bd-4743455ea774"
+   },
+   "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;3mYou are an assistant to a human, powered by a large language model trained by OpenAI.\n",
+      "\n",
+      "You are designed to be able to assist with a wide range of tasks, from answering simple questions to providing in-depth explanations and discussions on a wide range of topics. As a language model, you are able to generate human-like text based on the input you receive, allowing you to engage in natural-sounding conversations and provide responses that are coherent and relevant to the topic at hand.\n",
+      "\n",
+      "You are constantly learning and improving, and your capabilities are constantly evolving. You are able to process and understand large amounts of text, and can use this knowledge to provide accurate and informative responses to a wide range of questions. You have access to some personalized information provided by the human in the Context section below. Additionally, you are able to generate your own text based on the input you receive, allowing you to engage in discussions and provide explanations and descriptions on a wide range of topics.\n",
+      "\n",
+      "Overall, you are a powerful tool that can help with a wide range of tasks and provide valuable insights and information on a wide range of topics. Whether the human needs help with a specific question or just wants to have a conversation about a particular topic, you are here to assist.\n",
+      "\n",
+      "Context:\n",
+      "{'Deven': 'Deven is working on a hackathon project with Sam.', 'Sam': 'Sam is working on a hackathon project with Deven.'}\n",
+      "\n",
+      "Current conversation:\n",
+      "\n",
+      "Last line:\n",
+      "Human: Deven & Sam are working on a hackathon project\n",
+      "You:\u001b[0m\n",
+      "\n",
+      "\u001b[1m> Finished chain.\u001b[0m\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "' That sounds like a great project! What kind of project are they working on?'"
+      ]
+     },
+     "execution_count": 3,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "conversation.run(\"Deven & Sam are working on a hackathon project\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {
+    "colab": {
+     "base_uri": "https://localhost:8080/",
+     "height": 35
+    },
+    "id": "YsFE3hBjC6gl",
+    "outputId": "56ab5ca9-e343-41b5-e69d-47541718a9b4"
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "'Deven is working on a hackathon project with Sam.'"
+      ]
+     },
+     "execution_count": 4,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "conversation.memory.entity_store.get(\"Deven\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "'Sam is working on a hackathon project with Deven.'"
+      ]
+     },
+     "execution_count": 5,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "conversation.memory.entity_store.get(\"Sam\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "colab": {
+   "provenance": []
+  },
+  "kernelspec": {
+   "display_name": "venv",
+   "language": "python",
+   "name": "venv"
+  },
+  "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.11.3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 1
+}

langchain/memory/__init__.py

@@ -19,6 +19,7 @@
     ConversationEntityMemory,
     InMemoryEntityStore,
     RedisEntityStore,
+    SQLiteEntityStore,
 )
 from langchain.memory.kg import ConversationKGMemory
 from langchain.memory.readonly import ReadOnlySharedMemory
@@ -38,6 +39,7 @@
     "ConversationEntityMemory",
     "InMemoryEntityStore",
     "RedisEntityStore",
+    "SQLiteEntityStore",
     "ConversationSummaryMemory",
     "ChatMessageHistory",
     "ConversationStringBufferMemory",

langchain/memory/entity.py

@@ -148,6 +148,98 @@ def batched(iterable: Iterable[Any], batch_size: int) -> Iterable[Any]:
             self.redis_client.delete(*keybatch)
 
 
+class SQLiteEntityStore(BaseEntityStore):
+    """SQLite-backed Entity store"""
+
+    session_id: str = "default"
+    table_name: str = "memory_store"
+
+    def __init__(
+        self,
+        session_id: str = "default",
+        db_file: str = "entities.db",
+        table_name: str = "memory_store",
+        *args: Any,
+        **kwargs: Any,
+    ):
+        try:
+            import sqlite3
+        except ImportError:
+            raise ImportError(
+                "Could not import sqlite3 python package. "
+                "Please install it with `pip install sqlite3`."
+            )
+        super().__init__(*args, **kwargs)
+
+        self.conn = sqlite3.connect(db_file)
+        self.session_id = session_id
+        self.table_name = table_name
+        self._create_table_if_not_exists()
+
+    @property
+    def full_table_name(self) -> str:
+        return f"{self.table_name}_{self.session_id}"
+
+    def _create_table_if_not_exists(self) -> None:
+        create_table_query = f"""
+            CREATE TABLE IF NOT EXISTS {self.full_table_name} (
+                key TEXT PRIMARY KEY,
+                value TEXT
+            )
+        """
+        with self.conn:
+            self.conn.execute(create_table_query)
+
+    def get(self, key: str, default: Optional[str] = None) -> Optional[str]:
+        query = f"""
+            SELECT value
+            FROM {self.full_table_name}
+            WHERE key = ?
+        """
+        cursor = self.conn.execute(query, (key,))
+        result = cursor.fetchone()
+        if result is not None:
+            value = result[0]
+            return value
+        return default
+
+    def set(self, key: str, value: Optional[str]) -> None:
+        if not value:
+            return self.delete(key)
+        query = f"""
+            INSERT OR REPLACE INTO {self.full_table_name} (key, value)
+            VALUES (?, ?)
+        """
+        with self.conn:
+            self.conn.execute(query, (key, value))
+
+    def delete(self, key: str) -> None:
+        query = f"""
+            DELETE FROM {self.full_table_name}
+            WHERE key = ?
+        """
+        with self.conn:
+            self.conn.execute(query, (key,))
+
+    def exists(self, key: str) -> bool:
+        query = f"""
+            SELECT 1
+            FROM {self.full_table_name}
+            WHERE key = ?
+            LIMIT 1
+        """
+        cursor = self.conn.execute(query, (key,))
+        result = cursor.fetchone()
+        return result is not None
+
+    def clear(self) -> None:
+        query = f"""
+            DELETE FROM {self.full_table_name}
+        """
+        with self.conn:
+            self.conn.execute(query)
+
+
 class ConversationEntityMemory(BaseChatMemory):
     """Entity extractor & summarizer to memory."""