phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

(ci/cd) testing with team_id and /user/new

Browse source
This commit is contained in:
Ishaan Jaff 2024-04-20 11:09:34 -07:00
parent abeadadf3e
commit 7ebf2ca4d9
2 changed files with 50 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

7
litellm/proxy/proxy_server.py
View file

@ -6532,6 +6532,13 @@ async def team_member_add(
existing_team_row = await prisma_client.get_data( # type: ignore existing_team_row = await prisma_client.get_data( # type: ignore
team_id=data.team_id, table_name="team", query_type="find_unique" team_id=data.team_id, table_name="team", query_type="find_unique"
) )
if existing_team_row is None:
raise HTTPException(
status_code=404,
detail={
"error": f"Team not found for team_id={getattr(data, 'team_id', None)}"
},
)
new_member = data.member new_member = data.member

48
litellm/tests/test_key_generate_prisma.py
View file

@ -120,6 +120,15 @@ async def test_new_user_response(prisma_client):
await litellm.proxy.proxy_server.prisma_client.connect() await litellm.proxy.proxy_server.prisma_client.connect()
from litellm.proxy.proxy_server import user_api_key_cache from litellm.proxy.proxy_server import user_api_key_cache
await new_team(
NewTeamRequest(
team_id="ishaan-special-team",
),
user_api_key_dict=UserAPIKeyAuth(
user_role="proxy_admin", api_key="sk-1234", user_id="1234"
),
)
_response = await new_user( _response = await new_user(
data=NewUserRequest( data=NewUserRequest(
models=["azure-gpt-3.5"], models=["azure-gpt-3.5"],
@ -999,10 +1008,32 @@ def test_generate_and_update_key(prisma_client):
async def test(): async def test():
await litellm.proxy.proxy_server.prisma_client.connect() await litellm.proxy.proxy_server.prisma_client.connect()
# create team "litellm-core-infra@gmail.com""
print("creating team litellm-core-infra@gmail.com")
await new_team(
NewTeamRequest(
team_id="litellm-core-infra@gmail.com",
),
user_api_key_dict=UserAPIKeyAuth(
user_role="proxy_admin", api_key="sk-1234", user_id="1234"
),
)
await new_team(
NewTeamRequest(
team_id="ishaan-special-team",
),
user_api_key_dict=UserAPIKeyAuth(
user_role="proxy_admin", api_key="sk-1234", user_id="1234"
),
)
request = NewUserRequest( request = NewUserRequest(
metadata={"team": "litellm-team3", "project": "litellm-project3"}, metadata={"project": "litellm-project3"},
team_id="litellm-core-infra@gmail.com", team_id="litellm-core-infra@gmail.com",
) )
key = await new_user(request) key = await new_user(request)
print(key) print(key)
@ -1015,7 +1046,6 @@ def test_generate_and_update_key(prisma_client):
print("\n info for key=", result["info"]) print("\n info for key=", result["info"])
assert result["info"]["max_parallel_requests"] == None assert result["info"]["max_parallel_requests"] == None
assert result["info"]["metadata"] == { assert result["info"]["metadata"] == {
"team": "litellm-team3",
"project": "litellm-project3", "project": "litellm-project3",
} }
assert result["info"]["team_id"] == "litellm-core-infra@gmail.com" assert result["info"]["team_id"] == "litellm-core-infra@gmail.com"
@ -1037,7 +1067,7 @@ def test_generate_and_update_key(prisma_client):
# update the team id # update the team id
response2 = await update_key_fn( response2 = await update_key_fn(
request=Request, request=Request,
data=UpdateKeyRequest(key=generated_key, team_id="ishaan"), data=UpdateKeyRequest(key=generated_key, team_id="ishaan-special-team"),
) )
print("response2=", response2) print("response2=", response2)
@ -1048,11 +1078,10 @@ def test_generate_and_update_key(prisma_client):
print("\n info for key=", result["info"]) print("\n info for key=", result["info"])
assert result["info"]["max_parallel_requests"] == None assert result["info"]["max_parallel_requests"] == None
assert result["info"]["metadata"] == { assert result["info"]["metadata"] == {
"team": "litellm-team3",
"project": "litellm-project3", "project": "litellm-project3",
} }
assert result["info"]["models"] == ["ada", "babbage", "curie", "davinci"] assert result["info"]["models"] == ["ada", "babbage", "curie", "davinci"]
assert result["info"]["team_id"] == "ishaan" assert result["info"]["team_id"] == "ishaan-special-team"
# cleanup - delete key # cleanup - delete key
delete_key_request = KeyRequest(keys=[generated_key]) delete_key_request = KeyRequest(keys=[generated_key])
@ -1941,6 +1970,15 @@ async def test_master_key_hashing(prisma_client):
await litellm.proxy.proxy_server.prisma_client.connect() await litellm.proxy.proxy_server.prisma_client.connect()
from litellm.proxy.proxy_server import user_api_key_cache from litellm.proxy.proxy_server import user_api_key_cache
await new_team(
NewTeamRequest(
team_id="ishaan-special-team",
),
user_api_key_dict=UserAPIKeyAuth(
user_role="proxy_admin", api_key="sk-1234", user_id="1234"
),
)
_response = await new_user( _response = await new_user(
data=NewUserRequest( data=NewUserRequest(
models=["azure-gpt-3.5"], models=["azure-gpt-3.5"],