Refactor app.py: enhance BotConfig validation, streamline environment variable checks, and improve image generation handling for better clarity and functionality.

app.py

@@ -3,13 +3,12 @@ import threading
 import requests
 import logging
 import queue
-import re
 import json
 import time
 import random
 from concurrent.futures import ThreadPoolExecutor
 from fastapi import FastAPI, Request, HTTPException
-from fastapi.responses import PlainTextResponse, JSONResponse
+from fastapi.responses import PlainTextResponse
 from FLUX import generate_image
 from VoiceReply import generate_voice_reply
 from polLLM import generate_llm
@@ -17,19 +16,26 @@ from polLLM import generate_llm
 # --- Configuration and Client Classes ---
 
 class BotConfig:
-    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")
-    BOT_GROUP_CHAT        = "[email protected]"
-    IMAGE_DIR             = "/tmp/images"
-    AUDIO_DIR             = "/tmp/audio"
+    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")
+    BOT_GROUP_CHAT          = "[email protected]"
+    IMAGE_DIR               = "/tmp/images"
+    AUDIO_DIR               = "/tmp/audio"
+    DEFAULT_IMAGE_COUNT     = 4
 
     @classmethod
     def validate(cls):
-        if not all([cls.GREEN_API_URL, cls.GREEN_API_TOKEN, cls.GREEN_API_ID_INSTANCE, cls.WEBHOOK_AUTH_TOKEN]):
-            raise ValueError("Environment variables are not set properly")
+        missing = [name for name in (
+            "GREEN_API_URL",
+            "GREEN_API_TOKEN",
+            "GREEN_API_ID_INSTANCE",
+            "WEBHOOK_AUTH_TOKEN",
+        ) if not getattr(cls, name)]
+        if missing:
+            raise ValueError(f"Environment variables not set: {', '.join(missing)}")
 
 class BotClient:
     def __init__(self, cfg: BotConfig):
@@ -41,7 +47,12 @@ class BotClient:
         url = f"{self.cfg.GREEN_API_URL}/waInstance{self.cfg.GREEN_API_ID_INSTANCE}/{endpoint}/{self.cfg.GREEN_API_TOKEN}"
         for attempt in range(1, retries + 1):
             try:
-                resp = self.session.post(url, json=payload if files is None else None, data=None if files is None else payload, files=files)
+                resp = self.session.post(
+                    url,
+                    json=payload if files is None else None,
+                    data=None if files is None else payload,
+                    files=files
+                )
                 resp.raise_for_status()
                 return resp.json()
             except requests.RequestException as e:
@@ -52,11 +63,11 @@ class BotClient:
 
     def send_message(self, message_id: str, chat_id: str, text: str):
         payload = {"chatId": chat_id, "message": text, "quotedMessageId": message_id}
-        return self.send(f"sendMessage", payload)
+        return self.send("sendMessage", payload)
 
     def send_message_to(self, chat_id: str, text: str):
         payload = {"chatId": chat_id, "message": text}
-        return self.send(f"sendMessage", payload)
+        return self.send("sendMessage", payload)
 
     def send_media(self, message_id: str, chat_id: str, file_path: str, caption: str, media_type: str):
         endpoint = "sendFileByUpload"
@@ -70,35 +81,33 @@ class BotClient:
 BotConfig.validate()
 client = BotClient(BotConfig)
 
-# --- Queues, stores, and threading ---
+# --- Queues, stores, threading ---
 
 task_queue   = queue.Queue()
 trivia_store = {}
 polls        = {}
-
 last_message_time = time.time()
 
-# Inactivity monitor
 def inactivity_monitor():
     global last_message_time
     while True:
         time.sleep(60)
         if time.time() - last_message_time >= 300:
-            client.send_message_to(BotConfig.BOT_GROUP_CHAT,
-                                   "⏰ I haven't heard from you in a while! I'm still here if you need anything.")
+            client.send_message_to(
+                BotConfig.BOT_GROUP_CHAT,
+                "⏰ I haven't heard from you in a while! I'm still here if you need anything."
+            )
             last_message_time = time.time()
 
 threading.Thread(target=inactivity_monitor, daemon=True).start()
 
-# Worker pool
 executor = ThreadPoolExecutor(max_workers=4)
-
 def worker():
     while True:
         task = task_queue.get()
         try:
             if task["type"] == "image":
-                handle_image_generation(task["message_id"], task["chat_id"], task["prompt"])
+                handle_image_generation(task)
             elif task["type"] == "audio":
                 response_audio(task["message_id"], task["chat_id"], task["prompt"])
         except Exception as e:
