Update server.py

server.py

@@ -49,7 +49,8 @@ def chat():
     {
         "query": "user message text",
         "image_url": "optional image URL",
-        "session_id": "optional chat ID for conversation memory"
+        "session_id": "optional chat ID for conversation memory",
+        "user_id": "user identifier for authentication"
     }
     
     Response format:
@@ -72,17 +73,22 @@ def chat():
             user_message = data.get('query')
             chat_id = data.get('session_id')
             image_url = data.get('image_url')
+            user_id = data.get('user_id')
         else:
             # Fallback for form data (backward compatibility)
             data = request.form
             user_message = data.get('message') or data.get('query')
             chat_id = data.get('session_id')
             image_url = data.get('image_url')
+            user_id = data.get('user_id')
 
         # Validate input
         if not user_message and not image_url:
             return jsonify({"error": "Cannot process an empty message. Provide either 'query' text or 'image_url'."}), 400
         
+        if not user_id:
+            return jsonify({"error": "User ID is required for authentication."}), 400
+        
         # Clean up empty/null chat_id
         if chat_id in [None, '', 'null', 'undefined']:
             chat_id = None
@@ -90,6 +96,7 @@ def chat():
         # Process the query with the assistant
         result = assistant.process_query(
             user_message=user_message or "",
+            user_id=user_id,
             chat_id=chat_id,
             image_url=image_url
         )
@@ -119,7 +126,10 @@ def chat():
 @app.route('/chats', methods=['GET'])
 def get_all_chats():
     """
-    Get all chat sessions with basic information.
+    Get all chat sessions with basic information for a specific user.
+    
+    Query parameters:
+    - user_id: Required user identifier
     
     Response format:
     [
@@ -136,7 +146,11 @@ def get_all_chats():
         return jsonify({"error": "Conversation manager not available"}), 503
         
     try:
-        chats = assistant.get_all_chats()
+        user_id = request.args.get('user_id')
+        if not user_id:
+            return jsonify({"error": "User ID is required"}), 400
+            
+        chats = assistant.get_all_chats(user_id)
         return jsonify(chats)
         
     except Exception as e:
@@ -146,7 +160,10 @@ def get_all_chats():
 @app.route('/chat/<chat_id>/info', methods=['GET'])
 def get_chat_info(chat_id):
     """
