re edit images

app.py

@@ -5,6 +5,7 @@ Description: A comprehensive WhatsApp bot with a professional, class-based struc
              Features include image generation, image editing, voice replies,
              and various utility functions, all handled by an asynchronous task queue.
              Includes request logging and chat ID filtering.
+             Version 2.1.0: Added caching for bot-sent images to allow for editing via replies.
 """
 
 import os
@@ -299,9 +300,13 @@ class WhatsAppBot:
         self.conv_manager = ConversationManager(config.MAX_HISTORY_SIZE)
         self.intent_router = IntentRouter(self.conv_manager, self.logger)
         self.task_queue = queue.Queue()
-        # Cache to store recent image message IDs and their download URLs
+        
+        # Cache for user-sent image message IDs and their download URLs
         self.image_cache = defaultdict(lambda: deque(maxlen=50))
-        self.fastapi_app = FastAPI(title="WhatsApp Eve Bot", version="2.0.0")
+        # NEW: Cache for bot-sent image message IDs and their local file paths
+        self.sent_image_path_cache = defaultdict(lambda: deque(maxlen=50))
+
+        self.fastapi_app = FastAPI(title="WhatsApp Eve Bot", version="2.1.0")
         self._setup_routes()
         self._start_workers()
 
@@ -314,10 +319,8 @@ class WhatsAppBot:
             
             try:
                 payload = await request.json()
-                # Log the entire incoming request payload for debugging
                 self.logger.debug(f"Incoming webhook payload: {json.dumps(payload, indent=2)}")
 
-                # Process valid incoming messages in the background
                 if payload.get("typeWebhook") == "incomingMessageReceived":
                     executor.submit(self._process_incoming_message, payload)
                 else:
@@ -342,7 +345,6 @@ class WhatsAppBot:
         while True:
             task = self.task_queue.get()
             try:
-                # Set context for the worker thread
                 ThreadContext.set_context(task["chat_id"], task["message_id"])
                 handler = getattr(self, f"_task_{task['type']}", None)
                 if handler:
@@ -361,7 +363,6 @@ class WhatsAppBot:
             chat_id = payload["senderData"]["chatId"]
             message_id = payload["idMessage"]
             
-            # ** CHAT RESTRICTION LOGIC **
             if chat_id not in self.config.ALLOWED_CHATS:
                 self.logger.warning(f"Ignoring message from unauthorized chat ID: {chat_id}")
                 return
@@ -371,35 +372,26 @@ class WhatsAppBot:
             message_data = payload.get("messageData", {})
             type_message = message_data.get("typeMessage")
             
-            # --- Unified File & Text Extraction ---
-            text = ""
-            download_url = None
-            is_image_file = False
+            text, download_url, is_image_file = "", None, 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/")):
+                if type_message == "imageMessage" or 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.logger.debug(f"Caching incoming image 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", "").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()
 
-            # --- Action Dispatching ---
             if not text:
                 self.logger.debug(f"Received message type '{type_message}' without actionable text from {chat_id}.")
                 return
@@ -407,24 +399,18 @@ class WhatsAppBot:
             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
+                        "type": "edit_image", "chat_id": chat_id, "message_id": message_id,
+                        "prompt": intent.prompt, "input_source": download_url, "source_type": 'url',
                     })
-                    return # Task dispatched, we are done.
+                    return
 
-            # Handle slash commands
             if text.startswith('/'):
                 self._handle_command(chat_id, message_id, text, payload)
-            # Handle natural language replies or regular chat
             else:
                 self._handle_natural_language(chat_id, message_id, text, payload)
         
@@ -452,7 +438,7 @@ class WhatsAppBot:
             )
             self.api_client.send_message(chat_id, help_text, message_id)
         elif command == "/gen":
-            self.task_queue.put({"type": "generate_image", "chat_id": chat_id, "message_id": message_id, "prompt": args})
+            self.task_queue.put({"type": "generate_image", "chat_id": chat_id, "message_id": message_id, "prompt": args, "count": self.config.DEFAULT_IMAGE_COUNT})
         elif command == "/edit":
             self._dispatch_edit_image(chat_id, message_id, args, payload)
         elif command == "/joke":
@@ -467,128 +453,149 @@ class WhatsAppBot:
     def _handle_natural_language(self, chat_id, message_id, text, payload):
         """Processes natural language using the intent router."""
         intent = self.intent_router.get_intent(text, chat_id)
-        
-        task_data = {
-            "chat_id": chat_id, 
-            "message_id": message_id, 
-            **intent.model_dump()
-        }
+        task_data = {"chat_id": chat_id, "message_id": message_id, **intent.model_dump()}
         
         if intent.action == "edit_image":
-            # This action needs the original payload to find the replied-to image
             self._dispatch_edit_image(chat_id, message_id, intent.prompt, payload)
         elif hasattr(self, f"_task_{intent.action}"):
-            # Enqueue the task with its type and all necessary data
             self.task_queue.put({"type": intent.action, **task_data})
         else:
             self.logger.warning(f"No handler found for intent action: {intent.action}")
             self.api_client.send_message(chat_id, "Sorry, I'm not sure how to handle that.", message_id)
             
     def _dispatch_edit_image(self, chat_id, message_id, prompt, payload):
-        """Checks for a replied-to image and dispatches the edit task."""
+        """Finds a replied-to image from cache and dispatches the edit task."""
         message_data = payload.get("messageData", {})
-        
-        # 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:
-            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
-            
-        # Get the ID of the message being replied to
+        quoted_info = message_data.get("quotedMessage", {})
         quoted_message_id = quoted_info.get("stanzaId")
         if not quoted_message_id:
             self.logger.error(f"Could not find stanzaId in quoted message: {quoted_info}")
             self.api_client.send_message(chat_id, "Sorry, I couldn't identify which image to edit.", message_id)
             return
 
-        # Find the downloadUrl from our cache
-        download_url = None
-        if chat_id in self.image_cache:
-            # Search the cache for the matching message ID
-            for msg_id, url in self.image_cache[chat_id]:
+        input_source, source_type = None, None
+
+        # 1. Check cache for bot-sent images (local paths)
+        if chat_id in self.sent_image_path_cache:
+            for msg_id, path in reversed(self.sent_image_path_cache[chat_id]):
                 if msg_id == quoted_message_id:
-                    download_url = url
+                    if os.path.exists(path):
+                        input_source, source_type = path, 'path'
+                        self.logger.info(f"Found replied-to image in local path cache: {path}")
+                    else:
+                        self.logger.warning(f"Found path for message {msg_id} but file missing: {path}")
                     break
         
-        if not download_url:
-            self.logger.warning(f"Could not find download URL for message ID {quoted_message_id} in cache for chat {chat_id}.")
+        # 2. If not found, check cache for user-sent images (URLs)
+        if not input_source and chat_id in self.image_cache:
+            for msg_id, url in reversed(self.image_cache[chat_id]):
+                if msg_id == quoted_message_id:
+                    input_source, source_type = url, 'url'
+                    self.logger.info(f"Found replied-to image in URL cache: {url}")
+                    break
+
+        if not input_source:
+            self.logger.warning(f"Could not find any image source for message ID {quoted_message_id} in cache for chat {chat_id}.")
             self.api_client.send_message(chat_id, "I couldn't find the original image. It might be too old. Please send it again before editing.", message_id)
             return
 
-        self.logger.info(f"Found cached image URL for message {quoted_message_id}. Dispatching edit task.")
+        self.logger.info(f"Dispatching edit task for message {quoted_message_id} with source type '{source_type}'.")
         self.task_queue.put({
-            "type": "edit_image",
-            "chat_id": chat_id,
-            "message_id": message_id,
-            "prompt": prompt,
-            "download_url": download_url
+            "type": "edit_image", "chat_id": chat_id, "message_id": message_id,
+            "prompt": prompt, "input_source": input_source, "source_type": source_type,
         })
         
+    def _manage_cached_file(self, chat_id: str, message_id: str, file_path: str):
+        """Adds a file to the cache and cleans up the oldest file if the cache is full."""
+        # Check if the cache is full before adding a new item
+        if len(self.sent_image_path_cache[chat_id]) == self.sent_image_path_cache[chat_id].maxlen:
+            old_id, old_path = self.sent_image_path_cache[chat_id][0] # Oldest item is at the left
+            if os.path.exists(old_path):
+                self.logger.debug(f"Cache full. Removing oldest cached image file: {old_path}")
+                try:
+                    os.remove(old_path)
+                except OSError as e:
+                    self.logger.error(f"Error removing old cached file {old_path}: {e}")
+        
+        self.sent_image_path_cache[chat_id].append((message_id, file_path))
+
     # --- Task Handler Methods ---
 
     def _task_send_text(self, task: Dict[str, Any]):
         chat_id, message_id, message = task["chat_id"], task["message_id"], task["message"]
         self.api_client.send_message(chat_id, message, message_id)
         self.conv_manager.add_bot_message(chat_id, message)
-        # Asynchronously generate a voice reply for the sent text
         self.task_queue.put({"type": "voice_reply", "chat_id": chat_id, "message_id": message_id, "text": message})
 
     def _task_generate_image(self, task: Dict[str, Any]):
-        chat_id, mid, prompt, count = task["chat_id"], task["message_id"], task["prompt"], task.get("count", 1)
+        chat_id, mid, prompt = task["chat_id"], task["message_id"], task["prompt"]
+        count = task.get("count", self.config.DEFAULT_IMAGE_COUNT)
         self.api_client.send_message(chat_id, f"🎨 Generating {count} image(s) for: \"{prompt}\"...", mid)
         
         for i in range(count):
+            path = None
             try:
-                _, path, _, url = generate_image(prompt, mid, str(i), self.config.IMAGE_DIR, width=task.get("width"), height=task.get("height"))
+                _, path, _, _ = generate_image(prompt, mid, str(i), self.config.IMAGE_DIR, width=task.get("width"), height=task.get("height"))
                 caption = f"✨ Image {i+1}/{count}: {prompt}"
-                self.api_client.send_file(chat_id, path, caption, mid)
-                os.remove(path)
+                response = self.api_client.send_file(chat_id, path, caption, mid)
+
+                if response and response.get("idMessage"):
+                    sent_message_id = response["idMessage"]
+                    self.logger.info(f"Caching sent image path for message {sent_message_id}: {path}")
+                    self._manage_cached_file(chat_id, sent_message_id, path)
+                else:
+                    self.logger.warning(f"Could not get sent message ID for {path}. Deleting file.")
+                    os.remove(path)
             except Exception as e:
                 self.logger.error(f"Image generation {i+1} failed: {e}", exc_info=True)
                 self.api_client.send_message(chat_id, f"😢 Failed to generate image {i+1}.", mid)
+                if path and os.path.exists(path): os.remove(path)
     
     def _task_edit_image(self, task: Dict[str, Any]):
-        chat_id, mid, prompt, url = task["chat_id"], task["message_id"], task["prompt"], task["download_url"]
+        chat_id, mid, prompt = task["chat_id"], task["message_id"], task["prompt"]
+        input_source, source_type = task["input_source"], task["source_type"]
         self.api_client.send_message(chat_id, f"🎨 Editing image with prompt: \"{prompt}\"...", mid)
 
-        input_path, output_path = None, None
+        input_path_temp, output_path = None, None
         try:
-            image_data = self.api_client.download_file(url)
-            if not image_data:
-                raise ValueError("Failed to download image.")
-
             os.makedirs(self.config.TEMP_DIR, exist_ok=True)
-            input_path = os.path.join(self.config.TEMP_DIR, f"input_{mid}.jpg")
             output_path = os.path.join(self.config.TEMP_DIR, f"output_{mid}.jpg")
-            
-            with open(input_path, 'wb') as f:
-                f.write(image_data)
-            
-            flux_kontext_lib.generate_image(prompt, input_path, download_path=output_path)
+
+            if source_type == 'url':
+                image_data = self.api_client.download_file(input_source)
+                if not image_data: raise ValueError(f"Failed to download image from URL: {input_source}")
+                input_path_temp = os.path.join(self.config.TEMP_DIR, f"input_{mid}.jpg")
+                with open(input_path_temp, 'wb') as f: f.write(image_data)
+                edit_input_path = input_path_temp
+            elif source_type == 'path':
+                edit_input_path = input_source
+            else:
+                raise ValueError(f"Unknown source_type: {source_type}")
+
+            flux_kontext_lib.generate_image(prompt, edit_input_path, download_path=output_path)
             
             if os.path.exists(output_path):
                 caption = f"✨ Edited: {prompt}"
-                self.api_client.send_file(chat_id, output_path, caption, mid)
+                response = self.api_client.send_file(chat_id, output_path, caption, mid)
+                if response and response.get("idMessage"):
+                    sent_message_id = response["idMessage"]
+                    self.logger.info(f"Caching sent image path for message {sent_message_id}: {output_path}")
+                    self._manage_cached_file(chat_id, sent_message_id, output_path)
+                else:
+                    self.logger.warning(f"Could not get sent message ID for {output_path}. Deleting file.")
+                    os.remove(output_path)
             else:
                 raise ValueError("Edited image file not found.")
-
         except Exception as e:
             self.logger.error(f"Image editing task failed: {e}", exc_info=True)
             self.api_client.send_message(chat_id, "😢 Sorry, I failed to edit the image.", mid)
+            if output_path and os.path.exists(output_path): os.remove(output_path)
         finally:
-            for path in [input_path, output_path]:
-                if path and os.path.exists(path):
-                    os.remove(path)
+            if input_path_temp and os.path.exists(input_path_temp): os.remove(input_path_temp)
     
     def _task_voice_reply(self, task: Dict[str, Any]):
         text = task["text"]
@@ -635,7 +642,6 @@ class WhatsAppBot:
             os.makedirs(d, exist_ok=True)
             self.logger.debug(f"Ensured directory exists: {d}")
         
-        # Send a startup message to one of the allowed chats to confirm it's online
         self.api_client.send_message(
            "[email protected]",
            "🌟 Eve is online and ready to help! Type /help to see commands."
@@ -651,7 +657,6 @@ if __name__ == "__main__":
         bot = WhatsAppBot(bot_config)
         bot.run()
     except ValueError as e:
-        # Catch config validation errors
         logging.basicConfig()
         logging.getLogger().critical(f"❌ CONFIGURATION ERROR: {e}")
     except KeyboardInterrupt: