From fbfe178067449209be27c421490b917aa43fead8 Mon Sep 17 00:00:00 2001
From: Markov Andrey <anmarkov@pptrf.ru>
Date: Tue, 30 Jun 2026 16:15:00 +0000
Subject: [PATCH] Update 3 files

- /rag/intent_router.py
- /rag/query_processor.py
- /rag/history_manager.py
---
 rag/history_manager.py |  50 ++++++++++---------
 rag/intent_router.py   | 107 ++++++++++++++++++-----------------------
 rag/query_processor.py |  79 ++++++++++++++++--------------
 3 files changed, 118 insertions(+), 118 deletions(-)

diff --git a/rag/history_manager.py b/rag/history_manager.py
index 669b356..8e90c7d 100644
--- a/rag/history_manager.py
+++ b/rag/history_manager.py
@@ -11,13 +11,13 @@
 """
 
 import logging
-from typing import List, Dict, Optional, Any
+from typing import List, Dict, Optional, Any, Tuple
 
-from rag.services.postgres_service import PostgresService
-from rag.services.giga_client import GigaClient
-from rag.utils.text_utils import count_tokens
-from rag.functions.hierarchical_summarize import hierarchical_summarize
-from rag.config_models import AppConfig
+from core.services.postgres_service import PostgresService
+from core.services.giga_client import GigaClient
+from core.utils.text_utils import count_tokens
+from core.functions.hierarchical_summarize import hierarchical_summarize
+from core.config_models import AppConfig
 
 logger = logging.getLogger(__name__)
 
@@ -33,7 +33,7 @@ class HistoryManager:
         giga: GigaClient,
         config: AppConfig,
         default_prompts: Optional[Dict[str, str]] = None
-    ):
+    ) -> None:
         """
         Инициализация менеджера истории.
 
@@ -43,10 +43,14 @@ class HistoryManager:
             config: объект конфигурации (AppConfig)
             default_prompts: словарь промптов по умолчанию (нужен для hierarchical_summary)
         """
-        self.db = db
-        self.giga = giga
-        self.config = config
-        self.default_prompts = default_prompts or {}
+        self.db: PostgresService = db
+        self.giga: GigaClient = giga
+        self.config: AppConfig = config
+        self.default_prompts: Dict[str, str] = default_prompts or {}
+
+    # ------------------------------------------------------------------
+    # Получение и форматирование истории
+    # ------------------------------------------------------------------
 
     async def get_history(
         self,
@@ -65,7 +69,7 @@ class HistoryManager:
         Возвращает:
             List[Dict[str, str]]: список сообщений с ключами 'role' и 'content'
         """
-        raw_history = await self.db.get_history(user_jid, room_jid, limit=limit)
+        raw_history: List[Dict[str, Any]] = await self.db.get_history(user_jid, room_jid, limit=limit)
         # Преобразуем в формат, ожидаемый LLM
         return [
             {"role": rec['role'], "content": rec['content']}
@@ -95,8 +99,8 @@ class HistoryManager:
             return []
 
         # Подсчитываем общее количество токенов в истории
-        history_text = "\n".join([f"{msg['role']}: {msg['content']}" for msg in history])
-        total_tokens = count_tokens(history_text)
+        history_text: str = "\n".join([f"{msg['role']}: {msg['content']}" for msg in history])
+        total_tokens: int = count_tokens(history_text)
 
         if total_tokens <= max_tokens:
             logger.debug(f"История укладывается в {max_tokens} токенов (фактически {total_tokens})")
@@ -127,14 +131,14 @@ class HistoryManager:
                     )
 
         # Параметры сжатия из конфига
-        summarization_config = self.config.summarization
-        target_tokens = summarization_config.target_tokens_after_summary
-        chunk_size = summarization_config.chunk_size_tokens
-        max_depth = summarization_config.max_depth
-        temperature = self.config.ai.temperature
+        summarization_config: Dict[str, Any] = getattr(self.config, 'summarization', {})
+        target_tokens: int = summarization_config.get('target_tokens_after_summary', max_tokens)
+        chunk_size: int = summarization_config.get('chunk_size_tokens', 500)
+        max_depth: int = summarization_config.get('max_depth', 2)
+        temperature: float = getattr(self.config, 'ai_temperature', 0.1)
 
         try:
