Upload chat.py

chatbot/plugins/chat.py

@@ -51,7 +51,6 @@ async def geni_files_delete(name: str):
     return response.text
 
 spam_chats = []
-running_tasks = {}
 
 DISABLE_COMMAND = [
     "start",
@@ -80,21 +79,81 @@ I am ready to be a gemini bot developer
 - Command: /offchat (pm or group)
 """
 
+class TaskManager:
+    def __init__(self):
+        self.running_tasks = {}
+
+    async def add_task(self, chat_id, coro):
+        task = asyncio.create_task(coro)
+        self.running_tasks[chat_id] = task
+        try:
+            await task
+        finally:
+            self.running_tasks.pop(chat_id, None)
+
+    async def cancel_task(self, chat_id):
+        task = self.running_tasks.get(chat_id)
+        if task:
+            task.cancel()
+            del self.running_tasks[chat_id]
+            return True
+        return False
+
+    def list_tasks(self):
+        return list(self.running_tasks.keys())
+
+task_manager = TaskManager()
+
 @Client.on_callback_query(filters.regex("^cancels"))
 async def cancel_(client, callback_query):
-    try:
-        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 to cancel: {e}")
+    chat_id = callback_query.message.chat.id
+    if await task_manager.cancel_task(chat_id):
+        spam_chats.remove(chat_id)
+        await callback_query.edit_message_text("✅ Processing canceled.")
+    else:
+        await callback_query.edit_message_text("⚠️ No active task to cancel.")
 
+async def handle_video(client, message, model_):
+    chat_id = message.chat.id
+    async def process_video():
+        buttons = [
+            [
+                InlineKeyboardButton(
+                    text="Cancel",
+                    callback_data="cancels"
+                )
+            ],
+        ]
+        try:
+            spam_chats.append(chat_id)
+            video_file_name = await message.download(file_name="newvideo.mp4")
+            caption = message.caption or "What's this?"
+            model = genai.GenerativeModel(model_name=model_)
+
+            ai_reply = await message.reply_text(
+                "Uploading file..",
+                reply_markup=InlineKeyboardMarkup(buttons)
+            )
+            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}")
+
+            response = model.generate_content(
+                [video_file, caption],
+                request_options={"timeout": 600}
+            )
+            await ai_reply.edit_text(response.text)
+        except asyncio.CancelledError:
+            await ai_reply.edit_text("⚠️ Video processing was canceled.")
+        finally:
+            spam_chats.remove(chat_id)
+            task_manager.cancel_task(chat_id)
+
+    await task_manager.add_task(chat_id, process_video())
 
 @Client.on_message(
     ~filters.scheduled
@@ -324,81 +383,7 @@ async def chatbot_talk(client: Client, message: Message):
             return await ai_reply.edit_text(f"Error: {e}")
 
     if message.video:
-        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)
-            )
-                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,
-                    }
-                )
-                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
+        await handle_video(client, message, model_)
     if message.text:
         await client.send_chat_action(message.chat.id, enums.ChatAction.TYPING)
         await asyncio.sleep(1.5)