Update main.py

main.py

@@ -2,8 +2,7 @@ import asyncio
 import logging
 import os
 
-from pyrogram import Client, filters
-from pyrogram import *
+from pyrogram import Client, filters, idle
 from pyrogram.types import Message
 
 # Your other imports
@@ -56,14 +55,14 @@ async def chatbot_talk(client: Client, message: Message):
 
         # Handling Photo Messages
         if message.photo:
-            file_path = await message.download()
-            caption = message.caption or "What's this?"
-            x = GeminiLatest(api_keys=GOOGLE_API_KEY)
-            
-            # Send initial processing message
-            ai_reply = await message.reply_text("Processing...")
-            
             try:
+                file_path = await message.download()
+                caption = message.caption or "What's this?"
+                x = GeminiLatest(api_keys=GOOGLE_API_KEY)
+                
+                # Send initial processing message
+                ai_reply = await message.reply_text("Processing...")
+                
                 backup_chat = await db._get_chatbot_chat_from_db(message.from_user.id)
                 backup_chat.append({"role": "user", "parts": [{"text": caption}]})
                 
@@ -85,43 +84,41 @@ async def chatbot_talk(client: Client, message: Message):
                 await db._update_chatbot_chat_in_db(message.from_user.id, backup_chat)
                 
                 os.remove(file_path)
-                return
             except InvalidArgument as e:
                 await ai_reply.edit_text(f"Error: {e}")
-                return
             except Exception as e:
                 await ai_reply.edit_text(f"Error: {e}")
-                return
+            return
 
         # Handling Audio or Voice Messages
         if message.audio or message.voice:
-            ai_reply = await message.reply_text("Processing...")
-            audio_file_name = await message.download()
-            caption = message.caption or "What's this?"
-            model = genai.GenerativeModel(
-                model_name="gemini-1.5-flash",
-                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,
-                    genai.types.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: genai.types.HarmBlockThreshold.BLOCK_NONE,
-                    genai.types.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: genai.types.HarmBlockThreshold.BLOCK_NONE,
-                }
-            )
-            backup_chat = await db._get_chatbot_chat_from_db(message.from_user.id)
-            backup_chat.append({"role": "user", "parts": [{"text": caption}]})
-            
-            ai_reply.edit_text("Uploading file...")
-            audio_file = genai.upload_file(path=audio_file_name)
-            
-            while audio_file.state.name == "PROCESSING":
-                await asyncio.sleep(10)
-                audio_file = genai.get_file(audio_file.name)
-            
-            if audio_file.state.name == "FAILED":
-                await ai_reply.edit_text(f"Error: {audio_file.state.name}")
-                return
-            
             try:
+                ai_reply = await message.reply_text("Processing...")
+                audio_file_name = await message.download()
+                caption = message.caption or "What's this?"
+                model = genai.GenerativeModel(
+                    model_name="gemini-1.5-flash",
+                    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,
+                        genai.types.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: genai.types.HarmBlockThreshold.BLOCK_NONE,
+                        genai.types.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: genai.types.HarmBlockThreshold.BLOCK_NONE,
+                    }
+                )
+                backup_chat = await db._get_chatbot_chat_from_db(message.from_user.id)
+                backup_chat.append({"role": "user", "parts": [{"text": caption}]})
+                
+                await ai_reply.edit_text("Uploading file...")
+                audio_file = genai.upload_file(path=audio_file_name)
+                
+                while audio_file.state.name == "PROCESSING":
+                    await asyncio.sleep(10)
+                    audio_file = genai.get_file(audio_file.name)
+                
+                if audio_file.state.name == "FAILED":
+                    await ai_reply.edit_text(f"Error: {audio_file.state.name}")
+                    return
+                
                 response = model.generate_content(
                     [audio_file, caption],
                     request_options={"timeout": 600}
@@ -143,43 +140,41 @@ async def chatbot_talk(client: Client, message: Message):
                 
                 audio_file.delete()
                 os.remove(audio_file_name)
-                return
             except InvalidArgument as e:
                 await ai_reply.edit_text(f"Error: {e}")
-                return
             except Exception as e:
                 await ai_reply.edit_text(f"Error: {e}")
-                return
+            return
 
         # Handling Video Messages
         if message.video:
-            ai_reply = await message.reply_text("Processing...")
-            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",
-                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,
-                    genai.types.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: genai.types.HarmBlockThreshold.BLOCK_NONE,
-                    genai.types.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: genai.types.HarmBlockThreshold.BLOCK_NONE,
-                }
-            )
-            backup_chat = await db._get_chatbot_chat_from_db(message.from_user.id)
-            backup_chat.append({"role": "user", "parts": [{"text": caption}]})
-            
-            ai_reply.edit_text("Uploading file...")
-            video_file = genai.upload_file(path=video_file_name)
-            
-            while video_file.state.name == "PROCESSING":
-                await asyncio.sleep(10)
-                video_file = genai.get_file(video_file.name)
-            
-            if video_file.state.name == "FAILED":
-                await ai_reply.edit_text(f"Error: {video_file.state.name}")
-                return
-            
             try:
