scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-10 22:41:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 129

refactor(native): move chat eligibility checks into Rust

Browse source
This commit is contained in:
Yujong Lee 2026-09-05 12:59:51 -07:00 • committed by yujonglee
parent 12c02ddb04
commit a278ee0ef1
14 changed files with 182 additions and 75 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
litellm-rust/crates/core/src/chat_completions/mod.rs
View file

@ -45,6 +45,7 @@ pub fn chat_completions_decline_reason(
custom_llm_provider: Option<&str>,
messages: Value,
optional_params: &Map<String, Value>,
context: &LiteLlmRequestContext,
) -> Option<&'static str> {
let Ok((_, config)) = resolve_provider_config(model, custom_llm_provider) else {
return Some("provider is not on the rust chat completions path");
@ -56,7 +57,7 @@ pub fn chat_completions_decline_reason(
return Some("empty message list");
}
config
.unsupported_reason(&messages, optional_params)
.unsupported_reason(&messages, optional_params, context)
.map(|reason| reason.0)
}

4
litellm-rust/crates/core/src/chat_completions/prepare.rs
View file

@ -42,7 +42,7 @@ pub(super) fn parse_messages(messages: Value) -> Result<Vec<ChatMessage>, Error>
pub(super) fn resolve_request(
request: ChatCompletionsRequest<'_>,
options: RequestOptions,
_context: &LiteLlmRequestContext,
context: &LiteLlmRequestContext,
) -> Result<ResolvedChatCompletionsRequest, Error> {
let (model, config) =
resolve_provider_config(request.model, options.custom_llm_provider.as_deref())
@ -52,7 +52,7 @@ pub(super) fn resolve_request(
if messages.is_empty() {
return Err(Error::Declined("empty message list"));
}
if let Some(reason) = config.unsupported_reason(&messages, &request.optional_params) {
if let Some(reason) = config.unsupported_reason(&messages, &request.optional_params, context) {
return Err(Error::Declined(reason.0));
}
Ok(ResolvedChatCompletionsRequest {

84
litellm-rust/crates/core/src/chat_completions/tests.rs
View file

@ -513,7 +513,13 @@ fn decline_reason(
Value::Object(map) => map,
other => panic!("params must be an object, got {other}"),
};
super::chat_completions_decline_reason(model, provider, messages, &params)
super::chat_completions_decline_reason(
model,
provider,
messages,
&params,
&LiteLlmRequestContext::default(),
)
}
#[test]
@ -877,3 +883,79 @@ mod round_trip {
));
}
}
#[test]
fn preflight_and_execution_share_provider_metadata_eligibility() {
let messages = json!([{"role": "user", "content": "hi"}]);
let cases = [
("anthropic", json!({"user_id": "u-123"}), vec![], true),
("anthropic", json!({"user_id": null}), vec![], false),
("anthropic", json!({"trace_id": "t-1"}), vec![], false),
(
"anthropic",
json!({}),
vec!["user_api_key_team_id".into()],
false,
),
#[cfg(feature = "bedrock-auth")]
("bedrock", json!({"user_id": "u-123"}), vec![], false),
#[cfg(feature = "bedrock-auth")]
(
"bedrock",
json!({}),
vec!["user_api_key_team_id".into()],
true,
),
];
for (provider, metadata, request_metadata_fields, expected_decline) in cases {
let context = LiteLlmRequestContext {
metadata: metadata.as_object().cloned(),
request_metadata_fields,
..Default::default()
};
let params = Map::new();
let preflight = super::chat_completions_decline_reason(
"claude-sonnet-4-5",
Some(provider),
messages.clone(),
&params,
&context,
);
let execution = resolve_request(
ChatCompletionsRequest {
model: "claude-sonnet-4-5",
messages: messages.clone(),
optional_params: params,
options: RequestOptions {
custom_llm_provider: Some(provider.into()),
..Default::default()
},
},
&context,
);
assert_eq!(
preflight.is_some(),
expected_decline,
"{provider} preflight"
);
assert_eq!(execution.is_err(), expected_decline, "{provider} execution");
}
}
#[test]
fn anthropic_preflight_preserves_litellm_metadata_scope() {
let context = LiteLlmRequestContext {
litellm_metadata: json!({"user_id": "u-123"}).as_object().cloned(),
..Default::default()
};
assert_eq!(
super::chat_completions_decline_reason(
"claude-sonnet-4-5",
Some("anthropic"),
json!([{"role": "user", "content": "hi"}]),
&Map::new(),
&context,
),
None,
);
}

2
litellm-rust/crates/core/src/chat_completions/transformation.rs
View file

@ -1,4 +1,5 @@
use crate::Error;
use crate::request_context::LiteLlmRequestContext;
use serde_json::{Map, Value};
use super::types::{
@ -75,6 +76,7 @@ pub trait ChatCompletionsProviderConfig: Sync {
&self,
messages: &[ChatMessage],
optional_params: &Map<String, Value>,
_context: &LiteLlmRequestContext,
) -> Option<Unsupported> {
unsupported_param(
self.supported_openai_params(),

7
litellm-rust/crates/core/src/providers/anthropic/chat_completions/tests.rs
View file

@ -1,5 +1,6 @@
use super::*;
use crate::Error;
use crate::request_context::LiteLlmRequestContext;
use serde_json::json;
fn messages(value: Value) -> Vec<ChatMessage> {
@ -26,7 +27,11 @@ fn transform_response(body: Value) -> Result<ChatCompletionsResponse, Error> {
}
fn reason(msgs: Value, opts: Value) -> Option<Unsupported> {
ANTHROPIC_CHAT_COMPLETIONS_CONFIG.unsupported_reason(&messages(msgs), &params(opts))
ANTHROPIC_CHAT_COMPLETIONS_CONFIG.unsupported_reason(
&messages(msgs),
&params(opts),
&LiteLlmRequestContext::default(),
)
}
#[test]

10
litellm-rust/crates/core/src/providers/anthropic/chat_completions/transformation.rs
View file

@ -1,3 +1,4 @@
use crate::request_context::LiteLlmRequestContext;
use serde_json::{Map, Value, json};
use crate::chat_completions::conversation::{Conversation, build_conversation};
@ -126,9 +127,18 @@ impl ChatCompletionsProviderConfig for AnthropicChatCompletionsConfig {
&self,
messages: &[ChatMessage],
optional_params: &Map<String, Value>,
context: &LiteLlmRequestContext,
) -> Option<Unsupported> {
unsupported_param(self.supported_openai_params(), &[], optional_params)
.or_else(|| messages.iter().find_map(unsupported_message))
.or_else(|| {
context
.metadata
.as_ref()
.and_then(|metadata| metadata.get("user_id"))
.is_some_and(|value| !value.is_null())
.then_some(Unsupported("LiteLLM user metadata"))
})
// Anthropic rejects a request whose first turn is not a user turn.
// Python only repairs that under `litellm.modify_params`, which the
// core cannot observe, so decline instead of guessing.

7
litellm-rust/crates/core/src/providers/bedrock/chat_completions/tests.rs
View file

@ -1,5 +1,6 @@
use super::*;
use crate::Error;
use crate::request_context::LiteLlmRequestContext;
use serde_json::json;
fn messages(value: Value) -> Vec<ChatMessage> {
@ -32,7 +33,11 @@ fn transform_response(body: Value) -> Result<ChatCompletionsResponse, Error> {
}
fn reason(msgs: Value, opts: Value) -> Option<Unsupported> {
BEDROCK_CHAT_COMPLETIONS_CONFIG.unsupported_reason(&messages(msgs), &params(opts))
BEDROCK_CHAT_COMPLETIONS_CONFIG.unsupported_reason(
&messages(msgs),
&params(opts),
&LiteLlmRequestContext::default(),
)
}
#[test]

6
litellm-rust/crates/core/src/providers/bedrock/chat_completions/transformation.rs
View file

@ -1,3 +1,4 @@
use crate::request_context::LiteLlmRequestContext;
use serde_json::{Map, Value, json};
use crate::chat_completions::conversation::{Conversation, TurnRole, build_conversation};
@ -176,6 +177,7 @@ impl ChatCompletionsProviderConfig for BedrockChatCompletionsConfig {
&self,
messages: &[ChatMessage],
optional_params: &Map<String, Value>,
context: &LiteLlmRequestContext,
) -> Option<Unsupported> {
unsupported_param(
self.supported_openai_params(),
@ -183,6 +185,10 @@ impl ChatCompletionsProviderConfig for BedrockChatCompletionsConfig {
optional_params,
)
.or_else(|| messages.iter().find_map(unsupported_message))
.or_else(|| {
(!context.request_metadata_fields.is_empty())
.then_some(Unsupported("LiteLLM request metadata forwarding"))
})
// Python's Converse translation drops blank text blocks instead of
// substituting the placeholder the shared conversation builder
// applies, so decline blank text rather than diverge.

17
litellm-rust/crates/python-bridge/src/routes/chat_completions.rs
View file

@ -5,6 +5,7 @@ use litellm_core::chat_completions::chat_completions as run_route;
use litellm_core::chat_completions::chat_completions_decline_reason;
use litellm_core::chat_completions::types::{ChatCompletionsRequest, ChatCompletionsResponse};
use litellm_core::request_context::LiteLlmRequestContext;
use litellm_core::request_options::RequestOptions;
use pyo3::prelude::*;
use serde_json::{Map, Value};
use std::future::Future;
@ -40,13 +41,25 @@ fn prepare_chat_completions(
}
#[pyfunction]
#[pyo3(signature = (model, messages, optional_params=None, custom_llm_provider=None))]
#[pyo3(signature = (model, messages, optional_params=None, custom_llm_provider=None, *, options, context))]
fn chat_completions_decline(
model: String,
#[pyo3(from_py_with = litellm_python_interop::from_py)] messages: Value,
#[pyo3(from_py_with = litellm_python_interop::from_py)] optional_params: Option<Value>,
custom_llm_provider: Option<String>,
options: NativeRequestOptions,
context: NativeRequestContext,
) -> PyResult<Option<String>> {
if !matches!(
custom_llm_provider.as_deref(),
Some("anthropic" | "bedrock")
) {
return Ok(Some(
"provider is not on the rust chat completions path".into(),
));
}
let context: LiteLlmRequestContext = context.into();
let options: RequestOptions = options.into();
let optional_params = match optional_params {
None | Some(Value::Null) => Map::new(),
Some(Value::Object(params)) => params,
@ -61,6 +74,8 @@ fn chat_completions_decline(
custom_llm_provider.as_deref(),
messages,
&optional_params,
&options,
&context,
)
.map(str::to_string))
}

77
litellm/rust_bridge/chat_completions.py
View file

@ -10,16 +10,15 @@ from __future__ import annotations
import json
from collections.abc import Mapping, Sequence
from dataclasses import replace
from typing import TYPE_CHECKING, Final, Protocol
import httpx
from pydantic import TypeAdapter, ValidationError
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import (
convert_to_model_response_object,
)
from litellm.llms.bedrock.request_metadata import bedrock_request_metadata_is_owned
from litellm.llms.bedrock.request_metadata import get_bedrock_request_metadata_fields
from litellm.rust_bridge.bindings import UNCHANGED, NativeBinding, Unchanged
from litellm.rust_bridge.configuration import rust_enabled
from litellm.rust_bridge.protocols import (
@ -35,6 +34,8 @@ from litellm.rust_bridge.request import (
NativeRequestContext,
NativeRequestOptions,
PreparedNativeCall,
anthropic_options,
bedrock_options,
call_native,
with_capabilities,
)
@ -45,14 +46,6 @@ from litellm.types.utils import ModelResponse
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
# Providers whose `/chat/completions` deployments the Rust core can serve. A
# provider outside this set never reaches the bridge.
RUST_CHAT_COMPLETIONS_PROVIDERS: Final = frozenset({"anthropic", "bedrock"})
# `litellm_params` values are `object`, so validate the one this module reads
# rather than narrowing an unparameterized `Mapping` and typing the result Any.
_LITELLM_METADATA_ADAPTER: Final = TypeAdapter(Mapping[str, object])
RUST_RESPONSE_HEADER: Final = "x-litellm-rust"
@ -128,42 +121,21 @@ def set_rust_chat_completions(
_CHAT_PREFLIGHT.override(decline)
def _anthropic_user_id_reaches_the_body(litellm_params: Mapping[str, object] | None) -> bool:
metadata: Final = litellm_params.get("metadata") if litellm_params is not None else None
try:
entries: Final = _LITELLM_METADATA_ADAPTER.validate_python(metadata)
except ValidationError:
return False
return entries.get("user_id") is not None
def _litellm_metadata_reaches_the_provider(
custom_llm_provider: str | None, litellm_params: Mapping[str, object] | None
) -> bool:
"""Whether the Python transform would promote proxy-owned attribution into the
provider request, below this gate and inside the function the Rust route replaces.
`AnthropicConfig.transform_request` promotes a valid `metadata["user_id"]`
into the Messages body, so the core never sees the key and would send the
request to Anthropic with the abuse-detection attribution missing.
`AmazonConverseConfig` resolves proxy-owned `requestMetadata` onto the
Converse body whenever the operator armed `bedrock_request_metadata_fields`.
Owning that field also means evicting a caller-supplied one, which the core
cannot do either, so ownership alone is the condition rather than whether
anything resolved.
Deliberately a superset of Python's condition in both cases: declining a
request Python would not have attributed anyway costs only the Rust path,
while missing one loses the attribution silently.
"""
match custom_llm_provider:
case "anthropic":
return _anthropic_user_id_reaches_the_body(litellm_params)
case "bedrock":
return bedrock_request_metadata_is_owned()
case _:
return False
def _provider_eligibility_options(
provider: str | None,
litellm_params: Mapping[str, object] | None,
optional_params: Mapping[str, object],
) -> NativeRequestOptions:
bedrock: Final = (
replace(
bedrock_options(optional_params),
request_metadata_fields=get_bedrock_request_metadata_fields(),
)
if provider == "bedrock"
else None
)
anthropic: Final = anthropic_options(litellm_params) if provider == "anthropic" else None
return NativeRequestOptions(custom_llm_provider=provider, bedrock=bedrock, anthropic=anthropic)
def rust_chat_completions_accepts(
@ -182,13 +154,6 @@ def rust_chat_completions_accepts(
capability gate answers the second half; it resolves no credentials and
performs no I/O.
"""
if custom_llm_provider not in RUST_CHAT_COMPLETIONS_PROVIDERS:
return False
if stream:
return False
if _litellm_metadata_reaches_the_provider(custom_llm_provider, litellm_params):
verbose_logger.debug("Rust chat completions declined (litellm metadata user_id); using the Python path")
return False
if not rust_enabled():
return False
decline: Final = _CHAT_PREFLIGHT.load()
@ -200,6 +165,10 @@ def rust_chat_completions_accepts(
messages=messages,
optional_params=optional_params,
custom_llm_provider=custom_llm_provider,
options=_provider_eligibility_options(custom_llm_provider, litellm_params, optional_params),
context=NativeRequestContext(
capabilities=NativeRequestCapabilities(stream=bool(stream))
),
)
except Exception as error: # noqa: BLE001 # capability checks perform no provider I/O
verbose_logger.debug("Native chat acceptance check failed: %s", error)

3
litellm/rust_bridge/protocols.py
View file

@ -31,6 +31,9 @@ class RustChatCompletionsDecline(Protocol):
messages: Sequence[object],
optional_params: Mapping[str, object] | None,
custom_llm_provider: str | None,
*,
options: NativeRequestOptions,
context: NativeRequestContext,
) -> str | None: ...

2
tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py
View file

@ -2427,7 +2427,7 @@ class TestRustChatCompletionsHook:
)
except Exception:
pass
assert seen["gate"] == []
assert seen["call"] == []
def test_pre_call_logging_fires_exactly_once_on_the_rust_path(self):
from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion

5
tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py
View file

@ -181,11 +181,14 @@ def test_without_the_opt_in_the_core_is_never_consulted(monkeypatch):
def test_streaming_stays_on_the_python_path():
seen = _inject()
bridge.set_rust_chat_completions(
decline=pytest.importorskip("litellm.rust_bridge._native").chat_completions_decline
)
try:
_run(optional_params={"maxTokens": 16, "stream": True})
except Exception:
pass
assert seen["gate"] == []
assert seen["call"] == []
def test_a_declined_request_never_reaches_the_native_call():

30
tests/test_litellm/rust_bridge/test_chat_completions.py
View file

@ -176,25 +176,17 @@ class TestGate:
def test_declines_streaming_and_providers_off_the_path(self, monkeypatch):
monkeypatch.setenv("LITELLM_RUST", "1")
gate = _RecordingDecline()
gate = pytest.importorskip("litellm.rust_bridge._native").chat_completions_decline
bridge.set_rust_chat_completions(decline=gate)
assert _accepts(stream=True) is False
assert _accepts(custom_llm_provider="openai") is False
assert _accepts(custom_llm_provider=None) is False
assert gate.calls == []
def test_declines_an_anthropic_request_carrying_a_litellm_metadata_user_id(self, monkeypatch):
"""`AnthropicConfig.transform_request` copies a valid `user_id` into the Messages body.
It does that inside the function the Rust route replaces, and the core is
handed `optional_params` only, so accepting here would send the request
to Anthropic with the abuse-detection attribution silently missing.
"""
monkeypatch.setenv("LITELLM_RUST", "1")
gate = _RecordingDecline()
gate = pytest.importorskip("litellm.rust_bridge._native").chat_completions_decline
bridge.set_rust_chat_completions(decline=gate)
assert _accepts(litellm_params={"metadata": {"user_id": "u-123"}}) is False
assert gate.calls == [], "the core must not be consulted for a request it cannot see the key of"
# Bedrock's Converse transform reads no `user_id`, and an Anthropic request
# whose metadata carries none is one Python would not attribute either.
@ -202,6 +194,7 @@ class TestGate:
_accepts(
custom_llm_provider="bedrock",
model="bedrock/us-east-1/anthropic.claude-v2",
optional_params={"maxTokens": 16},
litellm_params={"metadata": {"user_id": "u-123"}},
)
is True
@ -209,6 +202,19 @@ class TestGate:
assert _accepts(litellm_params={"metadata": {"trace_id": "t-1"}}) is True
assert _accepts(litellm_params={"metadata": {"user_id": None}}) is True
assert _accepts(litellm_params={"metadata": None}) is True
assert _accepts(litellm_params={"litellm_metadata": {"user_id": "u-123"}}) is True
assert _accepts(litellm_params={"metadata": "invalid"}) is True
assert _accepts(litellm_params={"metadata": {"trace": object()}}) is True
assert _accepts(litellm_params={"metadata": {"user_id": object()}}) is False
assert (
_accepts(
custom_llm_provider="bedrock",
model="bedrock/us-east-1/anthropic.claude-v2",
optional_params={"maxTokens": 16},
litellm_params={"metadata": {"user_id": object()}},
)
is True
)
def test_declines_a_bedrock_request_while_the_proxy_owns_request_metadata(self, monkeypatch):
"""`AmazonConverseConfig` resolves proxy-owned `requestMetadata` onto the
@ -217,16 +223,16 @@ class TestGate:
who armed `bedrock_request_metadata_fields` keeps the Python path.
"""
monkeypatch.setenv("LITELLM_RUST", "1")
gate = _RecordingDecline()
gate = pytest.importorskip("litellm.rust_bridge._native").chat_completions_decline
bridge.set_rust_chat_completions(decline=gate)
bedrock = {
"custom_llm_provider": "bedrock",
"model": "bedrock/us-east-1/anthropic.claude-v2",
"optional_params": {"maxTokens": 16},
}
monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", ["user_api_key_team_id"])
assert _accepts(**bedrock) is False
assert gate.calls == [], "the core must not be consulted for a field it cannot write"
assert _accepts() is True, "arming Bedrock attribution must not decline Anthropic"
monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", None)