87 lines
2.6 KiB
Python
87 lines
2.6 KiB
Python
from typing import Optional
|
|
|
|
from aiogram import Router, F
|
|
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
|
|
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
|
|
]
|
|
|
|
|
|
@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()
|
|
|
|
if message.content_type == ContentType.TEXT and message.text.find('@' + bot_user.username) != -1:
|
|
message_text = message.text.replace('@' + bot_user.username, bot_user.first_name)
|
|
elif message.reply_to_message and message.reply_to_message.from_user.id == bot_user.id:
|
|
message_text = message.text
|
|
else:
|
|
return
|
|
|
|
if message.from_user.first_name and message.from_user.last_name:
|
|
user_name = "{} {}".format(message.from_user.first_name, message.from_user.last_name)
|
|
elif message.from_user.first_name:
|
|
user_name = message.from_user.first_name
|
|
elif message.from_user.username:
|
|
user_name = message.from_user.username
|
|
else:
|
|
user_name = str(message.from_user.id)
|
|
|
|
chat_prompt = chat['ai_prompt']
|
|
|
|
global agent
|
|
if agent is None:
|
|
# noinspection PyUnresolvedReferences
|
|
agent = AiAgent(message.bot.config['openrouter_token'])
|
|
|
|
await message.bot.send_chat_action(chat_id, 'typing')
|
|
await message.reply(await agent.get_reply(chat_id, chat_prompt, user_name, message_text))
|
|
|
|
|
|
def clear_ai_chat_context(chat_id: int):
|
|
global agent
|
|
if agent is not None:
|
|
agent.clear_chat_context(chat_id)
|