mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
fix(handler): stop extra_body from overriding the authorized model
The shared BaseLLMHTTPHandler merged caller-controlled extra_body over the transform-produced request body, so extra_body.model could swap the model that runs upstream while proxy model-access and budget checks only ever inspect the top-level model parameter. Any OpenAI-compatible provider that reads model from the body was affected. Merge through a small helper that re-asserts the transform's model after the extra_body merge, so the escape hatch can no longer change which model executes. Applied uniformly across the chat completion, responses API, and generate_content paths rather than special-casing one provider. Closes #31118
This commit is contained in:
parent
c546b58c09
commit
da7509d2c5
3 changed files with 57 additions and 5 deletions
|
|
@ -10,6 +10,7 @@ from typing import (
|
|||
Dict,
|
||||
List,
|
||||
Literal,
|
||||
Mapping,
|
||||
Optional,
|
||||
Tuple,
|
||||
Union,
|
||||
|
|
@ -194,6 +195,15 @@ def _responses_api_optional_request_param_names() -> frozenset[str]:
|
|||
return frozenset(get_type_hints(ResponsesAPIOptionalRequestParams).keys())
|
||||
|
||||
|
||||
def _merge_extra_body_preserving_model(
|
||||
transformed_body: Mapping[str, object], extra_body: Mapping[str, object]
|
||||
) -> dict[str, object]:
|
||||
merged = {**transformed_body, **extra_body}
|
||||
if "model" in transformed_body:
|
||||
return {**merged, "model": transformed_body["model"]}
|
||||
return merged
|
||||
|
||||
|
||||
def _custom_logger_callbacks(logging_obj: Any) -> list[Any]:
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.litellm_logging import (
|
||||
|
|
@ -496,7 +506,7 @@ class BaseLLMHTTPHandler:
|
|||
)
|
||||
|
||||
if extra_body is not None:
|
||||
data = {**data, **extra_body}
|
||||
data = _merge_extra_body_preserving_model(data, extra_body)
|
||||
|
||||
headers, signed_json_body = provider_config.sign_request(
|
||||
headers=headers,
|
||||
|
|
@ -2451,7 +2461,7 @@ class BaseLLMHTTPHandler:
|
|||
data = BaseResponsesAPIConfig.normalize_responses_api_request_dict(data)
|
||||
|
||||
if extra_body:
|
||||
data.update(extra_body)
|
||||
data = _merge_extra_body_preserving_model(data, extra_body)
|
||||
stream = bool(stream or data.get("stream"))
|
||||
|
||||
# Preserve the OpenAI-style request context (not sent to the provider) for streaming
|
||||
|
|
@ -2638,7 +2648,7 @@ class BaseLLMHTTPHandler:
|
|||
data = BaseResponsesAPIConfig.normalize_responses_api_request_dict(data)
|
||||
|
||||
if extra_body:
|
||||
data.update(extra_body)
|
||||
data = _merge_extra_body_preserving_model(data, extra_body)
|
||||
stream = bool(stream or data.get("stream"))
|
||||
|
||||
# Preserve the OpenAI-style request context (not sent to the provider) for streaming
|
||||
|
|
@ -11189,7 +11199,7 @@ class BaseLLMHTTPHandler:
|
|||
)
|
||||
|
||||
if extra_body:
|
||||
data.update(extra_body)
|
||||
data = _merge_extra_body_preserving_model(data, extra_body)
|
||||
|
||||
## LOGGING
|
||||
logging_obj.pre_call(
|
||||
|
|
@ -11304,7 +11314,7 @@ class BaseLLMHTTPHandler:
|
|||
)
|
||||
|
||||
if extra_body:
|
||||
data.update(extra_body)
|
||||
data = _merge_extra_body_preserving_model(data, extra_body)
|
||||
|
||||
## LOGGING
|
||||
logging_obj.pre_call(
|
||||
|
|
|
|||
|
|
@ -18,6 +18,7 @@ from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
|
|||
from litellm.llms.custom_httpx.llm_http_handler import (
|
||||
BaseLLMHTTPHandler,
|
||||
_google_genai_streaming_hidden_params,
|
||||
_merge_extra_body_preserving_model,
|
||||
)
|
||||
from litellm.types.llms.openai import ResponsesAPIResponse
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
|
|
@ -271,6 +272,25 @@ async def test_async_response_api_handler_streams_when_provider_transform_adds_s
|
|||
assert client.post.call_args.kwargs["json"]["stream"] is True
|
||||
|
||||
|
||||
def test_merge_extra_body_preserving_model_blocks_model_override():
|
||||
merged = _merge_extra_body_preserving_model(
|
||||
{"model": "authorized", "messages": []},
|
||||
{"model": "attacker", "passthrough": "kept"},
|
||||
)
|
||||
assert merged["model"] == "authorized"
|
||||
assert merged["passthrough"] == "kept"
|
||||
|
||||
|
||||
def test_merge_extra_body_preserving_model_passes_through_when_no_model():
|
||||
merged = _merge_extra_body_preserving_model(
|
||||
{"input": "hi"},
|
||||
{"model": "from-extra-body", "extra": "kept"},
|
||||
)
|
||||
assert merged["model"] == "from-extra-body"
|
||||
assert merged["extra"] == "kept"
|
||||
assert merged["input"] == "hi"
|
||||
|
||||
|
||||
def test_get_agentic_loop_settings_defaults_and_overrides():
|
||||
handler = BaseLLMHTTPHandler()
|
||||
|
||||
|
|
|
|||
|
|
@ -433,6 +433,28 @@ def test_custom_provider_with_extra_body():
|
|||
}
|
||||
|
||||
|
||||
def test_extra_body_cannot_override_authorized_model():
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
|
||||
client = HTTPHandler()
|
||||
with patch.object(client, "post") as mock_post:
|
||||
try:
|
||||
litellm.completion(
|
||||
model="deepseek/deepseek-chat",
|
||||
messages=[{"role": "user", "content": "Hello, how are you?"}],
|
||||
extra_body={"model": "deepseek/some-other-model", "passthrough": "kept"},
|
||||
api_key="fake-key-for-testing",
|
||||
client=client,
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
mock_post.assert_called_once()
|
||||
sent_body = json.loads(mock_post.call_args.kwargs["data"])
|
||||
assert sent_body["model"] == "deepseek-chat"
|
||||
assert sent_body["passthrough"] == "kept"
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def set_openrouter_api_key():
|
||||
original_api_key = os.environ.get("OPENROUTER_API_KEY")
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue