164 lines
9.2 KiB
Python
164 lines
9.2 KiB
Python
"""
|
||
Фоновый воркер для повторной индексации документов с флагом indexed=False.
|
||
Периодически проверяет базу данных и повторяет индексацию для документов,
|
||
которые не были успешно проиндексированы (например, из-за сбоя Qdrant).
|
||
"""
|
||
|
||
import asyncio
|
||
import logging
|
||
from typing import Optional
|
||
|
||
from rag.services.postgres_service import PostgresService
|
||
from rag.services.kb_service import KBService
|
||
from rag.config_models import AppConfig
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class ReindexWorker:
|
||
"""
|
||
Фоновый воркер для переиндексации документов с флагом indexed=False.
|
||
Запускается как асинхронная задача в rag_server.py.
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
db: PostgresService,
|
||
kb: KBService,
|
||
config: AppConfig,
|
||
interval_seconds: int = 300,
|
||
batch_size: int = 10,
|
||
max_retries: int = 3,
|
||
):
|
||
"""
|
||
Инициализация воркера.
|
||
|
||
Аргументы:
|
||
db: сервис PostgreSQL.
|
||
kb: сервис базы знаний (KBService) для повторной индексации.
|
||
config: глобальная конфигурация.
|
||
interval_seconds: интервал между проверками (сек).
|
||
batch_size: максимальное количество документов за один цикл.
|
||
max_retries: максимальное количество попыток переиндексации (пока не увеличим счётчик).
|
||
"""
|
||
self.db = db
|
||
self.kb = kb
|
||
self.config = config
|
||
self.interval_seconds = interval_seconds
|
||
self.batch_size = batch_size
|
||
self.max_retries = max_retries
|
||
self._running = False
|
||
self._task: Optional[asyncio.Task] = None
|
||
|
||
async def start(self) -> None:
|
||
"""Запускает воркер в фоновом режиме."""
|
||
if self._running:
|
||
logger.warning("ReindexWorker уже запущен")
|
||
return
|
||
self._running = True
|
||
self._task = asyncio.create_task(self._run())
|
||
logger.info("ReindexWorker запущен")
|
||
|
||
async def stop(self) -> None:
|
||
"""Останавливает воркер."""
|
||
if not self._running:
|
||
return
|
||
self._running = False
|
||
if self._task and not self._task.done():
|
||
self._task.cancel()
|
||
try:
|
||
await self._task
|
||
except asyncio.CancelledError:
|
||
pass
|
||
logger.info("ReindexWorker остановлен")
|
||
|
||
async def _run(self) -> None:
|
||
"""Основной цикл воркера."""
|
||
while self._running:
|
||
try:
|
||
await self._process_batch()
|
||
except asyncio.CancelledError:
|
||
break
|
||
except Exception as e:
|
||
logger.exception(f"Ошибка в цикле ReindexWorker: {e}")
|
||
await asyncio.sleep(self.interval_seconds)
|
||
|
||
async def _process_batch(self) -> None:
|
||
"""
|
||
Выбирает до batch_size документов с indexed=False и повторяет индексацию.
|
||
Для каждого документа проверяет, не превышено ли число попыток.
|
||
"""
|
||
# Получаем документы, у которых indexed=False и retry_count < max_retries
|
||
# (добавим колонку retry_count, если её нет, или используем существующую metadata)
|
||
# В текущей схеме нет retry_count, поэтому просто будем пытаться индексировать заново.
|
||
# Чтобы избежать бесконечных попыток, можно ограничить по времени последней ошибки.
|
||
# Для простоты будем считать, что документ помечен как неудачный, и мы пытаемся его
|
||
# переиндексировать, пока не удастся, но с ограничением по времени (например, не старше 7 дней).
|
||
# Реализуем простой вариант: берём документы с indexed=False, которые были созданы не более 7 дней назад.
|
||
import datetime
|
||
cutoff = datetime.datetime.now() - datetime.timedelta(days=7)
|
||
async with self.db.pool.acquire() as conn:
|
||
rows = await conn.fetch(
|
||
"""
|
||
SELECT id, source_name, owner_jid, room_jid, is_global, metadata, file_hash
|
||
FROM documents
|
||
WHERE indexed = FALSE
|
||
AND upload_time > $1
|
||
LIMIT $2
|
||
""",
|
||
cutoff,
|
||
self.batch_size,
|
||
)
|
||
if not rows:
|
||
return
|
||
|
||
logger.info(f"Найдено {len(rows)} документов для переиндексации")
|
||
|
||
for row in rows:
|
||
doc_id = row['id']
|
||
try:
|
||
# Восстанавливаем данные для индексации. Текст документа должен храниться
|
||
# либо в метаданных (как file_text), либо в отдельном поле, либо мы можем
|
||
# попытаться извлечь из Qdrant? Лучше хранить исходный текст в БД.
|
||
# Предположим, что у нас есть колонка file_text в documents (или мы её добавим).
|
||
# Поскольку в текущей схеме её нет, мы не можем переиндексировать без текста.
|
||
# Здесь нужно решение: либо добавить колонку, либо хранить файлы на диске.
|
||
# Для демонстрации реализуем заглушку: если в metadata есть 'file_text', используем его.
|
||
metadata = row.get('metadata') or {}
|
||
file_text = metadata.get('file_text')
|
||
if not file_text:
|
||
# Если текст не сохранён, пропускаем (логируем ошибку)
|
||
logger.error(f"Документ {doc_id} не содержит текста для переиндексации, пропускаем")
|
||
# Можно обновить indexed=False и добавить метку "невозможно переиндексировать"
|
||
continue
|
||
|
||
# Повторно индексируем документ
|
||
doc_id_new, chunk_count = await self.kb.add_document(
|
||
file_name=row['source_name'],
|
||
file_text=file_text,
|
||
user_jid=row['owner_jid'],
|
||
is_global=row['is_global'],
|
||
title=row['source_name'],
|
||
metadata=metadata,
|
||
room_jid=row['room_jid'],
|
||
file_hash=row['file_hash'],
|
||
update_if_exists=True, # перезаписать существующий документ
|
||
)
|
||
# Если индексация успешна, новая запись будет иметь indexed=True,
|
||
# а старая будет удалена (в add_document). Поэтому нам нужно убедиться,
|
||
# что старая запись удалена. Однако add_document создаёт новую запись,
|
||
# а затем удаляет старую, поэтому здесь мы можем просто отметить старую как удалённую.
|
||
# Но проще: если add_document вернул новый doc_id и chunk_count >0, считаем успех.
|
||
if chunk_count > 0:
|
||
logger.info(f"Документ {doc_id} успешно переиндексирован (новый doc_id={doc_id_new})")
|
||
# Удаляем старую запись (она уже должна быть удалена в add_document, но на всякий случай)
|
||
# await self.db.delete_documents_by_ids([doc_id])
|
||
else:
|
||
logger.error(f"Не удалось переиндексировать документ {doc_id}, chunk_count=0")
|
||
|
||
except Exception as e:
|
||
logger.exception(f"Ошибка переиндексации документа {doc_id}: {e}")
|
||
# Здесь можно увеличить счётчик ошибок, если он есть
|
||
# В противном случае просто продолжим
|
||
|
||
# После обработки пачки принудительно закроем соединение? Нет, пул сам управляет. |