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

1from __future__ import annotations 

2 

3from typing import TYPE_CHECKING, Final 

4 

5from litellm.types.guardrails import ( 

6 GuardrailEventHooks, 

7 Mode, 

8 SupportedGuardrailIntegrations, 

9) 

10 

11from .compresr import CompresrGuardrail 

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 _coerce_event_hook( 

18 mode: str | list[str] | Mode, 

19) -> GuardrailEventHooks | list[GuardrailEventHooks] | Mode: 

20 if isinstance(mode, Mode): 

21 return mode 

22 if isinstance(mode, list): 

23 return [GuardrailEventHooks(item) for item in mode] 

24 return GuardrailEventHooks(mode) 

25 

26 

27def _get_optional_value(litellm_params: LitellmParams, optional_params: object | None, attribute_name: str) -> object: 

28 if optional_params is not None: 

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

30 if value is not None: 

31 return value 

32 return getattr(litellm_params, attribute_name, None) 

33 

34 

35def initialize_guardrail(litellm_params: LitellmParams, guardrail: Guardrail) -> CompresrGuardrail: 

36 import litellm 

37 

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

39 

40 _callback: Final = CompresrGuardrail( 

41 api_base=litellm_params.api_base, 

42 api_key=litellm_params.api_key, 

43 model=litellm_params.model, 

44 target_compression_ratio=_get_optional_value(litellm_params, optional_params, "target_compression_ratio"), 

45 coarse=_get_optional_value(litellm_params, optional_params, "coarse"), 

46 min_chars_to_compress=_get_optional_value(litellm_params, optional_params, "min_chars_to_compress"), 

47 compress_tool_outputs=_get_optional_value(litellm_params, optional_params, "compress_tool_outputs"), 

48 compress_system=_get_optional_value(litellm_params, optional_params, "compress_system"), 

49 compress_history=_get_optional_value(litellm_params, optional_params, "compress_history"), 

50 compress_last_user=_get_optional_value(litellm_params, optional_params, "compress_last_user"), 

51 enable_retrieval=_get_optional_value(litellm_params, optional_params, "enable_retrieval"), 

52 max_bytes_per_call=_get_optional_value(litellm_params, optional_params, "max_bytes_per_call"), 

53 allow_bypass_header=_get_optional_value(litellm_params, optional_params, "allow_bypass_header"), 

54 dynamic=_get_optional_value(litellm_params, optional_params, "dynamic"), 

55 dynamic_min_ratio=_get_optional_value(litellm_params, optional_params, "dynamic_min_ratio"), 

56 dynamic_max_ratio=_get_optional_value(litellm_params, optional_params, "dynamic_max_ratio"), 

57 compression_params=_get_optional_value(litellm_params, optional_params, "compression_params"), 

58 guardrail_name=guardrail["guardrail_name"], 

59 event_hook=_coerce_event_hook(litellm_params.mode), 

60 default_on=litellm_params.default_on or False, 

61 unreachable_fallback=litellm_params.unreachable_fallback, 

62 ) 

63 litellm.logging_callback_manager.add_litellm_callback( # pyright: ignore[reportUnknownMemberType] # callback manager is untyped 

64 _callback 

65 ) 

66 return _callback 

67 

68 

69guardrail_initializer_registry: Final = { 

70 SupportedGuardrailIntegrations.COMPRESR.value: initialize_guardrail, 

71} 

72 

73guardrail_class_registry: Final = { 

74 SupportedGuardrailIntegrations.COMPRESR.value: CompresrGuardrail, 

75}