atlas/home-voice-agent/conversation/summarization/retention.py

"""
Conversation retention and deletion policies.
"""

import logging
from pathlib import Path
from typing import Optional, List
from datetime import datetime, timedelta
import sqlite3

logger = logging.getLogger(__name__)


class RetentionPolicy:
    """Defines retention policies for conversations."""
    
    def __init__(self,
                 max_age_days: int = 90,
                 max_sessions: int = 1000,
                 auto_delete: bool = False):
        """
        Initialize retention policy.
        
        Args:
            max_age_days: Maximum age in days before deletion
            max_sessions: Maximum number of sessions to keep
            auto_delete: Whether to auto-delete old sessions
        """
        self.max_age_days = max_age_days
        self.max_sessions = max_sessions
        self.auto_delete = auto_delete
    
    def should_delete(self, session_timestamp: datetime) -> bool:
        """
        Check if session should be deleted based on age.
        
        Args:
            session_timestamp: When session was created
            
        Returns:
            True if should be deleted
        """
        age = datetime.now() - session_timestamp
        return age.days > self.max_age_days


class ConversationRetention:
    """Manages conversation retention and deletion."""
    
    def __init__(self, db_path: Optional[Path] = None, policy: Optional[RetentionPolicy] = None):
        """
        Initialize retention manager.
        
        Args:
            db_path: Path to conversations database
            policy: Retention policy
        """
        if db_path is None:
            db_path = Path(__file__).parent.parent.parent / "data" / "conversations.db"
        
        self.db_path = db_path
        self.policy = policy or RetentionPolicy()
    
    def list_old_sessions(self) -> List[tuple]:
        """
        List sessions that should be deleted.
        
        Returns:
            List of (session_id, created_at) tuples
        """
        if not self.db_path.exists():
            return []
        
        conn = sqlite3.connect(str(self.db_path))
        conn.row_factory = sqlite3.Row
        cursor = conn.cursor()
        
        cutoff_date = datetime.now() - timedelta(days=self.policy.max_age_days)
        
        cursor.execute("""
            SELECT session_id, created_at 
            FROM sessions 
            WHERE created_at < ?
            ORDER BY created_at ASC
        """, (cutoff_date.isoformat(),))
        
        rows = cursor.fetchall()
        conn.close()
        
        return [(row["session_id"], row["created_at"]) for row in rows]
    
    def delete_session(self, session_id: str) -> bool:
        """
        Delete a session.
        
        Args:
            session_id: Session ID to delete
            
        Returns:
            True if deleted successfully
        """
        if not self.db_path.exists():
            return False
        
        conn = sqlite3.connect(str(self.db_path))
        cursor = conn.cursor()
        
        try:
            # Delete session
            cursor.execute("DELETE FROM sessions WHERE session_id = ?", (session_id,))
            
            # Delete messages
            cursor.execute("DELETE FROM messages WHERE session_id = ?", (session_id,))
            
            conn.commit()
            logger.info(f"Deleted session: {session_id}")
            return True
            
        except Exception as e:
            logger.error(f"Error deleting session {session_id}: {e}")
            conn.rollback()
            return False
        finally:
            conn.close()
    
    def cleanup_old_sessions(self) -> int:
        """
        Clean up old sessions based on policy.
        
        Returns:
            Number of sessions deleted
        """
        if not self.policy.auto_delete:
            return 0
        
        old_sessions = self.list_old_sessions()
        deleted_count = 0
        
        for session_id, _ in old_sessions:
            if self.delete_session(session_id):
                deleted_count += 1
        
        logger.info(f"Cleaned up {deleted_count} old sessions")
        return deleted_count
    
    def get_session_count(self) -> int:
        """
        Get total number of sessions.
        
        Returns:
            Number of sessions
        """
        if not self.db_path.exists():
            return 0
        
        conn = sqlite3.connect(str(self.db_path))
        cursor = conn.cursor()
        
        cursor.execute("SELECT COUNT(*) FROM sessions")
        count = cursor.fetchone()[0]
        conn.close()
        
        return count
    
    def enforce_max_sessions(self) -> int:
        """
        Enforce maximum session limit by deleting oldest sessions.
        
        Returns:
            Number of sessions deleted
        """
        current_count = self.get_session_count()
        
        if current_count <= self.policy.max_sessions:
            return 0
        
        # Get oldest sessions to delete
        conn = sqlite3.connect(str(self.db_path))
        conn.row_factory = sqlite3.Row
        cursor = conn.cursor()
        
        cursor.execute("""
            SELECT session_id 
            FROM sessions 
            ORDER BY created_at ASC 
            LIMIT ?
        """, (current_count - self.policy.max_sessions,))
        
        rows = cursor.fetchall()
        conn.close()
        
        deleted_count = 0
        for row in rows:
            if self.delete_session(row["session_id"]):
                deleted_count += 1
        
        logger.info(f"Enforced max sessions: deleted {deleted_count} sessions")
        return deleted_count


# Global retention manager
_retention = ConversationRetention()


def get_retention_manager() -> ConversationRetention:
    """Get the global retention manager instance."""
    return _retention