Coverage for presidio_analyzer / llm_utils / langextract_helper.py: 89%

1"""LangExtract helper utilities.""" 

2import logging 

3from typing import Dict, List, Optional 

4 

5from presidio_analyzer import AnalysisExplanation, RecognizerResult 

6 

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

8 

9try: 

10 import langextract as lx 

11 import langextract.factory as lx_factory 

12 

13 # Make sure builtin providers are pre-loaded to work around a bug in LangExtract 

14 # that fails to load them automatically if the provider name is specified for a 

15 # built-in provider rather than inferred from the model_id. 

16 lx.providers.load_builtins_once() 

17 lx.providers.load_plugins_once() 

18except ImportError: 

19 lx = None 

20 lx_factory = None 

21 

22__all__ = [ 

23 "lx", 

24 "lx_factory", 

25 "check_langextract_available", 

26 "extract_lm_config", 

27 "get_supported_entities", 

28 "create_reverse_entity_mapping", 

29 "calculate_extraction_confidence", 

30 "convert_langextract_to_presidio_results", 

31] 

32 

33 

34def check_langextract_available(): 

35 """Check if langextract is available and raise error if not.""" 

36 if not lx: 

37 raise ImportError( 

38 "LangExtract is not installed. " 

39 "Install it with: poetry install --extras langextract" 

40 ) 

41 

42 

43# Default alignment score mappings for LangExtract extractions 

44DEFAULT_ALIGNMENT_SCORES = { 

45 "MATCH_EXACT": 0.95, 

46 "MATCH_FUZZY": 0.80, 

47 "MATCH_LESSER": 0.70, 

48 "NOT_ALIGNED": 0.60, 

49} 

50 

51 

52def extract_lm_config(config: Dict) -> Dict: 

53 """Extract LM recognizer configuration section with default values. 

54 

55 :param config: Full configuration dictionary. 

56 :return: LM recognizer config with keys: supported_entities, min_score, 

57 labels_to_ignore, enable_generic_consolidation. 

58 """ 

59 lm_config_section = config.get("lm_recognizer", {}) 

60 

61 return { 

62 "supported_entities": lm_config_section.get("supported_entities"), 

63 "min_score": lm_config_section.get("min_score", 0.5), 

64 "labels_to_ignore": lm_config_section.get("labels_to_ignore", []), 

65 "enable_generic_consolidation": lm_config_section.get( 

66 "enable_generic_consolidation", True 

67 ), 

68 } 

69 

70 

71def get_supported_entities( 

72 lm_config: Dict, 

73 langextract_config: Dict 

74) -> Optional[List[str]]: 

75 """Get supported entities list, checking LM config first then LangExtract config. 

76 

77 :param lm_config: LM recognizer configuration dictionary. 

78 :param langextract_config: LangExtract configuration dictionary. 

79 :return: List of supported entity types, or None if not specified. 

80 """ 

81 return ( 

82 lm_config.get("supported_entities") 

83 or langextract_config.get("supported_entities") 

84 ) 

85 

86 

87def create_reverse_entity_mapping(entity_mappings: Dict) -> Dict: 

88 """Create reverse mapping from values to keys. 

89 

90 :param entity_mappings: Original entity mapping dictionary. 

91 :return: Reversed dictionary mapping values to keys. 

92 """ 

93 return {v: k for k, v in entity_mappings.items()} 

94 

95 

96def calculate_extraction_confidence( 

97 extraction, 

98 alignment_scores: Optional[Dict[str, float]] = None 

99) -> float: 

100 """Calculate confidence score based on extraction alignment status. 

101 

102 :param extraction: LangExtract extraction object with optional alignment_status. 

103 :param alignment_scores: Custom score mapping for alignment statuses (optional). 

104 :return: Confidence score between 0.0 and 1.0. 

105 """ 

106 default_score = 0.85 

107 

108 if alignment_scores is None: 

109 alignment_scores = DEFAULT_ALIGNMENT_SCORES 

110 

111 if not hasattr(extraction, "alignment_status") or not ( 

112 extraction.alignment_status 

113 ): 

114 return default_score 

115 

116 status = str(extraction.alignment_status).upper() 

117 for status_key, score in alignment_scores.items(): 

118 if status_key in status: 

119 return score 

120 

121 return default_score 

122 

123 

124def convert_langextract_to_presidio_results( 

125 langextract_result, 

126 entity_mappings: Dict, 

127 supported_entities: List[str], 

128 enable_generic_consolidation: bool, 

129 recognizer_name: str, 

130 alignment_scores: Optional[Dict[str, float]] = None 

131) -> List[RecognizerResult]: 

132 """Convert LangExtract extraction results to Presidio RecognizerResult objects. 

133 

134 :param langextract_result: LangExtract result object with extractions. 

135 :param entity_mappings: Mapping of extraction classes to Presidio entity types. 

136 :param supported_entities: List of supported Presidio entity types. 

137 :param enable_generic_consolidation: Whether to consolidate unknown entities. 

138 :param recognizer_name: Name of recognizer for result metadata. 

139 :param alignment_scores: Custom alignment score mappings (optional). 

140 :return: List of Presidio RecognizerResult objects. 

141 """ 

142 results = [] 

143 if not langextract_result or not langextract_result.extractions: 

144 return results 

145 

146 supported_entities_set = set(supported_entities) 

147 

148 for extraction in langextract_result.extractions: 

149 extraction_class = extraction.extraction_class 

150 

151 if extraction_class in supported_entities_set: 

152 entity_type = extraction_class 

153 else: 

154 extraction_class_lower = extraction_class.lower() 

155 entity_type = entity_mappings.get(extraction_class_lower) 

156 

157 if not entity_type: 

158 if enable_generic_consolidation: 

159 entity_type = extraction_class.upper() 

160 logger.debug( 

161 "Unknown extraction class '%s' will be consolidated to " 

162 "GENERIC_PII_ENTITY", 

163 extraction_class, 

164 ) 

165 else: 

166 logger.warning( 

167 "Unknown extraction class '%s' not found in entity " 

168 "mappings, skipping", 

169 extraction_class, 

170 ) 

171 continue 

172 

173 if not extraction.char_interval: 

174 logger.warning("Extraction missing char_interval, skipping") 

175 continue 

176 

177 confidence = calculate_extraction_confidence(extraction, alignment_scores) 

178 

179 metadata = {} 

180 if hasattr(extraction, 'attributes') and extraction.attributes: 

181 metadata['attributes'] = extraction.attributes 

182 if hasattr(extraction, 'alignment_status') and extraction.alignment_status: 

183 metadata['alignment'] = str(extraction.alignment_status) 

184 

185 explanation = AnalysisExplanation( 

186 recognizer=recognizer_name, 

187 original_score=confidence, 

188 textual_explanation=( 

189 f"LangExtract extraction with " 

190 f"{extraction.alignment_status} alignment" 

191 if hasattr(extraction, "alignment_status") 

192 and extraction.alignment_status 

193 else "LangExtract extraction" 

194 ), 

195 ) 

196 

197 result = RecognizerResult( 

198 entity_type=entity_type, 

199 start=extraction.char_interval.start_pos, 

200 end=extraction.char_interval.end_pos, 

201 score=confidence, 

202 analysis_explanation=explanation, 

203 recognition_metadata=metadata if metadata else None 

204 ) 

205 

206 results.append(result) 

207 

208 return results