Update conversation_manager.py

conversation_manager.py

@@ -2,6 +2,7 @@
 
 import os
 import uuid
+import hashlib
 import logging
 from typing import Dict, List, Any, Optional
 from datetime import datetime
@@ -16,7 +17,7 @@ load_dotenv()
 
 
 class ConversationManager:
-    """Enhanced ConversationManager with unique message IDs and flexible chat ID management"""
+    """Enhanced ConversationManager with user authentication and session isolation"""
     
     def __init__(self):
         """Initialize the conversation manager with Supabase client"""
@@ -35,15 +36,38 @@ class ConversationManager:
             logger.error(f"Failed to initialize ConversationManager: {e}")
             raise
     
-    def generate_chat_id(self) -> str:
-        """Generate a unique chat ID"""
-        return f"chat_{uuid.uuid4().hex[:12]}"
+    def generate_chat_id(self, user_id: str) -> str:
+        """Generate a unique chat ID for a specific user"""
+        timestamp = str(int(datetime.utcnow().timestamp()))
+        random_part = uuid.uuid4().hex[:8]
+        return f"chat_{user_id[:8]}_{timestamp}_{random_part}"
     
-    def generate_message_id(self, chat_id: str) -> int:
+    def generate_user_id(self, device_info: Optional[str] = None, ip_address: Optional[str] = None) -> str:
+        """
+        Generate a unique user ID based on device info and/or IP
+        This creates a consistent ID for the same device/IP combination
+        """
+        # Create a hash based on available info
+        identifier_string = ""
+        
+        if device_info:
+            identifier_string += device_info
+        if ip_address:
+            identifier_string += ip_address
+        
+        # Fallback to timestamp + random if no device info
+        if not identifier_string:
+            identifier_string = f"anonymous_{int(datetime.utcnow().timestamp())}_{uuid.uuid4().hex[:8]}"
+        
+        # Create consistent hash
+        user_hash = hashlib.sha256(identifier_string.encode()).hexdigest()[:16]
+        return f"user_{user_hash}"
+    
+    def generate_message_id(self, chat_id: str, user_id: str) -> int:
         """Generate the next message ID for a specific chat"""
         try:
             # Get current conversation to find the highest message ID
-            history = self.get_history(chat_id)
+            history = self.get_history(chat_id, user_id)
             
             if not history:
                 return 1
@@ -58,53 +82,57 @@ class ConversationManager:
             return max_id + 1
             
         except Exception as e:
-            logger.error(f"Error generating message ID for chat {chat_id}: {e}")
+            logger.error(f"Error generating message ID for chat {chat_id}, user {user_id}: {e}")
             return 1
     
-    def get_history(self, chat_id: str) -> List[Dict[str, Any]]:
+    def get_history(self, chat_id: str, user_id: str) -> List[Dict[str, Any]]:
         """
-        Retrieve conversation history for a specific chat ID
+        Retrieve conversation history for a specific chat ID and user
         
         Args:
             chat_id: The chat ID to retrieve history for
+            user_id: The user ID (for security/isolation)
             
         Returns:
-            List of message dictionaries, empty list if chat doesn't exist
+            List of message dictionaries, empty list if chat doesn't exist or unauthorized
         """
-        if not chat_id:
+        if not chat_id or not user_id:
+            logger.warning("Cannot get history: chat_id or user_id is empty")
             return []
         
         try:
-            response = self.supabase.table('conversations').select('history').eq('session_id', chat_id).execute()
+            response = self.supabase.table('conversations').select('history').eq('session_id', chat_id).eq('user_id', user_id).execute()
             
             if response.data and len(response.data) > 0:
                 history = response.data[0].get('history', [])
                 if isinstance(history, list):
+                    logger.info(f"Retrieved {len(history)} messages for chat {chat_id}, user {user_id}")
                     return history
                 else:
                     logger.warning(f"Invalid history format for chat {chat_id}")
                     return []
             else:
-                logger.info(f"No history found for chat {chat_id}")
+                logger.info(f"No history found for chat {chat_id}, user {user_id}")
                 return []
                 
         except Exception as e:
-            logger.error(f"Error retrieving history for chat {chat_id}: {e}")
+            logger.error(f"Error retrieving history for chat {chat_id}, user {user_id}: {e}")
             return []
     
