Coverage for presidio_analyzer / predefined_recognizers / third_party / azure_openai_provider.py: 100%

1"""Azure OpenAI Provider for LangExtract.""" 

2 

3import logging 

4import os 

5from typing import Optional 

6 

7try: 

8 import langextract as lx 

9 import openai 

10 from langextract.providers.openai import OpenAILanguageModel 

11 LANGEXTRACT_OPENAI_AVAILABLE = True 

12except ImportError: # pragma: no cover 

13 LANGEXTRACT_OPENAI_AVAILABLE = False 

14 lx = None 

15 OpenAILanguageModel = None 

16 openai = None 

17 

18try: 

19 from presidio_analyzer.llm_utils.azure_auth_helper import ( 

20 get_azure_credential, 

21 get_bearer_token_provider_for_scope, 

22 ) 

23 AZURE_IDENTITY_AVAILABLE = True 

24except ImportError: # pragma: no cover 

25 AZURE_IDENTITY_AVAILABLE = False 

26 get_azure_credential = None 

27 get_bearer_token_provider_for_scope = None 

28 

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

30 

31 

32if LANGEXTRACT_OPENAI_AVAILABLE: 

33 class AzureOpenAILanguageModel(OpenAILanguageModel): 

34 """ 

35 Custom LangExtract provider for Azure OpenAI. 

36 

37 This provider extends OpenAILanguageModel to support Azure-specific 

38 authentication and endpoint configuration. It reuses all inference logic 

39 from the parent class and only overrides client initialization. 

40 

41 Registered to handle model_id with "azure:" prefix (e.g., "azure:gpt-4o"). 

42 The recognizer adds this prefix automatically. 

43 """ 

44 

45 def __init__( 

46 self, 

47 model_id: str, 

48 api_key: Optional[str] = None, 

49 azure_endpoint: Optional[str] = None, 

50 api_version: Optional[str] = None, 

51 azure_deployment: Optional[str] = None, 

52 azure_ad_token_provider: Optional[any] = None, 

53 **kwargs 

54 ): 

55 """ 

56 Initialize Azure OpenAI provider. 

57 

58 :param model_id: Azure OpenAI deployment name or model identifier. 

59 :param api_key: Azure OpenAI API key (or set AZURE_OPENAI_API_KEY 

60 env var). If not provided, will automatically use managed 

61 identity (ChainedTokenCredential). 

62 :param azure_endpoint: Azure OpenAI endpoint URL 

63 (or set AZURE_OPENAI_ENDPOINT env var). 

64 :param api_version: Azure OpenAI API version 

65 (or set AZURE_OPENAI_API_VERSION env var). 

66 :param azure_deployment: Explicit deployment name (optional, 

67 defaults to model_id). 

68 :param azure_ad_token_provider: Custom Azure AD token provider function. 

69 :param kwargs: Additional parameters passed to parent class. 

70 """ 

71 # Strip 'azure:' prefix if present (added by recognizer) 

72 if model_id.lower().startswith("azure:"): 

73 clean_model_id = model_id[6:] # len("azure:") = 6 

74 else: 

75 clean_model_id = model_id 

76 

77 self.model_id = clean_model_id 

78 self.api_key = api_key or os.environ.get("AZURE_OPENAI_API_KEY") 

79 self.organization = None 

80 self.format_type = kwargs.get('format_type', lx.data.FormatType.JSON) 

81 self.temperature = kwargs.get('temperature', 0.0) 

82 self.max_workers = kwargs.get('max_workers', 10) 

83 self._extra_kwargs = kwargs 

84 

85 # Azure-specific configuration 

86 self.azure_endpoint = azure_endpoint or os.environ.get( 

87 "AZURE_OPENAI_ENDPOINT" 

88 ) 

89 self.api_version = api_version or os.environ.get( 

90 "AZURE_OPENAI_API_VERSION", "2024-02-15-preview" 

91 ) 

