Add new file

rag/intent_router.py

@@ -0,0 +1,393 @@
+# -*- coding: utf-8 -*-
+"""
+Маршрутизатор намерений.
+Обрабатывает специализированные запросы (METRICS, SUMMARY, CONTRADICTION,
+TEMPLATE_FILL, SURGICAL, GREETING, SPELLCHECK) и возвращает готовый ответ.
+
+Если намерение не распознано как специализированное, возвращает None,
+и тогда управление передаётся обычному RAG-пайплайну.
+"""
+
+import logging
+import re
+import asyncio
+import os
+from typing import Optional, Dict, List, Any
+
+from core.services.giga_client import GigaClient
+from core.services.kb_service import KBService
+from core.services.file_service import FileService
+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.check_spelling import check_spelling
+
+logger = logging.getLogger(__name__)
+
+
+class IntentRouter:
+    """
+    Маршрутизирует запросы по намерениям и выполняет специализированную обработку.
+    """
+
+    def __init__(
+        self,
+        giga: GigaClient,
+        kb: KBService,
+        files: FileService,
+        config,
+        default_prompts: Optional[Dict[str, str]] = None
+    ):
+        """
+        Инициализация маршрутизатора.
+
+        Аргументы:
+            giga: клиент GigaChat
+            kb: сервис базы знаний
+            files: сервис файлов
+            config: объект конфигурации
+            default_prompts: словарь промптов по умолчанию
+        """
+        self.giga = giga
+        self.kb = kb
+        self.files = files
+        self.config = config
+        self.default_prompts = default_prompts or {}
+
+    # ------------------------------------------------------------------
+    # Основной метод маршрутизации
+    # ------------------------------------------------------------------
+
+    async def route(
+        self,
+        intent: str,
+        query: str,
+        user_jid: str,
+        room_jid: Optional[str],
+        prompts: Dict[str, str],
+        last_file_path: Optional[str] = None,
+        last_file_text: Optional[str] = None,
+        history: Optional[List[Dict[str, str]]] = None,
+        system_prompt: Optional[str] = None,
+    ) -> Optional[Dict[str, Any]]:
+        """
+        Маршрутизирует запрос по намерению.
+
+        Аргументы:
+            intent: код намерения (METRICS, SUMMARY, ...)
+            query: текст запроса пользователя
+            user_jid: JID пользователя
+            room_jid: JID комнаты (None для личного чата)
+            prompts: словарь промптов (intent, metrics, summary, ...)
+            last_file_path: путь к последнему загруженному файлу
+            last_file_text: текст последнего загруженного файла
+            history: история диалога (для GREETING)
+            system_prompt: системный промпт (для GREETING)
+
+        Возвращает:
+            Словарь с ключами 'answer', 'context', 'sources' или None, если намерение не обработано.
+        """
+        # Обработка METRICS
+        if intent == "METRICS":
+            return await self._handle_metrics(query, user_jid, room_jid, prompts)
+
+        # Обработка SUMMARY
+        if intent == "SUMMARY":
+            return await self._handle_summary(query, last_file_text, prompts)
+
+        # Обработка CONTRADICTION
+        if intent == "CONTRADICTION":
+            return await self._handle_contradiction(query, user_jid, room_jid, prompts)
+
+        # Обработка TEMPLATE_FILL
+        if intent == "TEMPLATE_FILL":
+            return await self._handle_template_fill(
+                query, user_jid, room_jid, prompts,
+                last_file_path, last_file_text
+            )
+
+        # Обработка SURGICAL
+        if intent == "SURGICAL":
+            return await self._handle_surgical(query, last_file_path)
+
+        # Обработка SPELLCHECK
+        if intent == "SPELLCHECK":
+            return await self._handle_spellcheck(query, room_jid, last_file_text, last_file_path, prompts)
+
+        # Обработка GREETING
+        if intent == "GREETING":
+            return await self._handle_greeting(query, history, system_prompt, prompts)
+
+        # Намерение не обработано — возвращаем None
+        return None
+
+    # ------------------------------------------------------------------
+    # Обработчики конкретных намерений
+    # ------------------------------------------------------------------
+
+    async def _handle_metrics(
+        self,
+        query: str,
+        user_jid: str,
+        room_jid: Optional[str],
+        prompts: Dict[str, str]
+    ) -> Dict[str, Any]:
+        """Обработка METRICS: извлечение KPI из базы знаний."""
+        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:
+            return {"answer": "Не найдено данных для извлечения метрик.", "context": "", "sources": []}
+
+        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 = "Не удалось извлечь метрики."
+
+        return {"answer": answer, "context": context, "sources": []}
+
+    async def _handle_summary(
+        self,
+        query: str,
+        last_file_text: Optional[str],
+        prompts: Dict[str, str]
+    ) -> Dict[str, Any]:
+        """Обработка SUMMARY: суммаризация последнего документа."""
+        if not last_file_text:
+            return {"answer": "Нет документа для суммаризации.", "context": "", "sources": []}
+
+        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
+        )
+        return {"answer": answer, "context": "", "sources": []}
+
+    async def _handle_contradiction(
+        self,
+        query: str,
+        user_jid: str,
+        room_jid: Optional[str],
+        prompts: Dict[str, str]
+    ) -> Dict[str, Any]:
+        """Обработка 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:
+            return {"answer": "Недостаточно данных для проверки противоречий.", "context": "", "sources": []}
+
+        chunks = [c.strip() for c in context.split("\n\n") if c.strip()]
+        if len(chunks) < 2:
+            return {"answer": "Недостаточно фрагментов для проверки противоречий.", "context": context, "sources": []}
+
+        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 = "✅ Противоречий не обнаружено."
+
+        return {"answer": answer, "context": context, "sources": []}
+
+    async def _handle_template_fill(
+        self,
+        query: str,
+        user_jid: str,
+        room_jid: Optional[str],
+        prompts: Dict[str, str],
+        last_file_path: Optional[str],
+        last_file_text: Optional[str]
+    ) -> Dict[str, Any]:
+        """Обработка 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:
+            return {"answer": "❌ Нет шаблона документа для заполнения. Загрузите файл .docx.", "context": "", "sources": []}
+
+        # Обрезаем шаблон для экономии токенов
+        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 = "Заполни плейсхолдеры, используя базу знаний."
+
+        full_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=[],
+            query=full_prompt,
+            system_prompt=None,
+            file_id=None,
+            temperature=getattr(self.config, 'ai_temperature', 0.1)
+        )
+        return {"answer": answer, "context": context, "sources": []}
+
+    async def _handle_surgical(
+        self,
+        query: str,
+        last_file_path: Optional[str]
+    ) -> Dict[str, Any]:
+        """Обработка SURGICAL: хирургическая замена в DOCX."""
+        if not last_file_path:
+            return {"answer": "❌ Нет загруженного документа для замены.", "context": "", "sources": []}
+
+        # Парсим запрос: "замени слово на слово"
+        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:
+            return {"answer": "❌ Не удалось распознать, что на что заменять. Используйте формат: замени слово на слово", "context": "", "sources": []}
+
+        old_word = m.group(1).strip()
+        new_word = m.group(2).strip()
+
+        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)
+
+            # Кэш для разбора слов
+            parse_cache = {}
+            def get_new_form(new_word, old_form_text):
+                if old_form_text in parse_cache:
+                    parsed_old = parse_cache[old_form_text]
+                else:
+                    parsed_old = morph.parse(old_form_text)[0]
+                    parse_cache[old_form_text] = parsed_old
+                tags = set()
+                if parsed_old.tag.case:
+                    tags.add(parsed_old.tag.case)
+                if parsed_old.tag.number:
+                    tags.add(parsed_old.tag.number)
+                if parsed_old.tag.gender:
+                    tags.add(parsed_old.tag.gender)
+                parsed_new = morph.parse(new_word)[0]
+                inflected = parsed_new.inflect(tags)
+                return inflected.word if inflected else new_word
+
+            replacements = {}
+            for old_form in old_forms:
+                replacements[old_form] = get_new_form(new_word, old_form)
+
+            # Выполняем замену
+            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}'."
+            return {"answer": answer, "context": "", "sources": []}
+
+        except ImportError:
+            return {"answer": "❌ Библиотека mawo-pymorphy3 не установлена. Установите её для морфологической замены.", "context": "", "sources": []}
+        except Exception as e:
+            logger.error(f"Ошибка хирургической замены: {e}", exc_info=True)
+            return {"answer": f"❌ Ошибка при замене: {e}", "context": "", "sources": []}
+
+    async def _handle_spellcheck(
+        self,
+        query: str,
+        room_jid: Optional[str],
+        last_file_text: Optional[str],
+        last_file_path: Optional[str],
+        prompts: Dict[str, str]
+    ) -> Dict[str, Any]:
+        """Обработка SPELLCHECK: проверка орфографии в DOCX."""
+        if room_jid is not None:
+            return {"answer": "⚠️ Проверка орфографии доступна только в личном чате.", "context": "", "sources": []}
+
+        if not last_file_text or not last_file_path:
+            return {"answer": "❌ Нет документа для проверки. Сначала отправьте файл .docx.", "context": "", "sources": []}
+
+        if not last_file_path.lower().endswith('.docx'):
+            return {"answer": "❌ Проверка орфографии поддерживается только для файлов .docx.", "context": "", "sources": []}
+
+        spellcheck_prompt = prompts.get('spellcheck', '')
+        if not spellcheck_prompt:
+            return {"answer": "❌ Промпт проверки орфографии не загружен.", "context": "", "sources": []}
+
+        try:
+            replacements, changes = await check_spelling(
+                giga=self.giga,
+                original_text=last_file_text,
+                prompt_text=spellcheck_prompt,
+                config=self.config
+            )
+            if changes:
+                answer = "📝 **Исправления:**\n" + "\n".join(changes)
+            else:
+                answer = "✅ Ошибок не найдено."
+            return {"answer": answer, "context": "", "sources": []}
+        except Exception as e:
+            logger.error(f"Ошибка проверки орфографии: {e}", exc_info=True)
+            return {"answer": f"❌ Ошибка проверки орфографии: {e}", "context": "", "sources": []}
+
+    async def _handle_greeting(
+        self,
+        query: str,
+        history: Optional[List[Dict[str, str]]],
+        system_prompt: Optional[str],
+        prompts: Dict[str, str]
+    ) -> Dict[str, Any]:
+        """Обработка GREETING: простой ответ через GigaChat без поиска."""
+        # Используем synthesis промпт, если есть, или просто передаём запрос
+        synthesis_template = prompts.get('synthesis', '')
+        if synthesis_template:
+            full_query = synthesis_template.format(context="", query=query)
+        else:
+            full_query = query
+
+        answer = await self.giga.chat(
+            history=history or [],
+            query=full_query,
+            system_prompt=system_prompt,
+            file_id=None,
+            temperature=getattr(self.config, 'ai_temperature', 0.1)
+        )
+        return {"answer": answer, "context": "", "sources": []}