Add new file

core/workers/indexing_worker.py

@@ -0,0 +1,324 @@
+# -*- coding: utf-8 -*-
+"""
+Фоновый воркер (worker) для индексации файлов и URL через RAG-сервер.
+Читает задачи из asyncio.Queue и выполняет:
+- Для URL: скачивание, извлечение текста, индексация через HTTP-запрос к серверу.
+- Для файлов: извлечение текста, обработка изображений/аудио, индексация через HTTP-запрос.
+
+Все операции с БД и Qdrant выполняются на сервере. Воркер только отправляет данные.
+"""
+
+import asyncio
+import logging
+import os
+import hashlib
+import shutil
+from datetime import datetime
+
+logger = logging.getLogger(__name__)
+
+
+class IndexingWorker:
+    """
+    Фоновый воркер для индексации.
+    Получает задачи из очереди и выполняет их, отправляя данные на RAG-сервер.
+    """
+
+    def __init__(self, bot):
+        """
+        Аргументы:
+            bot: экземпляр Bot (содержит rag_client, config, file_handler и состояния)
+        """
+        self.bot = bot
+        logger.info("IndexingWorker инициализирован (режим: HTTP-клиент RAG-сервера)")
+
+    async def run(self):
+        """Бесконечный цикл обработки задач из очереди."""
+        while True:
+            task = await self.bot.indexing_queue.get()
+            try:
+                if task.get("type") == "url":
+                    await self._process_url(task)
+                else:
+                    await self._process_file(task)
+            except Exception as e:
+                logger.exception(f"Ошибка в воркере при обработке задачи: {e}")
+                self.bot.metrics_handler.increment("errors_total")
+            finally:
+                self.bot.indexing_queue.task_done()
+
+    # ==========================================================
+    # Обработка URL
+    # ==========================================================
+
+    async def _process_url(self, task):
+        """
+        Обрабатывает URL-задачу.
+        - Одиночный URL: загружает страницу/PDF, извлекает текст, индексирует через сервер.
+        - Рекурсивный обход: через crawl_url, затем индексирует каждую страницу.
+        """
+        from core.utils.web_utils import fetch_any_url, crawl_url
+
+        url = task["url"]
+        user_jid = task["user_jid"]
+        room_jid = task.get("room_jid")
+        custom_title = task.get("title")
+        is_global = task.get("is_global", False)
+        recursive = task.get("recursive", False)
+        depth = task.get("depth", 1)
+        max_pages = task.get("max_pages", self.bot.config.ws_max_pages)
+
+        try:
+            if recursive:
+                pages = await crawl_url(url, max_depth=depth, max_pages=max_pages, domain_restrict=True)
+                if not pages:
+                    await self._send_message(
+                        user_jid,
+                        f"❌ Не удалось загрузить ни одной страницы с {url}",
+                        room_jid
+                    )
+                    return
+                total = 0
+                for page_url, page_title, page_text in pages:
+                    if page_text:
+                        doc_title = custom_title if custom_title else (page_title if page_title else page_url)
+                        # Отправляем на сервер для индексации
+                        result = await self.bot.rag_client.index_document(
+                            file_name=page_url,
+                            file_text=page_text,
+                            user_jid=user_jid,
+                            room_jid=room_jid,
+                            is_global=is_global,
+                            title=doc_title,
+                            metadata={"url": page_url}
+                        )
+                        if result.get('doc_id'):
+                            total += 1
+                        else:
+                            logger.error(f"Ошибка индексации {page_url}: {result.get('error')}")
+                await self._send_message(
+                    user_jid,
+                    f"✅ Проиндексировано {total} страниц (рекурсивно).",
+                    room_jid
+                )
+            else:
+                title, text = await fetch_any_url(url)
+                if text:
+                    doc_title = custom_title if custom_title else (title if title else url)
+                    result = await self.bot.rag_client.index_document(
+                        file_name=url,
+                        file_text=text,
+                        user_jid=user_jid,
+                        room_jid=room_jid,
+                        is_global=is_global,
+                        title=doc_title,
+                        metadata={"url": url}
+                    )
+                    if result.get('doc_id'):
+                        scope = "глобальную" if is_global else ("комнату" if room_jid else "личную")
+                        await self._send_message(
+                            user_jid,
+                            f"✅ Страница '{doc_title}' добавлена в {scope} БЗ.",
+                            room_jid
+                        )
+                    else:
+                        await self._send_message(
+                            user_jid,
+                            f"❌ Ошибка индексации {url}: {result.get('error')}",
+                            room_jid
+                        )
+                else:
+                    await self._send_message(
+                        user_jid,
+                        f"❌ Не удалось загрузить или извлечь текст из {url}",
+                        room_jid
+                    )
+        except Exception as e:
+            logger.exception(f"Ошибка индексации URL {url}: {e}")
+            await self._send_message(
+                user_jid,
+                f"❌ Ошибка при обработке URL: {e}",
+                room_jid
+            )
+
+    # ==========================================================
+    # Обработка файлов
+    # ==========================================================
+
+    async def _process_file(self, task):
+        """
+        Обрабатывает задачу с файлом:
+        1. Извлекает текст (синхронно в потоке).
+        2. Обрабатывает изображения/аудио (Vision/распознавание) – если есть.
+        3. Сохраняет постоянную копию в data/ (локально, для истории и отладки).
+        4. Отправляет текст и метаданные на RAG-сервер для индексации.
+        5. Удаляет временный файл.
+        """
+        user_jid = task["user_jid"]
+        room_jid = task.get("room_jid")
+        file_path = task["file_path"]
+        original_fname = task["original_fname"]
+        title = task.get("title")
+        is_global = task.get("is_global", False)
+        mode = task.get("mode")
+        ext = task.get("ext", "").lower()
+
+        try:
+            # --- 1. Извлечение текста (синхронная операция в потоке) ---
+            result = await asyncio.to_thread(self.bot.files.process_any_file, file_path)
+            if isinstance(result, tuple) and len(result) == 2:
+                first, second = result
+                if isinstance(first, list):
+                    files_list = first
+                    total_count = second
+                else:
+                    files_list = [(original_fname, first)] if first else []
+                    total_count = second
+            else:
+                files_list = [(original_fname, result)] if result else []
+                total_count = 1 if result else 0
+
+            if total_count == 0:
+                await self._send_message(
+                    user_jid,
+                    f"⚠️ Не удалось извлечь текст из {original_fname}.",
+                    room_jid
+                )
+                return
+
+            # Сохраняем текст для локального использования (например, !summary)
+            full_text = "\n\n".join([t for _, t in files_list if t])
+            self.bot.last_file_texts[user_jid] = full_text
+
+            # --- 2. Сохранение постоянной копии в data/ (локально) ---
+            # Это нужно для команд !summary и хирургической замены, которые работают с файлом локально.
+            # В будущем можно перенести и их на сервер, но пока оставляем.
+            with open(file_path, 'rb') as f:
+                file_hash = hashlib.sha256(f.read()).hexdigest()
+            short_hash = file_hash[:12]
+            safe_user = user_jid.replace('@', '_at_').replace('/', '_')
+            safe_fname = original_fname.replace('/', '_').replace('\\', '_')
+            dest_filename = f"user_{short_hash}_{safe_user}_{safe_fname}"
+            dest_path = self.bot.config.data_dir / dest_filename
+
+            if not dest_path.exists():
+                shutil.copy2(file_path, dest_path)
+                logger.debug(f"Создана постоянная копия: {dest_path}")
+            else:
+                logger.debug(f"Постоянная копия уже существует: {dest_path}")
+
+            key = (user_jid, room_jid) if room_jid is not None else user_jid
+            self.bot._last_file_path[key] = str(dest_path)
+
+            # --- 3. Специальная обработка: изображения (Vision API) ---
+            # Если это изображение и включена опция vision, отправляем на анализ.
+            # Это можно делать отдельно, не блокируя индексацию.
+            if ext in ('.jpg', '.jpeg', '.png', '.bmp', '.tiff') and self.bot.config.vision:
+                file_id = await self.bot.giga.upload_file(file_path)  # giga всё ещё локально? нужно будет перенести и это.
+                # Это временно, пока мы не вынесли всё на сервер. Пока оставим как есть.
+                if file_id:
+                    answer = await self.bot.giga.chat(
+                        history=[],
+                        query="Проанализируй это изображение. Если есть текст – прочитай его.",
+                        system_prompt=None,
+                        file_id=file_id,
+                        temperature=self.bot.config.ai_temperature
+                    )
+                    await self._send_message(user_jid, answer, room_jid)
+                else:
+                    await self._send_message(user_jid, "⚠️ Не удалось передать изображение в GigaChat.", room_jid)
+                # После обработки изображения мы не индексируем его как документ (обычно OCR делает это),
+                # но можно и проиндексировать текст. Здесь мы просто возвращаемся.
+                return
+
+            # --- 4. Специальная обработка: аудио (голос → текст) ---
+            # Аналогично, используем локальный SaluteSpeech, пока не вынесли.
+            if ext in ('.ogg', '.wav', '.mp3', '.amr', '.m4a') and self.bot.config.voice_recognition:
+                first_text = files_list[0][1] if files_list else ""
+                if "--- РАСПОЗНАННЫЙ ГОЛОС ---" in first_text:
+                    voice_text = first_text.split("--- РАСПОЗНАННЫЙ ГОЛОС ---")[1].strip()
+                    await self._send_message(user_jid, f"🎤 Распознано: {voice_text}", room_jid)
+                    # Можно также ответить на голос, но это отдельный сценарий.
+                    # Мы не индексируем аудио как документ (разве что текст распознавания).
+                    # Если хотим индексировать текст голоса, нужно передать его как документ.
+                    # Например, можно создать виртуальный документ с текстом.
+                else:
+                    await self._send_message(user_jid, f"⚠️ {first_text}", room_jid)
+                return
+
+            # --- 5. Индексация документов через RAG-сервер ---
+            # Отправляем на сервер каждый извлечённый файл (включая вложенные из архивов)
+            if mode in ('personal', 'global') or (mode == 'room' and room_jid):
+                saved = 0
+                for sub_name, sub_text in files_list:
+                    if sub_text:
+                        result = await self.bot.rag_client.index_document(
+                            file_name=sub_name,
+                            file_text=sub_text,
+                            user_jid=user_jid,
+                            room_jid=room_jid,
+                            is_global=is_global,
+                            title=title,
+                            # file_hash=file_hash,  # можно передать, если нужно
+                            update_if_exists=True
+                        )
+                        if result.get('doc_id'):
+                            saved += 1
+                        else:
+                            logger.error(f"Ошибка индексации {sub_name}: {result.get('error')}")
+                if saved:
+                    if room_jid:
+                        scope = f"комнату {room_jid}"
+                    elif is_global:
+                        scope = "глобальную"
+                    else:
+                        scope = "личную"
+                    await self._send_message(
+                        user_jid,
+                        f"🎓 {saved} файлов добавлено в {scope} БЗ.",
+                        room_jid
+                    )
+                else:
+                    await self._send_message(
+                        user_jid,
+                        f"⚠️ Не удалось проиндексировать {original_fname}. Проверьте логи.",
+                        room_jid
+                    )
+            else:
+                # Если режим обучения не включён (или не совпадает), просто показываем текст
+                preview = full_text[:500] + ("..." if len(full_text) > 500 else "")
+                await self._send_message(
+                    user_jid,
+                    f"📄 Файл `{original_fname}` прочитан.\n\n{preview}",
+                    room_jid
+                )
+        except Exception as e:
+            logger.exception(f"Ошибка при обработке файла {original_fname}: {e}")
+            self.bot.metrics_handler.increment("errors_total")
+            await self._send_message(
+                user_jid,
+                f"❌ Ошибка при обработке файла '{original_fname}'. Проверьте логи.",
+                room_jid
+            )
+        finally:
+            # Удаляем временный файл (постоянная копия уже сохранена)
+            try:
+                if os.path.exists(file_path):
+                    os.unlink(file_path)
+                    logger.debug(f"Временный файл удалён: {file_path}")
+            except Exception as e:
+                logger.error(f"Не удалось удалить временный файл {file_path}: {e}")
+
+    # ==========================================================
+    # Вспомогательные методы
+    # ==========================================================
+
+    async def _send_message(self, user_jid: str, body: str, room_jid: str = None):
+        """
+        Отправляет сообщение пользователю (в личный чат или комнату).
+        Используется для уведомлений о результатах индексации.
+        """
+        if room_jid:
+            self.bot.send_message(mto=room_jid, mbody=body, mtype='groupchat')
+        else:
+            self.bot.send_message(mto=user_jid, mbody=body, mtype='chat')