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 

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 set_stream_callback(self, callback): 

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

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

92 self._handler.set_stream_callback(callback) 

93 

94 def analyze_initial( 

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

96 ) -> Dict[str, Any]: 

97 """Delegate to the configured handler.""" 

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

99 

100 def analyze_followup( 

101 self, 

102 question: str, 

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

104 previous_knowledge: str, 

105 nr_of_links: int, 

106 ) -> Dict[str, Any]: 

107 """Delegate to the configured handler.""" 

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

109 question, search_results, previous_knowledge, nr_of_links 

110 )