feat(proxy_server.py): enable cache controls per key + no-store cache flag

This commit is contained in:
Krrish Dholakia 2024-01-30 20:46:50 -08:00
parent 37de964da4
commit f9acad87dc
8 changed files with 108 additions and 42 deletions

View 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()