From a92ea2c90a87e2fa1e8e28fd2d47ee3903d4ff2e Mon Sep 17 00:00:00 2001 From: Markov Andrey Date: Tue, 30 Jun 2026 12:35:14 +0000 Subject: [PATCH] Add new file --- bots/handlers/file_handler.py | 173 ++++++++++++++++++++++++++++++++++ 1 file changed, 173 insertions(+) create mode 100644 bots/handlers/file_handler.py diff --git a/bots/handlers/file_handler.py b/bots/handlers/file_handler.py new file mode 100644 index 0000000..a016e51 --- /dev/null +++ b/bots/handlers/file_handler.py @@ -0,0 +1,173 @@ +# -*- coding: utf-8 -*- +""" +Обработчик файлов: скачивание файлов по URL (из OOB), сохранение во временную директорию, +постановка задачи в очередь индексации. + +Ключевое изменение для групповой работы: +- Если файл пришёл в групповой чат (room_jid не None), он автоматически индексируется + в базу знаний комнаты, без необходимости команды !learn. +- Для личных чатов – поведение зависит от режима обучения (learning_mode) и флага --global. +""" + +import os +import asyncio +import logging +import urllib.parse +import requests +from datetime import datetime +import re + +logger = logging.getLogger(__name__) + + +class FileHandler: + """ + Отвечает за скачивание файлов, переданных через ссылку в сообщении (XEP-0066 OOB). + """ + + def __init__(self, bot): + self.bot = bot + + async def _wait_for_queue(self): + """ + Ожидает создания очереди индексации (защита от гонки при старте бота). + + Проблема: очередь `self.bot.indexing_queue` создаётся в `session_start` после + подключения к БД. Если сообщение с файлом придёт ДО этого момента, + `indexing_queue` будет `None`, и `put()` вызовет исключение. + Решение: ждём, пока очередь станет доступна, проверяя каждые 0.5 секунды. + """ + while self.bot.indexing_queue is None: + logger.debug("Очередь индексации ещё не создана, ожидание...") + await asyncio.sleep(0.5) + + async def handle_file(self, msg, url: str, user_jid: str, room_jid: str = None): + """ + Скачивает файл и создаёт задачу для indexing_worker. + + Аргументы: + msg: stanza сообщения (нужна для получения user, ответных сообщений). + url (str): HTTP/HTTPS ссылка на файл (из атрибута url элемента oob). + user_jid (str): JID отправителя (без ресурса). + room_jid (str | None): JID комнаты, если сообщение из groupchat; иначе None. + + Логика: + 1. Извлечь имя файла из URL, декодировать, сгенерировать уникальное имя во временной папке. + 2. Скачать файл (синхронно, в отдельном потоке, с таймаутом из конфига). + 3. Определить режим индексации: + - Если room_jid не None -> mode = 'room', is_global = False (всегда). + - Если личный чат -> проверить learning_mode пользователя. + 4. Извлечь --title из тела сообщения (если есть). + 5. Положить задачу в очередь индексации (asyncio.Queue). + 6. Ответить пользователю, что файл принят в обработку. + """ + self.bot.metrics_handler.increment("files_processed") + try: + # --- 1. Извлечение имени файла из URL --- + raw_name = url.split('/')[-1].split('?')[0] + original_fname = urllib.parse.unquote(raw_name) + ext = os.path.splitext(original_fname)[1].lower() + fname = f"{msg['from'].user}_{int(datetime.now().timestamp())}{ext}" + target_path = os.path.join(self.bot.config.temp_dir, fname) + + # --- 2. Скачивание файла --- + # HEAD-запрос для получения размера + try: + head_resp = requests.head(url, timeout=self.bot.config.http_download_timeout) + content_length = head_resp.headers.get('Content-Length') + if content_length is not None and int(content_length) > self.bot.config.max_file_size_mb * 1024 * 1024: + reply = msg.reply(f"❌ Файл превышает максимальный размер {self.bot.config.max_file_size_mb} МБ.") + if reply: + reply.send() + return + except Exception: + pass + + # Скачивание с прерыванием по лимиту + def download_with_limit(): + r = requests.get(url, stream=True, timeout=self.bot.config.http_download_timeout) + if r.status_code != 200: + return False + downloaded = 0 + limit = self.bot.config.max_file_size_mb * 1024 * 1024 + with open(target_path, 'wb') as f: + for chunk in r.iter_content(8192): + f.write(chunk) + downloaded += len(chunk) + if downloaded > limit: + os.unlink(target_path) + raise Exception("File too large") + return True + + try: + success = await asyncio.to_thread(download_with_limit) + if not success: + reply = msg.reply("❌ Не удалось загрузить файл. Проверьте доступность ссылки.") + if reply: + reply.send() + return + except Exception as e: + if "File too large" in str(e): + reply = msg.reply(f"❌ Файл превышает максимальный размер {self.bot.config.max_file_size_mb} МБ.") + else: + reply = msg.reply("❌ Ошибка при скачивании файла.") + if reply: + reply.send() + return + + # --- 3. Определение режима индексации --- + if room_jid: + mode = 'room' + is_global = False + title = None + else: + mode_info = self.bot.learning_mode.get(user_jid) + if mode_info is None: + mode = None + is_global = False + title = None + else: + if isinstance(mode_info, tuple): + mode, saved_room = mode_info + is_global = (mode == 'global') + else: + mode = mode_info + is_global = (mode == 'global') + title = None + + # --- 4. Парсинг --title из тела сообщения --- + raw_body = msg['body'].strip() if msg['body'] else "" + if raw_body: + match = re.search(r'--title "([^"]+)"', raw_body) + if match: + title = match.group(1) + + # --- 5. Создание задачи для очереди --- + await self._wait_for_queue() + task = { + "user_jid": user_jid, + "room_jid": room_jid, + "file_path": target_path, + "original_fname": original_fname, + "title": title, + "is_global": is_global, + "mode": mode, + "ext": ext + } + await self.bot.indexing_queue.put(task) + + # --- 6. Уведомление пользователя --- + if room_jid: + reply = msg.reply(f"📥 Файл '{original_fname}' принят в обработку для КОМНАТНОЙ БЗ.") + else: + scope = "глобальную" if is_global else "личную" + reply = msg.reply(f"📥 Файл '{original_fname}' принят в обработку для {scope} БЗ.") + if reply: + reply.send() + + except Exception as e: + logger.exception(f"Ошибка при скачивании файла: {e}") + self.bot.metrics_handler.increment("errors_total") + reply = msg.reply("❌ Ошибка при обработке файла. Проверьте логи.") + if reply: + reply.send() \ No newline at end of file