-            compressed_text = await hierarchical_summarize(
+            compressed_text: str = await hierarchical_summarize(
                 text=history_text,
                 giga=self.giga,
                 prompt_template=prompt_template,
@@ -153,10 +157,10 @@ class HistoryManager:
         except Exception as e:
             logger.error(f"Ошибка при сжатии истории: {e}, выполняем простую обрезку")
             # Fallback: простая обрезка по токенам
-            truncated_history = []
-            total = 0
+            truncated_history: List[Dict[str, str]] = []
+            total: int = 0
             for msg in reversed(history):
-                tokens = count_tokens(msg['content'])
+                tokens: int = count_tokens(msg['content'])
                 if total + tokens <= max_tokens:
                     truncated_history.append(msg)
                     total += tokens
diff --git a/rag/intent_router.py b/rag/intent_router.py
index c06c460..805a57a 100644
--- a/rag/intent_router.py
+++ b/rag/intent_router.py
@@ -12,7 +12,7 @@ import logging
 import re
 import asyncio
 import os
-from typing import Optional, Dict, List, Any
+from typing import Optional, Dict, List, Any, Union, Tuple
 
 from core.services.giga_client import GigaClient
 from core.services.kb_service import KBService
@@ -21,6 +21,7 @@ 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
+from core.config_models import AppConfig   # предположим, что путь такой
 
 logger = logging.getLogger(__name__)
 
@@ -35,24 +36,17 @@ class IntentRouter:
         giga: GigaClient,
         kb: KBService,
         files: FileService,
-        config,
+        config: AppConfig,
         default_prompts: Optional[Dict[str, str]] = None
-    ):
+    ) -> 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 {}
+        self.giga: GigaClient = giga
+        self.kb: KBService = kb
+        self.files: FileService = files
+        self.config: AppConfig = config
+        self.default_prompts: Dict[str, str] = default_prompts or {}
 
     # ------------------------------------------------------------------
     # Основной метод маршрутизации
@@ -133,23 +127,23 @@ class IntentRouter:
         prompts: Dict[str, str]
     ) -> Dict[str, Any]:
         """Обработка METRICS: извлечение KPI из базы знаний."""
