Update chatbot/plugins/chat.py

chatbot/plugins/chat.py

@@ -659,6 +659,7 @@ async def chatbot_talk(client: Client, message: Message):
         await asyncio.sleep(1.5)
         query = message.text.strip()
         query_base = query
+        file_path = "gemini-native-image.png"
         try:
             if query_base in ["/", "/help"]:
                 await client.send_chat_action(message.chat.id, enums.ChatAction.CANCEL)
@@ -668,6 +669,42 @@ async def chatbot_talk(client: Client, message: Message):
                 await message.reply_text(CREDITS_DEVELOPER)
                 return
 
+            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:
+                            LOGS.info(f"Gemini Image: {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,
+                        reply_markup=keyboard_like
+                    )
+                    backup_chat.append({"role": "model", "parts": [{"text": "Nothing test image"}]})
+                    await db._update_chatbot_chat_in_db(message.from_user.id, backup_chat)
+                    await client.send_chat_action(message.chat.id, enums.ChatAction.CANCEL)
+                    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(
                     {"user_id": message.from_user.id},
@@ -726,7 +763,7 @@ async def chatbot_talk(client: Client, message: Message):
                     for part in text_parts:
                         await message.reply_text(part)
                 else:
-                    keyboard_like = create_keyboard()
+                    keyboard_like = create_keyboard(user_id=message.from_user.id)
                     await message.reply_text(
                         response.text,
                         disable_web_page_preview=True,