phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-26 19:24:27 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

(feat) working create budgets per key

Browse source
This commit is contained in:
ishaan-jaff 2024-01-22 11:58:20 -08:00
parent 770aa2e5c6
commit 1640a9d58e
2 changed files with 12 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

3
litellm/proxy/_types.py
View file

@ -128,6 +128,7 @@ class GenerateKeyRequest(LiteLLMBase):
aliases: Optional[dict] = {} aliases: Optional[dict] = {}
config: Optional[dict] = {} config: Optional[dict] = {}
spend: Optional[float] = 0 spend: Optional[float] = 0
max_budget: Optional[float] = None
user_id: Optional[str] = None user_id: Optional[str] = None
team_id: Optional[str] = None team_id: Optional[str] = None
max_parallel_requests: Optional[int] = None max_parallel_requests: Optional[int] = None
@ -145,6 +146,7 @@ class UpdateKeyRequest(LiteLLMBase):
aliases: Optional[dict] = None aliases: Optional[dict] = None
config: Optional[dict] = None config: Optional[dict] = None
spend: Optional[float] = None spend: Optional[float] = None
max_budget: Optional[float] = None
user_id: Optional[str] = None user_id: Optional[str] = None
max_parallel_requests: Optional[int] = None max_parallel_requests: Optional[int] = None
metadata: Optional[dict] = None metadata: Optional[dict] = None
@ -162,6 +164,7 @@ class UserAPIKeyAuth(LiteLLMBase): # the expected response object for user api
aliases: dict = {} aliases: dict = {}
config: dict = {} config: dict = {}
spend: Optional[float] = 0 spend: Optional[float] = 0
max_budget: Optional[float] = None
user_id: Optional[str] = None user_id: Optional[str] = None
max_parallel_requests: Optional[int] = None max_parallel_requests: Optional[int] = None
duration: str = "1h" duration: str = "1h"

10
litellm/proxy/proxy_server.py
View file

@ -1142,7 +1142,8 @@ async def generate_key_helper_fn(
aliases: dict, aliases: dict,
config: dict, config: dict,
spend: float, spend: float,
max_budget: Optional[float] = None, key_max_budget: Optional[float] = None, # key_max_budget is used to Budget Per key
max_budget: Optional[float] = None, # max_budget is used to Budget Per user
token: Optional[str] = None, token: Optional[str] = None,
user_id: Optional[str] = None, user_id: Optional[str] = None,
team_id: Optional[str] = None, team_id: Optional[str] = None,
@ -1215,6 +1216,7 @@ async def generate_key_helper_fn(
"aliases": aliases_json, "aliases": aliases_json,
"config": config_json, "config": config_json,
"spend": spend, "spend": spend,
"max_budget": key_max_budget,
"user_id": user_id, "user_id": user_id,
"team_id": team_id, "team_id": team_id,
"max_parallel_requests": max_parallel_requests, "max_parallel_requests": max_parallel_requests,
@ -2177,6 +2179,7 @@ async def generate_key_fn(
- aliases: Optional[dict] - Any alias mappings, on top of anything in the config.yaml model list. - https://docs.litellm.ai/docs/proxy/virtual_keys#managing-auth---upgradedowngrade-models - aliases: Optional[dict] - Any alias mappings, on top of anything in the config.yaml model list. - https://docs.litellm.ai/docs/proxy/virtual_keys#managing-auth---upgradedowngrade-models
- config: Optional[dict] - any key-specific configs, overrides config in config.yaml - 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 - 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.
- 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. - 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" } - metadata: Optional[dict] - Metadata for key, store information for key. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" }
@ -2201,6 +2204,11 @@ async def generate_key_fn(
) )
data_json = data.json() # type: ignore data_json = data.json() # type: ignore
# if we get max_budget passed to /key/generate, then use it as key_max_budget. Since generate_key_helper_fn is used to make new users
if "max_budget" in data_json:
data_json["key_max_budget"] = data_json.pop("max_budget", None)
response = await generate_key_helper_fn(**data_json) response = await generate_key_helper_fn(**data_json)
return GenerateKeyResponse( return GenerateKeyResponse(
key=response["token"], expires=response["expires"], user_id=response["user_id"] key=response["token"], expires=response["expires"], user_id=response["user_id"]