Редактировать file_handler.py

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}")