scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-24 05:45:11 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 80

refactor(model hub): build custom tags and metadata in the proxy instead of the router

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
jesus 2026-10-16 18:52:45 +00:00
parent 3512a825c0
commit cd2fd91f13
8 changed files with 118 additions and 96 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

27
litellm/router_utils/model_group_custom_fields.py → litellm/proxy/management_helpers/model_group_custom_fields.py
View file

@ -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),
}
)

15
litellm/proxy/proxy_server.py
View file

@ -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,

8
litellm/router.py
View file

@ -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,
}

3
litellm/types/proxy/management_endpoints/model_management_endpoints.py
View file

@ -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):

2
litellm/types/router.py
View file

@ -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():

7
tests/unit/proxy/public_endpoints/public_v1/test_model_hub.py
View file

@ -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,

79
tests/unit/proxy/test_proxy_utils.py
View file

@ -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 [

73
tests/unit/test_router/test_router.py
View file

@ -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=[