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

fix(router): bound auto-router classifier latency

Browse source
This commit is contained in:
moe-berri 2026-09-03 19:10:53 -07:00
parent 16db51e2cf
commit ddc5d8dc37
2 changed files with 88 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

27
litellm/router_strategy/complexity_router/complexity_router.py
View file

@ -1694,16 +1694,23 @@ class ComplexityRouter(CustomLogger):
}
}
response: Final[ModelResponse] = await self.litellm_router_instance.acompletion(
model=llm_config.model,
messages=messages_for_call,
response_format=response_format,
timeout=llm_config.timeout_ms / 1000,
metadata=metadata,
proxy_server_request=proxy_server_request,
turn_off_message_logging=turn_off_message_logging,
**classifier_call_params,
**_parent_session_kwargs(request_kwargs),
classifier_timeout_s: Final[float] = llm_config.timeout_ms / 1000
response: Final[ModelResponse] = await asyncio.wait_for(
self.litellm_router_instance.acompletion(
model=llm_config.model,
messages=messages_for_call,
stream=False,
response_format=response_format,
timeout=classifier_timeout_s,
num_retries=0,
disable_fallbacks=True,
metadata=metadata,
proxy_server_request=proxy_server_request,
turn_off_message_logging=turn_off_message_logging,
**classifier_call_params,
**_parent_session_kwargs(request_kwargs),
),
timeout=classifier_timeout_s,
)
content: Final = response.choices[0].message.content
if not content:

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

@ -1994,6 +1994,77 @@ class TestLLMClassifier:
assert outcome.cause == "llm_classifier"
assert outcome.classifier_cost == pytest.approx(1.35e-05)
@pytest.mark.asyncio
async def test_aclassify_timeout_does_not_inherit_router_retries_or_fallbacks(
self, llm_classifier_config
):
real_router = Router(
model_list=[
{
"model_name": "haiku-classifier",
"litellm_params": {
"model": "openai/mock-classifier",
"api_key": "mock-key",
"mock_timeout": True,
},
},
{
"model_name": "backup-classifier",
"litellm_params": {
"model": "openai/mock-backup-classifier",
"api_key": "mock-key",
"mock_response": '{"tier": "COMPLEX"}',
},
},
],
num_retries=2,
fallbacks=[{"haiku-classifier": ["backup-classifier"]}],
)
config = {
**llm_classifier_config,
"classifier_llm_config": {"model": "haiku-classifier", "timeout_ms": 10},
}
router = ComplexityRouter(
model_name="test-complexity-router",
litellm_router_instance=real_router,
complexity_router_config=config,
)
outcome = await router.aclassify("hi")
assert outcome.cause == "heuristic_scorer"
assert real_router.total_calls["openai/mock-classifier"] == 1
assert real_router.total_calls["openai/mock-backup-classifier"] == 0
@pytest.mark.asyncio
async def test_aclassify_enforces_total_classifier_deadline(
self, mock_router_instance, llm_classifier_config
):
cancelled = asyncio.Event()
async def slow_classifier(**_kwargs: object) -> None:
try:
await asyncio.sleep(1)
except asyncio.CancelledError:
cancelled.set()
raise
mock_router_instance.acompletion = AsyncMock(side_effect=slow_classifier)
config = {
**llm_classifier_config,
"classifier_llm_config": {"model": "haiku-classifier", "timeout_ms": 10},
}
router = ComplexityRouter(
model_name="test-complexity-router",
litellm_router_instance=mock_router_instance,
complexity_router_config=config,
)
outcome = await router.aclassify("hi")
assert outcome.cause == "heuristic_scorer"
assert cancelled.is_set()
@pytest.mark.asyncio
async def test_aclassify_classifier_cost_is_none_when_call_is_unpriced(
self, llm_complexity_router, mock_router_instance