Add new file

core/xmpp/client.py

@@ -0,0 +1,180 @@
+# -*- 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())