+                ai_reply = await message.reply_text("Processing...")
+                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",
+                    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,
+                        genai.types.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: genai.types.HarmBlockThreshold.BLOCK_NONE,
+                        genai.types.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: genai.types.HarmBlockThreshold.BLOCK_NONE,
+                    }
+                )
+                backup_chat = await db._get_chatbot_chat_from_db(message.from_user.id)
+                backup_chat.append({"role": "user", "parts": [{"text": caption}]})
+                
+                await ai_reply.edit_text("Uploading file...")
+                video_file = genai.upload_file(path=video_file_name)
+                
+                while video_file.state.name == "PROCESSING":
+                    await asyncio.sleep(10)
+                    video_file = genai.get_file(video_file.name)
+                
+                if video_file.state.name == "FAILED":
+                    await ai_reply.edit_text(f"Error: {video_file.state.name}")
+                    return
+                
                 response = model.generate_content(
                     [video_file, caption],
                     request_options={"timeout": 600}
@@ -201,29 +196,27 @@ async def chatbot_talk(client: Client, message: Message):
                 
                 video_file.delete()
                 os.remove(video_file_name)
-                return
             except InvalidArgument as e:
                 await ai_reply.edit_text(f"Error: {e}")
-                return
             except Exception as e:
                 await ai_reply.edit_text(f"Error: {e}")
-                return
+            return
 
         # Handling Text Messages
         if message.text:
-            query = message.text.strip()
-            match = re.search(r"\b(Randy|Rendi)\b(.*)", query, flags=re.IGNORECASE)
-            if match:
-                rest_of_sentence = match.group(2).strip()
-                query_base = rest_of_sentence if rest_of_sentence else query
-            else:
-                query_base = query
-            
-            parts = query.split(maxsplit=1)
-            command = parts[0].lower()
-            pic_query = parts[1].strip() if len(parts) > 1 else ""
-            
             try:
+                query = message.text.strip()
+                match = re.search(r"\b(Randy|Rendi)\b(.*)", query, flags=re.IGNORECASE)
+                if match:
+                    rest_of_sentence = match.group(2).strip()
+                    query_base = rest_of_sentence if rest_of_sentence else query
+                else:
+                    query_base = query
+                
+                parts = query.split(maxsplit=1)
+                command = parts[0].lower()
+                pic_query = parts[1].strip() if len(parts) > 1 else ""
+                
                 model_flash = genai.GenerativeModel(
                     model_name="gemini-1.5-flash"
                 )
@@ -258,9 +251,9 @@ async def main():
         LOGS.info("Connected to the database.")
         await client.start()   # Start the Pyrogram client
         LOGS.info("Bot started successfully.")
-        await idle()    # Keep the bot running until interrupted
+        await idle()            # Keep the bot running until interrupted
         LOGS.info("Bot stopping...")
-        await client.stop()    # Ensure the client stops gracefully
+        await client.stop()     # Ensure the client stops gracefully
     except Exception as e:
         LOGS.error(f"Unexpected error: {e}")
 
@@ -272,4 +265,3 @@ if __name__ == "__main__":
         LOGS.info("Bot has been terminated by the user.")
     except Exception as e:
         LOGS.error(f"Unexpected error: {e}")
-