atlas/home-voice-agent/memory/storage.py

"""
Memory storage using SQLite.
"""

import sqlite3
import json
import logging
from pathlib import Path
from typing import Optional, List
from datetime import datetime
from uuid import uuid4

from memory.schema import MemoryEntry, MemoryCategory, MemorySource

logger = logging.getLogger(__name__)


class MemoryStorage:
    """SQLite storage for memory entries."""
    
    def __init__(self, db_path: Optional[Path] = None):
        """
        Initialize memory storage.
        
        Args:
            db_path: Path to SQLite database. If None, uses default.
        """
        if db_path is None:
            db_path = Path(__file__).parent.parent / "data" / "memory.db"
        
        self.db_path = db_path
        self.db_path.parent.mkdir(parents=True, exist_ok=True)
        self._init_db()
    
    def _init_db(self):
        """Initialize database schema."""
        conn = sqlite3.connect(str(self.db_path))
        cursor = conn.cursor()
        
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS memory (
                id TEXT PRIMARY KEY,
                category TEXT NOT NULL,
                key TEXT NOT NULL,
                value TEXT NOT NULL,
                confidence REAL DEFAULT 0.5,
                source TEXT NOT NULL,
                timestamp TEXT NOT NULL,
                last_accessed TEXT,
                access_count INTEGER DEFAULT 0,
                tags TEXT,
                context TEXT,
                UNIQUE(category, key)
            )
        """)
        
        # Create indexes
        cursor.execute("""
            CREATE INDEX IF NOT EXISTS idx_category_key 
            ON memory(category, key)
        """)
        
        cursor.execute("""
            CREATE INDEX IF NOT EXISTS idx_category 
            ON memory(category)
        """)
        
        cursor.execute("""
            CREATE INDEX IF NOT EXISTS idx_last_accessed 
            ON memory(last_accessed)
        """)
        
        conn.commit()
        conn.close()
        logger.info(f"Memory database initialized at {self.db_path}")
    
    def store(self, entry: MemoryEntry) -> bool:
        """
        Store a memory entry.
        
        Args:
            entry: Memory entry to store
            
        Returns:
            True if stored successfully
        """
        conn = sqlite3.connect(str(self.db_path))
        cursor = conn.cursor()
        
        try:
            cursor.execute("""
                INSERT OR REPLACE INTO memory 
                (id, category, key, value, confidence, source, timestamp, 
                 last_accessed, access_count, tags, context)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            """, (
                entry.id,
                entry.category.value,
                entry.key,
                entry.value,
                entry.confidence,
                entry.source.value,
                entry.timestamp.isoformat(),
                entry.last_accessed.isoformat() if entry.last_accessed else None,
                entry.access_count,
                json.dumps(entry.tags),
                entry.context
            ))
            
            conn.commit()
            logger.info(f"Stored memory: {entry.category.value}/{entry.key} = {entry.value}")
            return True
            
        except Exception as e:
            logger.error(f"Error storing memory: {e}")
            conn.rollback()
            return False
        finally:
            conn.close()
    
    def get(self, category: MemoryCategory, key: str) -> Optional[MemoryEntry]:
        """
        Get a memory entry by category and key.
        
        Args:
            category: Memory category
            key: Memory key
            
        Returns:
            Memory entry or None
        """
        conn = sqlite3.connect(str(self.db_path))
        conn.row_factory = sqlite3.Row
        cursor = conn.cursor()
        
        cursor.execute("""
            SELECT * FROM memory 
            WHERE category = ? AND key = ?
        """, (category.value, key))
        
        row = cursor.fetchone()
        conn.close()
        
        if row:
            # Update access
            self._update_access(row["id"])
            return self._row_to_entry(row)
        return None
    
    def get_by_category(self, category: MemoryCategory, limit: Optional[int] = None) -> List[MemoryEntry]:
        """
        Get all memory entries in a category.
        
        Args:
            category: Memory category
            limit: Maximum number of entries to return
            
        Returns:
            List of memory entries
        """
        conn = sqlite3.connect(str(self.db_path))
        conn.row_factory = sqlite3.Row
        cursor = conn.cursor()
        
        query = "SELECT * FROM memory WHERE category = ? ORDER BY last_accessed DESC"
        if limit:
            query += f" LIMIT {limit}"
        
        cursor.execute(query, (category.value,))
        rows = cursor.fetchall()
        conn.close()
        
        entries = [self._row_to_entry(row) for row in rows]
        
        # Update access for all
        for entry in entries:
            self._update_access(entry.id)
        
        return entries
    
    def search(self, query: str, category: Optional[MemoryCategory] = None, limit: int = 10) -> List[MemoryEntry]:
        """
        Search memory entries by value or context.
        
        Args:
            query: Search query
            category: Optional category filter
            limit: Maximum results
            
        Returns:
            List of matching memory entries
        """
        conn = sqlite3.connect(str(self.db_path))
        conn.row_factory = sqlite3.Row
        cursor = conn.cursor()
        
        search_term = f"%{query.lower()}%"
        
        if category:
            cursor.execute("""
                SELECT * FROM memory 
                WHERE category = ? 
                AND (LOWER(value) LIKE ? OR LOWER(context) LIKE ?)
                ORDER BY confidence DESC, last_accessed DESC
                LIMIT ?
            """, (category.value, search_term, search_term, limit))
        else:
            cursor.execute("""
                SELECT * FROM memory 
                WHERE LOWER(value) LIKE ? OR LOWER(context) LIKE ?
                ORDER BY confidence DESC, last_accessed DESC
                LIMIT ?
            """, (search_term, search_term, limit))
        
        rows = cursor.fetchall()
        conn.close()
        
        entries = [self._row_to_entry(row) for row in rows]
        
        # Update access
        for entry in entries:
            self._update_access(entry.id)
        
        return entries
    
    def _update_access(self, entry_id: str):
        """Update access timestamp and count."""
        conn = sqlite3.connect(str(self.db_path))
        cursor = conn.cursor()
        
        cursor.execute("""
            UPDATE memory 
            SET last_accessed = ?, access_count = access_count + 1
            WHERE id = ?
        """, (datetime.now().isoformat(), entry_id))
        
        conn.commit()
        conn.close()
    
    def _row_to_entry(self, row: sqlite3.Row) -> MemoryEntry:
        """Convert database row to MemoryEntry."""
        tags = json.loads(row["tags"]) if row["tags"] else []
        
        return MemoryEntry(
            id=row["id"],
            category=MemoryCategory(row["category"]),
            key=row["key"],
            value=row["value"],
            confidence=row["confidence"],
            source=MemorySource(row["source"]),
            timestamp=datetime.fromisoformat(row["timestamp"]),
            last_accessed=datetime.fromisoformat(row["last_accessed"]) if row["last_accessed"] else None,
            access_count=row["access_count"],
            tags=tags,
            context=row["context"]
        )
    
    def delete(self, entry_id: str) -> bool:
        """
        Delete a memory entry.
        
        Args:
            entry_id: Entry ID to delete
            
        Returns:
            True if deleted successfully
        """
        conn = sqlite3.connect(str(self.db_path))
        cursor = conn.cursor()
        
        try:
            cursor.execute("DELETE FROM memory WHERE id = ?", (entry_id,))
            conn.commit()
            logger.info(f"Deleted memory entry: {entry_id}")
            return True
        except Exception as e:
            logger.error(f"Error deleting memory: {e}")
            conn.rollback()
            return False
        finally:
            conn.close()
    
    def update_confidence(self, entry_id: str, confidence: float) -> bool:
        """
        Update confidence of a memory entry.
        
        Args:
            entry_id: Entry ID
            confidence: New confidence value (0.0-1.0)
            
        Returns:
            True if updated successfully
        """
        conn = sqlite3.connect(str(self.db_path))
        cursor = conn.cursor()
        
        try:
            cursor.execute("""
                UPDATE memory 
                SET confidence = ?
                WHERE id = ?
            """, (confidence, entry_id))
            
            conn.commit()
            return True
        except Exception as e:
            logger.error(f"Error updating confidence: {e}")
            conn.rollback()
            return False
        finally:
            conn.close()