Coverage for agentos/prompts/few_shot.py: 37%

1""" 

2Few-Shot Example Management — intelligent few-shot selection strategies. 

3 

4Supports similarity-based, random, diversity-maximizing, and 

5custom selection algorithms for constructing optimal few-shot prompts. 

6""" 

7 

8import hashlib 

9import random 

10import re 

11from collections.abc import Iterable, Sequence 

12from dataclasses import dataclass, field 

13from enum import StrEnum 

14from typing import Any 

15 

16 

17class SelectionStrategy(StrEnum): 

18 """Strategy for selecting few-shot examples.""" 

19 

20 RANDOM = "random" 

21 SIMILARITY = "similarity" 

22 DIVERSITY = "diversity" 

23 RECENCY = "recency" 

24 LABEL_BALANCED = "label_balanced" 

25 ACTIVE_LEARNING = "active_learning" 

26 

27 

28@dataclass 

29class Example: 

30 """A single training example for few-shot learning.""" 

31 

32 input: str 

33 output: str 

34 id: str = "" 

35 label: str = "" 

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

37 score: float = 0.0 

38 

39 def __post_init__(self): 

40 if not self.id: 

41 self.id = hashlib.md5(f"{self.input}{self.output}".encode()).hexdigest()[:12] 

42 

43 

44class FewShotSelector: 

45 """Selects and formats the best few-shot examples for a given query. 

46 

47 Usage:: 

48 

49 examples = [ 

50 Example(input="What is 2+2?", output="4", label="math"), 

51 Example(input="Capital of France?", output="Paris", label="geo"), 

52 ] 

53 selector = FewShotSelector(examples, strategy=SelectionStrategy.SIMILARITY) 

54 prompt = selector.build_prompt("What is 3+5?", base_instruction="Answer:") 

55 """ 

56 

57 DEFAULT_FORMAT = "Q: {input}\nA: {output}" 

58 MAX_TOKEN_ESTIMATE = 4096 

59 

60 def __init__( 

61 self, 

62 examples: Sequence[Example], 

63 strategy: SelectionStrategy = SelectionStrategy.SIMILARITY, 

64 max_examples: int = 5, 

65 example_format: str = "", 

66 seed: int = 42, 

67 ): 

68 self.examples = list(examples) 

69 self.strategy = strategy 

70 self.max_examples = max_examples 

71 self.example_format = example_format or self.DEFAULT_FORMAT 

72 random.seed(seed) 

73 

74 def select(self, query: str, k: int | None = None) -> list[Example]: 

75 """Select top-k examples for the given query.""" 

76 k = k or self.max_examples 

77 if not self.examples: 

78 return [] 

79 

80 strategy_map = { 

81 SelectionStrategy.RANDOM: self._select_random, 

82 SelectionStrategy.SIMILARITY: self._select_similarity, 

83 SelectionStrategy.DIVERSITY: self._select_diversity, 

84 SelectionStrategy.RECENCY: self._select_recency, 

85 SelectionStrategy.LABEL_BALANCED: self._select_label_balanced, 

86 } 

87 selector_fn = strategy_map.get(self.strategy, self._select_similarity) 

88 return selector_fn(query, k) 

89 

90 def build_prompt( 

91 self, 

92 query: str, 

93 base_instruction: str = "", 

94 k: int | None = None, 

95 ) -> str: 

96 """Build a complete few-shot prompt string.""" 

97 selected = self.select(query, k) 

98 parts: list[str] = [] 

99 if base_instruction: 

100 parts.append(base_instruction) 

101 for ex in selected: 

102 parts.append(self.example_format.format(input=ex.input, output=ex.output)) 

103 parts.append(self.example_format.format(input=query, output="")) 

104 return "\n\n".join(parts) 

105 

106 def add_example(self, example: Example): 

107 """Add a new example to the pool.""" 

108 self.examples.append(example) 

109 

110 def remove_example(self, example_id: str): 

