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

rag/prompt_builder.py

@@ -2,12 +2,12 @@
 """
 Построитель промтов для RAG-системы.
 Реализует динамический подбор few-shot примеров и структуру «сэндвич».
-Добавлена полная типизация для всех локальных переменных.
+Добавлена санитизация пользовательских данных для защиты от инъекций ChatML.
 """
 
 import json
-import logging
 import re
+import logging
 from pathlib import Path
 from typing import List, Dict, Any, Optional
 
@@ -35,7 +35,22 @@ class PromptBuilder:
         if not examples_path.is_absolute():
             examples_path = config.profile_dir / examples_path
         self.examples_file: Path = examples_path
-        self._examples: Optional[List[Dict[str, Any]]] = None  # кэш загруженных примеров
+        self._examples: Optional[List[Dict[str, Any]]] = None
+
+    def _sanitize(self, text: str) -> str:
+        """
+        Удаляет или экранирует специальные токены ChatML,
+        чтобы предотвратить инъекцию промтов.
+
+        Аргументы:
+            text: исходный текст.
+
+        Возвращает:
+            str: очищенный текст.
+        """
+        if not text:
+            return ""
+        return text.replace('<|im_start|>', '').replace('<|im_end|>', '')
 
     def load_examples(self) -> List[Dict[str, Any]]:
         """
@@ -143,7 +158,9 @@ class PromptBuilder:
         """
         Форматирует один пример для вставки в промт.
         """
-        return f"Вопрос: {example['question']}\nОтвет: {example['answer']}"
+        question = self._sanitize(example['question'])
+        answer = self._sanitize(example['answer'])
+        return f"Вопрос: {question}\nОтвет: {answer}"
 
     def build_prompt(
         self,
@@ -168,26 +185,30 @@ class PromptBuilder:
         Возвращает:
             str: готовый промт для отправки в LLM.
         """
+        # Санитайзинг всех пользовательских данных
+        query = self._sanitize(query)
+        context = self._sanitize(context)
+
         parts: List[str] = []
 
         # 1. Системный промпт (в самом начале)
         if system_prompt:
-            parts.append(f"<|im_start|>system\n{system_prompt}\n<|im_end|>")
+            parts.append(f"<|im_start|>system\n{self._sanitize(system_prompt)}\n<|im_end|>")
 
         # 2. История диалога (если есть)
         if history:
             for msg in history:
                 role: str = msg.get('role', 'user')
-                content: str = msg.get('content', '')
+                content: str = self._sanitize(msg.get('content', ''))
                 if role == 'system':
-                    continue  # системное уже добавлено
+                    continue
                 parts.append(f"<|im_start|>{role}\n{content}\n<|im_end|>")
 
         # 3. Ключевые инструкции (краткое напоминание о задаче)
         instructions: str = "Ты — профессиональный методолог-аналитик. Отвечай строго на основе предоставленного контекста. Используй профессиональную терминологию. Если ответа нет в контексте, сообщи об этом."
         if extra_instructions:
             instructions += f"\n{extra_instructions}"
-        parts.append(f"<|im_start|>system\n{instructions}\n<|im_end|>")
+        parts.append(f"<|im_start|>system\n{self._sanitize(instructions)}\n<|im_end|>")
 
         # 4. Few-shot примеры (после инструкций, до основного контекста)
         examples: List[Dict[str, Any]] = self.select_examples(query, intent)
@@ -210,7 +231,7 @@ class PromptBuilder:
         if intent in ("CALCULATION", "PROCEDURE"):
             final_prompt += "\n\nПокажи пошаговое решение перед итоговым ответом."
 
-        parts.append(f"<|im_start|>user\n{final_prompt}\n<|im_end|>")
+        parts.append(f"<|im_start|>user\n{self._sanitize(final_prompt)}\n<|im_end|>")
 
         # 7. Переход к ответу (ассистент)
         parts.append("<|im_start|>assistant")