File size: 2,921 Bytes
6c15f74 29477ed 6c15f74 29477ed 6c15f74 29477ed 6c15f74 29477ed 6c15f74 29477ed 6c15f74 |
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 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 |
import sqlite3
import json
import logging
import os
from datetime import datetime
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
class MemoryManager:
def __init__(self, db_path="session_memory.db"):
self.db_path = db_path
self._init_db()
def _init_db(self):
try:
with sqlite3.connect(self.db_path) as conn:
cursor = conn.cursor()
cursor.execute("""
CREATE TABLE IF NOT EXISTS messages (
id INTEGER PRIMARY KEY AUTOINCREMENT,
session_id TEXT NOT NULL,
role TEXT NOT NULL,
content TEXT NOT NULL,
timestamp DATETIME DEFAULT CURRENT_TIMESTAMP
)
""")
cursor.execute("CREATE INDEX IF NOT EXISTS idx_session_id ON messages(session_id)")
conn.commit()
except Exception as e:
logger.error(f"Failed to init DB: {e}")
def add_message(self, session_id, role, content):
"""Append a message to the session's history."""
try:
with sqlite3.connect(self.db_path) as conn:
cursor = conn.cursor()
cursor.execute(
"INSERT INTO messages (session_id, role, content) VALUES (?, ?, ?)",
(session_id, role, content)
)
conn.commit()
except Exception as e:
logger.error(f"Failed to add message: {e}")
def get_history(self, session_id, limit=10):
"""Retrieve the last `limit` messages."""
try:
with sqlite3.connect(self.db_path) as conn:
conn.row_factory = sqlite3.Row
cursor = conn.cursor()
# Get last N messages ordered by time ASC
cursor.execute(
"""
SELECT role, content FROM (
SELECT role, content, timestamp
FROM messages
WHERE session_id = ?
ORDER BY id DESC
LIMIT ?
) ORDER BY id ASC
""",
(session_id, limit)
)
rows = cursor.fetchall()
return [{"role": r["role"], "content": r["content"]} for r in rows]
except Exception as e:
logger.error(f"Failed to retrieve history: {e}")
return []
def clear_history(self, session_id):
try:
with sqlite3.connect(self.db_path) as conn:
cursor = conn.cursor()
cursor.execute("DELETE FROM messages WHERE session_id = ?", (session_id,))
conn.commit()
except Exception as e:
logger.error(f"Failed to clear history: {e}")
|