From ce8b7a2a69d739bf2e08347b72ff46172ef60aa6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jose=20Ignacio=20Herv=C3=A1s=20D=C3=ADaz?= Date: Wed, 31 May 2023 03:39:47 +0200 Subject: [PATCH] SQLite-backed Entity Memory (#5129) # SQLite-backed Entity Memory Following the initiative of https://github.com/hwchase17/langchain/pull/2397 I think it would be helpful to be able to persist Entity Memory on disk by default Co-authored-by: Dev 2049 --- .../examples/entity_memory_with_sqlite.ipynb | 191 ++++++++++++++++++ langchain/memory/__init__.py | 2 + langchain/memory/entity.py | 92 +++++++++ 3 files changed, 285 insertions(+) create mode 100644 docs/modules/memory/examples/entity_memory_with_sqlite.ipynb diff --git a/docs/modules/memory/examples/entity_memory_with_sqlite.ipynb b/docs/modules/memory/examples/entity_memory_with_sqlite.ipynb new file mode 100644 index 00000000..f71e1b4f --- /dev/null +++ b/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 +} diff --git a/langchain/memory/__init__.py b/langchain/memory/__init__.py index e11316e8..69786273 100644 --- a/langchain/memory/__init__.py +++ b/langchain/memory/__init__.py @@ -19,6 +19,7 @@ from langchain.memory.entity import ( ConversationEntityMemory, InMemoryEntityStore, RedisEntityStore, + SQLiteEntityStore, ) from langchain.memory.kg import ConversationKGMemory from langchain.memory.readonly import ReadOnlySharedMemory @@ -38,6 +39,7 @@ __all__ = [ "ConversationEntityMemory", "InMemoryEntityStore", "RedisEntityStore", + "SQLiteEntityStore", "ConversationSummaryMemory", "ChatMessageHistory", "ConversationStringBufferMemory", diff --git a/langchain/memory/entity.py b/langchain/memory/entity.py index 381b2b18..a7808e15 100644 --- a/langchain/memory/entity.py +++ b/langchain/memory/entity.py @@ -148,6 +148,98 @@ class RedisEntityStore(BaseEntityStore): 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."""