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

1from typing import TYPE_CHECKING, Final 

2 

3from litellm.types.guardrails import SupportedGuardrailIntegrations 

4 

5from .hiddenlayer import HiddenlayerGuardrail, HiddenlayerGuardrailV2 

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 api_id: Final = litellm_params.api_id if hasattr(litellm_params, "api_id") else None 

15 auth_url: Final = litellm_params.auth_url if hasattr(litellm_params, "auth_url") else None 

16 version: Final[int | None] = litellm_params.version if hasattr(litellm_params, "version") else None 

17 

18 _hiddenlayer_callback: HiddenlayerGuardrail | HiddenlayerGuardrailV2 

19 if not version or version < 2: 

20 _hiddenlayer_callback = HiddenlayerGuardrail( 

21 api_base=litellm_params.api_base, 

22 api_id=api_id, 

23 api_key=litellm_params.api_key, 

24 auth_url=auth_url, 

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

26 event_hook=litellm_params.mode, 

27 default_on=litellm_params.default_on, 

28 ) 

29 else: 

30 _hiddenlayer_callback = HiddenlayerGuardrailV2( 

31 api_base=litellm_params.api_base, 

32 api_id=api_id, 

33 api_key=litellm_params.api_key, 

34 auth_url=auth_url, 

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

36 event_hook=litellm_params.mode, 

37 default_on=litellm_params.default_on, 

38 ) 

39 

40 litellm.logging_callback_manager.add_litellm_callback(_hiddenlayer_callback) 

41 return _hiddenlayer_callback 

42 

43 

44guardrail_initializer_registry: Final = { 

45 SupportedGuardrailIntegrations.HIDDENLAYER.value: initialize_guardrail, 

46} 

47 

48 

49guardrail_class_registry: Final = { 

50 SupportedGuardrailIntegrations.HIDDENLAYER.value: HiddenlayerGuardrail, 

51}