Upload chat.py

chatbot/plugins/chat.py

@@ -51,6 +51,7 @@ async def geni_files_delete(name: str):
     return response.text
 
 spam_chats = []
+running_tasks = {}
 
 DISABLE_COMMAND = [
     "start",
@@ -82,10 +83,18 @@ I am ready to be a gemini bot developer
 @Client.on_callback_query(filters.regex("^cancels"))
 async def cancel_(client, callback_query):
     try:
-        spam_chats.clear()
-        return await callback_query.edit_message_text("ok cancel")
+        chat_id = callback_query.message.chat.id
+        if chat_id in running_tasks:
+            running_tasks[chat_id].cancel()
+            del running_tasks[chat_id]
+            await callback_query.edit_message_text("✅ Processing canceled.")
+        else:
+            await callback_query.edit_message_text("⚠️ No active task to cancel.")
+    except asyncio.CancelledError:
+        await callback_query.edit_message_text("⚠️ Task already canceled.")
     except Exception as e:
-        await callback_query.answer(f"Nothing cancel: {e}")
+        await callback_query.answer(f"Nothing to cancel: {e}")
+
 
 @Client.on_message(
     ~filters.scheduled
@@ -252,12 +261,6 @@ async def chatbot_talk(client: Client, message: Message):
             return await ai_reply.edit_text(f"Error: {e}")
         except Exception as e:
             return await ai_reply.edit_text(f"Error: {e}")
-
-    if message.document:
-        is_check_plan = await db.is_gemini_plan(user_id=message.from_user.id)
-        if not is_check_plan:
-            return await message.reply_text("Only plan gemini")
-        return await message.reply_text("maintenance soon")
     if message.audio or message.voice:
         await client.send_chat_action(message.chat.id, enums.ChatAction.UPLOAD_AUDIO)
         await asyncio.sleep(1.5)
@@ -321,76 +324,81 @@ async def chatbot_talk(client: Client, message: Message):
             return await ai_reply.edit_text(f"Error: {e}")
 
     if message.video:
-        spam_chats.append(message.chat.id)
-        await client.send_chat_action(message.chat.id, enums.ChatAction.UPLOAD_VIDEO)
-        await asyncio.sleep(1.5)
-        buttons = [
-            [
-                InlineKeyboardButton(
-                    text="Cancel",
-                    callback_data="cancels"
-                ),
-            ],
-        ]
-        if client.me.is_premium:
-            ai_reply = await message.reply_text(f"{custom_loading}Processing...")
-        else:
-            ai_reply = await message.reply_text(
-                f"Processing...\n\nYou can use stop everything",
-                reply_markup=InlineKeyboardMarkup(buttons)
-            )
-        video_file_name = await message.download(file_name="newvideo.mp4")
-        caption = message.caption or "What's this?"
-        model = genai.GenerativeModel(
-            model_name=model_,
-            safety_settings={
-                genai.types.HarmCategory.HARM_CATEGORY_HATE_SPEECH: genai.types.HarmBlockThreshold.BLOCK_NONE,
-                genai.types.HarmCategory.HARM_CATEGORY_HARASSMENT: genai.types.HarmBlockThreshold.BLOCK_NONE,
-                genai.types.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: genai.types.HarmBlockThreshold.BLOCK_NONE,
-                genai.types.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: genai.types.HarmBlockThreshold.BLOCK_NONE,
-            }
-        )
-        backup_chat = await db._get_chatbot_chat_from_db(message.from_user.id)
-        backup_chat.append({"role": "user", "parts": [{"text": caption}]})
-        if client.me.is_premium:
-            await ai_reply.edit_text(f"{custom_loading}Uploading file..")
-        else:
-            await ai_reply.edit_text("Uploading file..")
-        video_file = genai.upload_file(path=video_file_name)
-        while video_file.state.name == "PROCESSING":
-            await asyncio.sleep(10)
-            video_file = genai.get_file(video_file.name)
-        if video_file.state.name == "FAILED":
-            return await ai_reply.edit_text(f"Error: {video_file.state.name}")
-        try:
-            await client.send_chat_action(message.chat.id, enums.ChatAction.TYPING)
-            await asyncio.sleep(1.5)
-            response = model.generate_content(
-                [video_file, caption],
-                request_options={"timeout": 600}
+        chat_id = message.chat.id
+        async def task_video():
+            try:
+                spam_chats.append()
+                await client.send_chat_action(message.chat.id, enums.ChatAction.UPLOAD_VIDEO)
+                await asyncio.sleep(1.5)
+                buttons = [
+                    [
+                        InlineKeyboardButton(
+                            text="Cancel",
+                            callback_data="cancels"
+                        ),
+                    ],
+                ]
+                ai_reply = await message.reply_text(
+                    f"Processing...\n\nYou can use stop everything",
+                    reply_markup=InlineKeyboardMarkup(buttons)
             )
-            if len(response.text) > 4096:
-                with open("chat.txt", "w+", encoding="utf8") as out_file:
-                    out_file.write(response.text)
-                await message.reply_document(
-                    document="chat.txt",
-                    disable_notification=True
+                await asyncio.sleep(5)
+                video_file_name = await message.download(file_name="newvideo.mp4")
+                caption = message.caption or "What's this?"
+                model = genai.GenerativeModel(
+                    model_name=model_,
+                    safety_settings={
+                        genai.types.HarmCategory.HARM_CATEGORY_HATE_SPEECH: genai.types.HarmBlockThreshold.BLOCK_NONE,
+                        genai.types.HarmCategory.HARM_CATEGORY_HARASSMENT: genai.types.HarmBlockThreshold.BLOCK_NONE,
+                        genai.types.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: genai.types.HarmBlockThreshold.BLOCK_NONE,
+                        genai.types.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: genai.types.HarmBlockThreshold.BLOCK_NONE,
+                    }
                 )
-                await ai_reply.delete()
-                os.remove("chat.txt")
-            else:
-                await ai_reply.edit_text(response.text)
-            backup_chat.append({"role": "model", "parts": [{"text": response.text}]})
-            await db._update_chatbot_chat_in_db(message.from_user.id, backup_chat)
-            await client.send_chat_action(message.chat.id, enums.ChatAction.CANCEL)
-            video_file.delete()
-            os.remove(video_file_name)
-            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}")
-
+                backup_chat = await db._get_chatbot_chat_from_db(message.from_user.id)
+                backup_chat.append({"role": "user", "parts": [{"text": caption}]})
+                await ai_reply.edit_text("Uploading file..")
+                video_file = genai.upload_file(path=video_file_name)
+                while video_file.state.name == "PROCESSING":
+                    await asyncio.sleep(10)
+                    video_file = genai.get_file(video_file.name)
+                if video_file.state.name == "FAILED":
+                    return await ai_reply.edit_text(f"Error: {video_file.state.name}")
+                await client.send_chat_action(message.chat.id, enums.ChatAction.TYPING)
+                await asyncio.sleep(1.5)
+                response = model.generate_content(
+                    [video_file, caption],
+                    request_options={"timeout": 600}
+                )
+                if len(response.text) > 4096:
+                    with open("chat.txt", "w+", encoding="utf8") as out_file:
+                        out_file.write(response.text)
+                    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.text)
+                backup_chat.append({"role": "model", "parts": [{"text": response.text}]})
+                await db._update_chatbot_chat_in_db(message.from_user.id, backup_chat)
+                await client.send_chat_action(message.chat.id, enums.ChatAction.CANCEL)
+                video_file.delete()
+                os.remove(video_file_name)
+                return
+            except asyncio.CancelledError:
+                await ai_reply.edit_text("✅ Video processed successfully!")
+            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}")
+            finally:
+                if chat_id in spam_chats:
+                    spam_chats.remove(chat_id)
+                if chat_id in running_tasks:
+                    del running_tasks[chat_id]
+        task = asyncio.create_task(task_video())
+        running_tasks[chat_id] = task
     if message.text:
         await client.send_chat_action(message.chat.id, enums.ChatAction.TYPING)
         await asyncio.sleep(1.5)
@@ -428,4 +436,5 @@ async def chatbot_talk(client: Client, message: Message):
             await client.send_chat_action(message.chat.id, enums.ChatAction.CANCEL)
             return
         except Exception as e:
-            return await message.reply_text(f"Error: {e}")
+            return await message.reply_text(f"Error: {e}")
+