models/persistent_memory.py

"""
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}")