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

163 statements  

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

15from dataclasses import dataclass, field 

16from datetime import datetime, timezone 

17from enum import Enum 

18from pathlib import Path 

19from typing import Any 

20 

21 

22# ── Enums & Data Classes ────────────────────────────────────────── 

23 

24class PromptType(str, Enum): 

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(str, Enum): 

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 version: int 

47 content: str 

48 rendered_example: str = "" 

49 created_at: str = "" 

50 author: str = "" 

51 change_summary: str = "" 

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

53 

54 def __post_init__(self): 

55 if not self.created_at: 

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

57 

58 

59@dataclass 

60class PromptTemplate: 

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

62 

63 Usage: 

64 tpl = PromptTemplate( 

65 name="code-review", 

66 type=PromptType.SYSTEM, 

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

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

69 ) 

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

71 """ 

72 

73 name: str 

74 type: PromptType 

75 content: str # Jinja2 template string 

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

77 description: str = "" 

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

79 current_version: int = 1 

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

81 

82 def __post_init__(self): 

83 if not self.versions: 

84 self.versions = [PromptVersion( 

85 version=1, 

86 content=self.content, 

87 change_summary="Initial version", 

88 )] 

89 

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

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

92 try: 

93 from jinja2 import Template, StrictUndefined 

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

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

96 except ImportError: 

97 # Fallback: simple {{ var }} substitution 

98 result = self.content 

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

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

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

102 return result 

103 

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

105 return { 

106 "name": self.name, 

107 "type": self.type.value, 

108 "content": self.content, 

109 "variables": self.variables, 

110 "description": self.description, 

111 "tags": self.tags, 

112 "current_version": self.current_version, 

113 } 

114 

115 def update( 

116 self, 

117 content: str, 

118 change_summary: str = "", 

119 rendered_example: str = "", 

120 author: str = "", 

121 ) -> PromptVersion: 

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

123 self.current_version += 1 

124 version = PromptVersion( 

125 version=self.current_version, 

126 content=content, 

127 rendered_example=rendered_example, 

128 author=author, 

129 change_summary=change_summary, 

130 ) 

131 self.content = content 

132 self.versions.append(version) 

133 return version 

134 

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

136 """Rollback to a previous version.""" 

137 for v in self.versions: 

138 if v.version == target_version: 

139 self.content = v.content 

140 return v 

141 return None 

142 

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

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

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

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

147 if not ver1 or not ver2: 

148 return "" 

149 

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

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

152 diff_lines = [] 

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

154 

155 for i in range(max_len): 

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

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

158 if l1 != l2: 

159 if l1: 

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

161 if l2: 

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

163 return "\n".join(diff_lines) 

164 

165 def hash(self) -> str: 

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

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

168 

169 

170# ── Prompt Hub ──────────────────────────────────────────────────── 

171 

172class PromptHub: 

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

174 

175 Usage: 

176 hub = PromptHub() 

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

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

179 """ 

180 

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

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

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

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

185 

186 # Load from storage if available 

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

188 self._load() 

189 

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

191 """Register a prompt template.""" 

192 self._prompts[template.name] = template 

193 

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

195 """Get a prompt by name.""" 

196 if name not in self._prompts: 

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

198 return self._prompts[name] 

199 

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

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

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

203 

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

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

206 q = query.lower() 

207 results = [] 

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

209 score = 0 

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

211 score += 10 

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

213 score += 5 

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

215 score += 3 

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

217 score += 2 

218 if score > 0: 

219 results.append((score, tpl)) 

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

221 

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

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

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

225 

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

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

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

229 

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

231 """Export all prompts to JSON.""" 

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

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

234 

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

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

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

238 count = 0 

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

240 tpl = PromptTemplate( 

241 name=name, 

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

243 content=pdata["content"], 

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

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

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

247 ) 

248 self.register(tpl) 

249 count += 1 

250 return count 

251 

252 def _load(self) -> None: 

253 """Load prompts from storage directory.""" 

254 if not self.storage_path: 

255 return 

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

257 self.import_json(fpath) 

258 

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

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

261 if not self.storage_path: 

262 return 

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

264 tpl = self._prompts.get(name) 

265 if tpl: 

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

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

268 ) 

269 

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

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

272 self._ab_active[prompt_name] = active 

273 

274 @property 

275 def count(self) -> int: 

276 return len(self._prompts) 

277 

278 

279# ── Built-in Prompt Presets ──────────────────────────────────────── 

280 

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

282 "system/reasoning": { 

283 "type": PromptType.SYSTEM, 

284 "content": ( 

285 "You are an expert reasoning assistant. " 

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

287 "1. Understand the problem\n" 

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

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

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

291 "{{ extra_instructions }}" 

292 ), 

293 "variables": {"extra_instructions": ""}, 

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

295 }, 

296 "system/code-assistant": { 

297 "type": PromptType.SYSTEM, 

298 "content": ( 

299 "You are a senior software engineer. " 

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

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

302 ), 

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

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

305 }, 

306 "few-shot/classification": { 

307 "type": PromptType.FEW_SHOT, 

308 "content": ( 

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

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

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

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

313 ), 

314 "variables": { 

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

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

317 "example_1_label": "positive", 

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

319 "example_2_label": "negative", 

320 "input_text": "", 

321 }, 

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

323 }, 

324 "cot/math": { 

325 "type": PromptType.CHAIN_OF_THOUGHT, 

326 "content": ( 

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

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

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

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

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

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

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

334 "Final Answer:" 

335 ), 

336 "variables": {"problem": ""}, 

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

338 }, 

339 "eval/accuracy": { 

340 "type": PromptType.EVAL, 

341 "content": ( 

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

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

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

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

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

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

348 "Justification:" 

349 ), 

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

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

352 }, 

353} 

354 

355 

356def create_default_hub() -> PromptHub: 

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

358 hub = PromptHub() 

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

360 hub.register(PromptTemplate( 

361 name=name, 

362 type=cfg["type"], 

363 content=cfg["content"], 

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

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

366 )) 

367 return hub 

368 

369 

370# ── Auto-generated compat stubs ── 

371 

372class BUILTIN_PROMPTS: pass