Implement task queue for sequential message processing and refactor audio/image handling

app.py

@@ -2,6 +2,7 @@ import os
 import threading
 import requests
 import logging
+import queue
 from fastapi import FastAPI, Request, HTTPException
 from fastapi.responses import PlainTextResponse, JSONResponse
 from FLUX import generate_image
@@ -11,20 +12,45 @@ from llm import generate_llm
 # Configure logging for debugging
 logging.basicConfig(level=logging.DEBUG, format="%(asctime)s [%(levelname)s] %(message)s")
 
+# Environment variables
 GREEN_API_URL = os.getenv("GREEN_API_URL")
 GREEN_API_MEDIA_URL = os.getenv("GREEN_API_MEDIA_URL", "https://api.green-api.com")
 GREEN_API_TOKEN = os.getenv("GREEN_API_TOKEN")
 GREEN_API_ID_INSTANCE = os.getenv("GREEN_API_ID_INSTANCE")
 WEBHOOK_AUTH_TOKEN = os.getenv("WEBHOOK_AUTH_TOKEN")
-PORT = 7860
+PORT = int(os.getenv("PORT", 7860))
 image_dir = "/tmp/images"
 audio_dir = "/tmp/audio"
 
 if not all([GREEN_API_URL, GREEN_API_TOKEN, GREEN_API_ID_INSTANCE, WEBHOOK_AUTH_TOKEN]):
     raise ValueError("Environment variables are not set properly")
 
+# Task queue for processing messages sequentially
+task_queue = queue.Queue()
+
 app = FastAPI()
 
+# Worker thread to process queued tasks one by one
+def worker():
+    while True:
+        task = task_queue.get()
+        try:
+            typ = task.get("type")
+            mid = task.get("message_id")
+            cid = task.get("chat_id")
+            if typ == "image":
+                handle_image_generation(mid, cid, task.get("prompt"))
+            elif typ == "audio":
+                response_audio(mid, cid, task.get("prompt"))
+        except Exception as e:
+            logging.error(f"Error processing task {task}: {e}")
+        finally:
+            task_queue.task_done()
+
+# Start the worker thread
+t = threading.Thread(target=worker, daemon=True)
+t.start()
+
 def send_message(message_id, to_number, message, retries=3):
     chat_id = to_number if to_number.endswith('@g.us') else to_number
     url = f"{GREEN_API_URL}/waInstance{GREEN_API_ID_INSTANCE}/sendMessage/{GREEN_API_TOKEN}"
@@ -44,6 +70,7 @@ def send_message(message_id, to_number, message, retries=3):
                 continue
             return {"error": str(e)}
 
+
 def send_image(message_id, to_number, image_path, retries=3):
     chat_id = to_number if to_number.endswith('@g.us') else to_number
     url = f"{GREEN_API_MEDIA_URL}/waInstance{GREEN_API_ID_INSTANCE}/sendFileByUpload/{GREEN_API_TOKEN}"
@@ -60,10 +87,8 @@ def send_image(message_id, to_number, image_path, retries=3):
                 continue
             return {"error": str(e)}
 
+
 def send_audio(message_id, to_number, audio_path, retries=3):
-    """
-    Send an audio file using the Green API similar to send_image.
-    """
     logging.debug("Entering send_audio")
     chat_id = to_number if to_number.endswith('@g.us') else to_number
 
@@ -91,6 +116,7 @@ def send_audio(message_id, to_number, audio_path, retries=3):
         logging.debug(f"Failed to open audio file: {e}")
         return {"error": str(e)}
 
+
 def response_text(message_id, chat_id, prompt):
     try:
         msg = generate_llm(prompt)
@@ -98,27 +124,28 @@ def response_text(message_id, chat_id, prompt):
     except Exception as e:
         send_message(message_id, chat_id, "There was an error processing your request.")
 
+
 def response_audio(message_id, chat_id, prompt):
     logging.debug("Entering response_audio with prompt: %s", prompt)
     try:
         result = generate_voice_reply(prompt, model="openai-audio", voice="coral", audio_dir=audio_dir)
         logging.debug("Result from generate_voice_reply: %s", result)
