From 12c6283fd34b74a23d344d31eabed11278631d0e 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/cache/selection.rs | 23 ++- .../python-bridge/src/routes/messages/host.rs | 2 +- .../python-bridge/src/routes/messages/mod.rs | 6 +- .../src/routes/messages/modify_params.rs | 152 ++++++++++++++++++ .../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 | 11 +- litellm/utils.py | 99 ++++-------- .../sdk/messages/test_native_parity.py | 60 ++++++- .../messages/test_dispatch.py | 42 ++++- .../rust_bridge/host/routes/test_messages.py | 26 +++ 12 files changed, 378 insertions(+), 80 deletions(-) create mode 100644 litellm-rust/crates/python-bridge/src/routes/messages/modify_params.rs diff --git a/litellm-rust/crates/python-bridge/src/cache/selection.rs b/litellm-rust/crates/python-bridge/src/cache/selection.rs index 66321a8ae87..a95b786f55a 100644 --- a/litellm-rust/crates/python-bridge/src/cache/selection.rs +++ b/litellm-rust/crates/python-bridge/src/cache/selection.rs @@ -51,15 +51,21 @@ fn selected_cache<'py>( kwargs: &Bound<'py, PyDict>, call_type: &str, ) -> PyResult>> { + if kwargs + .get_item("caching")? + .is_some_and(|value| value.is(pyo3::types::PyBool::new(py, false))) + { + return Ok(None); + } + configured_cache(py, call_type) +} + +fn configured_cache<'py>(py: Python<'py>, call_type: &str) -> PyResult>> { let configured = py .import(CACHE_HOST_MODULE)? .getattr("configured_cache")? .call0()?; - if configured.is_none() - || kwargs - .get_item("caching")? - .is_some_and(|value| value.is(pyo3::types::PyBool::new(py, false))) - { + if configured.is_none() { return Ok(None); } let supported = configured.getattr("supported_call_types")?; @@ -94,7 +100,7 @@ pub(crate) fn configure( ) -> PyResult { let call_type: String = function.getattr("__name__")?.extract()?; let selected = if python.asynchronous() { - selected_cache(py, arguments, &call_type)? + configured_cache(py, &call_type)? } else { None }; @@ -107,6 +113,9 @@ pub(crate) fn configure( let enabled = cache .call_method("should_use_cache", (), Some(arguments))? .extract::()?; + let skips_read = arguments + .get_item("caching")? + .is_some_and(|value| !value.is_none() && !value.is(pyo3::types::PyBool::new(py, true))); let controls = arguments .get_item("cache")? .filter(|value| !value.is_none()); @@ -122,7 +131,7 @@ pub(crate) fn configure( }; let options = CacheOptions { policy: CachePolicy { - no_cache: boolean("no-cache")?, + no_cache: boolean("no-cache")? || skips_read, no_store: boolean("no-store")?, ..CachePolicy::default() }, 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 ed63c8011a8..f84fe5b1bb9 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..c6f846478ad --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/messages/modify_params.rs @@ -0,0 +1,152 @@ +use litellm_host::hooks::CallHooks; +use litellm_host_python::{PythonOwned, PythonRuntime}; +use pyo3::{ + exceptions::PyException, + 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(()); + }; + let clamped = py + .import(ROUTE_HOST_MODULE) + .and_then(|module| module.getattr("modify_max_tokens")) + .and_then(|clamp| clamp.call1((model, args, arguments))); + match clamped { + Err(error) if !error.is_instance_of::(py) => Err(error), + _ => 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, ROUTE_HOST_MODULE}; + use litellm_host::hooks::CallHooks; + + #[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)); + }); + } + + #[rstest] + #[case::exception_leaves_max_tokens("ValueError", true)] + #[case::base_exception_propagates("KeyboardInterrupt", false)] + fn a_failing_clamp_behaves_like_the_client_wrappers_swallowed_try( + #[case] raised: &str, + #[case] swallowed: bool, + ) { + Python::initialize(); + Python::attach(|py| { + let locals = PyDict::new(py); + locals.set_item("name", ROUTE_HOST_MODULE).unwrap(); + locals.set_item("raised", raised).unwrap(); + py.run( + c" +import builtins, sys, types +for parent in ('litellm', 'litellm.rust_bridge', 'litellm.rust_bridge.host', 'litellm.rust_bridge.host.routes'): + sys.modules.setdefault(parent, types.ModuleType(parent)) +saved = sys.modules.get(name) +stub = types.ModuleType(name) +def modify_max_tokens(model, args, kwargs): + kwargs['max_tokens'] = 1 + raise getattr(builtins, raised)() +stub.modify_max_tokens = modify_max_tokens +sys.modules[name] = stub +", + Some(&locals), + Some(&locals), + ) + .unwrap(); + let args = PyTuple::empty(py); + let kwargs = PyDict::new(py); + kwargs + .set_item("model", "anthropic/claude-sonnet-4-5") + .unwrap(); + let mut hook = ModifyParams::new(&args, &kwargs).unwrap(); + let result = hook.arguments_prepared(py, &kwargs.clone().unbind()); + py.run( + c" +if saved is None: + sys.modules.pop(name, None) +else: + sys.modules[name] = saved +", + Some(&locals), + Some(&locals), + ) + .unwrap(); + assert_eq!(result.is_ok(), swallowed); + if let Err(error) = result { + assert!(error.is_instance_of::(py)); + } + }); + } +} 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 9ac52affddf..c1bfbdf1c85 100644 --- a/litellm/messages/dispatch.py +++ b/litellm/messages/dispatch.py @@ -139,7 +139,6 @@ _FEATURES: Final[detectors.Detectors] = MappingProxyType( { **{feature: detect for feature, detect in detectors.COMMON.items() if feature not in _NATIVE_COMMON_FEATURES}, Feature.SAFEGUARDS_CLASSIFIER: route_detectors.safeguards_classifier, - Feature.MODIFY_PARAMS: detectors.modify_params, Feature.WORKLOAD_IDENTITY: route_detectors.anthropic_workload_identity, } ) diff --git a/litellm/rust_bridge/catalog.py b/litellm/rust_bridge/catalog.py index c8d3f00166b..da011a69fcd 100644 --- a/litellm/rust_bridge/catalog.py +++ b/litellm/rust_bridge/catalog.py @@ -169,7 +169,6 @@ MESSAGES_PYTHON_FEATURES: Final = (COMMON_PYTHON_FEATURES - {Feature.TIMEOUT_VAR Feature.PRE_REQUEST_HOOK, Feature.WEBSEARCH_INTERCEPTION, Feature.MESSAGES_INTERCEPTOR, - Feature.MODIFY_PARAMS, Feature.WORKLOAD_IDENTITY, Feature.STRICT_TOOLS_RETRY, } diff --git a/litellm/rust_bridge/host/routes/messages.py b/litellm/rust_bridge/host/routes/messages.py index cc213566236..081359c3a5e 100644 --- a/litellm/rust_bridge/host/routes/messages.py +++ b/litellm/rust_bridge/host/routes/messages.py @@ -53,7 +53,8 @@ from litellm.rust_bridge.host import callbacks_legacy_python, http_transport from litellm.rust_bridge.public_call import NativeCall, known_params, optional_str from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse from litellm.types.router import GenericLiteLLMParams -from litellm.types.utils import ProviderSpecificHeader +from litellm.types.utils import CallTypes, ProviderSpecificHeader +from litellm.utils import clamp_max_tokens _METADATA: Final = TypeAdapter(dict[str, object]) _BODY_FIELDS: Final = frozenset( @@ -114,6 +115,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_fields(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 a3a1fe5ede4..74a9ef4d225 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2066,6 +2066,39 @@ def _is_litellm_router_call(kwargs: Mapping[str, object]) -> 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 @@ -2182,38 +2215,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() @@ -2478,38 +2480,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 02b29e81c38..6c0c65f8f42 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 @@ -166,6 +166,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==" @@ -329,7 +330,7 @@ class Case: gaps: Mapping[str, str | None] = field(default_factory=lambda: MappingProxyType({})) def params(self) -> Iterator[object]: - for aspect in ASPECTS: + for aspect in (*ASPECTS, "cached") if self.cache is not None else ASPECTS: gap = {**KNOWN_GAPS.get(self.name, {}), **self.gaps}.get(aspect) marks = () if not gap else (pytest.mark.xfail(reason=gap, strict=True),) yield pytest.param(self, aspect, id=f"{self.name}-{aspect}", marks=marks) @@ -911,7 +912,58 @@ 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_cache_write_without_read", + _kw(max_tokens=200_000, cache={"no-cache": True}), + upstream=(_ok(), _ok()), + settings=_MODIFY_PARAMS, + cache=Cache, + ), + Case("cache_write_without_read", _kw(cache={"no-cache": True}), upstream=(_ok(), _ok()), cache=Cache), + Case("modify_params_sync", _kw(max_tokens=200_000), sync=True, settings=_MODIFY_PARAMS), Case("request_timeout_kwarg", _kw(request_timeout=30)), Case("stream_timeout_kwarg", _kw(stream=True, stream_timeout=30), upstream=(_stream(TEXT_STREAM),)), ) @@ -1047,6 +1099,7 @@ class Leg: hidden_params: Mapping[str, object] logged: Mapping[str, object] callback_kwargs: Mapping[str, object] + cached: tuple[str, ...] = () @dataclass(slots=True) @@ -1325,6 +1378,7 @@ async def _leg(case: Case, server: RecordingServer, probe: FallbackProbe | None) hidden_params=hidden, logged=logged, callback_kwargs=callback_kwargs, + cached=() if litellm.cache is None else tuple(sorted(litellm.cache.cache.cache_dict)), ) @@ -1371,6 +1425,8 @@ async def test_native_messages_match_python( assert native.hidden_params == python.hidden_params case "logged": assert native.logged == python.logged + case "cached": + assert native.cached == python.cached case _: assert native.callback_kwargs.get(aspect) == python.callback_kwargs.get(aspect) diff --git a/tests/test_litellm_rust/messages/test_dispatch.py b/tests/test_litellm_rust/messages/test_dispatch.py index 40e141673db..c44b810b255 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 MESSAGES_PYTHON_FEATURES, RULES, Feature, Route, RouteRule, Rules, Target from litellm.rust_bridge.configuration import Rollout @@ -477,7 +478,6 @@ STRICT_TOOL: Final = {"name": "a", "input_schema": {"type": "object"}, "strict": (Feature.CALLER_CLIENT, "anthropic/claude-sonnet-4-5", {"client": object()}, {}), (Feature.AGENTIC_LOOP_HOOK, "anthropic/claude-sonnet-4-5", {}, {"callbacks": [_Agentic()]}), (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", @@ -676,6 +676,46 @@ async def test_a_configured_response_cache_stays_native( 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 0ad2bc356e3..ec5bb519aee 100644 --- a/tests/unit/rust_bridge/host/routes/test_messages.py +++ b/tests/unit/rust_bridge/host/routes/test_messages.py @@ -627,3 +627,29 @@ def test_body_params_prepare_prompt_caching_as_the_entry_that_ran( ) assert fields["system"] == expected_system + + +@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