Coverage for src/local_deep_research/advanced_search_system/questions/browsecomp_question.py: 95%
153 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-22 21:20 +0000
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-22 21:20 +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
12from ...utilities.llm_utils import invoke_llm_sync
15class BrowseCompQuestionGenerator(BaseQuestionGenerator):
16 """
17 Question generator optimized for BrowseComp-style queries.
19 Key features:
20 1. Extract concrete entities (dates, numbers, names, places)
21 2. Generate progressive search combinations
22 3. Start broad, then narrow systematically
23 4. Focus on verifiable facts
24 """
26 def __init__(
27 self,
28 model,
29 knowledge_truncate_length: int = 1500,
30 previous_searches_limit: int = 10,
31 ):
32 """Initialize the question generator.
34 Args:
35 model: The LLM model to use for generation
36 knowledge_truncate_length: Max chars for knowledge in prompts (None=unlimited)
37 previous_searches_limit: Max previous searches to show (None=unlimited)
38 """
39 super().__init__(model)
40 self.extracted_entities: Dict[str, List[str]] = {}
41 self.search_progression: List[str] = []
42 self.knowledge_truncate_length = knowledge_truncate_length
43 self.previous_searches_limit = previous_searches_limit
45 def generate_questions(
46 self,
47 current_knowledge: str,
48 query: str,
49 questions_per_iteration: int = 5,
50 questions_by_iteration: Optional[dict] = None,
51 results_by_iteration: Optional[dict] = None,
52 iteration: int = 1,
53 ) -> List[str]:
54 """Generate progressive search queries for BrowseComp problems."""
55 questions_by_iteration = questions_by_iteration or {}
57 # First iteration: Extract entities and create initial searches
58 if iteration == 1 or not self.extracted_entities:
59 self.extracted_entities = self._extract_entities(query)
60 return self._generate_initial_searches(
61 query, self.extracted_entities, questions_per_iteration
62 )
64 # Subsequent iterations: Progressive refinement
65 return self._generate_progressive_searches(
66 query,
67 current_knowledge,
68 self.extracted_entities,
69 questions_by_iteration,
70 results_by_iteration or {},
71 questions_per_iteration,
72 iteration,
73 )
75 def _extract_entities(self, query: str) -> Dict[str, List[str]]:
76 """Extract concrete entities from the query."""
77 prompt = f"""Extract ALL concrete, searchable entities from this query:
79Query: {query}
81Extract:
821. TEMPORAL: All years, dates, time periods (e.g., "2018", "between 1995 and 2006", "2023")
832. NUMERICAL: All numbers, statistics, counts (e.g., "300", "more than 3", "4-3", "84.5%")
843. NAMES: Partial names, name hints, proper nouns (e.g., "Dartmouth", "EMNLP", "Plastic Man")
854. LOCATIONS: Places, institutions, geographic features (e.g., "Pennsylvania", "Grand Canyon")
865. DESCRIPTORS: Key descriptive terms (e.g., "fourth wall", "ascetics", "decider game")
88For TEMPORAL entities, if there's a range (e.g., "between 2018-2023"), list EACH individual year.
90Format your response as:
91TEMPORAL: [entity1], [entity2], ...
92NUMERICAL: [entity1], [entity2], ...
93NAMES: [entity1], [entity2], ...
94LOCATIONS: [entity1], [entity2], ...
95DESCRIPTORS: [entity1], [entity2], ...
96"""
98 response = invoke_llm_sync(self.model, prompt)
99 content = get_llm_response_text(response)
101 entities: Dict[str, List[str]] = {
102 "temporal": [],
103 "numerical": [],
104 "names": [],
105 "locations": [],
106 "descriptors": [],
107 }
109 for line in content.strip().split("\n"):
110 line = line.strip()
111 if ":" in line:
112 category, values = line.split(":", 1)
113 category = category.strip().lower()
114 if category in entities:
115 # Parse comma-separated values
116 values = [v.strip() for v in values.split(",") if v.strip()]
117 entities[category].extend(values)
119 # Expand temporal ranges
120 entities["temporal"] = self._expand_temporal_ranges(
121 entities["temporal"]
122 )
124 logger.info(f"Extracted entities: {entities}")
125 return entities
127 def _expand_temporal_ranges(
128 self, temporal_entities: List[str]
129 ) -> List[str]:
130 """Expand year ranges into individual years."""
131 expanded = []
132 for entity in temporal_entities:
133 # Check for range patterns like "2018-2023" or "between 1995 and 2006"
134 range_match = re.search(
135 r"(\d{4})[-\s]+(?:to|and)?\s*(\d{4})", entity
136 )
137 if range_match:
138 start_year = int(range_match.group(1))
139 end_year = int(range_match.group(2))
140 for year in range(start_year, end_year + 1):
141 expanded.append(str(year))
142 else:
143 # Single year or other temporal entity
144 year_match = re.search(r"\d{4}", entity)
145 if year_match:
146 expanded.append(year_match.group())
147 else:
148 expanded.append(entity)
150 return list(set(expanded)) # Remove duplicates
152 def _generate_initial_searches(
153 self, query: str, entities: Dict[str, List[str]], num_questions: int
154 ) -> List[str]:
155 """Generate initial broad searches."""
156 searches = []
158 # 1. Original query (always include)
159 searches.append(query)
161 # If only 1 question requested, return just the original query
162 if num_questions <= 1:
163 return searches[:1]
165 # 2. Domain exploration searches (combine key entities)
166 if entities["names"] and len(searches) < num_questions:
167 for name in entities["names"][:2]: # Top 2 names
168 if len(searches) >= num_questions:
169 break
170 searches.append(f"{name}")
171 if entities["descriptors"] and len(searches) < num_questions:
172 searches.append(f"{name} {entities['descriptors'][0]}")
174 # 3. Temporal searches if years are important
175 if (
176 entities["temporal"]
177 and len(entities["temporal"]) <= 10
178 and len(searches) < num_questions
179 ):
180 # For small year ranges, search each year with a key term
181 key_term = (
182 entities["names"][0]
183 if entities["names"]
184 else entities["descriptors"][0]
185 if entities["descriptors"]
186 else ""
187 )
188 for year in entities["temporal"][:5]: # Limit to 5 years initially
189 if len(searches) >= num_questions: 189 ↛ 190line 189 didn't jump to line 190 because the condition on line 189 was never true
190 break
191 if key_term: 191 ↛ 188line 191 didn't jump to line 188 because the condition on line 191 was always true
192 searches.append(f"{key_term} {year}")
194 # 4. Location-based searches
195 if entities["locations"] and len(searches) < num_questions:
196 for location in entities["locations"][:2]:
197 if len(searches) >= num_questions: 197 ↛ 198line 197 didn't jump to line 198 because the condition on line 197 was never true
198 break
199 searches.append(f"{location}")
200 if entities["descriptors"] and len(searches) < num_questions:
201 searches.append(f"{location} {entities['descriptors'][0]}")
203 # Remove duplicates and limit to requested number
204 seen = set()
205 unique_searches = []
206 for s in searches:
207 if s.lower() not in seen:
208 seen.add(s.lower())
209 unique_searches.append(s)
210 if len(unique_searches) >= num_questions:
211 break
213 return unique_searches[:num_questions]
215 def _generate_progressive_searches(
216 self,
217 query: str,
218 current_knowledge: str,
219 entities: Dict[str, List[str]],
220 questions_by_iteration: dict,
221 results_by_iteration: dict,
222 num_questions: int,
223 iteration: int,
224 ) -> List[str]:
225 """Generate progressively more specific searches based on findings."""
227 # Only add strategy instructions if we have actual result data (adaptive mode)
228 strategy_instruction = ""
229 if results_by_iteration:
230 # Check if recent searches are failing (returning 0 results)
231 recent_iterations = list(range(max(1, iteration - 5), iteration))
232 zero_count = sum(
233 1
234 for i in recent_iterations
235 if results_by_iteration.get(i, 1) == 0
236 )
237 searches_failing = zero_count >= 3
239 # Adjust strategy based on success/failure
240 if searches_failing:
241 strategy_instruction = """
242IMPORTANT: Your recent searches are returning 0 results - they are TOO NARROW!
243- Use FEWER constraints (1-2 terms instead of 4-5)
244- Try BROADER, more general searches
245- Remove overly specific combinations
246- Focus on key concepts, not detailed entity combinations
247"""
248 else:
249 strategy_instruction = """
250Focus on finding the specific answer by combining entities systematically.
251"""
253 # Analyze what we've found so far
254 prompt = f"""Based on our search progress, generate targeted follow-up searches.
255{strategy_instruction}
257Original Query: {query}
259Entities Found:
260- Names/Terms: {", ".join(entities["names"][:5])}
261- Years: {", ".join(entities["temporal"][:5])}
262- Locations: {", ".join(entities["locations"][:3])}
263- Key Features: {", ".join(entities["descriptors"][:3])}
265Current Knowledge Summary:
266{current_knowledge[: self.knowledge_truncate_length] if self.knowledge_truncate_length else current_knowledge}
268Previous Searches:
269{self._format_previous_searches(questions_by_iteration, results_by_iteration)}
271Generate {num_questions} NEW search queries that:
2721. Combine 2-3 entities we haven't tried together
2732. If we found candidate names, search for them with other constraints
2743. For year ranges, systematically cover years we haven't searched
2754. Use quotes for exact phrases when beneficial
277Focus on finding the specific answer, not general information.
279Format: One search per line
280"""
282 response = invoke_llm_sync(self.model, prompt)
283 content = get_llm_response_text(response)
285 # Extract searches from response
286 searches = []
287 for line in content.strip().split("\n"):
288 line = line.strip()
289 if line and not line.endswith(":") and len(line) > 5: 289 ↛ 287line 289 didn't jump to line 287 because the condition on line 289 was always true
290 # Clean up common prefixes
291 for prefix in ["Q:", "Search:", "-", "*", "•"]:
292 if line.startswith(prefix):
293 line = line[len(prefix) :].strip()
294 if line: 294 ↛ 287line 294 didn't jump to line 287 because the condition on line 294 was always true
295 searches.append(line)
297 # Ensure we have enough searches, but respect the limit
298 while len(searches) < num_questions:
299 # Generate combinations programmatically
300 if iteration <= 5 and entities["temporal"]:
301 # Continue with year-based searches
302 added_any = False
303 for year in entities["temporal"]: 303 ↛ 312line 303 didn't jump to line 312 because the loop on line 303 didn't complete
304 if not self._was_searched(year, questions_by_iteration): 304 ↛ 303line 304 didn't jump to line 303 because the condition on line 304 was always true
305 base_term = (
306 entities["names"][0] if entities["names"] else ""
307 )
308 searches.append(f"{base_term} {year}".strip())
309 added_any = True
310 if len(searches) >= num_questions:
311 break
312 if not added_any: 312 ↛ 313line 312 didn't jump to line 313 because the condition on line 312 was never true
313 break # No more year searches to add
314 else:
315 # Combine multiple constraints
316 added_any = False
317 if entities["names"] and entities["descriptors"]:
318 for name in entities["names"]:
319 for desc in entities["descriptors"]:
320 combo = f"{name} {desc}"
321 if not self._was_searched(
322 combo, questions_by_iteration
323 ):
324 searches.append(combo)
325 added_any = True
326 if len(searches) >= num_questions:
327 break
328 if len(searches) >= num_questions:
329 break
330 if not added_any:
331 break # No more combinations to add
333 return searches[:num_questions]
335 def _format_previous_searches(
336 self, questions_by_iteration: dict, results_by_iteration: dict
337 ) -> str:
338 """Format previous searches for context with result counts."""
339 formatted = []
340 for iteration, questions in questions_by_iteration.items():
341 if isinstance(questions, list): 341 ↛ 340line 341 didn't jump to line 340 because the condition on line 341 was always true
342 result_count = results_by_iteration.get(iteration, "?")
343 # Limit questions per iteration (main uses 3)
344 questions_to_show = (
345 questions[:3] if self.previous_searches_limit else questions
346 )
347 for q in questions_to_show:
348 # Only show result counts if we have actual data (not "?")
349 if result_count != "?":
350 formatted.append(
351 f"Iteration {iteration}: {q} ({result_count} results)"
352 )
353 else:
354 formatted.append(f"Iteration {iteration}: {q}")
355 # Apply limit if configured (main uses last 10)
356 if self.previous_searches_limit:
357 return "\n".join(formatted[-self.previous_searches_limit :])
358 return "\n".join(formatted)
360 def _was_searched(self, term: str, questions_by_iteration: dict) -> bool:
361 """Check if a term was already searched."""
362 term_lower = term.lower()
363 for questions in questions_by_iteration.values():
364 if isinstance(questions, list):
365 for q in questions:
366 if term_lower in q.lower():
367 return True
368 return False