From 73938080f26eaa1f8b71c1c596cb3aace2eb920d Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Thu, 18 Jan 2024 13:16:25 -0800 Subject: [PATCH] (feat) track - api_key in spendLogs --- litellm/proxy/_types.py | 5 ++- litellm/proxy/schema.prisma | 3 +- litellm/proxy/utils.py | 82 +++++++++++++++++++++++++++++++++++++ litellm/utils.py | 2 + schema.prisma | 5 ++- 5 files changed, 92 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 9bc6b09b1..21629cb81 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -317,14 +317,15 @@ class LiteLLM_UserTable(LiteLLMBase): class LiteLLM_SpendLogs(LiteLLMBase): request_id: str + api_key: str + model: Optional[str] = "" call_type: str + spend: Optional[float] = 0.0 startTime: Union[str, datetime, None] endTime: Union[str, datetime, None] - model: Optional[str] = "" user: Optional[str] = "" modelParameters: Optional[Json] = {} messages: Optional[Json] = [] - spend: Optional[float] = 0.0 response: Optional[Json] = {} usage: Optional[Json] = {} metadata: Optional[Json] = {} diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 9049f953d..2e40a3204 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -36,13 +36,14 @@ model LiteLLM_Config { model LiteLLM_SpendLogs { request_id String @unique call_type String + api_key String @default ("") + spend Float @default(0.0) startTime DateTime // Assuming start_time is a DateTime field endTime DateTime // Assuming end_time is a DateTime field model String @default("") user String @default("") modelParameters Json @default("{}")// Assuming optional_params is a JSON field messages Json @default("[]") - spend Float @default(0.0) response Json @default("{}") usage Json @default("{}") metadata Json @default("{}") diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index f24412a84..23b66f22d 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -779,3 +779,85 @@ async def send_email(sender_name, sender_email, receiver_email, subject, html): except Exception as e: print_verbose("An error occurred while sending the email:", str(e)) + + +def hash_token(token: str): + import hashlib + + # Hash the string using SHA-256 + hashed_token = hashlib.sha256(token.encode()).hexdigest() + + return hashed_token + + +def get_logging_payload(kwargs, response_obj, start_time, end_time): + from litellm.proxy._types import LiteLLM_SpendLogs + from pydantic import Json + import uuid + + if kwargs == None: + kwargs = {} + # standardize this function to be used across, s3, dynamoDB, langfuse logging + litellm_params = kwargs.get("litellm_params", {}) + metadata = ( + litellm_params.get("metadata", {}) or {} + ) # if litellm_params['metadata'] == None + messages = kwargs.get("messages") + optional_params = kwargs.get("optional_params", {}) + call_type = kwargs.get("call_type", "litellm.completion") + cache_hit = kwargs.get("cache_hit", False) + usage = response_obj["usage"] + id = response_obj.get("id", str(uuid.uuid4())) + api_key = metadata.get("user_api_key", "") + if api_key is not None and type(api_key) == str: + # hash the api_key + api_key = hash_token(api_key) + + payload = { + "request_id": id, + "call_type": call_type, + "api_key": api_key, + "cache_hit": cache_hit, + "startTime": start_time, + "endTime": end_time, + "model": kwargs.get("model", ""), + "user": kwargs.get("user", ""), + "modelParameters": optional_params, + "messages": messages, + "response": response_obj, + "usage": usage, + "metadata": metadata, + } + + json_fields = [ + field + for field, field_type in LiteLLM_SpendLogs.__annotations__.items() + if field_type == Json or field_type == Optional[Json] + ] + str_fields = [ + field + for field, field_type in LiteLLM_SpendLogs.__annotations__.items() + if field_type == str or field_type == Optional[str] + ] + datetime_fields = [ + field + for field, field_type in LiteLLM_SpendLogs.__annotations__.items() + if field_type == datetime + ] + + for param in json_fields: + if param in payload and type(payload[param]) != Json: + if type(payload[param]) == litellm.ModelResponse: + payload[param] = payload[param].model_dump_json() + if type(payload[param]) == litellm.EmbeddingResponse: + payload[param] = payload[param].model_dump_json() + elif type(payload[param]) == litellm.Usage: + payload[param] = payload[param].model_dump_json() + else: + payload[param] = json.dumps(payload[param]) + + for param in str_fields: + if param in payload and type(payload[param]) != str: + payload[param] = str(payload[param]) + + return payload diff --git a/litellm/utils.py b/litellm/utils.py index 77f0c331e..e8740e8aa 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8442,10 +8442,12 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time): cache_hit = kwargs.get("cache_hit", False) usage = response_obj["usage"] id = response_obj.get("id", str(uuid.uuid4())) + api_key = metadata.get("user_api_key", "") payload = { "request_id": id, "call_type": call_type, + "api_key": api_key, "cache_hit": cache_hit, "startTime": start_time, "endTime": end_time, diff --git a/schema.prisma b/schema.prisma index a07dcad08..31eae05c2 100644 --- a/schema.prisma +++ b/schema.prisma @@ -35,16 +35,17 @@ model LiteLLM_Config { model LiteLLM_SpendLogs { request_id String @unique + api_key String @default ("") call_type String + spend Float @default(0.0) startTime DateTime // Assuming start_time is a DateTime field endTime DateTime // Assuming end_time is a DateTime field model String @default("") user String @default("") modelParameters Json @default("{}")// Assuming optional_params is a JSON field messages Json @default("[]") - spend Float @default(0.0) response Json @default("{}") usage Json @default("{}") metadata Json @default("{}") cache_hit String @default("") -} \ No newline at end of file +}