diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 268b244c86d..70ea821906b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -10918,10 +10918,20 @@ async def _apply_search_filter_to_models( search_lower = search.lower().strip() + def _model_matches_search(m: Dict[str, Any]) -> bool: + # Team BYOK models persist an internal `model_name` + # (e.g. `model_name_{team_id}_{uuid}`) and expose the user-facing + # name via `model_info.team_public_model_name`. Match both so the + # name shown in the UI is searchable. + if search_lower in (m.get("model_name") or "").lower(): + return True + team_public_model_name = (m.get("model_info") or {}).get( + "team_public_model_name" + ) or "" + return search_lower in team_public_model_name.lower() + # Filter models in router by search term - filtered_router_models = [ - m for m in all_models if search_lower in m.get("model_name", "").lower() - ] + filtered_router_models = [m for m in all_models if _model_matches_search(m)] # Separate filtered models into config vs db models, and track db model IDs filtered_config_models = [] @@ -10949,12 +10959,23 @@ async def _apply_search_filter_to_models( # Only query database if prisma_client is available if prisma_client is not None: try: - # Build where condition for database query + # Match either the persisted model_name or, for team BYOK rows, + # the user-facing team_public_model_name stored inside model_info. db_where_condition: Dict[str, Any] = { - "model_name": { - "contains": search_lower, - "mode": "insensitive", - } + "OR": [ + { + "model_name": { + "contains": search_lower, + "mode": "insensitive", + } + }, + { + "model_info": { + "path": ["team_public_model_name"], + "string_contains": search_lower, + } + }, + ] } # Exclude models already in router if we have any if db_model_ids_in_router: @@ -11102,6 +11123,15 @@ def _sort_models( model_info = model.get("model_info", {}) if sort_by == "model_name": + # Team BYOK models persist an internal `model_name` (e.g. + # `model_name_{team_id}_{uuid}`) and expose the user-facing + # name via `model_info.team_public_model_name` — same as the + # UI's getDisplayModelName. Sort by the displayed name so + # BYOK rows interleave alphabetically with non-BYOK rows + # instead of clumping at the end on their opaque IDs. + team_public_model_name = model_info.get("team_public_model_name") + if team_public_model_name: + return str(team_public_model_name).lower() return model.get("model_name", "").lower() elif sort_by == "created_at": @@ -11316,26 +11346,30 @@ async def _filter_models_by_team_id( team_object, team_id, prisma_client, llm_router ) - # Filter models based on direct_access or access_via_team_ids - # Models are already enriched with these fields before this function is called + # When filtering by a specific team we want exactly the models that team + # can use: its BYOK rows and the deployments resolved from team.models / + # access groups. `direct_access` describes the viewer's own permissions + # (the admin path sets it on every non-team model) and must NOT widen the + # team's visible set, otherwise selecting a team in the UI still shows + # every public model the admin can call. filtered_models = [] for _model in all_models: model_info = _model.get("model_info", {}) model_id = model_info.get("id", None) - # Include if direct_access is True - if model_info.get("direct_access", False): + # BYOK rows owned by this team are always accessible to it, even if + # they haven't been re-added to team.models for some reason. + if model_info.get("team_id") == team_id: filtered_models.append(_model) continue - # Include if team_id is in access_via_team_ids access_via_team_ids = model_info.get("access_via_team_ids", []) if isinstance(access_via_team_ids, list) and team_id in access_via_team_ids: filtered_models.append(_model) continue - # Also include if model_id is in team_accessible_model_ids (from config/db search) - # This catches models that might not have been enriched with access_via_team_ids yet + # Catches models resolved from team.models / access groups that + # weren't enriched with access_via_team_ids upstream. if model_id and model_id in team_accessible_model_ids: filtered_models.append(_model) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index bedc9c1f6e8..053d31e69c7 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -1169,6 +1169,146 @@ def test_add_team_models_to_all_models(): assert result == {"gpt-4-model-2": {"team1"}} +@pytest.mark.asyncio +async def test_apply_search_filter_matches_team_public_model_name(): + """ + Regression test: team BYOK models persist an internal model_name + (e.g. `model_name_{team_id}_{uuid}`) and surface the user-facing name + via `model_info.team_public_model_name`. The /v2/model/info search + filter must match that public name so BYOK rows appear in results. + """ + from litellm.proxy.proxy_server import _apply_search_filter_to_models + + byok_model = { + "model_name": "model_name_team-abc-123_4a6b8", + "litellm_params": {"model": "claude-sonnet-4-5"}, + "model_info": { + "id": "byok-id-1", + "team_id": "team-abc-123", + "team_public_model_name": "team-claude-sonnet", + "db_model": True, + }, + } + unrelated_model = { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "normal-id-1", "db_model": False}, + } + + # Search matching only team_public_model_name should still include BYOK + filtered, _ = await _apply_search_filter_to_models( + all_models=[byok_model, unrelated_model], + search="claude", + page=1, + size=50, + prisma_client=None, + proxy_config=MagicMock(), + sort_by=None, + ) + filtered_ids = {m["model_info"]["id"] for m in filtered} + assert "byok-id-1" in filtered_ids + assert "normal-id-1" not in filtered_ids + + # Search by internal model_name still matches as before + filtered, _ = await _apply_search_filter_to_models( + all_models=[byok_model, unrelated_model], + search="model_name_team-abc-123", + page=1, + size=50, + prisma_client=None, + proxy_config=MagicMock(), + sort_by=None, + ) + assert [m["model_info"]["id"] for m in filtered] == ["byok-id-1"] + + # Non-matching search returns nothing + filtered, _ = await _apply_search_filter_to_models( + all_models=[byok_model, unrelated_model], + search="gemini", + page=1, + size=50, + prisma_client=None, + proxy_config=MagicMock(), + sort_by=None, + ) + assert filtered == [] + + +@pytest.mark.asyncio +async def test_filter_models_by_team_id_excludes_viewer_direct_access(): + """ + Regression test: when the UI picks a specific team in the Current Team + selector, the model list must show only that team's BYOK rows + the + models assigned to the team. The admin viewer's `direct_access` flag + (set on every non-team model upstream) must NOT widen the team's + visible set, or selecting team-111 still shows every public model. + """ + from litellm.proxy.proxy_server import _filter_models_by_team_id + + public_model = { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4"}, + "model_info": { + "id": "public-id", + # admin viewer has direct_access on this public model + "direct_access": True, + # team-111 is NOT in access_via_team_ids -> shouldn't show for team-111 + "access_via_team_ids": ["team-222"], + }, + } + team111_byok = { + "model_name": "model_name_team-111_uuid", + "litellm_params": {"model": "claude-sonnet"}, + "model_info": { + "id": "byok-team-111", + "team_id": "team-111", + "team_public_model_name": "team-claude", + "access_via_team_ids": ["team-111"], + }, + } + team222_byok = { + "model_name": "model_name_team-222_uuid", + "litellm_params": {"model": "claude-haiku"}, + "model_info": { + "id": "byok-team-222", + "team_id": "team-222", + "team_public_model_name": "team-haiku", + "access_via_team_ids": ["team-222"], + }, + } + + prisma = MagicMock() + team_db = MagicMock() + team_db.model_dump.return_value = { + "team_id": "team-111", + "team_alias": "Team 111", + # specific models list that doesn't include the BYOK's internal name + "models": ["some-other-model"], + "access_group_ids": None, + } + prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_db) + prisma.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) + + router = MagicMock() + router.get_model_access_groups = MagicMock(return_value={}) + # team-111 only resolves "some-other-model", which has no deployments + router.get_model_list = MagicMock(return_value=[]) + + filtered = await _filter_models_by_team_id( + all_models=[public_model, team111_byok, team222_byok], + team_id="team-111", + prisma_client=prisma, + llm_router=router, + ) + visible_ids = sorted(m["model_info"]["id"] for m in filtered) + + assert "byok-team-111" in visible_ids, "team-111's own BYOK must always be visible" + assert "byok-team-222" not in visible_ids, "must not leak other teams' BYOK" + assert ( + "public-id" not in visible_ids + ), "viewer's direct_access must not widen the team's visible set" + + @pytest.mark.asyncio async def test_add_access_group_models_to_team_models(): """ @@ -6922,3 +7062,80 @@ class TestTransformRequestBannedParams: f"Expected 400 for banned param '{banned}', " f"got {response.status_code}: {response.json()}" ) + + +class TestSortModelsByDisplayName: + """Regression: team BYOK rows persist an internal `model_name` like + `model_name_{team_id}_{uuid}` and expose the user-facing name via + `model_info.team_public_model_name`. Sorting must use the displayed + name so BYOK rows interleave with non-BYOK rows alphabetically — + otherwise they clump at the end on their opaque IDs even though the + UI shows them under a normal-looking name. + """ + + def test_byok_models_sort_by_team_public_model_name(self): + from litellm.proxy.proxy_server import _sort_models + + models = [ + {"model_name": "claude-haiku-4-5", "model_info": {}}, + { + # Opaque internal name; UI displays team_public_model_name. + "model_name": "model_name_team-1_abc123", + "model_info": {"team_public_model_name": "anthropic/claude"}, + }, + {"model_name": "gpt-4o", "model_info": {}}, + ] + + sorted_models = _sort_models( + all_models=models, sort_by="model_name", sort_order="asc" + ) + displayed_order = [ + m["model_info"].get("team_public_model_name") or m["model_name"] + for m in sorted_models + ] + assert displayed_order == [ + "anthropic/claude", + "claude-haiku-4-5", + "gpt-4o", + ] + + def test_byok_models_sort_descending_by_display_name(self): + from litellm.proxy.proxy_server import _sort_models + + models = [ + {"model_name": "claude-haiku-4-5", "model_info": {}}, + { + "model_name": "model_name_team-1_zzz", + "model_info": {"team_public_model_name": "zeta/model"}, + }, + {"model_name": "gpt-4o", "model_info": {}}, + ] + + sorted_models = _sort_models( + all_models=models, sort_by="model_name", sort_order="desc" + ) + displayed_order = [ + m["model_info"].get("team_public_model_name") or m["model_name"] + for m in sorted_models + ] + assert displayed_order == [ + "zeta/model", + "gpt-4o", + "claude-haiku-4-5", + ] + + def test_empty_team_public_model_name_falls_back_to_model_name(self): + # Empty string for team_public_model_name (not None) must still + # fall back to model_name — otherwise BYOK rows with a blank + # display name would sort to the top. + from litellm.proxy.proxy_server import _sort_models + + models = [ + {"model_name": "alpha", "model_info": {"team_public_model_name": ""}}, + {"model_name": "beta", "model_info": {}}, + ] + + sorted_models = _sort_models( + all_models=models, sort_by="model_name", sort_order="asc" + ) + assert [m["model_name"] for m in sorted_models] == ["alpha", "beta"]