-    Get information about a specific chat.
+    Get information about a specific chat for a specific user.
+    
+    Query parameters:
+    - user_id: Required user identifier
     
     Response format:
     {
@@ -164,7 +181,11 @@ def get_chat_info(chat_id):
         if not chat_id:
             return jsonify({"error": "Chat ID is required"}), 400
             
-        info = assistant.get_chat_info(chat_id)
+        user_id = request.args.get('user_id')
+        if not user_id:
+            return jsonify({"error": "User ID is required"}), 400
+            
+        info = assistant.get_chat_info(chat_id, user_id)
         return jsonify(info)
         
     except Exception as e:
@@ -174,7 +195,10 @@ def get_chat_info(chat_id):
 @app.route('/chat/<chat_id>/messages', methods=['GET'])
 def get_chat_messages(chat_id):
     """
-    Get all messages for a specific chat.
+    Get all messages for a specific chat and user.
+    
+    Query parameters:
+    - user_id: Required user identifier
     
     Response format:
     [
@@ -200,7 +224,11 @@ def get_chat_messages(chat_id):
         if not chat_id:
             return jsonify({"error": "Chat ID is required"}), 400
             
-        messages = assistant.get_messages(chat_id)
+        user_id = request.args.get('user_id')
+        if not user_id:
+            return jsonify({"error": "User ID is required"}), 400
+            
+        messages = assistant.get_messages(chat_id, user_id)
         return jsonify(messages)
         
     except Exception as e:
@@ -210,7 +238,10 @@ def get_chat_messages(chat_id):
 @app.route('/chat/<chat_id>', methods=['DELETE'])
 def delete_chat(chat_id):
     """
-    Delete a specific chat and all its messages.
+    Delete a specific chat and all its messages for a specific user.
+    
+    Query parameters:
+    - user_id: Required user identifier
     
     Response format:
     {
@@ -226,7 +257,11 @@ def delete_chat(chat_id):
         if not chat_id:
             return jsonify({"error": "Chat ID is required"}), 400
             
-        success = assistant.clear_history(chat_id)
+        user_id = request.args.get('user_id')
+        if not user_id:
+            return jsonify({"error": "User ID is required"}), 400
+            
+        success = assistant.clear_history(chat_id, user_id)
         
         if success:
             return jsonify({
@@ -245,6 +280,162 @@ def delete_chat(chat_id):
         logger.error(f"Error deleting chat {chat_id}: {e}")
         return jsonify({"error": "Could not delete chat"}), 500
 
+# --- New User Management Endpoints ---
+
+@app.route('/users/<user_id>/stats', methods=['GET'])
+def get_user_stats(user_id):
+    """
+    Get statistics for a specific user.
+    
+    Response format:
+    {
+        "user_id": "user_abc123",
+        "total_sessions": 5,
+        "total_messages": 42,
+        "recent_sessions_24h": 2,
+        "average_messages_per_session": 8.4,
+        "first_session": "2024-01-01T10:00:00",
+        "last_activity": "2024-01-01T15:30:00"
+    }
+    """
+    if not conv_manager:
+        return jsonify({"error": "Conversation manager not available"}), 503
+        
+    try:
+        if not user_id:
+            return jsonify({"error": "User ID is required"}), 400
+            
+        stats = conv_manager.get_user_stats(user_id)
+        return jsonify(stats)
+        
+    except Exception as e:
+        logger.error(f"Error getting user stats for {user_id}: {e}")
+        return jsonify({"error": "Could not fetch user statistics"}), 500
+
+@app.route('/users/<user_id>/chats', methods=['DELETE'])
+def delete_all_user_chats(user_id):
+    """
+    Delete all chats for a specific user.
+    
+    Response format:
+    {
+        "status": "success",
+        "message": "All chats deleted successfully",
+        "user_id": "user_abc123"
+    }
+    """
+    if not conv_manager:
+        return jsonify({"error": "Conversation manager not available"}), 503
+        
+    try:
+        if not user_id:
+            return jsonify({"error": "User ID is required"}), 400
+            
+        success = conv_manager.delete_all_user_data(user_id)
+        
+        if success:
+            return jsonify({
+                "status": "success",
+                "message": f"All chats deleted for user {user_id}",
+                "user_id": user_id
+            })
+        else:
+            return jsonify({
+                "status": "error",
+                "message": f"Failed to delete chats for user {user_id}",
+                "user_id": user_id
+            }), 500
+            
+    except Exception as e:
+        logger.error(f"Error deleting all chats for user {user_id}: {e}")
+        return jsonify({"error": "Could not delete user chats"}), 500
+
+@app.route('/users/<user_id>/export', methods=['GET'])
+def export_user_data(user_id):
+    """
+    Export all data for a specific user.
+    
+    Response format:
+    {
+        "user_id": "user_abc123",
+        "export_date": "2024-01-01T10:00:00",
+        "stats": {...},
+        "sessions": [...]
+    }
+    """
+    if not conv_manager:
+        return jsonify({"error": "Conversation manager not available"}), 503
+        
+    try:
+        if not user_id:
+            return jsonify({"error": "User ID is required"}), 400
+            
+        # Get user stats and sessions
+        stats = conv_manager.get_user_stats(user_id)
+        sessions = conv_manager.get_all_chat_sessions(user_id)
+        
+        # Get all messages for each session
+        detailed_sessions = []
+        for session in sessions:
+            session_data = session.copy()
+            messages = conv_manager.get_history(session['session_id'], user_id)
+            session_data['messages'] = messages
+            detailed_sessions.append(session_data)
+        
+        export_data = {
+            "user_id": user_id,
+            "export_date": datetime.utcnow().isoformat(),
+            "stats": stats,
+            "sessions": detailed_sessions,
+            "total_exported_sessions": len(detailed_sessions),
+            "total_exported_messages": sum(len(s.get('messages', [])) for s in detailed_sessions)
+        }
+        
+        return jsonify(export_data)
+        
+    except Exception as e:
+        logger.error(f"Error exporting data for user {user_id}: {e}")
+        return jsonify({"error": "Could not export user data"}), 500
+
+@app.route('/users/<user_id>', methods=['DELETE'])
+def delete_user_account(user_id):
+    """
+    Delete a user account and all associated data.
+    
+    Response format:
+    {
+        "status": "success",
+        "message": "User account deleted successfully",
+        "user_id": "user_abc123"
+    }
+    """
+    if not conv_manager:
+        return jsonify({"error": "Conversation manager not available"}), 503
+        
+    try:
+        if not user_id:
+            return jsonify({"error": "User ID is required"}), 400
+            
+        # Delete all user data
+        success = conv_manager.delete_all_user_data(user_id)
+        
+        if success:
+            return jsonify({
+                "status": "success",
+                "message": f"User account {user_id} deleted successfully",
+                "user_id": user_id
+            })
+        else:
+            return jsonify({
+                "status": "error",
+                "message": f"Failed to delete user account {user_id}",
+                "user_id": user_id
+            }), 500
+            
+    except Exception as e:
+        logger.error(f"Error deleting user account {user_id}: {e}")
+        return jsonify({"error": "Could not delete user account"}), 500
+
 # --- Legacy Endpoints for Backward Compatibility ---
 
 @app.route('/history/sessions', methods=['GET'])