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 if 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 if 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 logger.warning( 

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

80 ) 

81 from .citation_handlers.standard_citation_handler import ( 

82 StandardCitationHandler, 

83 ) 

84 

85 return StandardCitationHandler( 

86 self.llm, settings_snapshot=self.settings_snapshot 

87 ) 

88 

89 def analyze_initial( 

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

91 ) -> Dict[str, Any]: 

92 """Delegate to the configured handler.""" 

93 return self._handler.analyze_initial(query, search_results) # type: ignore[no-any-return] 

94 

95 def analyze_followup( 

96 self, 

97 question: str, 

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

99 previous_knowledge: str, 

100 nr_of_links: int, 

101 ) -> Dict[str, Any]: 

102 """Delegate to the configured handler.""" 

103 return self._handler.analyze_followup( # type: ignore[no-any-return] 

104 question, search_results, previous_knowledge, nr_of_links 

105 )