Upload chat.py

chatbot/plugins/chat.py

@@ -113,6 +113,55 @@ async def cancel_(client, callback_query):
     else:
         await callback_query.edit_message_text("⚠️ No active task to cancel.")
 
+async def handle_photo(client, message):
+    chat_id = message.chat.id
+    async def process_photo():
+        buttons = [
+            [
+                InlineKeyboardButton(
+                    text="Cancel",
+                    callback_data="cancels"
+                )
+            ],
+        ]
+        try:
+            spam_chats.append(chat_id)
+            file_path = await message.download()
+            caption = message.caption or "What's this?"
+            x = GeminiLatest(api_keys=GOOGLE_API_KEY)
+            await client.send_chat_action(message.chat.id, enums.ChatAction.UPLOAD_PHOTO)
+            await asyncio.sleep(1.5)
+            ai_reply = await message.reply_text(
+                "Uploading file..",
+                reply_markup=InlineKeyboardMarkup(buttons)
+            )
+            backup_chat = await db._get_chatbot_chat_from_db(message.from_user.id)
+            backup_chat.append({"role": "user", "parts": [{"text": caption}]})
+            response_reads = x.get_response_image(caption, file_path)
+            if len(response_reads) > 4096:
+                with open("chat.txt", "w+", encoding="utf8") as out_file:
+                    out_file.write(response_reads)
+                await message.reply_document(
+                    document="chat.txt",
+                    disable_notification=True
+                )
+                await ai_reply.delete()
+                os.remove("chat.txt")
+            else:
+                await ai_reply.edit_text(response_reads)
+            backup_chat.append({"role": "model", "parts": [{"text": response_reads}]})
+            await db._update_chatbot_chat_in_db(message.from_user.id, backup_chat)
+            await client.send_chat_action(message.chat.id, enums.ChatAction.CANCEL)
+            os.remove(file_path)
+            return
+        except asyncio.CancelledError:
+            await ai_reply.edit_text("⚠️ Photo processing was canceled.")
+        finally:
+            spam_chats.remove(chat_id)
+            task_manager.cancel_task(chat_id)
+
+    await task_manager.add_task(chat_id, process_photo())
+
 async def handle_video(client, message, model_):
     chat_id = message.chat.id
     async def process_video():
@@ -129,7 +178,8 @@ async def handle_video(client, message, model_):
             video_file_name = await message.download(file_name="newvideo.mp4")
             caption = message.caption or "What's this?"
             model = genai.GenerativeModel(model_name=model_)
-
+            await client.send_chat_action(message.chat.id, enums.ChatAction.UPLOAD_PHOTO)
+            await asyncio.sleep(1.5)
             ai_reply = await message.reply_text(
                 "Uploading file..",
                 reply_markup=InlineKeyboardMarkup(buttons)
@@ -285,41 +335,7 @@ async def chatbot_talk(client: Client, message: Message):
         if message.reply_to_message.from_user.id != client.me.id:
             return
     if message.photo:
-        await client.send_chat_action(message.chat.id, enums.ChatAction.UPLOAD_PHOTO)
-        await asyncio.sleep(1.5)
-        file_path = await message.download()
-        caption = message.caption or "What's this?"
-        x = GeminiLatest(api_keys=GOOGLE_API_KEY)
-        if client.me.is_premium:
-            ai_reply = await message.reply_text(f"{custom_loading}Processing...")
-        else:
-            ai_reply = await message.reply_text(f"Processing...")
-        try:
-            await client.send_chat_action(message.chat.id, enums.ChatAction.TYPING)
-            await asyncio.sleep(1.5)
-            backup_chat = await db._get_chatbot_chat_from_db(message.from_user.id)
-            backup_chat.append({"role": "user", "parts": [{"text": caption}]})
-            response_reads = x.get_response_image(caption, file_path)
-            if len(response_reads) > 4096:
-                with open("chat.txt", "w+", encoding="utf8") as out_file:
-                    out_file.write(response_reads)
-                await message.reply_document(
-                    document="chat.txt",
-                    disable_notification=True
-                )
-                await ai_reply.delete()
-                os.remove("chat.txt")
-            else:
-                await ai_reply.edit_text(response_reads)
-            backup_chat.append({"role": "model", "parts": [{"text": response_reads}]})
-            await db._update_chatbot_chat_in_db(message.from_user.id, backup_chat)
-            await client.send_chat_action(message.chat.id, enums.ChatAction.CANCEL)
-            os.remove(file_path)
-            return
-        except InvalidArgument as e:
-            return await ai_reply.edit_text(f"Error: {e}")
-        except Exception as e:
-            return await ai_reply.edit_text(f"Error: {e}")
+        await handle_photo(client, message)
     if message.audio or message.voice:
         await client.send_chat_action(message.chat.id, enums.ChatAction.UPLOAD_AUDIO)
         await asyncio.sleep(1.5)