prompt_manager.py
1 # Copyright (c) 2024-2026 Tencent Zhuque Lab. All rights reserved. 2 # 3 # Licensed under the Apache License, Version 2.0 (the "License"); 4 # you may not use this file except in compliance with the License. 5 # You may obtain a copy of the License at 6 # 7 # http://www.apache.org/licenses/LICENSE-2.0 8 # 9 # Unless required by applicable law or agreed to in writing, software 10 # distributed under the License is distributed on an "AS IS" BASIS, 11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 # See the License for the specific language governing permissions and 13 # limitations under the License. 14 # 15 # Requirement: Any integration or derivative work must explicitly attribute 16 # Tencent Zhuque Lab (https://github.com/Tencent/AI-Infra-Guard) in its 17 # documentation or user interface, as detailed in the NOTICE file. 18 19 import os 20 import re 21 from datetime import datetime 22 from utils.config import base_dir 23 24 25 class PromptManager: 26 _instance = None 27 _templates = {} 28 29 def __new__(cls): 30 if cls._instance is None: 31 cls._instance = super(PromptManager, cls).__new__(cls) 32 return cls._instance 33 34 def load_template(self, name: str) -> str: 35 if name not in self._templates: 36 path = os.path.join(base_dir, "prompt", "system", f"{name}.md") 37 38 if not os.path.exists(path): 39 raise FileNotFoundError(f"Prompt template '{name}' not found.") 40 41 with open(path, "r", encoding="utf-8") as f: 42 content = f.read() 43 44 content = self._strip_frontmatter(content) 45 self._templates[name] = content 46 47 return self._templates[name] 48 49 def _strip_frontmatter(self, content: str) -> str: 50 """Remove YAML frontmatter (---...---) from template content.""" 51 frontmatter_pattern = r'^---\s*\n.*?\n---\s*\n' 52 content = re.sub(frontmatter_pattern, '', content, count=1, flags=re.DOTALL) 53 return content.strip() 54 55 def format_prompt(self, template_name: str, **kwargs) -> str: 56 template = self.load_template(template_name) 57 if "${NOWTIME}" in template: 58 kwargs.setdefault("NOWTIME", datetime.now().strftime("%Y-%m-%d %H:%M:%S")) 59 formatted = template 60 for key, value in kwargs.items(): 61 placeholder = "{" + key + "}" 62 if placeholder in formatted: 63 formatted = formatted.replace(placeholder, str(value)) 64 alt_placeholder = "${" + key + "}" 65 if alt_placeholder in formatted: 66 formatted = formatted.replace(alt_placeholder, str(value)) 67 68 return formatted 69 70 71 prompt_manager = PromptManager()