diff --git a/bots/workers/reindex_worker.py b/bots/workers/reindex_worker.py new file mode 100644 index 0000000..9f07428 --- /dev/null +++ b/bots/workers/reindex_worker.py @@ -0,0 +1,164 @@ +""" +Фоновый воркер для повторной индексации документов с флагом 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}") + # Здесь можно увеличить счётчик ошибок, если он есть + # В противном случае просто продолжим + + # После обработки пачки принудительно закроем соединение? Нет, пул сам управляет. \ No newline at end of file