mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
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:
parent
3fab335add
commit
f9519d9bae
4 changed files with 213 additions and 20 deletions
|
|
@ -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"))
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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 "
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue