scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-29 06:33:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 200

feat(proxy): accept request priority from the x-litellm-priority header (#47618)

Browse source 
* feat(proxy): accept request priority from the x-litellm-priority header

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(proxy): only apply x-litellm-priority on scheduler-queued routes

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(proxy): apply x-litellm-priority only where the router schedules the request

The header now maps onto priority only for chat and text completions that a
router serves, so the Cursor route's responses call and proxies started without
a config never forward it to the provider. /queue/chat/completions now honors
the header too, with the header winning over a body priority.

* test(proxy): use a real ProxyLogging in the priority header routing test

* fix(proxy): ignore a malformed priority header under drop_params instead of dropping the body priority

* fix(proxy): treat an empty priority header as unset and pass a valid one on as an integer

---------

Co-authored-by: Mubashir Osmani <mubashir@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
mubashir1osmani 2026-10-28 18:58:52 -04:00 • committed by GitHub
parent 6063187fea
commit 8bbce867d0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 257 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
litellm/proxy/common_request_processing.py
View file

@ -266,6 +266,8 @@ from litellm.proxy.anthropic_endpoints.streaming_model_restamp import (
AnthropicStreamModelRestamper,
)
from litellm.proxy.litellm_pre_call_utils import (
ROUTER_SCHEDULED_ROUTE_TYPES,
SCHEDULER_PRIORITY_HEADER,
LiteLLMProxyRequestSetup,
add_litellm_data_to_request,
get_metadata_variable_name,
@ -2478,14 +2480,20 @@ class ProxyBaseLLMRequestProcessing:
self.data["model"] = alias_target
aliased_model: Final = self.data.get("model")
if llm_router is not None and route_type in ROUTER_SCHEDULED_ROUTE_TYPES and isinstance(aliased_model, str):
self.data.update(
llm_router.priority_header_update(
model=aliased_model, kwargs=self.data, header=request.headers.get(SCHEDULER_PRIORITY_HEADER)
)
)
litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id"))
self.data["litellm_call_id"] = litellm_call_id
bind_telemetry_call_id(litellm_call_id)
DDSpanTagger.tag_call_id(self.data.get("litellm_call_id"))
DDSpanTagger.tag_call_id(litellm_call_id)
DDSpanTagger.tag_request(
user_api_key_dict=user_api_key_dict,
requested_model=self.data.get("model"),
requested_model=aliased_model,
)
### AUTO STREAM USAGE TRACKING ###

3
litellm/proxy/litellm_pre_call_utils.py
View file

@ -118,6 +118,9 @@ _SESSION_ID_VALUE_RE: Final = re.compile(r"^[a-zA-Z0-9_\-]{8,}$")
_SHA256_HEX_RE: Final = re.compile(r"^[0-9a-f]{64}$")
SCHEDULER_PRIORITY_HEADER: Final = "x-litellm-priority"
ROUTER_SCHEDULED_ROUTE_TYPES: Final = frozenset({"acompletion", "atext_completion"})
# W3C Trace Context traceparent header: https://www.w3.org/TR/trace-context/
# e.g. "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"
_TRACEPARENT_RE: Final = re.compile(r"^[0-9a-f]{2}-([0-9a-f]{32})-[0-9a-f]{16}-[0-9a-f]{2}$", re.IGNORECASE)

11
litellm/proxy/proxy_server.py
View file

@ -666,7 +666,11 @@ from litellm.proxy.list_api.common import (
problem_response,
request_validation_problem,
)
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup, add_litellm_data_to_request
from litellm.proxy.litellm_pre_call_utils import (
SCHEDULER_PRIORITY_HEADER,
LiteLLMProxyRequestSetup,
add_litellm_data_to_request,
)
from litellm.proxy.logging_endpoints.callback_logs_endpoints import (
rust_control_plane_router,
)
@ -18481,6 +18485,11 @@ async def async_queue_request(
raise HTTPException(status_code=500, detail={"error": CommonProxyErrors.no_llm_router.value})
router: Final = llm_router
data.update(
router.priority_header_update(
model=data["model"], kwargs=data, header=request.headers.get(SCHEDULER_PRIORITY_HEADER)
)
)
queued_priority, _ = router.split_request_priority(model=data["model"], kwargs=data)
if "stream" in data and data["stream"] is True: # use generate_responses to stream responses

15
litellm/router.py
View file

@ -359,6 +359,7 @@ from litellm.router_utils.request_priority import (
InvalidPriority,
parse_default_priority,
priority_error,
priority_header_fields,
request_drops_params,
resolve_request_priority,
unqueueable_priority_message,
@ -5411,14 +5412,22 @@ class Router:
request_priority: Final = resolve_request_priority(
requested=kwargs.get("priority"),
default_priority=self.default_priority,
drops_params=lambda: request_drops_params(
kwargs, self.default_litellm_params, self._request_deployment_params(model=model, kwargs=kwargs)
),
drops_params=lambda: self._drops_params(model=model, kwargs=kwargs),
)
if isinstance(request_priority, InvalidPriority):
raise priority_error(model=model, message=request_priority.message)
return request_priority, {key: value for key, value in kwargs.items() if key != "priority"}
def priority_header_update(
self, model: str, kwargs: Mapping[str, object], header: str | None
) -> Mapping[str, int | str]:
return priority_header_fields(header, drops_params=lambda: self._drops_params(model=model, kwargs=kwargs))
def _drops_params(self, model: str, kwargs: Mapping[str, object]) -> bool:
return request_drops_params(
kwargs, self.default_litellm_params, self._request_deployment_params(model=model, kwargs=kwargs)
)
async def _schedule_factory(
self,
model: str,

9
litellm/router_utils/request_priority.py
View file

@ -47,6 +47,15 @@ def resolve_request_priority(
return priority
def priority_header_fields(header: str | None, drops_params: Callable[[], bool]) -> Mapping[str, int | str]:
if not header:
return {}
priority: Final = _as_priority(header)
if isinstance(priority, InvalidPriority):
return {} if drops_params() else {"priority": header}
return {"priority": priority}
def parse_default_priority(value: object) -> int | None | InvalidPriority:
if value is None:
return None

2
tests/code_coverage_tests/router_code_coverage.py
View file

@ -175,6 +175,8 @@ ignored_function_names = [
"_responses_continuation", # Tested through the Responses mid-stream fallback tests in tests/unit/llms/base_llm/responses/test_transformation.py
"_responses_continuation_ceilings", # Tested through the Responses mid-stream fallback tests in tests/unit/llms/base_llm/responses/test_transformation.py
"_team_deployments_named", # Tested through the team lookups in the concurrent delete tests in test_router.py
"priority_header_update", # Tested through the priority header tests in test_common_request_processing.py
"_drops_params", # Read by split_request_priority and priority_header_update, tested through both
"_responses_event_announces_background_run", # Tested through the background Responses stream tests in test_router.py
"_aresponses_background_stream_can_release", # Tested through the background Responses stream tests in test_router.py
"_raise_released_stream_error", # Tested through the background Responses stream content-policy tests in test_router.py

48
tests/unit/proxy/proxy_server/test_routes_queue.py
View file

@ -191,6 +191,54 @@ def test_queue_chat_completions_rejects_a_request_it_cannot_queue_as_an_invalid_
assert add_request.await_count == 0
@pytest.mark.parametrize(
("default_priority", "priority_header", "priority_fields", "expected_priority"),
[
(5, "0", {"priority": 5}, 0),
(5, "0", {}, 0),
(None, "high", {"priority": 1, "drop_params": True}, 1),
(None, "", {"priority": 1}, 1),
],
ids=("header-beats-body", "header-only", "drop-params-keeps-body", "empty-header-keeps-body"),
)
def test_queue_chat_completions_queues_at_the_priority_header(
client,
auth_as,
monkeypatch: pytest.MonkeyPatch,
default_priority: int | None,
priority_header: str,
priority_fields: dict[str, object],
expected_priority: int,
) -> None:
router: Final = _queue_router(monkeypatch, default_priority=default_priority)
payload: Final = {"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}], **priority_fields}
with patch.object(router.scheduler, "add_request", wraps=router.scheduler.add_request) as add_request:
with auth_as():
response: Final = client.post(
"/queue/chat/completions", json=payload, headers={"x-litellm-priority": priority_header}
)
assert response.status_code == 200, response.text
assert [call.kwargs["request"].priority for call in add_request.await_args_list] == [expected_priority]
assert response.headers.get("x-litellm-priority") == str(expected_priority)
def test_queue_chat_completions_rejects_a_malformed_priority_header_like_a_malformed_body_value(
client, auth_as, monkeypatch: pytest.MonkeyPatch
) -> None:
router: Final = _queue_router(monkeypatch, default_priority=5)
payload: Final = {"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}], "priority": 1}
with patch.object(router.scheduler, "add_request", wraps=router.scheduler.add_request) as add_request:
with auth_as():
response: Final = client.post(
"/queue/chat/completions", json=payload, headers={"x-litellm-priority": "high"}
)
assert response.status_code == 400, response.text
error: Final = response.json()["error"]
assert (error["type"], error["param"]) == ("invalid_request_error", "priority")
assert "priority must be an integer" in error["message"]
assert add_request.await_count == 0
def test_streamed_queue_responses_report_the_priority_the_router_queued_at(
client, auth_as, monkeypatch: pytest.MonkeyPatch
) -> None:

163
tests/unit/proxy/test_common_request_processing.py
View file

@ -35,6 +35,7 @@ from litellm.constants import (
)
from litellm.integrations.custom_logger import CustomLogger
from litellm.integrations.opentelemetry import UserAPIKeyAuth
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.common_request_processing import (
ProxyBaseLLMRequestProcessing,
ProxyConfig,
@ -513,6 +514,168 @@ class TestProxyBaseLLMRequestProcessing:
assert exc_info.value.param == "model"
add_litellm_data_to_request.assert_not_awaited()
@pytest.mark.asyncio
@pytest.mark.parametrize(
("path", "route_type", "has_router", "priority_header", "body_priority", "expected_priority"),
[
("/v1/chat/completions", "acompletion", True, "0", 5, 0),
("/v1/completions", "atext_completion", True, "0", 5, 0),
("/v1/chat/completions", "acompletion", True, "1", None, 1),
("/v1/completions", "atext_completion", True, "1", None, 1),
("/v1/chat/completions", "acompletion", True, None, 5, 5),
("/v1/chat/completions", "acompletion", True, "", 5, 5),
("/cursor/chat/completions", "aresponses", True, "0", None, None),
("/v1/responses", "aresponses", True, "0", 5, 5),
("/v1/embeddings", "aembedding", True, "0", None, None),
("/v1/chat/completions", "acompletion", False, "0", None, None),
("/v1/chat/completions", "acompletion", False, "0", 5, 5),
],
ids=(
"chat-header-beats-body",
"text-completion-header-beats-body",
"chat-header-only",
"text-completion-header-only",
"chat-body-only",
"chat-empty-header-keeps-body",
"cursor-responses-route-ignores-header",
"responses-keeps-body",
"embeddings-ignores-header",
"routerless-ignores-header",
"routerless-keeps-body",
),
)
async def test_common_processing_pre_call_logic_applies_priority_header_only_where_the_router_schedules(
self,
path: str,
route_type: Literal["acompletion", "atext_completion", "aresponses", "aembedding"],
has_router: bool,
priority_header: str | None,
body_priority: int | None,
expected_priority: str | int | None,
) -> None:
from litellm.proxy.proxy_server import ProxyConfig as RealProxyConfig
header_entries: Final[tuple[tuple[bytes, bytes], ...]] = (
() if priority_header is None else ((b"x-litellm-priority", priority_header.encode()),)
)
request: Final = Request(
{
"type": "http",
"method": "POST",
"path": path,
"scheme": "http",
"server": ("localhost", 4000),
"client": ("127.0.0.1", 12345),
"query_string": b"",
"headers": ((b"content-type", b"application/json"), *header_entries),
}
)
router: Final = (
Router(
model_list=[
{
"model_name": "qa-model",
"litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-test", "mock_response": "hi"},
}
]
)
if has_router
else None
)
processing_obj: Final = ProxyBaseLLMRequestProcessing(
data={
"model": "qa-model",
"messages": [{"role": "user", "content": "hi"}],
**({"priority": body_priority} if body_priority is not None else {}),
}
)
returned_data, _ = await processing_obj.common_processing_pre_call_logic(
request=request,
general_settings={},
user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test"),
proxy_logging_obj=ProxyLogging(user_api_key_cache=UserApiKeyCache()),
proxy_config=RealProxyConfig(),
route_type=route_type,
llm_router=router,
)
assert ("priority" in returned_data, returned_data.get("priority")) == (
expected_priority is not None,
expected_priority,
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
("requested_model", "deployment_drops_params", "body_priority", "expected_priority"),
[
("qa-model", True, 5, 5),
("qa-model", True, None, None),
("key-alias", True, 5, 5),
("qa-model", False, 5, "high"),
],
ids=(
"drop-params-keeps-body",
"drop-params-without-body-leaves-default",
"drop-params-read-from-the-aliased-deployment",
"no-drop-params-passes-header-on-to-400",
),
)
async def test_common_processing_pre_call_logic_drops_only_a_malformed_priority_header_under_drop_params(
self,
requested_model: str,
deployment_drops_params: bool,
body_priority: int | None,
expected_priority: str | int | None,
) -> None:
from litellm.proxy.proxy_server import ProxyConfig as RealProxyConfig
request: Final = Request(
{
"type": "http",
"method": "POST",
"path": "/v1/chat/completions",
"scheme": "http",
"server": ("localhost", 4000),
"client": ("127.0.0.1", 12345),
"query_string": b"",
"headers": ((b"content-type", b"application/json"), (b"x-litellm-priority", b"high")),
}
)
router: Final = Router(
model_list=[
{
"model_name": "qa-model",
"litellm_params": {
"model": "openai/gpt-4o",
"api_key": "sk-test",
"mock_response": "hi",
"drop_params": deployment_drops_params,
},
}
]
)
processing_obj: Final = ProxyBaseLLMRequestProcessing(
data={
"model": requested_model,
"messages": [{"role": "user", "content": "hi"}],
**({"priority": body_priority} if body_priority is not None else {}),
}
)
returned_data, _ = await processing_obj.common_processing_pre_call_logic(
request=request,
general_settings={},
user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test", aliases={"key-alias": "qa-model"}),
proxy_logging_obj=ProxyLogging(user_api_key_cache=UserApiKeyCache()),
proxy_config=RealProxyConfig(),
route_type="acompletion",
llm_router=router,
)
assert ("priority" in returned_data, returned_data.get("priority")) == (
expected_priority is not None,
expected_priority,
)
@pytest.mark.asyncio
@pytest.mark.parametrize("safe_memory_mode", [False, True])
@pytest.mark.parametrize(