8fe355ed0c
- Updated README title to show OVERSPEND1 FORK - Changed maintainer credit to @overspend1 - Updated alive command to show @overspend1 as creator instead of Meliodas
158 lines
4.7 KiB
Python
158 lines
4.7 KiB
Python
import pickle
|
|
from io import BytesIO
|
|
|
|
from google.genai.chats import AsyncChat
|
|
from pyrogram.enums import ChatType, ParseMode
|
|
|
|
from app import BOT, Convo, Message, bot
|
|
from app.plugins.ai.gemini import AIConfig, Response, async_client
|
|
from app.plugins.ai.gemini.utils import create_prompts, run_basic_check
|
|
|
|
|
|
@bot.add_cmd(cmd="aic")
|
|
@run_basic_check
|
|
async def ai_chat(bot: BOT, message: Message):
|
|
"""
|
|
CMD: AICHAT
|
|
INFO: Have a Conversation with Gemini AI.
|
|
FLAGS:
|
|
"-s": use search
|
|
"-i": use image gen/edit mode
|
|
-a: audio output
|
|
-sp: multi speaker output
|
|
USAGE:
|
|
.aic hello
|
|
keep replying to AI responses with text | media [no need to reply in DM]
|
|
After 5 mins of Idle bot will export history and stop chat.
|
|
use .load_history to continue
|
|
|
|
"""
|
|
chat = async_client.chats.create(**AIConfig.get_kwargs(message.flags))
|
|
await do_convo(chat=chat, message=message)
|
|
|
|
|
|
@bot.add_cmd(cmd="lh")
|
|
@run_basic_check
|
|
async def history_chat(bot: BOT, message: Message):
|
|
"""
|
|
CMD: LOAD_HISTORY
|
|
INFO: Load a Conversation with Gemini AI from previous session.
|
|
USAGE:
|
|
.lh {question} [reply to history document]
|
|
"""
|
|
reply = message.replied
|
|
|
|
if not message.input:
|
|
await message.reply(f"Ask a question along with {message.trigger}{message.cmd}")
|
|
return
|
|
|
|
try:
|
|
assert reply.document.file_name == "AI_Chat_History.pkl"
|
|
except (AssertionError, AttributeError):
|
|
await message.reply("Reply to a Valid History file.")
|
|
return
|
|
|
|
resp = await message.reply("`Loading History...`")
|
|
|
|
doc = await reply.download(in_memory=True)
|
|
doc.seek(0)
|
|
pickle.load(doc)
|
|
|
|
await resp.edit("__History Loaded... Resuming chat__")
|
|
|
|
chat = async_client.chats.create(**AIConfig.get_kwargs(message.flags))
|
|
await do_convo(chat=chat, message=message)
|
|
|
|
|
|
CONVO_CACHE: dict[str, Convo] = {}
|
|
|
|
|
|
async def do_convo(chat: AsyncChat, message: Message):
|
|
chat_id = message.chat.id
|
|
|
|
old_conversation = CONVO_CACHE.get(message.unique_chat_user_id)
|
|
|
|
if old_conversation in Convo.CONVO_DICT[chat_id]:
|
|
Convo.CONVO_DICT[chat_id].remove(old_conversation)
|
|
|
|
if message.chat.type in (ChatType.PRIVATE, ChatType.BOT):
|
|
reply_to_user_id = None
|
|
else:
|
|
reply_to_user_id = message._client.me.id
|
|
|
|
conversation_object = Convo(
|
|
client=message._client,
|
|
chat_id=chat_id,
|
|
timeout=300,
|
|
check_for_duplicates=False,
|
|
from_user=message.from_user.id,
|
|
reply_to_user_id=reply_to_user_id,
|
|
)
|
|
|
|
CONVO_CACHE[message.unique_chat_user_id] = conversation_object
|
|
|
|
try:
|
|
async with conversation_object:
|
|
prompt = await create_prompts(message)
|
|
reply_to_id = message.id
|
|
|
|
while True:
|
|
ai_response = await chat.send_message(prompt)
|
|
prompt_message = await send_and_get_resp(
|
|
convo_obj=conversation_object,
|
|
response=ai_response,
|
|
reply_to_id=reply_to_id,
|
|
)
|
|
|
|
try:
|
|
prompt = await create_prompts(prompt_message, is_chat=True, check_size=False)
|
|
except Exception as e:
|
|
prompt_message = await send_and_get_resp(
|
|
conversation_object, str(e), reply_to_id=reply_to_id
|
|
)
|
|
prompt = await create_prompts(prompt_message, is_chat=True, check_size=False)
|
|
|
|
reply_to_id = prompt_message.id
|
|
|
|
except TimeoutError:
|
|
await export_history(chat, message)
|
|
finally:
|
|
CONVO_CACHE.pop(message.unique_chat_user_id, 0)
|
|
|
|
|
|
async def send_and_get_resp(
|
|
convo_obj: Convo,
|
|
response,
|
|
reply_to_id: int | None = None,
|
|
) -> Message:
|
|
|
|
response = Response(response)
|
|
|
|
if text := response.text():
|
|
await convo_obj.send_message(
|
|
text=f"**>\n•><**\n{text}",
|
|
reply_to_id=reply_to_id,
|
|
parse_mode=ParseMode.MARKDOWN,
|
|
disable_preview=True,
|
|
)
|
|
|
|
if response.image:
|
|
await convo_obj.send_photo(photo=response.image_file, reply_to_id=reply_to_id)
|
|
|
|
if response.audio:
|
|
await convo_obj.send_voice(
|
|
voice=response.audio_file,
|
|
waveform=response.audio_file.waveform,
|
|
reply_to_id=reply_to_id,
|
|
duration=response.audio_file.duration,
|
|
)
|
|
|
|
return await convo_obj.get_response()
|
|
|
|
|
|
async def export_history(chat: AsyncChat, message: Message):
|
|
doc = BytesIO(pickle.dumps(chat._curated_history))
|
|
doc.name = "AI_Chat_History.pkl"
|
|
caption = Response(await chat.send_message("Summarize our Conversation into one line.")).text()
|
|
await bot.send_document(chat_id=message.from_user.id, document=doc, caption=caption)
|