session_manager.py

"""
In-memory session management for chatbot conversations
"""
import uuid
from datetime import datetime, timedelta
from typing import Dict, List, Optional
from threading import Lock
from models import ChatMessage, MessageRole


class Session:
    def __init__(self, session_id: str):
        self.session_id = session_id
        self.messages: List[ChatMessage] = []
        self.created_at = datetime.utcnow()
        self.last_accessed = datetime.utcnow()
    
    def add_message(self, role: MessageRole, content: str) -> None:
        """Add a message to the session"""
        message = ChatMessage(role=role, content=content)
        self.messages.append(message)
        self.last_accessed = datetime.utcnow()
    
    def get_messages(self) -> List[ChatMessage]:
        """Get all messages in the session"""
        self.last_accessed = datetime.utcnow()
        return self.messages.copy()
    
    def is_expired(self, timeout_hours: int = 1) -> bool:
        """Check if session is expired based on last access time"""
        timeout_delta = timedelta(hours=timeout_hours)
        return datetime.utcnow() - self.last_accessed > timeout_delta


class SessionManager:
    def __init__(self, session_timeout_hours: int = 1):
        self.sessions: Dict[str, Session] = {}
        self.session_timeout_hours = session_timeout_hours
        self._lock = Lock()  # Thread safety for concurrent access
    
    def create_session(self) -> str:
        """Create a new session and return its ID"""
        session_id = str(uuid.uuid4())
        
        with self._lock:
            session = Session(session_id)
            self.sessions[session_id] = session
            
        return session_id
    
    def get_session(self, session_id: str) -> Optional[Session]:
        """Get a session by ID, return None if not found or expired"""
        with self._lock:
            session = self.sessions.get(session_id)
            
            if session is None:
                return None
            
            if session.is_expired(self.session_timeout_hours):
                # Remove expired session
                del self.sessions[session_id]
                return None
            
            return session
    
    def add_message_to_session(self, session_id: str, role: MessageRole, content: str) -> bool:
        """Add a message to a session, return True if successful"""
        session = self.get_session(session_id)
        if session is None:
            return False
        
        session.add_message(role, content)
        return True
    
    def get_session_messages(self, session_id: str) -> Optional[List[ChatMessage]]:
        """Get all messages from a session"""
        session = self.get_session(session_id)
        if session is None:
            return None
        
        return session.get_messages()
    
    def cleanup_expired_sessions(self) -> int:
        """Remove expired sessions and return count of removed sessions"""
        expired_sessions = []
        
        with self._lock:
            for session_id, session in self.sessions.items():
                if session.is_expired(self.session_timeout_hours):
                    expired_sessions.append(session_id)
            
            for session_id in expired_sessions:
                del self.sessions[session_id]
        
        return len(expired_sessions)
    
    def get_session_count(self) -> int:
        """Get current number of active sessions"""
        with self._lock:
            return len(self.sessions)
    
    def session_exists(self, session_id: str) -> bool:
        """Check if a session exists and is not expired"""
        return self.get_session(session_id) is not None