Upload database.py

Akeno/utils/database.py

@@ -4,8 +4,9 @@ import time
 from motor import motor_asyncio
 from motor.core import AgnosticClient
 
-from config import MONGO_URL
 from Akeno.utils.logger import LOGS
+from config import MONGO_URL
+
 
 class Database:
     def __init__(self, uri: str) -> None:
@@ -22,6 +23,7 @@ class Database:
         self.forcesub = self.db["forcesub"]
         self.gachabots = self.db["gachabots"]
         self.cohere = self.db["cohere"]
+        self.chatbot = self.db["chatbot"]
         self.antiarabic = self.db["antiarabic"]
         self.gban = self.db["gban"]
         self.gmute = self.db["gmute"]
@@ -618,4 +620,19 @@ class Database:
             return user_data.get("arabic", False)
         return False
 
+    async def add_chatbot(self, chat_id, user_id):
+        await self.chatbot.update_one(
+            {"chat_id": chat_id},
+            {"$set": {"user_id": user_id}},
+            upsert=True
+        )
+
+    async def get_chatbot(self, chat_id):
+        user_data = await self.chatbot.find_one({"chat_id": chat_id})
+        return user_data.get("user_id") if user_data else None
+
+    async def remove_chatbot(self, chat_id):
+        unset_data = {"user_id": None}
+        return await self.chatbot.update_one({"chat_id": chat_id}, {"$unset": unset_data})
+
 db = Database(MONGO_URL)