327 lines
17 KiB
Python
327 lines
17 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
Сервис базы знаний (Knowledge Base Service).
|
||
Реализует RAG-пайплайн:
|
||
- Индексация документов (разбиение на чанки, получение эмбеддингов, сохранение в Qdrant и PostgreSQL).
|
||
- Поиск релевантной информации по запросу (семантический поиск в Qdrant с фильтром по room_jid).
|
||
- Управление доступом к документам (личные / глобальные / комнатные).
|
||
- Поддержка обновления версий документов в комнатах (атомарное обновление: сначала создаётся новая версия, затем удаляется старая).
|
||
|
||
Добавлена транзакционность индексации через флаг indexed в PostgreSQL.
|
||
ВСЕ ВЫЗОВЫ QDRANT ОБЁРНУТЫ В asyncio.to_thread ДЛЯ ПРЕДОТВРАЩЕНИЯ БЛОКИРОВКИ EVENT LOOP.
|
||
"""
|
||
|
||
import asyncio
|
||
import logging
|
||
import re
|
||
from typing import List, Dict, Any, Optional, Tuple
|
||
|
||
# Импорт сервисов
|
||
from core.services.postgres_service import PostgresService
|
||
from core.services.qdrant_service import QdrantService
|
||
from core.services.embedding_service import EmbeddingService
|
||
from core.utils.text_utils import split_into_chunks
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class KBService:
|
||
"""
|
||
База знаний: индексация документов, поиск, управление.
|
||
Абстрагирует работу с PostgreSQL и Qdrant для RAG.
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
db: PostgresService,
|
||
qdrant: QdrantService,
|
||
embedding: EmbeddingService,
|
||
collection_name: str,
|
||
chunk_size_tokens: int,
|
||
overlap_tokens: int,
|
||
approx_chunk_chars: int,
|
||
approx_overlap_chars: int,
|
||
):
|
||
"""
|
||
Инициализация KBService.
|
||
|
||
Аргументы:
|
||
db: сервис PostgreSQL (asyncpg)
|
||
qdrant: сервис Qdrant
|
||
embedding: сервис эмбеддингов GigaChat
|
||
collection_name: имя коллекции в Qdrant
|
||
chunk_size_tokens: размер чанка в токенах (для tiktoken)
|
||
overlap_tokens: перекрытие между чанками в токенах
|
||
approx_chunk_chars: запасной размер чанка в символах (если tiktoken недоступен)
|
||
approx_overlap_chars: запасное перекрытие в символах
|
||
"""
|
||
self.db = db
|
||
self.qdrant = qdrant
|
||
self.embedding = embedding
|
||
self.collection_name = collection_name
|
||
self.chunk_size_tokens = chunk_size_tokens
|
||
self.overlap_tokens = overlap_tokens
|
||
self.approx_chunk_chars = approx_chunk_chars
|
||
self.approx_overlap_chars = approx_overlap_chars
|
||
|
||
# ------------------------------------------------------------------
|
||
# Поиск релевантной информации
|
||
# ------------------------------------------------------------------
|
||
|
||
async def find_relevant_info(self, query: str, user_jid: str, room_jid: str = None, top_k: int = 15) -> str:
|
||
"""
|
||
Находит релевантные фрагменты в БЗ.
|
||
ВСЕ ВЫЗОВЫ QDRANT ОБЁРНУТЫ В asyncio.to_thread.
|
||
"""
|
||
try:
|
||
# Определяем, является ли запрос "точным"
|
||
is_exact_query = bool(re.search(r'\d+\.\d+', query)) or '"' in query
|
||
|
||
if is_exact_query:
|
||
# Полнотекстовый поиск – обёрнуто в to_thread
|
||
search_results = await asyncio.to_thread(
|
||
self.qdrant.search_keywords,
|
||
query, user_jid, room_jid, top_k
|
||
)
|
||
else:
|
||
# Получаем dense-вектор запроса
|
||
query_vector = await self.embedding.embed(query)
|
||
if query_vector is None:
|
||
logger.error("Не удалось получить dense-эмбеддинг запроса, переключаемся на fallback")
|
||
return await self._fallback_search(query, user_jid, room_jid)
|
||
|
||
sparse_vector = await self.embedding.embed_sparse(query)
|
||
if sparse_vector is None:
|
||
logger.warning("Не удалось получить sparse-вектор запроса, используем только dense-поиск")
|
||
|
||
# Гибридный поиск – обёрнуто в to_thread
|
||
search_results = await asyncio.to_thread(
|
||
self.qdrant.search_hybrid,
|
||
dense_vector=query_vector,
|
||
sparse_vector=sparse_vector,
|
||
user_jid=user_jid,
|
||
room_jid=room_jid,
|
||
top_k=top_k
|
||
)
|
||
|
||
if not search_results:
|
||
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', {})
|
||
source_name = payload.get('source_name', 'Неизвестный источник')
|
||
chunk_text = payload.get('text', '')
|
||
if chunk_text:
|
||
context_parts.append(f"[источник: {source_name}]\n{chunk_text}")
|
||
|
||
context = "\n\n".join(context_parts)
|
||
logger.debug(f"Сформирован контекст длиной {len(context)} символов")
|
||
return context
|
||
|
||
except Exception as e:
|
||
logger.exception(f"Ошибка при поиске в RAG: {e}")
|
||
return await self._fallback_search(query, user_jid, room_jid)
|
||
|
||
async def _fallback_search(self, query: str, user_jid: str, room_jid: str = None) -> str:
|
||
"""
|
||
Полнотекстовый поиск по source_name в PostgreSQL (только indexed=TRUE).
|
||
Используется, если Qdrant недоступен или не вернул результатов.
|
||
"""
|
||
if not query or len(query.strip()) < 3:
|
||
return ""
|
||
|
||
clean_query = re.sub(r'[^\w\s]', ' ', query)
|
||
|
||
try:
|
||
rows = await self.db.fulltext_search(clean_query, user_jid, room_jid)
|
||
if not rows:
|
||
return ""
|
||
|
||
parts = [
|
||
f"[источник: {row['source_name']}]\n(Документ доступен, но текст не проиндексирован в Qdrant. Повторите запрос позже.)"
|
||
for row in rows
|
||
]
|
||
return "\n\n".join(parts)
|
||
except Exception as e:
|
||
logger.error(f"Ошибка при fallback поиске: {e}")
|
||
return ""
|
||
|
||
# ------------------------------------------------------------------
|
||
# Индексация документов (транзакционная, с атомарным обновлением)
|
||
# ------------------------------------------------------------------
|
||
|
||
async def add_document(
|
||
self,
|
||
file_name: str,
|
||
file_text: str,
|
||
user_jid: str,
|
||
is_global: bool,
|
||
title: Optional[str] = None,
|
||
metadata: Optional[Dict] = None,
|
||
room_jid: str = None,
|
||
file_hash: str = None,
|
||
update_if_exists: bool = True
|
||
) -> Tuple[int, int]:
|
||
"""
|
||
Индексирует документ в базу знаний с транзакционностью и атомарным обновлением:
|
||
1. Если документ уже существует в комнате и update_if_exists=True,
|
||
запоминаем старый doc_id, но НЕ удаляем его сразу.
|
||
2. Создаём новую запись с indexed=False и новым doc_id.
|
||
3. Разбиваем на чанки, получаем эмбеддинги, загружаем в Qdrant.
|
||
4. Если успешно – обновляем indexed=True у новой записи.
|
||
5. Удаляем старый документ (если он был) из Qdrant и PostgreSQL.
|
||
6. При ошибке – оставляем старый документ нетронутым, новую запись удаляем (или оставляем с indexed=False для диагностики).
|
||
|
||
Это гарантирует, что в любой момент времени либо старая версия доступна, либо новая.
|
||
"""
|
||
source_name = title if title else file_name
|
||
doc_meta = metadata or {}
|
||
if file_hash:
|
||
doc_meta['file_hash'] = file_hash
|
||
|
||
old_doc_id = None
|
||
|
||
# ---- 1. Поиск существующего документа в комнате (если нужно обновление) ----
|
||
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)
|
||
if existing:
|
||
old_doc_id = existing['id']
|
||
old_hash = existing.get('file_hash')
|
||
if file_hash and old_hash == file_hash:
|
||
logger.info(f"Документ {source_name} не изменился (хеш совпадает), пропускаем индексацию")
|
||
return old_doc_id, 0
|
||
logger.info(f"Обновление документа {source_name} в комнате {room_jid}: старая версия doc_id={old_doc_id}")
|
||
|
||
# ---- 2. Создание НОВОЙ записи в PostgreSQL с indexed=False ----
|
||
new_doc_id = await self.db.add_document(
|
||
source_name=source_name,
|
||
owner_jid=user_jid,
|
||
is_global=is_global,
|
||
collection_name=self.collection_name,
|
||
metadata=doc_meta,
|
||
room_jid=room_jid,
|
||
file_hash=file_hash,
|
||
indexed=False # начальное состояние – не проиндексирован
|
||
)
|
||
logger.info(f"Создана новая запись документа doc_id={new_doc_id} для {source_name}")
|
||
|
||
# ---- 3. Разбиение на чанки ----
|
||
chunks = split_into_chunks(
|
||
file_text,
|
||
chunk_size_tokens=self.chunk_size_tokens,
|
||
overlap_tokens=self.overlap_tokens,
|
||
approx_chunk_chars=self.approx_chunk_chars,
|
||
approx_overlap_chars=self.approx_overlap_chars,
|
||
strategy="recursive_split_by_sentences"
|
||
)
|
||
if not chunks:
|
||
# Если нет чанков – документ пустой. Удаляем новую запись и возвращаем старый doc_id (если был)
|
||
await self.db.delete_documents_by_ids([new_doc_id])
|
||
logger.warning(f"Документ {source_name} пуст, индексация отменена")
|
||
return old_doc_id or new_doc_id, 0
|
||
|
||
# ---- 4. Пакетное получение эмбеддингов ----
|
||
dense_vectors = await self.embedding.embed_batch(chunks)
|
||
if dense_vectors is None or len(dense_vectors) != len(chunks):
|
||
logger.error(f"Не удалось получить dense-эмбеддинги для документа {new_doc_id}, оставляем indexed=False")
|
||
# Можно удалить новую запись, но оставим для диагностики
|
||
return new_doc_id, 0
|
||
|
||
sparse_vectors = await self.embedding.embed_sparse_batch(chunks)
|
||
|
||
# ---- 5. Сохранение чанков в Qdrant (обёрнуто в to_thread) ----
|
||
chunk_count = 0
|
||
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]
|
||
|
||
payload = {
|
||
"doc_id": new_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={new_doc_id}: {e}")
|
||
success = False
|
||
|
||
# ---- 6. Если успешно – обновляем флаг indexed у НОВОЙ записи ----
|
||
if success and chunk_count > 0:
|
||
await self.db.update_document_indexed(new_doc_id, True)
|
||
# Для личных документов предоставляем доступ владельцу
|
||
if not is_global and not room_jid:
|
||
await self.db.add_access(new_doc_id, user_jid)
|
||
logger.info(f"Новый документ {source_name} (doc_id={new_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)
|
||
# ------------------------------------------------------------------
|
||
|
||
async def get_available_sources(self, user_jid: str, room_jid: str = None) -> List[str]:
|
||
"""Возвращает список названий документов (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]]:
|
||
"""Возвращает список документов с типом (только 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:
|
||
"""Удаляет все личные документы пользователя."""
|
||
doc_ids = await self.db.clear_user_kb(user_jid)
|
||
for doc_id in doc_ids:
|
||
await asyncio.to_thread(self.qdrant.delete_by_doc_id, doc_id)
|
||
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)
|
||
for doc_id in doc_ids:
|
||
await asyncio.to_thread(self.qdrant.delete_by_doc_id, doc_id)
|
||
logger.info(f"Комнатная БЗ {room_jid} очищена (документов: {len(doc_ids)})")
|
||
|
||
async def clear_global_kb(self) -> None:
|
||
"""Удаляет все глобальные документы."""
|
||
doc_ids = await self.db.clear_global_kb()
|
||
for doc_id in doc_ids:
|
||
await asyncio.to_thread(self.qdrant.delete_by_doc_id, doc_id)
|
||
logger.info("Глобальная БЗ очищена") |