92 self.azure_deployment = azure_deployment or clean_model_id 

93 

94 # Validate and initialize Azure OpenAI client 

95 self._client = self._create_azure_openai_client( 

96 azure_ad_token_provider 

97 ) 

98 

99 def _create_azure_openai_client( 

100 self, 

101 azure_ad_token_provider: Optional[any] = None 

102 ): 

103 """ 

104 Create and configure Azure OpenAI client with appropriate authentication. 

105 

106 :param azure_ad_token_provider: Optional custom token provider 

107 :return: Configured AzureOpenAI client instance 

108 :raises ValueError: If azure_endpoint is not provided 

109 :raises ImportError: If azure-identity is needed but not available 

110 """ 

111 if not self.azure_endpoint: 

112 raise ValueError( 

113 "Azure OpenAI endpoint is required. Set AZURE_OPENAI_ENDPOINT " 

114 "environment variable or pass azure_endpoint parameter." 

115 ) 

116 

117 if not self.api_key or azure_ad_token_provider: 

118 if not AZURE_IDENTITY_AVAILABLE and not azure_ad_token_provider: 

119 raise ImportError( 

120 "azure-identity is required for managed identity " 

121 "authentication. Install it with: pip install azure-identity" 

122 ) 

123 

124 if azure_ad_token_provider: 

125 token_provider = azure_ad_token_provider 

126 credential_type = "custom token provider" 

127 else: 

128 token_provider = get_bearer_token_provider_for_scope( 

129 "https://cognitiveservices.azure.com/.default" 

130 ) 

131 credential_type = ( 

132 "DefaultAzureCredential (development)" 

133 if os.getenv('ENV') == 'development' 

134 else "ChainedTokenCredential" 

135 ) 

136 

137 client = openai.AzureOpenAI( 

138 azure_ad_token_provider=token_provider, 

139 azure_endpoint=self.azure_endpoint, 

140 api_version=self.api_version 

141 ) 

142 

143 logger.debug( 

144 f"Initialized Azure OpenAI provider with {credential_type}" 

145 ) 

146 return client 

147 else: 

148 client = openai.AzureOpenAI( 

149 api_key=self.api_key, 

150 azure_endpoint=self.azure_endpoint, 

151 api_version=self.api_version 

152 ) 

153 

154 logger.debug( 

155 "Initialized Azure OpenAI provider with API key authentication" 

156 ) 

157 return client 

158 

159 def _get_client_model_id(self) -> str: 

160 """ 

161 Return the model/deployment identifier for API calls. 

162 

163 For Azure OpenAI, this is the deployment name, not the base model name. 

164 

165 :return: Azure deployment name. 

166 """ 

167 return self.azure_deployment 

168else: # pragma: no cover 

169 class AzureOpenAILanguageModel: 

170 """Placeholder when langextract is not installed.""" 

171 

172 def __init__(self, *args, **kwargs): 

173 raise ImportError( 

174 "LangExtract with OpenAI support is not installed. " 

175 "Install it with: pip install presidio-analyzer[langextract,openai] " 

176 "or: pip install langextract[openai]" 

177 ) 

178 

179 

180# Register the provider with LangExtract 

181if LANGEXTRACT_OPENAI_AVAILABLE: 

182 try: 

183 @lx.providers.registry.register( 

184 r'^azure:', 

185 priority=20 

186 ) 

187 class RegisteredAzureOpenAILanguageModel(AzureOpenAILanguageModel): 

188 """ 

189 Registered version of Azure OpenAI provider for LangExtract. 

190 

191 This class is automatically discovered by LangExtract's provider registry 

192 and used when model_id matches any of the registered patterns. 

193 """ 

194 

195 pass 

196 

197 logger.debug("Azure OpenAI provider registered with LangExtract") 

198 except Exception as e: # pragma: no cover 

199 logger.error(f"Failed to register Azure OpenAI provider: {e}") 

200 raise