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