scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-16 23:41:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

Merge pull request #41326 from BerriAI/litellm_forecast_classifier_entitlement

Browse source 
feat(router): limit unlicensed Capability and Fuse v2 routers to one each
This commit is contained in:
tin-berri 2026-09-15 17:36:22 -07:00 • committed by GitHub
commit 878716f806
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 153 additions and 15 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm/proxy/auth/litellm_license.py
View file

@ -155,8 +155,8 @@ class LicenseCheck:
def auto_router_capability_limit(self) -> int | None:
"""
How many auto-routers may claim each licensed capability (heuristic_v2, operator-defined
tier_definitions): unlimited (None) only when the signed license lists the auto_router
How many auto-routers may claim each gated classifier or customization capability:
unlimited (None) only when the signed license lists the auto_router
feature, otherwise one per capability. A license verified through the API carries no
feature list, so it does not lift the limit either.
"""

28
litellm/router_utils/auto_router_model_naming.py
View file

@ -218,9 +218,8 @@ class GatedAutoRouterCapability:
stored ``litellm_params`` (``{config}`` is the caller's expression for the normalized
``complexity_router_config`` jsonb, substituted as many times as the predicate needs); they live
on one record so they cannot drift apart. ``subject`` and ``remedy`` build the shared refusal
message. A validated config claims at most one capability, and the validator is what makes that
true: tier_definitions rejects every heuristic classifier_type, and it also rejects the
classifier system_prompt, which in turn only applies to the classifier types heuristic_v2 is not.
message. A validated config claims at most one capability: gated classifier types cannot be
combined with operator-defined tiers or classifier prompts.
"""
key: str
@ -238,6 +237,22 @@ HEURISTIC_V2_CAPABILITY: Final = GatedAutoRouterCapability(
sql_config_predicate="{config} ->> 'classifier_type' = 'heuristic_v2'",
)
CAPABILITY_CLASSIFIER_CAPABILITY: Final = GatedAutoRouterCapability(
key="capability",
subject="with classifier_type 'capability' (Capability)",
remedy="Use a different classifier or remove an existing Capability router.",
uses=lambda config: _mapping(config).get("classifier_type") == "capability",
sql_config_predicate="{config} ->> 'classifier_type' = 'capability'",
)
LLM_V2_CAPABILITY: Final = GatedAutoRouterCapability(
key="llm_v2",
subject="with classifier_type 'llm_v2' (Fuse v2)",
remedy="Use a different classifier or remove an existing Fuse v2 router.",
uses=lambda config: _mapping(config).get("classifier_type") == "llm_v2",
sql_config_predicate="{config} ->> 'classifier_type' = 'llm_v2'",
)
_OPERATOR_PROMPT_FIELDS_SQL: Final = " OR ".join(
f"{{config}} ->> '{field}' IS NOT NULL" for field in OPERATOR_CLASSIFIER_PROMPT_FIELDS
)
@ -258,7 +273,12 @@ CUSTOMIZATION_CAPABILITY: Final = GatedAutoRouterCapability(
),
)
GATED_AUTO_ROUTER_CAPABILITIES: Final = (HEURISTIC_V2_CAPABILITY, CUSTOMIZATION_CAPABILITY)
GATED_AUTO_ROUTER_CAPABILITIES: Final = (
HEURISTIC_V2_CAPABILITY,
CAPABILITY_CLASSIFIER_CAPABILITY,
LLM_V2_CAPABILITY,
CUSTOMIZATION_CAPABILITY,
)
def claimed_capability(complexity_router_config: object) -> GatedAutoRouterCapability | None:

45
tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py
View file

@ -4982,6 +4982,26 @@ class TestStrategyRouterWriteValidation:
_V2 = {"classifier_type": "heuristic_v2", "tiers": {"SIMPLE": "gpt-4o-mini"}}
_V1 = {"classifier_type": "heuristic", "tiers": {"SIMPLE": "gpt-4o-mini"}}
_FORECAST_BASE = {
"classifier_llm_config": {"model": "gpt-4o-mini"},
"tiers": {"SIMPLE": "gpt-4o-mini", "REASONING": "gpt-4o"},
}
_CAPABILITY = {
**_FORECAST_BASE,
"classifier_type": "capability",
"capability_classifier_config": {
"efficient_tier": "SIMPLE", "capable_tier": "REASONING", "base_threshold": 0.7,
},
}
_FUSE = {
**_FORECAST_BASE,
"classifier_type": "llm_v2",
"adaptive": False,
"llm_v2_config": {
"efficient_profile": "Small solver", "capable_profile": "Large solver",
"harness": "One attempt", "max_quality_gap": 0.05,
},
}
_CUSTOM_TIERS = {
"classifier_type": "llm",
"classifier_llm_config": {"model": "gpt-4o-mini"},
@ -5049,6 +5069,16 @@ class TestStrategyRouterWriteValidation:
@pytest.mark.parametrize(
"limit,effective_params,db_models,config_config,model_id,expected",
[
(1, {"model": "auto_router/complexity_router", "complexity_router_config": _CAPABILITY}, ["auto_router/complexity_router"], None, None, "refused"),
(1, {"model": "auto_router/complexity_router", "complexity_router_config": _CAPABILITY}, [], _CAPABILITY, None, "refused"),
(1, {"model": "auto_router/complexity_router", "complexity_router_config": _CAPABILITY}, [], _FUSE, None, "reserved"),
(1, {"model": "auto_router/complexity_router", "complexity_router_config": _CAPABILITY}, [], None, "held-id", "reserved"),
(None, {"model": "auto_router/complexity_router", "complexity_router_config": _CAPABILITY}, ["auto_router/complexity_router"], _CAPABILITY, None, "plain"),
(1, {"model": "auto_router/complexity_router", "complexity_router_config": _FUSE}, ["auto_router/complexity_router"], None, None, "refused"),
(1, {"model": "auto_router/complexity_router", "complexity_router_config": _FUSE}, [], _FUSE, None, "refused"),
(1, {"model": "auto_router/complexity_router", "complexity_router_config": _FUSE}, [], _CAPABILITY, None, "reserved"),
(1, {"model": "auto_router/complexity_router", "complexity_router_config": _FUSE}, [], None, "held-id", "reserved"),
(None, {"model": "auto_router/complexity_router", "complexity_router_config": _FUSE}, ["auto_router/complexity_router"], _FUSE, None, "plain"),
(1, {"model": "auto_router/complexity_router", "complexity_router_config": _V2}, ["auto_router/complexity_router"], None, None, "refused"),
(1, {"model": "auto_router/complexity_router", "complexity_router_config": _V2}, [], _V2, None, "refused"),
(1, {"model": "auto_router/complexity_router", "complexity_router_config": _V2}, [], None, None, "reserved"),
@ -5338,7 +5368,8 @@ class TestStrategyRouterWriteValidation:
assert events == ["slot-enter", "slot-exit", "team_model_add"]
@pytest.mark.asyncio
async def test_add_new_model_refuses_a_second_heuristic_v2_router_before_the_db_write(self) -> None:
@pytest.mark.parametrize("config", [_V2, _CAPABILITY, _FUSE])
async def test_add_new_model_refuses_a_second_gated_classifier_router_before_the_db_write(self, config: Mapping[str, object]) -> None:
from litellm.proxy._types import ProxyException
from litellm.proxy.management_endpoints.model_management_endpoints import (
add_new_model,
@ -5366,7 +5397,7 @@ class TestStrategyRouterWriteValidation:
await add_new_model(
model_params=Deployment(
model_name="second-v2",
litellm_params=LiteLLM_Params(model="auto_router/complexity_router", complexity_router_config=self._V2),
litellm_params=LiteLLM_Params(model="auto_router/complexity_router", complexity_router_config=config),
),
user_api_key_dict=admin,
)
@ -5463,7 +5494,8 @@ class TestStrategyRouterWriteValidation:
assert fake.litellm_proxymodeltable.update.await_count == 0
@pytest.mark.asyncio
async def test_patch_model_refuses_switching_another_router_to_heuristic_v2(self) -> None:
@pytest.mark.parametrize("config", [_V2, _CAPABILITY, _FUSE])
async def test_patch_model_refuses_switching_another_router_to_gated_classifier(self, config: Mapping[str, object]) -> None:
"""patch_model relays HTTPException as-is, so the license refusal reaches the client as a plain 403."""
from fastapi import HTTPException
@ -5498,7 +5530,7 @@ class TestStrategyRouterWriteValidation:
with pytest.raises(HTTPException) as exc_info:
await patch_model(
model_id=model_id,
patch_data=updateDeployment(litellm_params=updateLiteLLMParams(complexity_router_config=self._V2)),
patch_data=updateDeployment(litellm_params=updateLiteLLMParams(complexity_router_config=config)),
user_api_key_dict=admin,
)
assert exc_info.value.status_code == 403
@ -5506,7 +5538,8 @@ class TestStrategyRouterWriteValidation:
fake.litellm_proxymodeltable.update.assert_not_awaited()
@pytest.mark.asyncio
async def test_update_model_refuses_switching_another_router_to_heuristic_v2(self) -> None:
@pytest.mark.parametrize("config", [_V2, _CAPABILITY, _FUSE])
async def test_update_model_refuses_switching_another_router_to_gated_classifier(self, config: Mapping[str, object]) -> None:
from litellm.proxy._types import ProxyException
from litellm.proxy.management_endpoints.model_management_endpoints import (
update_model,
@ -5542,7 +5575,7 @@ class TestStrategyRouterWriteValidation:
with pytest.raises(ProxyException) as exc_info:
await update_model(
model_params=updateDeployment(
litellm_params=updateLiteLLMParams(complexity_router_config=self._V2),
litellm_params=updateLiteLLMParams(complexity_router_config=config),
model_info=ModelInfo(id=model_id),
),
user_api_key_dict=admin,

21
tests/test_litellm/proxy/proxy_server/test_proxy_config.py
View file

@ -317,13 +317,28 @@ _TWO_HEURISTIC_V2_ROUTERS_YAML = (
@pytest.mark.asyncio
@pytest.mark.parametrize("license_limit", [1, None])
async def test_ProxyConfig_load_config_takes_the_heuristic_v2_limit_from_the_license_only(
tmp_path, monkeypatch, license_limit: int | None
@pytest.mark.parametrize("classifier_type", ["heuristic_v2", "capability", "llm_v2"])
async def test_ProxyConfig_load_config_takes_the_classifier_limit_from_the_license_only(
tmp_path, monkeypatch, license_limit: int | None, classifier_type: str
) -> None:
"""`router_settings.auto_router_capability_limit` is managed outside config.yaml: an operator
cannot grant the entitlement by editing the config, and a licensed proxy boots both routers."""
f = tmp_path / "c.yaml"
f.write_text(_TWO_HEURISTIC_V2_ROUTERS_YAML)
forecast_settings = {
"capability": (
" classifier_llm_config: {model: gpt-4o-mini}\n"
" capability_classifier_config: {efficient_tier: SIMPLE, capable_tier: REASONING, base_threshold: 0.7}\n"
),
"llm_v2": (
" classifier_llm_config: {model: gpt-4o-mini}\n"
" adaptive: false\n"
" llm_v2_config: {efficient_profile: Small solver, capable_profile: Large solver, harness: One attempt, max_quality_gap: 0.05}\n"
),
}
config_yaml = _TWO_HEURISTIC_V2_ROUTERS_YAML.replace(
"classifier_type: heuristic_v2\n", f"classifier_type: {classifier_type}\n{forecast_settings.get(classifier_type, '')}"
).replace("tiers: {SIMPLE: gpt-4o-mini}", "tiers: {SIMPLE: gpt-4o-mini, REASONING: gpt-4o}")
f.write_text(config_yaml)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False)
monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False)

60
tests/test_litellm/router_strategy/test_complexity_router.py
View file

@ -1417,6 +1417,66 @@ class TestRouterComplexityDeploymentMethods:
router.init_complexity_router_deployment(deployment)
assert "auto_router/complexity_router/test-router" in router.complexity_routers
@staticmethod
def _forecast_row(model_name: str, model_id: str, classifier_type: str) -> dict[str, object]:
settings: Final = (
{"capability_classifier_config": {
"efficient_tier": "SIMPLE", "capable_tier": "REASONING", "base_threshold": 0.7,
}} if classifier_type == "capability" else {
"adaptive": False,
"llm_v2_config": {
"efficient_profile": "Small solver", "capable_profile": "Large solver",
"harness": "One attempt", "max_quality_gap": 0.05,
},
}
)
return {
"model_name": model_name,
"litellm_params": {
"model": "auto_router/complexity_router",
"complexity_router_config": {
"classifier_type": classifier_type,
"classifier_llm_config": {"model": "gpt-4o-mini"},
"tiers": {"SIMPLE": "gpt-4o-mini", "REASONING": "gpt-4o"},
**settings,
},
},
"model_info": {"id": model_id},
}
@pytest.mark.parametrize("classifier_type,sibling", [("capability", "llm_v2"), ("llm_v2", "capability")])
def test_forecast_cap_keeps_edits_and_refuses_extra_routers_and_type_switches(self, classifier_type: str, sibling: str) -> None:
router: Final = Router(
model_list=[
self._POOL,
self._forecast_row("held", "held-id", classifier_type),
self._forecast_row("sibling", "sibling-id", sibling),
self._router_row("other", "other-id", "heuristic_v2"),
self._custom_tier_row("custom", "custom-id"),
],
auto_router_capability_limit=lambda: 1,
ignore_invalid_deployments=True,
)
assert sorted(router.complexity_routers) == ["custom", "held", "other", "sibling"]
assert router.upsert_deployment(Deployment(**self._forecast_row("edited", "held-id", classifier_type))) is not None
assert router.upsert_deployment(Deployment(**self._forecast_row("second", "new-id", classifier_type))) is None
assert router.upsert_deployment(Deployment(**self._forecast_row("switched", "other-id", classifier_type))) is None
assert sorted(router.complexity_routers) == ["custom", "edited", "other", "sibling"]
assert router.upsert_deployment(Deployment(**self._router_row("released", "held-id", "heuristic"))) is not None
assert router.upsert_deployment(Deployment(**self._forecast_row("switched", "other-id", classifier_type))) is not None
assert sorted(router.complexity_routers) == ["custom", "released", "sibling", "switched"]
@pytest.mark.parametrize("classifier_type", ["capability", "llm_v2"])
@pytest.mark.parametrize("limit", [1, None])
def test_forecast_registration_applies_the_resolved_license_limit(self, classifier_type: str, limit: int | None) -> None:
rows: Final = [self._POOL, self._forecast_row("a", "id-a", classifier_type), self._forecast_row("b", "id-b", classifier_type)]
if limit is not None:
with pytest.raises(ValueError, match="At most 1 auto-router"):
Router(model_list=rows, auto_router_capability_limit=lambda: limit)
return
router: Final = Router(model_list=rows, auto_router_capability_limit=lambda: limit)
assert sorted(router.complexity_routers) == ["a", "b"]
@staticmethod
def _router_row(model_name: str, model_id: str, classifier_type: str) -> dict[str, object]:
return {

10
tests/test_litellm/router_utils/test_auto_router_model_naming.py
View file

@ -395,6 +395,8 @@ def test_placement_is_scoped_to_complexity_router_deployments(model, present_fie
_HV2_CONFIG: Mapping[str, object] = {"classifier_type": "heuristic_v2"}
_CAPABILITY_CONFIG: Mapping[str, object] = {"classifier_type": "capability"}
_FUSE_CONFIG: Mapping[str, object] = {"classifier_type": "llm_v2"}
_CUSTOM_TIER_CONFIG: Mapping[str, object] = {
"classifier_type": "llm",
"tier_definitions": [{"name": "routine", "description": "easy"}, {"name": "hard", "description": "hard"}],
@ -457,6 +459,10 @@ def test_is_complexity_router_model(model: str | None, expected: bool) -> None:
@pytest.mark.parametrize(
"litellm_params,expected_key",
[
({"model": "auto_router/complexity_router", "complexity_router_config": _CAPABILITY_CONFIG}, "capability"),
({"model": "auto_router/complexity_router-eu", "complexity_router_config": _FUSE_CONFIG}, "llm_v2"),
({"model": "openai/solver", "complexity_router_config": _CAPABILITY_CONFIG}, None),
({"model": "auto_router/quality_router", "complexity_router_config": _FUSE_CONFIG}, None),
({"model": "auto_router/complexity_router", "complexity_router_config": _HV2_CONFIG}, "heuristic_v2"),
({"model": "auto_router/complexity_router-eu", "complexity_router_config": _HV2_CONFIG}, "heuristic_v2"),
({"model": "auto_router/complexity_router", "complexity_router_config": _CUSTOM_TIER_CONFIG}, "tier_or_classifier_prompt"),
@ -493,6 +499,8 @@ def test_count_capability_routers_counts_only_its_own_capability(capability) ->
by_key = {
"heuristic_v2": (_HV2_CONFIG, _HV2_CONFIG),
"capability": (_CAPABILITY_CONFIG, _CAPABILITY_CONFIG),
"llm_v2": (_FUSE_CONFIG, _FUSE_CONFIG),
"tier_or_classifier_prompt": (_CUSTOM_TIER_CONFIG, _CUSTOM_PROMPT_CONFIG),
}
mine_first, mine_second = by_key[capability.key]
@ -545,6 +553,8 @@ def test_every_gated_capability_has_a_distinct_predicate_and_sql_spelling() -> N
"config",
[
_HV2_CONFIG,
_CAPABILITY_CONFIG,
_FUSE_CONFIG,
_CUSTOM_TIER_CONFIG,
_CUSTOM_PROMPT_CONFIG,
{"classifier_type": "heuristic"},