@@ -109,7 +118,7 @@ def worker():
 for _ in range(4):
     threading.Thread(target=worker, daemon=True).start()
 
-# --- Core response functions ---
+# --- Core responders ---
 
 def response_text(message_id, chat_id, prompt):
     try:
@@ -132,44 +141,55 @@ def response_audio(message_id, chat_id, prompt):
         logging.error(f"Audio error: {e}")
         client.send_message(message_id, chat_id, "Error generating audio. Try again later.")
 
-def handle_image_generation(message_id, chat_id, prompt):
-    attempts = 4
-    for attempt in range(1, attempts + 1):
-        try:
-            img, path, ret_prompt, url = generate_image(prompt, message_id, message_id, BotConfig.IMAGE_DIR)
-            if img:
+def handle_image_generation(task):
+    message_id = task["message_id"]
+    chat_id    = task["chat_id"]
+    prompt     = task["prompt"]
+    count      = task.get("num_images", BotConfig.DEFAULT_IMAGE_COUNT)
+
+    for i in range(1, count + 1):
+        for attempt in range(1, 4):
+            try:
+                img, path, ret_prompt, url = generate_image(
+                    prompt,
+                    message_id,
+                    f"{message_id}_{i}",
+                    BotConfig.IMAGE_DIR
+                )
+                if not img:
+                    raise RuntimeError("No image returned")
                 formatted = "\n\n".join(f"_{p.strip()}_" for p in ret_prompt.split("\n\n") if p.strip())
-                caption = f"✨ Image ready: {url}\n>{chr(8203)} {formatted}"
+                caption = f"✨ Image {i}/{count}: {url}\n>{chr(8203)} {formatted}"
                 client.send_media(message_id, chat_id, path, caption, media_type="image")
-                return
-            else:
-                raise RuntimeError("generate_image returned no image")
-        except Exception as e:
-            logging.warning(f"Image gen attempt {attempt}/{attempts} failed: {e}")
-            if attempt < attempts:
-                time.sleep(5)
-            else:
-                client.send_message(message_id, chat_id, "😢 Image generation failed after several attempts.")
+                os.remove(path)
+                break
+            except Exception as e:
+                logging.warning(f"Image {i}/{count} attempt {attempt} failed: {e}")
+                time.sleep(2 * attempt)
+        else:
+            client.send_message(message_id, chat_id, f"😢 Failed to generate image {i}.")
 
 # --- Startup ---
 
 def send_startup_message():
-    client.send_message_to(BotConfig.BOT_GROUP_CHAT,
-                           "🌟 Hi! I'm Eve, your friendly AI assistant. I'm now live and ready to help!")
+    client.send_message_to(
+        BotConfig.BOT_GROUP_CHAT,
+        "🌟 Hi! I'm Eve, your friendly AI assistant. I'm now live and ready to help!"
+    )
 
 send_startup_message()
 
 # --- FastAPI App & Webhook ---
 
 app = FastAPI()
