diff --git a/core/workers/indexing_worker.py b/core/workers/indexing_worker.py new file mode 100644 index 0000000..fef1dea --- /dev/null +++ b/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') \ No newline at end of file