diff --git a/rag/prompt_builder.py b/rag/prompt_builder.py index 3425c4a..25becc5 100644 --- a/rag/prompt_builder.py +++ b/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]]: """ @@ -54,8 +69,8 @@ class PromptBuilder: return self._examples else: logger.warning(f"Файл {self.examples_file} должен содержать список JSON-объектов") - except Exception as e: - logger.error(f"Ошибка загрузки few-shot примеров: {e}") + except Exception as e: + logger.error(f"Ошибка загрузки few-shot примеров: {e}") # Статические примеры (запасные) static_examples: 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")