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

1from typing import TYPE_CHECKING, Any, Final 

2 

3from litellm.types.guardrails import SupportedGuardrailIntegrations 

4 

5from .generic_guardrail_api import GenericGuardrailAPI 

6 

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

8 from litellm.types.guardrails import Guardrail, LitellmParams 

9 

10 

11def _get_config_value(litellm_params: "LitellmParams", optional_params: object, attribute_name: str) -> Any | None: 

12 if optional_params is not None: 

13 value: Final = ( 

14 optional_params.get(attribute_name) 

15 if isinstance(optional_params, dict) 

16 else getattr(optional_params, attribute_name, None) 

17 ) 

18 if value is not None: 

19 return value 

20 return getattr(litellm_params, attribute_name, None) 

21 

22 

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

24 import litellm 

25 

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

27 

28 _generic_guardrail_api_callback: Final = GenericGuardrailAPI( 

29 api_base=litellm_params.api_base, 

30 api_key=litellm_params.api_key, 

31 headers=getattr(litellm_params, "headers", None), 

32 additional_provider_specific_params=getattr(litellm_params, "additional_provider_specific_params", {}), 

33 unreachable_fallback=getattr(litellm_params, "unreachable_fallback", "fail_closed"), 

34 fail_on_error=getattr(litellm_params, "fail_on_error", True), 

35 extra_headers=getattr(litellm_params, "extra_headers", None), 

36 guardrail_name=guardrail.get("guardrail_name", ""), 

37 event_hook=litellm_params.mode, 

38 default_on=litellm_params.default_on, 

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

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

41 streaming_transform_mode=_get_config_value(litellm_params, optional_params, "streaming_transform_mode"), 

42 ) 

43 

44 litellm.logging_callback_manager.add_litellm_callback(_generic_guardrail_api_callback) 

45 return _generic_guardrail_api_callback 

46 

47 

48guardrail_initializer_registry: Final = { 

49 SupportedGuardrailIntegrations.GENERIC_GUARDRAIL_API.value: initialize_guardrail, 

50} 

51 

52guardrail_class_registry: Final = { 

53 SupportedGuardrailIntegrations.GENERIC_GUARDRAIL_API.value: GenericGuardrailAPI, 

54}