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

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

Browse Source
This commit is contained in:
Markov Andrey
2026-06-30 12:01:03 +00:00
parent 58ac79feac
commit 4b57abf269
1 changed files with 74 additions and 191 deletions
Download Patch File Download Diff File Expand all files Collapse all files

265
rag/services/kb_service.py
View File

@@ -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("Глобальная БЗ очищена")