Update Akeno/__main__.py

Akeno/__main__.py

@@ -1,83 +1,275 @@
-#!/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.
-
-# 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/>.
-
 import asyncio
-import importlib
 import logging
-import sys
-from contextlib import closing, suppress
-from importlib import import_module
+import os
 
-from pyrogram import idle
-from pyrogram.errors import *
-from uvloop import install
+from pyrogram import Client, filters
+from pyrogram import *
+from pyrogram.types import Message
 
-from Akeno import clients
-from Akeno.plugins import ALL_MODULES
-from Akeno.utils.database import db
-from Akeno.utils.logger import LOGS
+# 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
 
-logging.basicConfig(level=logging.INFO)
-logging.getLogger("pyrogram.syncer").setLevel(logging.WARNING)
-logging.getLogger("pyrogram.client").setLevel(logging.WARNING)
-loop = asyncio.get_event_loop()
+# 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")
 
-async def main():
+# 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:
-        await db.connect()
-        for module_name in ALL_MODULES:
-            imported_module = import_module(f"Akeno.plugins.{module_name}")
-        for cli in clients:
+        chat_user = await db.get_chatbot(message.chat.id)
+        if not chat_user:
+            return  # Exit if no chatbot is configured for the chat
+
+        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:
-                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)
+                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:
-                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")
+                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:
-                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!")
+                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__":
-    install()
-    with closing(loop):
-        with suppress(asyncio.exceptions.CancelledError):
-            loop.run_until_complete(main())
-        loop.run_until_complete(asyncio.sleep(3.0))
+    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}")