From 4821fa9201af5ee3da1cf337e3ba34c642ae6eee Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Thu, 18 Jan 2024 10:04:34 -0800 Subject: [PATCH 01/10] (v0) add schema.prisma --- litellm/proxy/_types.py | 16 ++++++++++++++++ litellm/proxy/schema.prisma | 15 +++++++++++++++ schema.prisma | 15 +++++++++++++++ 3 files changed, 46 insertions(+) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 6e858831411..220bf2f919a 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -194,6 +194,7 @@ class DynamoDBArgs(LiteLLMBase): user_table_name: str = "LiteLLM_UserTable" key_table_name: str = "LiteLLM_VerificationToken" config_table_name: str = "LiteLLM_Config" + spend_table_name: str = "LiteLLM_SpendLogs" class ConfigGeneralSettings(LiteLLMBase): @@ -312,3 +313,18 @@ class LiteLLM_UserTable(LiteLLMBase): if values.get("models") is None: values.update({"models", []}) return values + + +class LiteLLM_SpendLogs(LiteLLMBase): + id: str + call_type: str + startTime: Union[str, None] + endTime: Union[str, None] + model: str = "" + user: str = "" + modelParameters: Dict = {} + messages: List[str] = [] + call_cost: float = 0.0 + response: Dict = {} + usage: Dict = {} + metadata: Dict = {} diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index aa45a881865..d2e338bd4cf 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -31,4 +31,19 @@ model LiteLLM_VerificationToken { model LiteLLM_Config { param_name String @id param_value Json? +} + +model LiteLLM_SpendLogs { + request_id String @unique + call_type String + 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("{}") } \ No newline at end of file diff --git a/schema.prisma b/schema.prisma index 704ada42c98..df2c1d0b4b6 100644 --- a/schema.prisma +++ b/schema.prisma @@ -31,4 +31,19 @@ model LiteLLM_VerificationToken { model LiteLLM_Config { param_name String @id param_value Json? +} + +model LiteLLM_SpendLogs { + id String @unique + call_type String + 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("[]") + call_cost Float @default(0.0) + response Json @default("{}") + usage Json @default("{}") + metadata Json @default("{}") } \ No newline at end of file From 4a5f987512e94e46b542c1929c799dea99380aeb Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Thu, 18 Jan 2024 10:09:02 -0800 Subject: [PATCH 02/10] (feat) insert_data to spend table --- litellm/proxy/_types.py | 2 +- litellm/proxy/utils.py | 27 +++++++++++++++++++++++---- schema.prisma | 4 ++-- 3 files changed, 26 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 220bf2f919a..cf8350022b9 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -316,7 +316,7 @@ class LiteLLM_UserTable(LiteLLMBase): class LiteLLM_SpendLogs(LiteLLMBase): - id: str + request_id: str call_type: str startTime: Union[str, None] endTime: Union[str, None] diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index ab1fea463ee..f24412a84cf 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1,7 +1,12 @@ from typing import Optional, List, Any, Literal, Union import os, subprocess, hashlib, importlib, asyncio, copy, json, aiohttp, httpx import litellm, backoff -from litellm.proxy._types import UserAPIKeyAuth, DynamoDBArgs, LiteLLM_VerificationToken +from litellm.proxy._types import ( + UserAPIKeyAuth, + DynamoDBArgs, + LiteLLM_VerificationToken, + LiteLLM_SpendLogs, +) from litellm.caching import DualCache from litellm.proxy.hooks.parallel_request_limiter import MaxParallelRequestsHandler from litellm.proxy.hooks.max_budget_limiter import MaxBudgetLimiter @@ -316,7 +321,7 @@ class PrismaClient: self, key: str, value: Any, - table_name: Literal["users", "keys", "config"], + table_name: Literal["users", "keys", "config", "spend"], ): """ Generic implementation of get data @@ -334,6 +339,10 @@ class PrismaClient: response = await self.db.litellm_config.find_first( # type: ignore where={key: value} # type: ignore ) + elif table_name == "spend": + response = await self.db.l.find_first( # type: ignore + where={key: value} # type: ignore + ) return response except Exception as e: asyncio.create_task( @@ -417,7 +426,7 @@ class PrismaClient: on_backoff=on_backoff, # specifying the function to call on backoff ) async def insert_data( - self, data: dict, table_name: Literal["user", "key", "config"] + self, data: dict, table_name: Literal["user", "key", "config", "spend"] ): """ Add a key to the database. If it already exists, do nothing. @@ -473,8 +482,18 @@ class PrismaClient: ) tasks.append(updated_table_row) - await asyncio.gather(*tasks) + elif table_name == "spend": + db_data = self.jsonify_object(data=data) + new_spend_row = await self.db.litellm_spendlogs.upsert( + where={"request_id": data["request_id"]}, + data={ + "create": {**db_data}, # type: ignore + "update": {}, # don't do anything if it already exists + }, + ) + return new_spend_row + except Exception as e: print_verbose(f"LiteLLM Prisma Client Exception: {e}") asyncio.create_task( diff --git a/schema.prisma b/schema.prisma index df2c1d0b4b6..ed69f67a767 100644 --- a/schema.prisma +++ b/schema.prisma @@ -34,7 +34,7 @@ model LiteLLM_Config { } model LiteLLM_SpendLogs { - id String @unique + request_id String @unique call_type String startTime DateTime // Assuming start_time is a DateTime field endTime DateTime // Assuming end_time is a DateTime field @@ -42,7 +42,7 @@ model LiteLLM_SpendLogs { user String @default("") modelParameters Json @default("{}")// Assuming optional_params is a JSON field messages Json @default("[]") - call_cost Float @default(0.0) + spend Float @default(0.0) response Json @default("{}") usage Json @default("{}") metadata Json @default("{}") From d14d36af9add19434f96a6ef49482c18acc610b3 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Thu, 18 Jan 2024 11:54:15 -0800 Subject: [PATCH 03/10] (v0 ) working - writing /chat/completion spend tracking --- litellm/proxy/_types.py | 27 +++++++------- litellm/proxy/proxy_config.yaml | 4 +-- litellm/proxy/proxy_server.py | 35 ++++++++++++++++-- litellm/proxy/schema.prisma | 1 + litellm/utils.py | 63 +++++++++++++++++++++++++++++++++ schema.prisma | 1 + 6 files changed, 114 insertions(+), 17 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index cf8350022b9..9bc6b09b128 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1,8 +1,8 @@ -from pydantic import BaseModel, Extra, Field, root_validator +from pydantic import BaseModel, Extra, Field, root_validator, Json import enum -from typing import Optional, List, Union, Dict, Literal +from typing import Optional, List, Union, Dict, Literal, Any from datetime import datetime -import uuid, json +import uuid, json, sys, os class LiteLLMBase(BaseModel): @@ -318,13 +318,14 @@ class LiteLLM_UserTable(LiteLLMBase): class LiteLLM_SpendLogs(LiteLLMBase): request_id: str call_type: str - startTime: Union[str, None] - endTime: Union[str, None] - model: str = "" - user: str = "" - modelParameters: Dict = {} - messages: List[str] = [] - call_cost: float = 0.0 - response: Dict = {} - usage: Dict = {} - metadata: Dict = {} + 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] = {} + cache_hit: Optional[str] = "False" diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index 5b87ab775b2..8cd2fcec85f 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -61,8 +61,8 @@ litellm_settings: # setting callback class # callbacks: custom_callbacks.proxy_handler_instance # sets litellm.callbacks = [proxy_handler_instance] -# general_settings: - # master_key: sk-1234 +general_settings: + master_key: sk-1234 # database_type: "dynamo_db" # database_args: { # 👈 all args - https://github.com/BerriAI/litellm/blob/befbcbb7ac8f59835ce47415c128decf37aac328/litellm/proxy/_types.py#L190 # "billing_mode": "PAY_PER_REQUEST", diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d3667892bb1..fdc81f88e0e 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -510,6 +510,7 @@ async def track_cost_callback( global prisma_client, custom_db_client try: # check if it has collected an entire stream response + verbose_proxy_logger.debug(f"Proxy: In track_cost_callback for {kwargs}") verbose_proxy_logger.debug( f"kwargs stream: {kwargs.get('stream', None)} + complete streaming response: {kwargs.get('complete_streaming_response', None)}" ) @@ -546,13 +547,27 @@ async def track_cost_callback( prisma_client is not None or custom_db_client is not None ): await update_database( - token=user_api_key, response_cost=response_cost, user_id=user_id + token=user_api_key, + response_cost=response_cost, + user_id=user_id, + kwargs=kwargs, + completion_response=completion_response, + start_time=start_time, + end_time=end_time, ) except Exception as e: verbose_proxy_logger.debug(f"error in tracking cost callback - {str(e)}") -async def update_database(token, response_cost, user_id=None): +async def update_database( + token, + response_cost, + user_id=None, + kwargs=None, + completion_response=None, + start_time=None, + end_time=None, +): try: verbose_proxy_logger.debug( f"Enters prisma db call, token: {token}; user_id: {user_id}" @@ -622,9 +637,25 @@ async def update_database(token, response_cost, user_id=None): key=token, value={"spend": new_spend}, table_name="key" ) + async def _insert_spend_log_to_db(): + # Helper to generate payload to log + verbose_proxy_logger.debug("inserting spend log to db") + payload = litellm.utils.get_logging_payload( + kwargs=kwargs, + response_obj=completion_response, + start_time=start_time, + end_time=end_time, + ) + + payload["spend"] = response_cost + + if prisma_client is not None: + await prisma_client.insert_data(data=payload, table_name="spend") + tasks = [] tasks.append(_update_user_db()) tasks.append(_update_key_db()) + tasks.append(_insert_spend_log_to_db()) await asyncio.gather(*tasks) except Exception as e: verbose_proxy_logger.debug( diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index d2e338bd4cf..9049f953d75 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -46,4 +46,5 @@ model LiteLLM_SpendLogs { response Json @default("{}") usage Json @default("{}") metadata Json @default("{}") + cache_hit String @default("") } \ No newline at end of file diff --git a/litellm/utils.py b/litellm/utils.py index f7cc5d2a54a..b22a053ff50 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8423,3 +8423,66 @@ def print_args_passed_to_litellm(original_function, args, kwargs): except: # This should always be non blocking pass + + +def get_logging_payload(kwargs, response_obj, start_time, end_time): + from litellm.proxy._types import LiteLLM_SpendLogs + from pydantic import Json + + # 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())) + + payload = { + "request_id": id, + "call_type": call_type, + "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]) == ModelResponse: + payload[param] = payload[param].model_dump_json() + elif type(payload[param]) == 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/schema.prisma b/schema.prisma index ed69f67a767..a07dcad08e7 100644 --- a/schema.prisma +++ b/schema.prisma @@ -46,4 +46,5 @@ model LiteLLM_SpendLogs { response Json @default("{}") usage Json @default("{}") metadata Json @default("{}") + cache_hit String @default("") } \ No newline at end of file From b0e18e92b6cc5eaf9a511b71246459a508e8bb09 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Thu, 18 Jan 2024 12:05:08 -0800 Subject: [PATCH 04/10] (fix) when kwargs==None --- litellm/utils.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/litellm/utils.py b/litellm/utils.py index b22a053ff50..1ff9dbe1c81 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8429,6 +8429,8 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time): from litellm.proxy._types import LiteLLM_SpendLogs from pydantic import Json + if kwargs == None: + kwargs = {} # standardize this function to be used across, s3, dynamoDB, langfuse logging litellm_params = kwargs.get("litellm_params", {}) metadata = ( From 88cdfedf8429d3b705bd48301e4b0e783884854b Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Thu, 18 Jan 2024 12:21:56 -0800 Subject: [PATCH 05/10] (feat) track cost streaming --- litellm/proxy/proxy_server.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index fdc81f88e0e..10efd495b3a 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -531,7 +531,13 @@ async def track_cost_callback( prisma_client is not None or custom_db_client is not None ): await update_database( - token=user_api_key, response_cost=response_cost, user_id=user_id + token=user_api_key, + response_cost=response_cost, + user_id=user_id, + kwargs=kwargs, + completion_response=completion_response, + start_time=start_time, + end_time=end_time, ) elif kwargs["stream"] == False: # for non streaming responses response_cost = litellm.completion_cost( From 4eb9221121d357dafafeed0381fc58747d51ed46 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Thu, 18 Jan 2024 12:29:21 -0800 Subject: [PATCH 06/10] (fix) track EmbeddingResponse cost --- litellm/utils.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/litellm/utils.py b/litellm/utils.py index 1ff9dbe1c81..77f0c331e9b 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8478,6 +8478,8 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time): if param in payload and type(payload[param]) != Json: if type(payload[param]) == ModelResponse: payload[param] = payload[param].model_dump_json() + if type(payload[param]) == EmbeddingResponse: + payload[param] = payload[param].model_dump_json() elif type(payload[param]) == Usage: payload[param] = payload[param].model_dump_json() else: From 5b54bcc712c2d92f9709dc6961d8029ceedd3498 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Thu, 18 Jan 2024 12:39:11 -0800 Subject: [PATCH 07/10] (feat) spendLogs table DynamoDB --- litellm/proxy/db/dynamo_db.py | 21 ++++++++++++++++++++- litellm/proxy/proxy_server.py | 3 +++ 2 files changed, 23 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/db/dynamo_db.py b/litellm/proxy/db/dynamo_db.py index eb1c0852861..83cf6b15724 100644 --- a/litellm/proxy/db/dynamo_db.py +++ b/litellm/proxy/db/dynamo_db.py @@ -131,10 +131,27 @@ class DynamoDBWrapper(CustomDB): raise Exception( f"Failed to create table - {self.database_arguments.config_table_name}.\nPlease create a new table called {self.database_arguments.config_table_name}\nAND set `hash_key` as 'param_name'" ) + + ## Spend + try: + verbose_proxy_logger.debug("DynamoDB Wrapper - Creating Spend Table") + error_occurred = False + table = client.table(self.database_arguments.spend_table_name) + if not await table.exists(): + await table.create( + self.throughput_type, + KeySchema(hash_key=KeySpec("request_id", KeyType.string)), + ) + except Exception as e: + error_occurred = True + if error_occurred == True: + raise Exception( + f"Failed to create table - {self.database_arguments.key_table_name}.\nPlease create a new table called {self.database_arguments.key_table_name}\nAND set `hash_key` as 'token'" + ) verbose_proxy_logger.debug("DynamoDB Wrapper - Done connecting()") async def insert_data( - self, value: Any, table_name: Literal["user", "key", "config"] + self, value: Any, table_name: Literal["user", "key", "config", "spend"] ): from aiodynamo.client import Client from aiodynamo.credentials import Credentials, StaticCredentials @@ -166,6 +183,8 @@ class DynamoDBWrapper(CustomDB): table = client.table(self.database_arguments.key_table_name) elif table_name == "config": table = client.table(self.database_arguments.config_table_name) + elif table_name == "spend": + table = client.table(self.database_arguments.spend_table_name) for k, v in value.items(): if isinstance(v, datetime): diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 10efd495b3a..32e985113e6 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -658,6 +658,9 @@ async def update_database( if prisma_client is not None: await prisma_client.insert_data(data=payload, table_name="spend") + elif custom_db_client is not None: + await custom_db_client.insert_data(payload, table_name="spend") + tasks = [] tasks.append(_update_user_db()) tasks.append(_update_key_db()) From 73938080f26eaa1f8b71c1c596cb3aace2eb920d Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Thu, 18 Jan 2024 13:16:25 -0800 Subject: [PATCH 08/10] (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 9bc6b09b128..21629cb8168 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 9049f953d75..2e40a32045d 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 f24412a84cf..23b66f22d79 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 77f0c331e9b..e8740e8aaec 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 a07dcad08e7..31eae05c2ee 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 +} From 7bdf0a0a63f2733a38cab7e296cec762af82bfcf Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Thu, 18 Jan 2024 13:21:51 -0800 Subject: [PATCH 09/10] (chore) cleanup utils.py --- litellm/utils.py | 69 ------------------------------------------------ 1 file changed, 69 deletions(-) diff --git a/litellm/utils.py b/litellm/utils.py index e8740e8aaec..f7cc5d2a54a 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8423,72 +8423,3 @@ def print_args_passed_to_litellm(original_function, args, kwargs): except: # This should always be non blocking pass - - -def get_logging_payload(kwargs, response_obj, start_time, end_time): - from litellm.proxy._types import LiteLLM_SpendLogs - from pydantic import Json - - 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", "") - - 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]) == ModelResponse: - payload[param] = payload[param].model_dump_json() - if type(payload[param]) == EmbeddingResponse: - payload[param] = payload[param].model_dump_json() - elif type(payload[param]) == 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 From 4294657b99ad5939425adf2bfa339c75ffb2b2fa Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Thu, 18 Jan 2024 13:40:48 -0800 Subject: [PATCH 10/10] (fix) use get_logging_payload --- litellm/proxy/proxy_server.py | 3 ++- litellm/tests/test_key_generate_dynamodb.py | 4 ++++ 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 32e985113e6..8145d19aa36 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -72,6 +72,7 @@ from litellm.proxy.utils import ( ProxyLogging, _cache_user_row, send_email, + get_logging_payload, ) from litellm.proxy.secret_managers.google_kms import load_google_kms import pydantic @@ -646,7 +647,7 @@ async def update_database( async def _insert_spend_log_to_db(): # Helper to generate payload to log verbose_proxy_logger.debug("inserting spend log to db") - payload = litellm.utils.get_logging_payload( + payload = get_logging_payload( kwargs=kwargs, response_obj=completion_response, start_time=start_time, diff --git a/litellm/tests/test_key_generate_dynamodb.py b/litellm/tests/test_key_generate_dynamodb.py index 09f699af7d7..2cfa9c95312 100644 --- a/litellm/tests/test_key_generate_dynamodb.py +++ b/litellm/tests/test_key_generate_dynamodb.py @@ -179,6 +179,10 @@ def test_call_with_key_over_budget(custom_db_client): # 5. Make a call with a key over budget, expect to fail setattr(litellm.proxy.proxy_server, "custom_db_client", custom_db_client) setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") + from litellm._logging import verbose_proxy_logger + import logging + + verbose_proxy_logger.setLevel(logging.DEBUG) try: async def test():