Coverage for presidio_analyzer / lm_recognizer.py: 98%

1import logging 

2from abc import ABC, abstractmethod 

3from typing import TYPE_CHECKING, List, Optional 

4 

5from presidio_analyzer import RecognizerResult, RemoteRecognizer 

6from presidio_analyzer.llm_utils import ( 

7 consolidate_generic_entities, 

8 ensure_generic_entity_support, 

9 filter_results_by_entities, 

10 filter_results_by_labels, 

11 filter_results_by_score, 

12 skip_unmapped_entities, 

13 validate_result_positions, 

14) 

15 

16if TYPE_CHECKING: 

17 from presidio_analyzer.nlp_engine import NlpArtifacts 

18 

19logger = logging.getLogger("presidio-analyzer") 

20 

21 

22class LMRecognizer(RemoteRecognizer, ABC): 

23 """ 

24 Base class for language model-based PII recognizers. 

25 

26 Provides common functionality for LLM-based entity detection. 

27 Subclasses implement _call_llm() for specific LLM providers. 

28 """ 

29 

30 def __init__( 

31 self, 

32 supported_entities: Optional[List[str]] = None, 

33 supported_language: str = "en", 

34 name: Optional[str] = None, 

35 version: str = "1.0.0", 

36 model_id: Optional[str] = None, 

37 temperature: Optional[float] = None, 

38 min_score: float = 0.5, 

39 labels_to_ignore: Optional[List[str]] = None, 

40 enable_generic_consolidation: bool = True, 

41 ): 

42 """Initialize LM recognizer. 

43 

44 :param supported_entities: Entity types to detect. 

45 :param labels_to_ignore: Entity labels to skip. 

46 :param enable_generic_consolidation: Consolidate unknown 

47 entities to GENERIC_PII_ENTITY. 

48 """ 

49 if not supported_entities: 

50 raise ValueError( 

51 "LMRecognizer requires at least one entity in 'supported_entities'" 

52 ) 

53 

54 super().__init__( 

55 supported_entities=supported_entities, 

56 supported_language=supported_language, 

57 name=name, 

58 version=version, 

59 ) 

60 

61 self.model_id = model_id 

62 self.temperature = temperature 

63 self.min_score = min_score 

64 self.labels_to_ignore = [label.lower() for label in (labels_to_ignore or [])] 

65 self.enable_generic_consolidation = enable_generic_consolidation 

66 

67 self._generic_entities_logged = set() 

68 

69 self.supported_entities = ensure_generic_entity_support( 

70 self.supported_entities, enable_generic_consolidation 

71 ) 

72 

73 @abstractmethod 

74 def _call_llm( 

75 self, 

76 text: str, 

77 entities: List[str], 

78 **kwargs 

79 ) -> List[RecognizerResult]: 

80 """ 

81 Call LLM service and return RecognizerResult objects. 

82 

83 Subclasses implement this to integrate with specific LLM providers. 

84 

85 :param text: Text to analyze for PII. 

86 :param entities: Entity types to detect. 

87 :return: List of RecognizerResult objects. 

88 """ 

89 ... 

90 

91 def analyze( 

92 self, 

93 text: str, 

94 entities: Optional[List[str]] = None, 

95 nlp_artifacts: Optional["NlpArtifacts"] = None 

96 ) -> List[RecognizerResult]: 

97 """Analyze text for PII/PHI using LLM.""" 

98 if not text or not text.strip(): 

99 logger.debug("Empty text provided, returning empty results") 

100 return [] 

101 

102 if entities is None: 

103 requested_entities = self.supported_entities 

104 else: 

105 requested_entities = [e for e in entities if e in self.supported_entities] 

106 

107 if not requested_entities: 

108 logger.debug( 

109 "No requested entities (%s) match supported entities (%s)", 

110 entities, self.supported_entities 

111 ) 

112 return [] 

113 

114 results = self._call_llm(text, requested_entities) 

115 

116 filtered_results = self._filter_and_process_results( 

117 results, requested_entities 

118 ) 

119 

120 if filtered_results: 

121 logger.debug( 

122 "LLM recognizer found %d entities", 

123 len(filtered_results), 

124 ) 

125 

126 return filtered_results 

127 

128 def _filter_and_process_results( 

129 self, 

130 results: List[RecognizerResult], 

131 requested_entities: Optional[List[str]] = None 

132 ) -> List[RecognizerResult]: 

133 """Filter and process results.""" 

134 filtered_results = filter_results_by_labels(results, self.labels_to_ignore) 

135 

136 if self.enable_generic_consolidation: 

137 filtered_results = consolidate_generic_entities( 

138 filtered_results, 

139 self.supported_entities, 

140 self._generic_entities_logged 

141 ) 

142 else: 

143 filtered_results = skip_unmapped_entities( 

144 filtered_results, 

145 self.supported_entities 

146 ) 

147 

148 if requested_entities: 

149 filtered_results = filter_results_by_entities( 

150 filtered_results, 

151 requested_entities 

152 ) 

153 

154 filtered_results = validate_result_positions(filtered_results) 

155 filtered_results = filter_results_by_score(filtered_results, self.min_score) 

156 

157 return filtered_results 

158 

159 def get_supported_entities(self) -> List[str]: 

160 """Return list of supported PII entity types.""" 

161 return self.supported_entities