-        # Check result and also ensure the audio_file_path is not None or empty
         if result and result[0]:
-            audio_file_path, audio_data = result
+            audio_file_path, _ = result
             logging.debug("Audio file path generated: %s", audio_file_path)
             send_result = send_audio(message_id, chat_id, audio_file_path)
             logging.debug("Result from send_audio: %s", send_result)
             if os.path.exists(audio_file_path):
-                os.remove(audio_file_path)  # Clean up the file after sending
+                os.remove(audio_file_path)
                 logging.debug("Removed audio file: %s", audio_file_path)
         else:
-            logging.debug("generate_voice_reply returned None or empty audio file path, falling back to response_text")
+            logging.debug("Falling back to text response")
             response_text(message_id, chat_id, prompt)
     except Exception as e:
         logging.debug("Exception in response_audio: %s", e)
         send_message(message_id, chat_id, "There was an error generating the audio. Please try again later.")
 
+
 def handle_image_generation(message_id, chat_id, prompt):
     try:
         image, image_path, returned_prompt, image_url = generate_image(prompt, message_id, message_id, image_dir)
@@ -131,7 +158,7 @@ def handle_image_generation(message_id, chat_id, prompt):
             )
         else:
             send_message(message_id, chat_id, "Failed to generate image. Please try again later.")
-    except Exception as e:
+    except Exception:
         send_message(message_id, chat_id, "There was an error generating the image. Please try again later.")
 
 @app.get("/", response_class=PlainTextResponse)
@@ -143,23 +170,20 @@ async def whatsapp_webhook(request: Request):
     auth_header = request.headers.get('Authorization', '').strip()
     if auth_header != f"Bearer {WEBHOOK_AUTH_TOKEN}":
         raise HTTPException(status_code=403, detail="Unauthorized")
+
     try:
         data = await request.json()
     except Exception:
         return JSONResponse(content={"error": "Invalid JSON"}, status_code=400)
+
     if data.get('typeWebhook') != 'incomingMessageReceived':
         return {"success": True}
+
     try:
         chat_id = data['senderData']['chatId']
         message_id = data['idMessage']
         message_data = data.get('messageData', {})
 
-        # Ignore messages that are replies
-        if 'extendedTextMessageData' in message_data:
-            if message_data['extendedTextMessageData'].get('quotedMessageId'):
-                return {"success": True}
-        # If needed, add checks for other message types here
-
         if 'textMessageData' in message_data:
             body = message_data['textMessageData']['textMessage'].strip()
         elif 'extendedTextMessageData' in message_data:
@@ -168,20 +192,27 @@ async def whatsapp_webhook(request: Request):
             return {"success": True}
     except KeyError as e:
         return JSONResponse(content={"error": f"Missing key in data: {e}"}, status_code=200)
+
+    # Ignore replies between other users
+    if 'extendedTextMessageData' in message_data:
+        ctx = message_data['extendedTextMessageData'].get('contextInfo', {})
+        if ctx.get('quotedMessageId'):
+            logging.debug(f"Ignoring reply message (quotedMessageId={ctx['quotedMessageId']})")
+            return {"success": True}
+
+    # Enqueue tasks instead of spawning threads
     if body.lower().startswith('/imagine'):
         prompt = body.replace('/imagine', '').strip()
         if not prompt:
             send_message(message_id, chat_id, "Please provide a prompt after /imagine.")
         else:
             send_message(message_id, chat_id, "Generating...")
-            threading.Thread(target=handle_image_generation, args=(message_id, chat_id, prompt)).start()
+            task_queue.put({"type": "image", "message_id": message_id, "chat_id": chat_id, "prompt": prompt})
     else:
-        threading.Thread(target=response_audio, args=(message_id, chat_id, body)).start()
+        task_queue.put({"type": "audio", "message_id": message_id, "chat_id": chat_id, "prompt": body})
+
     return {"success": True}
 
-def main():
+if __name__ == '__main__':
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=PORT)
-
-if __name__ == '__main__':
-    main()