scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-01 07:03:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 129

Merge 83b8e792e4 into a392a39627

Browse source
This commit is contained in:
alexcastgonz 2026-10-31 14:36:54 +08:00 • committed by GitHub
commit 8ebf475fb8
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 177 additions and 26 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

51
litellm/proxy/management_endpoints/organization_endpoints.py
View file

@ -123,6 +123,8 @@ class _ObjectPermissionRow(Protocol):
class _UserTableClient(Protocol):
async def find_unique(self, where: Mapping[str, object]) -> "PrismaUserTable | None": ...
async def find_many(self, where: Mapping[str, object], take: int | None = None) -> "Sequence[PrismaUserTable]": ...
class _BudgetTableClient(Protocol):
async def create(self, data: Mapping[str, object]) -> "PrismaBudgetTable": ...
@ -1413,16 +1415,12 @@ async def find_member_if_email(user_email: str, prisma_client: PrismaClient) ->
"error": f"Unique user not found for user_email={user_email}. Potential duplicate OR non-existent user_email in LiteLLM_UserTable. Use 'user_id' instead."
},
)
try:
existing_user_email_row: Final = await UserRepository(prisma_client).table.find_unique(
where={"user_email": user_email}
)
except Exception:
existing_user_email_rows: Final = await _table(UserRepository(prisma_client)).find_many(
where={"user_email": user_email}, take=2
)
if len(existing_user_email_rows) != 1:
raise not_unique_user_email_error
if existing_user_email_row is None:
raise not_unique_user_email_error
existing_user_email_row_pydantic: Final = LiteLLM_UserTable.model_validate(existing_user_email_row.model_dump())
return existing_user_email_row_pydantic
return LiteLLM_UserTable.model_validate(existing_user_email_rows[0].model_dump())
@router.patch(
@ -1642,25 +1640,20 @@ async def add_member_to_organization(
try:
user_object: LiteLLM_UserTable | None = None
existing_user_id_row = None
existing_user_email_row = None
## Check if user exists in LiteLLM_UserTable - user exists - either the user_id or user_email is in LiteLLM_UserTable
if member.user_id is not None:
existing_user_id_row = await _table(UserRepository(prisma_client)).find_unique(
where={"user_id": member.user_id}
)
if existing_user_id_row is None and member.user_email is not None:
try:
existing_user_email_row = await UserRepository(prisma_client).table.find_unique(
where={"user_email": member.user_email}
)
except Exception as e:
raise ValueError(
f"Potential NON-Existent or Duplicate user email in DB: Error finding a unique instance of user_email={member.user_email} in LiteLLM_UserTable.: {e}"
)
existing_user_email_rows: Final = (
await _table(UserRepository(prisma_client)).find_many(where={"user_email": member.user_email}, take=2)
if existing_user_id_row is None and member.user_email is not None
else ()
)
## If user does not exist, create a new user
if existing_user_id_row is None and existing_user_email_row is None:
if existing_user_id_row is None and not existing_user_email_rows:
# Create a new user - since user does not exist
user_id: Final[str] = member.user_id or str(uuid.uuid4())
new_user_defaults: Final = get_new_internal_user_defaults(
@ -1671,16 +1664,20 @@ async def add_member_to_organization(
_returned_user = await prisma_client.insert_data(data=new_user_defaults, table_name="user")
if _returned_user is not None:
user_object = LiteLLM_UserTable.model_validate(_returned_user.model_dump())
elif existing_user_email_row is not None and (
len(existing_user_email_row) # pyright: ignore[reportArgumentType] # find_unique yields a row, not a list
> 1
):
elif len(existing_user_email_rows) > 1:
raise HTTPException(
status_code=400,
detail={"error": "Multiple users with this email found in db. Please use 'user_id' instead."},
)
elif existing_user_email_row is not None:
user_object = LiteLLM_UserTable.model_validate(existing_user_email_row.model_dump())
elif existing_user_email_rows and member.user_id is not None:
raise HTTPException(
status_code=400,
detail={
"error": f"user_email '{member.user_email}' and user_id '{member.user_id}' do not belong to the same user."
},
)
elif existing_user_email_rows:
user_object = LiteLLM_UserTable.model_validate(existing_user_email_rows[0].model_dump())
elif existing_user_id_row is not None:
user_object = LiteLLM_UserTable.model_validate(existing_user_id_row.model_dump())
else:
@ -1707,5 +1704,7 @@ async def add_member_to_organization(
)
return user_object, organization_membership
except HTTPException:
raise
except Exception as e:
raise ValueError(f"Error adding member={member} to organization={organization_id}: {e}")

152
tests/unit/proxy/management_endpoints/test_organization_endpoints.py
View file

@ -1575,6 +1575,158 @@ async def test_delete_organization_evicts_the_cache_of_the_keys_it_deletes(monke
assert jwt_table.rows == (kept_row,)
class _PrismaLikeUserTable:
"""Rejects find_unique on user_email like Prisma does, because the column is not @unique"""
def __init__(self, rows: tuple) -> None:
self.rows = rows
self.most_rows_read = 0
async def find_unique(self, where: Mapping[str, object]):
if set(where) != {"user_id"}:
raise ValueError(f"Could not find field at `findUniqueLiteLLM_UserTable.where.{next(iter(where))}`")
return next((row for row in self.rows if row.user_id == where["user_id"]), None)
async def find_many(self, where: Mapping[str, object], take: int | None = None):
matches = [row for row in self.rows if all(getattr(row, k) == v for k, v in where.items())]
read = matches if take is None else matches[:take]
self.most_rows_read = max(self.most_rows_read, len(read))
return read
def _org_prisma_with_users(monkeypatch, users: tuple):
from datetime import datetime
from litellm.proxy._types import LiteLLM_OrganizationMembershipTable, LiteLLM_UserTable
async def create_membership(data):
return LiteLLM_OrganizationMembershipTable(
user_id=data["user_id"],
organization_id=data["organization_id"],
user_role=data["user_role"],
created_at=datetime(2024, 1, 1),
updated_at=datetime(2024, 1, 1),
)
async def insert_user(data, table_name):
return LiteLLM_UserTable(user_id=data["user_id"], user_email=data.get("user_email"), user_role="internal_user")
mock_prisma = SimpleNamespace(
db=SimpleNamespace(
litellm_organizationtable=SimpleNamespace(find_unique=AsyncMock(return_value=SimpleNamespace())),
litellm_usertable=_PrismaLikeUserTable(users),
litellm_organizationmembership=SimpleNamespace(create=create_membership),
),
insert_data=insert_user,
)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
monkeypatch.setattr(
"litellm.proxy.management_endpoints.organization_endpoints._verify_org_access",
AsyncMock(),
)
return mock_prisma
async def _add_org_member(member: dict):
from litellm.proxy._types import LitellmUserRoles, OrganizationMemberAddRequest, UserAPIKeyAuth
from litellm.proxy.management_endpoints.organization_endpoints import organization_member_add
return await organization_member_add(
data=OrganizationMemberAddRequest(organization_id="org-1", member=member),
http_request=MagicMock(),
user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN),
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"member, expected_user_id, expected_email",
[
({"user_email": "alice@example.com", "role": "internal_user"}, "alice-id", "alice@example.com"),
({"user_id": "new-id", "user_email": "bob@example.com", "role": "org_admin"}, "new-id", "bob@example.com"),
({"user_email": "bob@example.com", "role": "internal_user"}, None, "bob@example.com"),
],
ids=["email-of-existing-user", "unknown-id-with-new-email", "new-email-only"],
)
async def test_organization_member_add_resolves_member_by_email(member, expected_user_id, expected_email, monkeypatch):
from litellm.proxy._types import LiteLLM_UserTable
alice: Final = LiteLLM_UserTable(user_id="alice-id", user_email="alice@example.com", user_role="internal_user")
_org_prisma_with_users(monkeypatch, (alice,))
response = await _add_org_member(member)
user = response.updated_users[0]
membership = response.updated_organization_memberships[0]
assert user.user_email == expected_email, response
assert membership.user_id == user.user_id, response
assert membership.user_role == member["role"], response
if expected_user_id is None:
assert user.user_id != alice.user_id, response
else:
assert user.user_id == expected_user_id, response
@pytest.mark.asyncio
@pytest.mark.parametrize(
"member, expected_error",
[
({"user_email": "dup@example.com", "role": "internal_user"}, "Multiple users with this email"),
(
{"user_id": "new-id", "user_email": "alice@example.com", "role": "internal_user"},
"user_email 'alice@example.com' and user_id 'new-id' do not belong to the same user",
),
],
ids=["email-shared-by-several-users", "unknown-id-with-another-users-email"],
)
async def test_organization_member_add_rejects_an_ambiguous_member_with_400(member, expected_error, monkeypatch):
from litellm.proxy._types import LiteLLM_UserTable, ProxyException
users: Final = (
LiteLLM_UserTable(user_id="alice-id", user_email="alice@example.com", user_role="internal_user"),
*(
LiteLLM_UserTable(user_id=user_id, user_email="dup@example.com", user_role="internal_user")
for user_id in ("dup-1", "dup-2", "dup-3")
),
)
prisma = _org_prisma_with_users(monkeypatch, users)
with pytest.raises(ProxyException) as exc:
await _add_org_member(member)
assert int(exc.value.code) == 400, exc.value.message
assert expected_error in str(exc.value.message)
assert prisma.db.litellm_usertable.most_rows_read <= 2, "an email lookup must not load every duplicate row"
@pytest.mark.asyncio
@pytest.mark.parametrize(
"rows, expected_user_id",
[
((("alice-id", "alice@example.com"),), "alice-id"),
((), None),
((("a", "alice@example.com"), ("b", "alice@example.com")), None),
],
ids=["unique", "missing", "duplicate"],
)
async def test_find_member_if_email_resolves_only_a_unique_email(rows, expected_user_id, monkeypatch):
from litellm.proxy._types import LiteLLM_UserTable
from litellm.proxy.management_endpoints.organization_endpoints import find_member_if_email
prisma = _org_prisma_with_users(
monkeypatch,
tuple(
LiteLLM_UserTable(user_id=user_id, user_email=email, user_role="internal_user") for user_id, email in rows
),
)
if expected_user_id is None:
with pytest.raises(HTTPException) as exc:
await find_member_if_email("alice@example.com", prisma)
assert exc.value.status_code == 400
else:
assert (await find_member_if_email("alice@example.com", prisma)).user_id == expected_user_id
@pytest.mark.asyncio
async def test_delete_organization_is_refused_while_one_of_its_teams_holds_a_ptu_share(monkeypatch):
"""Deleting an organization deletes its teams, so a team still holding a PTU share stops it the