Chandima Prabhath commited on
Commit
a9740d0
Β·
1 Parent(s): 6bc0800

Refactor handle_image_generation function for improved error handling and clarity; update system prompt in config.yaml with enhanced command descriptions.

Browse files
Files changed (2) hide show
  1. app.py +21 -24
  2. config.yaml +18 -16
app.py CHANGED
@@ -147,27 +147,24 @@ def handle_image_generation(task):
147
  prompt = task["prompt"]
148
  count = task.get("num_images", BotConfig.DEFAULT_IMAGE_COUNT)
149
 
 
150
  for i in range(1, count + 1):
151
- for attempt in range(1, 4):
152
- try:
153
- img, path, ret_prompt, url = generate_image(
154
- prompt,
155
- message_id,
156
- f"{message_id}_{i}",
157
- BotConfig.IMAGE_DIR
158
- )
159
- if not img:
160
- raise RuntimeError("No image returned")
161
- formatted = "\n\n".join(f"_{p.strip()}_" for p in ret_prompt.split("\n\n") if p.strip())
162
- caption = f"✨ Image {i}/{count}: {url}\n>{chr(8203)} {formatted}"
163
- client.send_media(message_id, chat_id, path, caption, media_type="image")
164
- os.remove(path)
165
- break
166
- except Exception as e:
167
- logging.warning(f"Image {i}/{count} attempt {attempt} failed: {e}")
168
- time.sleep(2 * attempt)
169
- else:
170
- client.send_message(message_id, chat_id, f"😒 Failed to generate image {i}.")
171
 
172
  # --- Startup ---
173
 
@@ -206,7 +203,7 @@ async def whatsapp_webhook(request: Request):
206
  global last_message_time
207
  last_message_time = time.time()
208
 
209
- # --- Auth & basic routing ---
210
  if request.headers.get("Authorization") != f"Bearer {BotConfig.WEBHOOK_AUTH_TOKEN}":
211
  raise HTTPException(403, "Unauthorized")
212
 
@@ -218,19 +215,19 @@ async def whatsapp_webhook(request: Request):
218
  md = data.get("messageData", {})
219
  mid = data["idMessage"]
220
 
221
- # --- Extract text + contextInfo in one go ---
222
  text_data = md.get("textMessageData") or md.get("extendedTextMessageData")
223
  if not text_data:
224
  return {"success": True}
225
  body = text_data.get("textMessage", text_data.get("text", "")).strip()
226
  ctx = text_data.get("contextInfo", {})
227
 
228
- # --- SKIP if this is a quoted/replied message ---
229
  if ctx.get("quotedMessageId") or ctx.get("quotedMessage"):
230
  logging.debug("Skipping quoted/reply message")
231
  return {"success": True}
232
 
233
- # --- SKIP if someone is @-mentioned ---
234
  if ctx.get("mentionedJidList"):
235
  logging.debug("Skipping mention")
236
  return {"success": True}
 
147
  prompt = task["prompt"]
148
  count = task.get("num_images", BotConfig.DEFAULT_IMAGE_COUNT)
149
 
150
+ # simple loop: one generate_image call per image
151
  for i in range(1, count + 1):
152
+ try:
153
+ img, path, ret_prompt, url = generate_image(
154
+ prompt,
155
+ message_id,
156
+ message_id,
157
+ BotConfig.IMAGE_DIR
158
+ )
159
+ if not img:
160
+ raise RuntimeError("generate_image returned no image")
161
+ formatted = "\n\n".join(f"_{p.strip()}_" for p in ret_prompt.split("\n\n") if p.strip())
162
+ caption = f"✨ Image {i}/{count}: {url}\n>{chr(8203)} {formatted}"
163
+ client.send_media(message_id, chat_id, path, caption, media_type="image")
164
+ os.remove(path)
165
+ except Exception as e:
166
+ logging.warning(f"Image {i}/{count} failed: {e}")
167
+ client.send_message(message_id, chat_id, f"😒 Failed to generate image {i}/{count}.")
 
 
 
 
168
 
169
  # --- Startup ---
170
 
 
203
  global last_message_time
204
  last_message_time = time.time()
205
 
206
+ # Auth
207
  if request.headers.get("Authorization") != f"Bearer {BotConfig.WEBHOOK_AUTH_TOKEN}":
208
  raise HTTPException(403, "Unauthorized")
209
 
 
215
  md = data.get("messageData", {})
216
  mid = data["idMessage"]
217
 
218
+ # Extract text + context
219
  text_data = md.get("textMessageData") or md.get("extendedTextMessageData")
220
  if not text_data:
221
  return {"success": True}
222
  body = text_data.get("textMessage", text_data.get("text", "")).strip()
223
  ctx = text_data.get("contextInfo", {})
224
 
225
+ # Skip replies/quotes
226
  if ctx.get("quotedMessageId") or ctx.get("quotedMessage"):
227
  logging.debug("Skipping quoted/reply message")
228
  return {"success": True}
229
 
230
+ # Skip @-mentions
231
  if ctx.get("mentionedJidList"):
232
  logging.debug("Skipping mention")
233
  return {"success": True}
config.yaml CHANGED
@@ -2,22 +2,24 @@ config:
2
  llm:
3
  model: koboldcpp/HF_SPACE_Tiefighter-13B
4
  system_prompt: |-
5
- You are {char}, a sweet and helpful AI assistant in Telegram and WhatsApp.
6
- You generate images and voice, text replies, and support these commands:
7
- β€’ /help β€” list all commands
8
- β€’ /gen <prompt> β€” generate an image
9
- β€’ /summarize <text>
10
- β€’ /translate <lang>|<text>
11
- β€’ /joke
12
- β€’ /weather <location>
13
- β€’ /inspire
14
- β€’ /trivia
15
- β€’ /answer
16
- β€’ /meme <text>
17
- β€’ /poll <Q>|<opt1>|<opt2>|…
18
- β€’ /results
19
- β€’ /endpoll
20
- Use a concise, friendly tone. If a command is malformed, gently ask the user to correct it. For any other message, respond via voice or text.
 
 
21
  char: Eve
22
 
23
  SD:
 
2
  llm:
3
  model: koboldcpp/HF_SPACE_Tiefighter-13B
4
  system_prompt: |-
5
+ You are {char}, a sweet and helpful AI assistant in Telegram and WhatsApp.
6
+ You generate images, voice and text replies, and support these commands:
7
+ β€’ /help β€” list all commands
8
+ β€’ /gen <prompt>|<count> β€” generate <count> images (default 4)
9
+ β€’ /summarize <text> β€” get a concise summary
10
+ β€’ /translate <lang>|<text> β€” translate text
11
+ β€’ /joke β€” tell a short joke
12
+ β€’ /weather <location> β€” short, creative weather report in Β°C
13
+ β€’ /weatherpoem <location> β€” poetic weather summary in Β°C
14
+ β€’ /inspire β€” inspirational quote
15
+ β€’ /trivia β€” start a trivia question
16
+ β€’ /answer β€” answer the current trivia
17
+ β€’ /meme <text> β€” generate a meme image
18
+ β€’ /poll <Q>|<opt1>|<opt2>|… β€” create a poll
19
+ β€’ /results β€” show poll results
20
+ β€’ /endpoll β€” end the poll
21
+ Use a concise, friendly tone. If a command is malformed, gently ask the user to correct it.
22
+ For any other message, respond via voice or text.
23
  char: Eve
24
 
25
  SD: