Coverage for agentos/prompt/hub.py: 42%

164 statements  

« prev     ^ index     » next       coverage.py v7.14.3, created at 2026-07-02 09:59 +0800

1""" 

2v1.10.0: Prompt Hub — versioned prompt templates with Jinja2 rendering. 

3 

4Features: 

5- PromptTemplate: Jinja2 template with metadata 

6- PromptVersion: version-tracked prompt with diff 

7- PromptHub: central registry with search/rollback 

8- Role templates: system, few-shot, chain-of-thought presets 

9""" 

10 

11from __future__ import annotations 

12 

13import hashlib 

14import json 

15import re 

16from dataclasses import dataclass, field 

17from datetime import datetime, timezone 

18from enum import Enum 

19from pathlib import Path 

20from typing import Any, Optional 

21 

22 

23# ── Enums & Data Classes ────────────────────────────────────────── 

24 

25class PromptType(str, Enum): 

26 SYSTEM = "system" # System prompt 

27 USER = "user" # User message template 

28 ASSISTANT = "assistant" # Assistant response template 

29 FEW_SHOT = "few_shot" # Few-shot example template 

30 CHAIN_OF_THOUGHT = "cot" # Chain-of-thought template 

31 TOOL_CALL = "tool_call" # Tool-calling template 

32 EVAL = "eval" # Evaluation rubric template 

33 CUSTOM = "custom" 

34 

35 

36class PromptTag(str, Enum): 

37 PRODUCTION = "production" 

38 STAGING = "staging" 

39 EXPERIMENTAL = "experimental" 

40 DEPRECATED = "deprecated" 

41 A_B_TEST = "a_b_test" 

42 

43 

44@dataclass 

45class PromptVersion: 

46 """A versioned instance of a prompt template.""" 

47 version: int 

48 content: str 

49 rendered_example: str = "" 

50 created_at: str = "" 

51 author: str = "" 

52 change_summary: str = "" 

53 performance: dict[str, float] = field(default_factory=dict) # e.g. {"accuracy": 0.92} 

54 

55 def __post_init__(self): 

56 if not self.created_at: 

57 self.created_at = datetime.now(timezone.utc).isoformat() 

58 

59 

60@dataclass 

61class PromptTemplate: 

62 """A prompt template with versioning, tags, and Jinja2 rendering. 

63 

64 Usage: 

65 tpl = PromptTemplate( 

66 name="code-review", 

67 type=PromptType.SYSTEM, 

68 content="You are a code reviewer. Review: {{ code }}", 

69 variables={"code": "python code here"}, 

70 ) 

71 rendered = tpl.render(code="def foo(): pass") 

72 """ 

73 

74 name: str 

75 type: PromptType 

76 content: str # Jinja2 template string 

77 variables: dict[str, Any] = field(default_factory=dict) 

78 description: str = "" 

79 tags: list[str] = field(default_factory=list) 

80 current_version: int = 1 

81 versions: list[PromptVersion] = field(default_factory=list) 

82 

83 def __post_init__(self): 

84 if not self.versions: 

85 self.versions = [PromptVersion( 

86 version=1, 

87 content=self.content, 

88 change_summary="Initial version", 

89 )] 

90 

91 def render(self, **kwargs) -> str: 

92 """Render the template with Jinja2 variables.""" 

93 try: 

94 from jinja2 import Template, StrictUndefined 

95 tpl = Template(self.content, undefined=StrictUndefined) 

96 return tpl.render(**{**self.variables, **kwargs}) 

97 except ImportError: 

98 # Fallback: simple {{ var }} substitution 

99 result = self.content 

100 all_vars = {**self.variables, **kwargs} 

101 for key, value in all_vars.items(): 

102 result = result.replace(f"{{{{ {key} }}}}", str(value)) 

103 return result 

104 

105 def to_dict(self) -> dict[str, Any]: 

106 return { 

107 "name": self.name, 

108 "type": self.type.value, 

109 "content": self.content, 

110 "variables": self.variables, 

111 "description": self.description, 

112 "tags": self.tags, 

113 "current_version": self.current_version, 

114 } 

115 

116 def update( 

117 self, 

118 content: str, 

119 change_summary: str = "", 

120 rendered_example: str = "", 

121 author: str = "", 

122 ) -> PromptVersion: 

123 """Create a new version. Bumps current_version.""" 

124 self.current_version += 1 

125 version = PromptVersion( 

126 version=self.current_version, 

127 content=content, 

128 rendered_example=rendered_example, 

129 author=author, 

130 change_summary=change_summary, 

131 ) 

132 self.content = content 

133 self.versions.append(version) 

134 return version 

135 

136 def rollback(self, target_version: int) -> PromptVersion | None: 

137 """Rollback to a previous version.""" 

