Update chatbot/plugins/chat.py

chatbot/plugins/chat.py

@@ -92,7 +92,6 @@ blacklist_patterns = [
     r"\bdeath\s*threat\b",
 ]
 
-
 def is_command_disabled(text):
   text_lower = text.lower()
   if text_lower.startswith("/") and not (text_lower.startswith("/setmodel") or text_lower.startswith("/start")):
@@ -1093,14 +1092,26 @@ async def chatbot_talk(client: Client, message: Message):
                         "user_id": message.from_user.id,
                         "credits_used": 0,
                         "last_reset": dt.now(),
-                        "is_premium": False
+                        "is_premium": False,
+                        "premium_expiry" = None
                     }
                     if dt.now() - user_data["last_reset"] > timedelta(days=1):
+                        await db.user_premium.update_one(
+                            {"user_id": message.from_user.id},
+                            {"$set": {"credits_used": 0, "last_reset": dt.now()}}
+                        )
                         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:
+
+                    if user_data.get("is_premium") and user_data["premium_expiry"] < dt.now():
+                        await db.user_premium.update_one(
+                            {"user_id": message.from_user.id},
+                            {"$set": {"is_premium": False}}
+                        )
+                        user_data["is_premium"] = False
+
+                    credit_limit = 50 if user_data["is_premium"] else 5
+                    remaining = credit_limit - user_data["credits_used"]
+                    if remaining <= 0:
                         time_left = (user_data["last_reset"] + timedelta(days=1) - dt.now())
                         hours = int(time_left.total_seconds() // 3600)
                         return await message.reply_text(
@@ -1112,6 +1123,7 @@ async def chatbot_talk(client: Client, message: Message):
                                 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())