File size: 2,117 Bytes
e3278e4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 |
# What this does?
## Checks if key is allowed to use the cache controls passed in to the completion() call
from fastapi import HTTPException
from litellm import verbose_logger
from litellm._logging import verbose_proxy_logger
from litellm.caching.caching import DualCache
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy._types import UserAPIKeyAuth
class _PROXY_CacheControlCheck(CustomLogger):
# Class variables or attributes
def __init__(self):
pass
async def async_pre_call_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict,
call_type: str,
):
try:
verbose_proxy_logger.debug("Inside Cache Control Check Pre-Call Hook")
allowed_cache_controls = user_api_key_dict.allowed_cache_controls
if data.get("cache", None) is None:
return
cache_args = data.get("cache", None)
if isinstance(cache_args, dict):
for k, v in cache_args.items():
if (
(allowed_cache_controls is not None)
and (isinstance(allowed_cache_controls, list))
and (
len(allowed_cache_controls) > 0
) # assume empty list to be nullable - https://github.com/prisma/prisma/issues/847#issuecomment-546895663
and k not in allowed_cache_controls
):
raise HTTPException(
status_code=403,
detail=f"Not allowed to set {k} as a cache control. Contact admin to change permissions.",
)
else: # invalid cache
return
except HTTPException as e:
raise e
except Exception as e:
verbose_logger.exception(
"litellm.proxy.hooks.cache_control_check.py::async_pre_call_hook(): Exception occured - {}".format(
str(e)
)
)
|