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

rag/services/kb_service.py

@@ -7,11 +7,7 @@
 - Управление доступом к документам (личные / глобальные / комнатные).
 - Поддержка обновления версий документов в комнатах (замена по хешу).
 
-ВНИМАНИЕ: Поиск работает в двух режимах:
+Добавлена транзакционность индексации через флаг indexed в PostgreSQL.
-1. Семантический (dense) — для общих вопросов, использует эмбеддинги GigaChat
-2. Полнотекстовый (keyword) — для точных запросов (пункт 1.1, фразы в кавычках)
-
-Использует PostgreSQL (asyncpg) через PostgresService и Qdrant через QdrantService.
 """
 
 import asyncio
@@ -74,31 +70,15 @@ class KBService:
     async def find_relevant_info(self, query: str, user_jid: str, room_jid: str = None, top_k: int = 15) -> str:
         """
         Находит релевантные фрагменты в БЗ.
-        
+        Фильтрация по indexed = TRUE уже происходит на уровне Qdrant (фильтр доступа),
-        Режимы поиска:
+        но дополнительно мы не возвращаем документы, которые не проиндексированы.
-        - Если запрос содержит номер пункта (например, "1.1") или кавычки — 
-          используется полнотекстовый поиск (search_keywords)
-        - Иначе — семантический поиск (dense через GigaChat эмбеддинги)
-        
-        При ошибке эмбеддинга или пустом результате использует fallback-поиск по source_name.
-
-        Аргументы:
-            query: поисковый запрос пользователя
-            user_jid: JID пользователя (для фильтрации личных документов)
-            room_jid: JID комнаты (None для личного чата) – для фильтрации комнатных документов
-            top_k: количество возвращаемых чанков
-
-        Возвращает:
-            str: строка с объединёнными фрагментами вида "[источник: имя]\nтекст"
         """
         try:
             # Определяем, является ли запрос "точным"
-            # Точный запрос: содержит номер пункта (1.1, 2.3.4) или фразу в кавычках
             is_exact_query = bool(re.search(r'\d+\.\d+', query)) or '"' in query
             
             if is_exact_query:
-                # РЕЖИМ 1: Полнотекстовый поиск для точных запросов
+                # Полнотекстовый поиск (использует фильтр indexed=TRUE в Qdrant)
-                logger.debug(f"Используем полнотекстовый поиск для запроса: {query}")
                 search_results = self.qdrant.search_keywords(query, user_jid, room_jid, top_k=top_k)
             else:
                 # Получаем dense-вектор запроса
@@ -107,12 +87,10 @@ class KBService:
                     logger.error("Не удалось получить dense-эмбеддинг запроса, переключаемся на fallback")
                     return await self._fallback_search(query, user_jid, room_jid)
 
-                # Получаем sparse-вектор запроса (опционально, может быть None)
                 sparse_vector = await self.embedding.embed_sparse(query)
                 if sparse_vector is None:
                     logger.warning("Не удалось получить sparse-вектор запроса, используем только dense-поиск")
 
-                # Гибридный поиск
                 search_results = self.qdrant.search_hybrid(
                     dense_vector=query_vector,
                     sparse_vector=sparse_vector,
@@ -125,7 +103,6 @@ class KBService:
                 logger.info(f"По запросу '{query[:50]}' ничего не найдено в Qdrant")
                 return await self._fallback_search(query, user_jid, room_jid)
 
-            # Формируем контекст из найденных чанков
             context_parts = []
             for res in search_results:
                 payload = res.get('payload', {})
@@ -144,21 +121,12 @@ class KBService:
 
     async def _fallback_search(self, query: str, user_jid: str, room_jid: str = None) -> str:
         """
-        Полнотекстовый поиск по source_name в PostgreSQL.
+        Полнотекстовый поиск по source_name в PostgreSQL (только indexed=TRUE).
         Используется, если Qdrant недоступен или не вернул результатов.
