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

1"""Cisco AI Defense Guardrail Integration for LiteLLM.""" 

2 

3from typing import TYPE_CHECKING, Final 

4 

5from litellm.types.guardrails import SupportedGuardrailIntegrations 

6 

7from .cisco_ai_defense import ( 

8 CiscoAIDefenseGuardrail, 

9 CiscoAIDefenseGuardrailAPIError, 

10 CiscoAIDefenseGuardrailMissingSecrets, 

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"): 

18 import litellm 

19 

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

21 if not guardrail_name: 

22 raise ValueError("Cisco AI Defense: guardrail_name is required") 

23 

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

25 

26 _callback: Final = CiscoAIDefenseGuardrail( 

27 guardrail_name=guardrail_name, 

28 api_key=litellm_params.api_key, 

29 api_base=litellm_params.api_base, 

30 inspection_type=_get_optional_value(litellm_params, optional_params, "inspection_type"), 

31 inspect_path=_get_optional_value(litellm_params, optional_params, "inspect_path"), 

32 enabled_rules=_get_optional_value(litellm_params, optional_params, "enabled_rules"), 

33 integration_profile_id=_get_optional_value(litellm_params, optional_params, "integration_profile_id"), 

34 integration_profile_version=_get_optional_value(litellm_params, optional_params, "integration_profile_version"), 

35 integration_tenant_id=_get_optional_value(litellm_params, optional_params, "integration_tenant_id"), 

36 integration_type=_get_optional_value(litellm_params, optional_params, "integration_type"), 

37 on_flagged_action=_get_optional_value(litellm_params, optional_params, "on_flagged_action"), 

38 fallback_on_error=_get_optional_value(litellm_params, optional_params, "fallback_on_error"), 

39 timeout=_get_optional_value(litellm_params, optional_params, "timeout"), 

40 event_hook=litellm_params.mode, 

41 default_on=litellm_params.default_on or False, 

42 ) 

43 litellm.logging_callback_manager.add_litellm_callback(_callback) 

44 

45 # MCP post-tool-call hooks are dispatched through success callbacks. 

46 litellm.logging_callback_manager.add_litellm_success_callback(_callback) 

47 

48 return _callback 

49 

50 

51def _get_optional_value(litellm_params, optional_params, attribute_name): 

52 """Resolve Cisco optional params without inheriting sibling defaults.""" 

53 if optional_params is not None: 

54 if isinstance(optional_params, dict): 

55 if attribute_name in optional_params: 

56 return optional_params[attribute_name] 

57 else: 

58 nested_fields_set: Final = getattr(optional_params, "model_fields_set", None) 

59 if nested_fields_set is None or attribute_name in nested_fields_set: 

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

61 if value is not None: 

62 return value 

63 

64 if litellm_params is None: 

65 return None 

66 # Only accept flattened values the caller explicitly set. 

67 fields_set: Final = getattr(litellm_params, "model_fields_set", None) 

68 if fields_set is None or attribute_name not in fields_set: 

69 return None 

70 return getattr(litellm_params, attribute_name, None) 

71 

72 

73guardrail_initializer_registry: Final = { 

74 SupportedGuardrailIntegrations.CISCO_AI_DEFENSE.value: initialize_guardrail, 

75} 

76 

77 

78guardrail_class_registry: Final = { 

79 SupportedGuardrailIntegrations.CISCO_AI_DEFENSE.value: CiscoAIDefenseGuardrail, 

80} 

81 

82 

83__all__ = [ 

84 "CiscoAIDefenseGuardrail", 

85 "CiscoAIDefenseGuardrailAPIError", 

86 "CiscoAIDefenseGuardrailMissingSecrets", 

87 "guardrail_class_registry", 

88 "guardrail_initializer_registry", 

89 "initialize_guardrail", 

90]