fix(complexity_router): give tool turns their own classifier breaker, floor adaptive escalation, share the tool context budget

A user_turn tool turn whose classifier call fails now opens a circuit breaker of its own, so it no longer stops classification of other sessions' new asks. Adaptive picks on a tool turn keep the held tier as a floor, held turns report the classifier failure reason and error type, and once a loop's tool traffic overflows classifier_context_budget_chars each entry is cut to a quarter of it so one large read cannot push the task out.
This commit is contained in:
mateo-berri 2026-10-28 18:12:21 -07:00
parent 3fab335add
commit f9519d9bae
4 changed files with 213 additions and 20 deletions

View file

@ -6577,8 +6577,9 @@ def tool_exchange_texts(message: Mapping[str, object]) -> tuple[tuple[ToolExchan
Chat completions puts calls on an assistant turn's `tool_calls` and each result on its own `tool` role,
while Anthropic Messages puts both inside content blocks, `tool_use` on the assistant turn and
`tool_result` on the next user turn, so this reads both places. The legacy `function_call` field and
`function` role, and Anthropic's server and MCP tool blocks, read the same way. Responses API input reaches
callers already resolved to the chat-completions shape by `resolve_structured_messages`.
`function` role, Anthropic's server tool calls, and its MCP tool blocks read the same way, while server tool
results such as `web_search_tool_result` are not read. Responses API input reaches callers already resolved
to the chat-completions shape by `resolve_structured_messages`.
"""
if message.get("role") in ("tool", "function"):
result: Final = _exchange_content_text(message.get("content"))

View file

@ -516,6 +516,7 @@ _CODEX_REMINDER_MARKERS: Final = _DEFAULT_REMINDER_MARKERS + (
_TRUNCATION_MARKER: Final = "..."
_TRUNCATION_HEAD_FRACTION: Final = 0.3
_MIN_QUOTED_TURN_CHARS: Final = 120
_TOOL_ENTRY_SHARES_OF_BUDGET: Final = 4
_CLASSIFIER_CIRCUIT_OPEN_SIGNAL: Final = "classifier-circuit-open"
@ -1064,6 +1065,10 @@ def _extract_agent_progress(
needs on a continuation turn. Only what follows the newest human ask counts, since tool traffic
before it served an earlier ask. Newest first is what the budget keeps, so the latest result always
reaches the classifier. Harness reminders are stripped from tool output the same way they are from asks.
A loop that overflows the budget has each entry cut to a quarter of it first, since one large file read
would otherwise fill the block alone and push out the result that described the task. A loop that fits
is quoted whole.
"""
from litellm.litellm_core_utils.prompt_templates.factory import tool_exchange_texts
@ -1083,7 +1088,13 @@ def _extract_agent_progress(
if per_turn_chars is None
else tuple((kind, _truncate(text, per_turn_chars)) for kind, text in newest_first)
)
return tuple(reversed(_turns_within_budget(clamped, budget_chars)))
overflows: Final = sum(len(text) for _, text in clamped) > budget_chars
shared: Final = (
tuple((kind, _truncate(text, budget_chars // _TOOL_ENTRY_SHARES_OF_BUDGET)) for kind, text in clamped)
if overflows
else clamped
)
return tuple(reversed(_turns_within_budget(shared, budget_chars)))
def _decision_is_pinnable(decision: StandardLoggingRoutingDecision | None) -> bool:
@ -1649,6 +1660,9 @@ class ComplexityRouter(CustomLogger):
self._classifier_circuit_breaker: _ClassifierCircuitBreaker | None = (
_ClassifierCircuitBreaker(circuit_breaker_cooldown) if circuit_breaker_cooldown is not None else None
)
self._tool_turn_circuit_breaker: _ClassifierCircuitBreaker | None = (
_ClassifierCircuitBreaker(circuit_breaker_cooldown) if circuit_breaker_cooldown is not None else None
)
self._tier_success_predictor: TierSuccessPredictor | None = (
TierSuccessPredictor(
resolve_tier_artifact(self.config.heuristic_v2_artifact),
@ -2089,6 +2103,8 @@ class ComplexityRouter(CustomLogger):
request_kwargs: dict[str, object] | None = None,
messages: Sequence[Mapping[str, object]] | None = None,
raw_messages: list[dict[str, object]] | None = None, # mutable-ok: same shape _run_routing_plugins receives
*,
tool_turn: bool = False,
) -> ClassificationOutcome:
"""
Classify a prompt by complexity, using the LLM classifier when configured.
@ -2099,26 +2115,32 @@ class ComplexityRouter(CustomLogger):
plugin fails, times out, or produces no usable tier, the configured fallback_tier wins on a
custom tier set, and classifier_fallback otherwise decides between the heuristic scorer and
default_model. The outcome's `cause` reports which path actually ran.
A user_turn tool turn (`tool_turn`) trips a circuit breaker of its own: its failure only keeps the held
tier, so it must not switch the classifier off for every other session's new asks.
"""
breaker: Final = self._tool_turn_circuit_breaker if tool_turn else self._classifier_circuit_breaker
if self.config.classifier_type == "heuristic_v2":
return self._classify_with_heuristic_v2(prompt)
if self.config.classifier_type == "custom":
return await self._classify_with_plugin(prompt, system_prompt, request_kwargs, raw_messages)
if self.config.classifier_type == "oss_classifier":
return await self._jev_classifier_outcome(prompt, system_prompt, request_kwargs, messages)
return await self._jev_classifier_outcome(prompt, system_prompt, request_kwargs, messages, breaker=breaker)
if self.config.classifier_type in ("heuristic_first", "hybrid") and _encrypted_classifier_task(
request_kwargs, self._reminder_markers_for_request(request_kwargs or EMPTY_MAPPING)
):
return await self._llm_classifier_outcome(prompt, system_prompt, request_kwargs, messages)
return await self._llm_classifier_outcome(prompt, system_prompt, request_kwargs, messages, breaker=breaker)
if self.config.classifier_type == "heuristic_first" and self.config.classifier_llm_config is not None:
return await self._classify_heuristic_first(prompt, system_prompt, request_kwargs, messages)
return await self._classify_heuristic_first(
prompt, system_prompt, request_kwargs, messages, breaker=breaker
)
if self.config.classifier_type == "hybrid" and self.config.classifier_llm_config is not None:
return await self._classify_hybrid(prompt, system_prompt, request_kwargs, messages)
return await self._classify_hybrid(prompt, system_prompt, request_kwargs, messages, breaker=breaker)
if self.config.classifier_type == "capability" and self.config.classifier_llm_config is not None:
return await self._capability_classifier_outcome(prompt, request_kwargs, messages)
return await self._capability_classifier_outcome(prompt, request_kwargs, messages, breaker=breaker)
if self.config.classifier_type not in ("llm", "llm_v2") or self.config.classifier_llm_config is None:
return self._classify_locally(prompt, system_prompt)
return await self._llm_classifier_outcome(prompt, system_prompt, request_kwargs, messages)
return await self._llm_classifier_outcome(prompt, system_prompt, request_kwargs, messages, breaker=breaker)
def _classify_locally(self, prompt: str, system_prompt: str | None) -> ClassificationOutcome:
if self._tier_success_predictor is not None:
@ -2175,6 +2197,8 @@ class ComplexityRouter(CustomLogger):
system_prompt: str | None,
request_kwargs: dict[str, object] | None, # mutable-ok: handed to _classify_with_llm as-is
messages: Sequence[Mapping[str, object]] | None,
*,
breaker: _ClassifierCircuitBreaker | None,
) -> ClassificationOutcome:
"""Score locally, and only pay for the classifier call when the scorer did not confidently
place the request at or below heuristic_first_max_tier.
@ -2201,7 +2225,9 @@ class ComplexityRouter(CustomLogger):
)
if decided_cheaply:
return scored._replace(cause="heuristic_first_short_circuit")
return await self._llm_classifier_outcome(prompt, system_prompt, request_kwargs, messages, scored=scored)
return await self._llm_classifier_outcome(
prompt, system_prompt, request_kwargs, messages, scored=scored, breaker=breaker
)
async def _classify_hybrid(
self,
@ -2209,6 +2235,8 @@ class ComplexityRouter(CustomLogger):
system_prompt: str | None,
request_kwargs: dict[str, object] | None, # mutable-ok: handed to _classify_with_llm as-is
messages: Sequence[Mapping[str, object]] | None,
*,
breaker: _ClassifierCircuitBreaker | None,
) -> ClassificationOutcome:
"""Score locally, and only pay for the classifier when the score sits near a tier boundary.
@ -2229,7 +2257,9 @@ class ComplexityRouter(CustomLogger):
)
if decided:
return scored._replace(cause="hybrid_short_circuit")
return await self._llm_classifier_outcome(prompt, system_prompt, request_kwargs, messages, scored=scored)
return await self._llm_classifier_outcome(
prompt, system_prompt, request_kwargs, messages, scored=scored, breaker=breaker
)
def _classifier_image_parts(
self, messages: Sequence[Mapping[str, object]] | None
@ -2263,9 +2293,10 @@ class ComplexityRouter(CustomLogger):
prompt: str,
request_kwargs: Mapping[str, object] | None,
messages: Sequence[Mapping[str, object]] | None,
*,
breaker: _ClassifierCircuitBreaker | None,
) -> ClassificationOutcome:
"""Forecast efficient-tier success, then apply the deterministic boundary policy."""
breaker: Final = self._classifier_circuit_breaker
permit: Final = breaker.acquire_permit() if breaker is not None else None
if breaker is not None and permit is None:
return self._capability_classifier_failure_outcome(
@ -2340,13 +2371,14 @@ class ComplexityRouter(CustomLogger):
request_kwargs: dict[str, object] | None, # mutable-ok: handed to _classify_with_llm as-is
messages: Sequence[Mapping[str, object]] | None,
scored: ClassificationOutcome | None = None,
*,
breaker: _ClassifierCircuitBreaker | None,
) -> ClassificationOutcome:
"""Call the LLM classifier and turn its verdict, or its failure, into an outcome.
`scored` is the heuristic outcome the caller already computed, which only "heuristic_first"
has. It is handed to the failure path so a classifier error does not re-run the scorer.
"""
breaker: Final = self._classifier_circuit_breaker
permit: Final = breaker.acquire_permit() if breaker is not None else None
if breaker is not None and permit is None:
return self._classifier_failure_outcome(
@ -2398,6 +2430,8 @@ class ComplexityRouter(CustomLogger):
system_prompt: str | None,
request_kwargs: Mapping[str, object] | None,
messages: Sequence[Mapping[str, object]] | None,
*,
breaker: _ClassifierCircuitBreaker | None,
) -> ClassificationOutcome:
config: Final = self.config.opensource_classifier_config
client: Final = self._jev_client
@ -2413,7 +2447,6 @@ class ComplexityRouter(CustomLogger):
system_prompt,
failure_reason="unsupported_input",
)
breaker: Final = self._classifier_circuit_breaker
permit: Final = breaker.acquire_permit() if breaker is not None else None
if breaker is not None and permit is None:
return self._classifier_failure_outcome(
@ -4756,6 +4789,12 @@ class ComplexityRouter(CustomLogger):
classifier_cost=tool_turn_escalation.classifier_cost
if tool_turn_escalation is not None
else None,
classifier_failure_reason=tool_turn_escalation.classifier_failure_reason
if tool_turn_escalation is not None
else None,
classifier_error_type=tool_turn_escalation.classifier_error_type
if tool_turn_escalation is not None
else None,
conversation_continuing=conversation_continuing,
tier_litellm_params=session_tier_litellm_params,
context_escalation_original_tier=pin_context_original_tier,
@ -4824,7 +4863,7 @@ class ComplexityRouter(CustomLogger):
if held_tier is None or ask is None:
return None
classification: Final = await self.aclassify(
ask, system_prompt, request_kwargs, resolved_messages, raw_messages=messages
ask, system_prompt, request_kwargs, resolved_messages, raw_messages=messages, tool_turn=True
)
if classification.classifier_failure_reason is not None or self._active_tier_severity(
classification.tier
@ -4839,6 +4878,7 @@ class ComplexityRouter(CustomLogger):
resolved_messages=resolved_messages,
context_fit=context_fit,
classification=classification,
held_tier=held_tier,
)
routed_tier: Final = (
routed.routing_decision.get("tier") if routed is not None and routed.routing_decision is not None else None
@ -4964,6 +5004,7 @@ class ComplexityRouter(CustomLogger):
resolved_messages: Sequence[Mapping[str, object]] | None = None,
context_fit: _RequestContextFit | None = None,
classification: ClassificationOutcome | None = None,
held_tier: ComplexityTier | str | None = None,
) -> PreRoutingHookResponse | None:
"""
Classifies the request by complexity and returns the appropriate model.
@ -4981,6 +5022,8 @@ class ComplexityRouter(CustomLogger):
caller does not have to.
classification: The classifier's verdict when the caller already ran it on this
request, so it is not asked twice.
held_tier: The session tier a user_turn tool turn escalates from, a hard floor for
the adaptive pick.
Returns:
PreRoutingHookResponse with the routed model, or None if no routing needed.
@ -5198,12 +5241,21 @@ class ComplexityRouter(CustomLogger):
housekeeping_ceiling: Final = tier if outcome.cause == "housekeeping" else None
# A context-escalated tier becomes the hard floor: a floor the bandit can slide
# under is not a floor.
adaptive_floor: Final = (
decided_floor: Final = (
tier
if context_original_tier is not None
or outcome.cause in ("capability_classifier", "capability_classifier_fallback")
else plan_floor
)
adaptive_floor: Final = (
held_tier
if decided_floor is None
or (
held_tier is not None
and self._active_tier_severity(held_tier) > self._active_tier_severity(decided_floor)
)
else decided_floor
)
adaptive_fit: Final = context_placement.holdable_models if context_placement is not None else None
sampled_model: Final = self._soft_floor_pick(
tier,

View file

@ -1376,10 +1376,15 @@ class ComplexityRouterConfig(LiteLLMBaseModel):
"classified this way. With classification_mode='user_turn' the tier pinned for the ask "
"is a floor for its tool turns: each one is classified, and only a higher tier moves the "
"session up and replaces the pin, so a loop never drops back mid-task, at the cost of a "
"classifier call per tool turn. Under heuristic_first and hybrid a turn with tool traffic "
"classifier call per tool turn, which the turn waits on, up to the classifier timeout, before it "
"routes. A tool turn whose classifier call fails keeps the held tier, and with the circuit breaker "
"on it opens a breaker of its own, so later tool turns replay the pin at once while new asks keep "
"classifying. Under heuristic_first and hybrid a turn with tool traffic "
"always goes to the classifier model, since the local scorer reads only the ask. The block is bounded "
"on its own by classifier_context_budget_chars and classifier_context_per_turn_chars, "
"newest first, but not by classifier_context_window_size, since one ask can run dozens "
"newest first, and once a loop overflows that budget each entry is cut to a quarter of it, so "
"one large file read cannot push out the rest. It is not bounded by "
"classifier_context_window_size, since one ask can run dozens "
"of tool calls; setting the window size to 0 still turns all context off. Tool output "
"such as file contents is sent to the classifier model, which may be a different "
"deployment or provider than the routed completion model. Off by default because "

View file

@ -11561,7 +11561,8 @@ class _TaskAwareJevClient:
class _ScriptedJevClient:
"""Answers the n-th classification with the n-th scripted tier, and fails it where the script says None."""
"""Answers the n-th classification with the n-th scripted tier, fails it where the script says None, and
times it out where the script says TIMEOUT."""
def __init__(self, tiers: tuple[str | None, ...]) -> None:
self._tiers: Final = tiers
@ -11574,6 +11575,8 @@ class _ScriptedJevClient:
tier: Final = self._tiers[len(self.states) - 1]
if tier is None:
raise httpx.ConnectError("classifier unreachable")
if tier == "TIMEOUT":
raise TimeoutError("classifier timed out")
return JevSystemOneResponse(
model="jev-1.13.0",
answers={"tier": JevChoiceAnswer(type="choice", choice=tier, probabilities={tier: 1.0}, confidence=0.9)},
@ -11699,6 +11702,37 @@ class TestAgentProgressReachesClassifier:
assert response.model == "strong"
assert "listing 0" not in client.states[-1]
@pytest.mark.asyncio
async def test_one_large_read_does_not_push_the_task_out(self) -> None:
"""A file read larger than the whole budget must not leave the classifier with that file alone and
drop the result that described the task."""
client: Final = _TaskAwareJevClient()
router: Final = self._jev_router(client, classifier_context_include_tool_turns=True)
large_read: Final = (
{
"role": "assistant",
"content": None,
"tool_calls": [
{
"id": "call_9",
"type": "function",
"function": {"name": "read_file", "arguments": '{"path": "ledger.py"}'},
}
],
},
{"role": "tool", "tool_call_id": "call_9", "content": "balance = 0\n" * 2000},
)
response: Final = await router.async_pre_routing_hook(
model="test-router",
request_kwargs={},
messages=[*_tool_loop_continuation("chat_completions"), *large_read],
)
assert response is not None
assert response.model == "strong"
assert "balance = 0" in client.states[-1]
@pytest.mark.asyncio
async def test_zero_context_window_turns_tool_context_off(self) -> None:
client: Final = _TaskAwareJevClient()
@ -11741,7 +11775,7 @@ class TestAgentProgressReachesClassifier:
async def test_user_turn_holds_the_ask_tier_as_a_floor_and_moves_up_on_tool_results(
self, include_tool_turns: bool, expected_models: str
) -> None:
"""The customer's ask: keep the current tier as a baseline and move up when the task turns out harder.
"""A user_turn loop keeps its ask's tier as a floor and moves up when the task turns out harder.
The third turn grades SIMPLE again, and the loop must stay where the second turn took it."""
client: Final = _ScriptedJevClient(("SIMPLE", "COMPLEX", "SIMPLE"))
router: Final = self._jev_router(
@ -11786,6 +11820,107 @@ class TestAgentProgressReachesClassifier:
assert [response.model for response in responses if response is not None] == ["cheap", "cheap"]
assert len(client.states) == 2
held_decision: Final = responses[1].routing_decision if responses[1] is not None else None
assert held_decision is not None
assert held_decision.get("classifier_failure_reason") == "classifier_error"
assert held_decision.get("classifier_error_type") == "ConnectError"
@pytest.mark.asyncio
async def test_user_turn_tool_turn_timeout_leaves_new_asks_classifying(self) -> None:
"""A tool turn's timeout only keeps its own held tier, so another session's new ask must still reach the
classifier, while the timed-out loop's next tool turn replays its pin without waiting on the call again."""
client: Final = _ScriptedJevClient(("SIMPLE", "TIMEOUT", "COMPLEX"))
router: Final = self._jev_router(
client, classification_mode="user_turn", classifier_context_include_tool_turns=True
)
loop: Final = _tool_loop_continuation("chat_completions")
new_ask: Final = [{"role": "user", "content": "Design a lock-free MPMC queue and prove linearizability"}]
responses: Final = [
await router.async_pre_routing_hook(
model="test-router", request_kwargs={"metadata": {"session_id": session}}, messages=turn
)
for session, turn in (
("loop-a", loop[:1]),
("loop-a", loop),
("loop-b", new_ask),
("loop-a", [*loop, *_RUN_TESTS_TURNS]),
)
]
assert [response.model for response in responses if response is not None] == [
"cheap",
"cheap",
"strong",
"cheap",
]
assert len(client.states) == 3
decisions: Final = [response.routing_decision if response is not None else None for response in responses]
assert decisions[1] is not None and decisions[1].get("classifier_failure_reason") == "timeout"
assert decisions[3] is not None and decisions[3].get("classifier_failure_reason") == "circuit_open"
@pytest.mark.asyncio
async def test_user_turn_escalation_never_samples_below_the_held_tier(self) -> None:
"""An adaptive router scoring every model must not let an escalating tool turn land below the tier the
session holds, however strongly the bandit favors the cheap model."""
from itertools import product
from litellm.router_strategy.adaptive_router.bandit import BanditCell
from litellm.types.router import RequestType
instance: Final = MagicMock(cache=DualCache())
instance.model_list = [
{
"model_name": name,
"litellm_params": {"model": f"openai/{name}", "input_cost_per_token": cost},
"model_info": {"adaptive_router_preferences": {"quality_tier": quality, "strengths": []}},
}
for name, cost, quality in (
("cheap", 0.00000015, 1),
("mid", 0.000001, 2),
("strong", 0.000005, 3),
("top", 0.000015, 3),
)
]
_answer_deployments_named(instance)
router: Final = ComplexityRouter(
"test-router",
instance,
{
"classifier_type": "jev",
"jev_classifier_config": {"api_key": "test", "timeout_ms": 100},
"tiers": _AGENT_TIERS,
"classification_mode": "user_turn",
"classifier_context_include_tool_turns": True,
"adaptive": True,
"tier_distance_penalty": 0.1,
},
derive_savings_baseline=False,
jev_client=_ScriptedJevClient(("MEDIUM", "COMPLEX")),
)
adaptive: Final = router._ensure_adaptive_router()
assert adaptive is not None
adaptive._cells.update(
{
(request_type, name): BanditCell(alpha=20.0, beta=1.0) if name == "cheap" else BanditCell(alpha=1.0, beta=20.0)
for request_type, name in product(RequestType, _AGENT_TIERS.values())
}
)
loop: Final = _tool_loop_continuation("chat_completions")
with patch(
"litellm.router_strategy.adaptive_router.bandit.thompson_sample",
side_effect=lambda cell, rng=None: cell.alpha / (cell.alpha + cell.beta),
):
responses: Final = [
await router.async_pre_routing_hook(
model="test-router", request_kwargs={"metadata": {"session_id": "agent-loop"}}, messages=turn
)
for turn in (loop[:1], loop)
]
assert [response.model for response in responses if response is not None][0] == "cheap"
assert responses[1] is not None and responses[1].model in ("mid", "strong", "top")
@pytest.mark.asyncio
@pytest.mark.parametrize(