overspend1/plain-ub-overfork
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

clean up .ids, -d in rename, do not log service texts. Beta AI cmds: .stt, .aim

Browse Source
This commit is contained in:
thedragonsinn
2024-04-20 19:13:19 +05:30
parent ef88f25c34
commit 1d036ded9c
10 changed files with 313 additions and 112 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
app/__init__.py
View File

@@ -1,11 +1 @@
from ub_core import (
BOT,
DB,
DB_CLIENT,
LOGGER,
Config,
Convo,
CustomDB,
Message,
bot,
)
from ub_core import BOT, DB, DB_CLIENT, LOGGER, Config, Convo, CustomDB, Message, bot

13
app/plugins/admin/fbans.py
View File

@@ -36,11 +36,7 @@ async def add_fed(bot: BOT, message: Message):
data = dict(name=message.input or message.chat.title, type=str(message.chat.type))
await FED_DB.add_data({"_id": message.chat.id, **data})
text = f"#FBANS\n<b>{data['name']}</b>: <code>{message.chat.id}</code> added to FED LIST."
await message.reply(
text=text,
del_in=5,
block=True,
)
await message.reply(text=text, del_in=5, block=True)
await bot.log_text(text=text, type="info")
@@ -64,13 +60,10 @@ async def remove_fed(bot: BOT, message: Message):
chat = chat.id
elif chat.lstrip("-").isdigit():
chat = int(chat)
deleted: bool | None = await FED_DB.delete_data(id=chat)
deleted: int = await FED_DB.delete_data(id=chat)
if deleted:
text = f"#FBANS\n<b>{name}</b><code>{chat}</code> removed from FED LIST."
await message.reply(
text=text,
del_in=8,
)
await message.reply(text=text, del_in=8)
await bot.log_text(text=text, type="info")
else:
await message.reply(text=f"<b>{name or chat}</b> not in FED LIST.", del_in=8)

187
app/plugins/ai/media_query.py Normal file
View File

