fckbot/rag/services/kb_service.py

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