106 lines
3.2 KiB
Python
106 lines
3.2 KiB
Python
from pyrogram import Client, enums, filters
|
|
from pyrogram.types import Message
|
|
|
|
from utils.misc import modules_help, prefix
|
|
from utils.config import cohere_key
|
|
from utils.db import db
|
|
from utils.scripts import format_exc, import_library, restart
|
|
|
|
cohere = import_library("cohere")
|
|
|
|
co = cohere.Client(cohere_key)
|
|
|
|
chatai_users = db.getaiusers()
|
|
|
|
|
|
@Client.on_message(filters.command("addai", prefix) & filters.me)
|
|
async def adduser(_, message: Message):
|
|
if len(message.command) > 1:
|
|
user_id = message.text.split(maxsplit=1)[1]
|
|
if user_id.isdigit():
|
|
user_id = int(user_id)
|
|
db.addaiuser(user_id)
|
|
await message.edit_text("<b>User ID Added</b>")
|
|
restart()
|
|
else:
|
|
await message.edit_text("<b>User ID is invalid.</b>")
|
|
return
|
|
else:
|
|
await message.edit_text(f"<b>Usage: </b><code>{prefix}addai [user_id]</code>")
|
|
return
|
|
|
|
|
|
@Client.on_message(filters.command("remai", prefix) & filters.me)
|
|
async def remuser(_, message: Message):
|
|
if len(message.command) > 1:
|
|
user_id = message.text.split(maxsplit=1)[1]
|
|
if user_id.isdigit():
|
|
user_id = int(user_id)
|
|
db.remaiuser(user_id)
|
|
await message.edit_text("<b>User ID Removed</b>")
|
|
restart()
|
|
else:
|
|
await message.edit_text("<b>User ID is invalid.</b>")
|
|
return
|
|
else:
|
|
await message.edit_text(f"<b>Usage: </b><code>{prefix}remai [user_id]</code>")
|
|
return
|
|
|
|
|
|
@Client.on_message(filters.user(users=chatai_users) & filters.text)
|
|
async def chatbot(_, message: Message):
|
|
user_id = message.chat.id
|
|
|
|
if user_id in chatai_users:
|
|
pass
|
|
else:
|
|
return
|
|
try:
|
|
await message.reply_chat_action(enums.ChatAction.TYPING)
|
|
|
|
chat_history = db.get_chat_history(user_id)
|
|
|
|
prompt = message.text
|
|
|
|
db.add_chat_history(user_id, {"role": "USER", "message": prompt})
|
|
|
|
response = co.chat(
|
|
chat_history=chat_history,
|
|
model="command-r-plus",
|
|
message=prompt,
|
|
temperature=0.3,
|
|
connectors=[{"id": "web-search", "options": {"site": "wikipedia.com"}}],
|
|
prompt_truncation="AUTO",
|
|
)
|
|
|
|
db.add_chat_history(user_id, {"role": "CHATBOT", "message": response.text})
|
|
|
|
await message.reply_text(
|
|
f"{response.text}", parse_mode=enums.ParseMode.MARKDOWN
|
|
)
|
|
|
|
except Exception as e:
|
|
await message.reply_text(f"An error occurred: {format_exc(e)}")
|
|
|
|
|
|
@Client.on_message(filters.command("chatoff", prefix) & filters.me)
|
|
async def chatoff(_, message: Message):
|
|
db.remove("core.chatbot", "chatai_users")
|
|
await message.reply_text("<b>ChatBot is off now</b>")
|
|
restart()
|
|
|
|
|
|
@Client.on_message(filters.command("listai", prefix) & filters.me)
|
|
async def listai(_, message: Message):
|
|
await message.edit_text(
|
|
f"<b>User ID's Currently in AI ChatBot List:</b>\n <code>{chatai_users}</code>"
|
|
)
|
|
|
|
|
|
modules_help["chatbot"] = {
|
|
"addai [user_id]*": "Add A user to AI ChatBot List",
|
|
"remai [user_id]*": "Remove A user from AI ChatBot List",
|
|
"listai": "List A user from AI ChatBot List",
|
|
"chatoff": "Turn off AI ChatBot",
|
|
}
|