Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/guardrails/guardrail_hooks/azure/__init__.py: 24%

1from typing import TYPE_CHECKING, Final 

2 

3from litellm.types.guardrails import SupportedGuardrailIntegrations 

4 

5from .prompt_shield import AzureContentSafetyPromptShieldGuardrail 

6from .text_moderation import AzureContentSafetyTextModerationGuardrail 

7 

8if TYPE_CHECKING: 8 ↛ 9line 8 didn't jump to line 9 because the condition on line 8 was never true

9 from litellm.types.guardrails import Guardrail, LitellmParams 

10 

11 

12def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"): 

13 import litellm 

14 

15 if not litellm_params.api_key: 

16 raise ValueError("Azure Content Safety: api_key is required") 

17 if not litellm_params.api_base: 

18 raise ValueError("Azure Content Safety: api_base is required") 

19 

20 azure_guardrail: Final = litellm_params.guardrail.split("/")[1] 

21 

22 guardrail_name: Final = guardrail.get("guardrail_name") 

23 if not guardrail_name: 

24 raise ValueError("Azure Content Safety: guardrail_name is required") 

25 

26 if azure_guardrail == "prompt_shield": 

27 azure_content_safety_guardrail: ( 

28 AzureContentSafetyPromptShieldGuardrail | AzureContentSafetyTextModerationGuardrail 

29 ) = AzureContentSafetyPromptShieldGuardrail( 

30 guardrail_name=guardrail_name, 

31 **{ 

32 **litellm_params.model_dump(exclude_none=True), 

33 "api_key": litellm_params.api_key, 

34 "api_base": litellm_params.api_base, 

35 "default_on": litellm_params.default_on, 

36 "event_hook": litellm_params.mode, 

37 }, 

38 ) 

39 elif azure_guardrail == "text_moderations": 

40 azure_content_safety_guardrail = AzureContentSafetyTextModerationGuardrail( 

41 guardrail_name=guardrail_name, 

42 **{ 

43 **litellm_params.model_dump(exclude_none=True), 

44 "api_key": litellm_params.api_key, 

45 "api_base": litellm_params.api_base, 

46 "default_on": litellm_params.default_on, 

47 "event_hook": litellm_params.mode, 

48 }, 

49 ) 

50 else: 

51 raise ValueError(f"Azure Content Safety: {azure_guardrail} is not a valid guardrail") 

52 

53 litellm.logging_callback_manager.add_litellm_callback(azure_content_safety_guardrail) 

54 return azure_content_safety_guardrail 

55 

56 

57guardrail_initializer_registry: Final = { 

58 SupportedGuardrailIntegrations.AZURE_PROMPT_SHIELD.value: initialize_guardrail, 

59 SupportedGuardrailIntegrations.AZURE_TEXT_MODERATIONS.value: initialize_guardrail, 

60} 

61 

62 

63guardrail_class_registry: Final = { 

64 SupportedGuardrailIntegrations.AZURE_PROMPT_SHIELD.value: AzureContentSafetyPromptShieldGuardrail, 

65 SupportedGuardrailIntegrations.AZURE_TEXT_MODERATIONS.value: AzureContentSafetyTextModerationGuardrail, 

66}