@@ -0,0 +1,187 @@
import asyncio
import glob
import mimetypes
import os
import shutil
import time
from io import BytesIO
import google.generativeai as genai
from google.ai import generativelanguage as glm
from ub_core.utils import run_shell_cmd
from app import BOT, Message, bot
from app.plugins.ai.models import (
IMAGE_MODEL,
MEDIA_MODEL,
TEXT_MODEL,
get_response_text,
)
CODE_EXTS = {
".java",
".c",
".cpp",
".cc",
".cxx",
".py",
".js",
".html",
".htm",
".css",
".rb",
".php",
".swift",
".go",
".sql",
".r",
".pl",
".kt",
}
PHOTO_EXTS = {".png", ".jpg", ".jpeg", ".webp"}
VIDEO_EXTS = {".mp4", ".mkv", ".webm", ".gif"}
AUDIO_EXTS = {".aac", ".mp3", ".opus", ".m4a", ".ogg"}
@bot.add_cmd(cmd="ocr")
async def photo_query(bot: BOT, message: Message):
"""
CMD: OCR
INFO: Ask a question to Gemini AI about replied image.
USAGE: .ocr [reply to a photo] explain the image.
"""
prompt = message.input
reply = message.replied
if not (prompt and reply and reply.photo):
await message.reply("Reply to an image and give a prompt.")
return
response_text = await handle_photo(prompt, reply)
await message.reply(response_text)
@bot.add_cmd(cmd="stt")
async def audio_to_text(bot: BOT, message: Message):
"""
CMD: STT
INFO: Convert Audio files to text.
USAGE: .stt [reply to audio file] summarise/transcribe the audio file.
"""
prompt = message.input
reply = message.replied
audio = message.audio or message.voice
if not (prompt and reply and audio):
await message.reply("Reply to an image and give a prompt.")
return
response_text = await handle_audio(prompt, reply)
await message.reply(response_text)
@bot.add_cmd(cmd="aim")
async def handle_document(bot: BOT, message: Message):
"""
CMD: AIM
INFO: Prompt Ai to perform task for documents containing pic, vid, code, audio.
USAGE: .aim [reply to audio file] convert this file to python | summarise the video, audio, picture.
"""
prompt = message.input
reply = message.replied
document = reply.document
if not (prompt and reply and document):
await message.reply("Reply to a document and give a prompt.")
return
file_name = document.file_name
if not file_name:
await message.reply("Unsupported file.")
return
name, ext = os.path.splitext(file_name)
if ext in PHOTO_EXTS:
response_text = await handle_photo(prompt, reply)
elif ext in AUDIO_EXTS:
response_text = await handle_audio(prompt, reply)
elif ext in CODE_EXTS:
response_text = await handle_code(prompt, reply)
elif ext in VIDEO_EXTS:
response_text = await handle_video(prompt, reply)
else:
await message.reply("Unsupported Media.")
return
await message.reply(response_text)
async def download_file(file_name: str, message: Message) -> tuple[str, str]:
download_dir = os.path.join("downloads", str(time.time()))
file_path = os.path.join(download_dir, file_name)
await message.download(file_path)
return file_path, download_dir
async def handle_audio(prompt: str, message: Message):
audio = message.document or message.audio or message.voice
file_name = audio.file_name or "audio.aac"
file_path, download_dir = await download_file(file_name, message)
file_response = genai.upload_file(path=file_path)
response = await MEDIA_MODEL.generate_content_async([prompt, file_response])
response_text = get_response_text(response)
genai.delete_file(name=file_response.name)
shutil.rmtree(file_path, ignore_errors=True)
return response_text
async def handle_code(prompt: str, message: Message):
file: BytesIO = await message.download(in_memory=True)
text = file.getvalue().decode("utf-8")
final_prompt = f"{text}\n\n{prompt}"
response = await TEXT_MODEL.generate_content_async(final_prompt)
return get_response_text(response)
async def handle_photo(prompt: str, message: Message):
file = await message.download(in_memory=True)
mime_type, _ = mimetypes.guess_type(file.name)
if mime_type is None:
mime_type = "image/unknown"
image_blob = glm.Blob(mime_type=mime_type, data=file.getvalue())
response = await IMAGE_MODEL.generate_content_async([prompt, image_blob])
return get_response_text(response)
async def handle_video(prompt: str, message: Message):
file_name = "v.mp4"
file_path, download_dir = await download_file(file_name, message)
output_path = os.path.join(download_dir, "output_frame_%04d.png")
ffmpeg_output_error = await run_shell_cmd(
f'ffmpeg -hide_banner -loglevel error -i {file_path} -vf "fps=1" {output_path}'
)
if ffmpeg_output_error:
return ffmpeg_output_error
extracted_frames = glob.glob(f"{download_dir}/*png")
uploaded_frames = []
for frame in extracted_frames:
uploaded_frame = await asyncio.to_thread(genai.upload_file, frame)
uploaded_frames.append(uploaded_frame)
response = await MEDIA_MODEL.generate_content_async([prompt, *uploaded_frames])
response_text = get_response_text(response)
for uploaded_frame in uploaded_frames:
genai.delete_file(name=uploaded_frame.name)
shutil.rmtree(download_dir, ignore_errors=True)
return response_text

55
app/plugins/ai/models.py Normal file
View File

