diff --git a/bots/xmpp/client.py b/bots/xmpp/client.py index 0cfc61c..dc52f05 100644 --- a/bots/xmpp/client.py +++ b/bots/xmpp/client.py @@ -7,6 +7,7 @@ XMPP клиент – главный класс Bot. ИСПОЛЬЗУЕТ: - AppConfig из rag.config_models для единой конфигурации. - RAGClient с передачей API-ключа для аутентификации. +- Комнаты загружаются с сервера при старте. """ import asyncio @@ -17,8 +18,7 @@ from slixmpp import ClientXMPP # Импорты из нового пакета rag для конфигурации from rag.config_models import AppConfig -# Остальные импорты остаются из core (старый код) – позже их нужно будет перенести в rag, -# но пока оставляем как есть для совместимости. +# Остальные импорты из core from core.services.file_service import FileService from core.handlers.metrics_handler import MetricsHandler from core.rag_client import RAGClient @@ -40,40 +40,19 @@ class Bot(ClientXMPP): Тонкий клиент: вся RAG-логика на сервере. """ - def __init__(self, config: AppConfig): # <-- изменён тип конфига - # Полный JID: user@domain/resource + def __init__(self, config: AppConfig): 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(): @@ -91,35 +70,29 @@ class Bot(ClientXMPP): 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 с передачей ключа + 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, # можно использовать таймаут из конфига + 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 = [] # комнаты, в которых бот состоит (загружается позже) + self.learning_mode = {} + self.last_file_texts = {} + self._last_file_path = {} + self.indexing_queue = None + self.rooms = [] # больше не используется, загружаем с сервера # ----- 5. Хендлеры ----- self.message_handler = MessageHandler(self) @@ -130,13 +103,13 @@ class Bot(ClientXMPP): 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 + 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() @@ -152,24 +125,18 @@ class Bot(ClientXMPP): 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}") @@ -182,31 +149,42 @@ class Bot(ClientXMPP): # Отправляем запросы на подписку администраторам 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}") + # Загружаем комнаты с сервера и входим в них + 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()) \ No newline at end of file + 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() \ No newline at end of file