fix(router)!: keep the deployment's extra_body fields when a request sends its own extra_body

The router builds each call as {**deployment litellm_params, **request
kwargs}, so a request's extra_body replaced the deployment's whole
extra_body. Any client that sent extra_body for one field silently dropped
every field the admin configured there, and could replace an admin field
with its own value

_update_kwargs_with_deployment now merges the two, request keys first and
the deployment's last, the same place deployment tools are merged with
request tools. A client can still add fields, but a key the deployment
sets keeps the deployment's value. A retry or fallback to another
deployment starts again from the request's own extra_body, so one
deployment's fields never reach the next

The regression tests send requests through Router.completion against a
mock transport and assert the outbound request body
This commit is contained in:
Mikalai Karytka 2026-11-03 09:06:04 +01:00
parent ecb8ef5ca8
commit bfdea7808a
No known key found for this signature in database
2 changed files with 94 additions and 1 deletions

View file

@ -678,6 +678,7 @@ _RESOLVED_RETRY_POLICY_ADAPTER: Final = TypeAdapter(RetryPolicy | None)
_ROUTING_KWARGS_ADAPTER: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None)
_FALLBACK_HOP_ADAPTER: Final = TypeAdapter(Mapping[str, object])
_REQUEST_KWARGS_ADAPTER: Final = TypeAdapter(Mapping[str, object])
_EXTRA_BODY_ADAPTER: Final = TypeAdapter(Mapping[str, object])
_DEPLOYMENT_PARAMS_ADAPTER: Final[TypeAdapter[Mapping[str, object] | GenericLiteLLMParams]] = TypeAdapter(
Mapping[str, object] | GenericLiteLLMParams
)
@ -749,6 +750,17 @@ def _drop_malformed_request_params_the_deployment_sets(
kwargs.pop(key)
def _merged_extra_body(
deployment_extra_body: object, request_extra_body: object
) -> dict[str, object] | None: # mutable-ok: provider and cache code reads extra_body only when it is a plain dict
if not isinstance(deployment_extra_body, Mapping) or not isinstance(request_extra_body, Mapping):
return None
return {
**_EXTRA_BODY_ADAPTER.validate_python(request_extra_body),
**_EXTRA_BODY_ADAPTER.validate_python(deployment_extra_body),
}
def _merged_tools(deployment_tools: Sequence[object], request_tools: object) -> list[object] | None:
if not isinstance(request_tools, (list, tuple)):
return None
@ -5303,6 +5315,15 @@ class PythonCore(DeprecatedRouterMembers):
if "tool_choice" not in kwargs and dep_params.get("tool_choice") is not None:
kwargs["tool_choice"] = dep_params["tool_choice"]
@staticmethod
def _merge_extra_body_from_deployment(
deployment_params: Mapping[str, object],
kwargs: dict[str, object], # mutable-ok: sets extra_body on the router's request kwargs in place
) -> None:
merged_extra_body: Final = _merged_extra_body(deployment_params.get("extra_body"), kwargs.get("extra_body"))
if merged_extra_body is not None:
kwargs["extra_body"] = merged_extra_body # rebind-ok: updates the router's request kwargs in place
def _update_kwargs_with_deployment(
self,
deployment: dict,
@ -5312,10 +5333,11 @@ class PythonCore(DeprecatedRouterMembers):
opens_responses_socket: bool = False,
) -> None:
"""
5 jobs:
6 jobs:
- Adds selected deployment, model_info and api_base to kwargs["metadata"] (used for logging)
- Adds default litellm params to kwargs, if set.
- Merges tools from deployment with request (proxy-configured tools + request tools).
- Merges extra_body from deployment with request, the deployment's value winning on a shared key.
- On a fallback hop, drops the encrypted reasoning this deployment cannot decrypt, keeping its summary.
- On a Responses continuation, sends the partial response as prefill where this deployment takes one.
"""
@ -5347,6 +5369,10 @@ class PythonCore(DeprecatedRouterMembers):
deployment_context,
defaults,
)
self._merge_extra_body_from_deployment(
deployment_params=deployment_context,
kwargs=kwargs, # pyright: ignore[reportUnknownArgumentType] # kwargs is the router's untyped request dict
)
declared_provider: Final = deployment_context.get("custom_llm_provider")
declared_model: Final = deployment_context.get("model")
deployment_provider: Final = (

View file

@ -10773,6 +10773,73 @@ def test_update_kwargs_with_deployment_merges_tools():
assert kwargs["tool_choice"] == "auto"
_EXTRA_BODY_CHAT_RESPONSE: Final = {
"id": "chatcmpl-extra-body",
"object": "chat.completion",
"created": 1,
"model": "gpt-x",
"choices": [{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}],
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
}
def _deployment_with_extra_body(model_name: str, api_base: str, extra_body: Mapping[str, object]) -> dict[str, object]:
return {
"model_name": model_name,
"litellm_params": {"model": "openai/gpt-x", "api_key": "k", "api_base": api_base, "extra_body": extra_body},
}
def test_router_keeps_deployment_extra_body_fields_when_the_request_sends_its_own(
respx_mock: respx.MockRouter,
) -> None:
route: Final = respx_mock.post("https://a.example/v1/chat/completions").respond(json=_EXTRA_BODY_CHAT_RESPONSE)
router: Final = Router(
model_list=[
_deployment_with_extra_body(
"g", "https://a.example/v1", {"set_by_admin": 1, "shared": "deployment"}
)
]
)
router.completion(
model="g",
messages=[{"role": "user", "content": "hi"}],
extra_body={"set_by_client": 0, "shared": "client"},
)
body: Final = json.loads(route.calls.last.request.content)
assert {k: body.get(k) for k in ("set_by_admin", "set_by_client", "shared")} == {
"set_by_admin": 1,
"set_by_client": 0,
"shared": "deployment",
}
def test_router_fallback_sends_only_the_next_deployments_extra_body(respx_mock: respx.MockRouter) -> None:
respx_mock.post("https://a.example/v1/chat/completions").respond(status_code=500, json={"error": {"message": "x"}})
fallback_route: Final = respx_mock.post("https://b.example/v1/chat/completions").respond(
json=_EXTRA_BODY_CHAT_RESPONSE
)
router: Final = Router(
model_list=[
_deployment_with_extra_body("g", "https://a.example/v1", {"only_on_a": 1}),
_deployment_with_extra_body("h", "https://b.example/v1", {"only_on_b": 2}),
],
num_retries=0,
fallbacks=[{"g": ["h"]}],
)
router.completion(model="g", messages=[{"role": "user", "content": "hi"}], extra_body={"set_by_client": 0})
body: Final = json.loads(fallback_route.calls.last.request.content)
assert {k: body.get(k) for k in ("only_on_a", "only_on_b", "set_by_client")} == {
"only_on_a": None,
"only_on_b": 2,
"set_by_client": 0,
}
def test_update_kwargs_with_deployment_merge_tools_deployment_only():
"""
Test that when only deployment has tools, they are applied to kwargs.