Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/hooks/__init__.py: 79%

1import os 

2from typing import Final, Literal 

3 

4from . import * 

5from .autorouter_baseline_cache import AutoRouterBaselineCache 

6from .cache_control_check import _PROXY_CacheControlCheck 

7from .litellm_skills import SkillsInjectionHook 

8from .max_budget_per_session_limiter import _PROXY_MaxBudgetPerSessionHandler 

9from .max_iterations_limiter import _PROXY_MaxIterationsHandler 

10from .parallel_request_limiter import _PROXY_MaxParallelRequestsHandler 

11from .parallel_request_limiter_v3 import _PROXY_MaxParallelRequestsHandler_v3 

12from .prompt_cache_prediction import PromptCacheObserver 

13from .responses_id_security import ResponsesIDSecurity 

14from .sensitive_data_routing import _PROXY_SensitiveDataRoutingHandler 

15 

16# List of all available hooks that can be enabled. 

17# Defined before the enterprise import below so that any module re-imported 

18# transitively through `enterprise.enterprise_hooks` can resolve `PROXY_HOOKS` 

19# and `get_proxy_hook` from this partially-initialized module without circling. 

20PROXY_HOOKS: Final = { 

21 "parallel_request_limiter": _PROXY_MaxParallelRequestsHandler_v3, 

22 "cache_control_check": _PROXY_CacheControlCheck, 

23 "responses_id_security": ResponsesIDSecurity, 

24 "litellm_skills": SkillsInjectionHook, 

25 "max_iterations_limiter": _PROXY_MaxIterationsHandler, 

26 "max_budget_per_session_limiter": _PROXY_MaxBudgetPerSessionHandler, 

27 "sensitive_data_routing": _PROXY_SensitiveDataRoutingHandler, 

28 "prompt_cache_prediction": PromptCacheObserver, 

29 "autorouter_baseline_cache": AutoRouterBaselineCache, 

30} 

31 

32## FEATURE FLAG HOOKS ## 

33if os.getenv("LEGACY_MULTI_INSTANCE_RATE_LIMITING", "false").lower() == "true": 33 ↛ 34line 33 didn't jump to line 34 because the condition on line 33 was never true

34 PROXY_HOOKS["parallel_request_limiter"] = _PROXY_MaxParallelRequestsHandler 

35 

36 

37def get_proxy_hook( 

38 hook_name: Literal["managed_files", "parallel_request_limiter", "cache_control_check"] | str, 

39): 

40 """ 

41 Factory method to get a proxy hook instance by name 

42 """ 

43 if hook_name not in PROXY_HOOKS: 43 ↛ 44line 43 didn't jump to line 44 because the condition on line 43 was never true

44 raise ValueError(f"Unknown hook: {hook_name}. Available hooks: {list(PROXY_HOOKS.keys())}") 

45 return PROXY_HOOKS[hook_name] 

46 

47 

48### CHECK IF ENTERPRISE HOOKS ARE AVAILABLE #### 

49 

50try: 

51 from enterprise.enterprise_hooks import ENTERPRISE_PROXY_HOOKS 

52except ImportError: 

53 ENTERPRISE_PROXY_HOOKS = {} 

54 

55 

56### update PROXY_HOOKS with ENTERPRISE_PROXY_HOOKS ### 

57 

58PROXY_HOOKS.update(ENTERPRISE_PROXY_HOOKS)