diff --git a/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py b/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py index af469376289..7b173d4d789 100644 --- a/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py @@ -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 diff --git a/tests/unit/proxy/management_endpoints/test_access_group_management.py b/tests/unit/proxy/management_endpoints/test_access_group_management.py index 59c2921e0d0..6f25c371578 100644 --- a/tests/unit/proxy/management_endpoints/test_access_group_management.py +++ b/tests/unit/proxy/management_endpoints/test_access_group_management.py @@ -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 == {}