Update 86 files
- /rag/commands/expert.py - /core/services/__init__.py - /core/services/embedding_service.py - /core/services/file_service.py - /core/services/giga_client.py - /core/services/kb_service.py - /core/services/postgres_service.py - /core/services/qdrant_service.py - /core/services/reranker_service.py - /core/functions/__init__.py - /core/functions/check_consistency.py - /core/functions/check_spelling.py - /core/functions/critique_answer.py - /core/functions/expand_query.py - /core/functions/extract_metrics.py - /core/functions/file_processor.py - /core/functions/generate_document.py - /core/functions/intent_classify.py - /core/functions/rerank_context.py - /core/functions/summarize_document.py - /core/utils/__init__.py - /core/utils/arg_parser.py - /core/utils/config_loader.py - /core/utils/layout_converter.py - /core/utils/logger.py - /core/utils/text_utils.py - /core/utils/web_utils.py - /bots/commands/__init__.py - /bots/commands/base.py - /bots/commands/create.py - /core/commands/global_remove.py - /core/commands/help.py - /core/commands/info.py - /core/commands/kb.py - /core/commands/learn.py - /core/commands/other.py - /core/commands/registry.py - /core/commands/stats.py - /core/commands/template.py - /core/handlers/__init__.py - /core/handlers/message_handler.py - /core/workers/__init__.py - /core/workers/indexing_worker.py - /core/xmpp/__init__.py - /core/xmpp/client.py - /bots/commands/global_remove.py - /bots/commands/help.py - /bots/commands/info.py - /bots/commands/kb.py - /bots/commands/registry.py - /bots/commands/other.py - /bots/commands/template.py - /bots/commands/learn.py - /bots/commands/stats.py - /bots/handlers/message_handler.py - /bots/handlers/__init__.py - /bots/workers/__init__.py - /bots/workers/indexing_worker.py - /bots/xmpp/__init__.py - /bots/xmpp/client.py - /rag/services/qdrant_service.py - /rag/services/giga_client.py - /rag/services/embedding_service.py - /rag/services/reranker_service.py - /rag/services/__init__.py - /rag/services/file_service.py - /rag/services/postgres_service.py - /rag/services/kb_service.py - /rag/functions/check_spelling.py - /rag/functions/__init__.py - /rag/functions/extract_metrics.py - /rag/functions/generate_document.py - /rag/functions/expand_query.py - /rag/functions/critique_answer.py - /rag/functions/file_processor.py - /rag/functions/check_consistency.py - /rag/functions/summarize_document.py - /rag/functions/rerank_context.py - /rag/functions/intent_classify.py - /rag/utils/__init__.py - /rag/utils/layout_converter.py - /rag/utils/text_utils.py - /rag/utils/arg_parser.py - /rag/utils/config_loader.py - /rag/utils/logger.py - /rag/utils/web_utils.py
This commit is contained in:
Markov Andrey
6
bots/workers/__init__.py
Normal file
6
bots/workers/__init__.py
Normal file
@@ -0,0 +1,6 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Пакет workers – фоновые асинхронные воркеры.
|
||||
Содержит:
|
||||
- IndexingWorker – индексация файлов и URL через HTTP-запросы к RAG-серверу.
|
||||
"""
|
||||
324
bots/workers/indexing_worker.py
Normal file
324
bots/workers/indexing_worker.py
Normal file
@@ -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')
|
||||
Reference in New Issue
Block a user