feat(proxy): dry-run a real request body on /auto_router/test_routing (#38590)

The endpoint built messages=[{"role": "user", "content": prompt}], so a dry run
could not carry prior turns, the caller's system prompt, or the tool definitions
a request advertises. A real agentic turn reduced to its last sentence classified
as trivial, which is why a config sweep reported savings for every configuration.

Accept messages, system and tools, and forward them to the same pre-routing hook
untranslated, with the raw-body snapshot built by the serving path's own owner,
refresh_proxy_server_request_body_snapshot. Loose types are deliberate: the hook
reads whatever dialect the surface produced, so validating against one surface's
schema would reject the others.

prompt stays as the single-ask shorthand, normalized into one user turn inside the
request model so the handler carries no mode branch.
This commit is contained in:
tin-berri 2026-08-27 18:35:10 -07:00 committed by GitHub
parent 10cd9259a3
commit 09b23742e7
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 353 additions and 43 deletions

View file

@ -11364,6 +11364,18 @@
"description": "Path to a JSON file containing ad-hoc recognizers for Presidio",
"title": "Presidio Ad Hoc Recognizers"
},
"presidio_analyze_chunk_size_bytes": {
"anyOf": [
{
"type": "integer"
},
{
"type": "null"
}
],
"description": "Maximum UTF-8 bytes of text sent in a single Presidio /analyze call. Longer texts are split into overlapping chunks of at most this size and the merged results are remapped onto the original text. Defaults to 500000; set it below your analyzer deployment's request body limit, leaving headroom for the rest of the analyze payload.",
"title": "Presidio Analyze Chunk Size Bytes"
},
"presidio_analyzer_api_base": {
"anyOf": [
{

View file

@ -1,7 +1,7 @@
"""
AUTO ROUTER MANAGEMENT ENDPOINTS
POST /auto_router/test_routing - Route one prompt through an unsaved complexity-router config
POST /auto_router/test_routing - Route one request through an unsaved complexity-router config
POST /auto_router/validate_complexity_router_config - Dry-run the complexity-router write gate without saving
"""
@ -32,7 +32,10 @@ from litellm.proxy.auth.auth_checks import (
)
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.db.autorouter_session_rollup import AUTOROUTER_BENCHMARKS_SQL
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
from litellm.proxy.litellm_pre_call_utils import (
LiteLLMProxyRequestSetup,
refresh_proxy_server_request_body_snapshot,
)
from litellm.repositories.base_repository import SupportsModelDump
from litellm.repositories.team_repository import TeamRepository
from litellm.router_strategy.complexity_router import ComplexityRouter
@ -285,19 +288,30 @@ async def preview_auto_router_routing(
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
) -> AutoRouterRoutingTestResponse:
"""
Route a single prompt through a complexity-router config and report where it landed.
Route a single request through a complexity-router config and report where it landed.
Answers "which model would this prompt get?" for a config that only exists in a form,
so an auto router can be checked before it is created. The prompt is classified by the
same pre-routing hook a live request runs, then dropped: nothing is sent to the model it
routed to, and no auto router is created. A heuristic config therefore spends nothing, while
an `llm` classifier or semantic keyword matching bills its classifier/embedding call to the
calling key, like Test Connection does.
Answers "which model would this request get?" for a config that only exists in a form,
so an auto router can be checked before it is created. The request is classified by the
same pre-routing hook a live request runs, over the same messages, system prompt and tool
definitions, then dropped: nothing is sent to the model it routed to, and no auto router is
created. A heuristic config therefore spends nothing, while an `llm` classifier or semantic
keyword matching bills its classifier/embedding call to the calling key, like Test Connection
does.
Send `messages` to classify a real turn, with `system` and `tools` beside it when the surface
carries them top level, as Anthropic /v1/messages does. `prompt` is the single-ask shorthand and
routes as one user turn with nothing around it.
**Example Request:**
```json
{
"prompt": "think step by step about how to shard this table",
"messages": [
{"role": "system", "content": "You are a database migration assistant"},
{"role": "user", "content": "the index is not unique"},
{"role": "assistant", "content": "Then two workers can both insert. Add a unique index"},
{"role": "user", "content": "ok do it"}
],
"tools": [{"type": "function", "function": {"name": "Bash", "description": "Run a command"}}],
"complexity_router_config": {
"tiers": {"SIMPLE": ["gpt-4o-mini"], "REASONING": ["o3"]},
"classifier_type": "heuristic"
@ -340,18 +354,21 @@ async def preview_auto_router_routing(
)
request_kwargs: Final = LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata(
data={"metadata": {}}, # mutable-ok: the request-metadata helper takes and returns request kwargs as a dict
data={ # mutable-ok: the request-metadata helper takes and returns request kwargs as a dict
**data.wire_body(),
"metadata": {}, # mutable-ok: the request-metadata helper writes the auth fields into this dict
"proxy_server_request": {"body": None}, # mutable-ok: the snapshot owner fills body in place
},
user_api_key_dict=user_api_key_dict,
_metadata_variable_name="metadata",
)
refresh_proxy_server_request_body_snapshot(request_kwargs)
try:
hook_response: Final = await complexity_router.async_pre_routing_hook(
model=data.router_name,
request_kwargs=request_kwargs,
messages=[ # mutable-ok: the routing hook's signature takes a list of message dicts
{"role": "user", "content": data.prompt}, # mutable-ok: a message is dict-shaped
],
messages=request_kwargs["messages"],
)
except Exception as e: # noqa: BLE001 -- surfaces any classifier/plugin failure to the caller as a 400 instead of a 500, since the config under test is caller input
verbose_proxy_logger.exception("Auto router routing test failed. Due to error - %s", e)

View file

@ -2,8 +2,9 @@
Types for auto-router management endpoints
"""
from collections.abc import Mapping
from collections.abc import Mapping, Sequence
from datetime import datetime, timezone
from types import MappingProxyType
from typing import Final, Literal, TypeAlias
from pydantic import BaseModel, Field, computed_field, field_validator, model_validator
@ -44,9 +45,30 @@ class ComplexityRouterConfigValidationResponse(BaseModel):
class AutoRouterRoutingTestRequest(BaseModel):
"""A single prompt to classify against a complexity-router config that need not be saved yet."""
"""A single request to classify against a complexity-router config that need not be saved yet.
prompt: str = Field(description="The prompt to route, as an end user would send it")
Carries the same fields the serving path carries, so a dry run classifies what a real turn
would classify. `messages`, `system` and `tools` are forwarded to the routing hook untranslated,
which is why they are typed loosely: the hook reads whatever dialect the surface produced, and
validating them against one surface's schema would reject the others.
"""
prompt: str | None = Field(
default=None,
description="A single ask to route, as an end user would send it. Mutually exclusive with messages",
)
messages: Sequence[Mapping[str, object]] | None = Field(
default=None,
description="The full message list to route, exactly as the serving path would receive it. Mutually exclusive with prompt",
)
system: str | Sequence[Mapping[str, object]] | None = Field(
default=None,
description="The top-level system prompt an Anthropic /v1/messages body carries beside its messages",
)
tools: Sequence[Mapping[str, object]] | None = Field(
default=None,
description="The tool definitions the request advertises, which decide whether the plan-mode floor applies",
)
complexity_router_config: RequestComplexityRouterConfig = Field(
description="The complexity router config to route against, in the shape /model/new accepts",
)
@ -63,13 +85,60 @@ class AutoRouterRoutingTestRequest(BaseModel):
description="Team the router is being created for. Required for a team admin, who may only test their own team's routers",
)
@field_validator("prompt")
@field_validator("messages")
@classmethod
def _require_non_blank_prompt(cls, value: str) -> str:
if not value.strip():
raise ValueError("prompt must not be blank")
def _reject_messages_no_surface_accepts(
cls, value: Sequence[Mapping[str, object]] | None
) -> Sequence[Mapping[str, object]] | None:
"""Reject what every supported surface rejects, and nothing beyond it.
A real request carrying a message with no string role, or with content that is neither text
nor a block list, is a 400 on the serving path, so answering it here with a routed tier
would promise a decision the request never gets. Only the two keys the dialects agree on
are constrained: anything else in a message stays untranslated and unread.
"""
if value is None:
return value
for index, message in enumerate(value):
if not isinstance(role := message.get("role"), str) or not role.strip():
raise ValueError(f"messages[{index}] needs a non-empty string role")
if (content := message.get("content")) is not None and not isinstance(content, str | list):
raise ValueError(f"messages[{index}] content must be a string, a list of blocks, or null")
return value
@model_validator(mode="after")
def _resolve_request_carrier(self) -> "AutoRouterRoutingTestRequest":
if self.prompt is not None and not self.prompt.strip():
raise ValueError("prompt must not be blank")
if self.messages is not None and not self.messages:
raise ValueError("messages must not be empty")
if (self.prompt is None) == (self.messages is None):
raise ValueError("provide exactly one of prompt or messages")
if self.messages is not None:
return self
return self.model_copy(
update={ # mutable-ok: model_copy types update as a plain dict
"messages": [ # mutable-ok: the routing hook's signature takes a list of message dicts
{"role": "user", "content": self.prompt} # mutable-ok: a message is dict-shaped
]
}
)
def wire_body(self) -> Mapping[str, object]:
"""The request kwargs a serving-path request would carry for this body.
Every value is handed out by identity rather than copied, so the messages the routing hook
classifies and the messages its raw-body plan-mode scan reads are one value, as they are on
the serving path.
"""
return MappingProxyType(
{ # mutable-ok: MappingProxyType needs a dict to wrap
key: value
for key, value in (("messages", self.messages), ("system", self.system), ("tools", self.tools))
if value is not None
}
)
class AutoRouterRoutingTestResponse(BaseModel):
"""Where one prompt would have been routed, and why."""

View file

@ -47,34 +47,87 @@ TIERS = {
}
ROUTER_MODEL_LIST = [
{"model_name": name, "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "fake-key"}}
for name in ("cheap-model", "mid-model", "strong-model", "reasoning-model")
]
def _router() -> Router:
return Router(
model_list=[
{"model_name": name, "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "fake-key"}}
for name in ("cheap-model", "mid-model", "strong-model", "reasoning-model")
]
)
return Router(model_list=ROUTER_MODEL_LIST)
def _request(prompt: str, **config_overrides: object) -> AutoRouterRoutingTestRequest:
class RecordingRouter(Router):
"""A real router that records the classifier calls the endpoint makes instead of sending them.
Injected at the same `proxy_server.llm_router` boundary the endpoint reads, so model resolution
and the key's model-access checks still run against a genuine Router.
"""
def __init__(self, classified_tier: str) -> None:
super().__init__(model_list=ROUTER_MODEL_LIST)
self.classified_tier = classified_tier
self.recorded_calls: list[dict] = []
async def acompletion(self, model, messages, stream=False, **kwargs):
self.recorded_calls.append({"model": model, "messages": messages, **kwargs})
return ModelResponse(
choices=[Choices(message=Message(content=f'{{"tier": "{self.classified_tier}"}}'))],
model=model,
)
def _request_from(body: Mapping[str, object], **config_overrides: object) -> AutoRouterRoutingTestRequest:
return AutoRouterRoutingTestRequest.model_validate(
{
"prompt": prompt,
**body,
"complexity_router_config": {"tiers": TIERS, "classifier_type": "heuristic", **config_overrides},
}
)
async def _route(prompt: str, monkeypatch: pytest.MonkeyPatch, **config_overrides: object):
def _request(prompt: str, **config_overrides: object) -> AutoRouterRoutingTestRequest:
return _request_from({"prompt": prompt}, **config_overrides)
async def _route_body(body: Mapping[str, object], monkeypatch: pytest.MonkeyPatch, **config_overrides: object):
import litellm.proxy.proxy_server as proxy_server
monkeypatch.setattr(proxy_server, "llm_router", _router())
return await preview_auto_router_routing(
data=_request(prompt, **config_overrides),
data=_request_from(body, **config_overrides),
user_api_key_dict=ADMIN,
)
async def _route(prompt: str, monkeypatch: pytest.MonkeyPatch, **config_overrides: object):
return await _route_body({"prompt": prompt}, monkeypatch, **config_overrides)
AGENTIC_MESSAGES = [
{"role": "system", "content": "You are a database migration assistant for a payments ledger"},
{"role": "user", "content": "duplicate ledger postings since the celery upgrade, same event_id twice"},
{"role": "assistant", "content": "The idempotency index is not unique, so two workers both insert"},
{"role": "user", "content": "ok do it"},
]
PLAN_MODE_TOOLS = [{"type": "function", "function": {"name": "exit_plan_mode", "description": "Leave plan mode"}}]
async def _classifier_user_payload(body: Mapping[str, object], monkeypatch: pytest.MonkeyPatch) -> str:
"""The variable half of the classifier call this body produces."""
from litellm.proxy import proxy_server
router = RecordingRouter("SIMPLE")
monkeypatch.setattr(proxy_server, "llm_router", router)
await preview_auto_router_routing(
data=_request_from(body, classifier_type="llm", classifier_llm_config={"model": "classifier-model"}),
user_api_key_dict=ADMIN,
)
return router.recorded_calls[0]["messages"][1]["content"]
@pytest.mark.asyncio
async def test_simple_prompt_routes_to_the_simple_tier(monkeypatch: pytest.MonkeyPatch):
response = await _route("what is 2+2", monkeypatch)
@ -160,6 +213,123 @@ async def test_llm_classifier_call_is_billed_to_the_calling_key(monkeypatch: pyt
assert calls[0]["metadata"]["user_api_key_user_id"] == ADMIN.user_id
@pytest.mark.asyncio
async def test_a_full_turn_is_classified_on_its_system_prompt_and_prior_turns(monkeypatch: pytest.MonkeyPatch):
"""A dry run over `messages` must produce the classifier call the serving path produces.
The `prompt` shorthand for the same final ask is the negative class: it carries neither the
caller's system prompt nor the conversation it continues, which is why a real agentic turn
reduced to its last sentence classifies as trivial.
"""
full_turn = await _classifier_user_payload({"messages": AGENTIC_MESSAGES}, monkeypatch)
last_sentence_only = await _classifier_user_payload({"prompt": "ok do it"}, monkeypatch)
assert "You are a database migration assistant for a payments ledger" in full_turn
assert "duplicate ledger postings since the celery upgrade" in full_turn
assert full_turn.endswith("Classify this message:\nok do it")
assert "database migration assistant" not in last_sentence_only
assert "duplicate ledger postings" not in last_sentence_only
assert last_sentence_only.endswith("Classify this message:\nok do it")
@pytest.mark.asyncio
async def test_a_top_level_system_prompt_is_not_classified_as_the_ask(monkeypatch: pytest.MonkeyPatch):
"""An Anthropic body carries `system` beside its messages, and the serving path leaves it
there: it reaches the raw-body scan, never the ask the classifier is asked to rate."""
payload = await _classifier_user_payload(
{"messages": [{"role": "user", "content": "ok do it"}], "system": "You migrate payment ledgers"},
monkeypatch,
)
assert payload.endswith("Classify this message:\nok do it")
assert "You migrate payment ledgers" not in payload
@pytest.mark.parametrize(
"body, expected_model",
[
pytest.param({"prompt": "what is 2+2", "tools": PLAN_MODE_TOOLS}, "strong-model", id="tools-carry-it"),
pytest.param(
{"prompt": "what is 2+2", "system": 'You are currently running in "Plan" mode.'},
"strong-model",
id="system-carries-it",
),
pytest.param({"prompt": "what is 2+2"}, "cheap-model", id="neither-carries-it"),
pytest.param(
{"prompt": "what is 2+2", "tools": [{"type": "function", "function": {"name": "Bash"}}]},
"cheap-model",
id="unrelated-tool",
),
],
)
@pytest.mark.asyncio
async def test_the_plan_mode_floor_sees_the_tools_and_system_the_request_carries(
monkeypatch: pytest.MonkeyPatch, body: dict, expected_model: str
):
response = await _route_body(body, monkeypatch, plan_mode_min_tier="COMPLEX")
assert response.routed_model == expected_model
def test_the_wire_body_hands_out_the_same_messages_the_hook_classifies():
"""The routing hook reads messages twice, as its own argument and through the raw-body scan.
One value, so the two can never disagree."""
request = _request_from({"messages": AGENTIC_MESSAGES})
assert request.wire_body()["messages"] is request.messages
def test_a_prompt_is_carried_as_one_user_turn():
assert _request_from({"prompt": "what is 2+2"}).messages == [{"role": "user", "content": "what is 2+2"}]
@pytest.mark.parametrize(
"message",
[
pytest.param({"content": "hi"}, id="no-role"),
pytest.param({"role": 123, "content": "hi"}, id="role-not-a-string"),
pytest.param({"role": " ", "content": "hi"}, id="blank-role"),
pytest.param({"role": "user", "content": {"weird": 1}}, id="content-neither-text-nor-blocks"),
],
)
def test_a_message_no_surface_would_accept_is_rejected(message: dict):
"""The serving path 400s on each of these, so a routed tier here would be a promise it breaks."""
with pytest.raises(ValidationError):
_request_from({"messages": [message]})
@pytest.mark.parametrize(
"message",
[
pytest.param({"role": "user", "content": "ok do it"}, id="text-content"),
pytest.param({"role": "user", "content": [{"type": "text", "text": "ok"}]}, id="block-content"),
pytest.param(
{"role": "assistant", "content": None, "tool_calls": [{"id": "c1", "type": "function"}]},
id="null-content-with-tool-calls",
),
pytest.param({"role": "user", "content": "hi", "cache_control": {"type": "ephemeral"}}, id="unknown-key"),
],
)
def test_a_message_a_serving_surface_accepts_is_kept(message: dict):
"""The serving path returns 200 for each of these, and none of their keys are translated."""
assert _request_from({"messages": [message]}).messages == [message]
@pytest.mark.parametrize(
"body",
[
pytest.param({}, id="neither"),
pytest.param({"prompt": "hi", "messages": [{"role": "user", "content": "hi"}]}, id="both"),
pytest.param({"prompt": " "}, id="blank-prompt"),
pytest.param({"messages": []}, id="empty-messages"),
],
)
def test_a_request_must_carry_exactly_one_usable_conversation(body: dict):
with pytest.raises(ValidationError):
_request_from(body)
@pytest.mark.parametrize(
"config_overrides",
[

View file

@ -1247,19 +1247,30 @@ export interface paths {
put?: never;
/**
* Preview Auto Router Routing
* @description Route a single prompt through a complexity-router config and report where it landed.
* @description Route a single request through a complexity-router config and report where it landed.
*
* Answers "which model would this prompt get?" for a config that only exists in a form,
* so an auto router can be checked before it is created. The prompt is classified by the
* same pre-routing hook a live request runs, then dropped: nothing is sent to the model it
* routed to, and no auto router is created. A heuristic config therefore spends nothing, while
* an `llm` classifier or semantic keyword matching bills its classifier/embedding call to the
* calling key, like Test Connection does.
* Answers "which model would this request get?" for a config that only exists in a form,
* so an auto router can be checked before it is created. The request is classified by the
* same pre-routing hook a live request runs, over the same messages, system prompt and tool
* definitions, then dropped: nothing is sent to the model it routed to, and no auto router is
* created. A heuristic config therefore spends nothing, while an `llm` classifier or semantic
* keyword matching bills its classifier/embedding call to the calling key, like Test Connection
* does.
*
* Send `messages` to classify a real turn, with `system` and `tools` beside it when the surface
* carries them top level, as Anthropic /v1/messages does. `prompt` is the single-ask shorthand and
* routes as one user turn with nothing around it.
*
* **Example Request:**
* ```json
* {
* "prompt": "think step by step about how to shard this table",
* "messages": [
* {"role": "system", "content": "You are a database migration assistant"},
* {"role": "user", "content": "the index is not unique"},
* {"role": "assistant", "content": "Then two workers can both insert. Add a unique index"},
* {"role": "user", "content": "ok do it"}
* ],
* "tools": [{"type": "function", "function": {"name": "Bash", "description": "Run a command"}}],
* "complexity_router_config": {
* "tiers": {"SIMPLE": ["gpt-4o-mini"], "REASONING": ["o3"]},
* "classifier_type": "heuristic"
@ -22850,7 +22861,12 @@ export interface components {
};
/**
* AutoRouterRoutingTestRequest
* @description A single prompt to classify against a complexity-router config that need not be saved yet.
* @description A single request to classify against a complexity-router config that need not be saved yet.
*
* Carries the same fields the serving path carries, so a dry run classifies what a real turn
* would classify. `messages`, `system` and `tools` are forwarded to the routing hook untranslated,
* which is why they are typed loosely: the hook reads whatever dialect the surface produced, and
* validating them against one surface's schema would reject the others.
*/
AutoRouterRoutingTestRequest: {
/** @description The complexity router config to route against, in the shape /model/new accepts */
@ -22861,21 +22877,42 @@ export interface components {
*/
default_model?: string | null;
/**
* Prompt
* @description The prompt to route, as an end user would send it
* Messages
* @description The full message list to route, exactly as the serving path would receive it. Mutually exclusive with prompt
*/
prompt: string;
messages?: {
[key: string]: unknown;
}[] | null;
/**
* Prompt
* @description A single ask to route, as an end user would send it. Mutually exclusive with messages
*/
prompt?: string | null;
/**
* Router Name
* @description Name reported as the router in the routing decision. Display only
* @default auto_router_routing_test
*/
router_name: string;
/**
* System
* @description The top-level system prompt an Anthropic /v1/messages body carries beside its messages
*/
system?: string | {
[key: string]: unknown;
}[] | null;
/**
* Team Id
* @description Team the router is being created for. Required for a team admin, who may only test their own team's routers
*/
team_id?: string | null;
/**
* Tools
* @description The tool definitions the request advertises, which decide whether the plan-mode floor applies
*/
tools?: {
[key: string]: unknown;
}[] | null;
};
/**
* AutoRouterRoutingTestResponse
@ -29937,6 +29974,11 @@ export interface components {
* @description Path to a JSON file containing ad-hoc recognizers for Presidio
*/
presidio_ad_hoc_recognizers?: string | null;
/**
* Presidio Analyze Chunk Size Bytes
* @description Maximum UTF-8 bytes of text sent in a single Presidio /analyze call. Longer texts are split into overlapping chunks of at most this size and the merged results are remapped onto the original text. Defaults to 500000; set it below your analyzer deployment's request body limit, leaving headroom for the rest of the analyze payload.
*/
presidio_analyze_chunk_size_bytes?: number | null;
/**
* Presidio Analyzer Api Base
* @description Base URL for the Presidio analyzer API