fckbot/bots/xmpp/client.py

# -*- coding: utf-8 -*-
"""
XMPP клиент – главный класс Bot.
Теперь это тонкий клиент, который только общается с RAG-сервером.
Не содержит БД, Qdrant, GigaChat – всё на сервере.
"""

import asyncio
import logging
from datetime import datetime
from slixmpp import ClientXMPP

from core.utils.config_loader import BotConfig
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: BotConfig):
        # Полный JID: user@domain/resource
        full_jid = f"{config.jid}/{config.xmpp_resource}"
        super().__init__(full_jid, config.xmpp_password)
        self.config = config

        # ----- 1. Локальные сервисы (только для файлов и метрик) -----
        # FileService нужен для извлечения текста из файлов перед отправкой на сервер
        self.files = FileService(config)
        # Метрики (счётчики сообщений, ошибок и т.д.)
        self.metrics_handler = MetricsHandler()

        # ----- 2. Загрузка промптов (остаются локально) -----
        # Промпты загружаются из файлов и передаются в каждом запросе к серверу.
        self.ai_prompts = {}
        prompt_files = {
            'system': config.system_prompt_file,
            'synthesis': config.synthesis_prompt_file,
            'quality': config.quality_prompt_file,
            'rerank': config.rerank_prompt_file,
            'expand': config.expand_prompt_file,
            'intent': config.intent_prompt_file,
            'metrics': config.metrics_prompt_file,
            'summary': config.summary_prompt_file,
            'consistency': config.consistency_prompt_file,
            'critique': config.critique_prompt_file,
            'spellcheck': config.spellcheck_prompt_file,
            'generate_document': config.generate_document_prompt_file,
        }
        for name, path in prompt_files.items():
            try:
                with open(path, 'r', encoding='utf-8') as f:
                    self.ai_prompts[name] = f.read()
                logger.debug(f"Загружен промпт {name} из {path}")
            except Exception as e:
                logger.error(f"Не удалось загрузить промпт {name} из {path}: {e}")
                self.ai_prompts[name] = ""
        # Сохраняем промпты как атрибут для передачи в запросах
        self.prompts = self.ai_prompts

        # ----- 3. HTTP-клиент для RAG-сервера -----
        # Адрес сервера берём из конфига (bot.conf) или используем значение по умолчанию
        rag_server_url = getattr(config, 'rag_server_url', 'http://localhost:8080')
        self.rag_client = RAGClient(server_url=rag_server_url)

        # ----- 4. Команды и состояния -----
        # Регистрируем все команды бота (из core/commands/)
        self.commands = register_commands(self)
        self._command_names = [cmd.name for cmd in self.commands.values() if cmd.name]

        # Временные состояния (хранятся в памяти)
        self.learning_mode = {}          # user_jid -> (mode, room_jid)
        self.last_file_texts = {}        # user_jid -> текст последнего файла
        self._last_file_path = {}        # key -> путь к последнему файлу
        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')   # Service Discovery
        self.register_plugin('xep_0199')   # Ping
        self.register_plugin('xep_0045')   # Multi-User Chat
        self.register_plugin('xep_0054')   # VCard
        self.register_plugin('xep_0172')   # User Nickname
        self.register_plugin('xep_0066')   # Out of Band Data
        self.register_plugin('xep_0363')   # HTTP File Upload

        # ----- 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:
        """Проверяет, зашифровано ли сообщение (OMEMO). Заглушка."""
        return False

    async def handle_ai_file_request(self, msg, ai_response: str, user_jid: str, room_jid: str = None, body: str = None):
        """
        Обрабатывает теги [SURGICAL_REPLACE], [REWRITE], [FILE] из ответа AI.
        """
        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:
            if admin_jid in self.roster:
                item = self.roster[admin_jid]
                if item and item.get('subscription') in ('both', 'from'):
                    continue
            self.send_presence(pto=admin_jid, ptype='subscribe')
            logger.info(f"Отправлен запрос подписки для администратора {admin_jid}")

        # Создаём очередь индексации и запускаем фоновый воркер
        self.indexing_queue = asyncio.Queue()
        asyncio.create_task(self.indexing_worker.run())

        # Загружаем комнаты из конфига или БД. Пока оставляем пустым.
        # Бот будет входить в комнаты только по приглашениям.
        self.rooms = []

        # Входим в комнаты (если есть)
        for room in self.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}")

        # Запускаем фоновые задачи:
        # - Переподключение к комнатам
        # - Очистка временных файлов
        asyncio.create_task(self.room_handler.rejoin_loop())
        asyncio.create_task(self.health_handler.cleanup_loop())