-    def save_history(self, chat_id: str, history: List[Dict[str, Any]], user_id: str) -> bool:
+    def save_history(self, chat_id: str, user_id: str, history: List[Dict[str, Any]]) -> bool:
         """
-        Save conversation history for a specific chat ID
+        Save conversation history for a specific chat ID and user
         
         Args:
             chat_id: The chat ID to save history for
+            user_id: The user ID (for security/isolation)
             history: List of message dictionaries
             
         Returns:
             True if successful, False otherwise
         """
-        if not chat_id:
-            logger.error("Cannot save history: chat_id is empty")
+        if not chat_id or not user_id:
+            logger.error("Cannot save history: chat_id or user_id is empty")
             return False
         
         if not isinstance(history, list):
@@ -112,40 +140,44 @@ class ConversationManager:
             return False
         
         try:
-            # Check if conversation exists
+            # Check if conversation exists for this user
             existing = self.supabase.table('conversations').select('session_id').eq('session_id', chat_id).eq('user_id', user_id).execute()
-
+            
+            current_time = datetime.utcnow().isoformat()
+            
             if existing.data and len(existing.data) > 0:
                 # Update existing conversation
                 response = self.supabase.table('conversations').update({
                     'history': history,
-                    'updated_at': datetime.utcnow().isoformat()
-                }).eq('session_id', chat_id).execute()
+                    'updated_at': current_time
+                }).eq('session_id', chat_id).eq('user_id', user_id).execute()
                 
-                logger.info(f"Updated history for chat {chat_id}")
+                logger.info(f"Updated history for chat {chat_id}, user {user_id}")
             else:
                 # Create new conversation
                 response = self.supabase.table('conversations').insert({
                     'session_id': chat_id,
+                    'user_id': user_id,
                     'history': history,
-                    'created_at': datetime.utcnow().isoformat(),
-                    'updated_at': datetime.utcnow().isoformat()
+                    'created_at': current_time,
+                    'updated_at': current_time
                 }).execute()
                 
-                logger.info(f"Created new conversation for chat {chat_id}")
+                logger.info(f"Created new conversation for chat {chat_id}, user {user_id}")
             
             return True
             
         except Exception as e:
-            logger.error(f"Error saving history for chat {chat_id}: {e}")
+            logger.error(f"Error saving history for chat {chat_id}, user {user_id}: {e}")
             return False
     
-    def add_message(self, chat_id: str, role: str, content: str, image_url: Optional[str] = None) -> Dict[str, Any]:
+    def add_message(self, chat_id: str, user_id: str, role: str, content: str, image_url: Optional[str] = None) -> Dict[str, Any]:
         """
         Add a single message to a chat and return the message with its ID
         
         Args:
             chat_id: The chat ID to add the message to
+            user_id: The user ID (for security/isolation)
             role: Message role ('user' or 'assistant')
             content: Message content
             image_url: Optional image URL
@@ -153,13 +185,13 @@ class ConversationManager:
         Returns:
             Dictionary containing the message with its assigned ID
         """
-        if not chat_id:
-            logger.error("Cannot add message: chat_id is empty")
+        if not chat_id or not user_id:
+            logger.error("Cannot add message: chat_id or user_id is empty")
             return {}
         
         try:
             # Generate message ID