111 """Remove an example by ID.""" 

112 self.examples = [e for e in self.examples if e.id != example_id] 

113 

114 def set_score(self, example_id: str, score: float): 

115 """Update the utility score for an example.""" 

116 for ex in self.examples: 

117 if ex.id == example_id: 

118 ex.score = score 

119 break 

120 

121 def _select_random(self, _query: str, k: int) -> list[Example]: 

122 return random.sample(self.examples, min(k, len(self.examples))) 

123 

124 def _select_similarity(self, query: str, k: int) -> list[Example]: 

125 """Jaccard-based token similarity for fast selection.""" 

126 query_tokens = set(_tokenize(query)) 

127 scored = [(ex, self._jaccard(query_tokens, ex)) for ex in self.examples] 

128 scored.sort(key=lambda x: x[1], reverse=True) 

129 return [ex for ex, _ in scored[:k]] 

130 

131 def _select_diversity(self, _query: str, k: int) -> list[Example]: 

132 """Maximize diversity via greedy farthest-first.""" 

133 if k >= len(self.examples): 

134 return list(self.examples) 

135 # Start with a random seed 

136 selected = [random.choice(self.examples)] 

137 remaining = [e for e in self.examples if e not in selected] 

138 while len(selected) < k and remaining: 

139 # Pick the example least similar to any already selected 

140 best = max( 

141 remaining, 

142 key=lambda ex: min(self._jaccard(set(_tokenize(ex.input)), s) for s in selected), 

143 ) 

144 selected.append(best) 

145 remaining.remove(best) 

146 return selected 

147 

148 def _select_recency(self, _query: str, k: int) -> list[Example]: 

149 """Most recent examples first (assumes append order = recency).""" 

150 return list(reversed(self.examples[-k:])) 

151 

152 def _select_label_balanced(self, _query: str, k: int) -> list[Example]: 

153 """Balance selection across unique labels.""" 

154 by_label: dict[str, list[Example]] = {} 

155 for ex in self.examples: 

156 by_label.setdefault(ex.label or "_unlabeled", []).append(ex) 

157 labels = list(by_label.keys()) 

158 result: list[Example] = [] 

159 idx = 0 

160 while len(result) < k and any(by_label.values()): 

161 label = labels[idx % len(labels)] 

162 pool = by_label[label] 

163 if pool: 

164 result.append(pool.pop(random.randrange(len(pool)))) 

165 idx += 1 

166 return result 

167 

168 @staticmethod 

169 def _jaccard(tokens_a: set[str], example: Example) -> float: 

170 tokens_b = set(_tokenize(example.input)) 

171 if not tokens_a or not tokens_b: 

172 return 0.0 

173 intersection = tokens_a & tokens_b 

174 union = tokens_a | tokens_b 

175 return len(intersection) / len(union) 

176 

177 

178def build_examples( 

179 pairs: Iterable[tuple[str, str]], 

180 labels: Iterable[str] | None = None, 

181 metadata: list[dict] | None = None, 

182) -> list[Example]: 

183 """Convenience factory to build a list of Example objects. 

184 

185 Args: 

186 pairs: Iterable of (input, output) tuples. 

187 labels: Optional labels for each example. 

188 metadata: Optional metadata dicts. 

189 

190 Returns: 

191 List of ``Example`` objects. 

192 """ 

193 examples: list[Example] = [] 

194 label_list = list(labels) if labels else [] 

195 meta_list = list(metadata) if metadata else [] 

196 for i, (inp, out) in enumerate(pairs): 

197 ex = Example( 

198 input=inp, 

199 output=out, 

200 label=label_list[i] if i < len(label_list) else "", 

201 metadata=meta_list[i] if i < len(meta_list) else {}, 

202 ) 

203 examples.append(ex) 

204 return examples 

205 

206 

207def _tokenize(text: str) -> list[str]: 

208 """Simple whitespace+punctuation tokenizer.""" 

209 return re.findall(r"\w+", text.lower())