handle document images

app.py

@@ -371,41 +371,61 @@ class WhatsAppBot:
             message_data = payload.get("messageData", {})
             type_message = message_data.get("typeMessage")
             
-            # --- Improved message processing logic ---
-            
-            # 1. Handle image caching
-            if type_message == "imageMessage":
-                download_url = message_data.get("fileMessageData", {}).get("downloadUrl")
-                if download_url:
-                    self.logger.debug(f"Caching image message {message_id} from {chat_id} with URL: {download_url}")
-                    self.image_cache[chat_id].append((message_id, download_url))
-
-            # 2. Extract text from various message types
+            # --- Unified File & Text Extraction ---
             text = ""
-            if type_message == "textMessage":
-                text = message_data.get("textMessageData", {}).get("textMessage", "")
-            elif type_message == "extendedTextMessage":
-                text = message_data.get("extendedTextMessageData", {}).get("text", "")
+            download_url = None
+            is_image_file = False
+
+            # Case 1 & 4: Image or Document with potential caption
+            if type_message in ["imageMessage", "documentMessage"]:
+                file_data = message_data.get("fileMessageData", {})
+                mime_type = file_data.get("mimeType", "")
+
+                if type_message == "imageMessage" or (type_message == "documentMessage" and mime_type.startswith("image/")):
+                    is_image_file = True
+                    download_url = file_data.get("downloadUrl")
+                    text = file_data.get("caption", "").strip()
+                    if download_url:
+                        self.logger.debug(f"Caching file message {message_id} from {chat_id}")
+                        self.image_cache[chat_id].append((message_id, download_url))
+
+            # Case 2 & 3: Reply to something (quoted message)
             elif type_message == "quotedMessage":
-                text = message_data.get("extendedTextMessageData", {}).get("text", "")
-            elif type_message == "imageMessage":
-                text = message_data.get("fileMessageData", {}).get("caption", "")
+                text = message_data.get("extendedTextMessageData", {}).get("text", "").strip()
+
+            # Fallback for other text types
+            elif type_message == "textMessage":
+                text = message_data.get("textMessageData", {}).get("textMessage", "").strip()
+            elif type_message == "extendedTextMessage":
+                 text = message_data.get("extendedTextMessageData", {}).get("text", "").strip()
 
-            text = text.strip()
+            # --- Action Dispatching ---
             if not text:
-                self.logger.debug(f"Received message type '{type_message}' without actionable text from {chat_id}. No action taken.")
+                self.logger.debug(f"Received message type '{type_message}' without actionable text from {chat_id}.")
                 return
-            
-            # --- End of improved logic ---
-            
-            self.logger.info(f"Processing message from {chat_id}: '{text}'")
-            self.conv_manager.add_user_message(chat_id, text)
 
-            # Handle direct commands
+            self.logger.info(f"Processing text from {chat_id}: '{text}'")
+            self.conv_manager.add_user_message(chat_id, text)
+            
+            # If it's a file with a caption, check if the caption is an edit command
+            if is_image_file and text:
+                intent = self.intent_router.get_intent(text, chat_id)
+                if isinstance(intent, EditImageIntent):
+                    self.logger.info("Detected direct edit command in image caption.")
+                    self.task_queue.put({
+                        "type": "edit_image",
+                        "chat_id": chat_id,
+                        "message_id": message_id,
+                        "prompt": intent.prompt,
+                        "download_url": download_url
+                    })
+                    return # Task dispatched, we are done.
+
+            # Handle slash commands
             if text.startswith('/'):
                 self._handle_command(chat_id, message_id, text, payload)
+            # Handle natural language replies or regular chat
             else:
-                # Handle natural language and replies
                 self._handle_natural_language(chat_id, message_id, text, payload)
         
         except KeyError as e:
@@ -468,13 +488,19 @@ class WhatsAppBot:
         """Checks for a replied-to image and dispatches the edit task."""
         message_data = payload.get("messageData", {})
         
-        # Ensure we are dealing with a reply
+        # This function is only for replies, so the top-level type must be quotedMessage
         if message_data.get("typeMessage") != "quotedMessage":
             self.api_client.send_message(chat_id, "To edit an image, please reply to an image with your instructions.", message_id)
             return
 
         quoted_info = message_data.get("quotedMessage")
-        if not quoted_info or quoted_info.get("typeMessage") != "imageMessage":
+        if not quoted_info:
+            self.api_client.send_message(chat_id, "Couldn't find the message you replied to.", message_id)
+            return
+
+        # Check if the replied-to message is an image OR a document
+        quoted_type = quoted_info.get("typeMessage")
+        if quoted_type not in ["imageMessage", "documentMessage"]:
             self.api_client.send_message(chat_id, "You must reply to an image to edit it.", message_id)
             return