/ agent-scan / utils / prompt_manager.py
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()