From a2f1d2ee526b098041601fac779eef145e72cf92 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Fri, 16 Feb 2024 15:44:34 -0800 Subject: [PATCH 1/7] (feat) set key-model budgets --- litellm/proxy/_types.py | 2 ++ litellm/proxy/proxy_server.py | 5 +++++ litellm/proxy/schema.prisma | 2 ++ schema.prisma | 2 ++ 4 files changed, 11 insertions(+) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 372b953e08d..b791f0dd928 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -383,6 +383,8 @@ class LiteLLM_VerificationToken(LiteLLMBase): budget_reset_at: Optional[datetime] = None allowed_cache_controls: Optional[list] = [] permissions: Dict = {} + model_spend: Dict = {} + model_max_budget: Dict = {} class UserAPIKeyAuth( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 37f55072e53..6866b142a51 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1635,6 +1635,7 @@ async def generate_key_helper_fn( key_alias: Optional[str] = None, allowed_cache_controls: Optional[list] = [], permissions: Optional[dict] = {}, + model_max_budget: Optional[dict] = {}, ): global prisma_client, custom_db_client, user_api_key_cache @@ -1668,6 +1669,8 @@ async def generate_key_helper_fn( config_json = json.dumps(config) permissions_json = json.dumps(permissions) metadata_json = json.dumps(metadata) + model_max_budget_json = json.dumps(model_max_budget) + user_id = user_id or str(uuid.uuid4()) user_role = user_role or "app_user" tpm_limit = tpm_limit @@ -1710,6 +1713,7 @@ async def generate_key_helper_fn( "budget_reset_at": key_reset_at, "allowed_cache_controls": allowed_cache_controls, "permissions": permissions_json, + "model_max_budget": model_max_budget_json, } if ( general_settings.get("allow_user_auth", False) == True @@ -3059,6 +3063,7 @@ async def generate_key_fn( - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. - metadata: Optional[dict] - Metadata for key, store information for key. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } - permissions: Optional[dict] - key-specific permissions. Currently just used for turning off pii masking (if connected). Example - {"pii": false} + - model_max_budget: Optional[dict] - key-specific model budget in USD. Example - {"text-davinci-002": 0.5, "gpt-3.5-turbo": 0.5}. IF null or {} then no model specific budget. Returns: - key: (str) The generated api key diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 5a57b880842..df840a9ee1b 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -64,6 +64,8 @@ model LiteLLM_VerificationToken { budget_duration String? budget_reset_at DateTime? allowed_cache_controls String[] @default([]) + model_spend Json @default("{}") + model_max_budget Json @default("{}") } // store proxy config.yaml diff --git a/schema.prisma b/schema.prisma index 5a57b880842..df840a9ee1b 100644 --- a/schema.prisma +++ b/schema.prisma @@ -64,6 +64,8 @@ model LiteLLM_VerificationToken { budget_duration String? budget_reset_at DateTime? allowed_cache_controls String[] @default([]) + model_spend Json @default("{}") + model_max_budget Json @default("{}") } // store proxy config.yaml From 2c9d142e420502ccb7505032063a6cf67e4fd78b Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Fri, 16 Feb 2024 16:15:24 -0800 Subject: [PATCH 2/7] (feat) track key spend per model --- litellm/proxy/proxy_server.py | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 6866b142a51..b49890fe73d 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -980,10 +980,22 @@ async def update_database( # Calculate the new cost by adding the existing cost and response_cost new_spend = existing_spend + response_cost - verbose_proxy_logger.debug(f"new cost: {new_spend}") + # track cost per model, for the given key + spend_per_model = existing_spend_obj.model_spend or {} + current_model = kwargs.get("model") + if current_model is not None and spend_per_model is not None: + if spend_per_model.get(current_model) is None: + spend_per_model[current_model] = response_cost + else: + spend_per_model[current_model] += response_cost + + verbose_proxy_logger.debug( + f"new cost: {new_spend}, new spend per model: {spend_per_model}" + ) # Update the cost column for the given token await prisma_client.update_data( - token=token, data={"spend": new_spend} + token=token, + data={"spend": new_spend, "model_spend": spend_per_model}, ) valid_token = user_api_key_cache.get_cache(key=token) From d65c6d38692979dbb4aaaa438f4fcfddee22c02e Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Fri, 16 Feb 2024 16:32:17 -0800 Subject: [PATCH 3/7] (feat) track spend key-model, user-model, team-model --- litellm/proxy/proxy_server.py | 24 +++++++++++++++++++++++- litellm/proxy/schema.prisma | 4 ++++ schema.prisma | 4 ++++ 3 files changed, 31 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index b49890fe73d..dda0d59c961 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -937,6 +937,16 @@ async def update_database( # Calculate the new cost by adding the existing cost and response_cost existing_spend_obj.spend = existing_spend + response_cost + # track cost per model, for the given user + spend_per_model = existing_spend_obj.model_spend or {} + current_model = kwargs.get("model") + if current_model is not None and spend_per_model is not None: + if spend_per_model.get(current_model) is None: + spend_per_model[current_model] = response_cost + else: + spend_per_model[current_model] += response_cost + existing_spend_obj.model_spend = spend_per_model + valid_token = user_api_key_cache.get_cache(key=id) if valid_token is not None and isinstance(valid_token, dict): user_api_key_cache.set_cache( @@ -1001,6 +1011,7 @@ async def update_database( valid_token = user_api_key_cache.get_cache(key=token) if valid_token is not None: valid_token.spend = new_spend + valid_token.model_spend = spend_per_model user_api_key_cache.set_cache(key=token, value=valid_token) elif custom_db_client is not None: # Fetch the existing cost for the given token @@ -1080,10 +1091,21 @@ async def update_database( # Calculate the new cost by adding the existing cost and response_cost new_spend = existing_spend + response_cost + # track cost per model, for the given team + spend_per_model = existing_spend_obj.model_spend or {} + current_model = kwargs.get("model") + if current_model is not None and spend_per_model is not None: + if spend_per_model.get(current_model) is None: + spend_per_model[current_model] = response_cost + else: + spend_per_model[current_model] += response_cost + verbose_proxy_logger.debug(f"new cost: {new_spend}") # Update the cost column for the given token await prisma_client.update_data( - team_id=team_id, data={"spend": new_spend}, table_name="team" + team_id=team_id, + data={"spend": new_spend, "model_spend": spend_per_model}, + table_name="team", ) elif custom_db_client is not None: diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index df840a9ee1b..101cf9b7f09 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -24,6 +24,8 @@ model LiteLLM_TeamTable { budget_reset_at DateTime? created_at DateTime @default(now()) @map("created_at") updated_at DateTime @default(now()) @updatedAt @map("updated_at") + model_spend Json @default("{}") + model_max_budget Json @default("{}") } // Track spend, rate limit, budget Users @@ -41,6 +43,8 @@ model LiteLLM_UserTable { budget_duration String? budget_reset_at DateTime? allowed_cache_controls String[] @default([]) + model_spend Json @default("{}") + model_max_budget Json @default("{}") } // Generate Tokens for Proxy diff --git a/schema.prisma b/schema.prisma index df840a9ee1b..101cf9b7f09 100644 --- a/schema.prisma +++ b/schema.prisma @@ -24,6 +24,8 @@ model LiteLLM_TeamTable { budget_reset_at DateTime? created_at DateTime @default(now()) @map("created_at") updated_at DateTime @default(now()) @updatedAt @map("updated_at") + model_spend Json @default("{}") + model_max_budget Json @default("{}") } // Track spend, rate limit, budget Users @@ -41,6 +43,8 @@ model LiteLLM_UserTable { budget_duration String? budget_reset_at DateTime? allowed_cache_controls String[] @default([]) + model_spend Json @default("{}") + model_max_budget Json @default("{}") } // Generate Tokens for Proxy From e8dcf8fa130c6af87922b7c8b9911e5e0868d66e Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Fri, 16 Feb 2024 17:00:23 -0800 Subject: [PATCH 4/7] (fix) setting model_max_budget --- litellm/proxy/proxy_server.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index dda0d59c961..7d30c959f17 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1763,6 +1763,11 @@ async def generate_key_helper_fn( saved_token["metadata"] = json.loads(saved_token["metadata"]) if isinstance(saved_token["permissions"], str): saved_token["permissions"] = json.loads(saved_token["permissions"]) + if isinstance(saved_token["model_max_budget"], str): + saved_token["model_max_budget"] = json.loads( + saved_token["model_max_budget"] + ) + if saved_token.get("expires", None) is not None and isinstance( saved_token["expires"], datetime ): From e76a3c5ce5f5729d4b06de710036fdf4b91f5f53 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Fri, 16 Feb 2024 17:47:22 -0800 Subject: [PATCH 5/7] (fix) _types for model_max_budget --- litellm/proxy/_types.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index b791f0dd928..be2cdd6ef5f 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -155,6 +155,9 @@ class GenerateKeyRequest(GenerateRequestBase): aliases: Optional[dict] = {} config: Optional[dict] = {} permissions: Optional[dict] = {} + model_max_budget: Optional[dict] = ( + {} + ) # {"gpt-4": 5.0, "gpt-3.5-turbo": 5.0}, defaults to {} class GenerateKeyResponse(GenerateKeyRequest): @@ -167,7 +170,13 @@ class GenerateKeyResponse(GenerateKeyRequest): def set_model_info(cls, values): if values.get("token") is not None: values.update({"key": values.get("token")}) - dict_fields = ["metadata", "aliases", "config", "permissions"] + dict_fields = [ + "metadata", + "aliases", + "config", + "permissions", + "model_max_budget", + ] for field in dict_fields: value = values.get(field) if value is not None and isinstance(value, str): From 62d0c54cfbc5749544c0d6ef08ac164d95b4893f Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Fri, 16 Feb 2024 18:18:35 -0800 Subject: [PATCH 6/7] (fix) issue with storing model max budget --- litellm/proxy/_types.py | 2 ++ litellm/proxy/proxy_server.py | 5 ++++- 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index be2cdd6ef5f..5d74b9dedeb 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -421,6 +421,8 @@ class LiteLLM_UserTable(LiteLLMBase): user_id: str max_budget: Optional[float] spend: float = 0.0 + model_max_budget: Optional[Dict] = {} + model_spend: Optional[Dict] = {} user_email: Optional[str] models: list = [] diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 7d30c959f17..72bb3b25af9 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -940,6 +940,7 @@ async def update_database( # track cost per model, for the given user spend_per_model = existing_spend_obj.model_spend or {} current_model = kwargs.get("model") + if current_model is not None and spend_per_model is not None: if spend_per_model.get(current_model) is None: spend_per_model[current_model] = response_cost @@ -953,7 +954,9 @@ async def update_database( key=id, value=existing_spend_obj.json() ) - verbose_proxy_logger.debug(f"new cost: {existing_spend_obj.spend}") + verbose_proxy_logger.debug( + f"user - new cost: {existing_spend_obj.spend}, user_id: {id}" + ) data_list.append(existing_spend_obj) # Update the cost column for the given user id From 1652e894c1cac5ce655fa2cc600479c6dbedcf2e Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Sat, 17 Feb 2024 15:34:55 -0800 Subject: [PATCH 7/7] (docs) set budget per model --- docs/my-website/docs/proxy/virtual_keys.md | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/docs/my-website/docs/proxy/virtual_keys.md b/docs/my-website/docs/proxy/virtual_keys.md index 2be4b95c1f4..83994701c49 100644 --- a/docs/my-website/docs/proxy/virtual_keys.md +++ b/docs/my-website/docs/proxy/virtual_keys.md @@ -93,6 +93,7 @@ Request Params: - `config`: *Optional[dict]* - any key-specific configs, overrides config in config.yaml - `spend`: *Optional[int]* - Amount spent by key. Default is 0. Will be updated by proxy whenever key is used. https://docs.litellm.ai/docs/proxy/virtual_keys#managing-auth---tracking-spend - `max_budget`: *Optional[float]* - Specify max budget for a given key. +- `model_max_budget`: *Optional[dict[str, float]]* - Specify max budget for each model, `model_max_budget={"gpt4": 0.5, "gpt-5": 0.01}` - `max_parallel_requests`: *Optional[int]* - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. - `metadata`: *Optional[dict]* - Metadata for key, store information for key. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } @@ -676,8 +677,6 @@ general_settings: ### [BETA] Dynamo DB -Only live in `v1.16.21.dev1`. - #### Step 1. Save keys to env ```shell