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

1from typing import TYPE_CHECKING, Final 

2 

3from litellm.types.guardrails import SupportedGuardrailIntegrations 

4 

5from .guardrails_ai import GuardrailsAI 

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 initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"): 

12 import litellm 

13 

14 if litellm_params.guard_name is None: 

15 raise Exception( 

16 "GuardrailsAIException - Please pass the Guardrails AI guard name via 'litellm_params::guard_name'" 

17 ) 

18 

19 _guardrails_ai_callback: Final = GuardrailsAI( 

20 api_base=litellm_params.api_base, 

21 api_key=litellm_params.api_key, 

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

23 event_hook=litellm_params.mode, 

24 default_on=litellm_params.default_on, 

25 guard_name=litellm_params.guard_name, 

26 guardrails_ai_api_input_format=getattr(litellm_params, "guardrails_ai_api_input_format", "llmOutput"), 

27 ) 

28 litellm.logging_callback_manager.add_litellm_callback(_guardrails_ai_callback) 

29 

30 return _guardrails_ai_callback 

31 

32 

33guardrail_initializer_registry: Final = { 

34 SupportedGuardrailIntegrations.GUARDRAILS_AI.value: initialize_guardrail, 

35} 

36 

37 

38guardrail_class_registry: Final = { 

39 SupportedGuardrailIntegrations.GUARDRAILS_AI.value: GuardrailsAI, 

40}