scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-22 00:31:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 51

refactor(router): Reorder async callback filter for deployments to improve request handling

Browse source
This commit is contained in:
DragonAssassin-one 2026-05-11 11:45:10 +08:00
parent 3dc2ada2b3
commit 07d32b87ad
2 changed files with 64 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

20
litellm/router.py
View file

@ -9262,16 +9262,6 @@ class Router:
)
healthy_deployments = _pre_cooldown_deployments
healthy_deployments = await self.async_callback_filter_deployments(
model=model,
healthy_deployments=healthy_deployments,
messages=(
cast(List[AllMessageValues], messages) if messages is not None else None
),
request_kwargs=request_kwargs,
parent_otel_span=parent_otel_span,
)
if self.enable_pre_call_checks and messages is not None:
healthy_deployments = self._pre_call_checks(
model=model,
@ -9290,6 +9280,16 @@ class Router:
),
)
healthy_deployments = await self.async_callback_filter_deployments(
model=model,
healthy_deployments=healthy_deployments,
messages=(
cast(List[AllMessageValues], messages) if messages is not None else None
),
request_kwargs=request_kwargs,
parent_otel_span=parent_otel_span,
)
## ORDER FILTERING ## -> if user set 'order' in deployments, return deployments with lowest order (e.g. order=1 > order=2)
_target_order = (request_kwargs or {}).pop("_target_order", None)
healthy_deployments = litellm.utils._get_order_filtered_deployments(

54
tests/test_litellm/router_strategy/test_router_tag_routing.py
View file

@ -25,6 +25,9 @@ from dotenv import load_dotenv
import litellm
from litellm import Router
from litellm._logging import verbose_logger
from litellm.router_utils.pre_call_checks.prompt_prefix_affinity_check import (
PromptPrefixAffinityCheck,
)
@pytest.mark.asyncio()
@ -219,6 +222,57 @@ async def test_default_tagged_deployments():
assert response_extra_info["model_id"] == "default-model"
@pytest.mark.asyncio()
async def test_prompt_prefix_affinity_does_not_bypass_default_tag_routing_for_untagged_requests():
"""
Ensure prompt-prefix affinity selection cannot bypass tag-based routing defaults.
For untagged requests, tag routing prefers deployments tagged "default" when present.
Prompt-prefix affinity must select only from that tag-filtered candidate set.
"""
router = litellm.Router(
model_list=[
{
"model_name": "gpt-4",
"litellm_params": {
"model": "gpt-4o",
"api_base": "https://exampleopenaiendpoint-production.up.railway.app/",
"tags": ["default"],
},
"model_info": {"id": "default-model"},
},
{
"model_name": "gpt-4",
"litellm_params": {
"model": "gpt-4o-mini",
"api_base": "https://exampleopenaiendpoint-production.up.railway.app/",
"tags": ["teamA"],
},
"model_info": {"id": "teamA-model"},
},
],
enable_tag_filtering=True,
optional_pre_call_checks=["prompt_prefix_affinity"],
prompt_prefix_affinity_tokens=64,
prompt_prefix_affinity_min_tokens=0,
)
# Force the affinity scorer to prefer the non-default deployment if it ever sees it.
def _prefer_team_a(self, prefix_hash: str, deployment_model_id: str) -> int: # noqa: ANN001
return 2 if deployment_model_id == "teamA-model" else 1
with patch.object(PromptPrefixAffinityCheck, "_get_prefix_hash", return_value="hash"):
with patch.object(PromptPrefixAffinityCheck, "_score_deployment", _prefer_team_a):
response = await router.acompletion(
model="gpt-4",
messages=[{"role": "user", "content": "Tell me a joke."}],
mock_response="Tell me a joke.",
)
assert response._hidden_params["model_id"] == "default-model"
for _ in range(5):
# requests tagged with "default", this should pick model with id == "default-model"
response = await router.acompletion(