Coverage for src/local_deep_research/advanced_search_system/questions/browsecomp_question.py: 95%

152 statements  

« prev     ^ index     » next       coverage.py v7.16.0, created at 2026-09-06 15:42 +0000

1""" 

2BrowseComp-specific question generation that creates progressive, entity-focused searches. 

3""" 

4 

5import re 

6from typing import Dict, List, Optional 

7 

8from loguru import logger 

9 

10from ...utilities.json_utils import get_llm_response_text 

11from .base_question import BaseQuestionGenerator 

12 

13 

14class BrowseCompQuestionGenerator(BaseQuestionGenerator): 

15 """ 

16 Question generator optimized for BrowseComp-style queries. 

17 

18 Key features: 

19 1. Extract concrete entities (dates, numbers, names, places) 

20 2. Generate progressive search combinations 

21 3. Start broad, then narrow systematically 

22 4. Focus on verifiable facts 

23 """ 

24 

25 def __init__( 

26 self, 

27 model, 

28 knowledge_truncate_length: int = 1500, 

29 previous_searches_limit: int = 10, 

30 ): 

31 """Initialize the question generator. 

32 

33 Args: 

34 model: The LLM model to use for generation 

35 knowledge_truncate_length: Max chars for knowledge in prompts (None=unlimited) 

36 previous_searches_limit: Max previous searches to show (None=unlimited) 

37 """ 

38 super().__init__(model) 

39 self.extracted_entities: Dict[str, List[str]] = {} 

40 self.search_progression: List[str] = [] 

41 self.knowledge_truncate_length = knowledge_truncate_length 

42 self.previous_searches_limit = previous_searches_limit 

43 

44 def generate_questions( 

45 self, 

46 current_knowledge: str, 

47 query: str, 

48 questions_per_iteration: int = 5, 

49 questions_by_iteration: Optional[dict] = None, 

50 results_by_iteration: Optional[dict] = None, 

51 iteration: int = 1, 

52 ) -> List[str]: 

53 """Generate progressive search queries for BrowseComp problems.""" 

54 questions_by_iteration = questions_by_iteration or {} 

55 

56 # First iteration: Extract entities and create initial searches 

57 if iteration == 1 or not self.extracted_entities: 

58 self.extracted_entities = self._extract_entities(query) 

59 return self._generate_initial_searches( 

60 query, self.extracted_entities, questions_per_iteration 

61 ) 

62 

63 # Subsequent iterations: Progressive refinement 

64 return self._generate_progressive_searches( 

65 query, 

66 current_knowledge, 

67 self.extracted_entities, 

68 questions_by_iteration, 

69 results_by_iteration or {}, 

70 questions_per_iteration, 

71 iteration, 

72 ) 

73 

74 def _extract_entities(self, query: str) -> Dict[str, List[str]]: 

75 """Extract concrete entities from the query.""" 

76 prompt = f"""Extract ALL concrete, searchable entities from this query: 

77 

78Query: {query} 

79 

80Extract: 

811. TEMPORAL: All years, dates, time periods (e.g., "2018", "between 1995 and 2006", "2023") 

822. NUMERICAL: All numbers, statistics, counts (e.g., "300", "more than 3", "4-3", "84.5%") 

833. NAMES: Partial names, name hints, proper nouns (e.g., "Dartmouth", "EMNLP", "Plastic Man") 

844. LOCATIONS: Places, institutions, geographic features (e.g., "Pennsylvania", "Grand Canyon") 

855. DESCRIPTORS: Key descriptive terms (e.g., "fourth wall", "ascetics", "decider game") 

86 

87For TEMPORAL entities, if there's a range (e.g., "between 2018-2023"), list EACH individual year. 

88 

89Format your response as: 

90TEMPORAL: [entity1], [entity2], ... 

91NUMERICAL: [entity1], [entity2], ... 

92NAMES: [entity1], [entity2], ... 

93LOCATIONS: [entity1], [entity2], ... 

94DESCRIPTORS: [entity1], [entity2], ... 

95""" 

96 

97 response = self.model.invoke(prompt) 

98 content = get_llm_response_text(response) 

99 

100 entities: Dict[str, List[str]] = { 

101 "temporal": [], 

102 "numerical": [], 

103 "names": [], 

104 "locations": [], 

105 "descriptors": [], 

106 } 

107 

108 for line in content.strip().split("\n"): 

109 line = line.strip() 

110 if ":" in line: 

111 category, values = line.split(":", 1) 

112 category = category.strip().lower() 

113 if category in entities: 

114 # Parse comma-separated values 

115 values = [v.strip() for v in values.split(",") if v.strip()] 

116 entities[category].extend(values) 

117 

118 # Expand temporal ranges 

119 entities["temporal"] = self._expand_temporal_ranges( 

120 entities["temporal"] 

121 ) 

122 

