Update file kb_service.py

This commit is contained in:
Markov Andrey
2026-06-30 15:52:37 +00:00
parent 5dfefad534
commit 3b1c4e2f56

View File

@@ -8,6 +8,7 @@
- Поддержка обновления версий документов в комнатах (замена по хешу). - Поддержка обновления версий документов в комнатах (замена по хешу).
Добавлена транзакционность индексации через флаг indexed в PostgreSQL. Добавлена транзакционность индексации через флаг indexed в PostgreSQL.
ВСЕ ВЫЗОВЫ QDRANT ОБЁРНУТЫ В asyncio.to_thread ДЛЯ ПРЕДОТВРАЩЕНИЯ БЛОКИРОВКИ EVENT LOOP.
""" """
import asyncio import asyncio
@@ -70,16 +71,18 @@ class KBService:
async def find_relevant_info(self, query: str, user_jid: str, room_jid: str = None, top_k: int = 15) -> str: async def find_relevant_info(self, query: str, user_jid: str, room_jid: str = None, top_k: int = 15) -> str:
""" """
Находит релевантные фрагменты в БЗ. Находит релевантные фрагменты в БЗ.
Фильтрация по indexed = TRUE уже происходит на уровне Qdrant (фильтр доступа), ВСЕ ВЫЗОВЫ QDRANT ОБЁРНУТЫ В asyncio.to_thread.
но дополнительно мы не возвращаем документы, которые не проиндексированы.
""" """
try: try:
# Определяем, является ли запрос "точным" # Определяем, является ли запрос "точным"
is_exact_query = bool(re.search(r'\d+\.\d+', query)) or '"' in query is_exact_query = bool(re.search(r'\d+\.\d+', query)) or '"' in query
if is_exact_query: if is_exact_query:
# Полнотекстовый поиск (использует фильтр indexed=TRUE в Qdrant) # Полнотекстовый поиск обёрнуто в to_thread
search_results = self.qdrant.search_keywords(query, user_jid, room_jid, top_k=top_k) search_results = await asyncio.to_thread(
self.qdrant.search_keywords,
query, user_jid, room_jid, top_k
)
else: else:
# Получаем dense-вектор запроса # Получаем dense-вектор запроса
query_vector = await self.embedding.embed(query) query_vector = await self.embedding.embed(query)
@@ -91,7 +94,9 @@ class KBService:
if sparse_vector is None: if sparse_vector is None:
logger.warning("Не удалось получить sparse-вектор запроса, используем только dense-поиск") logger.warning("Не удалось получить sparse-вектор запроса, используем только dense-поиск")
search_results = self.qdrant.search_hybrid( # Гибридный поиск обёрнуто в to_thread
search_results = await asyncio.to_thread(
self.qdrant.search_hybrid,
dense_vector=query_vector, dense_vector=query_vector,
sparse_vector=sparse_vector, sparse_vector=sparse_vector,
user_jid=user_jid, user_jid=user_jid,
@@ -165,6 +170,7 @@ class KBService:
2. Разбивает на чанки, получает эмбеддинги, загружает в Qdrant. 2. Разбивает на чанки, получает эмбеддинги, загружает в Qdrant.
3. Если успешно обновляет indexed = True. 3. Если успешно обновляет indexed = True.
4. При ошибке оставляет indexed = False (можно будет повторить). 4. При ошибке оставляет indexed = False (можно будет повторить).
ВСЕ ВЫЗОВЫ QDRANT ОБЁРНУТЫ В asyncio.to_thread.
""" """
source_name = title if title else file_name source_name = title if title else file_name
doc_meta = metadata or {} doc_meta = metadata or {}
@@ -172,6 +178,8 @@ class KBService:
doc_meta['file_hash'] = file_hash doc_meta['file_hash'] = file_hash
# --- ОБНОВЛЕНИЕ ДОКУМЕНТА В КОМНАТЕ (замена версии) --- # --- ОБНОВЛЕНИЕ ДОКУМЕНТА В КОМНАТЕ (замена версии) ---
# ⚠️ ВНИМАНИЕ: это место требует атомарного обновления (пункт 5).
# Пока оставляем как есть, но потом исправим.
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:
@@ -182,7 +190,8 @@ class KBService:
return old_doc_id, 0 return old_doc_id, 0
# Удаляем старый документ из Qdrant и PostgreSQL # Удаляем старый документ из 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}")
self.qdrant.delete_by_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]) await self.db.delete_documents_by_ids([old_doc_id])
@@ -217,12 +226,11 @@ class KBService:
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-эмбеддинги для документа {doc_id}, оставляем indexed=False")
# Можно также удалить запись, но оставим для возможности повтора
return doc_id, 0 return doc_id, 0
sparse_vectors = await self.embedding.embed_sparse_batch(chunks) sparse_vectors = await self.embedding.embed_sparse_batch(chunks)
# --- СОХРАНЕНИЕ ЧАНКОВ В QDRANT --- # --- СОХРАНЕНИЕ ЧАНКОВ В QDRANT (обёрнуто в to_thread) ---
chunk_count = 0 chunk_count = 0
success = True success = True
try: try:
@@ -246,6 +254,7 @@ 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,
@@ -267,7 +276,6 @@ class KBService:
else: else:
# Не удалось загрузить все чанки оставляем indexed=False # Не удалось загрузить все чанки оставляем indexed=False
logger.error(f"Документ {source_name} (doc_id={doc_id}) не проиндексирован (чанков: {chunk_count})") logger.error(f"Документ {source_name} (doc_id={doc_id}) не проиндексирован (чанков: {chunk_count})")
# Можно также удалить запись, но оставим для диагностики
return doc_id, chunk_count return doc_id, chunk_count
@@ -291,19 +299,20 @@ 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:
self.qdrant.delete_by_doc_id(doc_id) # Удаление из Qdrant обёрнуто в to_thread
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)})")
async def clear_room_kb(self, room_jid: str) -> None: async def clear_room_kb(self, room_jid: str) -> None:
"""Удаляет все документы указанной комнаты.""" """Удаляет все документы указанной комнаты."""
doc_ids = await self.db.clear_room_kb(room_jid) doc_ids = await self.db.clear_room_kb(room_jid)
for doc_id in doc_ids: for doc_id in doc_ids:
self.qdrant.delete_by_doc_id(doc_id) await asyncio.to_thread(self.qdrant.delete_by_doc_id, doc_id)
logger.info(f"Комнатная БЗ {room_jid} очищена (документов: {len(doc_ids)})") logger.info(f"Комнатная БЗ {room_jid} очищена (документов: {len(doc_ids)})")
async def clear_global_kb(self) -> None: async def clear_global_kb(self) -> None:
"""Удаляет все глобальные документы.""" """Удаляет все глобальные документы."""
doc_ids = await self.db.clear_global_kb() doc_ids = await self.db.clear_global_kb()
for doc_id in doc_ids: for doc_id in doc_ids:
self.qdrant.delete_by_doc_id(doc_id) await asyncio.to_thread(self.qdrant.delete_by_doc_id, doc_id)
logger.info("Глобальная БЗ очищена") logger.info("Глобальная БЗ очищена")