From 2130a61b6eabbb22f1013e71e724ff2ebedd97e8 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Wed, 24 Jan 2024 17:56:00 -0800 Subject: [PATCH 1/4] (feat) add cache_key in spend_log --- litellm/proxy/_types.py | 3 +-- litellm/proxy/schema.prisma | 1 + litellm/proxy/utils.py | 5 +++++ litellm/tests/test_key_generate_prisma.py | 4 ++++ schema.prisma | 1 + 5 files changed, 12 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 8a059c5077..670cefcf2f 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -343,8 +343,7 @@ class LiteLLM_SpendLogs(LiteLLMBase): endTime: Union[str, datetime, None] user: Optional[str] = "" modelParameters: Optional[Json] = {} - messages: Optional[Json] = [] - response: Optional[Json] = {} usage: Optional[Json] = {} metadata: Optional[Json] = {} cache_hit: Optional[str] = "False" + cache_key: Optional[str] = None diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 441c3515fb..f06d42ba5b 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -58,4 +58,5 @@ model LiteLLM_SpendLogs { usage Json @default("{}") metadata Json @default("{}") cache_hit String @default("") + cache_key String @default("") } \ No newline at end of file diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 15f230a6a4..d49ace138a 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -995,6 +995,10 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time): if api_key is not None and isinstance(api_key, str) and api_key.startswith("sk-"): # hash the api_key api_key = hash_token(api_key) + from litellm.caching import Cache + + c = Cache() + cache_key = c.get_cache_key(**kwargs) if "headers" in metadata and "authorization" in metadata["headers"]: metadata["headers"].pop( @@ -1013,6 +1017,7 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time): "modelParameters": optional_params, "usage": usage, "metadata": metadata, + "cache_key": cache_key, } json_fields = [ diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py index 49f091cd6f..f7f1d0a919 100644 --- a/litellm/tests/test_key_generate_prisma.py +++ b/litellm/tests/test_key_generate_prisma.py @@ -763,6 +763,10 @@ def test_call_with_key_over_budget(prisma_client): assert spend_log.request_id == request_id assert spend_log.spend == float("2e-05") assert spend_log.model == "chatgpt-v-2" + assert ( + spend_log.cache_key + == "a61ae14fe4a8b8014a61e6ae01a100c8bc6770ac37c293242afed954bc69207d" + ) # use generated key to auth in result = await user_api_key_auth(request=request, api_key=bearer_token) diff --git a/schema.prisma b/schema.prisma index 0dd11eb644..72d14e13bd 100644 --- a/schema.prisma +++ b/schema.prisma @@ -61,4 +61,5 @@ model LiteLLM_SpendLogs { usage Json @default("{}") metadata Json @default("{}") cache_hit String @default("") + cache_key String @default("") } \ No newline at end of file From bf851ef19a47aeb1c76a56f16fba7cb29d873f2a Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Wed, 24 Jan 2024 18:34:22 -0800 Subject: [PATCH 2/4] (fix) use litellm.cache for getting key --- litellm/proxy/proxy_server.py | 2 -- litellm/proxy/utils.py | 9 ++++----- 2 files changed, 4 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 512e956b04..eaa2373cc3 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1382,8 +1382,6 @@ async def initialize( verbose_proxy_logger.setLevel( level=logging.DEBUG ) # set proxy logs to debug - litellm.set_verbose = True - dynamic_config = {"general": {}, user_model: {}} if config: ( diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index d49ace138a..812157ca0e 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -995,15 +995,14 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time): if api_key is not None and isinstance(api_key, str) and api_key.startswith("sk-"): # hash the api_key api_key = hash_token(api_key) - from litellm.caching import Cache - - c = Cache() - cache_key = c.get_cache_key(**kwargs) - if "headers" in metadata and "authorization" in metadata["headers"]: metadata["headers"].pop( "authorization" ) # do not store the original `sk-..` api key in the db + if litellm.cache is not None: + cache_key = litellm.cache.get_cache_key(**kwargs) + else: + cache_key = "Cache OFF" payload = { "request_id": id, From 2f3765a03f5ff659c8292edf3e0491297e477fe7 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Wed, 24 Jan 2024 18:51:39 -0800 Subject: [PATCH 3/4] (fix) log cache hits on SpendLogs table --- litellm/proxy/proxy_server.py | 6 ++++++ litellm/proxy/utils.py | 4 ++++ 2 files changed, 10 insertions(+) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index eaa2373cc3..e1ca25e130 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -593,6 +593,12 @@ async def track_cost_callback( "user_api_key_user_id", None ) + if kwargs.get("cache_hit", False) == True: + response_cost = 0.0 + verbose_proxy_logger.info( + f"Cache Hit: response_cost {response_cost}, for user_id {user_id}" + ) + verbose_proxy_logger.info( f"response_cost {response_cost}, for user_id {user_id}" ) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 812157ca0e..25c5c82ce5 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1003,6 +1003,10 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time): cache_key = litellm.cache.get_cache_key(**kwargs) else: cache_key = "Cache OFF" + if cache_hit == True: + import time + + id = f"{id}_cache_hit{time.time()}" # SpendLogs does not allow duplicate request_id payload = { "request_id": id, From 6bc715cf85f9d9e18dc67ecf2bdffe90ef05d022 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Wed, 24 Jan 2024 18:54:23 -0800 Subject: [PATCH 4/4] (test) logging cache_key in spendLogs --- litellm/tests/test_key_generate_prisma.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py index f7f1d0a919..78fb756b25 100644 --- a/litellm/tests/test_key_generate_prisma.py +++ b/litellm/tests/test_key_generate_prisma.py @@ -716,6 +716,9 @@ def test_call_with_key_over_budget(prisma_client): # update spend using track_cost callback, make 2nd request, it should fail from litellm.proxy.proxy_server import track_cost_callback from litellm import ModelResponse, Choices, Message, Usage + from litellm.caching import Cache + + litellm.cache = Cache() import time request_id = f"chatcmpl-e41836bb-bb8b-4df2-8e70-8f3e160155ac{time.time()}"