From a66b96579e6e3948a18dbc3c9174fae8bf879d45 Mon Sep 17 00:00:00 2001 From: Filippo Menghi <113345637+Cyberfilo@users.noreply.github.com> Date: Fri, 15 May 2026 13:53:15 +0200 Subject: [PATCH] feat(proxy): add blocked flag to models for pause/resume from the UI (#27927) Squash-merged by litellm-agent from Cyberfilo's PR. --- .../migration.sql | 4 + .../litellm_proxy_extras/schema.prisma | 5 +- litellm/proxy/_types.py | 1 + .../model_management_endpoints.py | 17 ++ litellm/proxy/proxy_server.py | 13 ++ litellm/proxy/route_llm_request.py | 6 +- litellm/proxy/schema.prisma | 5 +- litellm/router.py | 102 +++++++++-- litellm/types/router.py | 4 + schema.prisma | 5 +- .../test_model_management_endpoints.py | 172 ++++++++++++++++++ tests/test_litellm/test_router.py | 145 +++++++++++++++ 12 files changed, 458 insertions(+), 21 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260514120000_add_blocked_to_proxy_model_table/migration.sql diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260514120000_add_blocked_to_proxy_model_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260514120000_add_blocked_to_proxy_model_table/migration.sql new file mode 100644 index 00000000000..3253b63a884 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260514120000_add_blocked_to_proxy_model_table/migration.sql @@ -0,0 +1,4 @@ +-- AlterTable +-- Adds the admin-toggleable pause flag used by the router's blocked filter and the +-- credential lookup helpers; defaults to false so existing rows behave unchanged. +ALTER TABLE "LiteLLM_ProxyModelTable" ADD COLUMN IF NOT EXISTS "blocked" BOOLEAN NOT NULL DEFAULT false; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index b53507abe6a..78143fe0411 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -48,9 +48,10 @@ model LiteLLM_CredentialsTable { // Models on proxy model LiteLLM_ProxyModelTable { model_id String @id @default(uuid()) - model_name String + model_name String litellm_params Json - model_info Json? + model_info Json? + blocked Boolean @default(false) created_at DateTime @default(now()) @map("created_at") created_by String updated_at DateTime @default(now()) @updatedAt @map("updated_at") diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index d4fa497698a..c989f5dff13 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -4549,6 +4549,7 @@ class PrismaCompatibleUpdateDBModel(TypedDict, total=False): model_name: str litellm_params: str model_info: str + blocked: bool updated_at: str updated_by: str diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index af84bc123ff..472306eb818 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -150,6 +150,9 @@ def update_db_model( model_info[key] = value.isoformat() prisma_compatible_model_dict["model_info"] = json.dumps(model_info) + if updated_patch.blocked is not None: + prisma_compatible_model_dict["blocked"] = updated_patch.blocked + return prisma_compatible_model_dict @@ -230,6 +233,20 @@ async def patch_model( premium_user=premium_user, ) + # Pause/resume (`blocked`) is a proxy-admin-only privilege. Team admins + # passed the auth check above for team-scoped models, but they must not + # be able to unblock (or block) a model their proxy admin has paused. + if ( + patch_data.blocked is not None + and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN + ): + raise ProxyException( + message="Only proxy admins can change a model's blocked flag.", + type=ProxyErrorTypes.auth_error.value, + code=status.HTTP_403_FORBIDDEN, + param="blocked", + ) + # Handle team model updates with proper alias management update_data = await _update_team_model_in_db( db_model=db_model, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 46ebff515b7..eeda5e83145 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -4728,6 +4728,7 @@ class ProxyConfig: if _id is not None: model.model_info["id"] = _id model.model_info["db_model"] = True + model.model_info["blocked"] = bool(getattr(model, "blocked", False)) if premium_user is True: # seeing "created_at", "updated_at", "created_by", "updated_by" is a LiteLLM Enterprise Feature @@ -8099,6 +8100,12 @@ async def model_list( only_model_access_groups=only_model_access_groups or False, ) + # Hide paused models from the public listing (admins manage them via /model/info) + if llm_router is not None: + blocked_names = llm_router.get_fully_blocked_model_names() + if blocked_names: + all_models = [m for m in all_models if m not in blocked_names] + # Build response data with all proxy models model_data = [] for model in all_models: @@ -8132,6 +8139,12 @@ async def model_list( user_api_key_cache=user_api_key_cache, ) + # Hide paused models from the public listing (admins manage them via /model/info) + if llm_router is not None: + blocked_names = llm_router.get_fully_blocked_model_names() + if blocked_names: + all_models = [m for m in all_models if m not in blocked_names] + # Build response data model_data = [] for model in all_models: diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index bfe6b8484fa..0dee21ee281 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -430,7 +430,11 @@ async def route_request( # noqa: PLR0915 - Complex routing function, refactorin deployment = llm_router.get_deployment_by_model_group_name( model_group_name=model ) - if deployment and deployment.litellm_params: + if ( + deployment + and deployment.litellm_params + and not llm_router._is_deployment_blocked(deployment) + ): deployment_creds = deployment.litellm_params.model_dump( exclude_none=True ) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index b53507abe6a..78143fe0411 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -48,9 +48,10 @@ model LiteLLM_CredentialsTable { // Models on proxy model LiteLLM_ProxyModelTable { model_id String @id @default(uuid()) - model_name String + model_name String litellm_params Json - model_info Json? + model_info Json? + blocked Boolean @default(false) created_at DateTime @default(now()) @map("created_at") created_by String updated_at DateTime @default(now()) @updatedAt @map("updated_at") diff --git a/litellm/router.py b/litellm/router.py index 3340e06627e..484fd816e3e 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -30,6 +30,7 @@ from typing import ( List, Literal, Optional, + Set, Tuple, Union, cast, @@ -6834,12 +6835,11 @@ class Router: unhealthy_deployments = _get_cooldown_deployments( litellm_router_instance=self, parent_otel_span=parent_otel_span ) - healthy_deployments: list = [] - for deployment in _all_deployments: - if deployment["model_info"]["id"] in unhealthy_deployments: - continue - else: - healthy_deployments.append(deployment) + unhealthy_set = set(unhealthy_deployments) + healthy_deployments: list = [ + d for d in _all_deployments if d["model_info"]["id"] not in unhealthy_set + ] + healthy_deployments = self._filter_blocked_deployments(healthy_deployments) return healthy_deployments, _all_deployments @@ -6867,10 +6867,12 @@ class Router: ) # Convert to set for O(1) lookup instead of O(n) unhealthy_deployments_set = set(unhealthy_deployments) - healthy_deployments: list = [] - for deployment in _all_deployments: - if deployment["model_info"]["id"] not in unhealthy_deployments_set: - healthy_deployments.append(deployment) + healthy_deployments: list = [ + d + for d in _all_deployments + if d["model_info"]["id"] not in unhealthy_deployments_set + ] + healthy_deployments = self._filter_blocked_deployments(healthy_deployments) return healthy_deployments, _all_deployments def routing_strategy_pre_call_checks(self, deployment: dict): @@ -8019,10 +8021,14 @@ class Router: def get_deployment_credentials(self, model_id: str) -> Optional[dict]: """ - Returns -> dict of credentials for a given model id + Returns -> dict of credentials for a given model id. + + Returns None if the deployment is paused via `LiteLLM_ProxyModelTable.blocked`, + so file/batch/passthrough callers that resolve credentials directly cannot keep + using a paused deployment. """ deployment = self.get_deployment(model_id=model_id) - if deployment is None: + if deployment is None or self._is_deployment_blocked(deployment): return None return CredentialLiteLLMParams( **deployment.litellm_params.model_dump(exclude_none=True) @@ -8067,7 +8073,9 @@ class Router: Returns: Dictionary containing api_key, api_base, custom_llm_provider, etc. - Returns None if model not found. + Returns None if model not found, or if the resolved deployment is + paused via `LiteLLM_ProxyModelTable.blocked` (so passthrough callers + cannot bypass an admin pause by resolving credentials directly). Example: credentials = router.get_deployment_credentials_with_provider("gpt-4o-litellm") @@ -8093,7 +8101,7 @@ class Router: elif isinstance(deployment_dict, Deployment): deployment = deployment_dict - if deployment is None: + if deployment is None or self._is_deployment_blocked(deployment): return None # Get basic credentials @@ -9120,6 +9128,29 @@ class Router: return model_names + def get_fully_blocked_model_names(self) -> Set[str]: + """ + Returns the set of model_names where every backing deployment has `blocked=True`. + + Used by `/v1/models` to hide paused models from client listings while still + surfacing them on admin endpoints (e.g. `/model/info`). A model with at least + one non-blocked deployment is still serviceable and remains visible. + """ + deployments = self.get_model_list() or [] + blocked_by_name: Dict[str, bool] = {} + for deployment in deployments: + name = deployment.get("model_name") or "" + if not name: + continue + is_blocked = (deployment.get("model_info") or {}).get("blocked") is True + if name in blocked_by_name: + blocked_by_name[name] = blocked_by_name[name] and is_blocked + else: + blocked_by_name[name] = is_blocked + return { + name for name, fully_blocked in blocked_by_name.items() if fully_blocked + } + def _get_team_specific_model( self, deployment: DeploymentTypedDict, team_id: Optional[str] = None ) -> Optional[str]: @@ -10006,6 +10037,8 @@ class Router: ) if isinstance(healthy_deployments, dict): + if (healthy_deployments.get("model_info") or {}).get("blocked") is True: + raise RouterRateLimitErrorBasic(model=model) return healthy_deployments # Health-check-based filtering (before cooldown) @@ -10039,6 +10072,8 @@ class Router: ) healthy_deployments = _pre_cooldown_deployments + healthy_deployments = self._filter_blocked_deployments(healthy_deployments) + healthy_deployments = await self.async_callback_filter_deployments( model=model, healthy_deployments=healthy_deployments, @@ -10419,6 +10454,8 @@ class Router: ) if isinstance(healthy_deployments, dict): + if (healthy_deployments.get("model_info") or {}).get("blocked") is True: + raise RouterRateLimitErrorBasic(model=model) return healthy_deployments parent_otel_span: Optional[Span] = _get_parent_otel_span_from_kwargs( @@ -10449,6 +10486,8 @@ class Router: ) healthy_deployments = _pre_cooldown_deployments + healthy_deployments = self._filter_blocked_deployments(healthy_deployments) + # filter pre-call checks if self.enable_pre_call_checks and messages is not None: healthy_deployments = self._pre_call_checks( @@ -10557,6 +10596,8 @@ class Router: # 2. If the returned is a specific deployment (Dict), verify and return directly if isinstance(healthy_deployments, dict): + if (healthy_deployments.get("model_info") or {}).get("blocked") is True: + raise RouterRateLimitErrorBasic(model=model) litellm_params = healthy_deployments.get("litellm_params", {}) if litellm_params.get("use_in_pass_through"): return healthy_deployments @@ -10596,6 +10637,9 @@ class Router: healthy_deployments=pass_through_deployments, cooldown_deployments=cooldown_deployments, ) + pass_through_deployments = self._filter_blocked_deployments( + pass_through_deployments + ) # 5. Apply pre-call checks (if enabled) if self.enable_pre_call_checks and messages is not None: @@ -10685,6 +10729,36 @@ class Router: if deployment["model_info"]["id"] not in cooldown_set ] + def _filter_blocked_deployments( + self, healthy_deployments: List[Dict] + ) -> List[Dict]: + """ + Filters out deployments that an admin has paused via `LiteLLM_ProxyModelTable.blocked`. + + Applied alongside the cooldown filter on every routing entry point that calls + `_common_checks_available_deployment` directly — the primary sync/async path, + the sync pass-through path, and the retry / health-check helpers — so paused + deployments never serve a request. The async pass-through path inherits this + filter through its delegation to `async_get_healthy_deployments`. + """ + return [ + deployment + for deployment in healthy_deployments + if (deployment.get("model_info") or {}).get("blocked") is not True + ] + + @staticmethod + def _is_deployment_blocked(deployment: "Deployment") -> bool: + """ + Returns True when a `Deployment` Pydantic instance carries the admin-paused + flag. Used by credential-lookup helpers so passthrough file / batch endpoints + cannot bypass the pause by resolving credentials directly. + """ + model_info = getattr(deployment, "model_info", None) + if model_info is None: + return False + return getattr(model_info, "blocked", None) is True + async def _async_filter_health_check_unhealthy_deployments( self, healthy_deployments: List[Dict], diff --git a/litellm/types/router.py b/litellm/types/router.py index 926815ba317..6601f552b52 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -133,6 +133,9 @@ class ModelInfo(BaseModel): # the model_name that can be used by the team when making LLM calls team_public_model_name: Optional[str] = None + # admin-toggled pause flag; mirrors LiteLLM_ProxyModelTable.blocked + blocked: Optional[bool] = None + def __init__(self, id: Optional[Union[str, int]] = None, **params): if id is None: id = str(uuid.uuid4()) # Generate a UUID if id is None or not provided @@ -323,6 +326,7 @@ class updateDeployment(BaseModel): model_name: Optional[str] = None litellm_params: Optional[updateLiteLLMParams] = None model_info: Optional[ModelInfo] = None + blocked: Optional[bool] = None model_config = ConfigDict(protected_namespaces=()) diff --git a/schema.prisma b/schema.prisma index b53507abe6a..78143fe0411 100644 --- a/schema.prisma +++ b/schema.prisma @@ -48,9 +48,10 @@ model LiteLLM_CredentialsTable { // Models on proxy model LiteLLM_ProxyModelTable { model_id String @id @default(uuid()) - model_name String + model_name String litellm_params Json - model_info Json? + model_info Json? + blocked Boolean @default(false) created_at DateTime @default(now()) @map("created_at") created_by String updated_at DateTime @default(now()) @updatedAt @map("updated_at") diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index f0bf4578636..b65f6305b77 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -1446,3 +1446,175 @@ class TestGetTeamDeployments: result = await _get_team_deployments(team_id, prisma_client) assert len(result) == 1 assert result[0] is dep1 + + +def _build_db_model_for_blocked_test(): + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + return Deployment( + model_name="gpt-4o", + litellm_params=LiteLLM_Params(model="openai/gpt-4o"), + model_info=ModelInfo(id="dep-0"), + ) + + +class TestUpdateDBModelBlocked: + """`update_db_model` must thread `blocked` through to the Prisma payload only + when the caller explicitly set it — PATCH semantics: an absent field means + "leave the stored value untouched".""" + + def test_update_db_model_passes_blocked_true_to_db(self): + from litellm.proxy.management_endpoints.model_management_endpoints import ( + update_db_model, + ) + + result = update_db_model( + db_model=_build_db_model_for_blocked_test(), + updated_patch=updateDeployment(blocked=True), + ) + assert result["blocked"] is True + + def test_update_db_model_passes_blocked_false_to_db(self): + from litellm.proxy.management_endpoints.model_management_endpoints import ( + update_db_model, + ) + + result = update_db_model( + db_model=_build_db_model_for_blocked_test(), + updated_patch=updateDeployment(blocked=False), + ) + assert result["blocked"] is False + + def test_update_db_model_omits_blocked_when_patch_is_none(self): + from litellm.proxy.management_endpoints.model_management_endpoints import ( + update_db_model, + ) + + result = update_db_model( + db_model=_build_db_model_for_blocked_test(), + updated_patch=updateDeployment(), + ) + assert "blocked" not in result + + +class TestGetModelInfoWithIdBlocked: + """`ProxyConfig.get_model_info_with_id` must propagate the DB-level `blocked` + column into the in-memory `model_info` dict so the router filter can read it.""" + + def test_get_model_info_with_id_propagates_blocked_true(self): + from litellm.proxy.proxy_server import ProxyConfig + + model = MagicMock() + model.model_id = "dep-1" + model.model_info = {} + model.blocked = True + info = ProxyConfig().get_model_info_with_id(model=model, db_model=True) + assert info.id == "dep-1" + assert getattr(info, "blocked") is True + + def test_get_model_info_with_id_defaults_blocked_to_false_when_missing(self): + from litellm.proxy.proxy_server import ProxyConfig + + model = MagicMock(spec=["model_id", "model_info"]) + model.model_id = "dep-2" + model.model_info = {} + info = ProxyConfig().get_model_info_with_id(model=model, db_model=True) + assert getattr(info, "blocked") is False + + +class TestPatchModelBlockedAuthGate: + """Only proxy admins may flip `blocked` — team admins authorized for + team-scoped models via `can_user_make_model_call` must still be rejected + when they attempt to toggle the pause flag.""" + + @pytest.mark.asyncio + async def test_team_admin_cannot_toggle_blocked(self): + from litellm.proxy.management_endpoints.model_management_endpoints import ( + patch_model, + ) + + non_admin = UserAPIKeyAuth( + user_id="team_admin", + user_role=LitellmUserRoles.INTERNAL_USER, + ) + existing_row = MagicMock() + existing_row.litellm_params = {"model": "openai/gpt-4o-mini"} + existing_row.model_dump.return_value = { + "model_name": "gpt-4o-mini", + "litellm_params": existing_row.litellm_params, + "model_info": {"id": "m1"}, + } + existing_row.model_dump_json.return_value = "{}" + + mock_prisma = MagicMock() + mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock( + return_value=existing_row + ) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch("litellm.proxy.proxy_server.llm_router", MagicMock()), + patch("litellm.proxy.proxy_server.store_model_in_db", True), + patch("litellm.proxy.proxy_server.premium_user", True), + patch( + "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", + new=AsyncMock(return_value=None), + ), + ): + with pytest.raises(Exception) as exc_info: + await patch_model( + model_id="m1", + patch_data=updateDeployment(blocked=True), + user_api_key_dict=non_admin, + ) + err = exc_info.value + assert getattr(err, "param", "") == "blocked" + assert "proxy admin" in getattr(err, "message", "").lower() + + @pytest.mark.asyncio + async def test_proxy_admin_can_toggle_blocked(self): + from litellm.proxy.management_endpoints.model_management_endpoints import ( + patch_model, + ) + + admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + existing_row = MagicMock() + existing_row.litellm_params = {"model": "openai/gpt-4o-mini"} + existing_row.model_dump.return_value = { + "model_name": "gpt-4o-mini", + "litellm_params": existing_row.litellm_params, + "model_info": {"id": "m1"}, + } + existing_row.model_dump_json.return_value = "{}" + updated_row = MagicMock() + updated_row.model_dump_json.return_value = "{}" + + mock_prisma = MagicMock() + mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock( + return_value=existing_row + ) + mock_prisma.db.litellm_proxymodeltable.update = AsyncMock( + return_value=updated_row + ) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch("litellm.proxy.proxy_server.llm_router", MagicMock()), + patch("litellm.proxy.proxy_server.store_model_in_db", True), + patch("litellm.proxy.proxy_server.premium_user", True), + patch( + "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", + new=AsyncMock(return_value=None), + ), + patch( + "litellm.proxy.management_endpoints.model_management_endpoints.clear_cache", + new=AsyncMock(return_value=None), + ), + ): + result = await patch_model( + model_id="m1", + patch_data=updateDeployment(blocked=True), + user_api_key_dict=admin, + ) + assert result is updated_row + mock_prisma.db.litellm_proxymodeltable.update.assert_awaited_once() diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 48facace528..02d24aded02 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -3697,3 +3697,148 @@ def test_try_early_resolve_deployments_for_model_not_in_names(): default_router.default_deployment["litellm_params"]["model"] == "openai/will-be-overridden" ) + + +def _router_with_two_deployments(blocked_flags): + import litellm + + model_list = [] + for idx, blocked in enumerate(blocked_flags): + model_list.append( + { + "model_name": "gpt-4o", + "litellm_params": {"model": f"openai/gpt-4o-{idx}"}, + "model_info": {"id": f"dep-{idx}", "blocked": blocked}, + } + ) + return litellm.Router(model_list=model_list) + + +def test_get_fully_blocked_model_names_marks_name_when_all_deployments_blocked(): + router = _router_with_two_deployments([True, True]) + assert router.get_fully_blocked_model_names() == {"gpt-4o"} + + +def test_get_fully_blocked_model_names_keeps_name_when_partial_blocked(): + router = _router_with_two_deployments([True, False]) + assert router.get_fully_blocked_model_names() == set() + + +def test_get_fully_blocked_model_names_treats_missing_key_as_unblocked(): + import litellm + + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4o", + "litellm_params": {"model": "openai/gpt-4o"}, + "model_info": {"id": "dep-0"}, + } + ] + ) + assert router.get_fully_blocked_model_names() == set() + + +@pytest.mark.asyncio +async def test_async_get_healthy_deployments_skips_blocked_deployment(): + router = _router_with_two_deployments([True, False]) + healthy, all_dep = await router._async_get_healthy_deployments( + model="gpt-4o", parent_otel_span=None + ) + healthy_ids = [d["model_info"]["id"] for d in healthy] + assert "dep-0" not in healthy_ids + assert "dep-1" in healthy_ids + assert len(all_dep) == 2 + + +def test_get_healthy_deployments_sync_skips_blocked_deployment(): + router = _router_with_two_deployments([False, True]) + healthy, all_dep = router._get_healthy_deployments( + model="gpt-4o", parent_otel_span=None + ) + healthy_ids = [d["model_info"]["id"] for d in healthy] + assert "dep-0" in healthy_ids + assert "dep-1" not in healthy_ids + assert len(all_dep) == 2 + + +def test_filter_blocked_deployments_drops_blocked_keeps_unblocked(): + router = _router_with_two_deployments([True, False]) + filtered = router._filter_blocked_deployments(router.get_model_list() or []) + ids = [d["model_info"]["id"] for d in filtered] + assert ids == ["dep-1"] + + +@pytest.mark.asyncio +async def test_public_async_get_healthy_deployments_skips_blocked_on_primary_path(): + router = _router_with_two_deployments([True, False]) + deployments = await router.async_get_healthy_deployments( + model="gpt-4o", request_kwargs={} + ) + assert isinstance(deployments, list) + ids = [d["model_info"]["id"] for d in deployments] + assert "dep-0" not in ids + assert "dep-1" in ids + + +def test_public_get_available_deployment_skips_blocked_on_primary_path(): + router = _router_with_two_deployments([True, False]) + deployment = router.get_available_deployment(model="gpt-4o", request_kwargs={}) + assert deployment["model_info"]["id"] == "dep-1" + + +def test_get_available_deployment_raises_when_addressed_dict_is_blocked(): + from litellm.types.router import RouterRateLimitErrorBasic + + router = _router_with_two_deployments([True, True]) + with pytest.raises(RouterRateLimitErrorBasic): + router.get_available_deployment(model="dep-0", request_kwargs={}) + + +def _router_with_two_pass_through_deployments(blocked_flags): + import litellm + + model_list = [] + for idx, blocked in enumerate(blocked_flags): + model_list.append( + { + "model_name": "gpt-4o", + "litellm_params": { + "model": f"openai/gpt-4o-{idx}", + "api_key": "sk-fake-for-tests", + "use_in_pass_through": True, + }, + "model_info": {"id": f"pt-{idx}", "blocked": blocked}, + } + ) + return litellm.Router(model_list=model_list) + + +def test_get_available_deployment_for_pass_through_skips_blocked(): + router = _router_with_two_pass_through_deployments([True, False]) + deployment = router.get_available_deployment_for_pass_through( + model="gpt-4o", request_kwargs={} + ) + assert deployment["model_info"]["id"] == "pt-1" + + +def test_get_available_deployment_for_pass_through_raises_when_dict_blocked(): + from litellm.types.router import RouterRateLimitErrorBasic + + router = _router_with_two_pass_through_deployments([True, True]) + with pytest.raises(RouterRateLimitErrorBasic): + router.get_available_deployment_for_pass_through( + model="pt-0", request_kwargs={} + ) + + +def test_get_deployment_credentials_returns_none_for_blocked_deployment(): + router = _router_with_two_deployments([True, False]) + assert router.get_deployment_credentials(model_id="dep-0") is None + assert router.get_deployment_credentials(model_id="dep-1") is not None + + +def test_get_deployment_credentials_with_provider_returns_none_for_blocked_deployment(): + router = _router_with_two_deployments([True, False]) + assert router.get_deployment_credentials_with_provider(model_id="dep-0") is None + assert router.get_deployment_credentials_with_provider(model_id="dep-1") is not None