Spaces:
Sleeping
Sleeping
File size: 6,429 Bytes
a2854ac | 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 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 | """
persistent_memory.py
~~~~~~~~~~~~~~~~~~~~
Production-ready persistent memory layer utilizing SQLite.
Stores user preferences, long-term summaries, and historical session logs.
Usage
-----
from models.persistent_memory import PersistentMemory
db = PersistentMemory()
db.save_preference("user_name", "Logesh")
db.save_summary("session_123", "User discussed deploying Qwen model on HF Spaces CPU.")
"""
from __future__ import annotations
import os
import sqlite3
from typing import Dict, Optional, List, Tuple
from models.logger_config import logger
DB_PATH = "logs/assistant_memory.db"
class PersistentMemory:
"""
SQLite-backed long-term storage for assistant state, profile preferences,
and semantic conversation summaries.
"""
def __init__(self, db_path: str = DB_PATH) -> None:
self.db_path = db_path
os.makedirs(os.path.dirname(self.db_path), exist_ok=True)
self._init_db()
def _get_connection(self) -> sqlite3.Connection:
conn = sqlite3.connect(self.db_path)
conn.row_factory = sqlite3.Row
return conn
def _init_db(self) -> None:
"""Create database tables if they do not exist."""
try:
with self._get_connection() as conn:
cursor = conn.cursor()
# 1. Preferences Table (Key-Value)
cursor.execute("""
CREATE TABLE IF NOT EXISTS preferences (
pref_key TEXT PRIMARY KEY,
pref_value TEXT,
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
)
""")
# 2. Conversation Summaries Table
cursor.execute("""
CREATE TABLE IF NOT EXISTS summaries (
session_id TEXT PRIMARY KEY,
summary TEXT,
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
)
""")
conn.commit()
logger.info(f"Initialized SQLite persistent memory database at {self.db_path}")
except Exception as e:
logger.error(f"Failed to initialize SQLite persistent memory: {e}")
# ── User Preference Storage ───────────────────────────────────────────────
def save_preference(self, key: str, value: str) -> None:
"""Store or update a user preference (e.g. name, custom theme, model preference)."""
try:
with self._get_connection() as conn:
cursor = conn.cursor()
cursor.execute("""
INSERT INTO preferences (pref_key, pref_value, updated_at)
VALUES (?, ?, CURRENT_TIMESTAMP)
ON CONFLICT(pref_key) DO UPDATE SET
pref_value=excluded.pref_value,
updated_at=CURRENT_TIMESTAMP
""", (key, value))
conn.commit()
logger.debug(f"Saved preference: {key} -> {value}")
except Exception as e:
logger.error(f"Failed to save preference '{key}': {e}")
def get_preference(self, key: str, default: Optional[str] = None) -> Optional[str]:
"""Retrieve a stored preference by key."""
try:
with self._get_connection() as conn:
cursor = conn.cursor()
cursor.execute("SELECT pref_value FROM preferences WHERE pref_key = ?", (key,))
row = cursor.fetchone()
if row:
return row["pref_value"]
except Exception as e:
logger.error(f"Failed to retrieve preference '{key}': {e}")
return default
def list_preferences(self) -> Dict[str, str]:
"""List all stored preferences."""
prefs = {}
try:
with self._get_connection() as conn:
cursor = conn.cursor()
cursor.execute("SELECT pref_key, pref_value FROM preferences")
for row in cursor.fetchall():
prefs[row["pref_key"]] = row["pref_value"]
except Exception as e:
logger.error(f"Failed to list preferences: {e}")
return prefs
# ── Session Summary Storage ────────────────────────────────────────────────
def save_summary(self, session_id: str, summary: str) -> None:
"""Save a cumulative summary of a specific chat session."""
try:
with self._get_connection() as conn:
cursor = conn.cursor()
cursor.execute("""
INSERT INTO summaries (session_id, summary, updated_at)
VALUES (?, ?, CURRENT_TIMESTAMP)
ON CONFLICT(session_id) DO UPDATE SET
summary=excluded.summary,
updated_at=CURRENT_TIMESTAMP
""", (session_id, summary))
conn.commit()
logger.debug(f"Saved conversation summary for session {session_id}")
except Exception as e:
logger.error(f"Failed to save session summary for '{session_id}': {e}")
def get_summary(self, session_id: str) -> Optional[str]:
"""Get the summary for a session."""
try:
with self._get_connection() as conn:
cursor = conn.cursor()
cursor.execute("SELECT summary FROM summaries WHERE session_id = ?", (session_id,))
row = cursor.fetchone()
if row:
return row["summary"]
except Exception as e:
logger.error(f"Failed to retrieve session summary for '{session_id}': {e}")
return None
def clear_all(self) -> None:
"""Clear all stored database records."""
try:
with self._get_connection() as conn:
cursor = conn.cursor()
cursor.execute("DELETE FROM preferences")
cursor.execute("DELETE FROM summaries")
conn.commit()
logger.info("Cleared all persistent memory records.")
except Exception as e:
logger.error(f"Failed to clear persistent memory: {e}")
|