From cd2fd91f130604eb86f527fd35c5a8133b45fb84 Mon Sep 17 00:00:00 2001 From: jesus Date: Fri, 16 Oct 2026 18:52:45 +0000 Subject: [PATCH] refactor(model hub): build custom tags and metadata in the proxy instead of the router Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../model_group_custom_fields.py | 27 +++++-- litellm/proxy/proxy_server.py | 15 ++-- litellm/router.py | 8 +- .../model_management_endpoints.py | 3 + litellm/types/router.py | 2 - .../public_v1/test_model_hub.py | 7 +- tests/unit/proxy/test_proxy_utils.py | 79 ++++++++++++++++++- tests/unit/test_router/test_router.py | 73 ----------------- 8 files changed, 118 insertions(+), 96 deletions(-) rename litellm/{router_utils => proxy/management_helpers}/model_group_custom_fields.py (61%) diff --git a/litellm/router_utils/model_group_custom_fields.py b/litellm/proxy/management_helpers/model_group_custom_fields.py similarity index 61% rename from litellm/router_utils/model_group_custom_fields.py rename to litellm/proxy/management_helpers/model_group_custom_fields.py index 4904c48e135..2be1c9076c1 100644 --- a/litellm/router_utils/model_group_custom_fields.py +++ b/litellm/proxy/management_helpers/model_group_custom_fields.py @@ -1,4 +1,4 @@ -"""Merge the admin-defined `custom_tags` and `custom_metadata` of a model group's deployments.""" +"""Merge the admin-defined `custom_tags` and `custom_metadata` from a model group's deployment `model_info`.""" import math from collections.abc import Iterator, Mapping, Sequence @@ -31,25 +31,40 @@ def _scalar_metadata(metadata: object) -> Iterator[tuple[str, CustomMetadataValu ) -def _custom_tag_values(model_infos: Sequence[Mapping[str, object]]) -> Iterator[str]: +def _custom_tag_values(model_infos: Sequence[Mapping[object, object]]) -> Iterator[str]: for model_info in model_infos: yield from _string_tags(model_info.get("custom_tags")) def _custom_metadata_entries( - model_infos: Sequence[Mapping[str, object]], + model_infos: Sequence[Mapping[object, object]], ) -> Iterator[tuple[str, CustomMetadataValue]]: for model_info in model_infos: yield from _scalar_metadata(model_info.get("custom_metadata")) -def merged_custom_tags(model_infos: Sequence[Mapping[str, object]]) -> tuple[str, ...] | None: +def _merged_custom_tags(model_infos: Sequence[Mapping[object, object]]) -> tuple[str, ...] | None: return tuple(dict.fromkeys(_custom_tag_values(model_infos))) or None -def merged_custom_metadata( - model_infos: Sequence[Mapping[str, object]], +def _merged_custom_metadata( + model_infos: Sequence[Mapping[object, object]], ) -> Mapping[str, CustomMetadataValue] | None: entries: Final = tuple(_custom_metadata_entries(model_infos)) first_seen: Final = MappingProxyType(dict(reversed(entries))) return MappingProxyType({key: first_seen[key] for key in dict.fromkeys(key for key, _ in entries)}) or None + + +def _model_info(deployment: Mapping[str, object]) -> Mapping[object, object]: + model_info: Final = deployment.get("model_info") + return _OBJECT_MAPPING_ADAPTER.validate_python(model_info) if isinstance(model_info, dict) else MappingProxyType({}) + + +def model_group_custom_fields(deployments: Sequence[Mapping[str, object]]) -> Mapping[str, object]: + model_infos: Final = tuple(_model_info(deployment) for deployment in deployments) + return MappingProxyType( + { + "custom_tags": _merged_custom_tags(model_infos), + "custom_metadata": _merged_custom_metadata(model_infos), + } + ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 3de3dccc532..44095646f11 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -147,6 +147,7 @@ from litellm.proxy.common_utils.realtime_utils import ( # noqa: F401, RUF100 # realtime_request_body, ) from litellm.proxy.management_helpers.auto_router_availability import AutoRouterCatalogEntry, build_auto_router_catalog +from litellm.proxy.management_helpers.model_group_custom_fields import model_group_custom_fields from litellm.router_utils.access_windows import access_windows_config_error from litellm.router_utils.add_retry_fallback_headers import ( get_fallback_errors_from_headers, @@ -17050,10 +17051,7 @@ def get_model_group_info( ) -> list[ModelGroupInfoProxy]: model_groups: Final[list[ModelGroupInfoProxy]] = [] - unique_models: Final = [] - for model in all_models_str: - if model not in unique_models: - unique_models.append(model) + unique_models: Final = tuple(dict.fromkeys(all_models_str)) for model in unique_models: if model_group is not None and model_group != model: @@ -17062,7 +17060,14 @@ def get_model_group_info( _model_group_info = llm_router.get_model_group_info(model_group=model) if _model_group_info is not None: - model_groups.append(ModelGroupInfoProxy.model_validate(_model_group_info.model_dump())) + model_groups.append( + ModelGroupInfoProxy.model_validate( + { + **_model_group_info.model_dump(), + **model_group_custom_fields(llm_router.get_model_list_of_routed_group(model)), + } + ) + ) else: model_group_info = ModelGroupInfoProxy( model_group=model, diff --git a/litellm/router.py b/litellm/router.py index eb599bb9397..4f2358f9008 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -263,7 +263,6 @@ from litellm.router_utils.handle_error import ( send_llm_exception_alert, ) from litellm.router_utils.health_state_cache import DeploymentHealthCache -from litellm.router_utils.model_group_custom_fields import merged_custom_metadata, merged_custom_tags from litellm.router_utils.pre_call_checks.continuation_prefill_check import ( MID_STREAM_CONTINUATION_MARKER, ContinuationOutput, @@ -11844,13 +11843,10 @@ class Router: reasoning_efforts_initialized = False reasoning_efforts_unknown = False model_list: Final = self.get_model_list_of_routed_group(model_group) - configured_model_infos: Final = tuple(_configured_model_info(m) for m in model_list) model_group_description: Final[str | None] = next( - (d for d in (info.get("description") for info in configured_model_infos) if isinstance(d, str)), + (d for d in (_configured_model_info(m).get("description") for m in model_list) if isinstance(d, str)), None, ) - model_group_custom_tags: Final = merged_custom_tags(configured_model_infos) - model_group_custom_metadata: Final = merged_custom_metadata(configured_model_infos) for model in model_list: is_match = False if ( @@ -11958,8 +11954,6 @@ class Router: "providers": [llm_provider], "description": model_group_description, **model_info, - "custom_tags": model_group_custom_tags, - "custom_metadata": model_group_custom_metadata, "supports_fast_mode": True, "supported_reasoning_efforts": None, } diff --git a/litellm/types/proxy/management_endpoints/model_management_endpoints.py b/litellm/types/proxy/management_endpoints/model_management_endpoints.py index 3cc85a9427e..6a4160c3c50 100644 --- a/litellm/types/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/types/proxy/management_endpoints/model_management_endpoints.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from datetime import datetime from typing import Any @@ -13,6 +14,8 @@ class ModelGroupInfoProxy(ModelGroupInfo): health_status: str | None = Field(default=None) health_response_time: float | None = Field(default=None) health_checked_at: str | None = Field(default=None) + custom_tags: tuple[str, ...] | None = None + custom_metadata: Mapping[str, str | int | float | bool] | None = None class UpdateUsefulLinksRequest(LiteLLMBaseModel): diff --git a/litellm/types/router.py b/litellm/types/router.py index ebe3d20ac53..0ca39e2400d 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -945,8 +945,6 @@ class ModelGroupInfo(LiteLLMBaseModel): supported_openai_params: list[str] | None = Field(default=[]) configurable_clientside_auth_params: CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS = None description: str | None = None - custom_tags: tuple[str, ...] | None = None - custom_metadata: Mapping[str, str | int | float | bool] | None = None def __init__(self, **data) -> None: for field_name, field_type in _resolved_annotations(self.__class__).items(): diff --git a/tests/unit/proxy/public_endpoints/public_v1/test_model_hub.py b/tests/unit/proxy/public_endpoints/public_v1/test_model_hub.py index 4267cdbbcc9..8c164ec860b 100644 --- a/tests/unit/proxy/public_endpoints/public_v1/test_model_hub.py +++ b/tests/unit/proxy/public_endpoints/public_v1/test_model_hub.py @@ -10,7 +10,7 @@ from fastapi.testclient import TestClient import litellm from litellm.proxy._types import LiteLLMRoutes from litellm.proxy.proxy_server import app -from litellm.types.router import ModelGroupInfo +from litellm.types.router import DeploymentTypedDict, ModelGroupInfo from tests._master_key import MASTER_KEY client = TestClient(app) @@ -21,13 +21,16 @@ LEGACY_MODEL_HUB_PATH = "/public/model_hub" @dataclass(frozen=True, slots=True) class _FakeRouter: - """Stands in for the running Router: `_get_model_group_info` only ever asks it this.""" + """Stands in for the running Router: `_get_model_group_info` only ever asks it these.""" infos: Mapping[str, ModelGroupInfo] def get_model_group_info(self, model_group: str) -> ModelGroupInfo | None: return self.infos.get(model_group) + def get_model_list_of_routed_group(self, model_group: str) -> list[DeploymentTypedDict]: + return [] + def _info( name: str, diff --git a/tests/unit/proxy/test_proxy_utils.py b/tests/unit/proxy/test_proxy_utils.py index 20a13bdc566..86240417f0b 100644 --- a/tests/unit/proxy/test_proxy_utils.py +++ b/tests/unit/proxy/test_proxy_utils.py @@ -1,7 +1,7 @@ import asyncio import json from datetime import datetime -from typing import Any, Dict, Final, List, Optional, Union +from typing import Any, Dict, Final, List, Optional, TypeAlias, Union from unittest.mock import Mock import pytest @@ -15,6 +15,7 @@ from litellm.types.guardrails import GuardrailEventHooks from unittest.mock import AsyncMock, MagicMock, patch import litellm +from litellm.router import Router from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.auth_utils import ( check_complete_credentials, @@ -1602,6 +1603,82 @@ def test_get_model_group_info(): assert len(model_list) == 1 +def _custom_metadata_router(*model_infos: dict[object, object]) -> Router: + return Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4", "api_key": f"fake-{index}"}, + "model_info": model_info, + } + for index, model_info in enumerate(model_infos) + ] + ) + + +_CustomFieldsCase: TypeAlias = tuple[ + tuple[dict[object, object], ...], tuple[str, ...] | None, dict[str, str | int | float | bool] | None +] +_CUSTOM_FIELDS_CASES: Final[tuple[_CustomFieldsCase, ...]] = ( + ((), None, None), + (({},), None, None), + ( + ({"custom_tags": ["deprecated", "ptu"], "custom_metadata": {"retirement_date": "2027-03-01"}},), + ("deprecated", "ptu"), + {"retirement_date": "2027-03-01"}, + ), + ( + ({}, {"custom_tags": ["ptu"], "custom_metadata": {"provisioned_ptu": 100, "gsu": False}}), + ("ptu",), + {"provisioned_ptu": 100, "gsu": False}, + ), + ( + ( + {"custom_tags": ["ptu", "eu"], "custom_metadata": {"provisioned_ptu": 100, "region": "eu"}}, + {"custom_tags": ["eu", "deprecated"], "custom_metadata": {"provisioned_ptu": 50, "max_tpm": 1.5}}, + ), + ("ptu", "eu", "deprecated"), + {"provisioned_ptu": 100, "region": "eu", "max_tpm": 1.5}, + ), + ( + ({"custom_tags": ["ok", 7, None, {"x": 1}], "custom_metadata": {"owner": "ml", "nested": {"a": 1}}},), + ("ok",), + {"owner": "ml"}, + ), + (({"custom_tags": "ptu", "custom_metadata": ["not", "a", "map"]},), None, None), + ( + ( + { + "custom_metadata": { + 2027: "retirement", + "nan": float("nan"), + "inf": float("inf"), + "neg_inf": float("-inf"), + "max_tpm": 1.5, + } + }, + ), + None, + {"max_tpm": 1.5}, + ), +) + + +@pytest.mark.parametrize(("model_infos", "expected_tags", "expected_metadata"), _CUSTOM_FIELDS_CASES) +def test_get_model_group_info_merges_custom_tags_and_metadata_across_deployments( + model_infos: tuple[dict[object, object], ...], + expected_tags: tuple[str, ...] | None, + expected_metadata: dict[str, str | int | float | bool] | None, +) -> None: + from litellm.proxy.proxy_server import get_model_group_info + + router: Final = _custom_metadata_router({"description": "first"}, *model_infos) + + (group,) = get_model_group_info(llm_router=router, all_models_str=["gpt-4"], model_group="gpt-4") + + assert (group.custom_tags, group.custom_metadata) == (expected_tags, expected_metadata) + + @pytest.fixture def mock_team_data(): return [ diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 1d501e2403e..6bc52ab3dfc 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -2372,79 +2372,6 @@ def test_model_group_info_description_from_model_info(): assert result.description == "State-of-the-art language model." -def _custom_metadata_router(*model_infos: dict[str, object]) -> Router: - return litellm.Router( - model_list=[ - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4", "api_key": f"fake-{index}"}, - "model_info": model_info, - } - for index, model_info in enumerate(model_infos) - ] - ) - - -@pytest.mark.parametrize( - ("model_infos", "expected_tags", "expected_metadata"), - [ - ((), None, None), - (({},), None, None), - ( - ({"custom_tags": ["deprecated", "ptu"], "custom_metadata": {"retirement_date": "2027-03-01"}},), - ("deprecated", "ptu"), - {"retirement_date": "2027-03-01"}, - ), - ( - ({}, {"custom_tags": ["ptu"], "custom_metadata": {"provisioned_ptu": 100, "gsu": False}}), - ("ptu",), - {"provisioned_ptu": 100, "gsu": False}, - ), - ( - ( - {"custom_tags": ["ptu", "eu"], "custom_metadata": {"provisioned_ptu": 100, "region": "eu"}}, - {"custom_tags": ["eu", "deprecated"], "custom_metadata": {"provisioned_ptu": 50, "max_tpm": 1.5}}, - ), - ("ptu", "eu", "deprecated"), - {"provisioned_ptu": 100, "region": "eu", "max_tpm": 1.5}, - ), - ( - ({"custom_tags": ["ok", 7, None, {"x": 1}], "custom_metadata": {"owner": "ml", "nested": {"a": 1}}},), - ("ok",), - {"owner": "ml"}, - ), - (({"custom_tags": "ptu", "custom_metadata": ["not", "a", "map"]},), None, None), - ( - ( - { - "custom_metadata": { - 2027: "retirement", - "nan": float("nan"), - "inf": float("inf"), - "neg_inf": float("-inf"), - "max_tpm": 1.5, - } - }, - ), - None, - {"max_tpm": 1.5}, - ), - ], -) -def test_model_group_info_merges_custom_tags_and_metadata_across_deployments( - model_infos: tuple[dict[str, object], ...], - expected_tags: tuple[str, ...] | None, - expected_metadata: dict[str, object] | None, -) -> None: - router: Final = _custom_metadata_router({"description": "first"}, *model_infos) - - result: Final = router.get_model_group_info("gpt-4") - - assert result is not None - assert result.custom_tags == expected_tags - assert result.custom_metadata == expected_metadata - - def _alias_cost_router() -> Router: return Router( model_list=[