scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

Merge pull request #23557 from BerriAI/litellm_ui_key_org-2

Browse source 
[Feature] Allow Setting organization_id on Key Update
This commit is contained in:
yuneng-jiang 2026-03-13 16:10:35 -07:00 • committed by GitHub
commit dca231e3f6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 321 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
litellm/proxy/_types.py
View file

@ -1007,6 +1007,7 @@ class UpdateKeyRequest(KeyRequestBase):
temp_budget_expiry: Optional[datetime] = None
auto_rotate: Optional[bool] = None
rotation_interval: Optional[str] = None
organization_id: Optional[str] = None
@model_validator(mode="after")
def validate_temp_budget(self) -> "UpdateKeyRequest":

48
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

@ -908,6 +908,11 @@ async def _check_team_key_limits(
keys = await prisma_client.db.litellm_verificationtoken.find_many(
where={"team_id": team_table.team_id},
)
# Exclude the key being updated to avoid double-counting its limits.
# key.token is the SHA-256 hash stored in DB; data.key is the raw key string.
if isinstance(data, UpdateKeyRequest):
hashed_key = hash_token(data.key)
keys = [key for key in keys if key.token != hashed_key]
check_team_key_model_specific_limits(
keys=keys,
team_table=team_table,
@ -1062,6 +1067,11 @@ async def _check_org_key_limits(
keys = await prisma_client.db.litellm_verificationtoken.find_many(
where={"organization_id": org_table.organization_id},
)
# Exclude the key being updated to avoid double-counting its limits.
# key.token is the SHA-256 hash stored in DB; data.key is the raw key string.
if isinstance(data, UpdateKeyRequest):
hashed_key = hash_token(data.key)
keys = [key for key in keys if key.token != hashed_key]
check_org_key_model_specific_limits(
keys=keys,
org_table=org_table,
@ -1811,6 +1821,7 @@ async def update_key_fn(
- user_id: Optional[str] - User ID associated with key
- team_id: Optional[str] - Team ID associated with key
- agent_id: Optional[str] - The agent id associated with the key.
- organization_id: Optional[str] - The organization id of the key.
- budget_id: Optional[str] - The budget id associated with the key. Created by calling `/budget/new`.
- models: Optional[list] - Model_name's a user is allowed to call
- tags: Optional[List[str]] - Tags for organizing keys (Enterprise only)
@ -1930,11 +1941,14 @@ async def update_key_fn(
user_api_key_cache=user_api_key_cache,
)
# Only check team limits if key has a team_id
# Check team limits if key has a team_id (from request or existing key)
team_obj: Optional[LiteLLM_TeamTableCachedObj] = None
if data.team_id is not None:
_team_id_to_check = data.team_id or getattr(
existing_key_row, "team_id", None
)
if _team_id_to_check is not None:
team_obj = await get_team_object(
team_id=data.team_id,
team_id=_team_id_to_check,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
check_db_only=True,
@ -1961,6 +1975,34 @@ async def update_key_fn(
user_api_key_cache=user_api_key_cache,
)
# Check org key limits only when throughput-related fields or organization_id change
_org_id_to_check = data.organization_id or getattr(
existing_key_row, "organization_id", None
)
_throughput_fields_changed = (
data.organization_id is not None
or data.tpm_limit is not None
or data.rpm_limit is not None
or data.tpm_limit_type is not None
or data.rpm_limit_type is not None
)
if _org_id_to_check is not None and _throughput_fields_changed:
org_table = await get_org_object(
org_id=_org_id_to_check,
user_api_key_cache=user_api_key_cache,
prisma_client=prisma_client,
)
if org_table is None:
raise HTTPException(
status_code=400,
detail=f"Organization not found for organization_id={_org_id_to_check}",
)
await _check_org_key_limits(
org_table=org_table,
data=data,
prisma_client=prisma_client,
)
# if team change - check if this is possible
if is_different_team(data=data, existing_key_row=existing_key_row):
if llm_router is None:

275
tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
View file

@ -1836,6 +1836,65 @@ async def test_check_team_key_limits_rpm_overallocation():
)
@pytest.mark.asyncio
async def test_check_team_key_limits_on_update_excludes_self():
"""
Test that _check_team_key_limits excludes the key being updated from the
allocated totals. Without this, the key's current limits would be
double-counted: once from find_many and once from data.tpm_limit/rpm_limit.
"""
from litellm.proxy._types import hash_token as _ht
# The key being updated is returned by find_many with its current limits.
# In the DB, token is stored as a SHA-256 hash of the raw key.
self_key = MagicMock()
self_key.token = _ht("sk-self-team-key")
self_key.tpm_limit = 6000
self_key.rpm_limit = 600
self_key.metadata = {}
# Another key in the team
other_key = MagicMock()
other_key.token = _ht("sk-other-team-key")
other_key.tpm_limit = 3000
other_key.rpm_limit = 300
other_key.metadata = {}
mock_prisma_client = AsyncMock()
mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
return_value=[self_key, other_key]
)
team_table = LiteLLM_TeamTableCachedObj(
team_id="test-team-self",
team_alias="test-team",
tpm_limit=10000,
rpm_limit=1000,
max_budget=100.0,
spend=0.0,
models=[],
blocked=False,
members_with_roles=[],
)
# Updating the key to 7000 TPM. Other key uses 3000, so total = 10000 <= 10000.
# Without the fix, this would be 6000 (self) + 3000 (other) + 7000 = 16000 > 10000.
data = UpdateKeyRequest(
key="sk-self-team-key",
tpm_limit=7000,
rpm_limit=700,
tpm_limit_type="guaranteed_throughput",
rpm_limit_type="guaranteed_throughput",
)
# Should not raise - the key's own limits should be excluded from the sum
await _check_team_key_limits(
team_table=team_table,
data=data,
prisma_client=mock_prisma_client,
)
@pytest.mark.asyncio
async def test_check_team_key_limits_no_team_limits():
"""
@ -6859,3 +6918,219 @@ class TestValidateKeyAliasFormat:
_validate_key_alias_format(alias)
assert str(exc.value.code) == "400"
assert "Invalid key_alias format" in str(exc.value.message)
@pytest.mark.asyncio
async def test_check_org_key_limits_on_update_within_bounds():
"""
Test that _check_org_key_limits works with UpdateKeyRequest when updating
a key's TPM/RPM limits within organization bounds.
"""
mock_prisma_client = AsyncMock()
mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
return_value=[]
)
org_table = LiteLLM_OrganizationTable(
organization_id="test-org-update-1",
organization_alias="test-org",
budget_id="budget-123",
models=["gpt-4"],
created_by="admin",
updated_by="admin",
litellm_budget_table=LiteLLM_BudgetTable(
budget_id="budget-123",
tpm_limit=20000,
rpm_limit=2000,
),
)
data = UpdateKeyRequest(
key="sk-test-key",
tpm_limit=10000,
rpm_limit=1000,
tpm_limit_type="guaranteed_throughput",
rpm_limit_type="guaranteed_throughput",
organization_id="test-org-update-1",
)
# Should not raise any exception
await _check_org_key_limits(
org_table=org_table,
data=data,
prisma_client=mock_prisma_client,
)
mock_prisma_client.db.litellm_verificationtoken.find_many.assert_called_once_with(
where={"organization_id": "test-org-update-1"}
)
@pytest.mark.asyncio
async def test_check_org_key_limits_on_update_overallocation():
"""
Test that _check_org_key_limits raises HTTPException when updating a key
would exceed organization TPM limits.
"""
from litellm.proxy._types import hash_token as _hash_token
existing_key = MagicMock()
existing_key.token = _hash_token("sk-other-key")
existing_key.tpm_limit = 15000
existing_key.rpm_limit = 1500
existing_key.metadata = {}
mock_prisma_client = AsyncMock()
mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
return_value=[existing_key]
)
org_table = LiteLLM_OrganizationTable(
organization_id="test-org-update-2",
organization_alias="test-org",
budget_id="budget-456",
models=["gpt-4"],
created_by="admin",
updated_by="admin",
litellm_budget_table=LiteLLM_BudgetTable(
budget_id="budget-456",
tpm_limit=20000,
rpm_limit=2000,
),
)
data = UpdateKeyRequest(
key="sk-test-key",
tpm_limit=10000, # 15000 + 10000 = 25000 > 20000
tpm_limit_type="guaranteed_throughput",
rpm_limit_type="guaranteed_throughput",
organization_id="test-org-update-2",
)
with pytest.raises(HTTPException) as exc:
await _check_org_key_limits(
org_table=org_table,
data=data,
prisma_client=mock_prisma_client,
)
assert exc.value.status_code == 400
assert "TPM limit" in str(exc.value.detail)
@pytest.mark.asyncio
async def test_check_org_key_limits_on_update_excludes_self():
"""
Test that _check_org_key_limits excludes the key being updated from the
allocated totals. Without this, the key's current limits would be
double-counted: once from find_many and once from data.tpm_limit/rpm_limit.
"""
from litellm.proxy._types import hash_token
# The key being updated is returned by find_many with its current limits.
# In the DB, token is stored as a SHA-256 hash of the raw key.
self_key = MagicMock()
self_key.token = hash_token("sk-test-key")
self_key.tpm_limit = 10000
self_key.rpm_limit = 1000
self_key.metadata = {}
# Another key in the org
other_key = MagicMock()
other_key.token = hash_token("sk-other-key")
other_key.tpm_limit = 5000
other_key.rpm_limit = 500
other_key.metadata = {}
mock_prisma_client = AsyncMock()
mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
return_value=[self_key, other_key]
)
org_table = LiteLLM_OrganizationTable(
organization_id="test-org-self",
organization_alias="test-org",
budget_id="budget-789",
models=["gpt-4"],
created_by="admin",
updated_by="admin",
litellm_budget_table=LiteLLM_BudgetTable(
budget_id="budget-789",
tpm_limit=20000,
rpm_limit=2000,
),
)
# Updating the key to 12000 TPM. Other key uses 5000, so total = 17000 < 20000.
# Without the fix, this would be 10000 (self) + 5000 (other) + 12000 = 27000 > 20000.
data = UpdateKeyRequest(
key="sk-test-key",
tpm_limit=12000,
rpm_limit=1200,
tpm_limit_type="guaranteed_throughput",
rpm_limit_type="guaranteed_throughput",
organization_id="test-org-self",
)
# Should not raise - the key's own limits should be excluded from the sum
await _check_org_key_limits(
org_table=org_table,
data=data,
prisma_client=mock_prisma_client,
)
def test_update_key_skips_org_check_when_no_throughput_fields_changed():
"""
Test that the org limit check guard condition correctly skips validation
when only non-throughput fields change on a key that belongs to an org.
This prevents blocking updates when the org has been deleted.
"""
def _check_throughput_changed(data: UpdateKeyRequest) -> bool:
return (
data.organization_id is not None
or data.tpm_limit is not None
or data.rpm_limit is not None
or data.tpm_limit_type is not None
or data.rpm_limit_type is not None
)
# Updating only key_alias — no throughput fields changed
data = UpdateKeyRequest(key="sk-test-key", key_alias="new-alias")
assert _check_throughput_changed(data) is False
# Updating tpm_limit — throughput field changed
data_with_tpm = UpdateKeyRequest(key="sk-test-key", tpm_limit=5000)
assert _check_throughput_changed(data_with_tpm) is True
# Updating organization_id — org change triggers check
data_with_org = UpdateKeyRequest(
key="sk-test-key", organization_id="new-org"
)
assert _check_throughput_changed(data_with_org) is True
# Updating tpm_limit_type — limit type change triggers check
data_with_tpm_type = UpdateKeyRequest(
key="sk-test-key", tpm_limit_type="guaranteed_throughput"
)
assert _check_throughput_changed(data_with_tpm_type) is True
# Updating rpm_limit_type — limit type change triggers check
data_with_rpm_type = UpdateKeyRequest(
key="sk-test-key", rpm_limit_type="guaranteed_throughput"
)
assert _check_throughput_changed(data_with_rpm_type) is True
def test_update_key_request_has_organization_id():
"""
Test that UpdateKeyRequest accepts organization_id field.
"""
data = UpdateKeyRequest(
key="sk-test-key",
organization_id="test-org-123",
)
assert data.organization_id == "test-org-123"
# Also verify it defaults to None
data_no_org = UpdateKeyRequest(key="sk-test-key")
assert data_no_org.organization_id is None