diff --git a/rag/services/kb_service.py b/rag/services/kb_service.py index ef99c82..7582762 100644 --- a/rag/services/kb_service.py +++ b/rag/services/kb_service.py @@ -7,11 +7,7 @@ - Управление доступом к документам (личные / глобальные / комнатные). - Поддержка обновления версий документов в комнатах (замена по хешу). -ВНИМАНИЕ: Поиск работает в двух режимах: -1. Семантический (dense) — для общих вопросов, использует эмбеддинги GigaChat -2. Полнотекстовый (keyword) — для точных запросов (пункт 1.1, фразы в кавычках) - -Использует PostgreSQL (asyncpg) через PostgresService и Qdrant через QdrantService. +Добавлена транзакционность индексации через флаг indexed в PostgreSQL. """ 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: """ Находит релевантные фрагменты в БЗ. - - Режимы поиска: - - Если запрос содержит номер пункта (например, "1.1") или кавычки — - используется полнотекстовый поиск (search_keywords) - - Иначе — семантический поиск (dense через GigaChat эмбеддинги) - - При ошибке эмбеддинга или пустом результате использует fallback-поиск по source_name. - - Аргументы: - query: поисковый запрос пользователя - user_jid: JID пользователя (для фильтрации личных документов) - room_jid: JID комнаты (None для личного чата) – для фильтрации комнатных документов - top_k: количество возвращаемых чанков - - Возвращает: - str: строка с объединёнными фрагментами вида "[источник: имя]\nтекст" + Фильтрация по indexed = TRUE уже происходит на уровне Qdrant (фильтр доступа), + но дополнительно мы не возвращаем документы, которые не проиндексированы. """ try: # Определяем, является ли запрос "точным" - # Точный запрос: содержит номер пункта (1.1, 2.3.4) или фразу в кавычках is_exact_query = bool(re.search(r'\d+\.\d+', query)) or '"' in query if is_exact_query: - # РЕЖИМ 1: Полнотекстовый поиск для точных запросов - logger.debug(f"Используем полнотекстовый поиск для запроса: {query}") + # Полнотекстовый поиск (использует фильтр indexed=TRUE в Qdrant) 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. Разбивает текст на чанки - 2. Получает эмбеддинги для каждого чанка через GigaChat - 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, количество_проиндексированных_чанков) + Индексирует документ в базу знаний с транзакционностью: + 1. Создаёт запись в PostgreSQL с indexed = False. + 2. Разбивает на чанки, получает эмбеддинги, загружает в Qdrant. + 3. Если успешно – обновляет indexed = True. + 4. При ошибке – оставляет indexed = False (можно будет повторить). """ 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 conn.execute("DELETE FROM documents WHERE id = $1", old_doc_id) + # Удаляем запись из БД (включая права доступа) + await self.db.delete_documents_by_ids([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): - dense_vec = dense_vectors[idx] if idx < len(dense_vectors) else None - if dense_vec is None: - logger.error(f"Нет dense-вектора для чанка {idx}, пропускаем") - continue + success = True + try: + for idx, chunk in enumerate(chunks): + dense_vec = dense_vectors[idx] if idx < len(dense_vectors) else None + if dense_vec is None: + logger.error(f"Нет dense-вектора для чанка {idx}, прерываем индексацию") + success = False + break - sparse_vec = None - if sparse_vectors and idx < len(sparse_vectors): - sparse_vec = sparse_vectors[idx] + sparse_vec = None + if sparse_vectors and idx < len(sparse_vectors): + sparse_vec = sparse_vectors[idx] - payload = { - "doc_id": doc_id, - "user_jid": user_jid, - "is_global": is_global, - "chunk_index": idx, - "text": chunk, - "source_name": source_name, - "room_jid": room_jid if room_jid else "", - } - await asyncio.to_thread( - self.qdrant.add_chunk, - dense=dense_vec, - sparse=sparse_vec, - payload=payload - ) - chunk_count += 1 + payload = { + "doc_id": doc_id, + "user_jid": user_jid, + "is_global": is_global, + "chunk_index": idx, + "text": chunk, + "source_name": source_name, + "room_jid": room_jid if room_jid else "", + } + await asyncio.to_thread( + self.qdrant.add_chunk, + dense=dense_vec, + sparse=sparse_vec, + payload=payload + ) + chunk_count += 1 + except Exception as e: + logger.error(f"Ошибка при добавлении чанков в Qdrant для doc_id={doc_id}: {e}") + success = False - # --- ДЛЯ ЛИЧНЫХ ДОКУМЕНТОВ ПРЕДОСТАВЛЯЕМ ДОСТУП ВЛАДЕЛЬЦУ --- - if not is_global and not room_jid: - await self.db.add_access(doc_id, user_jid) + if success and chunk_count > 0: + # Все чанки загружены успешно – обновляем флаг + 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), доступных пользователю в данном контексте. - - Аргументы: - 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] + """Возвращает список названий документов (source_name), доступных пользователю (только indexed=TRUE).""" + return await self.db.get_available_sources(user_jid, room_jid) async def get_sources_with_type(self, user_jid: str, room_jid: str = None) -> List[Dict[str, Any]]: - """ - Возвращает список документов с указанием, является ли документ глобальным. - Используется командой !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=TRUE).""" + return await self.db.get_sources_with_type(user_jid, room_jid) # ------------------------------------------------------------------ - # Очистка баз знаний + # Очистка баз знаний (удаление всех документов, независимо от indexed) # ------------------------------------------------------------------ async def clear_user_kb(self, user_jid: str) -> None: - """ - Удаляет все личные документы пользователя (не глобальные и не комнатные). - - Аргументы: - 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 + """Удаляет все личные документы пользователя.""" + doc_ids = await self.db.clear_user_kb(user_jid) 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: - """ - Удаляет все документы указанной комнаты. - - Аргументы: - 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] - + """Удаляет все документы указанной комнаты.""" + doc_ids = await self.db.clear_room_kb(room_jid) 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: - """ - Удаляет все глобальные документы. - """ - 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] - + """Удаляет все глобальные документы.""" + doc_ids = await self.db.clear_global_kb() 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("Глобальная БЗ очищена") \ No newline at end of file