Ok test

chatbot/plugins/chat.py

@@ -181,30 +181,6 @@ Type anything to begin.
 
 gen = genai.Client(api_key=GOOGLE_API_KEY)
 
-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()
 
 async def progress(current, total, message, start, type_of_ps, file_name=None):
     """Progress Bar For Showing Progress While Uploading / Downloading File - Normal"""
@@ -1113,6 +1089,29 @@ async def chatbot_talk(client: Client, message: Message):
             edit_image_now = re.match(r"--edit-images\s*(.*)", caption, re.IGNORECASE)
             if edit_image_now:
                 try:
+                    user_data = await db.user_premium.find_one({"user_id": message.from_user.id}) or {
+                        "user_id": message.from_user.id,
+                        "credits_used": 0,
+                        "last_reset": dt.now(),
+                        "is_premium": False
+                    }
+                    if dt.now() - user_data["last_reset"] > timedelta(days=1):
+                        user_data["credits_used"] = 0
+                        user_data["last_reset"] = dt.now()
+                        
+                    credit_limit = 50 if user_data.get("is_premium") else 5
+                    if user_data["credits_used"] >= credit_limit:
+                        time_left = (user_data["last_reset"] + timedelta(days=1) - dt.now())
+                        hours = int(time_left.total_seconds() // 3600)
+                        return await message.reply_text(
+                            f"🚫 Credit Limit Reached\n\n"
+                            f"• Used: {user_data['credits_used']}/{credit_limit} images today\n"
+                            f"• Resets in: {hours} hours\n\n"
+                            f"💎 Upgrade to Premium for 50 daily credits!",
+                            reply_markup=InlineKeyboardMarkup([[
+                                InlineKeyboardButton("Get Premium", callback_data="get_premium")
+                            ]])
+                        )
                     backup_chat = await db._get_chatbot_chat_from_db(message.from_user.id)
                     backup_chat.append({"role": "user", "parts": [{"text": caption}]})
                     images = Image.open(await message.download())
@@ -1142,6 +1141,15 @@ async def chatbot_talk(client: Client, message: Message):
                     )
                     backup_chat.append({"role": "model", "parts": [{"text": captions}]})
                     await db._update_chatbot_chat_in_db(message.from_user.id, backup_chat)
+                    await db.user_premium.update_one(
+                        {"user_id": user_id},
+                        {
+                            "$inc": {"credits_used": 1},
+                            "$set": {"last_reset": user_data["last_reset"]},
+                            "$setOnInsert": {"is_premium": False}
+                        },
+                        upsert=True
+                    )
                     await client.send_chat_action(message.chat.id, enums.ChatAction.CANCEL)
                     return
                 except Exception as e: