Coverage for presidio_analyzer / predefined_recognizers / third_party / langextract_recognizer.py: 98%

1import logging 

2from abc import ABC, abstractmethod 

3from typing import Any, Dict, List, Optional 

4 

5from presidio_analyzer.llm_utils import ( 

6 check_langextract_available, 

7 convert_langextract_to_presidio_results, 

8 convert_to_langextract_format, 

9 extract_lm_config, 

10 get_model_config, 

11 get_supported_entities, 

12 load_prompt_file, 

13 load_yaml_examples, 

14 load_yaml_file, 

15 lx, 

16 render_jinja_template, 

17 validate_config_fields, 

18) 

19from presidio_analyzer.lm_recognizer import LMRecognizer 

20 

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

22 

23 

24class LangExtractRecognizer(LMRecognizer, ABC): 

25 """ 

26 Base class for LangExtract-based PII recognizers. 

27 

28 Subclasses implement _call_langextract() for specific LLM providers. 

29 """ 

30 

31 def __init__( 

32 self, 

33 config_path: str, 

34 name: str = "LangExtract LLM PII", 

35 supported_language: str = "en", 

36 extract_params: Optional[Dict[str, Any]] = None, 

37 ): 

38 """Initialize LangExtract recognizer. 

39 

40 :param config_path: Path to configuration file. 

41 :param name: Name of the recognizer (provided by subclass). 

42 :param supported_language: Language this recognizer supports (default: "en"). 

43 :param extract_params: Dict with 'extract' and/or 'language_model' 

44 keys containing param defaults. 

45 """ 

46 check_langextract_available() 

47 

48 full_config = load_yaml_file(config_path) 

49 

50 lm_config = extract_lm_config(full_config) 

51 langextract_config = full_config.get("langextract", {}) 

52 

53 supported_entities = get_supported_entities(lm_config, langextract_config) 

54 

55 if not supported_entities: 

56 raise ValueError( 

57 "Configuration must contain 'supported_entities' in " 

58 "'lm_recognizer' or 'langextract'" 

59 ) 

60 

61 validate_config_fields( 

62 full_config, 

63 [ 

64 ("langextract",), 

65 ("langextract", "model"), 

66 ("langextract", "model", "model_id"), 

67 ("langextract", "entity_mappings"), 

68 ("langextract", "prompt_file"), 

69 ("langextract", "examples_file"), 

70 ] 

71 ) 

72 

73 self.config = langextract_config 

74 model_config = get_model_config( 

75 full_config, provider_key="langextract" 

76 ) 

77 

78 super().__init__( 

79 supported_entities=supported_entities, 

80 supported_language=supported_language, 

81 name=name, 

82 version="1.0.0", 

83 model_id=model_config["model_id"], 

84 temperature=model_config.get("temperature"), 

85 min_score=lm_config.get("min_score"), 

86 labels_to_ignore=lm_config.get("labels_to_ignore"), 

87 enable_generic_consolidation=lm_config.get( 

88 "enable_generic_consolidation" 

89 ), 

90 ) 

91 

92 examples_data = load_yaml_examples( 

93 langextract_config["examples_file"] 

94 ) 

95 self.examples = convert_to_langextract_format(examples_data) 

96 

97 prompt_template = load_prompt_file( 

98 langextract_config["prompt_file"] 

99 ) 

100 self.prompt_description = render_jinja_template( 

101 prompt_template, 

102 supported_entities=self.supported_entities, 

103 enable_generic_consolidation=self.enable_generic_consolidation, 

104 labels_to_ignore=self.labels_to_ignore, 

105 ) 

106 

107 self.entity_mappings = langextract_config["entity_mappings"] 

108 self.debug = langextract_config.get("debug", False) 

109 self._model_config = model_config 

110 

111 # Process extract params with config override 

112 self._extract_params = {} 

113 self._language_model_params = {} 

114 

115 if extract_params: 

116 if "extract" in extract_params: 

117 for param_name, default_value in extract_params["extract"].items(): 

118 self._extract_params[param_name] = self._model_config.get( 

119 param_name, default_value 

120 ) 

121 

122 if "language_model" in extract_params: 

123 for param_name, default_value in ( 

124 extract_params["language_model"].items() 

125 ): 

126 self._language_model_params[param_name] = ( 

127 self._model_config.get(param_name, default_value) 

128 ) 

129 

130 def _call_llm(self, text: str, entities: List[str], **kwargs): 

131 """Call LangExtract LLM.""" 

132 # Build extract params 

133 extract_params = { 

134 "text": text, 

135 "prompt": self.prompt_description, 

136 "examples": self.examples, 

137 "debug": self.debug, 

138 } 

139 

140 # Add temperature if configured 

141 if self.temperature is not None: 

142 extract_params["temperature"] = self.temperature 

143 

144 # Add any additional kwargs 

145 extract_params.update(kwargs) 

146 

147 langextract_result = self._call_langextract(**extract_params) 

148 

149 return convert_langextract_to_presidio_results( 

150 langextract_result=langextract_result, 

151 entity_mappings=self.entity_mappings, 

152 supported_entities=self.supported_entities, 

153 enable_generic_consolidation=self.enable_generic_consolidation, 

154 recognizer_name=self.__class__.__name__ 

155 ) 

156 

157 def _call_langextract(self, **kwargs): 

158 """Call LangExtract with configured parameters.""" 

159 try: 

160 extract_params = { 

161 "text_or_documents": kwargs.pop("text"), 

162 "prompt_description": kwargs.pop("prompt"), 

163 "examples": kwargs.pop("examples"), 

164 } 

165 

166 extract_params.update(self._get_provider_params()) 

167 extract_params.update(self._extract_params) 

168 if self._language_model_params: 

169 extract_params["language_model_params"] = self._language_model_params 

170 extract_params.update(kwargs) 

171 

172 return lx.extract(**extract_params) 

173 except Exception: 

174 logger.exception( 

175 "LangExtract extraction failed (model '%s')", 

176 self.model_id 

177 ) 

178 raise 

179 

180 @abstractmethod 

181 def _get_provider_params(self) -> Dict[str, Any]: 

182 """Return provider-specific params. 

183 

184 Examples: model_id, model_url, azure_endpoint, etc. 

185 """ 

186 ...