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

1""" 

2Pillar Security Guardrail Integration for LiteLLM 

3""" 

4 

5from typing import TYPE_CHECKING, Final 

6 

7from litellm.types.guardrails import SupportedGuardrailIntegrations 

8 

9from .pillar import ( 

10 PillarGuardrail, 

11 PillarGuardrailAPIError, 

12 PillarGuardrailMissingSecrets, 

13) 

14 

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

16 from litellm.types.guardrails import Guardrail, LitellmParams 

17 

18 

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

20 import litellm 

21 

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

23 if not guardrail_name: 

24 raise ValueError("Pillar guardrail name is required") 

25 

26 optional_params: Final = getattr(litellm_params, "optional_params", None) 

27 

28 _pillar_callback: Final = PillarGuardrail( 

29 guardrail_name=guardrail_name, 

30 api_key=litellm_params.api_key, 

31 api_base=litellm_params.api_base, 

32 on_flagged_action=getattr(litellm_params, "on_flagged_action", "monitor"), 

33 event_hook=litellm_params.mode, 

34 default_on=litellm_params.default_on, 

35 async_mode=_get_config_value(litellm_params, optional_params, "async_mode"), 

36 persist_session=_get_config_value(litellm_params, optional_params, "persist_session"), 

37 include_scanners=_get_config_value(litellm_params, optional_params, "include_scanners"), 

38 include_evidence=_get_config_value(litellm_params, optional_params, "include_evidence"), 

39 fallback_on_error=_get_config_value(litellm_params, optional_params, "fallback_on_error"), 

40 timeout=_get_config_value(litellm_params, optional_params, "timeout"), 

41 ) 

42 litellm.logging_callback_manager.add_litellm_callback(_pillar_callback) 

43 

44 return _pillar_callback 

45 

46 

47def _get_config_value(litellm_params, optional_params, attribute_name): 

48 """Return guardrail configuration value prioritising optional params when present.""" 

49 

50 if optional_params is not None: 

51 value: Final = getattr(optional_params, attribute_name, None) 

52 if value is not None: 

53 return value 

54 return getattr(litellm_params, attribute_name, None) 

55 

56 

57guardrail_initializer_registry: Final = { 

58 SupportedGuardrailIntegrations.PILLAR.value: initialize_guardrail, 

59} 

60 

61 

62guardrail_class_registry: Final = { 

63 SupportedGuardrailIntegrations.PILLAR.value: PillarGuardrail, 

64} 

65 

66__all__ = [ 

67 "PillarGuardrail", 

68 "PillarGuardrailAPIError", 

69 "PillarGuardrailMissingSecrets", 

70]