107 lines
3.9 KiB
Python
107 lines
3.9 KiB
Python
import re
|
||
from typing import Optional
|
||
|
||
from vkbottle import API
|
||
from vkbottle.bot import Message
|
||
from vkbottle.framework.labeler import BotLabeler
|
||
from vkbottle_types.codegen.objects import GroupsGroup
|
||
|
||
import utils
|
||
from ai_agent import AiAgent, AiMessage
|
||
import vk.vk_database as database
|
||
|
||
labeler = BotLabeler()
|
||
|
||
agent: Optional[AiAgent] = None
|
||
bot_user: Optional[GroupsGroup] = None
|
||
|
||
|
||
async def get_user_name_for_ai(api: API, user_id: int):
|
||
user = await api.users.get(user_ids=[user_id])
|
||
if len(user) == 1:
|
||
return "{} {}".format(user[0].first_name, user[0].last_name)
|
||
else:
|
||
return '@id' + str(user_id)
|
||
|
||
|
||
# Обычные сообщения (не команды и не действия)
|
||
@labeler.chat_message()
|
||
async def any_message_handler(message: Message):
|
||
chat_id = message.peer_id
|
||
chat = database.DB.create_chat_if_not_exists(chat_id)
|
||
if chat['active'] == 0:
|
||
return
|
||
|
||
# Игнорировать ботов
|
||
if message.from_id < 0:
|
||
return
|
||
|
||
# Не учитывать служебные сообщения
|
||
if message.action is not None:
|
||
return
|
||
|
||
user_id = message.from_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.ctx_api.groups.get_by_id()).groups[0]
|
||
|
||
ai_message = AiMessage()
|
||
ai_fwd_message: Optional[AiMessage] = None
|
||
|
||
# Ответ на сообщение бота
|
||
if message.reply_message and message.reply_message.from_id == -bot_user.id:
|
||
ai_message.text = message.text
|
||
else:
|
||
# Сообщение содержит @bot_username
|
||
bot_username_mention = '@' + bot_user.screen_name
|
||
if message.text is not None and message.text.find(bot_username_mention) != -1:
|
||
ai_message.text = message.text.replace(bot_username_mention, bot_user.name)
|
||
else:
|
||
# Сообщение содержит [club<bot_id>|<some_name>]
|
||
pattern = r"\[club" + str(bot_user.id) + r"\|(.+)]"
|
||
if message.text is not None and re.search(pattern, message.text) is not None:
|
||
ai_message.text = re.sub(pattern, r'\1', message.text)
|
||
else:
|
||
return
|
||
|
||
if message.reply_message and message.reply_message.text is not None:
|
||
ai_fwd_message = (
|
||
AiMessage(user_name=await get_user_name_for_ai(message.ctx_api, message.reply_message.from_id),
|
||
text=message.reply_message.text))
|
||
|
||
ai_message.user_name = await get_user_name_for_ai(message.ctx_api, message.from_id)
|
||
chat_prompt = chat['ai_prompt']
|
||
|
||
global agent
|
||
if agent is None:
|
||
# noinspection PyUnresolvedReferences
|
||
agent = AiAgent(message.ctx_api.config['openrouter_token'])
|
||
|
||
await message.ctx_api.messages.set_activity(peer_id=chat_id, type='typing')
|
||
await message.reply(await agent.get_reply(chat_id, chat_prompt, ai_message, ai_fwd_message))
|
||
|
||
|
||
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(api: API, chat_id, user_id: int, message: str, fwd_user_id: int, 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(api, user_id), text=message)
|
||
ai_fwd_message = AiMessage(user_name=await get_user_name_for_ai(api, fwd_user_id), text=fwd_message)
|
||
|
||
global agent
|
||
if agent is None:
|
||
# noinspection PyUnresolvedReferences
|
||
agent = AiAgent(api.config['openrouter_token'])
|
||
|
||
await api.messages.set_activity(peer_id=chat_id, type='typing')
|
||
return await agent.get_reply(chat_id, chat_prompt, ai_message, ai_fwd_message)
|