8a2d5ba2ff
- /rag/rag_server.py - /bots/rag_client.py - /bots/workers/indexing_worker.py - /bots/handlers/file_handler.py
328 lines
16 KiB
Python
328 lines
16 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
Фоновый воркер (worker) для индексации файлов и URL через RAG-сервер.
|
||
Читает задачи из asyncio.Queue и выполняет:
|
||
- Для URL: скачивание, извлечение текста, индексация через HTTP-запрос к серверу.
|
||
- Для файлов: извлечение текста, обработка изображений/аудио, индексация через HTTP-запрос.
|
||
|
||
Все операции с БД и Qdrant выполняются на сервере. Воркер только отправляет данные.
|
||
|
||
ДОБАВЛЕНО:
|
||
- Обработка изображений через RAG-сервер (OCR)
|
||
- Обработка аудио через RAG-сервер (транскрибация)
|
||
"""
|
||
|
||
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/распознавание) – через RAG-сервер.
|
||
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. Специальная обработка: изображения (OCR через RAG-сервер) ---
|
||
if ext in ('.jpg', '.jpeg', '.png', '.bmp', '.tiff') and self.bot.config.vision:
|
||
result = await self.bot.rag_client.vision(file_path)
|
||
if result.get('text'):
|
||
await self._send_message(
|
||
user_jid,
|
||
f"🖼️ **Текст на изображении:**\n{result['text']}",
|
||
room_jid
|
||
)
|
||
else:
|
||
error = result.get('error', 'неизвестная ошибка')
|
||
await self._send_message(
|
||
user_jid,
|
||
f"⚠️ Не удалось распознать текст: {error}",
|
||
room_jid
|
||
)
|
||
return
|
||
|
||
# --- 4. Специальная обработка: аудио (транскрибация через RAG-сервер) ---
|
||
if ext in ('.ogg', '.wav', '.mp3', '.amr', '.m4a') and self.bot.config.voice_recognition:
|
||
result = await self.bot.rag_client.transcribe(file_path)
|
||
if result.get('text'):
|
||
await self._send_message(
|
||
user_jid,
|
||
f"🎤 **Распознано:**\n{result['text']}",
|
||
room_jid
|
||
)
|
||
else:
|
||
error = result.get('error', 'неизвестная ошибка')
|
||
await self._send_message(
|
||
user_jid,
|
||
f"⚠️ Не удалось распознать речь: {error}",
|
||
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') |