Update chatbot/plugins/chat.py

chatbot/plugins/chat.py

@@ -35,6 +35,7 @@ from pyrogram.errors import *
 from RyuzakiLib import FaceAI, FullStackDev, GeminiLatest, RendyDevChat
 from config import *
 from database import db
+from database import users_collection
 from logger import LOGS
 
 import google.generativeai as genai
@@ -148,14 +149,6 @@ async def userstatus(client: Client, message: Message):
         )
     )
 
-@Client.on_message(
-    ~filters.scheduled
-    & filters.command(["testpay"])
-    & ~filters.forwarded
-)
-async def test_payment(client: Client, message: Message):
-    await db.add_gemini_plan(user_id=message.from_user.id)
-    return await message.reply_text("Successfully payment")
 
 @Client.on_message(
     ~filters.scheduled
@@ -205,6 +198,8 @@ async def rmchatbot_user(client: Client, message: Message):
 async def chatbot_talk(client: Client, message: Message):
     if message.from_user.id in NOT_ALLOWED_NON_PROGRAMMER:
         return
+    user = await users_collection.find_one({"user_id": message.from_user.id})
+    model_ = user.get("model") if user else ""
     genai.configure(api_key=GOOGLE_API_KEY)
     chat_user = await db.get_chatbot(message.chat.id)
     if not chat_user:
@@ -267,7 +262,7 @@ async def chatbot_talk(client: Client, message: Message):
             audio_file_name = await message.download()
         caption = message.caption or "What's this?"
         model = genai.GenerativeModel(
-            model_name="gemini-1.5-flash",
+            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,
@@ -326,7 +321,7 @@ async def chatbot_talk(client: Client, message: Message):
         video_file_name = await message.download(file_name="newvideo.mp4")
         caption = message.caption or "What's this?"
         model = genai.GenerativeModel(
-            model_name="gemini-1.5-pro",
+            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,
@@ -390,7 +385,7 @@ async def chatbot_talk(client: Client, message: Message):
         pic_query = parts[1].strip() if len(parts) > 1 else ""
         try:
             model_flash = genai.GenerativeModel(
-                model_name="gemini-1.5-flash"
+                model_name=model_
             )
             backup_chat = await db._get_chatbot_chat_from_db(message.from_user.id)
             backup_chat.append({"role": "user", "parts": [{"text": query_base}]})