Добавлена функциональность чата с ИИ через OpenRouter.
This commit is contained in:
parent
c9fd78e438
commit
651f755b2d
6 changed files with 216 additions and 0 deletions
|
|
@ -13,6 +13,7 @@ async def main() -> None:
|
||||||
print('Конфигурация загружена.')
|
print('Конфигурация загружена.')
|
||||||
|
|
||||||
bot = Bot(token=config['api_token'])
|
bot = Bot(token=config['api_token'])
|
||||||
|
bot.config = config
|
||||||
|
|
||||||
dp = Dispatcher()
|
dp = Dispatcher()
|
||||||
dp.include_router(handlers.router)
|
dp.include_router(handlers.router)
|
||||||
|
|
|
||||||
|
|
@ -1,12 +1,112 @@
|
||||||
|
from typing import Dict, List, Optional
|
||||||
|
|
||||||
from aiogram import Router, F
|
from aiogram import Router, F
|
||||||
from aiogram.types import Message
|
from aiogram.types import Message
|
||||||
|
from aiogram.types.user import User
|
||||||
from aiogram.enums.content_type import ContentType
|
from aiogram.enums.content_type import ContentType
|
||||||
|
|
||||||
|
from openrouter import OpenRouter
|
||||||
|
from dataclasses import dataclass
|
||||||
|
|
||||||
import utils
|
import utils
|
||||||
import tg.tg_database as database
|
import tg.tg_database as database
|
||||||
|
|
||||||
router = Router()
|
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] = [
|
ACCEPTED_CONTENT_TYPES: list[ContentType] = [
|
||||||
ContentType.TEXT,
|
ContentType.TEXT,
|
||||||
ContentType.ANIMATION,
|
ContentType.ANIMATION,
|
||||||
|
|
@ -42,3 +142,11 @@ async def any_message_handler(message: Message):
|
||||||
database.DB.create_user_if_not_exists(chat_id, 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_set_last_message(chat_id, user_id, utils.posix_time())
|
||||||
database.DB.user_increment_messages(chat_id, user_id)
|
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)
|
||||||
|
|
|
||||||
|
|
@ -11,6 +11,7 @@ class TgDatabase(database.BasicDatabase):
|
||||||
"active" INTEGER NOT NULL DEFAULT 0,
|
"active" INTEGER NOT NULL DEFAULT 0,
|
||||||
"rules" TEXT,
|
"rules" TEXT,
|
||||||
"greeting_join" TEXT,
|
"greeting_join" TEXT,
|
||||||
|
"ai_prompt" TEXT,
|
||||||
PRIMARY KEY("id"))
|
PRIMARY KEY("id"))
|
||||||
""")
|
""")
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -12,6 +12,7 @@ if __name__ == '__main__':
|
||||||
print('Конфигурация загружена.')
|
print('Конфигурация загружена.')
|
||||||
|
|
||||||
bot = VkBot(config['api_token'], labeler=handlers.labeler)
|
bot = VkBot(config['api_token'], labeler=handlers.labeler)
|
||||||
|
bot.api.config = config
|
||||||
bot.loop_wrapper.on_startup.append(tasks.startup_task(bot.api))
|
bot.loop_wrapper.on_startup.append(tasks.startup_task(bot.api))
|
||||||
bot.loop_wrapper.add_task(tasks.daily_maintenance_task(bot.api))
|
bot.loop_wrapper.add_task(tasks.daily_maintenance_task(bot.api))
|
||||||
bot.run_forever()
|
bot.run_forever()
|
||||||
|
|
|
||||||
|
|
@ -1,5 +1,11 @@
|
||||||
|
from typing import Dict, List, Optional
|
||||||
|
|
||||||
from vkbottle.bot import Message
|
from vkbottle.bot import Message
|
||||||
from vkbottle.framework.labeler import BotLabeler
|
from vkbottle.framework.labeler import BotLabeler
|
||||||
|
from vkbottle_types.codegen.objects import GroupsGroup
|
||||||
|
|
||||||
|
from openrouter import OpenRouter
|
||||||
|
from dataclasses import dataclass
|
||||||
|
|
||||||
import utils
|
import utils
|
||||||
import vk.vk_database as database
|
import vk.vk_database as database
|
||||||
|
|
@ -7,6 +13,96 @@ import vk.vk_database as database
|
||||||
labeler = BotLabeler()
|
labeler = BotLabeler()
|
||||||
|
|
||||||
|
|
||||||
|
@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[GroupsGroup] = 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.peer_id
|
||||||
|
|
||||||
|
# Extract user information and message content
|
||||||
|
user = await message.ctx_api.users.get(user_ids=[message.from_id])
|
||||||
|
if len(user) == 1:
|
||||||
|
user_name = "{} {}".format(user[0].first_name, user[0].last_name)
|
||||||
|
else:
|
||||||
|
user_name = '@id' + str(message.from_id)
|
||||||
|
|
||||||
|
bot_mention = '@' + bot_user.screen_name
|
||||||
|
message_text = message.text.replace(bot_mention, bot_user.name)
|
||||||
|
|
||||||
|
context = get_ai_chat_context(chat_id)
|
||||||
|
context.add_message(
|
||||||
|
role="user",
|
||||||
|
content=f"[{user_name}]: {message_text}"
|
||||||
|
)
|
||||||
|
|
||||||
|
# noinspection PyUnresolvedReferences
|
||||||
|
api_key = message.ctx_api.config['openrouter_token']
|
||||||
|
|
||||||
|
client = OpenRouter(api_key=api_key)
|
||||||
|
messages_for_api = context.get_messages_for_api()
|
||||||
|
|
||||||
|
await message.ctx_api.messages.set_activity(peer_id=chat_id, type='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("Извините, при обработке запроса произошла ошибка.")
|
||||||
|
|
||||||
|
|
||||||
# Обычные сообщения (не команды и не действия)
|
# Обычные сообщения (не команды и не действия)
|
||||||
@labeler.chat_message()
|
@labeler.chat_message()
|
||||||
async def any_message_handler(message: Message):
|
async def any_message_handler(message: Message):
|
||||||
|
|
@ -27,3 +123,11 @@ async def any_message_handler(message: Message):
|
||||||
database.DB.create_user_if_not_exists(chat_id, 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_set_last_message(chat_id, user_id, utils.posix_time())
|
||||||
database.DB.user_increment_messages(chat_id, user_id)
|
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]
|
||||||
|
|
||||||
|
bot_mention = '@' + bot_user.screen_name
|
||||||
|
if message.text is not None and message.text.find(bot_mention) != -1:
|
||||||
|
await ai_message_handler(message)
|
||||||
|
|
|
||||||
|
|
@ -13,6 +13,7 @@ class VkDatabase(database.BasicDatabase):
|
||||||
"greeting_join" TEXT,
|
"greeting_join" TEXT,
|
||||||
"greeting_rejoin" TEXT,
|
"greeting_rejoin" TEXT,
|
||||||
"birthday_message" TEXT,
|
"birthday_message" TEXT,
|
||||||
|
"ai_prompt" TEXT,
|
||||||
PRIMARY KEY("id"))
|
PRIMARY KEY("id"))
|
||||||
""")
|
""")
|
||||||
|
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue