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

1from typing import TYPE_CHECKING, Final, Literal, Optional 

2 

3import litellm 

4from litellm.proxy.guardrails.guardrail_hooks.mcp_security.mcp_security_guardrail import ( 

5 MCPSecurityGuardrail, 

6) 

7from litellm.types.guardrails import SupportedGuardrailIntegrations 

8 

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

10 from litellm import Router 

11 from litellm.types.guardrails import Guardrail, LitellmParams 

12 

13 

14def initialize_guardrail( 

15 litellm_params: "LitellmParams", 

16 guardrail: "Guardrail", 

17 llm_router: Optional["Router"] = None, 

18): 

19 guardrail_name: Final = guardrail.get("guardrail_name") 

20 if not guardrail_name: 

21 raise ValueError("MCP Security: guardrail_name is required") 

22 

23 on_violation: Final[Literal["block", "alert"]] = "block" if litellm_params.on_violation == "block" else "alert" 

24 

25 mcp_security_guardrail: Final = MCPSecurityGuardrail( 

26 guardrail_name=guardrail_name, 

27 event_hook=litellm_params.mode, 

28 default_on=litellm_params.default_on or False, 

29 on_violation=on_violation, 

30 ) 

31 

32 litellm.logging_callback_manager.add_litellm_callback(mcp_security_guardrail) 

33 return mcp_security_guardrail 

34 

35 

36guardrail_initializer_registry: Final = { 

37 SupportedGuardrailIntegrations.MCP_SECURITY.value: initialize_guardrail, 

38} 

39 

40guardrail_class_registry: Final = { 

41 SupportedGuardrailIntegrations.MCP_SECURITY.value: MCPSecurityGuardrail, 

42}