71 lines
2.3 KiB
Python
71 lines
2.3 KiB
Python
from functools import partial
|
||
|
||
from vkbottle.bot import Message
|
||
from vkbottle.dispatch.rules.base import RegexRule
|
||
from vkbottle.framework.labeler import BotLabeler
|
||
|
||
import ai_agent
|
||
import utils
|
||
from messages import *
|
||
|
||
import vk.vk_database as database
|
||
from vk.utils import *
|
||
|
||
labeler = BotLabeler()
|
||
|
||
|
||
@labeler.private_message(text="!старт")
|
||
async def start_handler(message: Message):
|
||
bot_id = get_bot_id(message.ctx_api)
|
||
chat_id = message.peer_id
|
||
database.DB.create_chat_if_not_exists(bot_id, chat_id)
|
||
database.DB.chat_update(bot_id, chat_id, active=1)
|
||
await message.answer("Привет!")
|
||
|
||
|
||
@labeler.private_message(RegexRule(r"^!личность ((?:.|\n)+)"))
|
||
async def set_prompt_handler(message: Message, match):
|
||
bot_id = get_bot_id(message.ctx_api)
|
||
chat_id = message.peer_id
|
||
database.DB.create_chat_if_not_exists(bot_id, chat_id)
|
||
|
||
database.DB.chat_update(bot_id, chat_id, ai_prompt=match[0])
|
||
await message.answer("Личность ИИ изменена.")
|
||
|
||
|
||
@labeler.private_message(text="!сброс")
|
||
async def reset_context_handler(message: Message):
|
||
bot_id = get_bot_id(message.ctx_api)
|
||
chat_id = message.peer_id
|
||
database.DB.create_chat_if_not_exists(bot_id, chat_id)
|
||
|
||
ai_agent.agent.clear_chat_context(bot_id, chat_id)
|
||
await message.answer("Контекст очищен.")
|
||
|
||
|
||
@labeler.private_message()
|
||
async def any_message_handler(message: Message):
|
||
bot_id = get_bot_id(message.ctx_api)
|
||
chat_id = message.peer_id
|
||
|
||
try:
|
||
ai_message = await create_ai_message(message)
|
||
except utils.UnsupportedContentType:
|
||
await message.answer(MESSAGE_UNSUPPORTED_CONTENT_TYPE)
|
||
return
|
||
|
||
answer: ai_agent.Message
|
||
success: bool
|
||
answer, success = await utils.run_with_progress(
|
||
partial(ai_agent.agent.get_private_chat_reply, bot_id, chat_id, ai_message),
|
||
partial(message.ctx_api.messages.set_activity, peer_id=chat_id, type='typing'),
|
||
interval=4)
|
||
|
||
if answer.image is not None:
|
||
photo = await upload_photo(answer.image, chat_id=chat_id, api=message.ctx_api)
|
||
answer_id = (await message.answer(answer.text, attachment=photo)).conversation_message_id
|
||
else:
|
||
answer_id = (await message.answer(answer.text)).message_id
|
||
|
||
if success:
|
||
ai_agent.agent.set_last_response_id(bot_id, chat_id, answer_id)
|