mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
feat(router): required-AND tag routing via & prefix
This commit is contained in:
Devin AI
parent
c274cf321c
commit
e858c8dd98
4 changed files with 300 additions and 41 deletions
|
|
@ -7,6 +7,7 @@ Use this to route requests between Teams
|
|||
"""
|
||||
|
||||
import re
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import TYPE_CHECKING, Any, Literal, Optional, Union
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -21,8 +22,8 @@ else:
|
|||
|
||||
|
||||
def _is_valid_deployment_tag_regex(
|
||||
tag_regexes: list[str],
|
||||
header_strings: list[str],
|
||||
tag_regexes: Sequence[str],
|
||||
header_strings: Sequence[str],
|
||||
) -> Optional[str]:
|
||||
"""
|
||||
Test compiled regex patterns against "Header-Name: value" strings.
|
||||
|
|
@ -43,7 +44,9 @@ def _is_valid_deployment_tag_regex(
|
|||
return None
|
||||
|
||||
|
||||
def is_valid_deployment_tag(deployment_tags: list[str], request_tags: list[str], match_any: bool = True) -> bool:
|
||||
def is_valid_deployment_tag(
|
||||
deployment_tags: Sequence[str], request_tags: Sequence[str], match_any: bool = True
|
||||
) -> bool:
|
||||
"""
|
||||
Check if a tag is valid, the matching can be either any or all based on `match_any` flag
|
||||
"""
|
||||
|
|
@ -71,10 +74,10 @@ def is_valid_deployment_tag(deployment_tags: list[str], request_tags: list[str],
|
|||
|
||||
def _match_deployment(
|
||||
deployment: Any,
|
||||
request_tags: Optional[list[str]],
|
||||
header_strings: list[str],
|
||||
request_tags: Sequence[str] | None,
|
||||
header_strings: Sequence[str],
|
||||
match_any: bool,
|
||||
) -> Optional[dict[str, str]]:
|
||||
) -> Mapping[str, str] | None:
|
||||
"""
|
||||
Determine whether *deployment* matches the current request.
|
||||
|
||||
|
|
@ -87,8 +90,8 @@ def _match_deployment(
|
|||
ran and failed, so the regex cannot override strict-tag policy.
|
||||
"""
|
||||
litellm_params = deployment.get("litellm_params", {})
|
||||
deployment_tags: Optional[list[str]] = litellm_params.get("tags")
|
||||
deployment_tag_regex: Optional[list[str]] = litellm_params.get("tag_regex")
|
||||
deployment_tags: Sequence[str] | None = litellm_params.get("tags")
|
||||
deployment_tag_regex: Sequence[str] | None = litellm_params.get("tag_regex")
|
||||
|
||||
# 1. Exact tag match (existing behaviour).
|
||||
if deployment_tags and request_tags:
|
||||
|
|
@ -114,26 +117,63 @@ def _match_deployment(
|
|||
return None
|
||||
|
||||
|
||||
def _split_tags(tags: list[str]) -> tuple[list[str], list[str]]:
|
||||
positive = [t for t in tags if not t.startswith("!")]
|
||||
def _split_tags(tags: Sequence[str]) -> tuple[Sequence[str], Sequence[str], Sequence[str]]:
|
||||
"""
|
||||
Split request tags into (required, positive, excluded) by prefix.
|
||||
|
||||
`&tag` is required-AND (deployment must carry every one of them), `!tag` is
|
||||
negation and a bare tag keeps the existing inclusion semantics. A lone "&" or
|
||||
"!" carries no value, so it is dropped.
|
||||
"""
|
||||
required = [tag[1:] for tag in tags if tag.startswith("&") and len(tag) > 1]
|
||||
excluded = [tag[1:] for tag in tags if tag.startswith("!") and len(tag) > 1]
|
||||
return positive, excluded
|
||||
positive = [t for t in tags if not t.startswith(("&", "!"))]
|
||||
return required, positive, excluded
|
||||
|
||||
|
||||
def _exclude_deployments(
|
||||
deployments: Union[list[Any], dict[Any, Any]],
|
||||
deployments: Union[Sequence[Any], Mapping[Any, Any]],
|
||||
excluded_set: frozenset[str],
|
||||
) -> list[Any]:
|
||||
) -> Sequence[Any]:
|
||||
if not excluded_set:
|
||||
return list(deployments)
|
||||
return [d for d in deployments if not excluded_set.intersection(d.get("litellm_params", {}).get("tags") or [])]
|
||||
|
||||
|
||||
def _require_all_tags(
|
||||
deployments: Sequence[Any],
|
||||
required_set: frozenset[str],
|
||||
) -> Sequence[Any]:
|
||||
if not required_set:
|
||||
return deployments
|
||||
return [d for d in deployments if required_set.issubset(d.get("litellm_params", {}).get("tags") or [])]
|
||||
|
||||
|
||||
def _default_pool(deployments: Sequence[Any]) -> Sequence[Any]:
|
||||
return [d for d in deployments if "default" in (d.get("litellm_params", {}).get("tags") or [])]
|
||||
|
||||
|
||||
def _tag_routing_metadata(
|
||||
deployment: Mapping[str, Any],
|
||||
matched_via: str,
|
||||
matched_value: str,
|
||||
request_tags: Sequence[str] | None,
|
||||
user_agent: str,
|
||||
) -> Mapping[str, Any]:
|
||||
return {
|
||||
"matched_deployment": deployment.get("model_name"),
|
||||
"matched_via": matched_via,
|
||||
"matched_value": matched_value,
|
||||
"request_tags": request_tags or [],
|
||||
"user_agent": user_agent,
|
||||
}
|
||||
|
||||
|
||||
def _require_candidates(
|
||||
candidates: list[Any],
|
||||
candidates: Sequence[Any],
|
||||
model: str,
|
||||
request_tags: Any,
|
||||
) -> list[Any]:
|
||||
) -> Sequence[Any]:
|
||||
if not candidates:
|
||||
raise ValueError(
|
||||
f"{RouterErrors.no_deployments_with_tag_routing.value}. Passed model={model} and tags={request_tags}"
|
||||
|
|
@ -142,8 +182,8 @@ def _require_candidates(
|
|||
|
||||
|
||||
def _ban_only_base_pool(
|
||||
deployments: Union[list[Any], dict[Any, Any]],
|
||||
) -> list[Any]:
|
||||
deployments: Union[Sequence[Any], Mapping[Any, Any]],
|
||||
) -> Sequence[Any]:
|
||||
# Mirrors untagged-request semantics so callers can't use !tags to escape the default pool.
|
||||
defaults = [d for d in deployments if "default" in (d.get("litellm_params", {}).get("tags") or [])]
|
||||
return defaults if defaults else list(deployments)
|
||||
|
|
@ -190,23 +230,39 @@ async def get_deployments_for_tag(
|
|||
# Build header strings for regex matching from what the proxy already stores.
|
||||
# Currently we match against User-Agent; format matches "^User-Agent: claude-code/..."
|
||||
user_agent = metadata.get("user_agent", "")
|
||||
header_strings: list[str] = [f"User-Agent: {user_agent}"] if user_agent else []
|
||||
header_strings: Sequence[str] = [f"User-Agent: {user_agent}"] if user_agent else []
|
||||
|
||||
positive_tags, excluded_patterns = _split_tags(request_tags or [])
|
||||
required_tags, positive_tags, excluded_patterns = _split_tags(request_tags or [])
|
||||
|
||||
excluded_set = frozenset(excluded_patterns)
|
||||
candidates = _exclude_deployments(healthy_deployments, excluded_set)
|
||||
required_set = frozenset(required_tags)
|
||||
allowed_deployments = _exclude_deployments(healthy_deployments, excluded_set)
|
||||
candidates = _require_all_tags(allowed_deployments, required_set)
|
||||
|
||||
default_deployments = _default_pool(allowed_deployments)
|
||||
|
||||
has_regex_deployments = any(d.get("litellm_params", {}).get("tag_regex") for d in candidates)
|
||||
has_tag_filter = bool(positive_tags) or (bool(header_strings) and has_regex_deployments)
|
||||
ban_only = bool(excluded_set) and not has_tag_filter
|
||||
ban_only = bool(excluded_set) and not required_set and not has_tag_filter
|
||||
|
||||
if ban_only:
|
||||
pool = _exclude_deployments(_ban_only_base_pool(healthy_deployments), excluded_set)
|
||||
return _require_candidates(pool, model, request_tags)
|
||||
|
||||
if required_set and not has_tag_filter:
|
||||
if candidates:
|
||||
if "tag_routing" not in metadata:
|
||||
metadata["tag_routing"] = _tag_routing_metadata(
|
||||
deployment=candidates[0],
|
||||
matched_via="required_tags",
|
||||
matched_value=",".join(required_tags),
|
||||
request_tags=request_tags,
|
||||
user_agent=user_agent,
|
||||
)
|
||||
return candidates
|
||||
return _require_candidates(default_deployments, model, request_tags)
|
||||
|
||||
new_healthy_deployments: list[Any] = []
|
||||
default_deployments: list[Any] = []
|
||||
|
||||
if has_tag_filter:
|
||||
verbose_logger.debug(
|
||||
|
|
@ -215,8 +271,6 @@ async def get_deployments_for_tag(
|
|||
user_agent,
|
||||
)
|
||||
for deployment in candidates:
|
||||
deployment_tags = deployment.get("litellm_params", {}).get("tags")
|
||||
|
||||
match_result = _match_deployment(
|
||||
deployment=deployment,
|
||||
request_tags=positive_tags,
|
||||
|
|
@ -232,18 +286,15 @@ async def get_deployments_for_tag(
|
|||
match_result["matched_value"],
|
||||
)
|
||||
if "tag_routing" not in metadata:
|
||||
metadata["tag_routing"] = {
|
||||
"matched_deployment": deployment.get("model_name"),
|
||||
"matched_via": match_result["matched_via"],
|
||||
"matched_value": match_result["matched_value"],
|
||||
"request_tags": request_tags or [],
|
||||
"user_agent": user_agent,
|
||||
}
|
||||
metadata["tag_routing"] = _tag_routing_metadata(
|
||||
deployment=deployment,
|
||||
matched_via=match_result["matched_via"],
|
||||
matched_value=match_result["matched_value"],
|
||||
request_tags=request_tags,
|
||||
user_agent=user_agent,
|
||||
)
|
||||
new_healthy_deployments.append(deployment)
|
||||
|
||||
if deployment_tags and "default" in deployment_tags:
|
||||
default_deployments.append(deployment)
|
||||
|
||||
if len(new_healthy_deployments) == 0 and len(default_deployments) == 0:
|
||||
raise ValueError(
|
||||
f"{RouterErrors.no_deployments_with_tag_routing.value}."
|
||||
|
|
|
|||
|
|
@ -363,6 +363,6 @@
|
|||
"limit": 105
|
||||
},
|
||||
"UP045": {
|
||||
"limit": 17824
|
||||
"limit": 17820
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -429,7 +429,8 @@ def test_get_tags_from_request_kwargs_various_inputs():
|
|||
def test_split_tags_positive_only():
|
||||
from litellm.router_strategy.tag_based_routing import _split_tags
|
||||
|
||||
positive, excluded = _split_tags(["paid", "teamA"])
|
||||
required, positive, excluded = _split_tags(["paid", "teamA"])
|
||||
assert required == []
|
||||
assert positive == ["paid", "teamA"]
|
||||
assert excluded == []
|
||||
|
||||
|
|
@ -437,7 +438,8 @@ def test_split_tags_positive_only():
|
|||
def test_split_tags_negation_only():
|
||||
from litellm.router_strategy.tag_based_routing import _split_tags
|
||||
|
||||
positive, excluded = _split_tags(["!provider:anthropic"])
|
||||
required, positive, excluded = _split_tags(["!provider:anthropic"])
|
||||
assert required == []
|
||||
assert positive == []
|
||||
assert excluded == ["provider:anthropic"]
|
||||
|
||||
|
|
@ -445,7 +447,10 @@ def test_split_tags_negation_only():
|
|||
def test_split_tags_mixed():
|
||||
from litellm.router_strategy.tag_based_routing import _split_tags
|
||||
|
||||
positive, excluded = _split_tags(["paid", "!provider:anthropic", "!inference:cerebras"])
|
||||
required, positive, excluded = _split_tags(
|
||||
["&reasoning_type:high", "paid", "!provider:anthropic", "!inference:cerebras"]
|
||||
)
|
||||
assert required == ["reasoning_type:high"]
|
||||
assert positive == ["paid"]
|
||||
assert len(excluded) == 2
|
||||
|
||||
|
|
@ -454,7 +459,26 @@ def test_split_tags_bare_bang_skipped():
|
|||
from litellm.router_strategy.tag_based_routing import _split_tags
|
||||
|
||||
# A bare "!" with nothing after it is not a valid negation tag; skip it
|
||||
positive, excluded = _split_tags(["paid", "!"])
|
||||
required, positive, excluded = _split_tags(["paid", "!"])
|
||||
assert required == []
|
||||
assert positive == ["paid"]
|
||||
assert excluded == []
|
||||
|
||||
|
||||
def test_split_tags_required_only():
|
||||
from litellm.router_strategy.tag_based_routing import _split_tags
|
||||
|
||||
required, positive, excluded = _split_tags(["&reasoning_type:high", "&provider:anthropic"])
|
||||
assert required == ["reasoning_type:high", "provider:anthropic"]
|
||||
assert positive == []
|
||||
assert excluded == []
|
||||
|
||||
|
||||
def test_split_tags_bare_ampersand_skipped():
|
||||
from litellm.router_strategy.tag_based_routing import _split_tags
|
||||
|
||||
required, positive, excluded = _split_tags(["paid", "&"])
|
||||
assert required == []
|
||||
assert positive == ["paid"]
|
||||
assert excluded == []
|
||||
|
||||
|
|
@ -462,7 +486,8 @@ def test_split_tags_bare_bang_skipped():
|
|||
def test_split_tags_empty():
|
||||
from litellm.router_strategy.tag_based_routing import _split_tags
|
||||
|
||||
positive, excluded = _split_tags([])
|
||||
required, positive, excluded = _split_tags([])
|
||||
assert required == []
|
||||
assert positive == []
|
||||
assert excluded == []
|
||||
|
||||
|
|
@ -1115,3 +1140,186 @@ async def test_request_level_enable_tag_filtering_false_cannot_disable_global():
|
|||
mock_response="hi",
|
||||
)
|
||||
assert response._hidden_params["model_id"] == "team-a-deployment"
|
||||
|
||||
|
||||
# --- get_deployments_for_tag required-AND (& prefix) integration tests ---
|
||||
|
||||
|
||||
def _deployment(deployment_id: str, tags=None, tag_regex=None):
|
||||
litellm_params = {
|
||||
"model": "gpt-4o",
|
||||
"api_base": "https://exampleopenaiendpoint-production.up.railway.app/",
|
||||
}
|
||||
if tags is not None:
|
||||
litellm_params["tags"] = tags
|
||||
if tag_regex is not None:
|
||||
litellm_params["tag_regex"] = tag_regex
|
||||
return {
|
||||
"model_name": "gpt-4",
|
||||
"litellm_params": litellm_params,
|
||||
"model_info": {"id": deployment_id},
|
||||
}
|
||||
|
||||
|
||||
def _tag_router(deployments, match_any: bool = True):
|
||||
return litellm.Router(
|
||||
model_list=deployments,
|
||||
enable_tag_filtering=True,
|
||||
tag_filtering_match_any=match_any,
|
||||
)
|
||||
|
||||
|
||||
async def _route(deployments, tags, user_agent=None, match_any: bool = True, metadata=None):
|
||||
from litellm.router_strategy.tag_based_routing import get_deployments_for_tag
|
||||
|
||||
request_metadata = metadata if metadata is not None else {}
|
||||
request_metadata["tags"] = tags
|
||||
if user_agent is not None:
|
||||
request_metadata["user_agent"] = user_agent
|
||||
result = await get_deployments_for_tag(
|
||||
llm_router_instance=_tag_router(deployments, match_any=match_any),
|
||||
model="gpt-4",
|
||||
healthy_deployments=deployments,
|
||||
request_kwargs={"metadata": request_metadata},
|
||||
)
|
||||
return [d["model_info"]["id"] for d in result]
|
||||
|
||||
|
||||
ANTHROPIC_HIGH = _deployment("anthropic-high", tags=["provider:anthropic", "reasoning_type:high"])
|
||||
ANTHROPIC_LOW = _deployment("anthropic-low", tags=["provider:anthropic", "reasoning_type:low"])
|
||||
OPENAI_HIGH = _deployment("openai-high", tags=["provider:openai", "reasoning_type:high"])
|
||||
|
||||
|
||||
@pytest.mark.asyncio()
|
||||
async def test_required_and_requires_every_ampersand_tag():
|
||||
"""
|
||||
"&reasoning_type:high,&provider:anthropic" must return only the deployment carrying
|
||||
both tags; under the existing OR semantics all three deployments would survive.
|
||||
"""
|
||||
deployments = [ANTHROPIC_HIGH, ANTHROPIC_LOW, OPENAI_HIGH]
|
||||
|
||||
assert await _route(deployments, ["&reasoning_type:high", "&provider:anthropic"]) == ["anthropic-high"]
|
||||
assert sorted(await _route(deployments, ["reasoning_type:high", "provider:anthropic"])) == [
|
||||
"anthropic-high",
|
||||
"anthropic-low",
|
||||
"openai-high",
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio()
|
||||
async def test_required_and_keeps_all_deployments_carrying_the_required_tag():
|
||||
deployments = [ANTHROPIC_HIGH, ANTHROPIC_LOW, OPENAI_HIGH]
|
||||
|
||||
assert sorted(await _route(deployments, ["&reasoning_type:high"])) == ["anthropic-high", "openai-high"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio()
|
||||
async def test_required_and_composes_with_positive_or_and_negation():
|
||||
"""
|
||||
The composed example from the feature request: must be high-reasoning, AND
|
||||
(anthropic OR openai), AND not cerebras-hosted.
|
||||
"""
|
||||
deployments = [
|
||||
_deployment("anthropic-high", tags=["reasoning_type:high", "provider:anthropic"]),
|
||||
_deployment("openai-high-cerebras", tags=["reasoning_type:high", "provider:openai", "inference:cerebras"]),
|
||||
_deployment("openai-high", tags=["reasoning_type:high", "provider:openai"]),
|
||||
_deployment("anthropic-low", tags=["reasoning_type:low", "provider:anthropic"]),
|
||||
_deployment("vertex-high", tags=["reasoning_type:high", "provider:vertex"]),
|
||||
]
|
||||
|
||||
selected = await _route(
|
||||
deployments,
|
||||
["&reasoning_type:high", "provider:anthropic", "provider:openai", "!inference:cerebras"],
|
||||
)
|
||||
|
||||
assert sorted(selected) == ["anthropic-high", "openai-high"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio()
|
||||
async def test_required_and_narrows_before_positive_or():
|
||||
"""A deployment matching a positive tag but missing a required tag must be dropped."""
|
||||
deployments = [ANTHROPIC_HIGH, ANTHROPIC_LOW]
|
||||
|
||||
assert await _route(deployments, ["&reasoning_type:high", "provider:anthropic"]) == ["anthropic-high"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio()
|
||||
async def test_required_and_falls_back_to_default_pool_when_nothing_matches():
|
||||
deployments = [ANTHROPIC_HIGH, _deployment("default-model", tags=["default"])]
|
||||
|
||||
assert await _route(deployments, ["&provider:vertex"]) == ["default-model"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio()
|
||||
async def test_required_and_raises_when_nothing_matches_and_no_default_pool():
|
||||
from litellm.types.router import RouterErrors
|
||||
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
await _route([ANTHROPIC_HIGH, OPENAI_HIGH], ["&provider:vertex"])
|
||||
|
||||
assert RouterErrors.no_deployments_with_tag_routing.value in str(exc_info.value)
|
||||
|
||||
|
||||
@pytest.mark.asyncio()
|
||||
async def test_required_and_blocks_tag_regex_deployment_missing_required_tag():
|
||||
"""A matching User-Agent must not resurrect a deployment that lacks a required tag."""
|
||||
deployments = [
|
||||
_deployment("claude-code-openai", tags=["provider:openai"], tag_regex=[r"^User-Agent: claude-code\/"]),
|
||||
ANTHROPIC_HIGH,
|
||||
]
|
||||
|
||||
selected = await _route(deployments, ["&reasoning_type:high"], user_agent="claude-code/1.2.3")
|
||||
|
||||
assert selected == ["anthropic-high"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio()
|
||||
async def test_required_and_matches_tag_regex_deployment_carrying_required_tag():
|
||||
deployments = [
|
||||
_deployment(
|
||||
"claude-code-high",
|
||||
tags=["reasoning_type:high"],
|
||||
tag_regex=[r"^User-Agent: claude-code\/"],
|
||||
),
|
||||
ANTHROPIC_LOW,
|
||||
]
|
||||
|
||||
selected = await _route(deployments, ["&reasoning_type:high"], user_agent="claude-code/1.2.3")
|
||||
|
||||
assert selected == ["claude-code-high"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio()
|
||||
async def test_required_and_records_tag_routing_metadata():
|
||||
metadata: dict = {}
|
||||
|
||||
await _route([ANTHROPIC_HIGH, ANTHROPIC_LOW], ["&reasoning_type:high"], metadata=metadata)
|
||||
|
||||
assert metadata["tag_routing"] == {
|
||||
"matched_deployment": "gpt-4",
|
||||
"matched_via": "required_tags",
|
||||
"matched_value": "reasoning_type:high",
|
||||
"request_tags": ["&reasoning_type:high"],
|
||||
"user_agent": "",
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio()
|
||||
async def test_required_and_end_to_end_routing():
|
||||
router = _tag_router([ANTHROPIC_HIGH, ANTHROPIC_LOW, OPENAI_HIGH])
|
||||
|
||||
for _ in range(5):
|
||||
response = await router.acompletion(
|
||||
model="gpt-4",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
metadata={"tags": ["&provider:anthropic", "&reasoning_type:high"]},
|
||||
mock_response="hi",
|
||||
)
|
||||
assert response._hidden_params["model_id"] == "anthropic-high"
|
||||
|
||||
|
||||
@pytest.mark.asyncio()
|
||||
async def test_bare_ampersand_tag_does_not_constrain_routing():
|
||||
deployments = [ANTHROPIC_HIGH, ANTHROPIC_LOW]
|
||||
|
||||
assert await _route(deployments, ["&", "reasoning_type:low"]) == ["anthropic-low"]
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
{
|
||||
"LIT001": {
|
||||
"limit": 23287
|
||||
"limit": 23265
|
||||
},
|
||||
"LIT002": {
|
||||
"limit": 27473
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue