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

164 statements  

« prev     ^ index     » next       coverage.py v7.14.3, created at 2026-07-08 13:14 +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 

15from dataclasses import dataclass, field 

16from datetime import UTC, datetime 

17from enum import StrEnum 

18from pathlib import Path 

19from typing import Any 

20 

21# ── Enums & Data Classes ────────────────────────────────────────── 

22 

23 

24class PromptType(StrEnum): 

25 SYSTEM = "system" # System prompt 

26 USER = "user" # User message template 

27 ASSISTANT = "assistant" # Assistant response template 

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

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

30 TOOL_CALL = "tool_call" # Tool-calling template 

31 EVAL = "eval" # Evaluation rubric template 

32 CUSTOM = "custom" 

33 

34 

35class PromptTag(StrEnum): 

36 PRODUCTION = "production" 

37 STAGING = "staging" 

38 EXPERIMENTAL = "experimental" 

39 DEPRECATED = "deprecated" 

40 A_B_TEST = "a_b_test" 

41 

42 

43@dataclass 

44class PromptVersion: 

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

46 

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(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 = [ 

86 PromptVersion( 

87 version=1, 

88 content=self.content, 

89 change_summary="Initial version", 

90 ) 

91 ] 

92 

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

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

95 try: 

96 from jinja2 import StrictUndefined, Template 

97 

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

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

100 except ImportError: 

101 # Fallback: simple {{ var }} substitution 

102 result = self.content 

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

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

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

106 return result 

107 

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

109 return { 

110 "name": self.name, 

111 "type": self.type.value, 

112 "content": self.content, 

113 "variables": self.variables, 

114 "description": self.description, 

115 "tags": self.tags, 

116 "current_version": self.current_version, 

117 } 

118 

119 def update( 

120 self, 

121 content: str, 

122 change_summary: str = "", 

123 rendered_example: str = "", 

124 author: str = "", 

125 ) -> PromptVersion: 

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

127 self.current_version += 1 

128 version = PromptVersion( 

129 version=self.current_version, 

130 content=content, 

131 rendered_example=rendered_example, 

132 author=author, 

133 change_summary=change_summary, 

134 ) 

135 self.content = content 

136 self.versions.append(version) 

137 return version 

138 

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

140 """Rollback to a previous version.""" 

141 for v in self.versions: 

142 if v.version == target_version: 

143 self.content = v.content 

144 return v 

145 return None 

146 

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

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

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

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

151 if not ver1 or not ver2: 

152 return "" 

153 

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

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

156 diff_lines = [] 

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

158 

159 for i in range(max_len): 

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

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

162 if l1 != l2: 

163 if l1: 

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

165 if l2: 

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

167 return "\n".join(diff_lines) 

168 

169 def hash(self) -> str: 

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

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

172 

173 

174# ── Prompt Hub ──────────────────────────────────────────────────── 

175 

176 

177class PromptHub: 

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

179 

180 Usage: 

181 hub = PromptHub() 

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

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

184 """ 

185 

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

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

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

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

190 

191 # Load from storage if available 

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

193 self._load() 

194 

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

196 """Register a prompt template.""" 

197 self._prompts[template.name] = template 

198 

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

200 """Get a prompt by name.""" 

201 if name not in self._prompts: 

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

203 return self._prompts[name] 

204 

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

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

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

208 

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

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

211 q = query.lower() 

212 results = [] 

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

214 score = 0 

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

216 score += 10 

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

218 score += 5 

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

220 score += 3 

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

222 score += 2 

223 if score > 0: 

224 results.append((score, tpl)) 

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

226 

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

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

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

230 

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

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

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

234 

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

236 """Export all prompts to JSON.""" 

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

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

239 

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

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

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

243 count = 0 

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

245 tpl = PromptTemplate( 

246 name=name, 

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

248 content=pdata["content"], 

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

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

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

252 ) 

253 self.register(tpl) 

254 count += 1 

255 return count 

256 

257 def _load(self) -> None: 

258 """Load prompts from storage directory.""" 

259 if not self.storage_path: 

260 return 

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

262 self.import_json(fpath) 

263 

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

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

266 if not self.storage_path: 

267 return 

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

269 tpl = self._prompts.get(name) 

270 if tpl: 

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

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

273 ) 

274 

275 def ab_test_set( 

276 self, prompt_name: str, variant_a: str, variant_b: str, active: str = "a" 

277 ) -> None: 

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

279 self._ab_active[prompt_name] = active 

280 

281 @property 

282 def count(self) -> int: 

283 return len(self._prompts) 

284 

285 

286# ── Built-in Prompt Presets ──────────────────────────────────────── 

287 

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

289 "system/reasoning": { 

290 "type": PromptType.SYSTEM, 

291 "content": ( 

292 "You are an expert reasoning assistant. " 

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

294 "1. Understand the problem\n" 

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

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

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

298 "{{ extra_instructions }}" 

299 ), 

300 "variables": {"extra_instructions": ""}, 

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

302 }, 

303 "system/code-assistant": { 

304 "type": PromptType.SYSTEM, 

305 "content": ( 

306 "You are a senior software engineer. " 

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

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

309 ), 

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

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

312 }, 

313 "few-shot/classification": { 

314 "type": PromptType.FEW_SHOT, 

315 "content": ( 

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

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

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

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

320 ), 

321 "variables": { 

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

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

324 "example_1_label": "positive", 

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

326 "example_2_label": "negative", 

327 "input_text": "", 

328 }, 

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

330 }, 

331 "cot/math": { 

332 "type": PromptType.CHAIN_OF_THOUGHT, 

333 "content": ( 

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

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

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

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

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

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

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

341 "Final Answer:" 

342 ), 

343 "variables": {"problem": ""}, 

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

345 }, 

346 "eval/accuracy": { 

347 "type": PromptType.EVAL, 

348 "content": ( 

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

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

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

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

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

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

355 "Justification:" 

356 ), 

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

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

359 }, 

360} 

361 

362 

363def create_default_hub() -> PromptHub: 

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

365 hub = PromptHub() 

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

367 hub.register( 

368 PromptTemplate( 

369 name=name, 

370 type=cfg["type"], 

371 content=cfg["content"], 

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

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

374 ) 

375 ) 

376 return hub 

377 

378 

379# ── Auto-generated compat stubs ── 

380 

381 

382class BUILTIN_PROMPTS: # noqa: F811,N801 

383 pass