From 9dc4127576c842416852154aa352ccdb6451805e Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 4 Apr 2024 10:11:18 -0700 Subject: [PATCH 1/3] v0 return cache key in responses --- litellm/proxy/proxy_server.py | 9 +++++---- litellm/utils.py | 2 ++ 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index a69f28d1089..e0668f466ee 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3432,10 +3432,10 @@ async def chat_completion( # Post Call Processing data["litellm_status"] = "success" # used for alerting - if hasattr(response, "_hidden_params"): - model_id = response._hidden_params.get("model_id", None) or "" - else: - model_id = "" + + hidden_params = getattr(response, "_hidden_params", {}) or {} + model_id = hidden_params.get("model_id", None) or "" + cache_key = hidden_params.get("cache_key", None) or "" if ( "stream" in data and data["stream"] == True @@ -3451,6 +3451,7 @@ async def chat_completion( ) fastapi_response.headers["x-litellm-model-id"] = model_id + fastapi_response.headers["x-litellm-cache-key"] = cache_key ### CALL HOOKS ### - modify outgoing data response = await proxy_logging_obj.post_call_success_hook( diff --git a/litellm/utils.py b/litellm/utils.py index e92f2068a9e..17a31751d96 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -3132,6 +3132,8 @@ def client(original_function): target=logging_obj.success_handler, args=(cached_result, start_time, end_time, cache_hit), ).start() + cache_key = kwargs.get("preset_cache_key", None) + cached_result._hidden_params["cache_key"] = cache_key return cached_result elif ( call_type == CallTypes.aembedding.value From c4cb0afa984be6253e25b88d6cdb1d24b06cddf2 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 4 Apr 2024 10:56:47 -0700 Subject: [PATCH 2/3] feat - delete cache key --- litellm/caching.py | 11 +++++++++ litellm/proxy/proxy_server.py | 45 +++++++++++++++++++++++++++++++++++ 2 files changed, 56 insertions(+) diff --git a/litellm/caching.py b/litellm/caching.py index 197dc0cc574..6bf645f77bd 100644 --- a/litellm/caching.py +++ b/litellm/caching.py @@ -347,6 +347,12 @@ class RedisCache(BaseCache): traceback.print_exc() raise e + async def delete_cache_keys(self, keys): + _redis_client = self.init_async_client() + # keys is a list, unpack it so it gets passed as individual elements to delete + async with _redis_client as redis_client: + await redis_client.delete(*keys) + def client_list(self): client_list = self.redis_client.client_list() return client_list @@ -1408,6 +1414,11 @@ class Cache: return await self.cache.ping() return None + async def delete_cache_keys(self, keys): + if hasattr(self.cache, "delete_cache_keys"): + return await self.cache.delete_cache_keys(keys) + return None + async def disconnect(self): if hasattr(self.cache, "disconnect"): await self.cache.disconnect() diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index e0668f466ee..75e9a4525aa 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -8202,6 +8202,51 @@ async def cache_ping(): ) +@router.post( + "/cache/delete", + tags=["caching"], + dependencies=[Depends(user_api_key_auth)], +) +async def cache_delete(request: Request): + """ + Endpoint for deleting a key from the cache. All responses from litellm proxy have `x-litellm-cache-key` in the headers + + Parameters: + - **keys**: *Optional[List[str]]* - A list of keys to delete from the cache. Example {"keys": ["key1", "key2"]} + + ```shell + curl -X POST "http://0.0.0.0:4000/cache/delete" \ + -H "Authorization: Bearer sk-1234" \ + -d '{"keys": ["key1", "key2"]}' + ``` + + """ + try: + if litellm.cache is None: + raise HTTPException( + status_code=503, detail="Cache not initialized. litellm.cache is None" + ) + + request_data = await request.json() + keys = request_data.get("keys", None) + + if litellm.cache.type == "redis": + await litellm.cache.delete_cache_keys(keys=keys) + return { + "status": "success", + } + else: + raise HTTPException( + status_code=500, + detail=f"Cache type {litellm.cache.type} does not support deleting a key. only `redis` is supported", + ) + except Exception as e: + raise HTTPException( + status_code=500, + detail=f"Cache Delete Failed({str(e)})", + ) + + @router.get( "/cache/redis/info", tags=["caching"], From 7e1d5c81b4633c13e60e1439f50b85320f561f8b Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 4 Apr 2024 11:00:00 -0700 Subject: [PATCH 3/3] return cache key in streming responses --- litellm/proxy/proxy_server.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 75e9a4525aa..25c6cd7361d 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3440,7 +3440,10 @@ async def chat_completion( if ( "stream" in data and data["stream"] == True ): # use generate_responses to stream responses - custom_headers = {"x-litellm-model-id": model_id} + custom_headers = { + "x-litellm-model-id": model_id, + "x-litellm-cache-key": cache_key, + } selected_data_generator = select_data_generator( response=response, user_api_key_dict=user_api_key_dict )