69 lines
2.6 KiB
Python
69 lines
2.6 KiB
Python
# admin_bot/main.py
|
||
import asyncio
|
||
import toml
|
||
import json
|
||
import nats
|
||
from nats.errors import TimeoutError
|
||
from loguru import logger
|
||
from aiogram import Bot, Dispatcher, types
|
||
from aiogram.types import Message
|
||
from aiogram.enums import ParseMode
|
||
from aiogram.fsm.storage.memory import MemoryStorage
|
||
from aiogram.client.default import DefaultBotProperties
|
||
|
||
# Загрузка конфигурации
|
||
config = toml.load("config.toml")
|
||
|
||
BOT_TOKEN = config["telegram"]["bot_token"]
|
||
ADMIN_ID = config["telegram"]["admin_id"]
|
||
NATS_URL = config["nats"]["url"]
|
||
NATS_SUBJECT = config["nats"].get("subject", "llm.inbox")
|
||
|
||
# Инициализация
|
||
bot = Bot(token=BOT_TOKEN, default=DefaultBotProperties(parse_mode=ParseMode.HTML))
|
||
dp = Dispatcher(storage=MemoryStorage())
|
||
|
||
# Подключение к NATS
|
||
nats_client = None
|
||
|
||
@dp.message()
|
||
async def handle_message(message: Message):
|
||
logger.info(f"Получено сообщение от {message.from_user.id}: {message.text}")
|
||
|
||
if message.from_user.id != ADMIN_ID:
|
||
logger.warning(f"Пользователь {message.from_user.id} не является админом. Доступ запрещён.")
|
||
await message.answer("⛔️ У вас нет прав на использование этого бота.")
|
||
return
|
||
|
||
if nats_client is None:
|
||
logger.error("Попытка отправки в NATS до инициализации клиента.")
|
||
await message.answer("🚫 NATS клиент не инициализирован.")
|
||
return
|
||
|
||
payload = {
|
||
"user_id": message.from_user.id,
|
||
"text": message.text
|
||
}
|
||
|
||
try:
|
||
response = await nats_client.request(NATS_SUBJECT.encode(), json.dumps(payload).encode(), timeout=30)
|
||
data = json.loads(response.data.decode())
|
||
await message.answer(data.get("response", "[no response]"))
|
||
except TimeoutError:
|
||
await message.answer("⌛ Превышен таймаут ожидания ответа.")
|
||
except Exception as e:
|
||
logger.exception("Ошибка при получении ответа от LLM")
|
||
await message.answer("❌ Ошибка при получении ответа.")
|
||
|
||
|
||
async def main():
|
||
global nats_client
|
||
logger.info("Подключение к NATS...")
|
||
nats_client = await nats.connect(servers=[NATS_URL])
|
||
logger.success("Подключение к NATS успешно.")
|
||
logger.info("Запуск polling Telegram-бота...")
|
||
await dp.start_polling(bot)
|
||
|
||
if __name__ == "__main__":
|
||
logger.info("Запуск бота...")
|
||
asyncio.run(main())
|