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

1from typing import TYPE_CHECKING, Final, Literal 

2 

3from pydantic import BaseModel 

4 

5import litellm 

6from litellm.types.guardrails import SupportedGuardrailIntegrations 

7 

8from .straiker import StraikerGuardrail 

9 

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

11 from litellm.types.guardrails import Guardrail, LitellmParams 

12 

13 

14class _V3Routing(BaseModel): 

15 api_version: Literal["v1", "v3"] | None = None 

16 agent_ref: str | None = None 

17 client: str | None = None 

18 format_hint: Literal["anthropic.messages", "openai.chat"] | None = None 

19 

20 

21_OPTIONAL_INIT_FIELDS: Final = ( 

22 "timeout", 

23 "max_retries", 

24 "initial_backoff", 

25 "max_backoff", 

26 "unreachable_fallback", 

27 "fail_on_error", 

28 "max_payload_bytes", 

29 "custom_headers", 

30 "metadata", 

31 "verbose", 

32) 

33 

34 

35def _get_config_value(litellm_params: "LitellmParams", optional_params: object, attribute_name: str) -> object: 

36 if optional_params is not None: 

37 if isinstance(optional_params, dict): 

38 value = optional_params.get(attribute_name) 

39 else: 

40 value = getattr(optional_params, attribute_name, None) 

41 if value is not None: 

42 return value 

43 return getattr(litellm_params, attribute_name, None) 

44 

45 

46def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"): 

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

48 api_key: Final = litellm_params.api_key 

49 if not api_key: 

50 raise ValueError("api_key is required for straiker") 

51 

52 api_base: Final = litellm_params.api_base or "https://api.prod.straiker.ai" 

53 default_app: Final = getattr(litellm_params, "default_app", None) or getattr(litellm_params, "source", None) 

54 source: Final = default_app if isinstance(default_app, str) and default_app else "LiteLLM Gateway" 

55 kwargs: Final[dict[str, object]] = { 

56 field: value 

57 for field in _OPTIONAL_INIT_FIELDS 

58 for value in [_get_config_value(litellm_params, optional_params, field)] 

59 if value is not None 

60 } 

61 routing: Final = _V3Routing.model_validate( 

62 { 

63 field: _get_config_value(litellm_params, optional_params, field) 

64 for field in ("api_version", "agent_ref", "client", "format_hint") 

65 } 

66 ) 

67 _callback: Final = StraikerGuardrail( 

68 api_key=api_key, 

69 api_base=api_base if isinstance(api_base, str) else "https://api.prod.straiker.ai", 

70 source=source, 

71 guardrail_name=guardrail.get("guardrail_name", "straiker"), 

72 event_hook=litellm_params.mode, 

73 default_on=litellm_params.default_on, 

74 api_version=routing.api_version, 

75 agent_ref=routing.agent_ref, 

76 client=routing.client, 

77 format_hint=routing.format_hint, 

78 **kwargs, 

79 ) 

80 

81 litellm.logging_callback_manager.add_litellm_callback(_callback) 

82 return _callback 

83 

84 

85guardrail_initializer_registry: Final = { 

86 SupportedGuardrailIntegrations.STRAIKER.value: initialize_guardrail, 

87} 

88 

89guardrail_class_registry: Final = { 

90 SupportedGuardrailIntegrations.STRAIKER.value: StraikerGuardrail, 

91}