-
-        Аргументы:
-            query: исходный запрос
-            user_jid: JID пользователя
-            room_jid: JID комнаты (None для личного чата)
-
-        Возвращает:
-            str: строка с информацией о найденных документах
         """
         if not query or len(query.strip()) < 3:
             return ""
         
-        # Очищаем запрос от спецсимволов для полнотекстового поиска
         clean_query = re.sub(r'[^\w\s]', ' ', query)
         
         try:
@@ -176,7 +144,7 @@ class KBService:
             return ""
 
     # ------------------------------------------------------------------
-    # Индексация документов
+    # Индексация документов (транзакционная)
     # ------------------------------------------------------------------
 
     async def add_document(
@@ -192,29 +160,11 @@ class KBService:
         update_if_exists: bool = True
     ) -> Tuple[int, int]:
         """
-        Индексирует документ в базу знаний.
+        Индексирует документ в базу знаний с транзакционностью:
-        
+        1. Создаёт запись в PostgreSQL с indexed = False.
-        Этапы индексации:
+        2. Разбивает на чанки, получает эмбеддинги, загружает в Qdrant.
-        1. Разбивает текст на чанки
+        3. Если успешно – обновляет indexed = True.
-        2. Получает эмбеддинги для каждого чанка через GigaChat
+        4. При ошибке – оставляет indexed = False (можно будет повторить).
-        3. Сохраняет метаданные в PostgreSQL
-        4. Сохраняет векторы и payload в Qdrant
-        5. Если update_if_exists=True и room_jid задан, проверяет наличие документа
-           с таким же source_name в комнате и при изменении хеша заменяет старую версию.
-
-        Аргументы:
-            file_name: исходное имя файла (или URL)
-            file_text: извлечённый текст документа
-            user_jid: JID владельца документа
-            is_global: глобальный ли документ
-            title: отображаемое название (если не указано, используется file_name)
-            metadata: дополнительные метаданные (сохраняются в JSONB)
-            room_jid: JID комнаты (если документ привязан к комнате)
-            file_hash: SHA-256 хеш содержимого (для проверки изменений)
-            update_if_exists: заменять ли существующий документ в комнате при изменении хеша
-
-        Возвращает:
-            Tuple[int, int]: (doc_id, количество_проиндексированных_чанков)
         """
         source_name = title if title else file_name
         doc_meta = metadata or {}
@@ -233,11 +183,10 @@ class KBService:
                 # Удаляем старый документ из Qdrant и PostgreSQL
                 logger.info(f"Обновление документа {source_name} в комнате {room_jid}: удаляем старую версию doc_id={old_doc_id}")
                 self.qdrant.delete_by_doc_id(old_doc_id)
-                async with self.db.pool.acquire() as conn:
+                # Удаляем запись из БД (включая права доступа)
-                    await conn.execute("DELETE FROM document_access WHERE doc_id = $1", old_doc_id)
+                await self.db.delete_documents_by_ids([old_doc_id])
-                    await conn.execute("DELETE FROM documents WHERE id = $1", old_doc_id)
 
-        # --- СОЗДАНИЕ ЗАПИСИ В POSTGRESQL ---
+        # --- СОЗДАНИЕ ЗАПИСИ В POSTGRESQL С indexed=False ---
         doc_id = await self.db.add_document(
             source_name=source_name,
             owner_jid=user_jid,
@@ -245,7 +194,8 @@ class KBService:
             collection_name=self.collection_name,
             metadata=doc_meta,
             room_jid=room_jid,
-            file_hash=file_hash
+            file_hash=file_hash,
+            indexed=False   # начальное состояние – не проиндексирован
         )
 
         # --- РАЗБИЕНИЕ НА ЧАНКИ ---
@@ -258,169 +208,102 @@ class KBService:
             strategy="recursive_split_by_sentences"
         )
         if not chunks:
+            # Если нет чанков – значит, документ пустой. Удаляем запись.
+            await self.db.delete_documents_by_ids([doc_id])
+            logger.warning(f"Документ {source_name} пуст, индексация отменена")
             return doc_id, 0
 
         # --- ПАКЕТНОЕ ПОЛУЧЕНИЕ ЭМБЕДДИНГОВ (dense + sparse) ---
         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}")
+            logger.error(f"Не удалось получить dense-эмбеддинги для документа {doc_id}, оставляем indexed=False")
+            # Можно также удалить запись, но оставим для возможности повтора
             return doc_id, 0
 
         sparse_vectors = await self.embedding.embed_sparse_batch(chunks)
-        # sparse_vectors может быть None или короче, но мы продолжим (sparse опционально)
 
         # --- СОХРАНЕНИЕ ЧАНКОВ В QDRANT ---
         chunk_count = 0
-        for idx, chunk in enumerate(chunks):
+        success = True
-            dense_vec = dense_vectors[idx] if idx < len(dense_vectors) else None
+        try:
-            if dense_vec is None:
+            for idx, chunk in enumerate(chunks):
-                logger.error(f"Нет dense-вектора для чанка {idx}, пропускаем")
+                dense_vec = dense_vectors[idx] if idx < len(dense_vectors) else None
-                continue
+                if dense_vec is None:
+                    logger.error(f"Нет dense-вектора для чанка {idx}, прерываем индексацию")
+                    success = False
+                    break
 
-            sparse_vec = None
+                sparse_vec = None
-            if sparse_vectors and idx < len(sparse_vectors):
+                if sparse_vectors and idx < len(sparse_vectors):
-                sparse_vec = sparse_vectors[idx]
+                    sparse_vec = sparse_vectors[idx]
 
-            payload = {
+                payload = {
-                "doc_id": doc_id,
+                    "doc_id": doc_id,
-                "user_jid": user_jid,
+                    "user_jid": user_jid,
-                "is_global": is_global,
+                    "is_global": is_global,
-                "chunk_index": idx,
+                    "chunk_index": idx,
-                "text": chunk,
+                    "text": chunk,
-                "source_name": source_name,
+                    "source_name": source_name,
-                "room_jid": room_jid if room_jid else "",
+                    "room_jid": room_jid if room_jid else "",
-            }
+                }
-            await asyncio.to_thread(
+                await asyncio.to_thread(
-                self.qdrant.add_chunk,
+                    self.qdrant.add_chunk,
-                dense=dense_vec,
+                    dense=dense_vec,
-                sparse=sparse_vec,
+                    sparse=sparse_vec,
-                payload=payload
+                    payload=payload
-            )
+                )
-            chunk_count += 1
+                chunk_count += 1
+        except Exception as e:
+            logger.error(f"Ошибка при добавлении чанков в Qdrant для doc_id={doc_id}: {e}")
+            success = False
 
-        # --- ДЛЯ ЛИЧНЫХ ДОКУМЕНТОВ ПРЕДОСТАВЛЯЕМ ДОСТУП ВЛАДЕЛЬЦУ ---
+        if success and chunk_count > 0:
-        if not is_global and not room_jid:
+            # Все чанки загружены успешно – обновляем флаг
-            await self.db.add_access(doc_id, user_jid)
+            await self.db.update_document_indexed(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})")
+            # Можно также удалить запись, но оставим для диагностики
 
-        logger.info(f"Документ {source_name} (doc_id={doc_id}) проиндексирован, чанков: {chunk_count}")
         return doc_id, chunk_count
 
     # ------------------------------------------------------------------
-    # Получение списков документов
+    # Получение списков документов (используют методы db с фильтром indexed=TRUE)
     # ------------------------------------------------------------------
 
     async def get_available_sources(self, user_jid: str, room_jid: str = None) -> List[str]:
-        """
+        """Возвращает список названий документов (source_name), доступных пользователю (только indexed=TRUE)."""
-        Возвращает список названий документов (source_name), доступных пользователю в данном контексте.
+        return await self.db.get_available_sources(user_jid, room_jid)
-        
-        Аргументы:
-            user_jid: JID пользователя
-            room_jid: JID комнаты (None для личного чата)
-
-        Возвращает:
-            List[str]: список названий
-        """
-        async with self.db.pool.acquire() as conn:
-            if room_jid:
-                rows = await conn.fetch(
-                    "SELECT source_name FROM documents WHERE room_jid = $1 ORDER BY source_name",
-                    room_jid
-                )
-            else:
-                rows = await conn.fetch(
-                    "SELECT source_name FROM documents WHERE (owner_jid = $1 OR is_global = TRUE) AND room_jid IS NULL ORDER BY source_name",
-                    user_jid
-                )
-            return [row['source_name'] for row in rows]
 
     async def get_sources_with_type(self, user_jid: str, room_jid: str = None) -> List[Dict[str, Any]]:
-        """
+        """Возвращает список документов с типом (только indexed=TRUE)."""
-        Возвращает список документов с указанием, является ли документ глобальным.
+        return await self.db.get_sources_with_type(user_jid, room_jid)
-        Используется командой !kb для раздельного отображения глобальных и локальных документов.
-
-        Аргументы:
-            user_jid: JID пользователя
-            room_jid: JID комнаты (None для личного чата)
-
-        Возвращает:
-            List[Dict[str, Any]]: список словарей [{"source_name": ..., "is_global": bool}]
-        """
-        async with self.db.pool.acquire() as conn:
-            if room_jid:
-                rows = await conn.fetch(
-                    "SELECT source_name, is_global FROM documents WHERE room_jid = $1 ORDER BY source_name",
-                    room_jid
-                )
-            else:
-                rows = await conn.fetch(
-                    "SELECT source_name, is_global FROM documents WHERE (owner_jid = $1 OR is_global = TRUE) AND room_jid IS NULL ORDER BY source_name",
-                    user_jid
-                )
-            return [{"source_name": row['source_name'], "is_global": row['is_global']} for row in rows]
 
     # ------------------------------------------------------------------
-    # Очистка баз знаний
+    # Очистка баз знаний (удаление всех документов, независимо от indexed)
     # ------------------------------------------------------------------
 
     async def clear_user_kb(self, user_jid: str) -> None:
-        """
+        """Удаляет все личные документы пользователя."""
-        Удаляет все личные документы пользователя (не глобальные и не комнатные).
+        doc_ids = await self.db.clear_user_kb(user_jid)
-        
-        Аргументы:
-            user_jid: JID пользователя
-        """
-        async with self.db.pool.acquire() as conn:
-            rows = await conn.fetch(
-                "SELECT id FROM documents WHERE owner_jid = $1 AND is_global = FALSE AND room_jid IS NULL",
-                user_jid
-            )
-            doc_ids = [row['id'] for row in rows]
-
-        # Удаляем из Qdrant
         for doc_id in doc_ids:
             self.qdrant.delete_by_doc_id(doc_id)
-
-        # Удаляем из PostgreSQL (используем ANY для массового удаления)
-        if doc_ids:
-            async with self.db.pool.acquire() as conn:
-                await conn.execute("DELETE FROM document_access WHERE doc_id = ANY($1)", doc_ids)
-                await conn.execute("DELETE FROM documents WHERE id = ANY($1)", doc_ids)
-
         logger.info(f"Личная БЗ пользователя {user_jid} очищена (документов: {len(doc_ids)})")
 
     async def clear_room_kb(self, room_jid: str) -> None:
-        """
+        """Удаляет все документы указанной комнаты."""
-        Удаляет все документы указанной комнаты.
+        doc_ids = await self.db.clear_room_kb(room_jid)
-        
-        Аргументы:
-            room_jid: JID комнаты
-        """
-        async with self.db.pool.acquire() as conn:
-            rows = await conn.fetch("SELECT id FROM documents WHERE room_jid = $1", room_jid)
-            doc_ids = [row['id'] for row in rows]
-
         for doc_id in doc_ids:
             self.qdrant.delete_by_doc_id(doc_id)
-
-        if doc_ids:
-            async with self.db.pool.acquire() as conn:
-                await conn.execute("DELETE FROM documents WHERE id = ANY($1)", doc_ids)
-
         logger.info(f"Комнатная БЗ {room_jid} очищена (документов: {len(doc_ids)})")
 
     async def clear_global_kb(self) -> None:
-        """
+        """Удаляет все глобальные документы."""
-        Удаляет все глобальные документы.
+        doc_ids = await self.db.clear_global_kb()
-        """
-        async with self.db.pool.acquire() as conn:
-            rows = await conn.fetch("SELECT id FROM documents WHERE is_global = TRUE")
-            doc_ids = [row['id'] for row in rows]
-
         for doc_id in doc_ids:
             self.qdrant.delete_by_doc_id(doc_id)
-
-        if doc_ids:
-            async with self.db.pool.acquire() as conn:
-                await conn.execute("DELETE FROM documents WHERE id = ANY($1)", doc_ids)
-
         logger.info("Глобальная БЗ очищена")