mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
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:
parent
ecb8ef5ca8
commit
bfdea7808a
2 changed files with 94 additions and 1 deletions
|
|
@ -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 = (
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue