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

feat(handle_jwt.py): add support for 'team_id_default

Browse source 
allows admin to set a default team id for spend-tracking + permissions
This commit is contained in:
Krrish Dholakia 2024-05-15 21:33:35 -07:00
parent da2ea0ba04
commit ed4315af38
4 changed files with 36 additions and 10 deletions
Download patch file Download diff file Expand all files Collapse all files

23
litellm/tests/test_jwt.py
View file

@ -368,9 +368,14 @@ async def test_team_token_output(prisma_client, audience):
@pytest.mark.parametrize("audience", [None, "litellm-proxy"])
@pytest.mark.parametrize("team_id_set", [True, False])
@pytest.mark.parametrize(
"team_id_set, default_team_id", [(True, None), (False, "1234")]
)
@pytest.mark.parametrize("user_id_upsert", [True, False])
@pytest.mark.asyncio
async def test_user_token_output(prisma_client, audience, team_id_set):
async def test_user_token_output(
prisma_client, audience, team_id_set, default_team_id, user_id_upsert
):
"""
- If user required, check if it exists
- fail initial request (when user doesn't exist)
@ -433,6 +438,8 @@ async def test_user_token_output(prisma_client, audience, team_id_set):
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth()
jwt_handler.litellm_jwtauth.user_id_jwt_field = "sub"
jwt_handler.litellm_jwtauth.team_id_default = default_team_id
if team_id_set:
jwt_handler.litellm_jwtauth.team_id_jwt_field = "client_id"
@ -515,6 +522,16 @@ async def test_user_token_output(prisma_client, audience, team_id_set):
),
user_api_key_dict=result,
)
if default_team_id is not None:
await new_team(
data=NewTeamRequest(
team_id=default_team_id,
tpm_limit=100,
rpm_limit=99,
models=["gpt-3.5-turbo", "gpt-4"],
),
user_api_key_dict=result,
)
except Exception as e:
pytest.fail(f"This should not fail - {str(e)}")
@ -555,7 +572,7 @@ async def test_user_token_output(prisma_client, audience, team_id_set):
## 6. ASSERT USER_API_KEY_AUTH format (used for tpm/rpm limiting in parallel_request_limiter.py AND cost tracking)
if team_id_set:
if team_id_set or default_team_id is not None:
assert team_result.team_tpm_limit == 100
assert team_result.team_rpm_limit == 99
assert team_result.team_models == ["gpt-3.5-turbo", "gpt-4"]