Coverage for presidio_analyzer / llm_utils / entity_mapper.py: 100%

1"""Entity mapping and filtering utilities for LLM recognizers.""" 

2import logging 

3from typing import List, Set 

4 

5from presidio_analyzer import RecognizerResult 

6 

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

8 

9GENERIC_PII_ENTITY = "GENERIC_PII_ENTITY" 

10 

11__all__ = [ 

12 "GENERIC_PII_ENTITY", 

13 "filter_results_by_labels", 

14 "filter_results_by_score", 

15 "filter_results_by_entities", 

16 "validate_result_positions", 

17 "consolidate_generic_entities", 

18 "skip_unmapped_entities", 

19 "ensure_generic_entity_support", 

20] 

21 

22 

23def filter_results_by_labels( 

24 results: List[RecognizerResult], 

25 labels_to_ignore: List[str] 

26) -> List[RecognizerResult]: 

27 """Filter out results with ignored entity labels. 

28 

29 :param results: List of recognizer results to filter. 

30 :param labels_to_ignore: Entity type labels to exclude (case-insensitive). 

31 :return: Filtered list of results without ignored labels. 

32 """ 

33 labels_to_ignore_lower = [label.lower() for label in labels_to_ignore] 

34 filtered = [] 

35 

36 for result in results: 

37 if not result.entity_type: 

38 logger.warning("LLM returned result without entity_type, skipping") 

39 continue 

40 

41 if result.entity_type.lower() in labels_to_ignore_lower: 

42 logger.debug( 

43 "Entity %s at [%d:%d] is in labels_to_ignore, skipping", 

44 result.entity_type, result.start, result.end 

45 ) 

46 continue 

47 

48 filtered.append(result) 

49 

50 return filtered 

51 

52 

53def filter_results_by_score( 

54 results: List[RecognizerResult], 

55 min_score: float 

56) -> List[RecognizerResult]: 

57 """Filter out results below minimum confidence score. 

58 

59 :param results: List of recognizer results to filter. 

60 :param min_score: Minimum confidence score threshold (0.0-1.0). 

61 :return: Filtered list of results meeting minimum score. 

62 """ 

63 filtered = [] 

64 

65 for result in results: 

66 if result.score < min_score: 

67 logger.debug( 

68 "Entity %s at [%d:%d] below min_score (%.2f < %.2f), skipping", 

69 result.entity_type, result.start, result.end, 

70 result.score, min_score 

71 ) 

72 continue 

73 

74 filtered.append(result) 

75 

76 return filtered 

77 

78 

79def filter_results_by_entities( 

80 results: List[RecognizerResult], 

81 requested_entities: List[str] 

82) -> List[RecognizerResult]: 

83 """Filter results to only include requested entity types. 

84 

85 :param results: List of recognizer results to filter. 

86 :param requested_entities: Entity types to include (empty list = include all). 

87 :return: Filtered list containing only requested entity types. 

88 """ 

89 if not requested_entities: 

90 return results 

91 

92 filtered = [] 

93 

94 for result in results: 

95 if result.entity_type not in requested_entities: 

96 logger.debug( 

97 "Entity %s at [%d:%d] not in requested entities %s, skipping", 

98 result.entity_type, result.start, result.end, requested_entities 

99 ) 

100 continue 

101 

102 filtered.append(result) 

103 

104 return filtered 

105 

106 

107def validate_result_positions( 

108 results: List[RecognizerResult] 

109) -> List[RecognizerResult]: 

110 """Filter out results with invalid or missing start/end positions. 

111 

112 :param results: List of recognizer results to validate. 

113 :return: Filtered list with only valid position ranges. 

114 """ 

115 filtered = [] 

116 

117 for result in results: 

118 if result.start is None or result.end is None: 

119 logger.warning( 

120 "LLM returned result without start/end positions, skipping: %s", result 

121 ) 

122 continue 

123 

124 filtered.append(result) 

125 

126 return filtered 

127 

128 

129def consolidate_generic_entities( 

130 results: List[RecognizerResult], 

131 supported_entities: List[str], 

132 generic_entities_logged: Set[str] 

133) -> List[RecognizerResult]: 

134 """Consolidate unmapped entity types to GENERIC_PII_ENTITY. 

135 

136 :param results: List of recognizer results to process. 

137 :param supported_entities: List of supported entity type names. 

138 :param generic_entities_logged: Set tracking logged generic entities 

139 (modified in-place). 

140 :return: Results with unmapped entities consolidated to 

141 GENERIC_PII_ENTITY. 

142 """ 

143 processed = [] 

144 

145 for result in results: 

146 if result.entity_type not in supported_entities: 

147 original_entity_type = result.entity_type 

148 result.entity_type = GENERIC_PII_ENTITY 

149 

150 if original_entity_type not in generic_entities_logged: 

151 logger.warning( 

152 "Detected unmapped entity '%s', " 

153 "consolidated to GENERIC_PII_ENTITY. " 

154 "To map or exclude, update " 

155 "'entity_mappings' or 'labels_to_ignore'.", 

156 original_entity_type, 

157 ) 

158 generic_entities_logged.add(original_entity_type) 

159 

160 if result.recognition_metadata is None: 

161 result.recognition_metadata = {} 

162 result.recognition_metadata["original_entity_type"] = original_entity_type 

163 

164 processed.append(result) 

165 

166 return processed 

167 

168 

169def skip_unmapped_entities( 

170 results: List[RecognizerResult], 

171 supported_entities: List[str] 

172) -> List[RecognizerResult]: 

173 """Skip unmapped entities instead of consolidating them. 

174 

175 :param results: List of recognizer results to filter. 

176 :param supported_entities: List of supported entity type names. 

177 :return: Filtered results excluding unmapped entity types. 

178 """ 

179 filtered = [] 

180 

181 for result in results: 

182 if result.entity_type not in supported_entities: 

183 logger.warning( 

184 "Detected unmapped entity '%s', skipped " 

185 "(enable_generic_consolidation=False). " 

186 "To map or exclude, update " 

187 "'entity_mappings' or 'labels_to_ignore'.", 

188 result.entity_type, 

189 ) 

190 continue 

191 

192 filtered.append(result) 

193 

194 return filtered 

195 

196 

197def ensure_generic_entity_support( 

198 supported_entities: List[str], 

199 enable_generic_consolidation: bool 

200) -> List[str]: 

201 """Ensure GENERIC_PII_ENTITY is in supported entities list if consolidation enabled. 

202 

203 :param supported_entities: Current list of supported entity types. 

204 :param enable_generic_consolidation: Whether generic consolidation is enabled. 

205 :return: Updated list including GENERIC_PII_ENTITY if needed. 

206 """ 

207 entities = supported_entities.copy() 

208 

209 if enable_generic_consolidation and GENERIC_PII_ENTITY not in entities: 

210 entities.append(GENERIC_PII_ENTITY) 

211 

212 return entities