File size: 1,884 Bytes
40050e3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 | import sqlite3
import json
from langchain.memory import ConversationBufferMemory
DB_FILE = ".conversation_memory.db"
class Memory:
def __init__(self, memory_key="chat_history"):
self.memory_key = memory_key
self.memory = ConversationBufferMemory(memory_key=memory_key, return_messages=True)
self._init_db()
self._load_history()
def _init_db(self):
conn = sqlite3.connect(DB_FILE)
cursor = conn.cursor()
cursor.execute("""
CREATE TABLE IF NOT EXISTS conversations (
id INTEGER PRIMARY KEY AUTOINCREMENT,
question TEXT,
answer TEXT
)
""")
conn.commit()
conn.close()
def _load_history(self):
conn = sqlite3.connect(DB_FILE)
cursor = conn.cursor()
cursor.execute("SELECT question, answer FROM conversations")
rows = cursor.fetchall()
conn.close()
for question, answer in rows:
self.memory.chat_memory.add_user_message(question)
self.memory.chat_memory.add_ai_message(answer)
def save_context(self, inputs, outputs):
"""Save both to memory and to SQLite"""
question = inputs.get("question")
answer = outputs.get("output")
if question and answer:
# Save to in-memory buffer
self.memory.save_context(inputs, outputs)
# Save to SQLite
conn = sqlite3.connect(DB_FILE)
cursor = conn.cursor()
cursor.execute(
"INSERT INTO conversations (question, answer) VALUES (?, ?)",
(question, answer)
)
conn.commit()
conn.close()
def load_memory_variables(self, inputs=None):
return self.memory.load_memory_variables(inputs)
|