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
« 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"""
5import re
6from typing import Dict, List, Optional
8from loguru import logger
10from ...utilities.json_utils import get_llm_response_text
11from .base_question import BaseQuestionGenerator
14class BrowseCompQuestionGenerator(BaseQuestionGenerator):
15 """
16 Question generator optimized for BrowseComp-style queries.
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 """
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.
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
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 {}
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 )
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 )
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:
78Query: {query}
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")
87For TEMPORAL entities, if there's a range (e.g., "between 2018-2023"), list EACH individual year.
89Format your response as:
90TEMPORAL: [entity1], [entity2], ...
91NUMERICAL: [entity1], [entity2], ...
92NAMES: [entity1], [entity2], ...
93LOCATIONS: [entity1], [entity2], ...
94DESCRIPTORS: [entity1], [entity2], ...
95"""
97 response = self.model.invoke(prompt)
98 content = get_llm_response_text(response)
100 entities: Dict[str, List[str]] = {
101 "temporal": [],
102 "numerical": [],
103 "names": [],
104 "locations": [],
105 "descriptors": [],
106 }
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)
118 # Expand temporal ranges
119 entities["temporal"] = self._expand_temporal_ranges(
120 entities["temporal"]
121 )
123 logger.info(f"Extracted entities: {entities}")
124 return entities
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)
149 return list(set(expanded)) # Remove duplicates
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 = []
157 # 1. Original query (always include)
158 searches.append(query)
160 # If only 1 question requested, return just the original query
161 if num_questions <= 1:
162 return searches[:1]
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]}")
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}")
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]}")
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
212 return unique_searches[:num_questions]
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."""
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
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"""
252 # Analyze what we've found so far
253 prompt = f"""Based on our search progress, generate targeted follow-up searches.
254{strategy_instruction}
256Original Query: {query}
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])}
264Current Knowledge Summary:
265{current_knowledge[: self.knowledge_truncate_length] if self.knowledge_truncate_length else current_knowledge}
267Previous Searches:
268{self._format_previous_searches(questions_by_iteration, results_by_iteration)}
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
276Focus on finding the specific answer, not general information.
278Format: One search per line
279"""
281 response = self.model.invoke(prompt)
282 content = get_llm_response_text(response)
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)
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
332 return searches[:num_questions]
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)
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