diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 72f22e0bab0..68dfde394e7 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -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) diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index 38411ef52ea..c8916f10965 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -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": "", "close": ""}]} 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(