Update chatbot/plugins/chat.py

chatbot/plugins/chat.py

@@ -261,6 +261,62 @@ async def deletemydata(client, callback):
     else:
         await callback.answer(delm, True)
 
+@Client.on_callback_query(filters.regex("^genprompt_(\d+)$"))
+async def geminigen_prompt(client, callback):
+    user_id = int(callback.matches[0].group(1))
+    captions = None
+    file_path = "gemini-native-image.png"
+    try:
+        backup_chat = await db._get_chatbot_chat_from_db(user_id)
+        data = await db.backup_chatbot.find_one({"user_id": user_id})
+        if not data:
+            return await callback.answer("Can't found user", True)
+        get_response = data.get("prompt_image", None)
+        if not get_response:
+            return await callback.answer("Server busy try again later", True)
+        await callback.message.delete()
+
+        backup_chat.append({"role": "user", "parts": [{"text": get_response}]})
+        response = await gen.aio.models.generate_content(
+            model="gemini-2.0-flash-exp-image-generation",
+            contents=get_response,
+            config=ty.GenerateContentConfig(
+                response_modalities=['TEXT', 'IMAGE']
+            )
+        )
+
+        for part in response.candidates[0].content.parts:
+            if part.text is not None:
+                captions += part.text
+            elif part.inline_data is not None:
+                image = Image.open(BytesIO((part.inline_data.data)))
+                image.save(file_path)
+
+        keyboard_like = create_keyboard(user_id=user_id)
+        await callback.message.reply_photo(
+            file_path,
+            caption=captions,
+            reply_markup=keyboard_like
+        )
+        await db.backup_chatbot.update_one(
+            {"user_id": user_id},
+            {"$set": {"translate_text": captions}},
+            upsert=True
+        )
+        backup_chat.append({"role": "model", "parts": [{"text": captions}]})
+        await db._update_chatbot_chat_in_db(user_id, backup_chat)
+        await client.send_chat_action(callback.message.chat.id, enums.ChatAction.CANCEL)
+        return
+    except Exception as e:
+        LOGS.error(f"geminigen_prompt failed: {str(e)}")
+        await callback.message.reply_text("Server busy try again later")
+    finally:
+        if file_path:
+            try:
+                os.remove(file_path)
+            except:
+                pass
+
 @Client.on_callback_query(filters.regex("^fluxprompt_(\d+)$"))
 async def flux_prompt(client, callback):
     user_id = int(callback.matches[0].group(1))
@@ -268,13 +324,14 @@ async def flux_prompt(client, callback):
     file_photo = None
     try:
         backup_chat = await db._get_chatbot_chat_from_db(user_id)
-        data = await db.backup_chatbot.find_one({"user_id": int(user_id)})
+        data = await db.backup_chatbot.find_one({"user_id": user_id})
         if not data:
             return await callback.answer("Can't found user", True)
-        get_response = data.get("prompt_flux", None)
+        get_response = data.get("prompt_image", None)
         if not get_response:
             return await callback.answer("Server busy try again later", True)
         await callback.message.delete()
+
         backup_chat.append({"role": "user", "parts": [{"text": get_response}]})
         response_js = await js.image.create(
             "black-forest-labs/flux-1-schnell",
@@ -287,7 +344,7 @@ async def flux_prompt(client, callback):
 
         file_photo = await gen.aio.files.upload(file=file_path)
         while file_photo.state.name == "PROCESSING":
-            await client.send_chat_action(message.chat.id, enums.ChatAction.UPLOAD_PHOTO)
+            await client.send_chat_action(callback.message.chat.id, enums.ChatAction.UPLOAD_PHOTO)
             await asyncio.sleep(10)
             file_photo = await gen.aio.files.get(name=file_photo.name)
         if file_photo.state.name == "FAILED":
@@ -311,9 +368,12 @@ async def flux_prompt(client, callback):
             upsert=True
         )
         backup_chat.append({"role": "model", "parts": [{"text": response.text}]})
