190 lines
8.9 KiB
Python
190 lines
8.9 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
XMPP клиент – главный класс Bot.
|
||
Теперь это тонкий клиент, который только общается с RAG-сервером.
|
||
Не содержит БД, Qdrant, GigaChat – всё на сервере.
|
||
|
||
ИСПОЛЬЗУЕТ:
|
||
- AppConfig из rag.config_models для единой конфигурации.
|
||
- RAGClient с передачей API-ключа для аутентификации.
|
||
- Комнаты загружаются с сервера при старте.
|
||
"""
|
||
|
||
import asyncio
|
||
import logging
|
||
from datetime import datetime
|
||
from slixmpp import ClientXMPP
|
||
|
||
# Импорты из нового пакета rag для конфигурации
|
||
from rag.config_models import AppConfig
|
||
|
||
# Остальные импорты из core
|
||
from core.services.file_service import FileService
|
||
from core.handlers.metrics_handler import MetricsHandler
|
||
from core.rag_client import RAGClient
|
||
from core.commands.registry import register_commands
|
||
from core.handlers.message_handler import MessageHandler
|
||
from core.handlers.file_handler import FileHandler
|
||
from core.handlers.room_handler import RoomHandler
|
||
from core.handlers.subscription_handler import SubscriptionHandler
|
||
from core.handlers.health_handler import HealthHandler
|
||
from core.workers.indexing_worker import IndexingWorker
|
||
from core.functions.file_processor import handle_ai_file_request as _handle_ai_file_request
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class Bot(ClientXMPP):
|
||
"""
|
||
Основной класс XMPP бота.
|
||
Тонкий клиент: вся RAG-логика на сервере.
|
||
"""
|
||
|
||
def __init__(self, config: AppConfig):
|
||
full_jid = f"{config.jid}/{config.xmpp_resource}"
|
||
super().__init__(full_jid, config.xmpp_password)
|
||
self.config = config
|
||
|
||
# ----- 1. Локальные сервисы (только для файлов и метрик) -----
|
||
self.files = FileService(config)
|
||
self.metrics_handler = MetricsHandler()
|
||
|
||
# ----- 2. Загрузка промптов (остаются локально) -----
|
||
self.ai_prompts = config.prompts_content.copy() if config.prompts_content else {}
|
||
if not self.ai_prompts:
|
||
logger.warning("Промпты не загружены из конфига, пробуем загрузить из файлов вручную...")
|
||
prompts_dir = config.prompts_dir_path
|
||
if prompts_dir and prompts_dir.exists():
|
||
for prompt_name, filename in config.ai.prompts.items():
|
||
file_path = prompts_dir / filename
|
||
if file_path.exists():
|
||
try:
|
||
with open(file_path, 'r', encoding='utf-8') as f:
|
||
self.ai_prompts[prompt_name] = f.read()
|
||
logger.debug(f"Загружен промпт {prompt_name} из {file_path}")
|
||
except Exception as e:
|
||
logger.error(f"Не удалось загрузить промпт {prompt_name} из {file_path}: {e}")
|
||
self.ai_prompts[prompt_name] = ""
|
||
else:
|
||
logger.warning(f"Файл промпта {filename} не найден в {prompts_dir}")
|
||
self.ai_prompts[prompt_name] = ""
|
||
else:
|
||
logger.error("Каталог промптов не найден, промпты не загружены")
|
||
self.prompts = self.ai_prompts
|
||
|
||
# ----- 3. HTTP-клиент для RAG-сервера -----
|
||
rag_server_url = config.rag_server_url
|
||
rag_api_key = config.rag_api_key
|
||
self.rag_client = RAGClient(
|
||
server_url=rag_server_url,
|
||
api_key=rag_api_key,
|
||
timeout=config.http_download_timeout,
|
||
max_retries=3,
|
||
retry_delay=1.0
|
||
)
|
||
|
||
# ----- 4. Команды и состояния -----
|
||
self.commands = register_commands(self)
|
||
self._command_names = [cmd.name for cmd in self.commands.values() if cmd.name]
|
||
|
||
# Временные состояния (хранятся в памяти)
|
||
self.learning_mode = {}
|
||
self.last_file_texts = {}
|
||
self._last_file_path = {}
|
||
self.indexing_queue = None
|
||
self.rooms = [] # больше не используется, загружаем с сервера
|
||
|
||
# ----- 5. Хендлеры -----
|
||
self.message_handler = MessageHandler(self)
|
||
self.file_handler = FileHandler(self)
|
||
self.room_handler = RoomHandler(self)
|
||
self.subscription_handler = SubscriptionHandler(self)
|
||
self.health_handler = HealthHandler(self)
|
||
self.indexing_worker = IndexingWorker(self)
|
||
|
||
# ----- 6. XMPP плагины -----
|
||
self.register_plugin('xep_0030')
|
||
self.register_plugin('xep_0199')
|
||
self.register_plugin('xep_0045')
|
||
self.register_plugin('xep_0054')
|
||
self.register_plugin('xep_0172')
|
||
self.register_plugin('xep_0066')
|
||
self.register_plugin('xep_0363')
|
||
|
||
# ----- 7. Обработчики событий slixmpp -----
|
||
self.roster_ready = asyncio.Event()
|
||
self.add_event_handler("session_start", self.session_start)
|
||
self.add_event_handler("message", self.message_handler.handle)
|
||
self.add_event_handler("groupchat_invite", self.room_handler.accept_invite)
|
||
self.add_event_handler("roster_subscription_request", self.subscription_handler.handle_subscription)
|
||
self.add_event_handler('roster_update', self._on_roster_ready)
|
||
self.add_event_handler("muc::kicked", self.room_handler.on_kicked)
|
||
self.add_event_handler("muc::banned", self.room_handler.on_banned)
|
||
|
||
# Время старта бота (для команды !stats)
|
||
self.start_time = datetime.now()
|
||
|
||
def _on_roster_ready(self, *args):
|
||
self.roster_ready.set()
|
||
|
||
def is_message_encrypted(self, msg) -> bool:
|
||
return False
|
||
|
||
async def handle_ai_file_request(self, msg, ai_response: str, user_jid: str, room_jid: str = None, body: str = None):
|
||
await _handle_ai_file_request(self, msg, ai_response, user_jid, room_jid, body=body)
|
||
|
||
async def session_start(self, event):
|
||
"""
|
||
Обработчик успешного подключения и авторизации.
|
||
"""
|
||
self.send_presence(pshow='chat', pstatus=self.config.status_message)
|
||
self.get_roster()
|
||
logger.info(f"Бот {self.config.name} запущен, JID: {self.boundjid}")
|
||
|
||
# Ждём загрузки ростера (максимум 10 секунд)
|
||
try:
|
||
await asyncio.wait_for(self.roster_ready.wait(), timeout=10.0)
|
||
except asyncio.TimeoutError:
|
||
logger.warning("Таймаут ожидания ростера")
|
||
|
||
# Отправляем запросы на подписку администраторам
|
||
for admin_jid in self.config.admin_jids:
|
||
self.send_presence(pto=admin_jid, ptype='subscribe')
|
||
|
||
# Создаём очередь индексации и запускаем фоновый воркер
|
||
self.indexing_queue = asyncio.Queue()
|
||
asyncio.create_task(self.indexing_worker.run())
|
||
|
||
# Загружаем комнаты с сервера и входим в них
|
||
try:
|
||
rooms = await self.rag_client.get_rooms()
|
||
for room in rooms:
|
||
try:
|
||
await self.plugin['xep_0045'].join_muc(room, self.boundjid.user)
|
||
logger.info(f"Вошли в комнату {room}")
|
||
except Exception as e:
|
||
logger.error(f"Ошибка входа в комнату {room}: {e}")
|
||
except Exception as e:
|
||
logger.error(f"Ошибка загрузки комнат с сервера: {e}")
|
||
|
||
# Запускаем фоновые задачи:
|
||
# - Переподключение к комнатам
|
||
# - Очистка временных файлов
|
||
# - Очистка словарей _last_file_path и last_file_texts
|
||
asyncio.create_task(self.room_handler.rejoin_loop())
|
||
asyncio.create_task(self.health_handler.cleanup_loop())
|
||
asyncio.create_task(self._cleanup_loop())
|
||
|
||
async def _cleanup_loop(self):
|
||
"""Периодическая очистка словарей _last_file_path и last_file_texts (раз в час)."""
|
||
while True:
|
||
await asyncio.sleep(3600)
|
||
if len(self._last_file_path) > 100:
|
||
self._last_file_path = dict(list(self._last_file_path.items())[-100:])
|
||
if len(self.last_file_texts) > 100:
|
||
self.last_file_texts = dict(list(self.last_file_texts.items())[-100:])
|
||
|
||
async def disconnect(self):
|
||
"""Переопределение для закрытия HTTP-сессии."""
|
||
await self.rag_client.close()
|
||
await super().disconnect() |