# -*- 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 (старый код) – позже их нужно будет перенести в rag, # но пока оставляем как есть для совместимости. 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): # <-- изменён тип конфига # Полный 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 = {} # В AppConfig теперь есть свойства для путей к файлам промптов, # но они называются так же, как в BotConfig, поэтому доступ через config.<имя> работает. # Однако в BotConfig они назывались system_prompt_file и т.д. – они есть в AppConfig? # В AppConfig мы добавили прокси-свойства только для основных параметров, # но не для путей к промптам. Поэтому здесь нужно либо использовать config.prompts_dir_path # и формировать пути вручную, либо оставить как было, если в AppConfig есть соответствующие свойства. # Для обратной совместимости мы добавили свойства в AppConfig, но не для prompt_file. # Поэтому временно оставим как было: будем использовать config.prompts_dir_path и имена файлов. # Лучше прочитать имена файлов из config.ai.prompts и загрузить их. # Но для простоты пока оставим старый способ, если свойства есть. # Однако в AppConfig нет system_prompt_file и т.д. – они были в BotConfig. # В BotConfig они были путями к файлам. В AppConfig мы не стали их добавлять, # чтобы не дублировать. Вместо этого мы загружаем все промпты в config.prompts_content. # Поэтому мы можем использовать config.prompts_content, который уже содержит все промпты. # Используем его: self.ai_prompts = config.prompts_content.copy() if config.prompts_content else {} if not self.ai_prompts: logger.warning("Промпты не загружены из конфига, пробуем загрузить из файлов вручную...") # Запасной вариант – загружаем из файлов, используя config.prompts_dir_path 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-сервера ----- # Адрес сервера и API-ключ берём из конфига rag_server_url = config.rag_server_url # теперь это свойство rag_api_key = config.rag_api_key # получаем API-ключ # Создаём RAGClient с передачей ключа 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. Команды и состояния ----- # Регистрируем все команды бота (из 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())