From cd5388bdb489373b72c85d83a8c7d6b9794918f8 Mon Sep 17 00:00:00 2001 From: narmstrong-berry Date: Tue, 10 Nov 2026 20:11:30 +0000 Subject: [PATCH] feat(messages): apply modify_params rewrites on rust --- .../python-bridge/src/routes/messages/host.rs | 2 +- .../python-bridge/src/routes/messages/mod.rs | 6 +- .../src/routes/messages/modify_params.rs | 92 +++++++++++++++++ .../crates/python-bridge/src/routes/mod.rs | 35 ++++++- litellm/messages/dispatch.py | 1 - litellm/rust_bridge/catalog.py | 1 - litellm/rust_bridge/host/routes/messages.py | 10 ++ litellm/utils.py | 99 +++++++------------ .../sdk/messages/test_native_parity.py | 46 ++++++++- .../messages/test_dispatch.py | 42 +++++++- .../rust_bridge/host/routes/test_messages.py | 26 +++++ 11 files changed, 289 insertions(+), 71 deletions(-) create mode 100644 litellm-rust/crates/python-bridge/src/routes/messages/modify_params.rs diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs index 2a5f79ca393..cfe7c16e3a4 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs @@ -28,7 +28,7 @@ use crate::{ routes::RequestView, }; -const ROUTE_HOST_MODULE: &str = "litellm.rust_bridge.host.routes.messages"; +pub(super) const ROUTE_HOST_MODULE: &str = "litellm.rust_bridge.host.routes.messages"; fn merge_headers( forwarded: Option>, diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs b/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs index a9ae01bd7ee..7c67db064e3 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs @@ -1,7 +1,9 @@ mod host; +mod modify_params; use host::MessagesPythonHost; use litellm_callbacks_legacy_python::LoggingOperation; +use modify_params::ModifyParams; use pyo3::prelude::*; use super::NativeCall; @@ -9,7 +11,8 @@ use super::NativeCall; fn run_messages(py: Python<'_>, call: NativeCall<'_>, asynchronous: bool) -> PyResult> { let (arguments, hooks) = crate::routes::call_hooks(py, LoggingOperation::Messages, &call, asynchronous)?; - crate::routes::run_public_call( + let modify_params = ModifyParams::new(&call.args, &call.kwargs)?; + crate::routes::run_public_call_after_policy( py, arguments, move |py, arguments, request| { @@ -46,6 +49,7 @@ fn run_messages(py: Python<'_>, call: NativeCall<'_>, asynchronous: bool) -> PyR }, MessagesPythonHost::new(call.view()?, asynchronous), hooks, + modify_params, asynchronous, ) } diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/modify_params.rs b/litellm-rust/crates/python-bridge/src/routes/messages/modify_params.rs new file mode 100644 index 00000000000..a0525dc2626 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/messages/modify_params.rs @@ -0,0 +1,92 @@ +use litellm_host::hooks::CallHooks; +use litellm_host_python::{PythonOwned, PythonRuntime}; +use pyo3::{ + gc::{PyTraverseError, PyVisit}, + prelude::*, + types::{PyDict, PyTuple}, +}; + +use super::host::ROUTE_HOST_MODULE; + +pub(super) struct ModifyParams { + call: Option<(Py, Py)>, +} + +impl ModifyParams { + pub(super) fn new(args: &Bound<'_, PyTuple>, kwargs: &Bound<'_, PyDict>) -> PyResult { + let model = match args.get_item(0) { + Ok(model) => model, + Err(_) => kwargs + .get_item("model")? + .unwrap_or_else(|| args.py().None().into_bound(args.py())), + }; + Ok(Self { + call: Some((model.unbind(), args.clone().unbind())), + }) + } +} + +impl CallHooks for ModifyParams { + fn arguments_prepared(&mut self, py: Python<'_>, arguments: &Py) -> PyResult<()> { + let Some((model, args)) = &self.call else { + return Ok(()); + }; + py.import(ROUTE_HOST_MODULE)? + .getattr("modify_max_tokens")? + .call1((model, args, arguments))?; + Ok(()) + } +} + +impl PythonOwned for ModifyParams { + fn close(&mut self, _: Python<'_>) { + self.call = None; + } + + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + if let Some((model, args)) = &self.call { + visit.call(model)?; + visit.call(args)?; + } + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use pyo3::{ + prelude::*, + types::{PyDict, PyTuple}, + }; + use rstest::rstest; + + use super::ModifyParams; + + #[rstest] + #[case::keyword_model( + "()", + "{'model': 'anthropic/claude-sonnet-4-5'}", + "'anthropic/claude-sonnet-4-5'" + )] + #[case::first_positional_wins("(64, [])", "{'model': 'anthropic/claude-sonnet-4-5'}", "64")] + #[case::no_model("()", "{}", "None")] + fn reads_the_model_the_client_wrapper_reads( + #[case] args: &str, + #[case] kwargs: &str, + #[case] expected: &str, + ) { + Python::initialize(); + Python::attach(|py| { + let eval = |source: &str| { + py.eval(&std::ffi::CString::new(source).unwrap(), None, None) + .unwrap() + }; + let args = eval(args).cast_into::().unwrap(); + let kwargs = eval(kwargs).cast_into::().unwrap(); + let hook = ModifyParams::new(&args, &kwargs).unwrap(); + let (model, captured) = hook.call.as_ref().unwrap(); + assert!(model.bind(py).eq(eval(expected)).unwrap()); + assert!(captured.bind(py).is(&args)); + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/routes/mod.rs b/litellm-rust/crates/python-bridge/src/routes/mod.rs index f51c755c33e..9fcd792bdaa 100644 --- a/litellm-rust/crates/python-bridge/src/routes/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/mod.rs @@ -166,6 +166,38 @@ fn run_public_call( hooks: impl PythonCallHooks + 'static, asynchronous: bool, ) -> PyResult> +where + H: PythonBinding + PythonHostCalls + 'static, + M: Machine + 'static, + M::Complete: Into::Response>>, +{ + run_public_call_after_policy( + py, + arguments, + start, + host, + hooks, + HookChain::new(), + asynchronous, + ) +} + +fn run_public_call_after_policy( + py: Python<'_>, + arguments: Py, + start: impl FnOnce( + Python<'_>, + &Bound<'_, PyDict>, + ::Request, + ) -> PyResult + + Send + + Sync + + 'static, + host: H, + hooks: impl PythonCallHooks + 'static, + after_policy: impl PythonCallHooks + 'static, + asynchronous: bool, +) -> PyResult> where H: PythonBinding + PythonHostCalls + 'static, M: Machine + 'static, @@ -179,7 +211,8 @@ where host, HookChain::new() .with(hooks) - .with(crate::preflight::SdkPolicy), + .with(crate::preflight::SdkPolicy) + .with(after_policy), arguments, CallOptions::new(asynchronous, execution_steps), ) diff --git a/litellm/messages/dispatch.py b/litellm/messages/dispatch.py index faa604a25ce..35e71ac1ccb 100644 --- a/litellm/messages/dispatch.py +++ b/litellm/messages/dispatch.py @@ -139,7 +139,6 @@ _FEATURES: Final[detectors.Detectors] = MappingProxyType( Feature.PROVIDER_AFFINITY: route_detectors.provider_affinity, Feature.SAFEGUARDS_CLASSIFIER: route_detectors.safeguards_classifier, Feature.RESPONSE_CACHE: detectors.response_cache, - Feature.MODIFY_PARAMS: detectors.modify_params, Feature.WORKLOAD_IDENTITY: route_detectors.anthropic_workload_identity, Feature.CACHE_CONTROL_INJECTION: detectors.cache_control_injection, Feature.STREAM_COST_INJECTION: route_detectors.stream_cost_injection, diff --git a/litellm/rust_bridge/catalog.py b/litellm/rust_bridge/catalog.py index c493617003d..90f27cf50a8 100644 --- a/litellm/rust_bridge/catalog.py +++ b/litellm/rust_bridge/catalog.py @@ -180,7 +180,6 @@ MESSAGES_PYTHON_FEATURES: Final = (COMMON_PYTHON_FEATURES - {Feature.TIMEOUT_VAR Feature.WEBSEARCH_INTERCEPTION, Feature.MESSAGES_INTERCEPTOR, Feature.RESPONSE_CACHE, - Feature.MODIFY_PARAMS, Feature.WORKLOAD_IDENTITY, Feature.CACHE_CONTROL_INJECTION, Feature.STREAM_COST_INJECTION, diff --git a/litellm/rust_bridge/host/routes/messages.py b/litellm/rust_bridge/host/routes/messages.py index c6568264edb..2b7197114b0 100644 --- a/litellm/rust_bridge/host/routes/messages.py +++ b/litellm/rust_bridge/host/routes/messages.py @@ -34,6 +34,8 @@ from litellm.rust_bridge import failures from litellm.rust_bridge.host import callbacks_legacy_python, http_transport from litellm.rust_bridge.public_call import known_params, optional_str from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse +from litellm.types.utils import CallTypes +from litellm.utils import clamp_max_tokens _METADATA: Final = TypeAdapter(dict[str, object]) _BODY_FIELDS: Final = frozenset( @@ -90,6 +92,14 @@ def inferred_provider(model: str, api_base: str | None) -> str | None: return provider +def modify_max_tokens( + model: object, + args: tuple[object, ...], + kwargs: dict[str, object], # mutable-ok: the call's keyword view, rewritten in place as `@client` does +) -> None: + clamp_max_tokens(model, args, kwargs, CallTypes.anthropic_messages.value) + + def body_params(request: Mapping[str, object]) -> Mapping[str, object]: params: Final = known_params(_BODY_FIELDS.difference(handler.additional_drop_paths(request)), request) if deployment_reasoning_effort_levels(request.get("model_info")) is None: diff --git a/litellm/utils.py b/litellm/utils.py index ca916e049d1..383a32cc32c 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2068,6 +2068,39 @@ def _is_litellm_router_call(kwargs: Mapping[str, object], *, is_async: bool) -> return any(isinstance(bucket, Mapping) and "model_group" in bucket for bucket in metadata_buckets) +def clamp_max_tokens( + model: object, + args: tuple[object, ...], + kwargs: dict[str, object], # mutable-ok: the @client wrapper rewrites the call's own kwargs + call_type: str, +) -> None: + if ( + kwargs.get("max_tokens", None) is not None + and model is not None + and litellm.modify_params is True + and call_type in (CallTypes.acompletion.value, CallTypes.completion.value, CallTypes.anthropic_messages.value) + ): + try: + base_model = model + if kwargs.get("hf_model_name", None) is not None: + base_model = f"huggingface/{kwargs.get('hf_model_name')}" + messages = None + if len(args) > 1: + messages = args[1] + elif kwargs.get("messages", None): + messages = kwargs["messages"] + kwargs["max_tokens"] = _get_modified_max_tokens()( + model=model, + base_model=base_model, + messages=messages, + user_max_tokens=kwargs.get("max_tokens"), + buffer_num=None, + buffer_perc=None, + ) + except Exception as e: + print_verbose(f"Error while checking max token limit: {e}") + + def client(original_function): from litellm.litellm_core_utils.core_helpers import max_retries_per_request_hit @@ -2184,38 +2217,7 @@ def client(original_function): verbose_logger.debug("Cache hit!") return caching_handler_response.cached_result - # CHECK MAX TOKENS - if ( - kwargs.get("max_tokens", None) is not None - and model is not None - and litellm.modify_params is True # user is okay with params being modified - and ( - call_type == CallTypes.acompletion.value - or call_type == CallTypes.completion.value - or call_type == CallTypes.anthropic_messages.value - ) - ): - try: - base_model = model - if kwargs.get("hf_model_name", None) is not None: - base_model = f"huggingface/{kwargs.get('hf_model_name')}" - messages = None - if len(args) > 1: - messages = args[1] - elif kwargs.get("messages", None): - messages = kwargs["messages"] - user_max_tokens: Final = kwargs.get("max_tokens") - modified_max_tokens: Final = _get_modified_max_tokens()( - model=model, - base_model=base_model, - messages=messages, - user_max_tokens=user_max_tokens, - buffer_num=None, - buffer_perc=None, - ) - kwargs["max_tokens"] = modified_max_tokens - except Exception as e: - print_verbose(f"Error while checking max token limit: {e}") + clamp_max_tokens(model, args, kwargs, call_type) # MODEL CALL result = original_function(*args, **kwargs) end_time = datetime.datetime.now() @@ -2480,38 +2482,7 @@ def client(original_function): if _llm_caching_handler.preset_cache_key is not None: logging_obj.litellm_params["preset_cache_key"] = _llm_caching_handler.preset_cache_key - # CHECK MAX TOKENS - if ( - kwargs.get("max_tokens", None) is not None - and model is not None - and litellm.modify_params is True # user is okay with params being modified - and ( - call_type == CallTypes.acompletion.value - or call_type == CallTypes.completion.value - or call_type == CallTypes.anthropic_messages.value - ) - ): - try: - base_model = model - if kwargs.get("hf_model_name", None) is not None: - base_model = f"huggingface/{kwargs.get('hf_model_name')}" - messages = None - if len(args) > 1: - messages = args[1] - elif kwargs.get("messages", None): - messages = kwargs["messages"] - user_max_tokens: Final = kwargs.get("max_tokens") - modified_max_tokens: Final = _get_modified_max_tokens()( - model=model, - base_model=base_model, - messages=messages, - user_max_tokens=user_max_tokens, - buffer_num=None, - buffer_perc=None, - ) - kwargs["max_tokens"] = modified_max_tokens - except Exception as e: - print_verbose(f"Error while checking max token limit: {e}") + clamp_max_tokens(model, args, kwargs, call_type) # MODEL CALL call_kwargs: Final = ( diff --git a/tests/rust-python-harness/strategies/e2e_parity/sdk/messages/test_native_parity.py b/tests/rust-python-harness/strategies/e2e_parity/sdk/messages/test_native_parity.py index 41008f0d57f..b3e7fe5053b 100644 --- a/tests/rust-python-harness/strategies/e2e_parity/sdk/messages/test_native_parity.py +++ b/tests/rust-python-harness/strategies/e2e_parity/sdk/messages/test_native_parity.py @@ -161,6 +161,7 @@ THINKING_REPLAY: Final[tuple[JsonValue, ...]] = ( {"role": "user", "content": "And Germany?"}, ) +_MODIFY_PARAMS: Final[Mapping[str, object]] = MappingProxyType({"modify_params": True}) NO_INPUT: Final[Mapping[str, JsonValue]] = MappingProxyType({}) PNG_PIXEL: Final = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" @@ -757,7 +758,50 @@ FALLBACK_CASES: Final = ( _ok(), ), ), - Case("modify_params", _kw(max_tokens=1), settings=MappingProxyType({"modify_params": True})), + Case("modify_params", _kw(max_tokens=1), settings=_MODIFY_PARAMS), + Case("modify_params_clamps_max_tokens", _kw(max_tokens=200_000), settings=_MODIFY_PARAMS), + Case( + "modify_params_clamps_streamed_max_tokens", + _kw(stream=True, max_tokens=200_000), + upstream=(_stream(TEXT_STREAM),), + settings=_MODIFY_PARAMS, + ), + Case( + "modify_params_clamps_tool_turns", + _kw(messages=list(TOOL_TURNS), tools=[dict(WEATHER_TOOL)], max_tokens=200_000), + upstream=(_ok(TOOL_USE_RESPONSE),), + settings=_MODIFY_PARAMS, + ), + Case( + "modify_params_clamps_thinking", + _kw(thinking={"type": "enabled", "budget_tokens": 2048}, max_tokens=200_000), + upstream=(_ok(THINKING_RESPONSE),), + settings=_MODIFY_PARAMS, + ), + Case( + "modify_params_image_content", + _kw( + messages=[ + { + "role": "user", + "content": [ + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": PNG_PIXEL}}, + {"type": "text", "text": "What is this?"}, + ], + } + ], + max_tokens=200_000, + ), + settings=_MODIFY_PARAMS, + ), + Case( + "modify_params_unknown_model", + _kw(max_tokens=200_000), + model="anthropic/claude-parity-unknown", + settings=_MODIFY_PARAMS, + ), + Case("modify_params_hf_model_name", _kw(max_tokens=200_000, hf_model_name="gpt2"), settings=_MODIFY_PARAMS), + Case("modify_params_sync", _kw(max_tokens=200_000), sync=True, settings=_MODIFY_PARAMS), Case( "stream_cost_injection", _kw(stream=True), diff --git a/tests/test_litellm_rust/messages/test_dispatch.py b/tests/test_litellm_rust/messages/test_dispatch.py index b580eca2e42..2aa99238942 100644 --- a/tests/test_litellm_rust/messages/test_dispatch.py +++ b/tests/test_litellm_rust/messages/test_dispatch.py @@ -17,6 +17,7 @@ from litellm.messages.dispatch import ( _DISPATCH, # pyright: ignore[reportPrivateUsage] # tests configured dispatch ) from litellm.proxy import proxy_server +from litellm.rust_bridge import catalog from litellm.rust_bridge.bindings import NativeBinding from litellm.rust_bridge.catalog import Feature, Route, RouteRule, Rules, Target from litellm.rust_bridge.configuration import Rollout @@ -478,7 +479,6 @@ STRICT_TOOL: Final = {"name": "a", "input_schema": {"type": "object"}, "strict": (Feature.AGENTIC_LOOP_HOOK, "anthropic/claude-sonnet-4-5", {}, {"callbacks": [_Agentic()]}), (Feature.RESPONSE_CACHE, "anthropic/claude-sonnet-4-5", {}, {"cache": Cache()}), (Feature.TLS_POLICY, "anthropic/claude-sonnet-4-5", {}, {"ssl_security_level": "DEFAULT@SECLEVEL=1"}), - (Feature.MODIFY_PARAMS, "anthropic/claude-sonnet-4-5", {}, {"modify_params": True}), ( Feature.WORKLOAD_IDENTITY, "anthropic/claude-sonnet-4-5", @@ -617,6 +617,46 @@ async def test_strict_tools_stay_native_unless_python_could_drop_strict( assert served is native_result +@pytest.mark.asyncio +async def test_modify_params_stays_native(monkeypatch: pytest.MonkeyPatch) -> None: + rules: Final[Rules] = ( + RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY, features=catalog.MESSAGES_PYTHON_FEATURES), + RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED), + ) + native_result: Final = response() + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setattr(litellm, "modify_params", True) + + def python(*call_args: object, **call_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: rejected fallback + pytest.fail("modify_params must not fall back to Python") + + async def apython(*call_args: object, **call_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: stub + pytest.fail("modify_params must not fall back to Python") + + async def native(request: NativeCall) -> AnthropicMessagesResponse: + return native_result + + call_args: Final = (200_000, MESSAGES, "anthropic/claude-sonnet-4-5") + served: Final = _DISPATCH.run( + call_args, + {}, + python=python, + binding=messages_binding(lambda request: native_result), + native=lambda hook, request, args, kwargs: hook(request), + rules=rules, + ) + aserved: Final = await _ADISPATCH.arun( + call_args, + {}, + python=apython, + binding=amessages_binding(native), + native=lambda hook, request, args, kwargs: hook(request), + rules=rules, + ) + + assert (served, aserved) == (native_result, native_result) + + ASYNC_ONLY_FEATURES: Final = ( pytest.param(Feature.PRE_REQUEST_HOOK, {}, {"callbacks": [_Rewrite()]}, id="pre-request-hook"), pytest.param( diff --git a/tests/unit/rust_bridge/host/routes/test_messages.py b/tests/unit/rust_bridge/host/routes/test_messages.py index aa7dcc74841..5d919abf7b9 100644 --- a/tests/unit/rust_bridge/host/routes/test_messages.py +++ b/tests/unit/rust_bridge/host/routes/test_messages.py @@ -445,3 +445,29 @@ def test_a_failure_before_any_response_maps_like_the_python_transport_error( f"litellm.InternalServerError: AnthropicException - {text}. Handle with `litellm.InternalServerError`." ) assert "This looks like a bug in LiteLLM" not in str(mapped) + + +@pytest.mark.parametrize( + ("modify_params", "args", "kwargs", "expected"), + [ + pytest.param(True, (), {"max_tokens": 200_000}, 64_000, id="clamped-to-model-output-limit"), + pytest.param(True, (), {"max_tokens": 1}, 1, id="under-the-limit"), + pytest.param(False, (), {"max_tokens": 200_000}, 200_000, id="modify-params-off"), + pytest.param("true", (), {"max_tokens": 200_000}, 200_000, id="truthy-but-not-true"), + pytest.param(True, (200_000,), {"max_tokens": 200_000}, 200_000, id="positional-model-is-not-a-model"), + ], +) +def test_modify_max_tokens_clamps_the_keyword_view_as_the_client_wrapper_does( + monkeypatch: pytest.MonkeyPatch, + modify_params: object, + args: tuple[object, ...], + kwargs: dict[str, object], + expected: int, +) -> None: + monkeypatch.setattr(litellm, "modify_params", modify_params) + model: Final = args[0] if args else "anthropic/claude-sonnet-4-5" + view: Final = {**kwargs, "messages": [{"role": "user", "content": "hi"}]} + + messages_host.modify_max_tokens(model, args, view) + + assert view["max_tokens"] == expected