-
 help_text = (
     "🤖 *Hi there, I'm Eve!* Here are the commands you can use:\n\n"
     "• */help* – _Show this help message._\n"
     "• */summarize <text>* – _Get a quick summary._\n"
     "• */translate <language>|<text>* – _Translate text._\n"
     "• */joke* – _Get a joke._\n"
-    "• */weather <location>* – _Get weather._\n"
+    "• */weather <location>* – _Short, creative weather report (°C)._ \n"
+    "• */weatherpoem <location>* – _Poetic weather summary._\n"
     "• */inspire* – _Get an inspirational quote._\n"
     "• */trivia* – _Start trivia._\n"
     "• */answer* – _Answer trivia._\n"
@@ -177,7 +197,7 @@ help_text = (
     "• */poll Q|A|B…* – _Create a poll._\n"
     "• */results* – _Show poll results._\n"
     "• */endpoll* – _End the poll._\n"
-    "• */gen <prompt>* – _Generate an image._\n\n"
+    "• */gen <prompt>|<count>* – _Generate images (default 4)._ \n\n"
     "Any other text → I'll send you a voice reply!"
 )
 
@@ -186,46 +206,49 @@ async def whatsapp_webhook(request: Request):
     global last_message_time
     last_message_time = time.time()
 
-    # Auth
-    auth = request.headers.get("Authorization", "")
-    if auth != f"Bearer {BotConfig.WEBHOOK_AUTH_TOKEN}":
+    # --- Auth & basic routing ---
+    if request.headers.get("Authorization") != f"Bearer {BotConfig.WEBHOOK_AUTH_TOKEN}":
         raise HTTPException(403, "Unauthorized")
 
     data = await request.json()
     chat_id = data.get("senderData", {}).get("chatId")
-    # Only respond in the group chat
-    if chat_id != BotConfig.BOT_GROUP_CHAT:
-        return {"success": True}
-
-    if data.get("typeWebhook") != "incomingMessageReceived":
+    if chat_id != BotConfig.BOT_GROUP_CHAT or data.get("typeWebhook") != "incomingMessageReceived":
         return {"success": True}
 
     md = data.get("messageData", {})
     mid = data["idMessage"]
 
-    # Extract text
-    if "textMessageData" in md:
-        body = md["textMessageData"]["textMessage"].strip()
-    elif "extendedTextMessageData" in md:
-        body = md["extendedTextMessageData"]["text"].strip()
-    else:
+    # --- Extract text + contextInfo in one go ---
+    text_data = md.get("textMessageData") or md.get("extendedTextMessageData")
+    if not text_data:
+        return {"success": True}
+    body = text_data.get("textMessage", text_data.get("text", "")).strip()
+    ctx  = text_data.get("contextInfo", {})
+
+    # --- SKIP if this is a quoted/replied message ---
+    if ctx.get("quotedMessageId") or ctx.get("quotedMessage"):
+        logging.debug("Skipping quoted/reply message")
+        return {"success": True}
+
+    # --- SKIP if someone is @-mentioned ---
+    if ctx.get("mentionedJidList"):
+        logging.debug("Skipping mention")
         return {"success": True}
 
     low = body.lower()
 
-    # Commands
+    # --- COMMANDS ---
     if low == "/help":
         client.send_message(mid, chat_id, help_text)
         return {"success": True}
 
     if low.startswith("/summarize "):
-        txt = body[len("/summarize "):].strip()
-        summary = generate_llm(f"Summarize this text in one short paragraph:\n\n{txt}")
+        summary = generate_llm(f"Summarize this text in one short paragraph:\n\n{body[11:].strip()}")
         client.send_message(mid, chat_id, summary)
         return {"success": True}
 
     if low.startswith("/translate "):
-        part = body[len("/translate "):]
+        part = body[11:]
         if "|" not in part:
             client.send_message(mid, chat_id, "Please use `/translate <language>|<text>`")
         else:
@@ -243,37 +266,53 @@ async def whatsapp_webhook(request: Request):
         return {"success": True}
 
     if low.startswith("/weather "):
-        loc = body[len("/weather "):].strip().replace(" ", "+")
+        loc = body[9:].strip().replace(" ", "+")
+        try:
+            raw = requests.get(f"http://sl.wttr.in/{loc}?format=4", timeout=5).text
+            prompt = (
+                f"Convert this weather report into Celsius and craft a short, creative, "
+                f"beautiful weather report with emojis:\n\n{raw}"
+            )
+            report = generate_llm(prompt)
+            client.send_message(mid, chat_id, report)
+            voice_prompt = f"Provide only the following weather report as speech: {report}"
+            task_queue.put({"type":"audio","message_id":mid,"chat_id":chat_id,"prompt":voice_prompt})
+        except:
+            client.send_message(mid, chat_id, "Could not fetch weather.")
+        return {"success": True}
+
+    if low.startswith("/weatherpoem "):
+        loc = body[13:].strip().replace(" ", "+")
         try:
-            w = requests.get(f"http://sl.wttr.in/{loc}?format=4", timeout=5).text
-            client.send_message(mid, chat_id, w)
+            raw = requests.get(f"http://sl.wttr.in/{loc}?format=4", timeout=5).text
+            poem = generate_llm(f"Write a short, poetic weather summary in Celsius based on this:\n\n{raw}")
+            client.send_message(mid, chat_id, poem)
+            voice_prompt = f"Speak only this poetic weather summary: {poem}"
+            task_queue.put({"type":"audio","message_id":mid,"chat_id":chat_id,"prompt":voice_prompt})
         except:
             client.send_message(mid, chat_id, "Could not fetch weather.")
         return {"success": True}
 
     if low == "/inspire":
-        quote = generate_llm("Give me a short inspirational quote.")
+        quote = generate_llm(f"Give me a short inspirational unique quote with seed {random.randint(0,1e6)}.")
         client.send_message(mid, chat_id, f"✨ {quote}")
         return {"success": True}
 
-    # Trivia
     if low == "/trivia":
-        seed = random.randint(0, 1_000_000)
         raw = generate_llm(
-            f"Generate a unique trivia Q&A in JSON based on seed {seed}:\n"
+            f"Generate a unique trivia Q&A in JSON:\n"
             '{"question":"...","answer":"..."}'
         )
         try:
             obj = json.loads(raw.strip().strip("```json").strip("```"))
             trivia_store[chat_id] = obj
-            client.send_message(mid, chat_id,
-                f"❓ {obj['question']}\nReply `/answer` or `/answer your guess`.")
+            client.send_message(mid, chat_id, f"❓ {obj['question']}\nReply `/answer` or `/answer your guess`.")
         except:
             client.send_message(mid, chat_id, "Failed to generate trivia.")
         return {"success": True}
 
     if low.startswith("/answer"):
-        resp = body[len("/answer"):].strip()
+        resp = body[7:].strip()
         if chat_id in trivia_store:
             qa = trivia_store.pop(chat_id)
             if resp:
@@ -287,22 +326,18 @@ async def whatsapp_webhook(request: Request):
             client.send_message(mid, chat_id, "No active trivia. `/trivia` to start.")
         return {"success": True}
 
-    # Meme
     if low.startswith("/meme "):
-        txt = body[len("/meme "):].strip()
         client.send_message(mid, chat_id, "🎨 Generating your meme...")
-        task_queue.put({"type":"image","message_id":mid,"chat_id":chat_id,"prompt":f"meme: {txt}"})
+        task_queue.put({"type":"image","message_id":mid,"chat_id":chat_id,"prompt":f"meme: {body[6:].strip()}"})
         return {"success": True}
 
-    # Polls
     if low.startswith("/poll "):
-        parts = body[len("/poll "):].split("|")
+        parts = body[6:].split("|")
         if len(parts) < 3:
             client.send_message(mid, chat_id, "Use `/poll Q|A|B`")
         else:
             q, *opts = [p.strip() for p in parts]
-            votes = {i+1:0 for i in range(len(opts))}
-            polls[chat_id] = {"question":q,"options":opts,"votes":votes,"voters":{}}
+            polls[chat_id] = {"question":q,"options":opts,"votes":{i+1:0 for i in range(len(opts))},"voters":{}}
             text = f"📊 *Poll:* {q}\n" + "\n".join(f"{i+1}. {o}" for i,o in enumerate(opts))
             client.send_message(mid, chat_id, text)
         return {"success": True}
@@ -322,7 +357,9 @@ async def whatsapp_webhook(request: Request):
     if low == "/results":
         if chat_id in polls:
             p = polls[chat_id]
-            text = f"📊 *Results:* {p['question']}\n" + "\n".join(f"{i}. {o}: {p['votes'][i]}" for i,o in enumerate(p["options"],1))
+            text = f"📊 *Results:* {p['question']}\n" + "\n".join(
+                f"{i}. {o}: {p['votes'][i]}" for i,o in enumerate(p["options"],1)
+            )
             client.send_message(mid, chat_id, text)
         else:
             client.send_message(mid, chat_id, "No active poll.")
@@ -331,23 +368,31 @@ async def whatsapp_webhook(request: Request):
     if low == "/endpoll":
         if chat_id in polls:
             p = polls.pop(chat_id)
-            text = f"📊 *Final Results:* {p['question']}\n" + "\n".join(f"{i}. {o}: {p['votes'][i]}" for i,o in enumerate(p["options"],1))
+            text = f"📊 *Final Results:* {p['question']}\n" + "\n".join(
+                f"{i}. {o}: {p['votes'][i]}" for i,o in enumerate(p["options"],1)
+            )
             client.send_message(mid, chat_id, text)
         else:
             client.send_message(mid, chat_id, "No active poll.")
         return {"success": True}
 
-    # Image gen
     if low.startswith("/gen"):
-        prompt = body[len("/gen"):].strip()
+        parts = body[4:].split("|", 1)
+        prompt = parts[0].strip()
+        count = BotConfig.DEFAULT_IMAGE_COUNT
+        if len(parts) > 1 and parts[1].strip().isdigit():
+            count = int(parts[1].strip())
         if not prompt:
-            client.send_message(mid, chat_id, "Use `/gen <prompt>`")
+            client.send_message(mid, chat_id, "Use `/gen <prompt>|<count>`")
         else:
-            client.send_message(mid, chat_id, "✨ Generating image...")
-            task_queue.put({"type":"image","message_id":mid,"chat_id":chat_id,"prompt":prompt})
+            client.send_message(mid, chat_id, f"✨ Generating {count} image(s)...")
+            task_queue.put({
+                "type":"image","message_id":mid,"chat_id":chat_id,
+                "prompt":prompt,"num_images":count
+            })
         return {"success": True}
 
-    # Fallback: voice reply
+    # Fallback → voice reply
     task_queue.put({"type":"audio","message_id":mid,"chat_id":chat_id,"prompt":body})
     return {"success": True}