From a3c5c0840421d8c1df6fce3032c563add2c5ce91 Mon Sep 17 00:00:00 2001 From: Markov Andrey Date: Tue, 30 Jun 2026 21:48:28 +0000 Subject: [PATCH] =?UTF-8?q?=D0=A0=D0=B5=D0=B4=D0=B0=D0=BA=D1=82=D0=B8?= =?UTF-8?q?=D1=80=D0=BE=D0=B2=D0=B0=D1=82=D1=8C=20file=5Fhandler.py?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- bots/handlers/file_handler.py | 195 ++++++++++++++-------------------- 1 file changed, 82 insertions(+), 113 deletions(-) diff --git a/bots/handlers/file_handler.py b/bots/handlers/file_handler.py index 9452089..f1735e1 100644 --- a/bots/handlers/file_handler.py +++ b/bots/handlers/file_handler.py @@ -9,8 +9,8 @@ - Для личных чатов – поведение зависит от режима обучения (learning_mode) и флага --global. ДОБАВЛЕНО: -- Удалён прямой вызов GigaChat для изображений (теперь это делает RAG-сервер) -- Удалён прямой вызов транскрибации (теперь это делает RAG-сервер) +- Гарантированное удаление временных файлов через tempfile.TemporaryDirectory. +- Обработка ошибок скачивания с очисткой. """ import os @@ -18,6 +18,8 @@ import asyncio import logging import urllib.parse import requests +import tempfile +import shutil from datetime import datetime import re @@ -35,11 +37,6 @@ class FileHandler: 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("Очередь индексации ещё не создана, ожидание...") @@ -48,126 +45,98 @@ class FileHandler: 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} МБ.") + # Создаём временную директорию (автоматически удалится) + with tempfile.TemporaryDirectory() as tmpdir: + fname = f"{msg['from'].user}_{int(datetime.now().timestamp())}{ext}" + target_path = os.path.join(tmpdir, fname) + + # Проверка размера через 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: + raise Exception("HTTP error") + limit = self.bot.config.max_file_size_mb * 1024 * 1024 + with open(target_path, 'wb') as f: + downloaded = 0 + for chunk in r.iter_content(8192): + f.write(chunk) + downloaded += len(chunk) + if downloaded > limit: + raise Exception("File too large") + + try: + await asyncio.to_thread(download_with_limit) + 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 - 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} МБ.") + # Определение режима индексации + if room_jid: + mode = 'room' + is_global = False else: - reply = msg.reply("❌ Ошибка при скачивании файла.") + mode_info = self.bot.learning_mode.get(user_jid) + if mode_info is None: + mode = None + is_global = False + 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 + 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) + + # Создаём задачу для фонового индексатора + 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) + + scope = "комнатную" if room_jid else ("глобальную" if is_global else "личную") + reply = msg.reply(f"📥 Файл '{original_fname}' принят в обработку для {scope} БЗ.") 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}")