From bf0c0c08c8b75d31e19ffcda122fed18cfb02f0f Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Sat, 25 May 2024 14:32:08 -0700 Subject: [PATCH 1/7] feat - set budget duration --- litellm/proxy/_types.py | 5 +++++ litellm/proxy/proxy_server.py | 5 +---- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index e8b3e6572..1f951605c 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -561,7 +561,12 @@ class TeamBase(LiteLLMBase): metadata: Optional[dict] = None tpm_limit: Optional[int] = None rpm_limit: Optional[int] = None + + # Budget fields max_budget: Optional[float] = None + budget_duration: Optional[str] = None + budget_reset_at: Optional[datetime] = None + models: list = [] blocked: bool = False diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 904a414c2..189b6cb88 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -8033,12 +8033,9 @@ async def new_team( _model_id = model_dict.id - ## ADD TO TEAM TABLE + ## ADD TO TEAM TABLE -> the user api key settings are overriding the data we passed i complete_team_data = LiteLLM_TeamTable( **data.json(), - max_parallel_requests=user_api_key_dict.max_parallel_requests, - budget_duration=user_api_key_dict.budget_duration, - budget_reset_at=user_api_key_dict.budget_reset_at, model_id=_model_id, ) From a03d13feddd92f67e40c97f20eaa7da6e595a5c7 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Sat, 25 May 2024 14:38:05 -0700 Subject: [PATCH 2/7] fix - allow setting budget duration and reset at --- litellm/proxy/proxy_server.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 189b6cb88..4bac6ff3c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -8033,7 +8033,16 @@ async def new_team( _model_id = model_dict.id - ## ADD TO TEAM TABLE -> the user api key settings are overriding the data we passed i + ## ADD TO TEAM TABLE + # Check budget_duration and budget_reset_at + if data.budget_duration: + duration_s = _duration_in_seconds(duration=data.budget_duration) + reset_at = datetime.now(timezone.utc) + timedelta(seconds=duration_s) + + # set the budget duration and budget_reset_at in DB + data.budget_duration = duration_s + data.budget_reset_at = reset_at + complete_team_data = LiteLLM_TeamTable( **data.json(), model_id=_model_id, From 6f65df0b1adc8e5495f3a675c9279a94f0983881 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Sat, 25 May 2024 14:57:29 -0700 Subject: [PATCH 3/7] feat - update budget duration --- litellm/proxy/_types.py | 2 ++ litellm/proxy/proxy_server.py | 11 ++++++++++- 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 1f951605c..7be1eb03f 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -612,6 +612,8 @@ class UpdateTeamRequest(LiteLLMBase): max_budget: Optional[float] = None models: Optional[list] = None blocked: Optional[bool] = None + budget_duration: Optional[str] = None + budget_reset_at: Optional[datetime] = None class DeleteTeamRequest(LiteLLMBase): diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 4bac6ff3c..136d37955 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -8035,7 +8035,7 @@ async def new_team( ## ADD TO TEAM TABLE # Check budget_duration and budget_reset_at - if data.budget_duration: + if data.budget_duration is not None: duration_s = _duration_in_seconds(duration=data.budget_duration) reset_at = datetime.now(timezone.utc) + timedelta(seconds=duration_s) @@ -8125,6 +8125,15 @@ async def update_team( detail={"error": f"Team not found, passed team_id={data.team_id}"}, ) + # Check budget_duration and budget_reset_at + if data.budget_duration is not None: + duration_s = _duration_in_seconds(duration=data.budget_duration) + reset_at = datetime.now(timezone.utc) + timedelta(seconds=duration_s) + + # set the budget duration and budget_reset_at in DB + data.budget_duration = duration_s + data.budget_reset_at = reset_at + updated_kv = data.json(exclude_none=True) team_row = await prisma_client.update_data( update_key_values=updated_kv, From 22d8eaab931ed07591ff2c2efa5dbd40f4379485 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Sat, 25 May 2024 15:20:49 -0700 Subject: [PATCH 4/7] test - new_team --- litellm/proxy/_types.py | 1 - litellm/proxy/proxy_server.py | 16 +++----- litellm/tests/test_key_generate_prisma.py | 48 +++++++++++++++++++++++ 3 files changed, 54 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 7be1eb03f..11e358a2f 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -565,7 +565,6 @@ class TeamBase(LiteLLMBase): # Budget fields max_budget: Optional[float] = None budget_duration: Optional[str] = None - budget_reset_at: Optional[datetime] = None models: list = [] blocked: bool = False diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 136d37955..6c6a01f45 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -8034,20 +8034,16 @@ async def new_team( _model_id = model_dict.id ## ADD TO TEAM TABLE - # Check budget_duration and budget_reset_at - if data.budget_duration is not None: - duration_s = _duration_in_seconds(duration=data.budget_duration) - reset_at = datetime.now(timezone.utc) + timedelta(seconds=duration_s) - - # set the budget duration and budget_reset_at in DB - data.budget_duration = duration_s - data.budget_reset_at = reset_at - complete_team_data = LiteLLM_TeamTable( **data.json(), - model_id=_model_id, ) + # If budget_duration is set, set `budget_reset_at` + if complete_team_data.budget_duration is not None: + duration_s = _duration_in_seconds(duration=complete_team_data.budget_duration) + reset_at = datetime.now(timezone.utc) + timedelta(seconds=duration_s) + complete_team_data.budget_reset_at = reset_at + team_row = await prisma_client.insert_data( data=complete_team_data.json(exclude_none=True), table_name="team" ) diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py index c1521e856..3619ce8f4 100644 --- a/litellm/tests/test_key_generate_prisma.py +++ b/litellm/tests/test_key_generate_prisma.py @@ -2137,3 +2137,51 @@ async def test_reset_spend_authentication(prisma_client): "Tried to access route=/global/spend/reset, which is only for MASTER KEY" in e.message ) + + +@pytest.mark.asyncio() +async def test_create_update_team(prisma_client): + """ + - Set max_budget, budget_duration, max_budget, tpm_limit, rpm_limit + - Assert response has correct values + + - Update max_budget, budget_duration, max_budget, tpm_limit, rpm_limit + - Assert response has correct values + + - Call team_info and assert response has correct values + """ + print("prisma client=", prisma_client) + + master_key = "sk-1234" + + setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) + setattr(litellm.proxy.proxy_server, "master_key", master_key) + import datetime + + await litellm.proxy.proxy_server.prisma_client.connect() + from litellm.proxy.proxy_server import user_api_key_cache + + _team_id = "test-team_{}".format(uuid.uuid4()) + response = await new_team( + NewTeamRequest( + team_id=_team_id, + max_budget=20, + budget_duration="30d", + tpm_limit=20, + rpm_limit=20, + ), + user_api_key_dict=UserAPIKeyAuth( + user_role="proxy_admin", api_key="sk-1234", user_id="1234" + ), + ) + + print("RESPONSE from new_team", response) + + assert response["team_id"] == _team_id + assert response["max_budget"] == 20 + assert response["tpm_limit"] == 20 + assert response["rpm_limit"] == 20 + assert response["budget_duration"] == "30d" + assert response["budget_reset_at"] is not None and isinstance( + response["budget_reset_at"], datetime.datetime + ) From e2d3c0c846ef11dda6daca592705b640642405d7 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Sat, 25 May 2024 15:31:25 -0700 Subject: [PATCH 5/7] fix - updating team --- litellm/proxy/_types.py | 1 - litellm/proxy/proxy_server.py | 8 +++--- litellm/tests/test_key_generate_prisma.py | 30 +++++++++++++++++++++++ 3 files changed, 34 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 11e358a2f..0ea6fff8e 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -612,7 +612,6 @@ class UpdateTeamRequest(LiteLLMBase): models: Optional[list] = None blocked: Optional[bool] = None budget_duration: Optional[str] = None - budget_reset_at: Optional[datetime] = None class DeleteTeamRequest(LiteLLMBase): diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 6c6a01f45..8bfc3314a 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -8121,16 +8121,16 @@ async def update_team( detail={"error": f"Team not found, passed team_id={data.team_id}"}, ) + updated_kv = data.json(exclude_none=True) + # Check budget_duration and budget_reset_at if data.budget_duration is not None: duration_s = _duration_in_seconds(duration=data.budget_duration) reset_at = datetime.now(timezone.utc) + timedelta(seconds=duration_s) - # set the budget duration and budget_reset_at in DB - data.budget_duration = duration_s - data.budget_reset_at = reset_at + # set the budget_reset_at in DB + updated_kv["budget_reset_at"] = reset_at - updated_kv = data.json(exclude_none=True) team_row = await prisma_client.update_data( update_key_values=updated_kv, data=updated_kv, diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py index 3619ce8f4..dda5c365c 100644 --- a/litellm/tests/test_key_generate_prisma.py +++ b/litellm/tests/test_key_generate_prisma.py @@ -52,6 +52,7 @@ from litellm.proxy.proxy_server import ( user_info, info_key_fn, new_team, + update_team, chat_completion, completion, embeddings, @@ -73,6 +74,7 @@ from litellm.proxy._types import ( UpdateKeyRequest, GenerateKeyRequest, NewTeamRequest, + UpdateTeamRequest, UserAPIKeyAuth, LiteLLM_UpperboundKeyGenerateParams, ) @@ -2185,3 +2187,31 @@ async def test_create_update_team(prisma_client): assert response["budget_reset_at"] is not None and isinstance( response["budget_reset_at"], datetime.datetime ) + + # updating team budget duration and reset at + + response = await update_team( + UpdateTeamRequest( + team_id=_team_id, + max_budget=30, + budget_duration="2d", + tpm_limit=30, + rpm_limit=30, + ), + user_api_key_dict=UserAPIKeyAuth( + user_role="proxy_admin", api_key="sk-1234", user_id="1234" + ), + ) + + print("RESPONSE from update_team", response) + _updated_info = response["data"] + _updated_info = dict(_updated_info) + + assert _updated_info["team_id"] == _team_id + assert _updated_info["max_budget"] == 30 + assert _updated_info["tpm_limit"] == 30 + assert _updated_info["rpm_limit"] == 30 + assert _updated_info["budget_duration"] == "2d" + assert _updated_info["budget_reset_at"] is not None and isinstance( + _updated_info["budget_reset_at"], datetime.datetime + ) From 59cb306d47c3fe8daedbb8911cbaf7c33528c8a3 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Sat, 25 May 2024 15:36:20 -0700 Subject: [PATCH 6/7] fix - __team_info --- litellm/tests/test_key_generate_prisma.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py index dda5c365c..528410ca6 100644 --- a/litellm/tests/test_key_generate_prisma.py +++ b/litellm/tests/test_key_generate_prisma.py @@ -50,6 +50,7 @@ from litellm.proxy.proxy_server import ( spend_key_fn, view_spend_logs, user_info, + team_info, info_key_fn, new_team, update_team, @@ -2215,3 +2216,20 @@ async def test_create_update_team(prisma_client): assert _updated_info["budget_reset_at"] is not None and isinstance( _updated_info["budget_reset_at"], datetime.datetime ) + + # now hit team_info + response = await team_info(team_id=_team_id) + + print("RESPONSE from team_info", response) + + _team_info = response["team_info"] + _team_info = dict(_team_info) + + assert _team_info["team_id"] == _team_id + assert _team_info["max_budget"] == 30 + assert _team_info["tpm_limit"] == 30 + assert _team_info["rpm_limit"] == 30 + assert _team_info["budget_duration"] == "2d" + assert _team_info["budget_reset_at"] is not None and isinstance( + _team_info["budget_reset_at"], datetime.datetime + ) From 6ec0efab502b8da0f5d8fa9dfc25124df6026e39 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Sat, 25 May 2024 15:42:52 -0700 Subject: [PATCH 7/7] fix - test team_alias --- litellm/proxy/proxy_server.py | 1 + 1 file changed, 1 insertion(+) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 8bfc3314a..7182444c4 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -8036,6 +8036,7 @@ async def new_team( ## ADD TO TEAM TABLE complete_team_data = LiteLLM_TeamTable( **data.json(), + model_id=_model_id, ) # If budget_duration is set, set `budget_reset_at`