# -*- 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())