Update main.py

main.py

@@ -101,32 +101,135 @@ GREETING_PATTERN = re.compile(r'^\s*(hi|hello|hola|hey|greetings|sawubona)\b.*$'
 # --- Duplicate Message Handling ---
 PROCESSED_MESSAGES_TTL_HOURS = 24 # Store processed message IDs for 24 hours
 
-def check_and_mark_processed(message_id: str) -> bool:
-    """Improved duplicate checking with atomic transactions"""
-    if not db or not message_id:
-        logger.warning("Firestore client or message ID missing")
-        return False
+# Add this new function to your code
 
-    doc_ref = db.collection("processed_messages").document(message_id)
+import threading
+from collections import OrderedDict
+from time import time
+
+class MessageDeduplicator:
+    """Thread-safe in-memory cache for message IDs with TTL and disk backup."""
     
-    try:
-        # Use transaction to prevent race conditions
-        @firestore.transactional
-        def check_and_mark_transaction(transaction):
-            doc = transaction.get(doc_ref)
-            if doc.exists:
+    def __init__(self, ttl_hours=24, max_cache_size=10000, db_client=None):
+        """
+        Initialize deduplicator with TTL in hours, max cache size, and optional DB client.
+        
+        Args:
+            ttl_hours: Hours to keep message IDs in memory/DB
+            max_cache_size: Maximum number of messages to keep in memory
+            db_client: Firestore client for persistent storage (optional)
+        """
+        self.ttl_seconds = ttl_hours * 3600
+        self.max_cache_size = max_cache_size
+        self.db_client = db_client
+        self.cache = OrderedDict()  # {message_id: timestamp}
+        self.lock = threading.RLock()
+        
+        # Start background cleanup thread
+        self._cleanup_thread = threading.Thread(
+            target=self._periodic_cleanup, 
+            daemon=True
+        )
+        self._cleanup_thread.start()
+    
+    def is_duplicate(self, message_id):
+        """
+        Check if message_id has been processed already.
+        
+        Args:
+            message_id: Unique identifier for the message
+            
+        Returns:
+            bool: True if message is a duplicate, False otherwise
+        """
+        if not message_id:
+            return False
+            
+        with self.lock:
+            # Check in-memory cache first (fastest)
+            if message_id in self.cache:
+                # Refresh the timestamp and move to end (most recently used)
+                self.cache.move_to_end(message_id)
                 return True
-            transaction.set(doc_ref, {
-                "processed_at": firestore.SERVER_TIMESTAMP,
-                "expires_at": datetime.now() + timedelta(hours=PROCESSED_MESSAGES_TTL_HOURS)
-            })
+                
+            # If not in memory, check database if available
+            if self.db_client:
+                try:
+                    doc_ref = self.db_client.collection("processed_messages").document(message_id)
+                    doc = doc_ref.get()
+                    if doc.exists:
+                        # Add to in-memory cache too
+                        self.cache[message_id] = time()
+                        # Maintain cache size limit
+                        if len(self.cache) > self.max_cache_size:
+                            self.cache.popitem(last=False)  # Remove oldest item
+                        return True
+                except Exception as e:
+                    logger.error(f"Database error in is_duplicate: {e}", exc_info=True)
+            
+            # Not a duplicate, mark as processed
+            self._mark_processed(message_id)
             return False
-
-        return check_and_mark_transaction(db.transaction())
+    
+    def _mark_processed(self, message_id):
+        """Mark message as processed in both cache and database."""
+        current_time = time()
         
-    except Exception as e:
-        logger.error(f"Firestore error: {e}", exc_info=True)
+        with self.lock:
+            # Add to in-memory cache
+            self.cache[message_id] = current_time
+            
+            # Maintain cache size limit
+            if len(self.cache) > self.max_cache_size:
+                self.cache.popitem(last=False)  # Remove oldest item
+            
+            # Add to database if available
+            if self.db_client:
+                try:
+                    expiry = datetime.now() + timedelta(hours=self.ttl_seconds/3600)
+                    doc_ref = self.db_client.collection("processed_messages").document(message_id)
+                    doc_ref.set({
+                        "processed_at": firestore.SERVER_TIMESTAMP,
+                        "expires_at": expiry
+                    })
+                except Exception as e:
+                    logger.error(f"Database error in _mark_processed: {e}", exc_info=True)
+    
+    def _periodic_cleanup(self):
+        """Background thread that removes expired entries periodically."""
+        while True:
+            try:
+                with self.lock:
+                    current_time = time()
+                    # Remove expired entries from memory
+                    expired_ids = [
+                        msg_id for msg_id, timestamp in list(self.cache.items())
+                        if current_time - timestamp > self.ttl_seconds
+                    ]
+                    for msg_id in expired_ids:
+                        self.cache.pop(msg_id, None)
+                
+                # Sleep for an hour between cleanups
+                threading.Event().wait(3600)
+            except Exception as e:
+                logger.error(f"Error in cleanup thread: {e}", exc_info=True)
+                # Sleep for a while before retrying
+                threading.Event().wait(300)
+
+# Initialize the deduplicator at the module level
+message_deduplicator = MessageDeduplicator(
+    ttl_hours=PROCESSED_MESSAGES_TTL_HOURS,
+    db_client=db
+)
+
+# Replace your check_and_mark_processed function with this:
+def check_and_mark_processed(message_id: str) -> bool:
+    """Check if a message has been processed already using the deduplicator."""
+    if not message_id:
+        logger.warning("Empty message ID provided to check_and_mark_processed")
         return False
+        
+    return message_deduplicator.is_duplicate(message_id)
 
 
 # --- Message Sending Wrappers (using whatsapp_client) ---
@@ -347,7 +450,7 @@ def process_audio_message(audio_id: str, mobile: str) -> None:
 # --- Webhook Endpoint ---
 @app.route("/", methods=["GET", "POST"])
 def webhook_handler():
-    """Handles incoming WhatsApp webhooks."""
+    """Handles incoming WhatsApp webhooks with improved duplicate detection."""
     if request.method == "GET":
         # Webhook Verification
         mode = request.args.get("hub.mode")
@@ -362,77 +465,89 @@ def webhook_handler():
 
     # --- POST Request Handling ---
     if request.method == "POST":
-        data = request.get_json()
-        if not data:
-            logger.warning("Received empty POST request body.")
-            return make_response("Bad Request: No data", 400)
+        try:
+            data = request.get_json()
+            if not data:
+                logger.warning("Received empty POST request body.")
+                return make_response("Bad Request: No data", 400)
 
-        # Log the raw incoming payload for debugging
-        logger.debug(f"Incoming Webhook Payload: {json.dumps(data, indent=2)}")
+            # Log the raw incoming payload for debugging (consider sampling for high traffic)
+            # logger.debug(f"Incoming Webhook Payload: {json.dumps(data, indent=2)}")
 
-        # Use the helper function from whatsapp_client to parse the payload
-        msg_details = wa_client.get_message_details(data)
+            # Use the helper function from whatsapp_client to parse the payload
+            msg_details = wa_client.get_message_details(data)
 
-        if not msg_details:
-            # Not a message we handle (e.g., status update, unsupported type) or parsing failed
-            logger.info("Webhook payload did not contain a processable message.")
-            return make_response("ok", 200) # Acknowledge receipt even if not processed
+            if not msg_details:
+                # Not a message we handle (e.g., status update, unsupported type) or parsing failed
+                logger.info("Webhook payload did not contain a processable message.")
+                return make_response("ok", 200)  # Acknowledge receipt even if not processed
 
-        message_id = msg_details.get("id")
-        mobile = msg_details.get("from")
-        message_type = msg_details.get("type")
+            message_id = msg_details.get("id")
+            mobile = msg_details.get("from")
+            message_type = msg_details.get("type")
 
-        # --- Duplicate Check ---
-        if message_id and check_and_mark_processed(message_id):
-             return make_response("ok - duplicate", 200) # Acknowledge receipt, but don't process
+            # Early return if no message ID or phone number
+            if not message_id or not mobile:
+                logger.warning("Missing message ID or sender phone number in payload")
+                return make_response("ok - invalid message", 200)
 
-        # --- Route based on message type ---
-        try:
+            # --- Improved Duplicate Check ---
+            if check_and_mark_processed(message_id):
+                logger.info(f"Ignoring duplicate message with ID {message_id} from {mobile}")
+                return make_response("ok - duplicate", 200)  # Acknowledge but don't process
+
+            # --- Route based on message type ---
             if message_type == "text":
                 message_text = msg_details.get("text")
-                if message_text and mobile:
+                if message_text:
                     process_text_message(message_text, mobile)
                 else:
-                    logger.warning(f"Missing text content or sender for text message: {msg_details}")
+                    logger.warning(f"Missing text content for text message: {msg_details}")
 
             elif message_type == "audio":
                 audio_id = msg_details.get("audio_id")
-                if audio_id and mobile:
+                if audio_id:
                     process_audio_message(audio_id, mobile)
                 else:
-                    logger.warning(f"Missing audio ID or sender for audio message: {msg_details}")
+                    logger.warning(f"Missing audio ID for audio message: {msg_details}")
 
             elif message_type == "interactive":
                 interactive_type = msg_details.get("interactive_type")
                 if interactive_type == "button_reply":
                     button_id = msg_details.get("button_reply_id")
-                    if button_id and mobile:
+                    if button_id:
                         handle_interactive_response(mobile, button_id)
                     else:
-                        logger.warning(f"Missing button ID or sender for button reply: {msg_details}")
-                # Add elif for list_reply etc. if needed
+                        logger.warning(f"Missing button ID for button reply: {msg_details}")
                 else:
-                     logger.info(f"Received unhandled interactive type: {interactive_type}")
-
+                    logger.info(f"Received unhandled interactive type: {interactive_type}")
             else:
-                # Handle other message types if needed (image, document, location etc.)
+                # Handle other message types if needed
                 logger.info(f"Received unhandled message type: {message_type} from {mobile}")
-                # Optionally send a message back
-                # wa_client.send_text_message(mobile, f"Sorry, I cannot handle '{message_type}' messages yet.")
 
         except Exception as e:
             # Generic error handler for safety
             logger.error(f"Unhandled exception during message processing: {e}", exc_info=True)
-            if mobile: # Try to notify the user if we know who they are
-                 wa_client.send_text_message(mobile, "Sorry, an unexpected error occurred while processing your message.")
-
-        # Acknowledge receipt of the webhook event
+            # Only try to notify if we can extract the mobile number
+            mobile = None
+            try:
+                msg_details = wa_client.get_message_details(data)
+                mobile = msg_details.get("from") if msg_details else None
+            except:
+                pass
+                
+            if mobile:
+                wa_client.send_text_message(
+                    mobile, 
+                    "Sorry, an unexpected error occurred while processing your message."
+                )
+
+        # Always acknowledge receipt of the webhook event
         return make_response("ok", 200)
 
     # Method Not Allowed
     return make_response("Method Not Allowed", 405)
 
-
 if __name__ == '__main__':
     port = int(os.environ.get("PORT", 7860)) # Use PORT env var if available (for deployment)
     debug_mode = os.environ.get("FLASK_DEBUG", "False").lower() == "true"