Add new file

core/rag_orchestrator.py

@@ -0,0 +1,506 @@
+# -*- coding: utf-8 -*-
+"""
+Главный оркестратор RAG-пайплайна.
+Принимает запрос пользователя, получает историю из БД, выполняет все этапы:
+классификацию, расширение, поиск, переранжирование, синтез, самокритику.
+Сохраняет историю в БД.
+
+Этот модуль не зависит от XMPP и может использоваться как в ботах,
+так и в отдельном RAG-ядре (API-сервисе).
+"""
+
+import asyncio
+import logging
+import os
+import re
+from typing import Optional, Dict, List, Any
+
+# Импорт сервисов и функций
+from core.services.postgres_service import PostgresService
+from core.services.qdrant_service import QdrantService
+from core.services.embedding_service import EmbeddingService
+from core.services.kb_service import KBService
+from core.services.giga_client import GigaClient
+from core.services.file_service import FileService
+from core.functions.intent_classify import classify_intent
+from core.functions.expand_query import expand_query
+from core.functions.extract_metrics import extract_metrics
+from core.functions.summarize_document import summarize_document
+from core.functions.check_consistency import check_consistency
+from core.functions.critique_answer import critique_answer
+from core.functions.rerank_context import rerank_context
+from core.functions.check_spelling import check_spelling
+from core.utils.text_utils import count_tokens
+
+logger = logging.getLogger(__name__)
+
+
+class RAGOrchestrator:
+    """
+    Оркестратор RAG-пайплайна.
+    Содержит ссылки на все сервисы и выполняет полный цикл обработки запроса.
+    """
+
+    def __init__(
+        self,
+        db: PostgresService,
+        qdrant: QdrantService,
+        embedding: EmbeddingService,
+        kb: KBService,
+        giga: GigaClient,
+        files: FileService,
+        config,
+        default_prompts: Optional[Dict[str, str]] = None
+    ):
+        """
+        Инициализация оркестратора.
+
+        Аргументы:
+            db: сервис PostgreSQL (для истории и метаданных)
+            qdrant: сервис Qdrant (векторный поиск)
+            embedding: сервис эмбеддингов (GigaChat)
+            kb: сервис базы знаний (индексация, поиск)
+            giga: клиент GigaChat (генерация)
+            files: сервис файлов (извлечение текста)
+            config: объект конфигурации (BotConfig)
+            default_prompts: словарь промптов по умолчанию
+        """
+        self.db = db
+        self.qdrant = qdrant
+        self.embedding = embedding
+        self.kb = kb
+        self.giga = giga
+        self.files = files
+        self.config = config
+        self.default_prompts = default_prompts or {}
+        logger.info("RAGOrchestrator инициализирован")
+
+    async def process_query(
+        self,
+        query: str,
+        user_jid: str,
+        room_jid: Optional[str],
+        prompts: Optional[Dict[str, str]] = None,
+        intent_override: Optional[str] = None,
+        last_file_path: Optional[str] = None,
+        last_file_text: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        """
+        Основной метод обработки запроса.
+
+        ИСТОРИЯ ДИАЛОГА ПОЛУЧАЕТСЯ ИЗ БД, а не из запроса.
+        Это обеспечивает единый контекст для всех клиентов (XMPP, Telegram и т.д.).
+
+        Аргументы:
+            query (str): текст запроса пользователя
+            user_jid (str): JID пользователя (без ресурса)
+            room_jid (Optional[str]): JID комнаты (None для личного чата)
+            prompts (Optional[Dict[str, str]]): словарь промптов для текущего запроса.
+                Если не передан, используются default_prompts.
+            intent_override (Optional[str]): принудительное переопределение намерения
+            last_file_path (Optional[str]): путь к последнему загруженному файлу
+            last_file_text (Optional[str]): текст последнего загруженного файла
+
+        Возвращает:
+            Dict[str, Any]: словарь с ключами:
+                - answer (str): итоговый ответ
+                - intent (str): распознанное намерение
+                - context (str): использованный контекст (для отладки)
+                - sources (List[str]): список источников
+                - confidence (float): оценка уверенности (если есть)
+        """
+        # ----- 1. Подготовка промптов -----
+        # Если промпты не переданы, используем стандартные (загруженные из файлов)
+        if prompts is None:
+            prompts = self.default_prompts.copy()
+        # Базовый шаблон синтеза: подставляет контекст и запрос
+        synthesis_template = prompts.get('synthesis', '')
+        if not synthesis_template:
+            synthesis_template = "{context}\n\n{query}\n\nОтвет:"
+        system_prompt = prompts.get('system', None)
+
+        # ----- 2. Получение истории диалога из БД -----
+        # Запрашиваем последние 100 сообщений для данного пользователя/комнаты
+        history = await self.db.get_history(user_jid, room_jid, limit=100)
+
+        # ----- 3. Обрезка истории по токенам (чтобы не превысить лимит модели) -----
+        max_context_tokens = getattr(self.config, 'max_context_tokens', 3000)
+        total_tokens = 0
+        truncated_history = []
+        # Идём с конца (свежие сообщения) к началу, чтобы сохранить самые последние
+        for record in reversed(history):
+            tokens = count_tokens(record['content'])
+            if total_tokens + tokens <= max_context_tokens:
+                truncated_history.append(record)
+                total_tokens += tokens
+            else:
+                break
+        # Восстанавливаем хронологический порядок (от старого к новому)
+        truncated_history.reverse()
+        # Преобразуем в формат, ожидаемый GigaChat: [{"role": "user/assistant", "content": "..."}]
+        formatted_history = truncated_history
+
+        logger.debug(
+            f"История из БД: {len(truncated_history)} сообщений, "
+            f"{total_tokens} токенов из {max_context_tokens} допустимых"
+        )
+
+        # ----- 4. Классификация намерений (если не переопределена) -----
+        intent = intent_override
+        if intent is None and getattr(self.config, 'enable_intent_classification', True):
+            intent_prompt = prompts.get('intent', '')
+            if intent_prompt:
+                intent = await classify_intent(
+                    giga=self.giga,
+                    query=query,
+                    prompt_text=intent_prompt,
+                    bot_config=self.config
+                )
+            else:
+                intent = "GENERAL"
+        else:
+            intent = intent or "GENERAL"
+
+        # ----- 5. Принудительная установка SURGICAL по ключевым словам -----
+        # Если пользователь явно просит заменить текст и есть загруженный файл
+        keywords = getattr(self.config, 'surgical_keywords', [])
+        if any(kw in query.lower() for kw in keywords) and last_file_path:
+            intent = "SURGICAL"
+            logger.info(f"Принудительный Intent: SURGICAL (есть файл и ключевые слова)")
+
+        # ----- 6. Обработка специализированных намерений -----
+        answer = None
+        context = None
+        ctx_for_critique = None
+        synthesis_template_for_critique = None
+
+        # --- 6.1. METRICS: извлечение числовых показателей ---
+        if intent == "METRICS":
+            context = await self.kb.find_relevant_info(
+                query, user_jid, room_jid,
+                top_k=getattr(self.config, 'rag_metrics_top_k', 30)
+            )
+            if not context:
+                answer = "Не найдено данных для извлечения метрик."
+            else:
+                metrics_prompt = prompts.get('metrics', '')
+                metrics = await extract_metrics(
+                    giga=self.giga,
+                    context=context,
+                    prompt_text=metrics_prompt,
+                    bot_config=self.config
+                )
+                if metrics:
+                    lines = [f"- {m.get('metric_name')}: {m.get('value')} {m.get('unit', '')}" for m in metrics[:10]]
+                    answer = "📊 **Извлечённые метрики:**\n" + "\n".join(lines)
+                else:
+                    answer = "Не удалось извлечь метрики."
+
+        # --- 6.2. SUMMARY: суммаризация документа ---
+        elif intent == "SUMMARY":
+            if not last_file_text:
+                answer = "Нет документа для суммаризации."
+            else:
+                summary_prompt = prompts.get('summary', '')
+                answer = await summarize_document(
+                    giga=self.giga,
+                    text=last_file_text,
+                    title="Ваш документ",
+                    prompt_text=summary_prompt,
+                    bot_config=self.config
+                )
+
+        # --- 6.3. CONTRADICTION: проверка противоречий ---
+        elif intent == "CONTRADICTION":
+            context = await self.kb.find_relevant_info(
+                query, user_jid, room_jid,
+                top_k=getattr(self.config, 'rag_contradiction_top_k', 10)
+            )
+            if not context:
+                answer = "Недостаточно данных для проверки противоречий."
+            else:
+                chunks = [c.strip() for c in context.split("\n\n") if c.strip()]
+                if len(chunks) < 2:
+                    answer = "Недостаточно фрагментов."
+                else:
+                    consistency_prompt = prompts.get('consistency', '')
+                    consistency = await check_consistency(
+                        giga=self.giga,
+                        chunks=chunks,
+                        query=query,
+                        prompt_text=consistency_prompt,
+                        bot_config=self.config
+                    )
+                    if "[CONFLICT]" in consistency:
+                        answer = f"⚠️ **Обнаружены противоречия:**\n{consistency}"
+                    else:
+                        answer = "✅ Противоречий не обнаружено."
+
+        # --- 6.4. TEMPLATE_FILL: заполнение шаблона документа ---
+        elif intent == "TEMPLATE_FILL":
+            template_text = last_file_text or ""
+            if not template_text and last_file_path and os.path.exists(last_file_path):
+                # Если текст не передан, извлекаем его из файла
+                result = await asyncio.to_thread(self.files.process_any_file, last_file_path)
+                if isinstance(result, tuple) and len(result) == 2:
+                    template_text = result[0]
+                else:
+                    template_text = str(result)
+            if not template_text:
+                answer = "❌ Нет шаблона документа для заполнения. Загрузите файл .docx."
+            else:
+                # Обрезаем шаблон до 5000 символов, чтобы не перегружать контекст
+                truncated_template = template_text[:5000]
+                search_query = f"{query}\n{truncated_template}"
+                context = await self.kb.find_relevant_info(search_query, user_jid, room_jid, top_k=15)
+                fill_prompt = prompts.get('generate_document', '')
+                if not fill_prompt:
+                    fill_prompt = "Заполни плейсхолдеры, используя базу знаний."
+                prompt = (
+                    f"Перед тобой шаблон документа. Заполни плейсхолдеры, используя ТОЛЬКО базу знаний.\n\n"
+                    f"[ТЕКСТ ШАБЛОНА]:\n{truncated_template}\n\n"
+                    f"[ДАННЫЕ ИЗ БАЗЫ ЗНАНИЙ]:\n{context}\n\n"
+                    f"Инструкция по заполнению:\n{fill_prompt}\n\n"
+                    f"Формат ответа: [SURGICAL_REPLACE]\nинструкция_из_шаблона ||| текст_из_БЗ\n[/SURGICAL_REPLACE]"
+                )
+                answer = await self.giga.chat(
+                    history=formatted_history,
+                    query=prompt,
+                    system_prompt=system_prompt,
+                    file_id=None,
+                    temperature=getattr(self.config, 'ai_temperature', 0.1)
+                )
+
+        # --- 6.5. SPELLCHECK: проверка орфографии ---
+        # Внимание: эта функция требует доступа к промпту spellcheck и локальному file_service
+        # Пока оставлена заглушкой, так как требует отдельного рефакторинга.
+        # В будущем check_spelling должен быть адаптирован для работы через сервер.
+        elif intent == "SPELLCHECK":
+            if room_jid is not None:
+                answer = "⚠️ Проверка орфографии доступна только в личном чате."
+            else:
+                if not last_file_text or not last_file_path:
+                    answer = "❌ Нет документа для проверки. Сначала отправьте файл .docx."
+                elif not last_file_path.lower().endswith('.docx'):
+                    answer = "❌ Проверка орфографии поддерживается только для файлов .docx."
+                else:
+                    # TODO: адаптировать check_spelling для работы через HTTP
+                    answer = "⚠️ Проверка орфографии требует доработки (передача промптов)."
+
+        # --- 6.6. GREETING: приветствие (без контекста) ---
+        elif intent == "GREETING":
+            answer = await self.giga.chat(
+                history=formatted_history,
+                query=query,
+                system_prompt=system_prompt,
+                file_id=None,
+                temperature=getattr(self.config, 'ai_temperature', 0.1)
+            )
+
+        # --- 6.7. SURGICAL: хирургическая замена в документе ---
+        elif intent == "SURGICAL":
+            if not last_file_path:
+                answer = "❌ Нет загруженного документа для замены."
+            else:
+                # Парсим запрос вида "замени X на Y"
+                import re
+                q = re.sub(r'\s+', ' ', query.lower()).strip()
+                m = re.search(r'замени\s+(.+?)\s+на\s+(.+)', q)
+                if not m:
+                    m = re.search(r'заменить\s+(.+?)\s+на\s+(.+)', q)
+                if not m:
+                    answer = "❌ Не удалось распознать, что на что заменять. Используйте формат: замени слово на слово"
+                else:
+                    old_word = m.group(1).strip()
+                    new_word = m.group(2).strip()
+                    # Морфологическая замена с помощью pymorphy3
+                    try:
+                        from mawo_pymorphy3 import create_analyzer
+                        morph = create_analyzer()
+                        # Генерируем все формы старого слова
+                        old_forms = set()
+                        parsed_old = morph.parse(old_word)[0]
+                        old_forms.add(old_word)
+                        old_forms.add(parsed_old.normal_form)
+                        cases = ['nomn', 'gent', 'datv', 'accs', 'ablt', 'loct']
+                        numbers = ['sing', 'plur']
+                        for number in numbers:
+                            for case in cases:
+                                inflected = parsed_old.inflect({case, number})
+                                if inflected:
+                                    old_forms.add(inflected.word)
+                        # Сопоставляем с формами нового слова
+                        replacements = {}
+                        for old_form in old_forms:
+                            parsed_old_form = morph.parse(old_form)[0]
+                            tags = set()
+                            if parsed_old_form.tag.case:
+                                tags.add(parsed_old_form.tag.case)
+                            if parsed_old_form.tag.number:
+                                tags.add(parsed_old_form.tag.number)
+                            if parsed_old_form.tag.gender:
+                                tags.add(parsed_old_form.tag.gender)
+                            parsed_new = morph.parse(new_word)[0]
+                            new_form = parsed_new.inflect(tags)
+                            replacements[old_form] = new_form.word if new_form else new_word
+                        # Добавляем исходные формы
+                        replacements[old_word] = new_word
+                        replacements[parsed_old.normal_form] = new_word
+
+                        # Выполняем замену (синхронно в потоке)
+                        new_path = await asyncio.to_thread(self.files.surgical_replace, last_file_path, replacements)
+                        if new_path:
+                            answer = f"Замена '{old_word}' → '{new_word}' выполнена. Файл сохранён: {new_path}"
+                            # Можно также вернуть файл клиенту, но пока только сообщаем
+                        else:
+                            answer = f"❌ Ошибка при замене '{old_word}' → '{new_word}'."
+                    except ImportError:
+                        answer = "❌ Библиотека mawo-pymorphy3 не установлена. Установите её для морфологической замены."
+
+        # ----- 7. Обычный RAG (для GENERAL, FACT, PROCEDURE, COMPARISON, CALCULATION) -----
+        if answer is None:
+            # 7.1. Расширение запроса (добавляем синонимы и аббревиатуры)
+            expanded = await expand_query(
+                giga=self.giga,
+                query=query,
+                prompt_text=prompts.get('expand', ''),
+                bot_config=self.config
+            )
+            search_query = expanded if expanded and expanded != query else query
+
+            # 7.2. Поиск релевантного контекста в базе знаний
+            context = await self.kb.find_relevant_info(
+                search_query, user_jid, room_jid,
+                top_k=getattr(self.config, 'rag_default_top_k', 30)
+            )
+            logger.info(f"Найден контекст длиной {len(context)} символов (room={room_jid})")
+
+            # 7.3. Переранжирование контекста (отбрасываем нерелевантные чанки)
+            rerank_min_length = getattr(self.config, 'rerank_min_length', 5000)
+            if intent != "FACT" and len(context) > rerank_min_length:
+                context = await rerank_context(
+                    bot=None,  # передаём None, так как используем кросс-энкодер
+                    query=query,
+                    context=context,
+                    prompt_text=None,
+                    bot_config=self.config
+                )
+
+            # 7.4. Синтез ответа с добавлением цепочки рассуждений (CoT) для расчётов и инструкций
+            synthesis_template = synthesis_template
+            if intent in ("CALCULATION", "PROCEDURE"):
+                cot_instruction = (
+                    "\n\nПожалуйста, покажи пошаговое решение перед итоговым ответом. "
+                    "Опиши каждый шаг вычислений или действий в логической последовательности. "
+                    "После всех шагов дай итоговый ответ."
+                )
+                synthesis_template += cot_instruction
+
+            # Подставляем контекст и запрос в шаблон
+            full_query = synthesis_template.format(context=context, query=query)
+            logger.debug(f"Полный запрос к GigaChat (первые 500 символов): {full_query[:500]}")
+
+            # 7.5. Генерация ответа через GigaChat
+            answer = await self.giga.chat(
+                history=formatted_history,
+                query=full_query,
+                system_prompt=system_prompt,
+                file_id=None,
+                temperature=getattr(self.config, 'ai_temperature', 0.1)
+            )
+
+            # Сохраняем контекст для самокритики
+            ctx_for_critique = context
+            synthesis_template_for_critique = synthesis_template
+
+        # ----- 8. Самокритика (если включена) -----
+        # Проверяем ответ на галлюцинации и соответствие контексту
+        if (getattr(self.config, 'enable_self_critique', False) and
+            intent not in ("METRICS", "SUMMARY", "CONTRADICTION") and
+            ctx_for_critique is not None and
+            answer is not None):
+            critique_prompt = prompts.get('critique', '')
+            if critique_prompt:
+                logger.debug("Запуск самокритики")
+                is_ok = await critique_answer(
+                    giga=self.giga,
+                    query=query,
+                    context=ctx_for_critique,
+                    answer=answer,
+                    prompt_text=critique_prompt,
+                    bot_config=self.config
+                )
+                if not is_ok:
+                    logger.warning("Ответ не прошёл самокритику, перегенерация")
+                    # Перегенерируем с тем же контекстом
+                    full_query_retry = synthesis_template_for_critique.format(
+                        context=ctx_for_critique,
+                        query=query
+                    )
+                    answer = await self.giga.chat(
+                        history=formatted_history,
+                        query=full_query_retry,
+                        system_prompt=system_prompt,
+                        file_id=None,
+                        temperature=getattr(self.config, 'ai_temperature', 0.1)
+                    )
+                    # Повторная проверка после перегенерации
+                    if not await critique_answer(
+                        giga=self.giga,
+                        query=query,
+                        context=ctx_for_critique,
+                        answer=answer,
+                        prompt_text=critique_prompt,
+                        bot_config=self.config
+                    ):
+                        answer = "⚠️ Извините, я не уверен в точности ответа. Проверьте данные."
+
+        # ----- 9. Сохранение истории диалога в БД -----
+        # Сохраняем запрос пользователя и ответ бота
+        await self.db.add_history(user_jid, "user", query, room_jid)
+        if answer:
+            await self.db.add_history(user_jid, "assistant", answer, room_jid)
+
+        # ----- 10. Извлечение источников из контекста для ответа -----
+        sources = []
+        if context:
+            for match in re.finditer(r'\[источник:\s*([^\]]+)\]', context):
+                sources.append(match.group(1))
+
+        # ----- 11. Формирование результата -----
+        return {
+            "answer": answer,
+            "intent": intent,
+            "context": context,
+            "sources": list(set(sources)),  # удаляем дубликаты
+            "confidence": None
+        }
+
+    async def index_document(
+        self,
+        file_name: str,
+        file_text: str,
+        user_jid: str,
+        room_jid: Optional[str],
+        is_global: bool = False,
+        title: Optional[str] = None,
+        metadata: Optional[Dict] = None,
+        file_hash: Optional[str] = None,
+        update_if_exists: bool = True
+    ) -> Dict[str, Any]:
+        """
+        Индексация документа в базу знаний.
+        Вызывается из HTTP-эндпоинта /rag/index.
+        """
+        doc_id, chunk_count = await self.kb.add_document(
+            file_name=file_name,
+            file_text=file_text,
+            user_jid=user_jid,
+            is_global=is_global,
+            title=title,
+            metadata=metadata,
+            room_jid=room_jid,
+            file_hash=file_hash,
+            update_if_exists=update_if_exists
+        )
+        return {"doc_id": doc_id, "chunk_count": chunk_count}