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

1"""Rubrik guardrail integration for LiteLLM.""" 

2 

3from typing import TYPE_CHECKING, Final 

4 

5from litellm.integrations.rubrik import RubrikLogger 

6from litellm.types.guardrails import SupportedGuardrailIntegrations 

7 

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

9 from litellm.types.guardrails import Guardrail, LitellmParams 

10 

11 

12def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail") -> RubrikLogger: 

13 """Create and register a RubrikLogger instance. 

14 

15 The ``mode`` field in the guardrail config controls which surfaces are 

16 moderated: 

17 - ``pre_call`` (or a mode that includes it): prompt moderation via the 

18 ``/v1/before_prompt/openai/v1`` webhook. 

19 - ``post_call`` (the default when ``mode`` is omitted): response and tool 

20 call moderation via the ``/v1/after_completion/openai/v1`` webhook. 

21 

22 Both hooks are active when ``mode`` covers both ``pre_call`` and 

23 ``post_call``. 

24 """ 

25 import litellm 

26 

27 rubrik_callback: Final = RubrikLogger( 

28 api_key=litellm_params.api_key, 

29 api_base=litellm_params.api_base, 

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

31 event_hook=litellm_params.mode, 

32 default_on=litellm_params.default_on, 

33 ) 

34 

35 litellm.logging_callback_manager.add_litellm_callback(rubrik_callback) 

36 return rubrik_callback 

37 

38 

39guardrail_initializer_registry: Final = { 

40 SupportedGuardrailIntegrations.RUBRIK.value: initialize_guardrail, 

41} 

42 

43guardrail_class_registry: Final = { 

44 SupportedGuardrailIntegrations.RUBRIK.value: RubrikLogger, 

45}