Coverage for agentos/prompts/optimizer.py: 58%

1""" 

2Prompt Optimizer — DSPy-inspired automatic prompt improvement via 

3iterative refinement, few-shot bootstrapping, and multi-strategy optimization. 

4""" 

5 

6import random 

7from collections.abc import Callable 

8from dataclasses import dataclass, field 

9from enum import StrEnum 

10from typing import Any 

11 

12 

13class OptimizationStrategy(StrEnum): 

14 """Available optimization approaches.""" 

15 

16 BOOTSTRAP_FEWSHOT = "bootstrap_fewshot" 

17 MIPRO = "mipro" # Multi-prompt instruction proposal 

18 GRADIENT_FREE = "gradient_free" 

19 ENSEMBLE = "ensemble" 

20 CHAIN_OF_THOUGHT = "chain_of_thought" 

21 

22 

23@dataclass 

24class OptimizerConfig: 

25 """Configuration for prompt optimization runs.""" 

26 

27 strategy: OptimizationStrategy = OptimizationStrategy.BOOTSTRAP_FEWSHOT 

28 max_iterations: int = 10 

29 candidates_per_iteration: int = 4 

30 eval_samples: int = 20 

31 target_metric: str = "accuracy" 

32 target_threshold: float = 0.90 

33 temperature_range: tuple[float, float] = (0.1, 0.9) 

34 keep_top_k: int = 3 

35 early_stop_patience: int = 3 

36 seed: int = 42 

37 

38 

39@dataclass 

40class PromptCandidate: 

41 """A single prompt variant under evaluation.""" 

42 

43 id: str 

44 text: str 

45 score: float = 0.0 

46 metrics: dict[str, float] = field(default_factory=dict) 

47 generation: int = 0 

48 parent_id: str = "" 

49 

50 

51@dataclass 

52class OptimizationResult: 

53 """Final result after optimization converges or exhausts budget.""" 

54 

55 best_prompt: str 

56 best_score: float 

57 iterations: int 

58 candidates_evaluated: int 

59 strategy: OptimizationStrategy 

60 history: list[PromptCandidate] = field(default_factory=list) 

61 metadata: dict[str, Any] = field(default_factory=dict) 

62 

63 

64class PromptOptimizer: 

65 """Iteratively refines prompts using a pluggable scoring function. 

66 

67 Usage:: 

68 

69 def score(prompt: str) -> float: 

70 # run your LLM eval and return metric 

71 return measure(prompt) 

72 

73 opt = PromptOptimizer(config) 

74 result = opt.optimize(base_prompt, score_fn=score) 

75 print(result.best_prompt) 

76 """ 

77 

78 SEED_TEMPLATES = { 

79 OptimizationStrategy.BOOTSTRAP_FEWSHOT: [ 

80 "{base}\n\nHere are some examples:\n{examples}", 

81 "Task: {base}\n\nIllustrative examples:\n{examples}", 

82 "{base}\n\nDemonstrations:\n{examples}", 

83 ], 

84 OptimizationStrategy.CHAIN_OF_THOUGHT: [ 

85 "{base}\n\nLet's think step by step.", 

86 "{base}\n\nBreak this down logically:", 

87 "Solve step-by-step:\n{base}", 

88 ], 

89 OptimizationStrategy.ENSEMBLE: [ 

90 "Consider multiple perspectives:\n{base}", 

91 "Review from different angles:\n{base}", 

92 "Analyze comprehensively:\n{base}", 

93 ], 

94 } 

95 

96 def __init__(self, config: OptimizerConfig | None = None): 

97 self.config = config or OptimizerConfig() 

98 random.seed(self.config.seed) 

99 

100 def optimize( 

101 self, 

102 base_prompt: str, 

103 score_fn: Callable[[str], float], 

104 few_shot_examples: list[str] | None = None, 

105 ) -> OptimizationResult: 

106 """Run optimization and return the best prompt found.""" 

107 best = PromptCandidate( 

108 id="base", 

109 text=base_prompt, 

110 score=score_fn(base_prompt), 

111 generation=0, 

112 ) 

113 history = [best] 

114 no_improve = 0 

115 

116 for iteration in range(1, self.config.max_iterations + 1): 

117 candidates = self._generate_candidates(best.text, iteration, few_shot_examples) 

118 for c in candidates: 

119 c.score = score_fn(c.text) 

120 history.append(c) 

121 

122 # Select best from this iteration 

123 iteration_best = max(candidates, key=lambda c: c.score) 

124 if iteration_best.score > best.score: 

125 best = iteration_best 

126 no_improve = 0 

127 else: 

128 no_improve += 1 

129 

130 # Keep top-K across all generations 

131 history.sort(key=lambda c: c.score, reverse=True) 

132 history = history[: self.config.keep_top_k * 2] 

133 

134 if best.score >= self.config.target_threshold: 

135 break 

136 if no_improve >= self.config.early_stop_patience: 

137 break 

138 

139 return OptimizationResult( 

140 best_prompt=best.text, 

141 best_score=best.score, 

142 iterations=iteration, 

143 candidates_evaluated=len(history), 

144 strategy=self.config.strategy, 

145 history=history[: self.config.keep_top_k], 

146 ) 

147 

148 def _generate_candidates( 

149 self, 

150 base: str, 

151 generation: int, 

152 examples: list[str] | None, 

153 ) -> list[PromptCandidate]: 

154 templates = self.SEED_TEMPLATES.get( 

155 self.config.strategy, 

156 self.SEED_TEMPLATES[OptimizationStrategy.BOOTSTRAP_FEWSHOT], 

157 ) 

158 candidates: list[PromptCandidate] = [] 

159 

160 for i in range(self.config.candidates_per_iteration): 

161 tmpl = random.choice(templates) 

162 text = tmpl.format( 

163 base=base, 

164 examples=self._format_examples(examples) if examples else "", 

165 ) 

166 # Add small perturbations 

167 if random.random() < 0.3 and generation > 1: 

168 text = self._perturb(text) 

169 

170 candidates.append( 

171 PromptCandidate( 

172 id=f"gen{generation}_{i}", 

173 text=text, 

174 generation=generation, 

175 parent_id="base" if generation == 1 else f"gen{generation-1}_0", 

176 ) 

177 ) 

178 

179 return candidates 

180 

181 def _format_examples(self, examples: list[str]) -> str: 

182 return "\n".join(f"- {e}" for e in examples[:5]) 

183 

184 def _perturb(self, text: str) -> str: 

185 """Apply minor random perturbations.""" 

186 perturbations = [ 

187 lambda t: t.replace(".", ". Please be thorough."), 

188 lambda t: "Carefully: " + t, 

189 lambda t: t + "\nBe precise and concise.", 

190 lambda t: t.replace(":", ":\n"), 

191 lambda t: t.replace("the ", "the relevant "), 

192 ] 

193 return random.choice(perturbations)(text)