""" Фоновый воркер для повторной индексации документов с флагом 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}") # Здесь можно увеличить счётчик ошибок, если он есть # В противном случае просто продолжим # После обработки пачки принудительно закроем соединение? Нет, пул сам управляет.