Coverage for src/local_deep_research/citation_handlers/precision_extraction_handler.py: 97%
213 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-22 21:20 +0000
1"""
2Precision Extraction Citation Handler
4This handler focuses on extracting precise, complete answers for SimpleQA-style questions.
5It includes specialized extractors for:
6- Full names (including middle names)
7- Single answers when only one is requested
8- Dimension-aware measurements
9- Specific entities without extra information
10"""
12import re
13from datetime import datetime, timezone
14from typing import Any, Dict, List, Union
16from loguru import logger
18from .base_citation_handler import BaseCitationHandler
21class PrecisionExtractionHandler(BaseCitationHandler):
22 """Citation handler optimized for precise answer extraction."""
24 # Keep the base's terse source format: regex extractors scan the source
25 # text as evidence, so titles and URLs could introduce false matches.
27 def __init__(self, *args, **kwargs):
28 super().__init__(*args, **kwargs)
30 # Answer type patterns
31 self.answer_patterns = {
32 "full_name": re.compile(
33 r"\b([A-Z][a-z]+(?:\s+[A-Z][a-z]+){1,4})\b"
34 ),
35 "year": re.compile(r"\b(19\d{2}|20\d{2})\b"),
36 "number": re.compile(r"\b(\d+(?:\.\d+)?)\b"),
37 "dimension": re.compile(
38 r"(\d+(?:\.\d+)?)\s*(meters?|feet|inches|cm|km|miles?|m|ft|kg|pounds?|lbs?)",
39 re.I,
40 ),
41 "score": re.compile(r"(\d+)\s*[-–]\s*(\d+)"),
42 "percentage": re.compile(r"(\d+(?:\.\d+)?)\s*%"),
43 "location": re.compile(r"\b([A-Z][a-z]+(?:\s+[A-Z][a-z]+)*)\b"),
44 }
46 def analyze_initial(
47 self, query: str, search_results: Union[str, List[Dict]]
48 ) -> Dict[str, Any]:
49 """Initial analysis with precision extraction."""
50 documents = self._create_documents(search_results)
51 formatted_sources = self._format_sources(documents)
53 # Determine question type for targeted extraction
54 question_type = self._identify_question_type(query)
56 current_timestamp = datetime.now(timezone.utc).strftime(
57 "%Y-%m-%d %H:%M"
58 )
60 output_prefix = self._get_output_instruction_prefix()
62 prompt = f"""{output_prefix}Analyze the following information and provide a PRECISE answer to the question. Include citations using numbers in square brackets [1], [2], etc.
64Question: {query}
65Question Type: {question_type}
67Sources:
68{formatted_sources}
70Current time is {current_timestamp} UTC for verifying temporal references in sources.
72PRECISION INSTRUCTIONS:
731. Extract the EXACT answer as it appears in the sources
742. For names: Include FULL names with all parts (first, middle, last)
753. For numbers: Include exact values with units if present
764. For single-answer questions: Provide ONLY ONE answer, not multiple options
775. For dimensions: Specify the exact measurement type (height, length, width)
786. Citations should support the specific answer given
80Format: Start with the direct, precise answer, then explain with citations."""
82 response = self._invoke_with_streaming(prompt)
84 # Apply precision extraction if needed
85 response = self._apply_precision_extraction(
86 response, query, question_type, formatted_sources
87 )
89 return {"content": response, "documents": documents}
91 def analyze_followup(
92 self,
93 question: str,
94 search_results: Union[str, List[Dict]],
95 previous_knowledge: str,
96 nr_of_links: int,
97 ) -> Dict[str, Any]:
98 """Follow-up analysis with precision extraction."""
99 documents = self._create_documents(
100 search_results, nr_of_links=nr_of_links
101 )
102 formatted_sources = self._format_sources(documents)
104 question_type = self._identify_question_type(question)
106 # Extract key facts from previous knowledge
107 key_facts = self._extract_key_facts(previous_knowledge, question_type)
109 current_timestamp = datetime.now(timezone.utc).strftime(
110 "%Y-%m-%d %H:%M"
111 )
113 output_prefix = self._get_output_instruction_prefix()
115 prompt = f"""{output_prefix}Using the previous knowledge and new sources, provide a PRECISE answer to the question.
117Previous Key Facts:
118{key_facts}
120Question: {question}
121Question Type: {question_type}
123New Sources:
124{formatted_sources}
126Current time is {current_timestamp} UTC for verifying temporal references in sources.
128PRECISION REQUIREMENTS:
1291. Build on previous knowledge to provide the MOST COMPLETE answer
1302. If a full name was partially found before, complete it now
1313. If multiple candidates exist, select the one with the MOST evidence
1324. For measurements, ensure units and dimension types match the question
1335. Reconcile any conflicts by choosing the most frequently cited answer
135Provide the precise answer with citations. Do not create the bibliography, it will be provided automatically."""
137 content = self._invoke_with_streaming(prompt)
139 # Apply precision extraction
140 content = self._apply_precision_extraction(
141 content, question, question_type, formatted_sources
142 )
144 return {"content": content, "documents": documents}
146 def _identify_question_type(self, query: str) -> str:
147 """Identify the type of question for targeted extraction."""
148 query_lower = query.lower()
150 # Name questions
151 if any(
152 phrase in query_lower
153 for phrase in ["full name", "name of", "who was", "who is"]
154 ):
155 if "full name" in query_lower:
156 return "full_name"
157 return "name"
159 # Location questions
160 if any(
161 phrase in query_lower
162 for phrase in ["where", "location", "city", "country", "place"]
163 ):
164 return "location"
166 # Temporal questions
167 if any(phrase in query_lower for phrase in ["when", "year", "date"]):
168 return "temporal"
170 # Numerical questions
171 if any(
172 phrase in query_lower
173 for phrase in ["how many", "how much", "number", "count"]
174 ):
175 return "number"
177 # Score/result questions
178 if any(
179 phrase in query_lower
180 for phrase in ["score", "result", "final", "outcome"]
181 ):
182 return "score"
184 # Dimension questions
185 if any(
186 phrase in query_lower
187 for phrase in [
188 "height",
189 "length",
190 "width",
191 "size",
192 "tall",
193 "long",
194 "wide",
195 ]
196 ):
197 return "dimension"
199 # Single answer questions
200 if query_lower.startswith("which") and "one" in query_lower:
201 return "single_choice"
203 return "general"
205 def _apply_precision_extraction(
206 self, content: str, query: str, question_type: str, sources: str
207 ) -> str:
208 """Apply precision extraction based on question type."""
210 if question_type == "full_name":
211 return self._extract_full_name(content, query, sources)
212 if question_type == "name":
213 return self._extract_best_name(content, query, sources)
214 if question_type == "single_choice":
215 return self._extract_single_answer(content, query, sources)
216 if question_type == "dimension":
217 return self._extract_dimension(content, query, sources)
218 if question_type == "score":
219 return self._extract_score(content, query, sources)
220 if question_type == "temporal":
221 return self._extract_temporal(content, query, sources)
222 if question_type == "number":
223 return self._extract_number(content, query, sources)
225 return content
227 def _extract_full_name(self, content: str, query: str, sources: str) -> str:
228 """Extract complete full names."""
229 # First, use LLM to identify all name variations
230 extraction_prompt = f"""Find ALL variations of the person's name mentioned in the sources.
232Question: {query}
234Content: {content[:2000]}
235Sources: {sources[:2000]}
237List all name variations found:
2381. Shortest version:
2392. Longest/most complete version:
2403. Most frequently mentioned version:
242Which is the FULL name (including middle name if present)?"""
244 try:
245 extraction = self._invoke_text(extraction_prompt)
246 extraction_lower = extraction.lower()
248 # Extract the identified full name
249 if "full name" in extraction_lower: 249 ↛ 265line 249 didn't jump to line 265 because the condition on line 249 was always true
250 lines = extraction.split("\n")
251 for line in lines:
252 if "full name" in line.lower() or "longest" in line.lower(): 252 ↛ 251line 252 didn't jump to line 251 because the condition on line 252 was always true
253 # Extract name from this line
254 matches = self.answer_patterns["full_name"].findall(
255 line
256 )
257 if matches:
258 # Choose the longest match
259 full_name = max(
260 matches, key=lambda x: len(x.split())
261 )
262 return f"{full_name}. {content}"
264 # Fallback: find all names and pick the longest
265 all_names = self.answer_patterns["full_name"].findall(
266 content + " " + sources
267 )
268 if all_names: 268 ↛ 290line 268 didn't jump to line 290 because the condition on line 268 was always true
269 # Group similar names and pick the longest variant
270 name_groups: Dict[str, List[str]] = {}
271 for name in all_names:
272 last_word = name.split()[-1]
273 if last_word not in name_groups:
274 name_groups[last_word] = []
275 name_groups[last_word].append(name)
277 # Find the group with the most complete name
278 best_name = ""
279 for group in name_groups.values():
280 longest_in_group = max(group, key=lambda x: len(x.split()))
281 if len(longest_in_group.split()) > len(best_name.split()): 281 ↛ 279line 281 didn't jump to line 279 because the condition on line 281 was always true
282 best_name = longest_in_group
284 if best_name: 284 ↛ 290line 284 didn't jump to line 290 because the condition on line 284 was always true
285 return f"{best_name}. {content}"
287 except Exception:
288 logger.exception("Error in full name extraction")
290 return content
292 def _extract_single_answer(
293 self, content: str, query: str, sources: str
294 ) -> str:
295 """Extract a single answer when multiple options might be present."""
296 extraction_prompt = f"""The question asks for ONE specific answer. Extract ONLY that answer.
298Question: {query}
299Content: {content[:1500]}
301Rules:
3021. If multiple items are listed, identify which ONE actually answers the question
3032. Look for the PRIMARY or FIRST mentioned item
3043. Do not include alternatives or additional options
306The single answer is:"""
308 try:
309 answer = self._invoke_text(extraction_prompt)
310 if not answer:
311 return content
313 # Clean up the answer
314 answer = answer.split(",")[
315 0
316 ].strip() # Take only first if comma-separated
317 answer = answer.split(" and ")[
318 0
319 ].strip() # Take only first if "and"-separated
320 answer = answer.split(" or ")[
321 0
322 ].strip() # Take only first if "or"-separated
324 return f"{answer}. {content}"
326 except Exception:
327 logger.exception("Error in single answer extraction")
329 return content
331 def _extract_dimension(self, content: str, query: str, sources: str) -> str:
332 """Extract specific dimensions with correct units and context awareness."""
333 # Enhanced dimension type detection
334 dimension_types = {
335 "height": ["height", "tall", "high", "elevation", "altitude"],
336 "length": ["length", "long", "distance", "reach", "span"],
337 "width": ["width", "wide", "breadth", "diameter"],
338 "depth": ["depth", "deep", "thickness"],
339 "weight": ["weight", "weigh", "heavy", "mass"],
340 "speed": ["speed", "fast", "velocity", "mph", "kmh"],
341 "area": ["area", "square"],
342 "volume": ["volume", "cubic"],
343 }
345 query_lower = query.lower()
346 dimension_type = None
347 dimension_keywords = []
349 # Find the most specific dimension type
350 for dim_type, keywords in dimension_types.items():
351 matching_keywords = [kw for kw in keywords if kw in query_lower]
352 if matching_keywords:
353 dimension_type = dim_type
354 dimension_keywords = matching_keywords
355 break
357 extraction_prompt = f"""Extract the EXACT measurement that answers this question.
359Question: {query}
360Content: {content[:1500]}
362Rules:
3631. Find the specific {dimension_type or "dimension"} measurement
3642. Return ONLY the number and unit (e.g., "20 meters", "5.5 feet")
3653. Distinguish between different types of measurements:
366 - Height/tall: vertical measurements
367 - Length/long: horizontal distance
368 - Width/wide: horizontal breadth
3694. Look for context clues near the measurement
3705. If multiple measurements, choose the one that matches the question type
372The exact {dimension_type or "dimension"} is:"""
374 try:
375 answer = self._invoke_text(extraction_prompt)
377 # Clean and validate the answer
378 import re
380 measurement_match = re.search(
381 r"(\d+(?:\.\d+)?)\s*([a-zA-Z/°]+)", answer
382 )
383 if measurement_match:
384 number, unit = measurement_match.groups()
385 clean_answer = f"{number} {unit}"
386 return f"{clean_answer}. {content}"
388 # Fallback: intelligent pattern matching
389 all_dimensions = self.answer_patterns["dimension"].findall(
390 content + " " + sources
391 )
392 if all_dimensions:
393 # Score dimensions based on context and dimension type
394 scored_dimensions = []
396 for dim in all_dimensions:
397 number, unit = dim
398 dim_str = f"{number} {unit}"
399 score = 0
401 # Find the dimension in content
402 pos = content.find(dim_str)
403 if pos >= 0: 403 ↛ 455line 403 didn't jump to line 455 because the condition on line 403 was always true
404 # Get context around this measurement
405 context = content[max(0, pos - 100) : pos + 100].lower()
407 # Score based on dimension keywords in context
408 for keyword in dimension_keywords:
409 if keyword in context:
410 score += 10
412 # Score based on unit appropriateness
413 unit_lower = unit.lower()
414 if (
415 (
416 dimension_type == "height"
417 and any(
418 u in unit_lower
419 for u in ["m", "meter", "ft", "feet", "cm"]
420 )
421 )
422 or (
423 dimension_type == "length"
424 and any(
425 u in unit_lower
426 for u in ["m", "meter", "km", "mile", "ft"]
427 )
428 )
429 or (
430 dimension_type == "weight"
431 and any(
432 u in unit_lower
433 for u in [
434 "kg",
435 "lb",
436 "pound",
437 "gram",
438 "ton",
439 ]
440 )
441 )
442 or (
443 dimension_type == "speed"
444 and any(
445 u in unit_lower
446 for u in ["mph", "kmh", "km/h", "m/s"]
447 )
448 )
449 ):
450 score += 5
452 # Prefer measurements closer to the beginning (more likely to be primary)
453 score += max(0, 5 - (pos // 100))
455 scored_dimensions.append((score, dim_str))
457 # Return the highest scoring dimension
458 if scored_dimensions: 458 ↛ 464line 458 didn't jump to line 464 because the condition on line 458 was always true
459 scored_dimensions.sort(key=lambda x: x[0], reverse=True)
460 best_dimension = scored_dimensions[0][1]
461 return f"{best_dimension}. {content}"
463 # Final fallback: first dimension
464 return (
465 f"{all_dimensions[0][0]} {all_dimensions[0][1]}. {content}"
466 )
468 except Exception:
469 logger.exception("Error in dimension extraction")
471 return content
473 def _extract_score(self, content: str, query: str, sources: str) -> str:
474 """Extract game scores or results."""
475 # Find all score patterns
476 scores = self.answer_patterns["score"].findall(content + " " + sources)
478 if scores:
479 # Use LLM to identify the correct score
480 extraction_prompt = f"""Which score/result answers this question?
482Question: {query}
483Found scores: {scores}
484Context: {content[:1000]}
486The answer is:"""
488 try:
489 answer = self._invoke_text(extraction_prompt)
490 if not answer:
491 logger.warning(
492 "Empty LLM answer in score extraction; "
493 "falling back to first matched score"
494 )
495 return f"{scores[0][0]}-{scores[0][1]}. {content}"
496 return f"{answer}. {content}"
497 except Exception:
498 # Return first score found if LLM extraction fails
499 return f"{scores[0][0]}-{scores[0][1]}. {content}"
501 return content
503 def _extract_temporal(self, content: str, query: str, sources: str) -> str:
504 """Extract dates or years."""
505 # Find all year patterns
506 years = self.answer_patterns["year"].findall(content + " " + sources)
508 if years:
509 # Use LLM to pick the right one
510 extraction_prompt = f"""Which date/year specifically answers this question?
512Question: {query}
513Found years: {set(years)}
514Context: {content[:1000]}
516The answer is:"""
518 try:
519 answer = self._invoke_text(extraction_prompt)
520 if not answer:
521 logger.warning(
522 "Empty LLM answer in temporal extraction; "
523 "falling back to first matched year"
524 )
525 return f"{years[0]}. {content}"
526 # Clean to just the year/date
527 year_match = self.answer_patterns["year"].search(answer)
528 if year_match:
529 return f"{year_match.group()}. {content}"
530 return f"{answer}. {content}"
531 except Exception:
532 # Fallback to first found year if LLM extraction fails
533 return f"{years[0]}. {content}"
535 return content
537 def _extract_number(self, content: str, query: str, sources: str) -> str:
538 """Extract specific numbers."""
539 # Find all numbers
540 numbers = self.answer_patterns["number"].findall(
541 content + " " + sources
542 )
544 if numbers:
545 extraction_prompt = f"""Which number specifically answers this question?
547Question: {query}
548Found numbers: {numbers[:10]}
549Context: {content[:1000]}
551The answer is:"""
553 try:
554 answer = self._invoke_text(extraction_prompt)
555 if not answer:
556 logger.warning(
557 "Empty LLM answer in number extraction; "
558 "falling back to first matched number"
559 )
560 return f"{numbers[0]}. {content}"
561 return f"{answer}. {content}"
562 except Exception:
563 # Fallback to first found number if LLM extraction fails
564 return f"{numbers[0]}. {content}"
566 return content
568 def _extract_best_name(self, content: str, query: str, sources: str) -> str:
569 """Extract the best matching name (not necessarily full)."""
570 # Find all potential names
571 names = self.answer_patterns["full_name"].findall(
572 content + " " + sources
573 )
575 if names:
576 # Count frequency
577 name_counts: Dict[str, int] = {}
578 for name in names:
579 name_counts[name] = name_counts.get(name, 0) + 1
581 # Get most frequent
582 best_name = max(name_counts.items(), key=lambda x: x[1])[0]
583 return f"{best_name}. {content}"
585 return content
587 def _extract_key_facts(
588 self, previous_knowledge: str, question_type: str
589 ) -> str:
590 """Extract key facts from previous knowledge."""
591 extraction_prompt = f"""Extract key facts related to a {question_type} question from this knowledge:
593{previous_knowledge[:1500]}
595List the most important facts (names, numbers, dates) found:"""
597 try:
598 facts = self._invoke_text(extraction_prompt)
599 return facts[:500]
600 except Exception:
601 # Fallback to truncated previous knowledge if LLM extraction fails
602 return previous_knowledge[:500]