scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-12 08:42:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 207

feat(messages): apply modify_params rewrites on rust

Browse source
This commit is contained in:
narmstrong-berry 2026-11-10 20:11:30 +00:00
parent f0d234c4c6
commit cd5388bdb4
No known key found for this signature in database
11 changed files with 289 additions and 71 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm-rust/crates/python-bridge/src/routes/messages/host.rs
View file

@ -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<Map<String, Value>>,

6
litellm-rust/crates/python-bridge/src/routes/messages/mod.rs
View file

@ -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<Py<PyAny>> {
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,
)
}

92
litellm-rust/crates/python-bridge/src/routes/messages/modify_params.rs Normal file
View file

@ -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<PyAny>, Py<PyTuple>)>,
}
impl ModifyParams {
pub(super) fn new(args: &Bound<'_, PyTuple>, kwargs: &Bound<'_, PyDict>) -> PyResult<Self> {
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<PythonRuntime> for ModifyParams {
fn arguments_prepared(&mut self, py: Python<'_>, arguments: &Py<PyDict>) -> 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::<PyTuple>().unwrap();
let kwargs = eval(kwargs).cast_into::<PyDict>().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));
});
}
}

35
litellm-rust/crates/python-bridge/src/routes/mod.rs
View file

@ -166,6 +166,38 @@ fn run_public_call<H, M>(
hooks: impl PythonCallHooks + 'static,
asynchronous: bool,
) -> PyResult<Py<PyAny>>
where
H: PythonBinding + PythonHostCalls<H::Protocol> + 'static,
M: Machine<Protocol = H::Protocol> + 'static,
M::Complete: Into<HostedCompletion<<H::Protocol as Protocol>::Response>>,
{
run_public_call_after_policy(
py,
arguments,
start,
host,
hooks,
HookChain::new(),
asynchronous,
)
}
fn run_public_call_after_policy<H, M>(
py: Python<'_>,
arguments: Py<PyDict>,
start: impl FnOnce(
Python<'_>,
&Bound<'_, PyDict>,
<H::Protocol as Protocol>::Request,
) -> PyResult<M>
+ Send
+ Sync
+ 'static,
host: H,
hooks: impl PythonCallHooks + 'static,
after_policy: impl PythonCallHooks + 'static,
asynchronous: bool,
) -> PyResult<Py<PyAny>>
where
H: PythonBinding + PythonHostCalls<H::Protocol> + 'static,
M: Machine<Protocol = H::Protocol> + '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),
)

1
litellm/messages/dispatch.py
View file

@ -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,

1
litellm/rust_bridge/catalog.py
View file

@ -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,

10
litellm/rust_bridge/host/routes/messages.py
View file

@ -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:

99
litellm/utils.py
View file

@ -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 = (

46
tests/rust-python-harness/strategies/e2e_parity/sdk/messages/test_native_parity.py
View file

@ -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),

42
tests/test_litellm_rust/messages/test_dispatch.py
View file

@ -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(

26
tests/unit/rust_bridge/host/routes/test_messages.py
View file

@ -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