7ba6bc44f2
- ruff check --fix: 210 auto-fixed errors (import sorting, trailing whitespace, unused imports, f-string fixups, deprecated annotations) - ruff format: 104 files reformatted to consistent style - 268 non-auto-fixable issues remain (S113 requests timeout, etc.)
101 lines
3.1 KiB
Python
101 lines
3.1 KiB
Python
from pyrogram import Client, enums, filters
|
|
from pyrogram.types import Message
|
|
from utils.config import cohere_key
|
|
from utils.db import db
|
|
from utils.misc import modules_help, prefix
|
|
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',
|
|
}
|