scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-04 07:27:51 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 299

Merge f6a80150e7 into e16d79d635

Browse source
This commit is contained in:
Praising Harris Ratnam 2026-10-27 02:24:31 +00:00 • committed by GitHub
commit 39043798bc
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 186 additions and 69 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

53
litellm/proxy/management_endpoints/model_access_group_management_endpoints.py
View file

@ -627,11 +627,8 @@ async def create_model_group(
detail={"error": "Either model_names or model_ids must be provided and non-empty"},
)
# If model_ids is provided, use it (more precise targeting)
use_model_ids: Final = has_model_ids
# Validate model_names exist in router (only if using model_names path)
if not use_model_ids and has_model_names:
# Validate model_names exist in router whenever model_names is provided
if has_model_names:
assert data.model_names is not None
missing_models: Final = await _missing_models_after_read_through(
model_names=data.model_names,
@ -663,21 +660,27 @@ async def create_model_group(
},
)
# Update deployments using the appropriate method
if use_model_ids:
# Update deployments using provided model_ids and/or model_names
updated_dict: Final[dict[str, object]] = {} # mutable-ok: accumulator for unique deployment updates
if has_model_ids:
assert data.model_ids is not None
updated_pairs = await update_specific_deployments_with_access_group(
for pair in await update_specific_deployments_with_access_group(
model_ids=data.model_ids,
access_group=data.access_group,
prisma_client=prisma_client,
)
else:
):
updated_dict[pair[0]] = pair[1]
if has_model_names:
assert data.model_names is not None
updated_pairs = await update_deployments_with_access_group(
for pair in await update_deployments_with_access_group(
model_names=data.model_names,
access_group=data.access_group,
prisma_client=prisma_client,
)
):
updated_dict[pair[0]] = pair[1]
updated_pairs: Final = tuple(updated_dict.items())
models_updated: Final = len(updated_pairs)
live_before_reload: Final = live_model_ids_snapshot()
@ -879,8 +882,6 @@ async def update_access_group(
detail={"error": "Either model_names or model_ids must be provided and non-empty"},
)
use_model_ids: Final = has_model_ids
# Validation: Check if access group exists
try:
access_groups_map: Final = await get_all_access_groups_from_db(prisma_client=prisma_client)
@ -897,8 +898,8 @@ async def update_access_group(
detail={"error": f"Failed to check access group existence: {e}"},
)
# Validation: Check if all new models exist (only if using model_names path)
if not use_model_ids and has_model_names:
# Validation: Check if all new models exist whenever model_names is provided
if has_model_names:
assert data.model_names is not None
missing_models: Final = await _missing_models_after_read_through(
model_names=data.model_names,
@ -926,21 +927,27 @@ async def update_access_group(
]
stripped_pairs: Final = tuple(pair for pair in stripped if pair is not None)
# Step 2: Add access group using the appropriate method
if use_model_ids:
# Step 2: Add access group using provided model_ids and/or model_names
updated_dict: Final[dict[str, object]] = {} # mutable-ok: accumulator for unique deployment updates
if has_model_ids:
assert data.model_ids is not None
updated_pairs = await update_specific_deployments_with_access_group(
for pair in await update_specific_deployments_with_access_group(
model_ids=data.model_ids,
access_group=access_group,
prisma_client=prisma_client,
)
else:
):
updated_dict[pair[0]] = pair[1]
if has_model_names:
assert data.model_names is not None
updated_pairs = await update_deployments_with_access_group(
for pair in await update_deployments_with_access_group(
model_names=data.model_names,
access_group=access_group,
prisma_client=prisma_client,
)
):
updated_dict[pair[0]] = pair[1]
updated_pairs: Final = tuple(updated_dict.items())
models_updated: Final = len(updated_pairs)
# Clear cache and reload models to pick up the access group changes

202
tests/unit/proxy/management_endpoints/test_access_group_management.py
View file

@ -10,7 +10,6 @@ from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from litellm import Router
from litellm.proxy.management_endpoints.model_management_endpoints import (
ReconcileOutcome,
@ -74,7 +73,6 @@ async def test_create_duplicate_access_group_fails():
patch("litellm.proxy.proxy_server.llm_router", mock_router),
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
):
# Should raise 409 Conflict
with pytest.raises(HTTPException) as exc_info:
await create_model_group(data=request_data, user_api_key_dict=mock_user)
@ -103,9 +101,7 @@ async def test_create_access_group_with_model_ids_tags_only_specific_deployments
mock_prisma = MagicMock()
mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[])
mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(
return_value=deploy_a
)
mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=deploy_a)
mock_prisma.db.litellm_proxymodeltable.update = AsyncMock()
mock_user = UserAPIKeyAuth(
@ -126,15 +122,11 @@ async def test_create_access_group_with_model_ids_tags_only_specific_deployments
new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)),
),
):
response = await create_model_group(
data=request_data, user_api_key_dict=mock_user
)
response = await create_model_group(data=request_data, user_api_key_dict=mock_user)
assert response.models_updated == 1
assert response.model_ids == ["deploy-A"]
mock_prisma.db.litellm_proxymodeltable.find_unique.assert_called_once_with(
where={"model_id": "deploy-A"}
)
mock_prisma.db.litellm_proxymodeltable.find_unique.assert_called_once_with(where={"model_id": "deploy-A"})
assert mock_prisma.db.litellm_proxymodeltable.update.call_count == 1
update_call = mock_prisma.db.litellm_proxymodeltable.update.call_args
assert update_call.kwargs["where"] == {"model_id": "deploy-A"}
@ -169,9 +161,7 @@ async def test_create_access_group_with_model_names_tags_all_deployments():
)
mock_prisma = MagicMock()
mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(
side_effect=[[], [deploy_a, deploy_b, deploy_c]]
)
mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(side_effect=[[], [deploy_a, deploy_b, deploy_c]])
mock_prisma.db.litellm_proxymodeltable.update = AsyncMock()
mock_user = UserAPIKeyAuth(
@ -179,9 +169,7 @@ async def test_create_access_group_with_model_names_tags_all_deployments():
user_role=LitellmUserRoles.PROXY_ADMIN,
)
request_data = NewModelGroupRequest(
access_group="production-models", model_names=["gpt-4o"]
)
request_data = NewModelGroupRequest(access_group="production-models", model_names=["gpt-4o"])
with (
patch("litellm.proxy.proxy_server.llm_router", mock_router),
@ -191,9 +179,7 @@ async def test_create_access_group_with_model_names_tags_all_deployments():
new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)),
),
):
response = await create_model_group(
data=request_data, user_api_key_dict=mock_user
)
response = await create_model_group(data=request_data, user_api_key_dict=mock_user)
assert response.models_updated == 3
assert response.model_names == ["gpt-4o"]
@ -201,9 +187,10 @@ async def test_create_access_group_with_model_names_tags_all_deployments():
@pytest.mark.asyncio
async def test_create_access_group_model_ids_takes_priority_over_model_names():
async def test_create_access_group_applies_both_model_names_and_model_ids():
"""
Test that when both model_ids and model_names are provided, model_ids is used.
Test that when both model_ids and model_names are provided, both are applied with deduplication.
Fixes: #47012
"""
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.management_endpoints.model_access_group_management_endpoints import (
@ -214,12 +201,26 @@ async def test_create_access_group_model_ids_takes_priority_over_model_names():
)
deploy_a = MagicMock(model_id="deploy-A", model_name="gpt-4o", model_info={})
deploy_b = MagicMock(model_id="deploy-B", model_name="gpt-4o", model_info={})
mock_router = Router(
model_list=[
{
"model_name": "gpt-4o",
"litellm_params": {"model": "gpt-4o", "api_key": "fake-key"},
"model_info": {"id": "deploy-A", "db_model": True},
},
{
"model_name": "gpt-4o",
"litellm_params": {"model": "gpt-4o", "api_key": "fake-key"},
"model_info": {"id": "deploy-B", "db_model": True},
},
]
)
mock_prisma = MagicMock()
mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[])
mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(
return_value=deploy_a
)
mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(side_effect=[[], [deploy_a, deploy_b]])
mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=deploy_a)
mock_prisma.db.litellm_proxymodeltable.update = AsyncMock()
mock_user = UserAPIKeyAuth(
@ -227,6 +228,7 @@ async def test_create_access_group_model_ids_takes_priority_over_model_names():
user_role=LitellmUserRoles.PROXY_ADMIN,
)
# deploy-A is in model_ids AND in model_names (gpt-4o), deploy-B is only from model_names
request_data = NewModelGroupRequest(
access_group="production-models",
model_names=["gpt-4o"],
@ -234,21 +236,134 @@ async def test_create_access_group_model_ids_takes_priority_over_model_names():
)
with (
patch("litellm.proxy.proxy_server.llm_router", MagicMock()),
patch("litellm.proxy.proxy_server.llm_router", mock_router),
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
patch(
"litellm.proxy.management_endpoints.model_access_group_management_endpoints.clear_cache",
new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)),
),
):
response = await create_model_group(
data=request_data, user_api_key_dict=mock_user
response = await create_model_group(data=request_data, user_api_key_dict=mock_user)
# Both deploy-A and deploy-B should be updated without duplicate pairs
assert response.models_updated == 2
assert response.model_names == ["gpt-4o"]
assert response.model_ids == ["deploy-A"]
@pytest.mark.asyncio
async def test_update_access_group_not_found_raises_404():
"""
Test that updating a non-existent access group raises HTTPException 404.
"""
from fastapi import HTTPException
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.management_endpoints.model_access_group_management_endpoints import (
update_access_group,
)
from litellm.types.proxy.management_endpoints.model_management_endpoints import (
UpdateModelGroupRequest,
)
mock_prisma = MagicMock()
# find_many returns empty list -> access group not found
mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[])
mock_user = UserAPIKeyAuth(
user_id="test_admin",
user_role=LitellmUserRoles.PROXY_ADMIN,
)
request_data = UpdateModelGroupRequest(
model_names=["gpt-4o"],
)
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma):
with pytest.raises(HTTPException) as exc_info:
await update_access_group(
access_group="non-existent-group",
data=request_data,
user_api_key_dict=mock_user,
)
assert exc_info.value.status_code == 404
assert "not found" in str(exc_info.value.detail).lower()
@pytest.mark.asyncio
async def test_update_access_group_applies_both_model_names_and_model_ids():
"""
Test that update_access_group updates deployments for both model_names and model_ids.
"""
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.management_endpoints.model_access_group_management_endpoints import (
update_access_group,
)
from litellm.types.proxy.management_endpoints.model_management_endpoints import (
UpdateModelGroupRequest,
)
existing_deploy = MagicMock(
model_id="deploy-existing",
model_name="gpt-4o",
model_info={"access_groups": ["my-group"]},
)
deploy_a = MagicMock(model_id="deploy-A", model_name="gpt-4o", model_info={})
deploy_b = MagicMock(model_id="deploy-B", model_name="gpt-4o", model_info={})
mock_router = Router(
model_list=[
{
"model_name": "gpt-4o",
"litellm_params": {"model": "gpt-4o", "api_key": "fake-key"},
"model_info": {"id": "deploy-A", "db_model": True},
},
{
"model_name": "gpt-4o",
"litellm_params": {"model": "gpt-4o", "api_key": "fake-key"},
"model_info": {"id": "deploy-B", "db_model": True},
},
]
)
mock_prisma = MagicMock()
# 1. get_all_access_groups_from_db -> [existing_deploy] (has "my-group")
# 2. Step 1: find_many for stripping -> [existing_deploy]
# 3. Step 2: update_deployments_with_access_group -> [deploy_a, deploy_b]
mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(
side_effect=[[existing_deploy], [existing_deploy], [deploy_a, deploy_b]]
)
mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=deploy_a)
mock_prisma.db.litellm_proxymodeltable.update = AsyncMock()
mock_user = UserAPIKeyAuth(
user_id="test_admin",
user_role=LitellmUserRoles.PROXY_ADMIN,
)
request_data = UpdateModelGroupRequest(
model_names=["gpt-4o"],
model_ids=["deploy-A"],
)
with (
patch("litellm.proxy.proxy_server.llm_router", mock_router),
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
patch(
"litellm.proxy.management_endpoints.model_access_group_management_endpoints.clear_cache",
new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)),
),
):
response = await update_access_group(
access_group="my-group",
data=request_data,
user_api_key_dict=mock_user,
)
assert response.models_updated == 1
mock_prisma.db.litellm_proxymodeltable.find_unique.assert_called_once_with(
where={"model_id": "deploy-A"}
)
assert response.models_updated == 2
assert response.model_names == ["gpt-4o"]
assert response.model_ids == ["deploy-A"]
@pytest.mark.asyncio
@ -257,6 +372,7 @@ async def test_create_access_group_requires_model_names_or_model_ids():
Test that creating an access group without model_names or model_ids fails.
"""
from fastapi import HTTPException
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.management_endpoints.model_access_group_management_endpoints import (
create_model_group,
@ -288,6 +404,7 @@ async def test_create_access_group_invalid_model_id_returns_400():
Test that passing a non-existent model_id returns 400 error.
"""
from fastapi import HTTPException
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.management_endpoints.model_access_group_management_endpoints import (
create_model_group,
@ -367,7 +484,6 @@ async def test_create_access_group_surfaces_dropped_models():
assert "deploy-A" in str(exc_info.value.detail)
@pytest.mark.asyncio
async def test_create_access_group_trusts_reload_snapshot_over_post_lock_fresh_read():
"""A concurrent reconcile sampled after the lock is released must not make this
@ -396,9 +512,7 @@ async def test_create_access_group_trusts_reload_snapshot_over_post_lock_fresh_r
patch(
"litellm.proxy.management_endpoints.model_access_group_management_endpoints.clear_cache",
new=AsyncMock(
return_value=ReconcileOutcome(
still_desired=frozenset({"deploy-A"}), live_after=frozenset({"deploy-A"})
)
return_value=ReconcileOutcome(still_desired=frozenset({"deploy-A"}), live_after=frozenset({"deploy-A"}))
),
),
):
@ -484,6 +598,8 @@ async def test_delete_access_group_ignores_models_that_were_already_dead():
async def test_create_access_group_read_through_recovers_model_created_on_sibling_replica():
"""Regression: an access group referencing a model that another replica just wrote
to the DB must be created instead of 400ing until the periodic config reload."""
from types import SimpleNamespace
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.management_endpoints.model_access_group_management_endpoints import (
create_model_group,
@ -492,8 +608,6 @@ async def test_create_access_group_read_through_recovers_model_created_on_siblin
NewModelGroupRequest,
)
from types import SimpleNamespace
model_name = "e2e-ag-sibling-replica-model"
db_row = SimpleNamespace(
model_id=f"{model_name}-id",
@ -575,6 +689,7 @@ async def test_create_access_group_model_missing_everywhere_still_400s():
assert exc_info.value.status_code == 400
assert model_name in str(exc_info.value.detail)
@dataclass
class _FakeBudgetRow:
budget_id: str
@ -701,7 +816,6 @@ class _FakeModelTable:
async def update(self, where, data, include=None):
self.journal.append(f"model_table.update:{where['model_id']}")
self.updates.append((dict(where), dict(data)))
return None
class _FakePrismaClient:
@ -1128,9 +1242,7 @@ async def test_deleting_the_access_group_strips_deployments_before_dropping_the_
cache = _FakeAuthCache()
with _proxy_with_stubbed_reload(prisma):
response = await delete_access_group(
access_group="prod-models", user_api_key_dict=_admin(), auth_cache=cache
)
response = await delete_access_group(access_group="prod-models", user_api_key_dict=_admin(), auth_cache=cache)
assert response.models_updated == 1
assert prisma.access_group_budget_table.rows == {}
@ -1282,9 +1394,7 @@ async def test_deleting_an_access_group_that_never_had_a_budget_still_evicts():
cache = _FakeAuthCache(journal)
with _proxy_with_stubbed_reload(prisma):
response = await delete_access_group(
access_group="prod-models", user_api_key_dict=_admin(), auth_cache=cache
)
response = await delete_access_group(access_group="prod-models", user_api_key_dict=_admin(), auth_cache=cache)
assert response.models_updated == 1
assert prisma.access_group_budget_table.rows == {}