markov/fckbot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Update file kb_service.py

Browse Source
This commit is contained in:
Markov Andrey
2026-06-30 16:07:12 +00:00
parent 6713e52fa6
commit bfacbad0de
1 changed files with 48 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files

87
rag/services/kb_service.py
View File

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