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 <dev.dev2049@gmail.com>
This commit is contained in:
Jose Ignacio Hervás Díaz 2023-05-31 03:39:47 +02:00 committed by GitHub
parent 46e181aa8b
commit ce8b7a2a69
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 285 additions and 0 deletions

View File

@ -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
}

View File

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

View File

@ -148,6 +148,98 @@ class RedisEntityStore(BaseEntityStore):
self.redis_client.delete(*keybatch) 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): class ConversationEntityMemory(BaseChatMemory):
"""Entity extractor & summarizer to memory.""" """Entity extractor & summarizer to memory."""