Редактировать reindex_worker.py

2026-06-30 21:58:40 +00:00
parent 45e41937f7
commit 517618ba00
1 changed files with 0 additions and 0 deletions
--- a/rag/workers/reindex_worker.py
+++ b/rag/workers/reindex_worker.py
@@ -0,0 +1,140 @@
+"""
+Фоновый воркер для повторной индексации документов с флагом 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}")