atlas/home-voice-agent/conversation/session_manager.py

"""
Session Manager - Manages multi-turn conversations.

Handles session context, message history, and context window management.
"""

import sqlite3
import uuid
from datetime import datetime, timedelta
from pathlib import Path
from typing import Any, Dict, List, Optional
from dataclasses import dataclass, asdict
import json

# Database file location
DB_PATH = Path(__file__).parent.parent / "data" / "conversations.db"

# Context window settings
MAX_CONTEXT_MESSAGES = 20  # Keep last N messages in context
MAX_CONTEXT_TOKENS = 8000  # Approximate token limit (conservative)
SESSION_EXPIRY_HOURS = 24  # Sessions expire after 24 hours of inactivity


@dataclass
class Message:
    """Represents a single message in a conversation."""
    role: str  # "user", "assistant", "system"
    content: str
    timestamp: datetime
    tool_calls: Optional[List[Dict[str, Any]]] = None
    tool_results: Optional[List[Dict[str, Any]]] = None


@dataclass
class Session:
    """Represents a conversation session."""
    session_id: str
    agent_type: str  # "work" or "family"
    created_at: datetime
    last_activity: datetime
    messages: List[Message]
    summary: Optional[str] = None


class SessionManager:
    """Manages conversation sessions."""
    
    def __init__(self, db_path: Path = DB_PATH):
        """Initialize session manager with database."""
        self.db_path = db_path
        self.db_path.parent.mkdir(parents=True, exist_ok=True)
        self._init_db()
        self._active_sessions: Dict[str, Session] = {}
    
    def _init_db(self):
        """Initialize database schema."""
        conn = sqlite3.connect(str(self.db_path))
        cursor = conn.cursor()
        
        # Sessions table
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS sessions (
                session_id TEXT PRIMARY KEY,
                agent_type TEXT NOT NULL,
                created_at TEXT NOT NULL,
                last_activity TEXT NOT NULL,
                summary TEXT
            )
        """)
        
        # Messages table
        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 TEXT NOT NULL,
                tool_calls TEXT,
                tool_results TEXT,
                FOREIGN KEY (session_id) REFERENCES sessions(session_id)
            )
        """)
        
        conn.commit()
        conn.close()
    
    def create_session(self, agent_type: str) -> str:
        """Create a new conversation session."""
        session_id = str(uuid.uuid4())
        now = datetime.now()
        
        session = Session(
            session_id=session_id,
            agent_type=agent_type,
            created_at=now,
            last_activity=now,
            messages=[]
        )
        
        # Store in database
        conn = sqlite3.connect(str(self.db_path))
        cursor = conn.cursor()
        cursor.execute("""
            INSERT INTO sessions (session_id, agent_type, created_at, last_activity)
            VALUES (?, ?, ?, ?)
        """, (session_id, agent_type, now.isoformat(), now.isoformat()))
        conn.commit()
        conn.close()
        
        # Cache in memory
        self._active_sessions[session_id] = session
        
        return session_id
    
    def get_session(self, session_id: str) -> Optional[Session]:
        """Get session by ID, loading from DB if not in cache."""
        # Check cache first
        if session_id in self._active_sessions:
            session = self._active_sessions[session_id]
            # Check if expired
            if datetime.now() - session.last_activity > timedelta(hours=SESSION_EXPIRY_HOURS):
                self._active_sessions.pop(session_id)
                return None
            return session
        
        # Load from database
        conn = sqlite3.connect(str(self.db_path))
        conn.row_factory = sqlite3.Row
        cursor = conn.cursor()
        
        cursor.execute("""
            SELECT * FROM sessions WHERE session_id = ?
        """, (session_id,))
        session_row = cursor.fetchone()
        
        if not session_row:
            conn.close()
            return None
        
        # Load messages
        cursor.execute("""
            SELECT * FROM messages 
            WHERE session_id = ? 
            ORDER BY timestamp ASC
        """, (session_id,))
        message_rows = cursor.fetchall()
        
        conn.close()
        
        # Reconstruct session
        messages = []
        for row in message_rows:
            tool_calls = json.loads(row['tool_calls']) if row['tool_calls'] else None
            tool_results = json.loads(row['tool_results']) if row['tool_results'] else None
            messages.append(Message(
                role=row['role'],
                content=row['content'],
                timestamp=datetime.fromisoformat(row['timestamp']),
                tool_calls=tool_calls,
                tool_results=tool_results
            ))
        
        session = Session(
            session_id=session_row['session_id'],
            agent_type=session_row['agent_type'],
            created_at=datetime.fromisoformat(session_row['created_at']),
            last_activity=datetime.fromisoformat(session_row['last_activity']),
            messages=messages,
            summary=session_row['summary']
        )
        
        # Cache if not expired
        if datetime.now() - session.last_activity <= timedelta(hours=SESSION_EXPIRY_HOURS):
            self._active_sessions[session_id] = session
        
        return session
    
    def add_message(self, session_id: str, role: str, content: str, 
                   tool_calls: Optional[List[Dict[str, Any]]] = None,
                   tool_results: Optional[List[Dict[str, Any]]] = None):
        """Add a message to a session."""
        session = self.get_session(session_id)
        if not session:
            raise ValueError(f"Session not found: {session_id}")
        
        message = Message(
            role=role,
            content=content,
            timestamp=datetime.now(),
            tool_calls=tool_calls,
            tool_results=tool_results
        )
        
        session.messages.append(message)
        session.last_activity = datetime.now()
        
        # Store in database
        conn = sqlite3.connect(str(self.db_path))
        cursor = conn.cursor()
        cursor.execute("""
            INSERT INTO messages (session_id, role, content, timestamp, tool_calls, tool_results)
            VALUES (?, ?, ?, ?, ?, ?)
        """, (
            session_id,
            role,
            content,
            message.timestamp.isoformat(),
            json.dumps(tool_calls) if tool_calls else None,
            json.dumps(tool_results) if tool_results else None
        ))
        cursor.execute("""
            UPDATE sessions SET last_activity = ? WHERE session_id = ?
        """, (session.last_activity.isoformat(), session_id))
        conn.commit()
        conn.close()
    
    def get_context_messages(self, session_id: str, max_messages: int = MAX_CONTEXT_MESSAGES) -> List[Dict[str, Any]]:
        """
        Get messages for LLM context, keeping only recent messages.
        
        Returns messages in OpenAI chat format.
        """
        session = self.get_session(session_id)
        if not session:
            return []
        
        # Get recent messages
        recent_messages = session.messages[-max_messages:]
        
        # Convert to OpenAI format
        context = []
        for msg in recent_messages:
            message_dict = {
                "role": msg.role,
                "content": msg.content
            }
            
            # Add tool calls if present
            if msg.tool_calls:
                message_dict["tool_calls"] = msg.tool_calls
            
            # Add tool results if present
            if msg.tool_results:
                message_dict["tool_results"] = msg.tool_results
            
            context.append(message_dict)
        
        return context
    
    def summarize_old_messages(self, session_id: str, keep_recent: int = 10):
        """
        Summarize old messages to reduce context size.
        
        This is a placeholder - actual summarization would use an LLM.
        """
        session = self.get_session(session_id)
        if not session or len(session.messages) <= keep_recent:
            return
        
        # For now, just keep recent messages
        # TODO: Implement actual summarization using LLM
        old_messages = session.messages[:-keep_recent]
        recent_messages = session.messages[-keep_recent:]
        
        # Create summary placeholder
        summary = f"Previous conversation had {len(old_messages)} messages. Key topics discussed."
        
        # Update session
        session.messages = recent_messages
        session.summary = summary
        
        # Update database
        conn = sqlite3.connect(str(self.db_path))
        cursor = conn.cursor()
        cursor.execute("""
            UPDATE sessions SET summary = ? WHERE session_id = ?
        """, (summary, session_id))
        
        # Delete old messages
        cursor.execute("""
            DELETE FROM messages 
            WHERE session_id = ? AND timestamp < ?
        """, (session_id, recent_messages[0].timestamp.isoformat()))
        
        conn.commit()
        conn.close()
    
    def delete_session(self, session_id: str):
        """Delete a session and all its messages."""
        conn = sqlite3.connect(str(self.db_path))
        cursor = conn.cursor()
        cursor.execute("DELETE FROM messages WHERE session_id = ?", (session_id,))
        cursor.execute("DELETE FROM sessions WHERE session_id = ?", (session_id,))
        conn.commit()
        conn.close()
        
        # Remove from cache
        self._active_sessions.pop(session_id, None)
    
    def cleanup_expired_sessions(self):
        """Remove expired sessions."""
        expiry_time = datetime.now() - timedelta(hours=SESSION_EXPIRY_HOURS)
        
        conn = sqlite3.connect(str(self.db_path))
        cursor = conn.cursor()
        
        # Find expired sessions
        cursor.execute("""
            SELECT session_id FROM sessions 
            WHERE last_activity < ?
        """, (expiry_time.isoformat(),))
        
        expired_sessions = [row[0] for row in cursor.fetchall()]
        
        # Delete expired sessions
        for session_id in expired_sessions:
            cursor.execute("DELETE FROM messages WHERE session_id = ?", (session_id,))
            cursor.execute("DELETE FROM sessions WHERE session_id = ?", (session_id,))
            self._active_sessions.pop(session_id, None)
        
        conn.commit()
        conn.close()


# Global session manager instance
_session_manager = SessionManager()


def get_session_manager() -> SessionManager:
    """Get the global session manager instance."""
    return _session_manager