Coverage for src / local_deep_research / citation_handler.py: 100%

1# citation_handler.py 

2 

3from typing import Any, Dict, List, Optional, Union 

4 

5from loguru import logger 

6 

7 

8class CitationHandler: 

9 """ 

10 Configurable citation handler that delegates to specific implementations. 

11 Maintains backward compatibility while allowing strategy-specific handlers. 

12 """ 

13 

14 def __init__( 

15 self, llm, handler_type: Optional[str] = None, settings_snapshot=None 

16 ): 

17 self.llm = llm 

18 self.settings_snapshot = settings_snapshot or {} 

19 

20 # Determine which handler to use 

21 if handler_type is None: 

22 # Try to get from settings snapshot, default to standard 

23 if "citation.handler_type" in self.settings_snapshot: 

24 value = self.settings_snapshot["citation.handler_type"] 

25 handler_type = ( 

26 value["value"] 

27 if isinstance(value, dict) and "value" in value 

28 else value 

29 ) 

30 else: 

31 handler_type = "standard" 

32 

33 # Import and instantiate the appropriate handler 

34 self._handler = self._create_handler(handler_type) 

35 

36 # For backward compatibility, expose internal methods 

37 self._create_documents = self._handler._create_documents 

38 self._format_sources = self._handler._format_sources 

39 

40 def _create_handler(self, handler_type: str): 

41 """Create the appropriate citation handler based on type.""" 

42 handler_type = handler_type.lower() 

43 

44 if handler_type == "standard": 

45 from .citation_handlers.standard_citation_handler import ( 

46 StandardCitationHandler, 

47 ) 

48 

49 logger.info("Using StandardCitationHandler") 

50 return StandardCitationHandler( 

51 self.llm, settings_snapshot=self.settings_snapshot 

52 ) 

53 

54 elif handler_type in ["forced", "forced_answer", "browsecomp"]: 

55 from .citation_handlers.forced_answer_citation_handler import ( 

56 ForcedAnswerCitationHandler, 

57 ) 

58 

59 logger.info( 

60 "Using ForcedAnswerCitationHandler for better benchmark performance" 

61 ) 

62 return ForcedAnswerCitationHandler( 

63 self.llm, settings_snapshot=self.settings_snapshot 

64 ) 

65 

66 elif handler_type in ["precision", "precision_extraction", "simpleqa"]: 

67 from .citation_handlers.precision_extraction_handler import ( 

68 PrecisionExtractionHandler, 

69 ) 

70 

71 logger.info( 

72 "Using PrecisionExtractionHandler for precise answer extraction" 

73 ) 

74 return PrecisionExtractionHandler( 

75 self.llm, settings_snapshot=self.settings_snapshot 

76 ) 

77 

78 else: 

79 logger.warning( 

80 f"Unknown citation handler type: {handler_type}, falling back to standard" 

81 ) 

82 from .citation_handlers.standard_citation_handler import ( 

83 StandardCitationHandler, 

84 ) 

85 

86 return StandardCitationHandler( 

87 self.llm, settings_snapshot=self.settings_snapshot 

88 ) 

89 

90 def analyze_initial( 

91 self, query: str, search_results: Union[str, List[Dict]] 

92 ) -> Dict[str, Any]: 

93 """Delegate to the configured handler.""" 

94 return self._handler.analyze_initial(query, search_results) 

95 

96 def analyze_followup( 

97 self, 

98 question: str, 

99 search_results: Union[str, List[Dict]], 

100 previous_knowledge: str, 

101 nr_of_links: int, 

102 ) -> Dict[str, Any]: 

103 """Delegate to the configured handler.""" 

104 return self._handler.analyze_followup( 

105 question, search_results, previous_knowledge, nr_of_links 

106 )