From 528edf3d22c61cc4f5cf9c1dda85affcdd8d58e3 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 21 Sep 2026 01:34:46 +0000 Subject: [PATCH] fix(complexity_router): escalate short turns in long conversations past heuristic_first Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../complexity_router/README.md | 6 + .../complexity_router/complexity_router.py | 14 +- .../complexity_router/config.py | 21 ++ .../test_heuristic_first_long_context_e2e.py | 198 ++++++++++++++++++ .../router_strategy/test_complexity_router.py | 82 ++++++++ .../add_model/ClassificationMethodConfig.tsx | 28 +++ .../add_model/ComplexityRouterConfig.tsx | 2 + .../add_model/add_auto_router_tab.tsx | 1 + .../build_complexity_router_config.test.ts | 10 + .../build_complexity_router_config.ts | 9 + .../classifier_type_transition.test.ts | 3 + .../add_model/classifier_type_transition.ts | 2 + ...d_updated_complexity_router_config.test.ts | 1 + .../edit_auto_router_modal.tsx | 6 + 14 files changed, 382 insertions(+), 1 deletion(-) create mode 100644 tests/e2e/router/test_heuristic_first_long_context_e2e.py diff --git a/litellm/router_strategy/complexity_router/README.md b/litellm/router_strategy/complexity_router/README.md index 6505746bca1..eca0f666147 100644 --- a/litellm/router_strategy/complexity_router/README.md +++ b/litellm/router_strategy/complexity_router/README.md @@ -470,6 +470,7 @@ model_list: complexity_router_config: classifier_type: heuristic_first heuristic_first_max_tier: SIMPLE + heuristic_first_max_context_tokens: 8000 classifier_llm_config: model: gpt-5-mini reasoning_effort: low @@ -510,6 +511,11 @@ A request short-circuits, meaning it routes on the scorer's own tier with no cla two things hold: the scorer landed at or below `heuristic_first_max_tier`, and it produced at least one signal. Everything else goes to the classifier, which then decides as it normally would. +Set `heuristic_first_max_context_tokens` to veto that shortcut when the estimated whole conversation +exceeds the limit. The estimate counts all message text at approximately four characters per token, +so a short newest nudge in a long agentic session still reaches the classifier. Leave it unset to +keep the scorer's tier in control at any conversation size + The signal requirement is what keeps this from quietly routing everything to your cheapest model. A prompt where no dimension fires scores exactly 0.0, which is below `simple_medium`, so the score to tier mapping calls it SIMPLE by default rather than by evidence. Around half of general traffic diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index a3d6ccbd437..e2587b14b83 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -460,6 +460,10 @@ def _message_text(content: object) -> str: return content if isinstance(content, str) else "" +def _estimated_conversation_tokens(messages: Sequence[Mapping[str, object]] | None) -> int: + return sum(len(_message_text(message.get("content"))) // 4 for message in messages or ()) + + def _reminder_block_spans(lowered: str, open_marker: str, close_marker: str) -> Iterator[tuple[int, int]]: """Span of each complete reminder block for one marker pair, left to right. @@ -1925,6 +1929,9 @@ class ComplexityRouter(CustomLogger): A turn carrying images the classifier would see is never decided cheaply: the scorer reads text alone, so its confidence describes a request it has only partly seen, and a trivial caption beside a screenshot is exactly the misrouting vision classification exists to stop. + + A configured conversation-size limit also vetoes the cheap decision because a short newest + turn can conceal a complex task in the preceding agentic context. """ tier, score, signals, cause = self._score_and_classify(prompt, system_prompt) scored: Final = ClassificationOutcome(tier=tier, score=score, signals=signals, cause=cause) @@ -1933,12 +1940,17 @@ class ComplexityRouter(CustomLogger): threshold is not None and bool(signals) and not self._classifier_image_parts(messages) + and not self._exceeds_heuristic_first_context(messages) and self._active_tier_severity(tier) <= self._active_tier_severity(threshold) ) if decided_cheaply: return ClassificationOutcome(tier=tier, score=score, signals=signals, cause="heuristic_first_short_circuit") return await self._llm_classifier_outcome(prompt, system_prompt, request_kwargs, messages, scored=scored) + def _exceeds_heuristic_first_context(self, messages: Sequence[Mapping[str, object]] | None) -> bool: + limit: Final = self.config.heuristic_first_max_context_tokens + return limit is not None and _estimated_conversation_tokens(messages) > limit + async def _classify_hybrid( self, prompt: str, @@ -2691,7 +2703,7 @@ class ComplexityRouter(CustomLogger): else () ) - cumulative_tokens: Final = sum(len(_message_text(msg.get("content"))) // 4 for msg in messages or ()) + cumulative_tokens: Final = _estimated_conversation_tokens(messages) trajectory_block: Final = ( (f"\nConversation so far: ~{cumulative_tokens} tokens across the request",) if has_prior_conversation diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index aa39dff8c53..baed837cf67 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -1065,6 +1065,18 @@ class ComplexityRouterConfig(BaseModel): "may not name the highest one, since that would make the LLM classifier unreachable." ), ) + heuristic_first_max_context_tokens: int | None = Field( + default=None, + gt=0, + description=( + "The estimated size of the whole conversation, counting all message text at approximately four " + "characters per token, above which the local scorer may not decide cheaply and the request goes to " + "the LLM classifier even when the newest turn scores at or below heuristic_first_max_tier. The " + "newest turn in a long agentic session is usually a short nudge such as 'run the tests' or 'why did " + "that fail?' whose token-count signal says nothing about the task living in the conversation. None " + "keeps the scorer's tier at any conversation size." + ), + ) hybrid_boundary_margin: float | None = Field( default=None, ge=0, @@ -1745,6 +1757,15 @@ class ComplexityRouterConfig(BaseModel): ) return self + @model_validator(mode="after") + def _validate_heuristic_first_max_context_tokens(self) -> "ComplexityRouterConfig": + if self.classifier_type != "heuristic_first" and self.heuristic_first_max_context_tokens is not None: + raise ValueError( + f"heuristic_first_max_context_tokens is set but classifier_type is {self.classifier_type!r}; " + "set classifier_type 'heuristic_first' or remove heuristic_first_max_context_tokens" + ) + return self + @model_validator(mode="after") def _validate_hybrid_boundary_margin(self) -> "ComplexityRouterConfig": if self.classifier_type != "hybrid": diff --git a/tests/e2e/router/test_heuristic_first_long_context_e2e.py b/tests/e2e/router/test_heuristic_first_long_context_e2e.py new file mode 100644 index 00000000000..735dd1a3cdc --- /dev/null +++ b/tests/e2e/router/test_heuristic_first_long_context_e2e.py @@ -0,0 +1,198 @@ +"""Live e2e repros for heuristic-first classification of short turns with context.""" + +from __future__ import annotations + +from typing import Final + +import pytest +from complexity_router_client import ComplexityRouterClient +from e2e_config import unique_marker +from e2e_http import StreamingResponse, require_successful_call +from lifecycle import ResourceManager +from models import ( + ChatAssistantTurn, + ChatBody, + ChatMessage, + ChatToolResultTurn, + KeyGenerateBody, + LiteLLMParamsBody, + ToolCall, + ToolCallFunction, +) +from proxy_client import ProxyClient + +pytestmark = pytest.mark.e2e + +ROUTER_BACKENDS: Final = ("gpt-5.5", "claude-haiku-4-5") +SIMPLE_MODELS: Final = frozenset(("openai/gpt-5.5", "gpt-5.5")) + + +@pytest.fixture(scope="module") +def heuristic_first_router(proxy: ProxyClient, request: pytest.FixtureRequest) -> str: + router_name: Final = f"e2e-heuristic-first-router-{unique_marker()}" + model_id: Final = proxy.create_model( + router_name, + LiteLLMParamsBody( + model="auto_router/complexity_router", + complexity_router_config={ + "classifier_type": "heuristic_first", + "heuristic_first_max_tier": "MEDIUM", + "heuristic_first_max_context_tokens": 8000, + "classifier_fallback": "heuristic", + "classifier_llm_config": {"model": "gpt-5.5"}, + "tiers": { + "SIMPLE": "gpt-5.5", + "MEDIUM": "claude-haiku-4-5", + "COMPLEX": "claude-haiku-4-5", + "REASONING": "claude-haiku-4-5", + }, + }, + ), + ) + request.addfinalizer(lambda: proxy.delete_model(model_id)) + return router_name + + +@pytest.fixture +def heuristic_first_key( + resources: ResourceManager, + client: ComplexityRouterClient, + heuristic_first_router: str, +) -> str: + key: Final = client.proxy.generate_key( + KeyGenerateBody( + models=[heuristic_first_router, *ROUTER_BACKENDS], + user_id=f"e2e-heuristic-first-{unique_marker()}", + ) + ) + resources.defer(lambda: client.proxy.delete_key(key)) + return key + + +def _agentic_messages(marker: str) -> tuple[ChatMessage | ChatAssistantTurn | ChatToolResultTurn, ...]: + system: Final = ChatMessage( + role="system", + content=( + f"You are a coding agent operating on a repository. Preserve the marker {marker}. " + "Use tools to inspect files, run tests, and diagnose failures. Keep track of prior " + "commands and their outputs before proposing a fix. Never discard relevant logs or " + "assume that a failed command was unrelated to the current change." + ), + ) + rounds: Final = tuple( + turn + for round_index in range(5) + for turn in ( + ChatMessage( + role="user", + content=( + f"Inspect the repository state for debugging round {round_index} using marker {marker}. " + "Run the relevant checks and report every warning, traceback, and changed file." + ), + ), + ChatAssistantTurn( + content=None, + tool_calls=[ + ToolCall( + id=f"{marker}-call-{round_index}", + type="function", + function=ToolCallFunction( + name="run_tests", + arguments=f'{{"round": {round_index}, "marker": "{marker}"}}', + ), + ) + ], + ), + ChatToolResultTurn( + tool_call_id=f"{marker}-call-{round_index}", + content="\n".join( + f"{marker} round={round_index} line={line_index} " + "synthetic test output records a failing assertion, a retry, a provider " + "response, a stack frame, and the captured repository state for diagnosis" + for line_index in range(120) + ), + ), + ) + ) + return (system, *rounds, ChatMessage(role="user", content=f"why did that fail? {marker}")) + + +def _send( + client: ComplexityRouterClient, + key: str, + body: ChatBody, +) -> StreamingResponse: + response: Final = client.proxy.transport.send( + "/chat/completions", + headers=client.proxy.transport.bearer(key), + json=body, + stream=False, + ) + require_successful_call(response) + return response + + +def _assert_classifier_consulted(response: StreamingResponse, context: str) -> None: + classifier_cost: Final = response.headers.get("x-litellm-classifier-cost") + assert classifier_cost is not None, ( + f"{context}: classifier header missing; observed headers={response.headers!r}; body={response.body[:300]!r}" + ) + try: + parsed_cost: Final = float(classifier_cost) + except ValueError as exc: + raise AssertionError( + f"{context}: classifier header was not parseable as float: {classifier_cost!r}; " + f"observed headers={response.headers!r}" + ) from exc + assert parsed_cost >= 0, f"{context}: classifier cost was negative: {parsed_cost}; headers={response.headers!r}" + + +@pytest.mark.covers("reliability.routing.complexity_heuristic.scores_current_ask_only") +class TestHeuristicFirstLongContext: + def test_short_turn_in_long_agentic_conversation_consults_classifier( + self, + client: ComplexityRouterClient, + heuristic_first_key: str, + heuristic_first_router: str, + ) -> None: + marker: Final = unique_marker() + response: Final = _send( + client, + heuristic_first_key, + ChatBody( + model=heuristic_first_router, + messages=_agentic_messages(marker), + max_tokens=16, + ), + ) + _assert_classifier_consulted( + response, + f"long agentic context marker={marker}", + ) + + def test_short_single_turn_stays_on_heuristic_path( + self, + client: ComplexityRouterClient, + heuristic_first_key: str, + heuristic_first_router: str, + ) -> None: + marker: Final = unique_marker() + response: Final = _send( + client, + heuristic_first_key, + ChatBody( + model=heuristic_first_router, + messages=[ChatMessage(role="user", content=f"why did that fail? {marker}")], + max_tokens=16, + ), + ) + assert "x-litellm-classifier-cost" not in response.headers, ( + f"single-turn heuristic path unexpectedly consulted classifier; " + f"observed headers={response.headers!r}; body={response.body[:300]!r}" + ) + rows: Final = client.proxy.poll_logs_for_key(heuristic_first_key, min_rows=1) + served: Final = tuple(row.model for row in rows if row.model is not None) + assert len(served) == 1 and served[0] in SIMPLE_MODELS, ( + f"single-turn heuristic path should serve SIMPLE backend {sorted(SIMPLE_MODELS)!r}; " + f"observed spend-log models={served!r}; headers={response.headers!r}" + ) diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index ecd25ff654f..b53d5e764b9 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -49,6 +49,7 @@ from litellm.router_strategy.complexity_router.complexity_router import ( KeywordOverride, _built_in_prompt, _ClassifierCircuitBreaker, + _estimated_conversation_tokens, _is_classifier_timeout, _matched_plan_mode_sentinel, classification_system_prompt, @@ -13444,6 +13445,35 @@ class TestHeuristicFirstConfig: assert config.uses_llm_classifier is True assert ComplexityRouterConfig(tiers=dict(HEURISTIC_FIRST_TIERS)).uses_llm_classifier is False + def test_context_limit_is_accepted_on_heuristic_first(self): + config = ComplexityRouterConfig( + tiers=dict(HEURISTIC_FIRST_TIERS), + classifier_type="heuristic_first", + heuristic_first_max_tier="SIMPLE", + heuristic_first_max_context_tokens=8000, + classifier_llm_config={"model": "haiku-classifier"}, + ) + assert config.heuristic_first_max_context_tokens == 8000 + + def test_context_limit_is_rejected_on_llm(self): + with pytest.raises(ValidationError, match="heuristic_first_max_context_tokens is set but classifier_type"): + ComplexityRouterConfig( + tiers=dict(HEURISTIC_FIRST_TIERS), + classifier_type="llm", + heuristic_first_max_context_tokens=8000, + classifier_llm_config={"model": "haiku-classifier"}, + ) + + def test_context_limit_rejects_zero(self): + with pytest.raises(ValidationError, match="greater than 0"): + ComplexityRouterConfig( + tiers=dict(HEURISTIC_FIRST_TIERS), + classifier_type="heuristic_first", + heuristic_first_max_tier="SIMPLE", + heuristic_first_max_context_tokens=0, + classifier_llm_config={"model": "haiku-classifier"}, + ) + class TestHeuristicFirst: """Behavior of the heuristic-first chain: when the classifier call is skipped, and when it is not.""" @@ -13461,6 +13491,58 @@ class TestHeuristicFirst: assert outcome.signals assert outcome.classifier_cost is None + @pytest.mark.asyncio + async def test_long_context_vetoes_cheap_short_turn(self, mock_router_instance): + mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "MEDIUM"}')) + router = _heuristic_first_router( + mock_router_instance, + heuristic_first_max_tier="MEDIUM", + heuristic_first_max_context_tokens=10, + ) + messages = [ + {"role": "system", "content": "x" * 80}, + {"role": "user", "content": "why did that fail?"}, + ] + + outcome = await router.aclassify("why did that fail?", messages=messages) + + mock_router_instance.acompletion.assert_awaited_once() + assert outcome.cause != "heuristic_first_short_circuit" + assert outcome.cause == "llm_classifier" + + @pytest.mark.asyncio + @pytest.mark.parametrize("context_limit", [None, 100]) + async def test_short_context_keeps_cheap_short_turn(self, mock_router_instance, context_limit): + mock_router_instance.acompletion = AsyncMock() + router = _heuristic_first_router( + mock_router_instance, + heuristic_first_max_tier="MEDIUM", + heuristic_first_max_context_tokens=context_limit, + ) + messages = [{"role": "user", "content": "why did that fail?"}] + + outcome = await router.aclassify("why did that fail?", messages=messages) + + mock_router_instance.acompletion.assert_not_called() + assert outcome.cause == "heuristic_first_short_circuit" + + @pytest.mark.parametrize( + "messages, expected", + [ + (None, 0), + ( + [ + {"role": "system", "content": "abcd"}, + {"role": "user", "content": [{"type": "text", "text": "efghij"}, {"type": "image_url"}]}, + {"role": "assistant", "content": "klmnopqr"}, + ], + 4, + ), + ], + ) + def test_estimated_conversation_tokens_counts_text_parts(self, messages, expected): + assert _estimated_conversation_tokens(messages) == expected + @pytest.mark.asyncio async def test_no_signal_prompt_escalates_even_though_it_scores_simple(self, mock_router_instance): """The core guard. This prompt scores 0.0 and the mapping calls it SIMPLE, which is at the diff --git a/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx index 64b08fc9ed1..97f0965c49a 100644 --- a/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx @@ -55,6 +55,7 @@ const CLASSIFIER_TIMEOUT_ID = "classifier-timeout-ms"; const CLASSIFIER_CONTEXT_WINDOW_SIZE_ID = "classifier-context-window-size"; const CLASSIFIER_CONTEXT_BUDGET_CHARS_ID = "classifier-context-budget-chars"; const HYBRID_BOUNDARY_MARGIN_ID = "hybrid-boundary-margin"; +const HEURISTIC_FIRST_MAX_CONTEXT_TOKENS_ID = "heuristic-first-max-context-tokens"; const CUSTOM_PROMPT_WITH_HEURISTIC_FALLBACK = "This router classifies with your own prompt, so the tier comes from whatever rubric it states. The four tier " + @@ -268,6 +269,18 @@ const ClassificationMethodConfig: React.FC = ({ onChange({ ...value, heuristic_first_max_tier: tier }); }; + const handleHeuristicFirstMaxContextTokensChange = (raw: string) => { + setDraft({ id: HEURISTIC_FIRST_MAX_CONTEXT_TOKENS_ID, raw }); + if (raw.trim() === "") { + onChange({ ...value, heuristic_first_max_context_tokens: undefined }); + return; + } + const parsed: number = Number(raw); + if (Number.isFinite(parsed)) { + onChange({ ...value, heuristic_first_max_context_tokens: Math.max(1, Math.round(parsed)) }); + } + }; + const handleHybridBoundaryMarginChange = (raw: string) => { setDraft({ id: HYBRID_BOUNDARY_MARGIN_ID, raw }); const parsed = Number(raw); @@ -425,6 +438,21 @@ const ClassificationMethodConfig: React.FC = ({ ))} + + handleHeuristicFirstMaxContextTokensChange(event.target.value)} + onBlur={() => setDraft(null)} + className="w-full" + />

A request the scorer places at or below this tier routes there without a classifier call. Anything the scorer places higher, and anything it found no signal for at all, goes to the classifier instead diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx index f6b50ce20bc..ad4b952147a 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx @@ -388,6 +388,8 @@ export interface ComplexityRouterConfigValue { classification_examples?: string; /** Highest tier the scorer may decide alone under heuristic_first. Required by that type, rejected by the others. */ heuristic_first_max_tier?: string; + /** Conversation token estimate above which heuristic_first defers to the classifier. */ + heuristic_first_max_context_tokens?: number; /** How near a tier boundary a score may land before hybrid defers to the classifier. Required by that type, rejected by the others. */ hybrid_boundary_margin?: number; classification_mode?: ClassificationMode; diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx index 126d9ba2311..63eb027590e 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx @@ -401,6 +401,7 @@ const AddAutoRouterTab: React.FC = ({ classificationPrompt: complexityRouterConfig.classification_prompt, classificationExamples: complexityRouterConfig.classification_examples, heuristicFirstMaxTier: complexityRouterConfig.heuristic_first_max_tier, + heuristicFirstMaxContextTokens: complexityRouterConfig.heuristic_first_max_context_tokens, hybridBoundaryMargin: complexityRouterConfig.hybrid_boundary_margin, classificationMode: complexityRouterConfig.classification_mode, tierLabels: complexityRouterConfig.tier_labels, diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts index 2990878d086..a0b29bc36cc 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts @@ -970,6 +970,7 @@ describe("heuristic_first", () => { ...baseParams, classifierType: "heuristic_first", heuristicFirstMaxTier: "SIMPLE", + heuristicFirstMaxContextTokens: 8000, classifierLlmConfig: { model: "gpt-4o-mini", timeout_ms: 3000 }, classifierContextWindowSize: 5, classifierContextBudgetChars: 4000, @@ -980,6 +981,15 @@ describe("heuristic_first", () => { const config = buildComplexityRouterConfig(heuristicFirstParams); expect(config.classifier_type).toBe("heuristic_first"); expect(config.heuristic_first_max_tier).toBe("SIMPLE"); + expect(config.heuristic_first_max_context_tokens).toBe(8000); + }); + + it("omits heuristic_first_max_context_tokens when empty", () => { + const config = buildComplexityRouterConfig({ + ...heuristicFirstParams, + heuristicFirstMaxContextTokens: undefined, + }); + expect(config.heuristic_first_max_context_tokens).toBeUndefined(); }); it("keeps every classifier key the operator set, since heuristic_first still calls the classifier", () => { diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts index 8a377c17ad7..87dcbdbf2f1 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts @@ -141,6 +141,7 @@ export interface StoredComplexityRouterConfig { classification_prompt?: unknown; classification_examples?: unknown; heuristic_first_max_tier?: unknown; + heuristic_first_max_context_tokens?: unknown; hybrid_boundary_margin?: unknown; tier_labels?: unknown; classifier_type?: ClassifierType; @@ -259,6 +260,7 @@ export interface ComplexityRouterConfigPayload { classification_prompt?: string; classification_examples?: string; heuristic_first_max_tier?: string; + heuristic_first_max_context_tokens?: number; hybrid_boundary_margin?: number; classification_mode: ClassificationMode; session_affinity: boolean; @@ -506,6 +508,7 @@ const classifierWireFields = ( classifierLlmConfig, classifierFallback, heuristicFirstMaxTier, + heuristicFirstMaxContextTokens, hybridBoundaryMargin, classifierContextWindowSize, classifierContextBudgetChars, @@ -515,6 +518,7 @@ const classifierWireFields = ( | "classifierLlmConfig" | "classifierFallback" | "heuristicFirstMaxTier" + | "heuristicFirstMaxContextTokens" | "hybridBoundaryMargin" | "classifierContextWindowSize" | "classifierContextBudgetChars" @@ -532,6 +536,9 @@ const classifierWireFields = ( ...(supportsFallback && classifierFallback !== undefined && { classifier_fallback: classifierFallback }), ...(effectiveType === "heuristic_first" && heuristicFirstMaxTier?.trim() && { heuristic_first_max_tier: heuristicFirstMaxTier }), + ...(effectiveType === "heuristic_first" && + heuristicFirstMaxContextTokens !== undefined && + { heuristic_first_max_context_tokens: heuristicFirstMaxContextTokens }), ...(effectiveType === "hybrid" && hybridBoundaryMargin !== undefined && { hybrid_boundary_margin: hybridBoundaryMargin }), ...(usesLlmClassifier(effectiveType) && @@ -567,6 +574,7 @@ export const buildComplexityRouterConfig = ({ classificationPrompt, classificationExamples, heuristicFirstMaxTier, + heuristicFirstMaxContextTokens, hybridBoundaryMargin, classificationMode, sessionAffinity, @@ -620,6 +628,7 @@ export const buildComplexityRouterConfig = ({ classifierLlmConfig, classifierFallback, heuristicFirstMaxTier, + heuristicFirstMaxContextTokens, hybridBoundaryMargin, classifierContextWindowSize, classifierContextBudgetChars, diff --git a/ui/litellm-dashboard/src/components/add_model/classifier_type_transition.test.ts b/ui/litellm-dashboard/src/components/add_model/classifier_type_transition.test.ts index e3fe00d2bc8..1a4589df0d8 100644 --- a/ui/litellm-dashboard/src/components/add_model/classifier_type_transition.test.ts +++ b/ui/litellm-dashboard/src/components/add_model/classifier_type_transition.test.ts @@ -9,6 +9,7 @@ const standard: ComplexityRouterConfigValue = { classifier_context_budget_chars: 16000, classifier_context_include_assistant_turns: true, classifier_fallback: "default_model", + heuristic_first_max_context_tokens: 8000, tiers: { SIMPLE: ["efficient"], MEDIUM: ["middle"], COMPLEX: [], REASONING: ["capable"] }, }; @@ -22,6 +23,7 @@ describe("transitionClassifierType", () => { classifier_context_budget_chars: 16000, classifier_context_include_assistant_turns: true, classifier_fallback: "default_model", + ...(target === "heuristic_first" && { heuristic_first_max_context_tokens: 8000 }), }; expect(result).toMatchObject(expectedSettings); }); @@ -30,6 +32,7 @@ describe("transitionClassifierType", () => { const result = transitionClassifierType(standard, target); expect(result.classifier_llm_config).toEqual({ model: "judge", timeout_ms: 20000 }); expect(result.classifier_fallback).toBeUndefined(); + expect(result.heuristic_first_max_context_tokens).toBeUndefined(); if (target === "capability") { expect(result.capability_classifier_config?.base_threshold).toBeNaN(); } else { diff --git a/ui/litellm-dashboard/src/components/add_model/classifier_type_transition.ts b/ui/litellm-dashboard/src/components/add_model/classifier_type_transition.ts index df87e2854e3..9442d447428 100644 --- a/ui/litellm-dashboard/src/components/add_model/classifier_type_transition.ts +++ b/ui/litellm-dashboard/src/components/add_model/classifier_type_transition.ts @@ -42,6 +42,8 @@ export const transitionClassifierType = ( classifierType === "heuristic_first" ? value.heuristic_first_max_tier ?? DEFAULT_HEURISTIC_FIRST_MAX_TIER : undefined, + heuristic_first_max_context_tokens: + classifierType === "heuristic_first" ? value.heuristic_first_max_context_tokens : undefined, hybrid_boundary_margin: classifierType === "hybrid" ? value.hybrid_boundary_margin ?? DEFAULT_HYBRID_BOUNDARY_MARGIN : undefined, ...nonReasoningTierFields(classifierType, value), diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts index 4ae6efbb12d..f75eeffb654 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts +++ b/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts @@ -670,6 +670,7 @@ describe("managed keys survive an untouched open-and-save", () => { tier_labels: { SIMPLE: "Cheap" }, classifier_type: "heuristic_first", heuristic_first_max_tier: "SIMPLE", + heuristic_first_max_context_tokens: 8000, classifier_llm_config: { model: "gpt-4o-mini", timeout_ms: 3000, reasoning_effort: "low" }, classifier_context_window_size: 5, classifier_context_budget_chars: 4000, diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx index e25c7f07dd7..f4676de71ab 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx @@ -158,6 +158,10 @@ export const hydrateComplexityRouterConfig = ( typeof parsedConfig.heuristic_first_max_tier === "string" && parsedConfig.heuristic_first_max_tier.trim() !== "" ? parsedConfig.heuristic_first_max_tier : undefined, + heuristic_first_max_context_tokens: + typeof parsedConfig.heuristic_first_max_context_tokens === "number" + ? parsedConfig.heuristic_first_max_context_tokens + : undefined, hybrid_boundary_margin: typeof parsedConfig.hybrid_boundary_margin === "number" ? parsedConfig.hybrid_boundary_margin : undefined, classification_mode: @@ -226,6 +230,7 @@ export const MANAGED_COMPLEXITY_ROUTER_KEYS = new Set([ "classification_prompt", "classification_examples", "heuristic_first_max_tier", + "heuristic_first_max_context_tokens", "hybrid_boundary_margin", "classification_mode", "session_affinity", @@ -325,6 +330,7 @@ export const buildUpdatedComplexityRouterConfig = ( classificationPrompt: value.classification_prompt, classificationExamples: value.classification_examples, heuristicFirstMaxTier: value.heuristic_first_max_tier, + heuristicFirstMaxContextTokens: value.heuristic_first_max_context_tokens, hybridBoundaryMargin: value.hybrid_boundary_margin, classificationMode: value.classification_mode, tierLabels: value.tier_labels,