Update conversation_manager.py

conversation_manager.py

@@ -1,70 +1,381 @@
-# conversation_manager.py
-
-import os
-from supabase import create_client, Client
-from typing import List, Dict, Any, Optional
-import logging
-
-logger = logging.getLogger(__name__)
-
-class ConversationManager:
-    def __init__(self):
-        """Initializes the Supabase client."""
-        supabase_url = os.getenv("SUPABASE_URL")
-        supabase_key = os.getenv("SUPABASE_KEY")
-        
-        if not supabase_url or not supabase_key:
-            raise ValueError("Supabase URL and Key must be set in environment variables.")
-            
-        self.supabase: Client = create_client(supabase_url, supabase_key)
-        self.table_name = "conversations"
-        logger.info("ConversationManager initialized with Supabase client.")
-
-    def get_history(self, session_id: str) -> List[Dict[str, Any]]:
-        """
-        Retrieves conversation history for a given session_id.
-        Returns an empty list if no history is found.
-        """
-        try:
-            response = self.supabase.table(self.table_name)\
-                .select("history")\
-                .eq("session_id", session_id)\
-                .limit(1)\
-                .execute()
-            
-            if response.data:
-                return response.data[0].get("history", [])
-            return []
-        except Exception as e:
-            logger.error(f"Error fetching history for session {session_id}: {e}")
-            return []
-
-    def save_history(self, session_id: str, history: List[Dict[str, Any]]) -> None:
-        """
-        Saves or updates the conversation history for a session_id.
-        Uses 'upsert' to create a new record or update an existing one.
-        """
-        try:
-            self.supabase.table(self.table_name).upsert({
-                "session_id": session_id,
-                "history": history
-            }).execute()
-            logger.info(f"History for session {session_id} saved successfully.")
-        except Exception as e:
-            logger.error(f"Error saving history for session {session_id}: {e}")
-
-    def delete_history(self, session_id: str) -> bool:
-        """
-        Deletes the conversation history for a given session_id.
-        Returns True on success, False on failure.
-        """
-        try:
-            self.supabase.table(self.table_name)\
-                .delete()\
-                .eq("session_id", session_id)\
-                .execute()
-            logger.info(f"History for session {session_id} deleted successfully.")
-            return True
-        except Exception as e:
-            logger.error(f"Error deleting history for session {session_id}: {e}")
-            return False
+# conversation_manager.py
+
+import os
+import uuid
+import logging
+from typing import Dict, List, Any, Optional
+from datetime import datetime
+from supabase import create_client, Client
+from dotenv import load_dotenv
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+load_dotenv()
+
+
+class ConversationManager:
+    """Enhanced ConversationManager with unique message IDs and flexible chat ID management"""
+    
+    def __init__(self):
+        """Initialize the conversation manager with Supabase client"""
+        try:
+            # Initialize Supabase client
+            supabase_url = os.getenv("SUPABASE_URL")
+            supabase_key = os.getenv("SUPABASE_ANON_KEY")
+            
+            if not supabase_url or not supabase_key:
+                raise ValueError("Missing Supabase credentials in environment variables")
+            
+            self.supabase: Client = create_client(supabase_url, supabase_key)
+            logger.info("ConversationManager initialized successfully with Supabase")
+            
+        except Exception as e:
+            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_message_id(self, chat_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)
+            
+            if not history:
+                return 1
+            
+            # Find the highest message_id
+            max_id = 0
+            for message in history:
+                msg_id = message.get('message_id', 0)
+                if isinstance(msg_id, int) and msg_id > max_id:
+                    max_id = msg_id
+            
+            return max_id + 1
+            
+        except Exception as e:
+            logger.error(f"Error generating message ID for chat {chat_id}: {e}")
+            return 1
+    
+    def get_history(self, chat_id: str) -> List[Dict[str, Any]]:
+        """
+        Retrieve conversation history for a specific chat ID
+        
+        Args:
+            chat_id: The chat ID to retrieve history for
+            
+        Returns:
+            List of message dictionaries, empty list if chat doesn't exist
+        """
+        if not chat_id:
+            return []
+        
+        try:
+            response = self.supabase.table('conversations').select('history').eq('session_id', chat_id).execute()
+            
+            if response.data and len(response.data) > 0:
+                history = response.data[0].get('history', [])
+                if isinstance(history, list):
+                    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}")
+                return []
+                
+        except Exception as e:
+            logger.error(f"Error retrieving history for chat {chat_id}: {e}")
+            return []
+    
+    def save_history(self, chat_id: str, history: List[Dict[str, Any]]) -> bool:
+        """
+        Save conversation history for a specific chat ID
+        
+        Args:
+            chat_id: The chat ID to save history for
+            history: List of message dictionaries
+            
+        Returns:
+            True if successful, False otherwise
+        """
+        if not chat_id:
+            logger.error("Cannot save history: chat_id is empty")
+            return False
+        
+        if not isinstance(history, list):
+            logger.error("Cannot save history: history must be a list")
+            return False
+        
+        try:
+            # Check if conversation exists
+            existing = self.supabase.table('conversations').select('session_id').eq('session_id', chat_id).execute()
+            
+            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()
+                
+                logger.info(f"Updated history for chat {chat_id}")
+            else:
+                # Create new conversation
+                response = self.supabase.table('conversations').insert({
+                    'session_id': chat_id,
+                    'history': history,
+                    'created_at': datetime.utcnow().isoformat(),
+                    'updated_at': datetime.utcnow().isoformat()
+                }).execute()
+                
+                logger.info(f"Created new conversation for chat {chat_id}")
+            
+            return True
+            
+        except Exception as e:
+            logger.error(f"Error saving history for chat {chat_id}: {e}")
+            return False
+    
+    def add_message(self, chat_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
+            role: Message role ('user' or 'assistant')
+            content: Message content
+            image_url: Optional image URL
+            
+        Returns:
+            Dictionary containing the message with its assigned ID
+        """
+        if not chat_id:
+            logger.error("Cannot add message: chat_id is empty")
+            return {}
+        
+        try:
+            # Generate message ID
+            message_id = self.generate_message_id(chat_id)
+            
+            # Create message object
+            message = {
+                "message_id": message_id,
+                "role": role,
+                "content": content,
+                "timestamp": datetime.utcnow().isoformat()
+            }
+            
+            # Add image URL if provided
+            if image_url:
+                message["imageUrl"] = image_url
+            
+            # Get current history and add the new message
+            current_history = self.get_history(chat_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}")
+                return message
+            else:
+                logger.error(f"Failed to save message to chat {chat_id}")
+                return {}
+                
+        except Exception as e:
+            logger.error(f"Error adding message to chat {chat_id}: {e}")
+            return {}
+    
+    def delete_history(self, chat_id: str) -> bool:
+        """
+        Delete conversation history for a specific chat ID
+        
+        Args:
+            chat_id: The chat ID to delete
+            
+        Returns:
+            True if successful, False otherwise
+        """
+        if not chat_id:
+            logger.error("Cannot delete history: chat_id is empty")
+            return False
+        
+        try:
+            response = self.supabase.table('conversations').delete().eq('session_id', chat_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}")
+                return True
+            else:
+                logger.warning(f"No conversation found to delete for chat {chat_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}")
+            return False
+    
+    def get_all_chat_sessions(self) -> List[Dict[str, Any]]:
+        """
+        Get all chat sessions with basic information
+        
+        Returns:
+            List of dictionaries with chat session information
+        """
+        try:
+            response = self.supabase.table('conversations').select('session_id, history, created_at, updated_at').execute()
+            
+            sessions = []
+            for conv in response.data:
+                session_id = conv.get('session_id')
+                history = conv.get('history', [])
+                created_at = conv.get('created_at')
+                updated_at = conv.get('updated_at')
+                
+                # Generate title from first user message
+                title = "New Chat"
+                if history and len(history) > 0:
+                    first_user_msg = None
+                    for msg in history:
+                        if msg.get('role') == 'user':
+                            first_user_msg = msg
+                            break
+                    
+                    if first_user_msg:
+                        content = first_user_msg.get('content', '')
+                        if content:
+                            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,
+                    "message_count": len(history),
+                    "created_at": created_at,
+                    "updated_at": updated_at
+                })
+            
+            # Sort by updated_at (most recent first)
+            sessions.sort(key=lambda x: x.get('updated_at', ''), reverse=True)
+            
+            return sessions
+            
+        except Exception as e:
+            logger.error(f"Error fetching chat sessions: {e}")
+            return []
+    
+    def get_chat_info(self, chat_id: str) -> Dict[str, Any]:
+        """
+        Get information about a specific chat
+        
+        Args:
+            chat_id: The chat ID to get information for
+            
+        Returns:
+            Dictionary with chat information
+        """
+        if not chat_id:
+            return {}
+        
+        try:
+            response = self.supabase.table('conversations').select('*').eq('session_id', chat_id).execute()
+            
+            if response.data and len(response.data) > 0:
+                conv = response.data[0]
+                history = conv.get('history', [])
+                
+                return {
+                    "session_id": conv.get('session_id'),
+                    "message_count": len(history),
+                    "created_at": conv.get('created_at'),
+                    "updated_at": conv.get('updated_at'),
+                    "exists": True
+                }
+            else:
+                return {
+                    "session_id": chat_id,
+                    "message_count": 0,
+                    "exists": False
+                }
+                
+        except Exception as e:
+            logger.error(f"Error getting chat info for {chat_id}: {e}")
+            return {
+                "session_id": chat_id,
+                "error": str(e),
+                "exists": False
+            }
+    
+    def chat_exists(self, chat_id: str) -> bool:
+        """
+        Check if a chat ID exists in the database
+        
+        Args:
+            chat_id: The chat ID to check
+            
+        Returns:
+            True if chat exists, False otherwise
+        """
+        if not chat_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)
+            
+        except Exception as e:
+            logger.error(f"Error checking if chat {chat_id} exists: {e}")
+            return False
+
+
+# Example usage and testing
+if __name__ == "__main__":
+    print("=== ConversationManager 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}")
+        
+        # Add some messages
+        msg1 = manager.add_message(test_chat_id, "user", "Hello, I need help with farming")
+        print(f"Added user message: {msg1}")
+        
+        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}")
+        
+        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}")
+        
+        # Test 2: Retrieve history
+        print(f"\n2. Retrieved history for {test_chat_id}:")
+        history = manager.get_history(test_chat_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 4: Auto-generated chat ID
+        auto_chat_id = manager.generate_chat_id()
+        print(f"\n4. Auto-generated chat ID: {auto_chat_id}")
+        
+        # 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)")
+        
+        print("\n✅ All tests completed successfully!")
+        
+    except Exception as e:
+        print(f"❌ Test failed: {e}")