diff --git a/rag/intent_router.py b/rag/intent_router.py new file mode 100644 index 0000000..c06c460 --- /dev/null +++ b/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": []} \ No newline at end of file