-        context = await self.kb.find_relevant_info(
+        context: str = 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(
+        metrics_prompt: str = prompts.get('metrics', '')
+        metrics: List[Dict] = 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)
+            lines: List[str] = [f"- {m.get('metric_name')}: {m.get('value')} {m.get('unit', '')}" for m in metrics[:10]]
+            answer: str = "📊 **Извлечённые метрики:**\n" + "\n".join(lines)
         else:
             answer = "Не удалось извлечь метрики."
 
@@ -165,8 +159,8 @@ class IntentRouter:
         if not last_file_text:
             return {"answer": "Нет документа для суммаризации.", "context": "", "sources": []}
 
-        summary_prompt = prompts.get('summary', '')
-        answer = await summarize_document(
+        summary_prompt: str = prompts.get('summary', '')
+        answer: str = await summarize_document(
             giga=self.giga,
             text=last_file_text,
             title="Ваш документ",
@@ -183,19 +177,19 @@ class IntentRouter:
         prompts: Dict[str, str]
     ) -> Dict[str, Any]:
         """Обработка CONTRADICTION: проверка противоречий в базе знаний."""
-        context = await self.kb.find_relevant_info(
+        context: str = 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()]
+        chunks: List[str] = [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(
+        consistency_prompt: str = prompts.get('consistency', '')
+        consistency: str = await check_consistency(
             giga=self.giga,
             chunks=chunks,
             query=query,
@@ -203,7 +197,7 @@ class IntentRouter:
             bot_config=self.config
         )
         if "[CONFLICT]" in consistency:
-            answer = f"⚠️ **Обнаружены противоречия:**\n{consistency}"
+            answer: str = f"⚠️ **Обнаружены противоречия:**\n{consistency}"
         else:
             answer = "✅ Противоречий не обнаружено."
 
@@ -219,10 +213,9 @@ class IntentRouter:
         last_file_text: Optional[str]
     ) -> Dict[str, Any]:
         """Обработка TEMPLATE_FILL: заполнение шаблона документа."""
-        # Получаем текст шаблона
-        template_text = last_file_text or ""
+        template_text: str = 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)
+            result: Any = 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:
@@ -231,16 +224,15 @@ class IntentRouter:
         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)
+        truncated_template: str = template_text[:5000]
+        search_query: str = f"{query}\n{truncated_template}"
+        context: str = await self.kb.find_relevant_info(search_query, user_jid, room_jid, top_k=15)
 
-        fill_prompt = prompts.get('generate_document', '')
+        fill_prompt: str = prompts.get('generate_document', '')
         if not fill_prompt:
             fill_prompt = "Заполни плейсхолдеры, используя базу знаний."
 
-        full_prompt = (
+        full_prompt: str = (
             f"Перед тобой шаблон документа. Заполни плейсхолдеры, используя ТОЛЬКО базу знаний.\n\n"
             f"[ТЕКСТ ШАБЛОНА]:\n{truncated_template}\n\n"
             f"[ДАННЫЕ ИЗ БАЗЫ ЗНАНИЙ]:\n{context}\n\n"
@@ -248,7 +240,7 @@ class IntentRouter:
             f"Формат ответа: [SURGICAL_REPLACE]\nинструкция_из_шаблона ||| текст_из_БЗ\n[/SURGICAL_REPLACE]"
         )
 
-        answer = await self.giga.chat(
+        answer: str = await self.giga.chat(
             history=[],
             query=full_prompt,
             system_prompt=None,
@@ -266,42 +258,39 @@ class IntentRouter:
         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)
+        q: str = re.sub(r'\s+', ' ', query.lower()).strip()
+        m: Optional[re.Match] = 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()
+        old_word: str = m.group(1).strip()
+        new_word: str = m.group(2).strip()
 
         try:
             from mawo_pymorphy3 import create_analyzer
             morph = create_analyzer()
-            # Генерируем все формы старого слова
-            old_forms = set()
+            old_forms: set = 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']
+            cases: List[str] = ['nomn', 'gent', 'datv', 'accs', 'ablt', 'loct']
+            numbers: List[str] = ['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):
+            parse_cache: Dict[str, Any] = {}
+            def get_new_form(new_word: str, old_form_text: str) -> str:
                 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()
+                tags: set = set()
                 if parsed_old.tag.case:
                     tags.add(parsed_old.tag.case)
                 if parsed_old.tag.number:
@@ -312,14 +301,13 @@ class IntentRouter:
                 inflected = parsed_new.inflect(tags)
                 return inflected.word if inflected else new_word
 
-            replacements = {}
+            replacements: Dict[str, str] = {}
             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)
+            new_path: Optional[str] = await asyncio.to_thread(self.files.surgical_replace, last_file_path, replacements)
             if new_path:
-                answer = f"Замена '{old_word}' → '{new_word}' выполнена. Файл сохранён: {new_path}"
+                answer: str = f"Замена '{old_word}' → '{new_word}' выполнена. Файл сохранён: {new_path}"
             else:
                 answer = f"❌ Ошибка при замене '{old_word}' → '{new_word}'."
             return {"answer": answer, "context": "", "sources": []}
@@ -348,7 +336,7 @@ class IntentRouter:
         if not last_file_path.lower().endswith('.docx'):
             return {"answer": "❌ Проверка орфографии поддерживается только для файлов .docx.", "context": "", "sources": []}
 
-        spellcheck_prompt = prompts.get('spellcheck', '')
+        spellcheck_prompt: str = prompts.get('spellcheck', '')
         if not spellcheck_prompt:
             return {"answer": "❌ Промпт проверки орфографии не загружен.", "context": "", "sources": []}
 
@@ -360,7 +348,7 @@ class IntentRouter:
                 config=self.config
             )
             if changes:
-                answer = "📝 **Исправления:**\n" + "\n".join(changes)
+                answer: str = "📝 **Исправления:**\n" + "\n".join(changes)
             else:
                 answer = "✅ Ошибок не найдено."
             return {"answer": answer, "context": "", "sources": []}
@@ -376,14 +364,13 @@ class IntentRouter:
         prompts: Dict[str, str]
     ) -> Dict[str, Any]:
         """Обработка GREETING: простой ответ через GigaChat без поиска."""
-        # Используем synthesis промпт, если есть, или просто передаём запрос
-        synthesis_template = prompts.get('synthesis', '')
+        synthesis_template: str = prompts.get('synthesis', '')
         if synthesis_template:
-            full_query = synthesis_template.format(context="", query=query)
+            full_query: str = synthesis_template.format(context="", query=query)
         else:
             full_query = query
 
-        answer = await self.giga.chat(
+        answer: str = await self.giga.chat(
             history=history or [],
             query=full_query,
             system_prompt=system_prompt,
diff --git a/rag/query_processor.py b/rag/query_processor.py
index fd5495c..73368b6 100644
--- a/rag/query_processor.py
+++ b/rag/query_processor.py
@@ -9,14 +9,14 @@ import logging
 import re
 from typing import Optional, Dict, List, Any
 
-from .services.giga_client import GigaClient
-from .services.kb_service import KBService
-from .functions.expand_query import expand_query
-from .functions.rerank_context import rerank_context
-from .functions.critique_answer import critique_answer
-from .utils.text_utils import count_tokens
-from .prompt_builder import PromptBuilder
-from .config_models import AppConfig
+from core.services.giga_client import GigaClient
+from core.services.kb_service import KBService
+from core.functions.expand_query import expand_query
+from core.functions.rerank_context import rerank_context
+from core.functions.critique_answer import critique_answer
+from core.utils.text_utils import count_tokens
+from core.prompt_builder import PromptBuilder
+from core.config_models import AppConfig
 
 logger = logging.getLogger(__name__)
 
@@ -32,12 +32,12 @@ class QueryProcessor:
         kb: KBService,
         config: AppConfig,
         default_prompts: Optional[Dict[str, str]] = None
-    ):
-        self.giga = giga
-        self.kb = kb
-        self.config = config
-        self.default_prompts = default_prompts or {}
-        self.prompt_builder = PromptBuilder(config)
+    ) -> None:
+        self.giga: GigaClient = giga
+        self.kb: KBService = kb
+        self.config: AppConfig = config
+        self.default_prompts: Dict[str, str] = default_prompts or {}
+        self.prompt_builder: PromptBuilder = PromptBuilder(config)
 
     async def process(
         self,
@@ -52,18 +52,31 @@ class QueryProcessor:
     ) -> Dict[str, Any]:
         """
         Выполняет полный RAG-пайплайн для обычного запроса.
+
+        Аргументы:
+            query: текст запроса пользователя
+            user_jid: JID пользователя
+            room_jid: JID комнаты (None для личного чата)
+            prompts: словарь промптов (expand, synthesis, critique, ...)
+            intent: код намерения (для выбора стратегии)
+            history: история диалога (уже сжатая, если нужно)
+            system_prompt: системный промпт
+            available_tokens_for_context: сколько токенов доступно для контекста
+
+        Возвращает:
+            Словарь с ключами 'answer', 'context', 'sources', 'confidence'
         """
         # 1. Расширение запроса
-        expanded = await expand_query(
+        expanded: str = 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
+        search_query: str = expanded if expanded and expanded != query else query
 
         # 2. Поиск релевантного контекста в базе знаний
-        context = await self.kb.find_relevant_info(
+        context: str = await self.kb.find_relevant_info(
             search_query, user_jid, room_jid,
             top_k=getattr(self.config, 'rag_default_top_k', 30)
         )
@@ -71,20 +84,20 @@ class QueryProcessor:
 
         # 3. Обрезка контекста по токенам
         if context:
-            context_tokens = count_tokens(context)
+            context_tokens: int = count_tokens(context)
             if context_tokens > available_tokens_for_context:
                 logger.warning(
                     f"Контекст слишком длинный ({context_tokens} токенов), "
                     f"обрезаем до {available_tokens_for_context}"
                 )
-                max_context_chars = int(available_tokens_for_context * 3.5)
+                max_context_chars: int = int(available_tokens_for_context * 3.5)
                 if max_context_chars > 0:
                     context = context[:max_context_chars]
                 else:
                     context = ""
 
         # 4. Переранжирование контекста (если включено и контекст достаточно длинный)
-        rerank_min_length = getattr(self.config, 'rerank_min_length', 5000)
+        rerank_min_length: int = getattr(self.config, 'rerank_min_length', 5000)
         if intent != "FACT" and len(context) > rerank_min_length:
             context = await rerank_context(
                 bot=None,
@@ -94,8 +107,8 @@ class QueryProcessor:
                 bot_config=self.config
             )
 
-        # 5. Формирование промта с помощью PromptBuilder (динамические few-shot, сэндвич)
-        extra_instructions = ""
+        # 5. Формирование промта с помощью PromptBuilder
+        extra_instructions: str = ""
         if intent in ("CALCULATION", "PROCEDURE"):
             extra_instructions = (
                 "Пожалуйста, покажи пошаговое решение перед итоговым ответом. "
@@ -103,13 +116,11 @@ class QueryProcessor:
                 "После всех шагов дай итоговый ответ."
             )
 
-        # Если synthesis_template задан, добавим его в extra_instructions (но PromptBuilder уже использует стандартные инструкции)
-        synthesis_template = prompts.get('synthesis', '')
+        synthesis_template: str = prompts.get('synthesis', '')
         if synthesis_template and not extra_instructions:
-            # Если есть кастомный шаблон, используем его как дополнительную инструкцию
             extra_instructions = synthesis_template.format(context=context, query=query) if '{context}' in synthesis_template else synthesis_template
 
-        prompt = self.prompt_builder.build_prompt(
+        prompt: str = self.prompt_builder.build_prompt(
             query=query,
             intent=intent,
             context=context,
@@ -119,21 +130,21 @@ class QueryProcessor:
         )
         logger.debug(f"Сформированный промт (первые 500 символов): {prompt[:500]}")
 
-        # 6. Генерация ответа (без отдельной передачи истории, она уже в промте)
-        answer = await self.giga.chat(
-            history=[],  # история уже в промте
+        # 6. Генерация ответа
+        answer: str = await self.giga.chat(
+            history=[],
             query=prompt,
-            system_prompt=None,  # системный промпт тоже в промте
+            system_prompt=None,
             file_id=None,
             temperature=getattr(self.config, 'ai_temperature', 0.1)
         )
 
         # 7. Самокритика (если включена)
         if getattr(self.config, 'enable_self_critique', False) and context:
-            critique_prompt = prompts.get('critique', '')
+            critique_prompt: str = prompts.get('critique', '')
             if critique_prompt:
                 logger.debug("Запуск самокритики")
-                is_ok = await critique_answer(
+                is_ok: bool = await critique_answer(
                     giga=self.giga,
                     query=query,
                     context=context,
@@ -143,7 +154,6 @@ class QueryProcessor:
                 )
                 if not is_ok:
                     logger.warning("Ответ не прошёл самокритику, перегенерация")
-                    # Перегенерируем с тем же промтом
                     answer = await self.giga.chat(
                         history=[],
                         query=prompt,
@@ -151,7 +161,6 @@ class QueryProcessor:
                         file_id=None,
                         temperature=getattr(self.config, 'ai_temperature', 0.1)
                     )
-                    # Повторная проверка после перегенерации
                     if not await critique_answer(
                         giga=self.giga,
                         query=query,
@@ -163,7 +172,7 @@ class QueryProcessor:
                         answer = "⚠️ Извините, я не уверен в точности ответа. Проверьте данные."
 
         # 8. Извлечение источников
-        sources = []
+        sources: List[str] = []
         if context:
             for match in re.finditer(r'\[источник:\s*([^\]]+)\]', context):
                 sources.append(match.group(1))