@@ -0,0 +1,55 @@
import google.generativeai as genai
from app import Message, extra_config
async def init_task():
if extra_config.GEMINI_API_KEY:
genai.configure(api_key=extra_config.GEMINI_API_KEY)
GENERATION_CONFIG = {"temperature": 0.69, "max_output_tokens": 2048}
SAFETY_SETTINGS = [
{"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_ONLY_HIGH"},
{"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_ONLY_HIGH"},
{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_ONLY_HIGH"},
{"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_ONLY_HIGH"},
]
TEXT_MODEL = genai.GenerativeModel(
model_name="gemini-pro",
generation_config=GENERATION_CONFIG,
safety_settings=SAFETY_SETTINGS,
)
IMAGE_MODEL = genai.GenerativeModel(
model_name="gemini-pro-vision",
generation_config=GENERATION_CONFIG,
safety_settings=SAFETY_SETTINGS,
)
MEDIA_MODEL = genai.GenerativeModel(
model_name="models/gemini-1.5-pro-latest",
generation_config=GENERATION_CONFIG,
safety_settings=SAFETY_SETTINGS,
)
async def basic_check(message: Message):
if not extra_config.GEMINI_API_KEY:
await message.reply(
"Gemini API KEY not found."
"\nGet it <a href='https://makersuite.google.com/app/apikey'>HERE</a> "
"and set GEMINI_API_KEY var."
)
return
if not message.input:
await message.reply("Ask a Question.")
return
return 1
def get_response_text(response):
return "\n".join([part.text for part in response.parts])

84
app/plugins/misc/gemini.py → app/plugins/ai/text_query.py
View File

@@ -1,54 +1,11 @@
import mimetypes
import pickle
from io import BytesIO
import google.generativeai as genai
from google.ai import generativelanguage as glm
from pyrogram import filters
from pyrogram.enums import ParseMode
from app import BOT, Convo, Message, bot, extra_config
GENERATION_CONFIG = {"temperature": 0.69, "max_output_tokens": 2048}
SAFETY_SETTINGS = [
{"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_ONLY_HIGH"},
{"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_ONLY_HIGH"},
{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_ONLY_HIGH"},
{"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_ONLY_HIGH"},
]
TEXT_MODEL = genai.GenerativeModel(
model_name="gemini-pro",
generation_config=GENERATION_CONFIG,
safety_settings=SAFETY_SETTINGS,
)
VISION_MODEL = genai.GenerativeModel(
model_name="gemini-pro-vision",
generation_config=GENERATION_CONFIG,
safety_settings=SAFETY_SETTINGS,
)
async def init_task():
if extra_config.GEMINI_API_KEY:
genai.configure(api_key=extra_config.GEMINI_API_KEY)
async def basic_check(message: Message):
if not extra_config.GEMINI_API_KEY:
await message.reply(
"Gemini API KEY not found."
"\nGet it <a href='https://makersuite.google.com/app/apikey'>HERE</a> "
"and set GEMINI_API_KEY var."
)
return
if not message.input:
await message.reply("Ask a Question.")
return
return 1
from app import BOT, Convo, Message, bot
from app.plugins.ai.models import TEXT_MODEL, basic_check, get_response_text
@bot.add_cmd(cmd="ai")
@@ -59,25 +16,15 @@ async def question(bot: BOT, message: Message):
USAGE: .ai what is the meaning of life.
"""
if not (await basic_check(message)): # fmt:skip
if not await basic_check(message):
return
prompt = message.input
reply = message.replied
if reply and reply.photo:
file = await reply.download(in_memory=True)
mime_type, _ = mimetypes.guess_type(file.name)
if mime_type is None:
mime_type = "image/unknown"
image_blob = glm.Blob(mime_type=mime_type, data=file.getvalue())
response = await VISION_MODEL.generate_content_async([prompt, image_blob])
else:
response = await TEXT_MODEL.generate_content_async(prompt)
response = await TEXT_MODEL.generate_content_async(prompt)
response_text = get_response_text(response)
if not isinstance(message, Message):
await message.edit(
text=f"```\n{prompt}```**GEMINI AI**:\n{response_text.strip()}",
@@ -103,7 +50,7 @@ async def ai_chat(bot: BOT, message: Message):
After 5 mins of Idle bot will export history and stop chat.
use .load_history to continue
"""
if not (await basic_check(message)): # fmt:skip
if not await basic_check(message):
return
chat = TEXT_MODEL.start_chat(history=[])
try:
@@ -113,14 +60,14 @@ async def ai_chat(bot: BOT, message: Message):
@bot.add_cmd(cmd="load_history")
async def ai_chat(bot: BOT, message: Message):
async def history_chat(bot: BOT, message: Message):
"""
CMD: LOAD_HISTORY
INFO: Load a Conversation with Gemini AI from previous session.
USAGE:
.load_history {question} [reply to history document]
"""
if not (await basic_check(message)): # fmt:skip
if not await basic_check(message):
return
reply = message.replied
if (
@@ -142,15 +89,11 @@ async def ai_chat(bot: BOT, message: Message):
await export_history(chat, message)
def get_response_text(response):
return "\n".join([part.text for part in response.parts])
async def do_convo(chat, message: Message, history: bool = False):
prompt = message.input
reply_to_message_id = message.id
async with Convo(
client=bot,
client=message._client,
chat_id=message.chat.id,
filters=generate_filter(message),
timeout=300,
@@ -177,7 +120,7 @@ def generate_filter(message: Message):
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
or msg.reply_to_message.from_user.id != message._client.me.id
):
return False
return True
@@ -188,6 +131,7 @@ def generate_filter(message: Message):
async def export_history(chat, message: Message):
doc = BytesIO(pickle.dumps(chat.history))
doc.name = "AI_Chat_History.pkl"
await bot.send_document(
chat_id=message.from_user.id, document=doc, caption=message.text
caption = get_response_text(
await chat.send_message_async("Summarize our Conversation into one line.")
)
await bot.send_document(chat_id=message.from_user.id, document=doc, caption=caption)

6
app/plugins/files/download.py
View File

@@ -66,8 +66,10 @@ async def down_load(bot: BOT, message: Message):
async def telegram_download(
message: Message, response: Message, path: str, file_name: str | None = None
) -> DownloadedFile:
message: Message,
response: Message,
path: str,
file_name: str | None = None) -> DownloadedFile:
"""
:param message: Message Containing Media
:param response: Response to Edit

20
app/plugins/files/rename.py
View File

@@ -8,7 +8,7 @@ from ub_core.utils.helpers import progress
from app import BOT, Message, bot
from app.plugins.files.download import telegram_download
from app.plugins.files.upload import FILE_TYPE_MAP
from app.plugins.files.upload import FILE_TYPE_MAP, MediaType
@bot.add_cmd(cmd="rename")
@@ -31,10 +31,7 @@ async def rename(bot: BOT, message: Message):
await response.edit("Input verified....Starting Download...")
if message.replied:
download_coro = telegram_download(
message=message.replied,
path=dl_path,
file_name=input,
response=response,
message=message.replied, path=dl_path, file_name=input, response=response
)
else:
url, file_name = input.split(maxsplit=1)
@@ -44,9 +41,16 @@ async def rename(bot: BOT, message: Message):
download_coro = dl_obj.download()
try:
downloaded_file: DownloadedFile = await download_coro
media: dict = await FILE_TYPE_MAP[downloaded_file.type](
downloaded_file, has_spoiler="-s" in message.flags
)
if "-d" in message.flags:
media: dict = await FILE_TYPE_MAP[MediaType.DOCUMENT](
downloaded_file, has_spoiler="-s" in message.flags
)
else:
media: dict = await FILE_TYPE_MAP[downloaded_file.type](
downloaded_file, has_spoiler="-s" in message.flags
)
progress_args = (
response,
"Uploading...",

29
app/plugins/tg_tools/chat.py
View File

@@ -1,7 +1,8 @@
import asyncio
import os
from pyrogram.enums import ChatType
from pyrogram.errors import BadRequest
from ub_core.utils import get_name
from app import BOT, Message, bot
@@ -10,19 +11,26 @@ from app import BOT, Message, bot
async def get_ids(bot: BOT, message: Message) -> None:
reply: Message = message.replied
if reply:
ids: str = ""
reply_forward = reply.forward_from_chat
reply_user = reply.from_user
ids += f"<b>Chat</b>: `{reply.chat.id}`\n"
resp_str: str = ""
reply_user, reply_forward = reply.forward_from_chat, reply.from_user
resp_str += f"<b>{get_name(reply.chat)}</b>: <code>{reply.chat.id}</code>\n"
if reply_forward:
ids += f"<b>Replied {'Channel' if reply_forward.type == ChatType.CHANNEL else 'Chat'}</b>: `{reply_forward.id}`\n"
resp_str += (
f"<b>{get_name(reply_forward)}</b>: <code>{reply_forward.id}</code>\n"
)
if reply_user:
ids += f"<b>User</b>: {reply.from_user.id}"
resp_str += f"<b>{get_name(reply_user)}</b>: <code>{reply_user.id}</code>"
elif message.input:
ids: int = (await bot.get_chat(message.input[1:])).id
resp_str: int = (await bot.get_chat(message.input[1:])).id
else:
ids: str = f"<b>Chat</b>:`{message.chat.id}`"
await message.reply(ids)
resp_str: str = (
f"<b>{get_name(message.chat)}</b>: <code>{message.chat.id}</code>"
)
await message.reply(resp_str)
@bot.add_cmd(cmd="join")
@@ -50,6 +58,7 @@ async def leave_chat(bot: BOT, message: Message) -> None:
del_in=5,
block=True,
)
await asyncio.sleep(2)
try:
await bot.leave_chat(chat)
except Exception as e:

15
app/plugins/tg_tools/pm_n_tag_logger.py
View File

@@ -64,6 +64,7 @@ basic_filters = (
& ~filters.service
& ~filters.chat(chats=[bot.me.id])
& ~filters.me
& ~filters.create(lambda _, __, m: m.chat.is_support)
)
@@ -134,31 +135,43 @@ async def runner():
if not (extra_config.TAG_LOGGER or extra_config.PM_LOGGER):
return
last_pm_logged_id = 0
while True:
cached_keys = list(MESSAGE_CACHE.keys())
if not cached_keys:
await asyncio.sleep(5)
continue
first_key = cached_keys[0]
cached_list = MESSAGE_CACHE.copy()[first_key]
if not cached_list:
MESSAGE_CACHE.pop(first_key)
for idx, msg in enumerate(cached_list):
if msg.chat.type == ChatType.PRIVATE:
if last_pm_logged_id != first_key:
last_pm_logged_id = first_key
log_info = True
else:
log_info = False
coro = log_pm(message=msg, log_info=log_info)
else:
coro = log_chat(message=msg)
try:
await coro
except BaseException:
pass
MESSAGE_CACHE[first_key].remove(msg)
await asyncio.sleep(5)
await asyncio.sleep(15)
@@ -205,7 +218,7 @@ async def log_chat(message: Message):
try:
logged = await message.forward(extra_config.MESSAGE_LOGGER_CHAT)
await logged.reply(
text=f"#TAG\n{mention} [{u_id}]\nMessage: \n<a href='{message.link}'>{message.chat.title}</a> ({message.chat.id})",
text=f"#TAG\n{mention} [{u_id}]\nMessage: \n<a href='{message.link}'>{message.chat.title}</a> ({message.chat.id})"
)
except MessageIdInvalid:
await message.copy(extra_config.MESSAGE_LOGGER_CHAT, caption=notice)

4
app/plugins/tg_tools/pm_permit.py
View File

@@ -41,6 +41,10 @@ async def handle_new_pm(bot: BOT, message: Message):
type="info",
)
RECENT_USERS[user_id] += 1
if message.chat.is_support:
return
if RECENT_USERS[user_id] >= 5:
await message.reply("You've been blocked for spamming.")
await bot.block_user(user_id)