From bfacbad0ded321e22e753cdca215ef74b9f7a56a Mon Sep 17 00:00:00 2001 From: Markov Andrey Date: Tue, 30 Jun 2026 16:07:12 +0000 Subject: [PATCH] Update file kb_service.py --- rag/services/kb_service.py | 87 +++++++++++++++++++++----------------- 1 file changed, 48 insertions(+), 39 deletions(-) diff --git a/rag/services/kb_service.py b/rag/services/kb_service.py index 9ff2c20..0212d14 100644 --- a/rag/services/kb_service.py +++ b/rag/services/kb_service.py @@ -5,7 +5,7 @@ - Индексация документов (разбиение на чанки, получение эмбеддингов, сохранение в Qdrant и PostgreSQL). - Поиск релевантной информации по запросу (семантический поиск в Qdrant с фильтром по room_jid). - Управление доступом к документам (личные / глобальные / комнатные). -- Поддержка обновления версий документов в комнатах (замена по хешу). +- Поддержка обновления версий документов в комнатах (атомарное обновление: сначала создаётся новая версия, затем удаляется старая). Добавлена транзакционность индексации через флаг indexed в PostgreSQL. ВСЕ ВЫЗОВЫ QDRANT ОБЁРНУТЫ В asyncio.to_thread ДЛЯ ПРЕДОТВРАЩЕНИЯ БЛОКИРОВКИ EVENT LOOP. @@ -149,7 +149,7 @@ class KBService: return "" # ------------------------------------------------------------------ - # Индексация документов (транзакционная) + # Индексация документов (транзакционная, с атомарным обновлением) # ------------------------------------------------------------------ async def add_document( @@ -165,21 +165,25 @@ class KBService: update_if_exists: bool = True ) -> Tuple[int, int]: """ - Индексирует документ в базу знаний с транзакционностью: - 1. Создаёт запись в PostgreSQL с indexed = False. - 2. Разбивает на чанки, получает эмбеддинги, загружает в Qdrant. - 3. Если успешно – обновляет indexed = True. - 4. При ошибке – оставляет indexed = False (можно будет повторить). - ВСЕ ВЫЗОВЫ QDRANT ОБЁРНУТЫ В asyncio.to_thread. + Индексирует документ в базу знаний с транзакционностью и атомарным обновлением: + 1. Если документ уже существует в комнате и update_if_exists=True, + запоминаем старый doc_id, но НЕ удаляем его сразу. + 2. Создаём новую запись с indexed=False и новым doc_id. + 3. Разбиваем на чанки, получаем эмбеддинги, загружаем в Qdrant. + 4. Если успешно – обновляем indexed=True у новой записи. + 5. Удаляем старый документ (если он был) из Qdrant и PostgreSQL. + 6. При ошибке – оставляем старый документ нетронутым, новую запись удаляем (или оставляем с indexed=False для диагностики). + + Это гарантирует, что в любой момент времени либо старая версия доступна, либо новая. """ source_name = title if title else file_name doc_meta = metadata or {} if file_hash: doc_meta['file_hash'] = file_hash - # --- ОБНОВЛЕНИЕ ДОКУМЕНТА В КОМНАТЕ (замена версии) --- - # ⚠️ ВНИМАНИЕ: это место требует атомарного обновления (пункт 5). - # Пока оставляем как есть, но потом исправим. + old_doc_id = None + + # ---- 1. Поиск существующего документа в комнате (если нужно обновление) ---- if update_if_exists and room_jid and not is_global: existing = await self.db.find_document_by_name_and_room(source_name, room_jid) if existing: @@ -188,15 +192,10 @@ class KBService: if file_hash and old_hash == file_hash: logger.info(f"Документ {source_name} не изменился (хеш совпадает), пропускаем индексацию") return old_doc_id, 0 - # Удаляем старый документ из Qdrant и PostgreSQL - logger.info(f"Обновление документа {source_name} в комнате {room_jid}: удаляем старую версию doc_id={old_doc_id}") - # Удаление из Qdrant – обёрнуто в to_thread - await asyncio.to_thread(self.qdrant.delete_by_doc_id, old_doc_id) - # Удаляем запись из БД (включая права доступа) - await self.db.delete_documents_by_ids([old_doc_id]) + logger.info(f"Обновление документа {source_name} в комнате {room_jid}: старая версия doc_id={old_doc_id}") - # --- СОЗДАНИЕ ЗАПИСИ В POSTGRESQL С indexed=False --- - doc_id = await self.db.add_document( + # ---- 2. Создание НОВОЙ записи в PostgreSQL с indexed=False ---- + new_doc_id = await self.db.add_document( source_name=source_name, owner_jid=user_jid, is_global=is_global, @@ -206,8 +205,9 @@ class KBService: file_hash=file_hash, indexed=False # начальное состояние – не проиндексирован ) + logger.info(f"Создана новая запись документа doc_id={new_doc_id} для {source_name}") - # --- РАЗБИЕНИЕ НА ЧАНКИ --- + # ---- 3. Разбиение на чанки ---- chunks = split_into_chunks( file_text, chunk_size_tokens=self.chunk_size_tokens, @@ -217,20 +217,21 @@ class KBService: strategy="recursive_split_by_sentences" ) if not chunks: - # Если нет чанков – значит, документ пустой. Удаляем запись. - await self.db.delete_documents_by_ids([doc_id]) + # Если нет чанков – документ пустой. Удаляем новую запись и возвращаем старый doc_id (если был) + await self.db.delete_documents_by_ids([new_doc_id]) logger.warning(f"Документ {source_name} пуст, индексация отменена") - return doc_id, 0 + return old_doc_id or new_doc_id, 0 - # --- ПАКЕТНОЕ ПОЛУЧЕНИЕ ЭМБЕДДИНГОВ (dense + sparse) --- + # ---- 4. Пакетное получение эмбеддингов ---- dense_vectors = await self.embedding.embed_batch(chunks) if dense_vectors is None or len(dense_vectors) != len(chunks): - logger.error(f"Не удалось получить dense-эмбеддинги для документа {doc_id}, оставляем indexed=False") - return doc_id, 0 + logger.error(f"Не удалось получить dense-эмбеддинги для документа {new_doc_id}, оставляем indexed=False") + # Можно удалить новую запись, но оставим для диагностики + return new_doc_id, 0 sparse_vectors = await self.embedding.embed_sparse_batch(chunks) - # --- СОХРАНЕНИЕ ЧАНКОВ В QDRANT (обёрнуто в to_thread) --- + # ---- 5. Сохранение чанков в Qdrant (обёрнуто в to_thread) ---- chunk_count = 0 success = True try: @@ -246,7 +247,7 @@ class KBService: sparse_vec = sparse_vectors[idx] payload = { - "doc_id": doc_id, + "doc_id": new_doc_id, "user_jid": user_jid, "is_global": is_global, "chunk_index": idx, @@ -254,7 +255,6 @@ class KBService: "source_name": source_name, "room_jid": room_jid if room_jid else "", } - # ВЫЗОВ QDRANT ОБЁРНУТ В asyncio.to_thread await asyncio.to_thread( self.qdrant.add_chunk, dense=dense_vec, @@ -263,21 +263,31 @@ class KBService: ) chunk_count += 1 except Exception as e: - logger.error(f"Ошибка при добавлении чанков в Qdrant для doc_id={doc_id}: {e}") + logger.error(f"Ошибка при добавлении чанков в Qdrant для doc_id={new_doc_id}: {e}") success = False + # ---- 6. Если успешно – обновляем флаг indexed у НОВОЙ записи ---- if success and chunk_count > 0: - # Все чанки загружены успешно – обновляем флаг - await self.db.update_document_indexed(doc_id, True) + await self.db.update_document_indexed(new_doc_id, True) # Для личных документов предоставляем доступ владельцу if not is_global and not room_jid: - await self.db.add_access(doc_id, user_jid) - logger.info(f"Документ {source_name} (doc_id={doc_id}) проиндексирован, чанков: {chunk_count}") - else: - # Не удалось загрузить все чанки – оставляем indexed=False - logger.error(f"Документ {source_name} (doc_id={doc_id}) не проиндексирован (чанков: {chunk_count})") + await self.db.add_access(new_doc_id, user_jid) + logger.info(f"Новый документ {source_name} (doc_id={new_doc_id}) проиндексирован, чанков: {chunk_count}") - return doc_id, chunk_count + # ---- 7. Удаляем СТАРЫЙ документ (если был) ---- + if old_doc_id is not None: + logger.info(f"Удаление старой версии документа doc_id={old_doc_id}") + # Удаление из Qdrant + await asyncio.to_thread(self.qdrant.delete_by_doc_id, old_doc_id) + # Удаление из PostgreSQL (включая права доступа) + await self.db.delete_documents_by_ids([old_doc_id]) + logger.info(f"Старая версия doc_id={old_doc_id} удалена") + else: + # Не удалось загрузить все чанки – оставляем новую запись с indexed=False для диагностики + logger.error(f"Документ {source_name} (doc_id={new_doc_id}) не проиндексирован (чанков: {chunk_count})") + # Старый документ (если был) остаётся нетронутым + + return new_doc_id, chunk_count # ------------------------------------------------------------------ # Получение списков документов (используют методы db с фильтром indexed=TRUE) @@ -299,7 +309,6 @@ class KBService: """Удаляет все личные документы пользователя.""" doc_ids = await self.db.clear_user_kb(user_jid) for doc_id in doc_ids: - # Удаление из Qdrant – обёрнуто в to_thread await asyncio.to_thread(self.qdrant.delete_by_doc_id, doc_id) logger.info(f"Личная БЗ пользователя {user_jid} очищена (документов: {len(doc_ids)})")