Create main.py

main.py

@@ -0,0 +1,272 @@
+import asyncio
+import logging
+import os
+
+from pyrogram import Client, filters
+from pyrogram import *
+from pyrogram.types import Message
+
+# Your other imports
+from dotenv import load_dotenv
+from database import db
+from logger import LOGS
+from RyuzakiLib import GeminiLatest  # and other imports as needed
+import google.generativeai as genai
+from google.api_core.exceptions import InvalidArgument
+
+# Load environment variables
+load_dotenv()
+API_ID = os.getenv("API_ID")
+API_HASH = os.getenv("API_HASH")
+BOT_TOKEN = os.getenv("BOT_TOKEN")
+GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
+
+# Validate essential environment variables
+if not all([API_ID, API_HASH, BOT_TOKEN, GOOGLE_API_KEY]):
+    LOGS.critical("Missing one or more essential environment variables.")
+    exit(1)
+
+# Initialize Pyrogram Client
+client = Client(
+    "chatbotai",
+    api_id=API_ID,
+    api_hash=API_HASH,
+    bot_token=BOT_TOKEN
+)
+
+# Define your handler
+@client.on_message(
+    filters.incoming
+    & (
+        filters.text
+        | filters.photo
+        | filters.video
+        | filters.audio
+        | filters.voice
+        | filters.regex(r"\b(Randy|Rendi)\b(.*)", flags=re.IGNORECASE)
+    )
+    & filters.private
+    & ~filters.via_bot
+    & ~filters.forwarded,
+    group=2,
+)
+async def chatbot_talk(client: Client, message: Message):
+    try:
+        genai.configure(api_key=GOOGLE_API_KEY)
+
+        # 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:
+                backup_chat = await db._get_chatbot_chat_from_db(message.from_user.id)
+                backup_chat.append({"role": "user", "parts": [{"text": caption}]})
+                
+                response_reads = x.get_response_image(caption, file_path)
+                
+                if len(response_reads) > 4096:
+                    with open("chat.txt", "w+", encoding="utf8") as out_file:
+                        out_file.write(response_reads)
+                    await message.reply_document(
+                        document="chat.txt",
+                        disable_notification=True
+                    )
+                    await ai_reply.delete()
+                    os.remove("chat.txt")
+                else:
+                    await ai_reply.edit_text(response_reads)
+                
+                backup_chat.append({"role": "model", "parts": [{"text": response_reads}]})
+                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
+
+        # 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:
+                response = model.generate_content(
+                    [audio_file, caption],
+                    request_options={"timeout": 600}
+                )
+                if len(response.text) > 4096:
+                    with open("chat.txt", "w+", encoding="utf8") as out_file:
+                        out_file.write(response.text)
+                    await message.reply_document(
+                        document="chat.txt",
+                        disable_notification=True
+                    )
+                    await ai_reply.delete()
+                    os.remove("chat.txt")
+                else:
+                    await ai_reply.edit_text(response.text)
+                
+                backup_chat.append({"role": "model", "parts": [{"text": response.text}]})
+                await db._update_chatbot_chat_in_db(message.from_user.id, backup_chat)
+                
+                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
+
+        # 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:
+                response = model.generate_content(
+                    [video_file, caption],
+                    request_options={"timeout": 600}
+                )
+                if len(response.text) > 4096:
+                    with open("chat.txt", "w+", encoding="utf8") as out_file:
+                        out_file.write(response.text)
+                    await message.reply_document(
+                        document="chat.txt",
+                        disable_notification=True
+                    )
+                    await ai_reply.delete()
+                    os.remove("chat.txt")
+                else:
+                    await ai_reply.edit_text(response.text)
+                
+                backup_chat.append({"role": "model", "parts": [{"text": response.text}]})
+                await db._update_chatbot_chat_in_db(message.from_user.id, backup_chat)
+                
+                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
+
+        # 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:
+                model_flash = genai.GenerativeModel(
+                    model_name="gemini-1.5-flash"
+                )
+                backup_chat = await db._get_chatbot_chat_from_db(message.from_user.id)
+                backup_chat.append({"role": "user", "parts": [{"text": query_base}]})
+                
+                chat_session = model_flash.start_chat(history=backup_chat)
+                response_data = chat_session.send_message(query_base)
+                output = response_data.text
+                
+                if len(output) > 4096:
+                    with open("chat.txt", "w+", encoding="utf8") as out_file:
+                        out_file.write(output)
+                    await message.reply_document(
+                        document="chat.txt",
+                        disable_notification=True
+                    )
+                    os.remove("chat.txt")
+                else:
+                    await message.reply_text(output)
+                
+                backup_chat.append({"role": "model", "parts": [{"text": output}]})
+                await db._update_chatbot_chat_in_db(message.from_user.id, backup_chat)
+            except Exception as e:
+                await message.reply_text(str(e))
+    # End of handler
+
+# Define the main coroutine
+async def main():
+    await db.connect()      # Connect to your database
+    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
+    LOGS.info("Bot stopping...")
+    await client.stop()    # Ensure the client stops gracefully
+
+# Entry point
+if __name__ == "__main__":
+    try:
+        asyncio.run(main())
+    except (KeyboardInterrupt, SystemExit):
+        LOGS.info("Bot has been terminated by the user.")
+    except Exception as e:
+        LOGS.error(f"Unexpected error: {e}")
+