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

1# citation_handler.py 

2 

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

4 

5from loguru import logger 

6 

7from .utilities.type_utils import unwrap_setting 

8 

9 

10class CitationHandler: 

11 """ 

12 Configurable citation handler that delegates to specific implementations. 

13 Maintains backward compatibility while allowing strategy-specific handlers. 

14 """ 

15 

16 def __init__( 

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

18 ): 

19 self.llm = llm 

20 self.settings_snapshot = settings_snapshot or {} 

21 

22 # Determine which handler to use 

23 if handler_type is None: 

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

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

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

27 handler_type = unwrap_setting(value) 

28 else: 

29 handler_type = "standard" 

30 

31 # Import and instantiate the appropriate handler 

32 self._handler = self._create_handler(handler_type) 

33 

34 # For backward compatibility, expose internal methods 

35 self._create_documents = self._handler._create_documents 

36 self._format_sources = self._handler._format_sources 

37 

38 def _create_handler(self, handler_type: str): 

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

40 handler_type = handler_type.lower() 

41 

42 if handler_type == "standard": 

43 from .citation_handlers.standard_citation_handler import ( 

44 StandardCitationHandler, 

45 ) 

46 

47 logger.info("Using StandardCitationHandler") 

48 return StandardCitationHandler( 

49 self.llm, settings_snapshot=self.settings_snapshot 

50 ) 

51 

52 if handler_type in ["forced", "forced_answer", "browsecomp"]: 

53 from .citation_handlers.forced_answer_citation_handler import ( 

54 ForcedAnswerCitationHandler, 

55 ) 

56 

57 logger.info( 

58 "Using ForcedAnswerCitationHandler for better benchmark performance" 

59 ) 

60 return ForcedAnswerCitationHandler( 

61 self.llm, settings_snapshot=self.settings_snapshot 

62 ) 

63 

64 if handler_type in ["precision", "precision_extraction", "simpleqa"]: 

65 from .citation_handlers.precision_extraction_handler import ( 

66 PrecisionExtractionHandler, 

67 ) 

68 

69 logger.info( 

70 "Using PrecisionExtractionHandler for precise answer extraction" 

71 ) 

72 return PrecisionExtractionHandler( 

73 self.llm, settings_snapshot=self.settings_snapshot 

74 ) 

75 

76 logger.warning( 

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

78 ) 

79 from .citation_handlers.standard_citation_handler import ( 

80 StandardCitationHandler, 

81 ) 

82 

83 return StandardCitationHandler( 

84 self.llm, settings_snapshot=self.settings_snapshot 

85 ) 

86 

87 def set_stream_callback(self, callback): 

88 """Set a streaming callback on the underlying handler.""" 

89 if hasattr(self._handler, "set_stream_callback"): 89 ↛ exitline 89 didn't return from function 'set_stream_callback' because the condition on line 89 was always true

90 self._handler.set_stream_callback(callback) 

91 

92 def analyze_initial( 

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

94 ) -> Dict[str, Any]: 

95 """Delegate to the configured handler.""" 

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

97 

98 def analyze_followup( 

99 self, 

100 question: str, 

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

102 previous_knowledge: str, 

103 nr_of_links: int, 

104 ) -> Dict[str, Any]: 

105 """Delegate to the configured handler.""" 

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

107 question, search_results, previous_knowledge, nr_of_links 

108 )