Spaces:
Runtime error
Runtime error
File size: 1,659 Bytes
e8d69f7 |
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 |
import sqlite3
import json
from typing import List, Dict
class MemoryManager:
def __init__(self, db_path: str = "memory.db"):
self.db_path = db_path
self._init_db()
def _init_db(self):
with sqlite3.connect(self.db_path) as conn:
conn.execute("""
CREATE TABLE IF NOT EXISTS sessions (
session_id TEXT PRIMARY KEY,
history TEXT
)
""")
conn.commit()
def get_history(self, session_id: str) -> List[Dict[str, str]]:
with sqlite3.connect(self.db_path) as conn:
cursor = conn.execute("SELECT history FROM sessions WHERE session_id = ?", (session_id,))
row = cursor.fetchone()
if row:
return json.loads(row[0])
return []
def save_message(self, session_id: str, role: str, content: str):
history = self.get_history(session_id)
history.append({"role": role, "content": content})
# Keep only last 10 messages for context window management if needed,
# but the prompt says persistent history.
with sqlite3.connect(self.db_path) as conn:
conn.execute("""
INSERT OR REPLACE INTO sessions (session_id, history)
VALUES (?, ?)
""", (session_id, json.dumps(history)))
conn.commit()
def clear_session(self, session_id: str):
with sqlite3.connect(self.db_path) as conn:
conn.execute("DELETE FROM sessions WHERE session_id = ?", (session_id,))
conn.commit()
memory_manager = MemoryManager()
|