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