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

1from __future__ import annotations 

2 

3from typing import TYPE_CHECKING, Final 

4 

5from pydantic import BaseModel 

6 

7from litellm.types.guardrails import ( 

8 GuardrailEventHooks, 

9 Mode, 

10 SupportedGuardrailIntegrations, 

11) 

12from litellm.types.proxy.guardrails.guardrail_hooks.typesafe import ( 

13 TypeSafeGuardrailOptionalParams, 

14) 

15 

16from .typesafe import TypeSafeGuardrail 

17 

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

19 from litellm.types.guardrails import Guardrail, LitellmParams 

20 

21 

22def _coerce_event_hook( 

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

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

25 if isinstance(mode, Mode): 

26 return mode 

27 if isinstance(mode, list): 

28 return [ # mutable-ok: CustomGuardrail event_hook contract wants a list 

29 GuardrailEventHooks(item) for item in mode 

30 ] 

31 return GuardrailEventHooks(mode) 

32 

33 

34def _optional_params(litellm_params: LitellmParams) -> TypeSafeGuardrailOptionalParams: 

35 value: Final = litellm_params.optional_params 

36 if isinstance(value, TypeSafeGuardrailOptionalParams): 

37 return value 

38 if isinstance(value, BaseModel): 

39 return TypeSafeGuardrailOptionalParams.model_validate(value.model_dump()) 

40 return TypeSafeGuardrailOptionalParams() 

41 

42 

43def initialize_guardrail(litellm_params: LitellmParams, guardrail: Guardrail) -> TypeSafeGuardrail: 

44 import litellm 

45 

46 optional_params: Final = _optional_params(litellm_params) 

47 

48 _callback: Final = TypeSafeGuardrail( 

49 api_base=litellm_params.api_base, 

50 api_key=litellm_params.api_key, 

51 model=litellm_params.model, 

52 relevance_threshold=optional_params.relevance_threshold, 

53 min_chars_to_evaluate=optional_params.min_chars_to_evaluate, 

54 max_result_chars_in_state=optional_params.max_result_chars_in_state, 

55 guardrail_name=guardrail["guardrail_name"], 

56 event_hook=_coerce_event_hook(litellm_params.mode), 

57 default_on=litellm_params.default_on or False, 

58 unreachable_fallback=( 

59 litellm_params.unreachable_fallback if "unreachable_fallback" in litellm_params.model_fields_set else None 

60 ), 

61 ) 

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

63 _callback 

64 ) 

65 return _callback 

66 

67 

68guardrail_initializer_registry: Final = { # mutable-ok: guardrail_registry discovery checks isinstance(registry, dict) 

69 SupportedGuardrailIntegrations.TYPESAFE.value: initialize_guardrail, 

70} 

71 

72guardrail_class_registry: Final = { # mutable-ok: guardrail_registry discovery checks isinstance(registry, dict) 

73 SupportedGuardrailIntegrations.TYPESAFE.value: TypeSafeGuardrail, 

74}