Редактировать prompt_builder.py

rag/prompt_builder.py

@@ -2,6 +2,7 @@
 """
 Построитель промтов для RAG-системы.
 Реализует динамический подбор few-shot примеров и структуру «сэндвич».
+Добавлена полная типизация для всех локальных переменных.
 """
 
 import json
@@ -21,20 +22,20 @@ class PromptBuilder:
     и структуру «сэндвич» (важное в начале и в конце).
     """
 
-    def __init__(self, config: AppConfig):
+    def __init__(self, config: AppConfig) -> None:
         """
         Инициализация.
 
         Аргументы:
             config: объект AppConfig (содержит пути и настройки few-shot).
         """
-        self.config = config
+        self.config: AppConfig = config
-        self.max_examples = config.fewshot.max_examples
+        self.max_examples: int = config.fewshot.max_examples
-        examples_path = Path(config.fewshot.examples_file)
+        examples_path: Path = Path(config.fewshot.examples_file)
         if not examples_path.is_absolute():
             examples_path = config.profile_dir / examples_path
-        self.examples_file = examples_path
+        self.examples_file: Path = examples_path
-        self._examples = None  # кэш загруженных примеров
+        self._examples: Optional[List[Dict[str, Any]]] = None  # кэш загруженных примеров
 
     def load_examples(self) -> List[Dict[str, Any]]:
         """
@@ -46,7 +47,7 @@ class PromptBuilder:
         if self.examples_file and self.examples_file.exists():
             try:
                 with open(self.examples_file, 'r', encoding='utf-8') as f:
-                    data = json.load(f)
+                    data: Any = json.load(f)
                     if isinstance(data, list):
                         self._examples = data
                         logger.info(f"Загружено {len(self._examples)} few-shot примеров из {self.examples_file}")
@@ -57,7 +58,7 @@ class PromptBuilder:
                     logger.error(f"Ошибка загрузки few-shot примеров: {e}")
 
         # Статические примеры (запасные)
-        static_examples = [
+        static_examples: List[Dict[str, Any]] = [
             {
                 "keywords": ["OEE", "эффективность", "оборудование"],
                 "intents": ["PROCEDURE", "CALCULATION"],
@@ -98,33 +99,33 @@ class PromptBuilder:
         Возвращает:
             Список словарей (примеры), не более max_examples.
         """
-        all_examples = self.load_examples()
+        all_examples: List[Dict[str, Any]] = self.load_examples()
         if not all_examples:
             return []
 
         # Нормализуем запрос
-        query_lower = query.lower()
+        query_lower: str = query.lower()
-        words = set(re.findall(r'\w+', query_lower))
+        words: set = set(re.findall(r'\w+', query_lower))
 
         # Вычисляем релевантность для каждого примера
-        scored = []
+        scored: List[tuple] = []
         for ex in all_examples:
             # Проверка по намерению (если указано)
-            intent_match = (intent in ex.get('intents', [])) if ex.get('intents') else False
+            intent_match: bool = (intent in ex.get('intents', [])) if ex.get('intents') else False
 
             # Проверка по ключевым словам
-            keywords = ex.get('keywords', [])
+            keywords: List[str] = ex.get('keywords', [])
-            keyword_match = any(kw.lower() in query_lower for kw in keywords) if keywords else False
+            keyword_match: bool = any(kw.lower() in query_lower for kw in keywords) if keywords else False
 
             # Количество совпадающих слов из запроса (дополнительно)
             if keywords:
-                kw_set = set(kw.lower() for kw in keywords)
+                kw_set: set = set(kw.lower() for kw in keywords)
-                word_overlap = len(words.intersection(kw_set))
+                word_overlap: int = len(words.intersection(kw_set))
             else:
                 word_overlap = 0
 
             # Суммарный балл
-            score = 0
+            score: float = 0
             if intent_match:
                 score += 2
             if keyword_match:
@@ -135,7 +136,7 @@ class PromptBuilder:
 
         # Сортируем по убыванию балла и берём top N
         scored.sort(key=lambda x: x[0], reverse=True)
-        selected = [ex for _, ex in scored[:self.max_examples] if ex.get('question') and ex.get('answer')]
+        selected: List[Dict[str, Any]] = [ex for _, ex in scored[:self.max_examples] if ex.get('question') and ex.get('answer')]
         return selected
 
     def format_example(self, example: Dict[str, Any]) -> str:
@@ -167,7 +168,7 @@ class PromptBuilder:
         Возвращает:
             str: готовый промт для отправки в LLM.
         """
-        parts = []
+        parts: List[str] = []
 
         # 1. Системный промпт (в самом начале)
         if system_prompt:
@@ -176,22 +177,22 @@ class PromptBuilder:
         # 2. История диалога (если есть)
         if history:
             for msg in history:
-                role = msg.get('role', 'user')
+                role: str = msg.get('role', 'user')
-                content = msg.get('content', '')
+                content: str = msg.get('content', '')
                 if role == 'system':
                     continue  # системное уже добавлено
                 parts.append(f"<|im_start|>{role}\n{content}\n<|im_end|>")
 
         # 3. Ключевые инструкции (краткое напоминание о задаче)
-        instructions = "Ты — профессиональный методолог-аналитик. Отвечай строго на основе предоставленного контекста. Используй профессиональную терминологию. Если ответа нет в контексте, сообщи об этом."
+        instructions: str = "Ты — профессиональный методолог-аналитик. Отвечай строго на основе предоставленного контекста. Используй профессиональную терминологию. Если ответа нет в контексте, сообщи об этом."
         if extra_instructions:
             instructions += f"\n{extra_instructions}"
         parts.append(f"<|im_start|>system\n{instructions}\n<|im_end|>")
 
         # 4. Few-shot примеры (после инструкций, до основного контекста)
-        examples = self.select_examples(query, intent)
+        examples: List[Dict[str, Any]] = self.select_examples(query, intent)
         if examples:
-            example_text = "\n\n".join(self.format_example(ex) for ex in examples)
+            example_text: str = "\n\n".join(self.format_example(ex) for ex in examples)
             parts.append(f"<|im_start|>system\nВот несколько примеров ответов на похожие вопросы:\n\n{example_text}\n<|im_end|>")
 
         # 5. Основной контекст (источники)
@@ -199,7 +200,12 @@ class PromptBuilder:
             parts.append(f"<|im_start|>system\n[ДАННЫЕ ИЗ БАЗЫ ЗНАНИЙ]\n{context}\n<|im_end|>")
 
         # 6. Повторение ключевого вопроса и требований (в конце – важен для фокуса)
-        final_prompt = f"{query}\n\nПожалуйста, дай ответ на основе предоставленного контекста. Укажи источники в формате [Источник: название_файла]. Не используй LaTeX, формулы пиши текстом. Будь краток, но точен."
+        final_prompt: str = (
+            f"{query}\n\n"
+            "Пожалуйста, дай ответ на основе предоставленного контекста. "
+            "Укажи источники в формате [Источник: название_файла]. "
+            "Не используй LaTeX, формулы пиши текстом. Будь краток, но точен."
+        )
 
         if intent in ("CALCULATION", "PROCEDURE"):
             final_prompt += "\n\nПокажи пошаговое решение перед итоговым ответом."