Update Akeno/__main__.py

Akeno/__main__.py

@@ -1,272 +1,83 @@
-import asyncio
-import logging
-import os
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# Copyright 2020-2023 (c) Randy W @xtdevs, @xtsea
+#
+# from : https://github.com/TeamKillerX
+# Channel : @RendyProjects
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
 
-from pyrogram import Client, filters
-from pyrogram import *
-from pyrogram.types import Message
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
 
-# 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
+import asyncio
+import importlib
+import logging
+import sys
+from contextlib import closing, suppress
+from importlib import import_module
 
-# 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")
+from pyrogram import idle
+from pyrogram.errors import *
+from uvloop import install
 
-# 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)
+from Akeno import clients
+from Akeno.plugins import ALL_MODULES
+from Akeno.utils.database import db
+from Akeno.utils.logger import LOGS
 
-# Initialize Pyrogram Client
-client = Client(
-    "chatbotai",
-    api_id=API_ID,
-    api_hash=API_HASH,
-    bot_token=BOT_TOKEN
-)
+logging.basicConfig(level=logging.INFO)
+logging.getLogger("pyrogram.syncer").setLevel(logging.WARNING)
+logging.getLogger("pyrogram.client").setLevel(logging.WARNING)
+loop = asyncio.get_event_loop()
 
-# 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):
+async def main():
     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
-            
+        await db.connect()
+        for module_name in ALL_MODULES:
+            imported_module = import_module(f"Akeno.plugins.{module_name}")
+        for cli in clients:
             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
+                await cli.start()
+            except SessionExpired as e:
+                LOGS.info(f"Error {e}")
+                sys.exit(1)
+            except ApiIdInvalid as e:
+                LOGS.info(f"Error {e}")
+                sys.exit(1)
+            except UserDeactivated as e:
+                LOGS.info(f"Error {e}")
+                sys.exit(1)
+            except AuthKeyDuplicated as e:
+                LOGS.info(f"Error {e}")
+                sys.exit(1)
             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
-            
+                LOGS.info(f"Error starting userbot: {e}")
+            ex = await cli.get_me()
+            LOGS.info(f"Started {ex.first_name}")
+            await cli.send_message("me", "Starting Akeno Userbot")
             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
+                await cli.join_chat("RendyProjects")
+            except UserIsBlocked:
+                return LOGS.info("You have been blocked. Please support @xtdevs")
+        await idle()
+    except Exception as e:
+        LOGS.info(f"Error in main: {e}")
+    finally:
+        for task in asyncio.all_tasks():
+            task.cancel()
+        LOGS.info("All tasks completed successfully!")
 
-# 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}")
+    install()
+    with closing(loop):
+        with suppress(asyncio.exceptions.CancelledError):
+            loop.run_until_complete(main())
+        loop.run_until_complete(asyncio.sleep(3.0))