140 lines
6.0 KiB
Python
140 lines
6.0 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 и повторяет индексацию.
|
||
Для каждого документа проверяет, не превышено ли число попыток.
|
||
"""
|
||
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:
|
||
# Восстанавливаем текст документа. Предполагаем, что в metadata есть поле 'file_text'
|
||
# или мы добавили колонку file_text в таблицу. Для демонстрации используем metadata.
|
||
metadata = row.get('metadata') or {}
|
||
file_text = metadata.get('file_text')
|
||
if not file_text:
|
||
logger.error(f"Документ {doc_id} не содержит текста для переиндексации, пропускаем")
|
||
# Можно обновить метаданные, чтобы отметить как невосстановимый
|
||
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,
|
||
)
|
||
if chunk_count > 0:
|
||
logger.info(f"Документ {doc_id} успешно переиндексирован (новый doc_id={doc_id_new})")
|
||
else:
|
||
logger.error(f"Не удалось переиндексировать документ {doc_id}, chunk_count=0")
|
||
|
||
except Exception as e:
|
||
logger.exception(f"Ошибка переиндексации документа {doc_id}: {e}") |