Add new file

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()