123 logger.info(f"Extracted entities: {entities}") 

124 return entities 

125 

126 def _expand_temporal_ranges( 

127 self, temporal_entities: List[str] 

128 ) -> List[str]: 

129 """Expand year ranges into individual years.""" 

130 expanded = [] 

131 for entity in temporal_entities: 

132 # Check for range patterns like "2018-2023" or "between 1995 and 2006" 

133 range_match = re.search( 

134 r"(\d{4})[-\s]+(?:to|and)?\s*(\d{4})", entity 

135 ) 

136 if range_match: 

137 start_year = int(range_match.group(1)) 

138 end_year = int(range_match.group(2)) 

139 for year in range(start_year, end_year + 1): 

140 expanded.append(str(year)) 

141 else: 

142 # Single year or other temporal entity 

143 year_match = re.search(r"\d{4}", entity) 

144 if year_match: 

145 expanded.append(year_match.group()) 

146 else: 

147 expanded.append(entity) 

148 

149 return list(set(expanded)) # Remove duplicates 

150 

151 def _generate_initial_searches( 

152 self, query: str, entities: Dict[str, List[str]], num_questions: int 

153 ) -> List[str]: 

154 """Generate initial broad searches.""" 

155 searches = [] 

156 

157 # 1. Original query (always include) 

158 searches.append(query) 

159 

160 # If only 1 question requested, return just the original query 

161 if num_questions <= 1: 

162 return searches[:1] 

163 

164 # 2. Domain exploration searches (combine key entities) 

165 if entities["names"] and len(searches) < num_questions: 

166 for name in entities["names"][:2]: # Top 2 names 

167 if len(searches) >= num_questions: 

168 break 

169 searches.append(f"{name}") 

170 if entities["descriptors"] and len(searches) < num_questions: 

171 searches.append(f"{name} {entities['descriptors'][0]}") 

172 

173 # 3. Temporal searches if years are important 

174 if ( 

175 entities["temporal"] 

176 and len(entities["temporal"]) <= 10 

177 and len(searches) < num_questions 

178 ): 

179 # For small year ranges, search each year with a key term 

180 key_term = ( 

181 entities["names"][0] 

182 if entities["names"] 

183 else entities["descriptors"][0] 

184 if entities["descriptors"] 

185 else "" 

186 ) 

187 for year in entities["temporal"][:5]: # Limit to 5 years initially 

188 if len(searches) >= num_questions: 188 ↛ 189line 188 didn't jump to line 189 because the condition on line 188 was never true

189 break 

190 if key_term: 190 ↛ 187line 190 didn't jump to line 187 because the condition on line 190 was always true

191 searches.append(f"{key_term} {year}") 

192 

193 # 4. Location-based searches 

194 if entities["locations"] and len(searches) < num_questions: 

195 for location in entities["locations"][:2]: 

196 if len(searches) >= num_questions: 196 ↛ 197line 196 didn't jump to line 197 because the condition on line 196 was never true

197 break 

198 searches.append(f"{location}") 

199 if entities["descriptors"] and len(searches) < num_questions: 

200 searches.append(f"{location} {entities['descriptors'][0]}") 

201 

202 # Remove duplicates and limit to requested number 

203 seen = set() 

204 unique_searches = [] 

205 for s in searches: 

206 if s.lower() not in seen: 

207 seen.add(s.lower()) 

208 unique_searches.append(s) 

209 if len(unique_searches) >= num_questions: 

210 break 

211 

212 return unique_searches[:num_questions] 

213 

214 def _generate_progressive_searches( 

215 self, 

216 query: str, 

217 current_knowledge: str, 

218 entities: Dict[str, List[str]], 

219 questions_by_iteration: dict, 

220 results_by_iteration: dict, 

221 num_questions: int, 

222 iteration: int, 

223 ) -> List[str]: 

224 """Generate progressively more specific searches based on findings.""" 

225 

226 # Only add strategy instructions if we have actual result data (adaptive mode) 

227 strategy_instruction = "" 

228 if results_by_iteration: 

229 # Check if recent searches are failing (returning 0 results) 

230 recent_iterations = list(range(max(1, iteration - 5), iteration)) 

231 zero_count = sum( 

232 1 

233 for i in recent_iterations 

234 if results_by_iteration.get(i, 1) == 0 

235 ) 

236 searches_failing = zero_count >= 3 

237 

238 # Adjust strategy based on success/failure 

239 if searches_failing: 

240 strategy_instruction = """ 

241IMPORTANT: Your recent searches are returning 0 results - they are TOO NARROW! 

242- Use FEWER constraints (1-2 terms instead of 4-5) 

243- Try BROADER, more general searches 

244- Remove overly specific combinations 

245- Focus on key concepts, not detailed entity combinations 

246""" 

247 else: 

248 strategy_instruction = """ 

249Focus on finding the specific answer by combining entities systematically. 

250""" 

