122 lines
4.1 KiB
Python
122 lines
4.1 KiB
Python
from functools import partial
|
|
from typing import Optional
|
|
|
|
from aiogram import Router, F, Bot
|
|
from aiogram.types import Message
|
|
from aiogram.types.user import User
|
|
from aiogram.enums.content_type import ContentType
|
|
|
|
import utils
|
|
from ai_agent import AiAgent, AiMessage
|
|
import tg.tg_database as database
|
|
|
|
router = Router()
|
|
|
|
agent: Optional[AiAgent] = None
|
|
bot_user: Optional[User] = None
|
|
|
|
ACCEPTED_CONTENT_TYPES: list[ContentType] = [
|
|
ContentType.TEXT,
|
|
ContentType.ANIMATION,
|
|
ContentType.AUDIO,
|
|
ContentType.DOCUMENT,
|
|
ContentType.PAID_MEDIA,
|
|
ContentType.PHOTO,
|
|
ContentType.STICKER,
|
|
ContentType.STORY,
|
|
ContentType.VIDEO,
|
|
ContentType.VIDEO_NOTE,
|
|
ContentType.VOICE,
|
|
ContentType.CHECKLIST,
|
|
ContentType.CONTACT,
|
|
ContentType.POLL,
|
|
ContentType.VENUE,
|
|
ContentType.LOCATION
|
|
]
|
|
|
|
|
|
async def get_user_name_for_ai(user: User):
|
|
if user.first_name and user.last_name:
|
|
return "{} {}".format(user.first_name, user.last_name)
|
|
elif user.first_name:
|
|
return user.first_name
|
|
elif user.username:
|
|
return user.username
|
|
else:
|
|
return str(user.id)
|
|
|
|
|
|
@router.message(F.content_type.in_(ACCEPTED_CONTENT_TYPES))
|
|
async def any_message_handler(message: Message):
|
|
chat_id = message.chat.id
|
|
chat = database.DB.create_chat_if_not_exists(chat_id)
|
|
if chat['active'] == 0:
|
|
return
|
|
|
|
# Игнорировать ботов
|
|
if message.from_user.is_bot:
|
|
return
|
|
|
|
user_id = message.from_user.id
|
|
database.DB.create_user_if_not_exists(chat_id, user_id)
|
|
database.DB.user_set_last_message(chat_id, user_id, utils.posix_time())
|
|
database.DB.user_increment_messages(chat_id, user_id)
|
|
|
|
global bot_user
|
|
if bot_user is None:
|
|
bot_user = await message.bot.get_me()
|
|
|
|
ai_message = AiMessage()
|
|
ai_fwd_messages: list[AiMessage] = []
|
|
|
|
# Ответ на сообщение бота
|
|
if message.reply_to_message and message.reply_to_message.from_user.id == bot_user.id:
|
|
ai_message.text = message.text
|
|
else:
|
|
# Сообщение содержит @bot_username
|
|
bot_username_mention = '@' + bot_user.username
|
|
if message.content_type == ContentType.TEXT and message.text.find(bot_username_mention) != -1:
|
|
ai_message.text = message.text.replace(bot_username_mention, bot_user.first_name)
|
|
else:
|
|
return
|
|
|
|
if message.reply_to_message and message.reply_to_message.content_type == ContentType.TEXT:
|
|
ai_fwd_messages = [
|
|
AiMessage(user_name=await get_user_name_for_ai(message.reply_to_message.from_user),
|
|
text=message.reply_to_message.text)]
|
|
|
|
ai_message.user_name = await get_user_name_for_ai(message.from_user)
|
|
chat_prompt = chat['ai_prompt']
|
|
|
|
global agent
|
|
if agent is None:
|
|
# noinspection PyUnresolvedReferences
|
|
agent = AiAgent(message.bot.config['openrouter_token'])
|
|
|
|
await message.reply(
|
|
await utils.run_with_progress(partial(agent.get_reply, chat_id, chat_prompt, ai_message, ai_fwd_messages),
|
|
partial(message.bot.send_chat_action, chat_id, 'typing'),
|
|
interval=4))
|
|
|
|
|
|
def clear_ai_chat_context(chat_id: int):
|
|
global agent
|
|
if agent is not None:
|
|
agent.clear_chat_context(chat_id)
|
|
|
|
|
|
async def get_ai_reply(bot: Bot, chat_id, user: User, message: str, fwd_user: User, fwd_message: str) -> str:
|
|
chat = database.DB.create_chat_if_not_exists(chat_id)
|
|
chat_prompt = chat['ai_prompt']
|
|
|
|
ai_message = AiMessage(user_name=await get_user_name_for_ai(user), text=message)
|
|
ai_fwd_messages = [AiMessage(user_name=await get_user_name_for_ai(fwd_user), text=fwd_message)]
|
|
|
|
global agent
|
|
if agent is None:
|
|
# noinspection PyUnresolvedReferences
|
|
agent = AiAgent(bot.config['openrouter_token'])
|
|
|
|
return await utils.run_with_progress(partial(agent.get_reply, chat_id, chat_prompt, ai_message, ai_fwd_messages),
|
|
partial(bot.send_chat_action(chat_id, 'typing')),
|
|
interval=4)
|