From f74cb640353785f380e6667c136d994ee1a7bfec Mon Sep 17 00:00:00 2001 From: Markov Andrey Date: Tue, 30 Jun 2026 09:29:03 +0000 Subject: [PATCH] Add new file --- core/xmpp/client.py | 180 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 180 insertions(+) create mode 100644 core/xmpp/client.py diff --git a/core/xmpp/client.py b/core/xmpp/client.py new file mode 100644 index 0000000..46714ad --- /dev/null +++ b/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()) \ No newline at end of file