From 41247e834df88c0d6c1ec0e6f461345f54aaf1cd Mon Sep 17 00:00:00 2001
From: thedragonsinn <98635854+thedragonsinn@users.noreply.github.com>
Date: Mon, 12 Feb 2024 14:57:54 +0530
Subject: [PATCH] cooking almost done.
---
app/config.py | 68 ++++++++++++++---------------
app/plugins/misc/gemini.py | 87 ++++++++++++++++++++++++++++++++++++++
sample-config.env | 5 +++
3 files changed, 124 insertions(+), 36 deletions(-)
create mode 100644 app/plugins/misc/gemini.py
diff --git a/app/config.py b/app/config.py
index d98ce29..13fc0cc 100644
--- a/app/config.py
+++ b/app/config.py
@@ -9,63 +9,59 @@ from app.utils import Str
class Cmd(Str):
def __init__(self, cmd: str, func: Callable, path: str, sudo: bool):
- self.cmd: str = cmd
- self.func: Callable = func
- self.path: str = path
- self.dirname: str = os.path.basename(os.path.dirname(path))
- self.doc: str = func.__doc__ or "Not Documented."
- self.sudo: bool = sudo
+ cmd: str = cmd
+ func: Callable = func
+ path: str = path
+ dirname: str = os.path.basename(os.path.dirname(path))
+ doc: str = func.__doc__ or "Not Documented."
+ sudo: bool = sudo
-class _Config(Str):
+class Config:
CMD = Cmd
- def __init__(self):
- self.CMD_DICT: dict[str, Cmd] = {}
+ CMD_DICT: dict[str, Cmd] = {}
- self.CMD_TRIGGER: str = os.environ.get("CMD_TRIGGER", ".")
+ CMD_TRIGGER: str = os.environ.get("CMD_TRIGGER", ".")
- self.DEV_MODE: int = int(os.environ.get("DEV_MODE", 0))
+ DEV_MODE: int = int(os.environ.get("DEV_MODE", 0))
- self.DISABLED_SUPERUSERS: list[int] = []
+ DISABLED_SUPERUSERS: list[int] = []
- self.FBAN_LOG_CHANNEL: int = int(
- os.environ.get("FBAN_LOG_CHANNEL", os.environ.get("LOG_CHAT"))
- )
+ FBAN_LOG_CHANNEL: int = int(
+ os.environ.get("FBAN_LOG_CHANNEL", os.environ.get("LOG_CHAT"))
+ )
- self.INIT_TASKS: list[Coroutine] = []
+ GEMINI_API_KEY: str = os.environ.get("GEMINI_API_KEY")
- self.LOG_CHAT: int = int(os.environ.get("LOG_CHAT"))
+ INIT_TASKS: list[Coroutine] = []
- self.MESSAGE_LOGGER_CHAT: int = int(
- os.environ.get("MESSAGE_LOGGER_CHAT", self.LOG_CHAT)
- )
+ LOG_CHAT: int = int(os.environ.get("LOG_CHAT"))
- self.MESSAGE_LOGGER_TASK: asyncio.Task | None = None
+ MESSAGE_LOGGER_CHAT: int = int(os.environ.get("MESSAGE_LOGGER_CHAT", LOG_CHAT))
- self.OWNER_ID: int = int(os.environ.get("OWNER_ID"))
+ MESSAGE_LOGGER_TASK: asyncio.Task | None = None
- self.PM_GUARD: bool = False
+ OWNER_ID: int = int(os.environ.get("OWNER_ID"))
- self.PM_LOGGER: bool = False
+ PM_GUARD: bool = False
- self.REPO: Repo = Repo(".")
+ PM_LOGGER: bool = False
- self.SUDO: bool = False
+ REPO: Repo = Repo(".")
- self.SUDO_TRIGGER: str = os.environ.get("SUDO_TRIGGER", "!")
+ SUDO: bool = False
- self.SUDO_CMD_LIST: list[str] = []
+ SUDO_TRIGGER: str = os.environ.get("SUDO_TRIGGER", "!")
- self.SUDO_USERS: list[int] = []
+ SUDO_CMD_LIST: list[str] = []
- self.SUPERUSERS: list[int] = []
+ SUDO_USERS: list[int] = []
- self.TAG_LOGGER: bool = False
+ SUPERUSERS: list[int] = []
- self.UPSTREAM_REPO: str = os.environ.get(
- "UPSTREAM_REPO", "https://github.com/thedragonsinn/plain-ub"
- )
+ TAG_LOGGER: bool = False
-
-Config = _Config()
+ UPSTREAM_REPO: str = os.environ.get(
+ "UPSTREAM_REPO", "https://github.com/thedragonsinn/plain-ub"
+ )
diff --git a/app/plugins/misc/gemini.py b/app/plugins/misc/gemini.py
new file mode 100644
index 0000000..aa3a203
--- /dev/null
+++ b/app/plugins/misc/gemini.py
@@ -0,0 +1,87 @@
+import google.generativeai as genai
+from pyrogram import filters
+
+from app import BOT, Message, bot, Config, Convo
+
+MODEL = genai.GenerativeModel("gemini-pro")
+
+
+async def init_task():
+ if Config.GEMINI_API_KEY:
+ genai.configure(api_key=Config.GEMINI_API_KEY)
+
+
+async def basic_check(message: Message):
+ if not Config.GEMINI_API_KEY:
+ await message.reply(
+ "Gemini API KEY not found."
+ "\nGet it HERE "
+ "and set GEMINI_API_KEY var."
+ )
+ return
+ if not message.input:
+ await message.reply("Ask a Question.")
+ return
+ return 1
+
+
+@bot.add_cmd(cmd="ai")
+async def question(bot: BOT, message: Message):
+ """
+ CMD: AI
+ INFO: Ask a question to Gemini AI.
+ USAGE: .ai what is the meaning of life.
+ """
+ if not (await basic_check(message)): # fmt:skip
+ return
+ response = await MODEL.generate_content_async(message.input)
+ await message.reply(response)
+
+
+@bot.add_cmd(cmd="aichat")
+async def ai_chat(bot: BOT, message: Message):
+ """
+ CMD: AICHAT
+ INFO: Have a Conversation with Gemini AI.
+ USAGE:
+ .aichat hello
+ keep replying to AI responses
+ """
+ if not (await basic_check(message)): # fmt:skip
+ return
+ try:
+ await do_convo(message)
+ except TimeoutError:
+ await message.reply("AI Chat TimeOut.")
+
+
+async def do_convo(message: Message):
+ chat = MODEL.start_chat(history=[])
+ prompt = message.input
+ async with Convo(
+ client=bot,
+ chat_id=message.chat.id,
+ filters=generate_filter(message),
+ timeout=600,
+ ) as convo:
+ while True:
+ ai_response = (await chat.send_message_async(prompt)).text
+ _, prompt = await convo.send_message(
+ text=f"GEMINI AI:\n\n{ai_response}", get_response=True
+ )
+
+
+def generate_filter(message: Message):
+ async def _filter(_, __, msg: Message):
+ if (
+ not msg.text
+ or not msg.from_user
+ or msg.from_user.id != message.from_user.id
+ or not msg.reply_to_message
+ or not msg.reply_to_message.from_user
+ or msg.reply_to_message.from_user.id == bot.me.id
+ ):
+ return False
+ return True
+
+ return filters.create(_filter)
diff --git a/sample-config.env b/sample-config.env
index 4aceb6d..cb73c67 100644
--- a/sample-config.env
+++ b/sample-config.env
@@ -23,6 +23,11 @@ DB_URL=
FBAN_LOG_CHANNEL=
# Optional FedBan Proof and logs.
+GEMINI_API_KEY=
+# Optional API Key
+# Get from https://ai.google.dev/
+
+
LOG_CHAT=
# Bot logs chat/channel