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 96

fix(router): extract capability tasks with request scoped markers

Browse source
This commit is contained in:
Tin Chi Lo 2026-09-15 11:52:20 -07:00
parent e62f0d0376
commit 896f35c751
2 changed files with 37 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -2210,7 +2210,8 @@ class ComplexityRouter(CustomLogger):
if capability is None or classifier_system_prompt is None:
raise ValueError("capability classifier is not configured")
asks_newest_first: Final = tuple(_iter_human_asks_newest_first(messages or (), self._reminder_markers))
markers: Final = self._reminder_markers_for_request(request_kwargs or EMPTY_MAPPING)
asks_newest_first: Final = tuple(_iter_human_asks_newest_first(messages or (), markers))
opening_task: Final = asks_newest_first[-1] if asks_newest_first else prompt
latest_follow_up: Final = asks_newest_first[0] if len(asks_newest_first) > 1 else None
task_messages: list[AllMessageValues] = [ # mutable-ok: the latest message gains optional image parts below
@ -2239,9 +2240,7 @@ class ComplexityRouter(CustomLogger):
messages_for_call,
request_kwargs,
max_output_tokens=capability.max_output_tokens,
encrypted_task=_encrypted_classifier_task(
request_kwargs, self._reminder_markers_for_request(request_kwargs or EMPTY_MAPPING)
),
encrypted_task=_encrypted_classifier_task(request_kwargs, markers),
)
verdict: Final = parse_capability_classifier_verdict(content)
threshold: Final = verdict.routing_threshold(capability.base_threshold, capability.threshold_step)

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

@ -2588,6 +2588,40 @@ class TestCapabilityClassifier:
complexity_router_config=_capability_router_config(**overrides),
)
@pytest.mark.asyncio
@pytest.mark.parametrize("custom_markers", (False, True))
async def test_task_forecast_uses_request_scoped_codex_markers(
self, mock_router_instance: MagicMock, custom_markers: bool
) -> None:
completion: Final = AsyncMock(return_value=_llm_response(_capability_reply(p_solve=0.8)))
mock_router_instance.acompletion = completion
router: Final = self._router(
mock_router_instance,
escalation_keywords=[],
**({"reminder_markers": [{"open": "<custom>", "close": "</custom>"}]} if custom_markers else {}),
)
envelope: Final = "\n".join(_CODEX_ENVELOPES)
opening: Final = f"{envelope}\nFix nested behavior"
messages: Final = [
{"role": "user", "content": opening},
{"role": "user", "content": "Preserve empty inputs"},
{"role": "user", "content": envelope},
]
original: Final = deepcopy(messages)
for user_agent in ("codex-tui", "curl/8.7.1", "codex_cli_rs/0.62.0"):
result: Final = await router.async_pre_routing_hook(
model="capability-router", messages=messages, request_kwargs={"metadata": {"user_agent": user_agent}}
)
assert result is not None and result.model == "efficient-model"
sent: Final = completion.call_args.kwargs["messages"]
if user_agent.startswith("codex") and not custom_markers:
assert [message["content"] for message in sent[1:]] == ["Fix nested behavior", "Preserve empty inputs"]
else:
assert [message["content"] for message in sent[1:]] == [opening, envelope]
assert result.messages == original
assert completion.await_count == 3
assert messages == original
@pytest.mark.asyncio
@pytest.mark.parametrize("p_solve,expected_model", ((0.95, "capable-model"), (0.98, "efficient-model")))
async def test_fitted_probability_controls_routing_and_preserves_raw_score(