From e1a3fca8fa1320c522d97b008da6a5b40533c291 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Tue, 16 Jun 2026 17:01:11 -0700 Subject: [PATCH] refactor(proxy): simplify team model listing translation --- litellm/proxy/proxy_server.py | 47 +++++++--------- .../test_team_model_name_translation.py | 54 +++++++++++++------ 2 files changed, 55 insertions(+), 46 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 73352c357d5..acd219d4236 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -15,7 +15,7 @@ import threading import time import traceback import warnings -from collections.abc import Mapping, Sequence +from collections.abc import Mapping from datetime import datetime, timedelta, timezone from typing import ( TYPE_CHECKING, @@ -8346,10 +8346,7 @@ async def model_list( # Surface the public name for team-scoped rows by default. Operators # that need legacy internal routing keys can explicitly disable this. - if _should_use_team_public_model_name(): - all_models = _translate_model_names_for_listing( - all_models, _get_team_model_deployments_for_listing(llm_router) - ) + all_models = _translate_team_model_names_for_listing(all_models, llm_router) # Build response data with all proxy models model_data = [] @@ -8390,10 +8387,7 @@ async def model_list( # Surface the public name for team-scoped rows by default. Operators that # need legacy internal routing keys can explicitly disable this. - if _should_use_team_public_model_name(): - all_models = _translate_model_names_for_listing( - all_models, _get_team_model_deployments_for_listing(llm_router) - ) + all_models = _translate_team_model_names_for_listing(all_models, llm_router) # Build response data model_data = [] @@ -12604,23 +12598,9 @@ def _translate_model_name_for_response(model: dict) -> dict: return {**model, "model_name": team_public} -def _should_use_team_public_model_name() -> bool: - settings = cast(dict[str, object], general_settings) # any-ok: legacy settings - use_public_name = settings.get("use_team_public_model_name", True) - return use_public_name is not False - - -def _get_team_model_deployments_for_listing( - llm_router: Router | None, -) -> Sequence[Mapping[str, object]]: - if llm_router is None: - return () - return cast(Sequence[Mapping[str, object]], llm_router.get_model_list() or ()) - - -def _translate_model_names_for_listing( +def _translate_team_model_names_for_listing( model_names: list[str], - model_deployments: Sequence[Mapping[str, object]], + llm_router: Router | None, ) -> list[str]: """Swap internal team routing keys for their public names in list-style responses (e.g. `/v1/models`, `/models`). @@ -12632,17 +12612,26 @@ def _translate_model_names_for_listing( unchanged (see issue #28382). Sibling deployments collapse to one public name, so the result is de-duplicated while preserving order. """ - if not model_deployments: + settings = cast(dict[str, object], general_settings) # any-ok: legacy settings + if settings.get("use_team_public_model_name", True) is False or llm_router is None: return model_names + + router_model_list = llm_router.get_model_list() + if not isinstance(router_model_list, list): + return model_names + internal_to_public: dict[str, str] = {} - for model in model_deployments: - model_info_raw = model.get("model_info") + for model in router_model_list: + if not isinstance(model, dict): + continue + model_dict = cast(dict[str, object], model) # any-ok: checked + model_info_raw: object = model_dict.get("model_info") if not isinstance(model_info_raw, Mapping): continue model_info = cast(Mapping[str, object], model_info_raw) # any-ok: checked team_id = model_info.get("team_id") team_public = model_info.get("team_public_model_name") - name = model.get("model_name") + name = model_dict.get("model_name") if ( isinstance(team_id, str) and isinstance(team_public, str) diff --git a/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py b/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py index 59d9f46d1b6..f1056ee201d 100644 --- a/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py +++ b/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py @@ -716,13 +716,13 @@ async def test_v1_models_translates_team_model_with_metadata(monkeypatch): ] -def test_translate_model_names_for_listing_swaps_and_dedupes(): +def test_translate_team_model_names_for_listing_swaps_and_dedupes(monkeypatch): """Internal team routing keys -> public name; sibling deployments sharing a public name collapse to one entry (order preserved); globals untouched.""" - from litellm.proxy.proxy_server import _translate_model_names_for_listing + from litellm.proxy.proxy_server import _translate_team_model_names_for_listing router = MagicMock() - router.model_list = [ + router.get_model_list.return_value = [ { "model_name": "model_name_teamX_uuidA", "model_info": { @@ -739,34 +739,54 @@ def test_translate_model_names_for_listing_swaps_and_dedupes(): }, {"model_name": "gpt-4o", "model_info": {"db_model": False}}, ] + monkeypatch.setattr(ps, "general_settings", {}) - out = _translate_model_names_for_listing( + out = _translate_team_model_names_for_listing( ["model_name_teamX_uuidA", "model_name_teamX_uuidB", "gpt-4o"], - router.model_list, + router, ) assert out == ["tushar-gpt-4.1", "gpt-4o"] -def test_translate_model_names_for_listing_leaves_unmapped_names(): +def test_translate_team_model_names_for_listing_leaves_unmapped_names(monkeypatch): """Names with no team mapping (globals, access-group keys) pass through.""" - from litellm.proxy.proxy_server import _translate_model_names_for_listing + from litellm.proxy.proxy_server import _translate_team_model_names_for_listing router = MagicMock() - router.model_list = [{"model_name": "gpt-4o", "model_info": {"db_model": False}}] - assert _translate_model_names_for_listing( - ["gpt-4o", "beta-group"], router.model_list + router.get_model_list.return_value = [ + {"model_name": "gpt-4o", "model_info": {"db_model": False}} + ] + monkeypatch.setattr(ps, "general_settings", {}) + + assert _translate_team_model_names_for_listing( + ["gpt-4o", "beta-group"], router ) == ["gpt-4o", "beta-group"] -def test_translate_model_names_for_listing_none_router(): +def test_translate_team_model_names_for_listing_none_router(monkeypatch): """No router -> return the input list unchanged.""" - from litellm.proxy.proxy_server import _translate_model_names_for_listing + from litellm.proxy.proxy_server import _translate_team_model_names_for_listing - assert _translate_model_names_for_listing(["a", "b"], ()) == ["a", "b"] + monkeypatch.setattr(ps, "general_settings", {}) + assert _translate_team_model_names_for_listing(["a", "b"], None) == ["a", "b"] -def test_get_team_model_deployments_for_listing_none_router(): - """No router -> no deployment metadata to translate against.""" - from litellm.proxy.proxy_server import _get_team_model_deployments_for_listing +def test_translate_team_model_names_for_listing_respects_legacy_flag(monkeypatch): + """Operators can keep returning the legacy internal routing key.""" + from litellm.proxy.proxy_server import _translate_team_model_names_for_listing - assert _get_team_model_deployments_for_listing(None) == () + router = MagicMock() + router.get_model_list.return_value = [ + { + "model_name": "model_name_teamX_uuidA", + "model_info": { + "team_id": "teamX", + "team_public_model_name": "tushar-gpt-4.1", + }, + } + ] + monkeypatch.setattr(ps, "general_settings", {"use_team_public_model_name": False}) + + assert _translate_team_model_names_for_listing( + ["model_name_teamX_uuidA"], router + ) == ["model_name_teamX_uuidA"]