-            message_id = self.generate_message_id(chat_id)
+            message_id = self.generate_message_id(chat_id, user_id)
             
             # Create message object
             message = {
@@ -174,59 +206,68 @@ class ConversationManager:
                 message["imageUrl"] = image_url
             
             # Get current history and add the new message
-            current_history = self.get_history(chat_id)
+            current_history = self.get_history(chat_id, user_id)
             current_history.append(message)
             
             # Save updated history
-            if self.save_history(chat_id, current_history):
-                logger.info(f"Added message {message_id} to chat {chat_id}")
+            if self.save_history(chat_id, user_id, current_history):
+                logger.info(f"Added message {message_id} to chat {chat_id}, user {user_id}")
                 return message
             else:
-                logger.error(f"Failed to save message to chat {chat_id}")
+                logger.error(f"Failed to save message to chat {chat_id}, user {user_id}")
                 return {}
                 
         except Exception as e:
-            logger.error(f"Error adding message to chat {chat_id}: {e}")
+            logger.error(f"Error adding message to chat {chat_id}, user {user_id}: {e}")
             return {}
     
-    def delete_history(self, chat_id: str) -> bool:
+    def delete_history(self, chat_id: str, user_id: str) -> bool:
         """
-        Delete conversation history for a specific chat ID
+        Delete conversation history for a specific chat ID and user
         
         Args:
             chat_id: The chat ID to delete
+            user_id: The user ID (for security/isolation)
             
         Returns:
             True if successful, False otherwise
         """
-        if not chat_id:
-            logger.error("Cannot delete history: chat_id is empty")
+        if not chat_id or not user_id:
+            logger.error("Cannot delete history: chat_id or user_id is empty")
             return False
         
         try:
-            response = self.supabase.table('conversations').delete().eq('session_id', chat_id).execute()
+            response = self.supabase.table('conversations').delete().eq('session_id', chat_id).eq('user_id', user_id).execute()
             
             # Check if any rows were affected
             if hasattr(response, 'data') and response.data is not None:
-                logger.info(f"Deleted conversation history for chat {chat_id}")
+                logger.info(f"Deleted conversation history for chat {chat_id}, user {user_id}")
                 return True
             else:
-                logger.warning(f"No conversation found to delete for chat {chat_id}")
+                logger.warning(f"No conversation found to delete for chat {chat_id}, user {user_id}")
                 return True  # Consider it successful if nothing to delete
                 
         except Exception as e:
-            logger.error(f"Error deleting history for chat {chat_id}: {e}")
+            logger.error(f"Error deleting history for chat {chat_id}, user {user_id}: {e}")
             return False
     
     def get_all_chat_sessions(self, user_id: str) -> List[Dict[str, Any]]:
         """
-        Get all chat sessions with basic information
+        Get all chat sessions for a specific user
         
+        Args:
+            user_id: The user ID to get sessions for
+            
         Returns:
             List of dictionaries with chat session information
         """
+        if not user_id:
+            logger.error("Cannot get chat sessions: user_id is empty")
+            return []
+        
         try:
-            response = self.supabase.table('conversations').select('session_id, history, created_at, updated_at').eq('user_id', user_id).execute()    
+            response = self.supabase.table('conversations').select('session_id, history, created_at, updated_at').eq('user_id', user_id).execute()
+            
             sessions = []
             for conv in response.data:
                 session_id = conv.get('session_id')
@@ -234,7 +275,7 @@ class ConversationManager:
                 created_at = conv.get('created_at')
                 updated_at = conv.get('updated_at')
                 
-                # Generate title from first user message
+                # Generate title from first user message in the history
                 title = "New Chat"
                 if history and len(history) > 0:
                     first_user_msg = None
@@ -249,7 +290,7 @@ class ConversationManager:
                             title = content[:35] + '...' if len(content) > 35 else content
                         elif first_user_msg.get('imageUrl'):
                             title = "Image Query"
-                
+
                 sessions.append({
                     "session_id": session_id,
                     "title": title,
@@ -261,27 +302,29 @@ class ConversationManager:
             # Sort by updated_at (most recent first)
             sessions.sort(key=lambda x: x.get('updated_at', ''), reverse=True)
             
+            logger.info(f"Retrieved {len(sessions)} sessions for user {user_id}")
             return sessions
             
         except Exception as e:
-            logger.error(f"Error fetching chat sessions: {e}")
+            logger.error(f"Error fetching chat sessions for user {user_id}: {e}")
             return []
     
-    def get_chat_info(self, chat_id: str) -> Dict[str, Any]:
+    def get_chat_info(self, chat_id: str, user_id: str) -> Dict[str, Any]:
         """
-        Get information about a specific chat
+        Get information about a specific chat for a specific user
         
         Args:
             chat_id: The chat ID to get information for
+            user_id: The user ID (for security/isolation)
             
         Returns:
             Dictionary with chat information
         """
-        if not chat_id:
+        if not chat_id or not user_id:
             return {}
         
         try:
-            response = self.supabase.table('conversations').select('*').eq('session_id', chat_id).execute()
+            response = self.supabase.table('conversations').select('*').eq('session_id', chat_id).eq('user_id', user_id).execute()
             
             if response.data and len(response.data) > 0:
                 conv = response.data[0]
@@ -289,6 +332,7 @@ class ConversationManager:
                 
                 return {
                     "session_id": conv.get('session_id'),
+                    "user_id": conv.get('user_id'),
                     "message_count": len(history),
                     "created_at": conv.get('created_at'),
                     "updated_at": conv.get('updated_at'),
@@ -297,82 +341,205 @@ class ConversationManager:
             else:
                 return {
                     "session_id": chat_id,
+                    "user_id": user_id,
                     "message_count": 0,
                     "exists": False
                 }
                 
         except Exception as e:
-            logger.error(f"Error getting chat info for {chat_id}: {e}")
+            logger.error(f"Error getting chat info for {chat_id}, user {user_id}: {e}")
             return {
                 "session_id": chat_id,
+                "user_id": user_id,
                 "error": str(e),
                 "exists": False
             }
     
-    def chat_exists(self, chat_id: str) -> bool:
+    def chat_exists(self, chat_id: str, user_id: str) -> bool:
         """
-        Check if a chat ID exists in the database
+        Check if a chat ID exists for a specific user
         
         Args:
             chat_id: The chat ID to check
+            user_id: The user ID (for security/isolation)
             
         Returns:
-            True if chat exists, False otherwise
+            True if chat exists and belongs to user, False otherwise
         """
-        if not chat_id:
+        if not chat_id or not user_id:
             return False
         
         try:
-            response = self.supabase.table('conversations').select('session_id').eq('session_id', chat_id).execute()
-            return bool(response.data and len(response.data) > 0)
+            response = self.supabase.table('conversations').select('session_id').eq('session_id', chat_id).eq('user_id', user_id).execute()
+            exists = bool(response.data and len(response.data) > 0)
+            
+            if exists:
+                logger.info(f"Chat {chat_id} exists for user {user_id}")
+            
+            return exists
             
         except Exception as e:
-            logger.error(f"Error checking if chat {chat_id} exists: {e}")
+            logger.error(f"Error checking if chat {chat_id} exists for user {user_id}: {e}")
+            return False
+    
+    def get_user_stats(self, user_id: str) -> Dict[str, Any]:
+        """
+        Get statistics for a specific user
+        
+        Args:
+            user_id: The user ID to get stats for
+            
+        Returns:
+            Dictionary with user statistics
+        """
+        if not user_id:
+            return {}
+        
+        try:
+            sessions = self.get_all_chat_sessions(user_id)
+            total_sessions = len(sessions)
+            total_messages = sum(session.get('message_count', 0) for session in sessions)
+            
+            # Calculate recent activity (last 24 hours)
+            from datetime import datetime, timedelta
+            now = datetime.utcnow()
+            yesterday = now - timedelta(days=1)
+            
+            recent_sessions = 0
+            for session in sessions:
+                updated_at = session.get('updated_at')
+                if updated_at:
+                    try:
+                        session_time = datetime.fromisoformat(updated_at.replace('Z', '+00:00'))
+                        if session_time > yesterday:
+                            recent_sessions += 1
+                    except:
+                        pass
+            
+            return {
+                "user_id": user_id,
+                "total_sessions": total_sessions,
+                "total_messages": total_messages,
+                "recent_sessions_24h": recent_sessions,
+                "average_messages_per_session": round(total_messages / total_sessions, 2) if total_sessions > 0 else 0,
+                "first_session": min(session.get('created_at', '') for session in sessions) if sessions else None,
+                "last_activity": max(session.get('updated_at', '') for session in sessions) if sessions else None
+            }
+            
+        except Exception as e:
+            logger.error(f"Error getting user stats for {user_id}: {e}")
+            return {"user_id": user_id, "error": str(e)}
+    
+    def delete_all_user_data(self, user_id: str) -> bool:
+        """
+        Delete all conversation data for a specific user
+        
+        Args:
+            user_id: The user ID to delete all data for
+            
+        Returns:
+            True if successful, False otherwise
+        """
+        if not user_id:
+            logger.error("Cannot delete user data: user_id is empty")
+            return False
+        
+        try:
+            response = self.supabase.table('conversations').delete().eq('user_id', user_id).execute()
+            
+            logger.info(f"Deleted all conversation data for user {user_id}")
+            return True
+            
+        except Exception as e:
+            logger.error(f"Error deleting all user data for {user_id}: {e}")
+            return False
+    
+    def migrate_anonymous_sessions(self, old_session_ids: List[str], new_user_id: str) -> bool:
+        """
+        Migrate anonymous sessions to a specific user ID
+        Useful when user decides to create an account
+        
+        Args:
+            old_session_ids: List of session IDs to migrate
+            new_user_id: The new user ID to assign these sessions to
+            
+        Returns:
+            True if successful, False otherwise
+        """
+        if not old_session_ids or not new_user_id:
+            logger.error("Cannot migrate sessions: missing session_ids or user_id")
+            return False
+        
+        try:
+            migrated_count = 0
+            
+            for session_id in old_session_ids:
+                try:
+                    # Update the user_id for this session
+                    response = self.supabase.table('conversations').update({
+                        'user_id': new_user_id,
+                        'updated_at': datetime.utcnow().isoformat()
+                    }).eq('session_id', session_id).execute()
+                    
+                    migrated_count += 1
+                    
+                except Exception as e:
+                    logger.error(f"Failed to migrate session {session_id}: {e}")
+                    continue
+            
+            logger.info(f"Migrated {migrated_count} sessions to user {new_user_id}")
+            return migrated_count > 0
+            
+        except Exception as e:
+            logger.error(f"Error migrating sessions to user {new_user_id}: {e}")
             return False
 
 
 # Example usage and testing
 if __name__ == "__main__":
-    print("=== ConversationManager Test ===")
+    print("=== ConversationManager with User Authentication Test ===")
     
     try:
         # Initialize manager
         manager = ConversationManager()
         
-        # Test 1: Create a new chat
-        test_chat_id = "test_chat_123"
-        print(f"\n1. Testing with chat ID: {test_chat_id}")
+        # Test user ID generation
+        test_user_id = manager.generate_user_id("test_device_info", "192.168.1.100")
+        print(f"\n1. Generated user ID: {test_user_id}")
         
-        # Add some messages
-        msg1 = manager.add_message(test_chat_id, "user", "Hello, I need help with farming")
-        print(f"Added user message: {msg1}")
+        # Create a new chat for this user
+        test_chat_id = manager.generate_chat_id(test_user_id)
+        print(f"2. Generated chat ID: {test_chat_id}")
         
-        msg2 = manager.add_message(test_chat_id, "assistant", "Hello! I'd be happy to help you with farming. What specific questions do you have?")
-        print(f"Added assistant message: {msg2}")
+        # Add some messages
+        msg1 = manager.add_message(test_chat_id, test_user_id, "user", "Hello, I need help with farming")
+        print(f"3. Added user message: {msg1}")
         
-        msg3 = manager.add_message(test_chat_id, "user", "What crops should I grow?", "https://example.com/image.jpg")
-        print(f"Added user message with image: {msg3}")
+        msg2 = manager.add_message(test_chat_id, test_user_id, "assistant", "Hello! I'd be happy to help you with farming. What specific questions do you have?")
+        print(f"4. Added assistant message: {msg2}")
         
-        # Test 2: Retrieve history
-        print(f"\n2. Retrieved history for {test_chat_id}:")
-        history = manager.get_history(test_chat_id)
+        # Test retrieving history
+        print(f"\n5. Retrieved history for {test_chat_id}:")
+        history = manager.get_history(test_chat_id, test_user_id)
         for msg in history:
             print(f"  Message {msg.get('message_id')}: {msg.get('role')} - {msg.get('content')[:50]}...")
         
-        # Test 3: Get chat info
-        print(f"\n3. Chat info:")
-        info = manager.get_chat_info(test_chat_id)
-        print(f"  Exists: {info.get('exists')}, Messages: {info.get('message_count')}")
+        # Test getting all sessions for this user
+        print(f"\n6. All sessions for user {test_user_id}:")
+        sessions = manager.get_all_chat_sessions(test_user_id)
+        for session in sessions:
+            print(f"  {session.get('session_id')}: {session.get('title')} ({session.get('message_count')} messages)")
         
-        # Test 4: Auto-generated chat ID
-        auto_chat_id = manager.generate_chat_id()
-        print(f"\n4. Auto-generated chat ID: {auto_chat_id}")
+        # Test user stats
+        print(f"\n7. User statistics:")
+        stats = manager.get_user_stats(test_user_id)
+        for key, value in stats.items():
+            print(f"  {key}: {value}")
         
-        # Test 5: Get all sessions
-        print(f"\n5. All chat sessions:")
-        sessions = manager.get_all_chat_sessions()
-        for session in sessions[:3]:  # Show first 3
-            print(f"  {session.get('session_id')}: {session.get('title')} ({session.get('message_count')} messages)")
+        # Test security - try to access with wrong user ID
+        print(f"\n8. Security test - trying to access with wrong user ID:")
+        wrong_user_history = manager.get_history(test_chat_id, "wrong_user_id")
+        print(f"  History with wrong user ID: {len(wrong_user_history)} messages (should be 0)")
         
         print("\n✅ All tests completed successfully!")