251 

252 # Analyze what we've found so far 

253 prompt = f"""Based on our search progress, generate targeted follow-up searches. 

254{strategy_instruction} 

255 

256Original Query: {query} 

257 

258Entities Found: 

259- Names/Terms: {", ".join(entities["names"][:5])} 

260- Years: {", ".join(entities["temporal"][:5])} 

261- Locations: {", ".join(entities["locations"][:3])} 

262- Key Features: {", ".join(entities["descriptors"][:3])} 

263 

264Current Knowledge Summary: 

265{current_knowledge[: self.knowledge_truncate_length] if self.knowledge_truncate_length else current_knowledge} 

266 

267Previous Searches: 

268{self._format_previous_searches(questions_by_iteration, results_by_iteration)} 

269 

270Generate {num_questions} NEW search queries that: 

2711. Combine 2-3 entities we haven't tried together 

2722. If we found candidate names, search for them with other constraints 

2733. For year ranges, systematically cover years we haven't searched 

2744. Use quotes for exact phrases when beneficial 

275 

276Focus on finding the specific answer, not general information. 

277 

278Format: One search per line 

279""" 

280 

281 response = self.model.invoke(prompt) 

282 content = get_llm_response_text(response) 

283 

284 # Extract searches from response 

285 searches = [] 

286 for line in content.strip().split("\n"): 

287 line = line.strip() 

288 if line and not line.endswith(":") and len(line) > 5: 288 ↛ 286line 288 didn't jump to line 286 because the condition on line 288 was always true

289 # Clean up common prefixes 

290 for prefix in ["Q:", "Search:", "-", "*", "•"]: 

291 if line.startswith(prefix): 

292 line = line[len(prefix) :].strip() 

293 if line: 293 ↛ 286line 293 didn't jump to line 286 because the condition on line 293 was always true

294 searches.append(line) 

295 

296 # Ensure we have enough searches, but respect the limit 

297 while len(searches) < num_questions: 

298 # Generate combinations programmatically 

299 if iteration <= 5 and entities["temporal"]: 

300 # Continue with year-based searches 

301 added_any = False 

302 for year in entities["temporal"]: 302 ↛ 311line 302 didn't jump to line 311 because the loop on line 302 didn't complete

303 if not self._was_searched(year, questions_by_iteration): 303 ↛ 302line 303 didn't jump to line 302 because the condition on line 303 was always true

304 base_term = ( 

305 entities["names"][0] if entities["names"] else "" 

306 ) 

307 searches.append(f"{base_term} {year}".strip()) 

308 added_any = True 

309 if len(searches) >= num_questions: 

310 break 

311 if not added_any: 311 ↛ 312line 311 didn't jump to line 312 because the condition on line 311 was never true

312 break # No more year searches to add 

313 else: 

314 # Combine multiple constraints 

315 added_any = False 

316 if entities["names"] and entities["descriptors"]: 

317 for name in entities["names"]: 

318 for desc in entities["descriptors"]: 

319 combo = f"{name} {desc}" 

320 if not self._was_searched( 

321 combo, questions_by_iteration 

322 ): 

323 searches.append(combo) 

324 added_any = True 

325 if len(searches) >= num_questions: 

326 break 

327 if len(searches) >= num_questions: 

328 break 

329 if not added_any: 

330 break # No more combinations to add 

331 

332 return searches[:num_questions] 

333 

334 def _format_previous_searches( 

335 self, questions_by_iteration: dict, results_by_iteration: dict 

336 ) -> str: 

337 """Format previous searches for context with result counts.""" 

338 formatted = [] 

339 for iteration, questions in questions_by_iteration.items(): 

340 if isinstance(questions, list): 340 ↛ 339line 340 didn't jump to line 339 because the condition on line 340 was always true

341 result_count = results_by_iteration.get(iteration, "?") 

342 # Limit questions per iteration (main uses 3) 

343 questions_to_show = ( 

344 questions[:3] if self.previous_searches_limit else questions 

345 ) 

346 for q in questions_to_show: 

347 # Only show result counts if we have actual data (not "?") 

348 if result_count != "?": 

349 formatted.append( 

350 f"Iteration {iteration}: {q} ({result_count} results)" 

351 ) 

352 else: 

353 formatted.append(f"Iteration {iteration}: {q}") 

354 # Apply limit if configured (main uses last 10) 

355 if self.previous_searches_limit: 

356 return "\n".join(formatted[-self.previous_searches_limit :]) 

357 return "\n".join(formatted) 

358 

359 def _was_searched(self, term: str, questions_by_iteration: dict) -> bool: 

360 """Check if a term was already searched.""" 

361 term_lower = term.lower() 

362 for questions in questions_by_iteration.values(): 

363 if isinstance(questions, list): 

364 for q in questions: 

365 if term_lower in q.lower(): 

366 return True 

367 return False