mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-25 10:44:24 +00:00
feat(proxy_server.py): enable cache controls per key + no-store cache flag
This commit is contained in:
parent
37de964da4
commit
f9acad87dc
8 changed files with 108 additions and 42 deletions
55
litellm/proxy/hooks/cache_control_check.py
Normal file
55
litellm/proxy/hooks/cache_control_check.py
Normal file
|
@ -0,0 +1,55 @@
|
|||
# What this does?
|
||||
## Checks if key is allowed to use the cache controls passed in to the completion() call
|
||||
|
||||
from typing import Optional
|
||||
import litellm
|
||||
from litellm.caching import DualCache
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from fastapi import HTTPException
|
||||
import json, traceback
|
||||
|
||||
|
||||
class CacheControlCheck(CustomLogger):
|
||||
# Class variables or attributes
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def print_verbose(self, print_statement):
|
||||
if litellm.set_verbose is True:
|
||||
print(print_statement) # noqa
|
||||
|
||||
async def async_pre_call_hook(
|
||||
self,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
cache: DualCache,
|
||||
data: dict,
|
||||
call_type: str,
|
||||
):
|
||||
try:
|
||||
self.print_verbose(f"Inside Cache Control Check Pre-Call Hook")
|
||||
allowed_cache_controls = user_api_key_dict.allowed_cache_controls
|
||||
|
||||
if (allowed_cache_controls is None) or (
|
||||
len(allowed_cache_controls) == 0
|
||||
): # assume empty list to be nullable - https://github.com/prisma/prisma/issues/847#issuecomment-546895663
|
||||
return
|
||||
|
||||
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 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:
|
||||
traceback.print_exc()
|
Loading…
Add table
Add a link
Reference in a new issue