+        await db._update_chatbot_chat_in_db(user_id, backup_chat)
+        await client.send_chat_action(callback.message.chat.id, enums.ChatAction.CANCEL)
+        return
     except Exception as e:
         LOGS.error(f"flux_prompt failed: {str(e)}")
-        await callback.message.reply_text("Server error")
+        await callback.message.reply_text("Server busy try again later")
     finally:
         if file_path and file_photo:
             try:
@@ -346,7 +406,7 @@ async def remover_bg(client, callback):
         await callback.answer()
     except Exception as e:
         LOGS.error(f"remover_bg failed: {str(e)}")
-        await callback.answer("Server error", show_alert=True)
+        await callback.answer("Server busy try again later", show_alert=True)
 
 @Client.on_callback_query(filters.regex("^refreshch$"))
 async def reshch(client, callback):
@@ -602,7 +662,7 @@ async def chatbot_talk(client: Client, message: Message):
                             os.remove(file_path)
                         except:
                             pass
-    
+
             if re.findall(r"\b(pro:editimage)\b", caption, re.IGNORECASE):
                 await db.backup_chatbot.update_one(
                     {"user_id": message.from_user.id},
@@ -789,45 +849,40 @@ async def chatbot_talk(client: Client, message: Message):
 
             if re.findall(r"\b(image)\b", query_base, re.IGNORECASE):
                 try:
-                    backup_chat = await db._get_chatbot_chat_from_db(message.from_user.id)
-                    backup_chat.append({"role": "user", "parts": [{"text": query_base}]})
-                    response = await gen.aio.models.generate_content(
-                        model="gemini-2.0-flash-exp-image-generation",
-                        contents=query_base,
-                        config=ty.GenerateContentConfig(
-                            response_modalities=['TEXT', 'IMAGE']
-                        )
-                    )
-                    for part in response.candidates[0].content.parts:
-                        if part.text is not None:
-                            captions += part.text
-                        elif part.inline_data is not None:
-                            image = Image.open(BytesIO((part.inline_data.data)))
-                            image.save(file_path)
-                    keyboard_like = create_keyboard(user_id=message.from_user.id)
-                    await message.reply_photo(
-                        file_path,
-                        caption=captions,
-                        reply_markup=keyboard_like
-                    )
+                    buttons = [
+                        [
+                            InlineKeyboardButton(
+                                text="🎨 Gemini AI Generate Image",
+                                callback_data=f"genprompt_{message.from_user.id}"
+                            )
+                        ],
+                        [
+                            InlineKeyboardButton(
+                                text="🖌️ FLUX AI Generate Image",
+                                callback_data=f"fluxprompt_{message.from_user.id}"
+                            )
+                        ]
+                        [
+                            InlineKeyboardButton(
+                                text="❌ Cancel",
+                                callback_data="closedd"
+                            )
+                        ]
+                    ]
                     await db.backup_chatbot.update_one(
                         {"user_id": message.from_user.id},
-                        {"$set": {"translate_text": captions}},
+                        {"$set": {"prompt_image": query_base}},
                         upsert=True
                     )
-                    backup_chat.append({"role": "model", "parts": [{"text": captions}]})
-                    await db._update_chatbot_chat_in_db(message.from_user.id, backup_chat)
-                    await client.send_chat_action(message.chat.id, enums.ChatAction.CANCEL)
+                    await message.reply_text(
+                        "Are you sure you want to prompt this Image generate?\n\n"
+                        f"Your answer: `{query_base}`",
+                        reply_markup=InlineKeyboardMarkup(buttons)
+                    )
                     return
                 except Exception as e:
                     LOGS.error(f"Error: Gemini Image: {str(e)}")
                     return await message.reply_text("Server busy try again later")
-                finally:
-                    if file_path:
-                        try:
-                            os.remove(file_path)
-                        except:
-                            pass
 
             if re.findall(r"\b(enabled:chatsystem)\b", query_base, re.IGNORECASE):
                 await db.backup_chatbot.update_one(