152 lines
4.9 KiB
Python
152 lines
4.9 KiB
Python
from typing import Dict, List, Optional
|
||
|
||
from aiogram import Router, F
|
||
from aiogram.types import Message
|
||
from aiogram.types.user import User
|
||
from aiogram.enums.content_type import ContentType
|
||
|
||
from openrouter import OpenRouter
|
||
from dataclasses import dataclass
|
||
|
||
import utils
|
||
import tg.tg_database as database
|
||
|
||
router = Router()
|
||
|
||
|
||
@dataclass
|
||
class ChatContext:
|
||
def __init__(self, max_messages: int):
|
||
self.max_messages: int = max_messages
|
||
self.messages: List[Dict[str, str]] = []
|
||
|
||
def add_message(self, role: str, content: str):
|
||
if len(self.messages) == self.max_messages:
|
||
# Всегда сохраняем в контексте системное сообщение
|
||
self.messages.pop(1)
|
||
self.messages.append({"role": role, "content": content})
|
||
|
||
def get_messages_for_api(self) -> List[Dict[str, str]]:
|
||
return self.messages
|
||
|
||
|
||
chat_contexts: Dict[int, ChatContext] = {}
|
||
bot_user: Optional[User] = None
|
||
|
||
system_prompt = """
|
||
Ты - помощник в групповом чате Telegram.
|
||
Отвечай на вопросы и поддерживай контекст беседы.
|
||
Ты не можешь обсуждать политику и религию.
|
||
Сообщения пользователей будут приходить в следующем формате: '[Имя]: текст сообщения'
|
||
При ответе НЕ нужно указывать пользователя, которому он предназначен.
|
||
"""
|
||
|
||
|
||
def get_ai_chat_context(chat_id: int) -> ChatContext:
|
||
"""Get or create chat context for a specific chat"""
|
||
if chat_id not in chat_contexts:
|
||
chat_contexts[chat_id] = ChatContext(10)
|
||
chat = database.DB.get_chat(chat_id)
|
||
prompt = system_prompt
|
||
if chat['ai_prompt'] is not None:
|
||
prompt += '\n\n' + chat['ai_prompt']
|
||
chat_contexts[chat_id].add_message(role="system", content=prompt)
|
||
return chat_contexts[chat_id]
|
||
|
||
|
||
async def ai_message_handler(message: Message):
|
||
chat_id = message.chat.id
|
||
|
||
# Extract user information and message content
|
||
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)
|
||
|
||
bot_mention = '@' + bot_user.username
|
||
message_text = message.text.replace(bot_mention, bot_user.first_name)
|
||
|
||
context = get_ai_chat_context(chat_id)
|
||
context.add_message(
|
||
role="user",
|
||
content=f"[{user_name}]: {message_text}"
|
||
)
|
||
|
||
# noinspection PyUnresolvedReferences
|
||
api_key = message.bot.config['openrouter_token']
|
||
|
||
client = OpenRouter(api_key=api_key)
|
||
messages_for_api = context.get_messages_for_api()
|
||
|
||
await message.bot.send_chat_action(chat_id, 'typing')
|
||
|
||
try:
|
||
# Get response from OpenRouter
|
||
response = await client.chat.send_async(
|
||
model="meta-llama/llama-3.3-70b-instruct:free",
|
||
messages=messages_for_api,
|
||
max_tokens=500,
|
||
temperature=0.7
|
||
)
|
||
|
||
# Extract AI response
|
||
ai_response = response.choices[0].message.content
|
||
|
||
# Add AI response to context
|
||
context.add_message(role="assistant", content=ai_response)
|
||
|
||
# Send response back to chat
|
||
await message.reply(ai_response)
|
||
|
||
except Exception as e:
|
||
print(f"Error processing message: {e}")
|
||
await message.reply("Извините, при обработке запроса произошла ошибка.")
|
||
|
||
|
||
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()
|
||
|
||
bot_mention = '@' + bot_user.username
|
||
if message.content_type == ContentType.TEXT and message.text.find(bot_mention) != -1:
|
||
await ai_message_handler(message)
|