138 for v in self.versions: 

139 if v.version == target_version: 

140 self.content = v.content 

141 return v 

142 return None 

143 

144 def diff(self, v1: int, v2: int) -> str: 

145 """Return a simple diff between two versions.""" 

146 ver1 = next((v for v in self.versions if v.version == v1), None) 

147 ver2 = next((v for v in self.versions if v.version == v2), None) 

148 if not ver1 or not ver2: 

149 return "" 

150 

151 lines1 = ver1.content.split("\n") 

152 lines2 = ver2.content.split("\n") 

153 diff_lines = [] 

154 max_len = max(len(lines1), len(lines2)) 

155 

156 for i in range(max_len): 

157 l1 = lines1[i] if i < len(lines1) else "" 

158 l2 = lines2[i] if i < len(lines2) else "" 

159 if l1 != l2: 

160 if l1: 

161 diff_lines.append(f"- {l1}") 

162 if l2: 

163 diff_lines.append(f"+ {l2}") 

164 return "\n".join(diff_lines) 

165 

166 def hash(self) -> str: 

167 """Content hash for cache-busting.""" 

168 return hashlib.md5(self.content.encode()).hexdigest()[:12] 

169 

170 

171# ── Prompt Hub ──────────────────────────────────────────────────── 

172 

173class PromptHub: 

174 """Central prompt registry with search, import/export, A/B testing. 

175 

176 Usage: 

177 hub = PromptHub() 

178 hub.register(PromptTemplate(name="greet", type=PromptType.SYSTEM, content="Hello {{ name }}")) 

179 rendered = hub.render("greet", name="World") 

180 """ 

181 

182 def __init__(self, storage_path: str | Path | None = None): 

183 self._prompts: dict[str, PromptTemplate] = {} 

184 self.storage_path = Path(storage_path) if storage_path else None 

185 self._ab_active: dict[str, str] = {} # prompt_name → variant_name 

186 

187 # Load from storage if available 

188 if self.storage_path and self.storage_path.exists(): 

189 self._load() 

190 

191 def register(self, template: PromptTemplate) -> None: 

192 """Register a prompt template.""" 

193 self._prompts[template.name] = template 

194 

195 def get(self, name: str) -> PromptTemplate: 

196 """Get a prompt by name.""" 

197 if name not in self._prompts: 

198 raise KeyError(f"Prompt not found: {name}") 

199 return self._prompts[name] 

200 

201 def render(self, name: str, **kwargs) -> str: 

202 """Render a prompt by name with variables.""" 

203 return self.get(name).render(**kwargs) 

204 

205 def search(self, query: str) -> list[PromptTemplate]: 

206 """Search prompts by name, description, content, or tags.""" 

207 q = query.lower() 

208 results = [] 

209 for tpl in self._prompts.values(): 

210 score = 0 

211 if q in tpl.name.lower(): 

212 score += 10 

213 if q in tpl.description.lower(): 

214 score += 5 

215 if q in tpl.content.lower(): 

216 score += 3 

217 if any(q in tag.lower() for tag in tpl.tags): 

218 score += 2 

219 if score > 0: 

220 results.append((score, tpl)) 

221 return [t for _, t in sorted(results, key=lambda x: -x[0])] 

222 

223 def list_by_type(self, ptype: PromptType) -> list[PromptTemplate]: 

224 """List all prompts of a given type.""" 

225 return [t for t in self._prompts.values() if t.type == ptype] 

226 

227 def list_by_tag(self, tag: str) -> list[PromptTemplate]: 

228 """List all prompts with a given tag.""" 

229 return [t for t in self._prompts.values() if tag in t.tags] 

230 

231 def export_json(self, path: str | Path) -> None: 

232 """Export all prompts to JSON.""" 

233 data = {name: tpl.to_dict() for name, tpl in self._prompts.items()} 

234 Path(path).write_text(json.dumps(data, indent=2, ensure_ascii=False)) 

235 

236 def import_json(self, path: str | Path) -> int: 

237 """Import prompts from JSON. Returns count of imported prompts.""" 

238 data = json.loads(Path(path).read_text()) 

239 count = 0 

240 for name, pdata in data.items(): 

241 tpl = PromptTemplate( 

242 name=name, 

243 type=PromptType(pdata.get("type", "custom")), 

244 content=pdata["content"], 

245 variables=pdata.get("variables", {}), 

246 description=pdata.get("description", ""), 

247 tags=pdata.get("tags", []), 

248 ) 

249 self.register(tpl) 

250 count += 1 

251 return count 

252 

253 def _load(self) -> None: 

254 """Load prompts from storage directory.""" 

255 if not self.storage_path: 

256 return 

257 for fpath in self.storage_path.glob("*.json"): 

