From b506aad02de8040fd3fd0f3abaadc40797b40f98 Mon Sep 17 00:00:00 2001
From: thedragonsinn <98635854+thedragonsinn@users.noreply.github.com>
Date: Mon, 10 Mar 2025 18:24:29 +0530
Subject: [PATCH] adjust according to core

---
 README.md                               | 2 +-
 app/__init__.py                         | 2 --
 app/plugins/admin/fbans.py              | 2 +-
 app/plugins/ai/models.py                | 3 +--
 app/plugins/sudo/commands.py            | 2 +-
 app/plugins/sudo/users.py               | 4 ++--
 app/plugins/tg_tools/pm_n_tag_logger.py | 2 +-
 app/plugins/tg_tools/pm_permit.py       | 4 ++--
 8 files changed, 9 insertions(+), 12 deletions(-)

diff --git a/README.md b/README.md
index 4b25457..57cc18f 100644
--- a/README.md
+++ b/README.md
@@ -42,7 +42,7 @@ async def test_function(bot: BOT, message: Message):
 ```python
 from app import BOT, bot, Message, CustomDB
 
-TEST_COLLECTION = CustomDB("TEST_COLLECTION")
+TEST_COLLECTION = CustomDB["TEST_COLLECTION"]
 
 @bot.add_cmd(cmd="add_data")
 async def test_function(bot: BOT, message: Message):
diff --git a/app/__init__.py b/app/__init__.py
index 366199d..f2878f5 100644
--- a/app/__init__.py
+++ b/app/__init__.py
@@ -1,7 +1,5 @@
 from ub_core import (
     BOT,
-    DATABASE,
-    DATABASE_CLIENT,
     LOGGER,
     Config,
     Convo,
diff --git a/app/plugins/admin/fbans.py b/app/plugins/admin/fbans.py
index fc5ab9d..fe4be31 100644
--- a/app/plugins/admin/fbans.py
+++ b/app/plugins/admin/fbans.py
@@ -10,7 +10,7 @@ from app import BOT, Config, CustomDB, Message, bot, extra_config
 
 FBAN_TASK_LOCK = asyncio.Lock()
 
-FED_DB = CustomDB("FED_LIST")
+FED_DB = CustomDB["FED_LIST"]
 
 BASIC_FILTER = filters.user([609517172, 2059887769]) & ~filters.service
 
diff --git a/app/plugins/ai/models.py b/app/plugins/ai/models.py
index c7d10ec..f8b2e8f 100644
--- a/app/plugins/ai/models.py
+++ b/app/plugins/ai/models.py
@@ -9,13 +9,12 @@ from google.genai.types import (
     SafetySetting,
     Tool,
 )
-from pyrogram import filters
 
 from app import BOT, CustomDB, Message, extra_config
 
 logging.getLogger("google_genai.models").setLevel(logging.WARNING)
 
-DB_SETTINGS = CustomDB("COMMON_SETTINGS")
+DB_SETTINGS = CustomDB["COMMON_SETTINGS"]
 
 try:
     client: Client = Client(api_key=extra_config.GEMINI_API_KEY)
diff --git a/app/plugins/sudo/commands.py b/app/plugins/sudo/commands.py
index fcb22da..677136f 100644
--- a/app/plugins/sudo/commands.py
+++ b/app/plugins/sudo/commands.py
@@ -1,6 +1,6 @@
 from app import BOT, Config, CustomDB, Message
 
-DB = CustomDB("SUDO_CMD_LIST")
+DB = CustomDB["SUDO_CMD_LIST"]
 
 
 async def init_task():
diff --git a/app/plugins/sudo/users.py b/app/plugins/sudo/users.py
index 53decc2..13c817c 100644
--- a/app/plugins/sudo/users.py
+++ b/app/plugins/sudo/users.py
@@ -3,8 +3,8 @@ from ub_core.utils.helpers import extract_user_data, get_name
 
 from app import BOT, Config, CustomDB, Message
 
-SUDO = CustomDB("COMMON_SETTINGS")
-SUDO_USERS = CustomDB("SUDO_USERS")
+SUDO = CustomDB["COMMON_SETTINGS"]
+SUDO_USERS = CustomDB["SUDO_USERS"]
 
 
 async def init_task():
diff --git a/app/plugins/tg_tools/pm_n_tag_logger.py b/app/plugins/tg_tools/pm_n_tag_logger.py
index 1e0e298..8098166 100644
--- a/app/plugins/tg_tools/pm_n_tag_logger.py
+++ b/app/plugins/tg_tools/pm_n_tag_logger.py
@@ -8,7 +8,7 @@ from ub_core.utils.helpers import get_name
 
 from app import BOT, Config, CustomDB, Message, bot, extra_config
 
-LOGGER = CustomDB("COMMON_SETTINGS")
+LOGGER = CustomDB["COMMON_SETTINGS"]
 
 MESSAGE_CACHE: dict[int, list[Message]] = defaultdict(list)
 FLOOD_LIST: list[int] = []
diff --git a/app/plugins/tg_tools/pm_permit.py b/app/plugins/tg_tools/pm_permit.py
index 9a3dc8a..2781781 100644
--- a/app/plugins/tg_tools/pm_permit.py
+++ b/app/plugins/tg_tools/pm_permit.py
@@ -7,8 +7,8 @@ from ub_core.utils.helpers import get_name
 
 from app import BOT, CustomDB, Message, bot, extra_config
 
-PM_USERS = CustomDB("PM_USERS")
-PM_GUARD = CustomDB("COMMON_SETTINGS")
+PM_USERS = CustomDB["PM_USERS"]
+PM_GUARD = CustomDB["COMMON_SETTINGS"]
 
 ALLOWED_USERS: list[int] = []
 RECENT_USERS: dict = defaultdict(int)