diff --git a/core/services/kb_service.py b/core/services/kb_service.py new file mode 100644 index 0000000..ef99c82 --- /dev/null +++ b/core/services/kb_service.py @@ -0,0 +1,426 @@ +# -*- coding: utf-8 -*- +""" +Сервис базы знаний (Knowledge Base Service). +Реализует RAG-пайплайн: +- Индексация документов (разбиение на чанки, получение эмбеддингов, сохранение в Qdrant и PostgreSQL). +- Поиск релевантной информации по запросу (семантический поиск в Qdrant с фильтром по room_jid). +- Управление доступом к документам (личные / глобальные / комнатные). +- Поддержка обновления версий документов в комнатах (замена по хешу). + +ВНИМАНИЕ: Поиск работает в двух режимах: +1. Семантический (dense) — для общих вопросов, использует эмбеддинги GigaChat +2. Полнотекстовый (keyword) — для точных запросов (пункт 1.1, фразы в кавычках) + +Использует PostgreSQL (asyncpg) через PostgresService и Qdrant через QdrantService. +""" + +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: + """ + Находит релевантные фрагменты в БЗ. + + Режимы поиска: + - Если запрос содержит номер пункта (например, "1.1") или кавычки — + используется полнотекстовый поиск (search_keywords) + - Иначе — семантический поиск (dense через GigaChat эмбеддинги) + + При ошибке эмбеддинга или пустом результате использует fallback-поиск по source_name. + + Аргументы: + query: поисковый запрос пользователя + user_jid: JID пользователя (для фильтрации личных документов) + room_jid: JID комнаты (None для личного чата) – для фильтрации комнатных документов + top_k: количество возвращаемых чанков + + Возвращает: + str: строка с объединёнными фрагментами вида "[источник: имя]\nтекст" + """ + 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}") + search_results = self.qdrant.search_keywords(query, user_jid, room_jid, top_k=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-вектор запроса (опционально, может быть 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, + 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. + Используется, если 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: + 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. Разбивает текст на чанки + 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, количество_проиндексированных_чанков) + """ + source_name = title if title else file_name + doc_meta = metadata or {} + if file_hash: + doc_meta['file_hash'] = file_hash + + # --- ОБНОВЛЕНИЕ ДОКУМЕНТА В КОМНАТЕ (замена версии) --- + 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 + # Удаляем старый документ из 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) + + # --- СОЗДАНИЕ ЗАПИСИ В POSTGRESQL --- + 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 + ) + + # --- РАЗБИЕНИЕ НА ЧАНКИ --- + 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: + 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}") + 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 + + 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 + + # --- ДЛЯ ЛИЧНЫХ ДОКУМЕНТОВ ПРЕДОСТАВЛЯЕМ ДОСТУП ВЛАДЕЛЬЦУ --- + 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}") + return doc_id, chunk_count + + # ------------------------------------------------------------------ + # Получение списков документов + # ------------------------------------------------------------------ + + 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] + + 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] + + # ------------------------------------------------------------------ + # Очистка баз знаний + # ------------------------------------------------------------------ + + 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 + 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] + + 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] + + 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