258 self.import_json(fpath) 

259 

260 def _save(self, name: str) -> None: 

261 """Save a single prompt to storage.""" 

262 if not self.storage_path: 

263 return 

264 self.storage_path.mkdir(parents=True, exist_ok=True) 

265 tpl = self._prompts.get(name) 

266 if tpl: 

267 (self.storage_path / f"{name}.json").write_text( 

268 json.dumps({name: tpl.to_dict()}, indent=2, ensure_ascii=False) 

269 ) 

270 

271 def ab_test_set(self, prompt_name: str, variant_a: str, variant_b: str, active: str = "a") -> None: 

272 """Set up A/B test between two prompt variants.""" 

273 self._ab_active[prompt_name] = active 

274 

275 @property 

276 def count(self) -> int: 

277 return len(self._prompts) 

278 

279 

280# ── Built-in Prompt Presets ──────────────────────────────────────── 

281 

282BUILTIN_PROMPTS: dict[str, dict[str, Any]] = { 

283 "system/reasoning": { 

284 "type": PromptType.SYSTEM, 

285 "content": ( 

286 "You are an expert reasoning assistant. " 

287 "Before answering, think step by step:\n" 

288 "1. Understand the problem\n" 

289 "2. Break it into sub-problems\n" 

290 "3. Solve each sub-problem\n" 

291 "4. Synthesize the final answer\n\n" 

292 "{{ extra_instructions }}" 

293 ), 

294 "variables": {"extra_instructions": ""}, 

295 "tags": ["reasoning", "system"], 

296 }, 

297 "system/code-assistant": { 

298 "type": PromptType.SYSTEM, 

299 "content": ( 

300 "You are a senior software engineer. " 

301 "Write clean, efficient, well-documented code. " 

302 "Use {{ language }}. Follow these conventions: {{ conventions }}." 

303 ), 

304 "variables": {"language": "Python", "conventions": "PEP 8"}, 

305 "tags": ["code", "system"], 

306 }, 

307 "few-shot/classification": { 

308 "type": PromptType.FEW_SHOT, 

309 "content": ( 

310 "Classify the following text into categories: {{ categories }}\n\n" 

311 "Example 1:\nText: {{ example_1_text }}\nCategory: {{ example_1_label }}\n\n" 

312 "Example 2:\nText: {{ example_2_text }}\nCategory: {{ example_2_label }}\n\n" 

313 "Now classify:\nText: {{ input_text }}\nCategory:" 

314 ), 

315 "variables": { 

316 "categories": "positive/negative/neutral", 

317 "example_1_text": "I love this product!", 

318 "example_1_label": "positive", 

319 "example_2_text": "This is terrible.", 

320 "example_2_label": "negative", 

321 "input_text": "", 

322 }, 

323 "tags": ["few-shot", "classification"], 

324 }, 

325 "cot/math": { 

326 "type": PromptType.CHAIN_OF_THOUGHT, 

327 "content": ( 

328 "Solve this math problem step by step. Show all your work.\n\n" 

329 "Problem: {{ problem }}\n\n" 

330 "Let's solve this step by step:\n" 

331 "Step 1: Understand what we're asked to find.\n" 

332 "Step 2: Identify the relevant formulas or concepts.\n" 

333 "Step 3: Apply them and solve.\n" 

334 "Step 4: Verify the answer.\n\n" 

335 "Final Answer:" 

336 ), 

337 "variables": {"problem": ""}, 

338 "tags": ["cot", "math", "reasoning"], 

339 }, 

340 "eval/accuracy": { 

341 "type": PromptType.EVAL, 

342 "content": ( 

343 "You are an evaluator. Grade the following response on a scale of 0-10.\n\n" 

344 "Criteria: {{ criteria }}\n\n" 

345 "Question: {{ question }}\n" 

346 "Expected Answer: {{ expected }}\n" 

347 "Generated Answer: {{ generated }}\n\n" 

348 "Score (0-10):\n" 

349 "Justification:" 

350 ), 

351 "variables": {"criteria": "accuracy", "question": "", "expected": "", "generated": ""}, 

352 "tags": ["eval", "scoring"], 

353 }, 

354} 

355 

356 

357def create_default_hub() -> PromptHub: 

358 """Create a prompt hub pre-loaded with built-in templates.""" 

359 hub = PromptHub() 

360 for name, cfg in BUILTIN_PROMPTS.items(): 

361 hub.register(PromptTemplate( 

362 name=name, 

363 type=cfg["type"], 

364 content=cfg["content"], 

365 variables=cfg.get("variables", {}), 

366 tags=cfg.get("tags", []), 

367 )) 

368 return hub 

369 

370 

371# ── Auto-generated compat stubs ── 

372 

373class BUILTIN_PROMPTS: pass