# -*- 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')