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

1"""Gray Swan Cygnal guardrail integration for LiteLLM.""" 

2 

3from typing import TYPE_CHECKING, Final 

4 

5from litellm.types.guardrails import SupportedGuardrailIntegrations 

6 

7from .grayswan import ( 

8 GraySwanGuardrail, 

9 GraySwanGuardrailAPIError, 

10 GraySwanGuardrailMissingSecrets, 

11) 

12 

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

14 from litellm.types.guardrails import Guardrail, LitellmParams 

15 

16 

17def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail") -> GraySwanGuardrail: 

18 import litellm 

19 

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

21 if not guardrail_name: 

22 raise ValueError("Gray Swan guardrail requires a guardrail_name") 

23 

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

25 

26 grayswan_guardrail: Final = GraySwanGuardrail( 

27 guardrail_name=guardrail_name, 

28 api_key=litellm_params.api_key, 

29 api_base=litellm_params.api_base, 

30 on_flagged_action=_get_config_value(litellm_params, optional_params, "on_flagged_action"), 

31 violation_threshold=_get_config_value(litellm_params, optional_params, "violation_threshold"), 

32 reasoning_mode=_get_config_value(litellm_params, optional_params, "reasoning_mode"), 

33 categories=_get_config_value(litellm_params, optional_params, "categories"), 

34 policy_id=_get_config_value(litellm_params, optional_params, "policy_id"), 

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

36 or False, 

37 streaming_sampling_rate=_get_config_value(litellm_params, optional_params, "streaming_sampling_rate") or 5, 

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

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

40 event_hook=litellm_params.mode, 

41 default_on=litellm_params.default_on, 

42 ) 

43 

44 litellm.logging_callback_manager.add_litellm_callback(grayswan_guardrail) 

45 return grayswan_guardrail 

46 

47 

48def _get_config_value(litellm_params, optional_params, attribute_name): 

49 if optional_params is not None: 

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

51 if value is not None: 

52 return value 

53 return getattr(litellm_params, attribute_name, None) 

54 

55 

56guardrail_initializer_registry: Final = { 

57 SupportedGuardrailIntegrations.GRAYSWAN.value: initialize_guardrail, 

58} 

59 

60 

61guardrail_class_registry: Final = { 

62 SupportedGuardrailIntegrations.GRAYSWAN.value: GraySwanGuardrail, 

63} 

64 

65 

66__all__ = [ 

67 "GraySwanGuardrail", 

68 "GraySwanGuardrailAPIError", 

69 "GraySwanGuardrailMissingSecrets", 

70 "initialize_guardrail", 

71]