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): serve anthropic messages on rust only (#51119)

Browse source 
* feat(messages): apply modify_params rewrites on rust

* feat(messages): mint anthropic workload identity tokens on rust

* feat(messages): apply prompt caching injection on rust

* feat(messages): serve the easy python-only detectors on rust

* feat(messages): retry rejected thinking signatures on rust

* feat(messages): serve response cache reads on rust

* feat(messages): retry rejected strict tools on rust

* refactor(messages): delete the python path for anthropic messages

Provider anthropic on /v1/messages is now RUST_REQUIRED, so it runs on Rust
without LITELLM_RUST or litellm.rust(True) and never falls back to Python.
A call the rules would send to Python (a MESSAGES_PYTHON_FEATURES feature or a
detector that raises) gets UnsupportedParamsError naming the feature, raised
from one place in litellm/messages/dispatch.py. Other providers keep the
Python fallback. A call whose arguments the handler signature rejects still
gets Python's TypeError.

Deleted Python code that only served anthropic Messages:
- workload identity in AnthropicMessagesConfig.validate/avalidate (subclasses
  were never eligible; the Rust host mints the token itself)
- the anthropic Foundry beta-row and affinity-base branch in
  async_anthropic_messages_handler, and the "forward" mode of the beta filter
  that only that branch used

The proxy MCP tools detector now uses the light tool precheck so a function
tool never imports the gateway handler, since a raising detector now refuses.

* feat(messages): forward message threads on rust

The Rust path for anthropic /v1/messages dropped the thread body field, so
a create or continue turn reached Anthropic without its thread or the
message-threads-2026-08-12 beta. It now matches the Python handler from
#50976:

- the host forwards thread with the other body params
- a non-null thread adds the message-threads beta
- a well-formed continue turn keeps advisor blocks in history and derives
  no compaction, per-turn-control or tool-change beta from its messages
- the recorded httpx_response carries the request that was sent, for JSON
  and streamed calls, so the proxy's thread ownership hook resends the
  create turn's betas on a continue as it does for Python

Upstream error.details already reached the exception message unchanged, so
thread_not_found relays and keeps the deployment out of cooldown on both
paths. Parity cases cover create, continue, streaming, sync, retries,
thread_not_found and the ownership hook round trip.

* feat(messages): serve mock responses on rust

* feat(messages): send caller clients and unsupported tls through python on rust

* feat(messages): run pre-request hooks on rust

* feat(messages): run agentic loop hooks on rust

* feat(messages): run the safeguards classifier on rust

* feat(messages): run registered interceptors on rust

* fix(messages): match python stream cache gate, logging and interceptor order on rust

Store a streamed Messages response in the cache only when Python would: a
message_stop event line, no provider error frame and at least one frame with
output. Leave dict results untouched by mark_rust_response so hidden params
never reach result dicts, logging objects or hook digests. Ask interceptors
whether they handle the call after the cache lookup, not at decode time, so a
cache hit never calls can_handle. Emit the track_llm_api_timing service event
for the provider exchange. Bill and log a stream superseded by a router
fallback as the failure of its error frame, with the partial usage warning.
Hand litellm_logging_obj to deployment pre-call hooks as @client does, so the
autorouter baseline cache sees retried requests.

* fix(messages): forward python-only json values to anthropic as python does

Accept any max_tokens, sampling scalar and messages collection that
Python main accepts, and carry values serde_json cannot hold (ints beyond
64 bits, NaN and infinities, strings with lone surrogates) as validated
raw JSON leaves written back verbatim on the wire, so the upstream body
matches Python's json.dumps

* fix(messages): match python's body reads, error bodies, encoded heads and tls warnings on rust

* fix(messages): log python's plain provider header names on rust streams

* fix(messages): match python's request params, logged params and validation on rust

Read a non-bool stream by Python truthiness and send it as given, log the
caller's stream flag as Python's handler does, raise the raw Python errors
main raises for malformed messages, tools, model_info and hook-set
litellm_params or providers, read the tool custom-field drop from
drops_unsupported_params, log the caller's nested fields that
additional_drop_params strips and the Python values of raw JSON leaves,
keep a missing stop_reason absent, and map provider-config rejections
(missing max_tokens, unknown reasoning effort) as Python's HTTP handler
maps them, with Python's text, effort order, model and provider

* fix(messages): match python's tls curve names, redirects, auth token, header encoding and timeouts on rust

* fix(messages): match python on hooked streams, empty cache writes and mcp cache hits

* fix(messages): join already buffered stream reads and trim the env api key as python does on rust

* fix(messages): deliver stream bytes read before a cut, keep wide reply numbers and follow python's redirect limit on rust

A stream cut by a transport error now hands the caller the bytes already
buffered before raising, as aiohttp does for a caller that reads them before
the event loop sees the cut, so the frames and their usage reach the caller and
the failure log. Provider replies keep integers wider than 64 bits exactly and
read numbers that overflow a float as infinity, through the response, the
response cache and the logged result. The provider client follows up to 20
redirects as httpx does, and running out raises Python's text. A
provider_specific_header value the ascii codec cannot encode fails before the
send like a caller extra header. The deployment success hook gets the stream as
AnthropicMessagesStreamingResponse, as Python's handler hands it.

* feat(messages): serve proxy mcp tools on rust

* test(messages): cover hook, agentic and interceptor edge cases in the parity suite

* fix(messages): keep provider raw-json marked strings, drop non-string tool_result texts and key clamped streams as python does on rust

* fix(messages): match python's api key encoding, malformed api_base and string provider header errors on rust

A header value Python's client cannot encode as ASCII now fails before the send
whichever header carries it, so a caller api_key with a non-ASCII character
raises Python's InternalServerError instead of a raw UnicodeEncodeError after
the reply. A caller api_base without an http(s) scheme, or with whitespace,
control or non-ASCII characters, is sent through Python's client at the URL
Python builds, so its relative-URL and invalid-URL errors match, and a base
urlsplit cannot split raises Python's bare ValueError before the send. A
provider_specific_header of the wrong shape is scoped through Python's
ProviderSpecificHeaderUtils first, so a plain string raises Python's
AttributeError.

* fix(messages): keep integer temperature and top_p digits a float cannot hold and raise python's typeerror for a string max_tokens while capping the thinking budget on rust

temperature and top_p now deserialize through recognized::exact_float, which keeps an integer
where float(n) != n (beyond 2**53) as an unrecognized value with its exact digits, so it is sent
upstream and logged as the integer python forwards, and comparisons still go through scalar.rs

cap_thinking_budget_to_max_tokens compares max_tokens <= 1024 first, so a max_tokens python
cannot order against an int (a string, list or dict) raises TypeError "'<=' not supported between
instances of 'str' and 'int'" in both the reasoning effort translation and the adaptive effort
reshape. rust now returns ErrorDetail::UnorderableTypes there, the bridge raises it as a python
TypeError marked as a late rejection, and map_failure hands it to the caller unmapped

* test(messages): format the non-string tool_result parity helper

* fix(messages): flag and log the call before python's pre-send raises on rust

The bridge now hands settings() and map_failure() whether the call is async.
A provider_specific_header of the wrong shape sets the aanthropic_messages
flag before re-raising, as Python's handler has already flagged the call. A
base urlsplit cannot split also logs the resolved model, provider and optional
params first, as Python raises it only after its logger is updated. The
TypeError raised while capping the thinking budget for a string max_tokens
sets the flag too

* feat(messages): serve the web search short-circuit override on rust through the intercept host call

* fix(messages): give the web search entry kwargs its mutable-ok reason

* test(messages): drop the body hook gap the rust short-circuit closes

* fix(messages): name the partially unknown handler types on the short-circuit host call
This commit is contained in:
Nate Armstrong 2026-11-11 16:07:20 -08:00 • committed by GitHub
parent a0e6cafbd8
commit 59538ff1d9
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
175 changed files with 19176 additions and 2747 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
litellm-rust/Cargo.lock generated
View file

@ -4013,6 +4013,7 @@ dependencies = [
"litellm-coroutine",
"litellm-secret-value",
"rstest",
"serde",
"serde_json",
"thiserror 2.0.19",
"tokio",
@ -4155,9 +4156,12 @@ name = "litellm-inference-messages"
version = "0.1.0"
dependencies = [
"bytes",
"flate2",
"futures-util",
"http 1.4.2",
"litellm-auth",
"litellm-auth-aws",
"litellm-cache",
"litellm-cache-memory",
"litellm-cache-response",
"litellm-core-utils",
@ -4178,6 +4182,7 @@ dependencies = [
"rstest",
"serde",
"serde_json",
"thiserror 2.0.19",
"tokio",
"tracing",
"wiremock",
@ -4292,6 +4297,7 @@ dependencies = [
"base64 0.22.1",
"bytes",
"data-url",
"fancy-regex 0.19.2",
"futures-util",
"litellm-auth",
"litellm-auth-aws",
@ -4307,6 +4313,7 @@ dependencies = [
"litellm-router-types",
"litellm-secret-value",
"litellm-secrets",
"litellm-tracing",
"reqwest",
"rstest",
"serde",

4
litellm-rust/crates/cache-response/src/service.rs
View file

@ -27,6 +27,10 @@ impl Default for ResponseCacheConfig {
pub trait ResponseCacheService: Send + Sync {
fn config(&self) -> &ResponseCacheConfig;
fn keyed_by_caller(&self) -> bool {
false
}
fn lookup<'a>(
&'a self,
request: &'a ResponseCacheRequest,

28
litellm-rust/crates/callbacks-legacy-python/python_contract.json
View file

@ -11,7 +11,8 @@
"logger",
"kwargs",
"start_time",
"end_time"
"end_time",
"cache_hit"
],
"update_logging": [
"logger",
@ -41,13 +42,24 @@
"logger",
"status",
"headers",
"body"
"body",
"sent"
],
"retried_request": [
"logger",
"logged_body",
"body",
"reasons"
],
"record_api_window": [
"logger",
"start",
"end"
],
"times_provider_exchange": [
"client",
"timeout"
],
"defers_async_logging": [
"logger"
],
@ -89,6 +101,12 @@
"start",
"end"
],
"cache_hit_success": [
"logger",
"response",
"start",
"end"
],
"enqueue_logging": [
"coroutine"
],
@ -113,14 +131,16 @@
],
"stream_opened": [
"logger",
"head"
"head",
"keeps_requested"
],
"stamp_stream_head": [
"head",
"logger",
"kwargs",
"start_time",
"end_time"
"end_time",
"cache_hit"
],
"stream_success": [
"logger",

392
litellm-rust/crates/callbacks-legacy-python/src/adapter.rs
View file

@ -6,10 +6,10 @@ use crate::LoggingOperation;
use litellm_host_python::PythonOwned;
use litellm_host::{
interceptors::{RawResponse, RequestContext, WireRequest},
interceptors::{RawResponse, RequestContext, RetriedRequest, RetryReason, WireRequest},
lifecycle::{FailureOrigin, Timing, epoch_seconds},
};
use litellm_host_python::{Error as HostError, HookStep, from_py, to_py};
use litellm_host_python::{Error as HostError, HookStep, from_py, json_from_py, json_to_py, to_py};
use pyo3::{
exceptions::{PyBaseException, PyException},
gc::{PyTraverseError, PyVisit},
@ -19,7 +19,7 @@ use pyo3::{
use serde_json::Value;
use crate::{
DeploymentHooks, LegacyCallbacks, PublicCall, PythonLogger,
DeploymentHooks, LegacyCallbacks, PublicCall, PythonLogger, SentRequest,
deferred::{PendingLogging, PendingSuccess},
finalize, is_internal_call,
python::{Logging, Streaming},
@ -41,6 +41,7 @@ struct DeliveredStream {
struct LoggedRequest {
body: Py<PyDict>,
headers: Py<PyDict>,
url: String,
context: RequestContext,
}
@ -59,6 +60,16 @@ pub struct LegacyLogging {
cache_key: Option<String>,
/// When the provider request went out, until its window is recorded (Messages only).
api_started: Option<f64>,
host_answered: bool,
}
fn retry_reason(reason: &RetryReason) -> Value {
match reason {
RetryReason::ThinkingSignature => serde_json::json!({"kind": "thinking_signature"}),
RetryReason::StrictTools { dropped } => {
serde_json::json!({"kind": "strict_tools", "dropped": dropped})
}
}
}
fn datetime(py: Python<'_>, epoch_seconds: f64) -> PyResult<Py<PyAny>> {
@ -90,6 +101,7 @@ impl LegacyLogging {
internal: false,
cache_key: None,
api_started: None,
host_answered: false,
}
}
@ -171,6 +183,14 @@ impl LegacyLogging {
Ok(context)
}
fn times_provider_exchange(&self, py: Python<'_>) -> PyResult<bool> {
let client = self.call.lookup(py, "client")?;
let timeout = self.call.lookup(py, "timeout")?;
Logging::TimesProviderExchange
.call(py, (client, timeout))?
.extract()
}
/// Python times the provider exchange inside the async `AsyncHTTPHandler.post`; its sync
/// client is untimed.
fn close_api_window(&mut self, py: Python<'_>, end: f64) -> PyResult<()> {
@ -264,6 +284,7 @@ impl LegacyLogging {
self.call.kwargs(),
&self.start,
&self.end,
self.messages_cache_hit(),
)?;
self.response
.as_ref()
@ -271,6 +292,60 @@ impl LegacyLogging {
.ok_or_else(|| HostError::MissingState.into())
}
fn messages_cache_hit(&self) -> bool {
self.operation == LoggingOperation::Messages && self.cache_key.is_some()
}
fn flag_unsent_messages_call(&self, py: Python<'_>) -> PyResult<()> {
if self.operation != LoggingOperation::Messages
|| self.request.is_some()
|| self.cache_key.is_some()
|| self.host_answered
{
return Ok(());
}
let details = self.logger()?.object(py).getattr("model_call_details")?;
details
.call_method1("setdefault", ("litellm_params", PyDict::new(py)))?
.set_item("aanthropic_messages", self.asynchronous)
}
fn unstreamed_stream_request(&self, py: Python<'_>) -> PyResult<bool> {
if self.operation != LoggingOperation::Messages
|| self.stream.is_some()
|| self.messages_cache_hit()
{
return Ok(false);
}
let Some(stream) = self.call.kwargs().bind(py).get_item("stream")? else {
return Ok(false);
};
if !stream.is_exact_instance_of::<pyo3::types::PyBool>() || !stream.is_truthy()? {
return Ok(false);
}
let logger = self.logger()?.object(py);
logger.setattr("stream", true)?;
logger
.getattr("model_call_details")?
.set_item("stream", true)?;
Ok(true)
}
fn dispatch_cache_hit(&self, py: Python<'_>) -> PyResult<()> {
let logger = self.logger()?;
let logged = Logging::CacheHitSuccess.call(
py,
(logger.object(py), &self.response, &self.start, &self.end),
);
match logged {
Err(error) if error.is_instance_of::<PyException>(py) => {
error.write_unraisable(py, Some(logger.object(py)));
Ok(())
}
result => result.map(|_| ()),
}
}
fn dispatch_success(&self, py: Python<'_>) -> PyResult<()> {
match self.try_dispatch_success(py) {
Err(error) if error.is_instance_of::<PyException>(py) => {
@ -489,6 +564,10 @@ impl LegacyLogging {
self.logger = Some(result.logger()?);
self.call.set_kwargs(result.kwargs()?);
if self.runs_deployment_hooks() {
self.call
.kwargs()
.bind(py)
.set_item("litellm_logging_obj", self.logger()?.object(py))?;
return Ok(HookStep::Await(
DeploymentHooks::before_call(py, self.call.kwargs(), self.call_type())?,
Self::resume_begin,
@ -512,9 +591,13 @@ impl LegacyLogging {
context,
self.logged_request_params(context).as_ref(),
)?;
let body = to_py(py, &wire.body)?
.into_bound(py)
.cast_into::<PyDict>()?;
let lossless = self.operation == LoggingOperation::Messages;
let body = if lossless {
json_to_py(py, &wire.body)?
} else {
to_py(py, &wire.body)?.into_bound(py)
}
.cast_into::<PyDict>()?;
for (name, sent) in wire.body.as_object().into_iter().flatten() {
if let Some(value) = self.caller_value(py, name, sent)? {
body.set_item(name, value)?;
@ -534,6 +617,7 @@ impl LegacyLogging {
self.request = Some(LoggedRequest {
body: body.clone().unbind(),
headers: headers.clone().unbind(),
url: wire.url.clone(),
context: context.clone(),
});
let input = self.logged_input(py, &body)?;
@ -545,7 +629,10 @@ impl LegacyLogging {
&headers,
&wire.url,
)?;
if self.operation == LoggingOperation::Messages && self.asynchronous {
if self.operation == LoggingOperation::Messages
&& self.asynchronous
&& self.times_provider_exchange(py)?
{
self.api_started = Some(epoch_seconds());
}
let headers = headers
@ -553,7 +640,11 @@ impl LegacyLogging {
.map(|(name, value)| Ok((name.extract::<String>()?, value.extract::<String>()?)))
.collect::<PyResult<Vec<_>>>()?;
Ok(HookStep::Ready(Box::new(WireRequest {
body: from_py(&body)?,
body: if lossless {
json_from_py(body.as_any())?
} else {
from_py(&body)?
},
headers,
..*wire
})))
@ -567,7 +658,11 @@ impl LegacyLogging {
) -> PyResult<HookStep<Self, Py<PyAny>>> {
self.end = Some(datetime(py, timing.end_time)?);
self.response = Some(response);
if self.runs_deployment_hooks() {
self.flag_unsent_messages_call(py)?;
if self.unstreamed_stream_request(py)? {
return self.finalize(py);
}
if self.runs_deployment_hooks() && !self.messages_cache_hit() {
return Ok(HookStep::Await(
DeploymentHooks::after_success(
py,
@ -588,7 +683,21 @@ impl LegacyLogging {
) -> PyResult<HookStep<Self, ()>> {
use litellm_host::interceptors::{Cost, ResultSource};
self.cache_key = match &facts.source {
ResultSource::Provider => None,
ResultSource::Cache { key } => Some(key.clone()),
ResultSource::Host => {
self.host_answered = true;
return Ok(HookStep::Ready(()));
}
};
let logger = self.logger()?.object(py);
if self.messages_cache_hit() {
logger
.getattr("model_call_details")?
.set_item("cache_hit", true)?;
return Ok(HookStep::Ready(()));
}
let params = logger
.getattr("litellm_params")?
.cast_into::<PyDict>()?
@ -616,10 +725,6 @@ impl LegacyLogging {
};
logger.setattr("_provider_reported_cost", reported_cost)?;
details.call_method1("pop", ("response_cost", py.None()))?;
self.cache_key = match &facts.source {
ResultSource::Provider => None,
ResultSource::Cache { key } => Some(key.clone()),
};
if self.cache_key.is_some() {
details.set_item("cache_hit", true)?;
} else {
@ -644,8 +749,13 @@ impl LegacyLogging {
// response, and logging replays it through `AnthropicConfig.transform_response`.
self.close_api_window(py, epoch_seconds())?;
if let Some(head) = &raw.head {
let sent = self.request.as_ref().map(|request| SentRequest {
url: &request.url,
headers: &request.headers,
body: &request.body,
});
self.logger()?
.record_response(py, head.status, &head.headers, &raw.body)?;
.record_response(py, head.status, &head.headers, &raw.body, sent)?;
}
return Ok(HookStep::Ready(()));
}
@ -664,6 +774,40 @@ impl LegacyLogging {
Ok(HookStep::Ready(()))
}
pub(crate) fn request_retried(
&mut self,
py: Python<'_>,
retry: &RetriedRequest,
) -> PyResult<HookStep<Self, ()>> {
let Some(request) = &self.request else {
return Ok(HookStep::Ready(()));
};
let reasons: Vec<Value> = retry.reasons.iter().map(retry_reason).collect();
let retried = Logging::RetriedRequest.call(
py,
(
self.logger()?.object(py),
request.body.bind(py),
to_py(py, &retry.body)?,
to_py(py, &reasons)?,
),
)?;
let now = epoch_seconds();
if self.api_started.is_some() {
self.close_api_window(py, now)?;
self.api_started = Some(now);
}
Ok(HookStep::Await(retried.unbind(), Self::resume_retried))
}
fn resume_retried(
&mut self,
_py: Python<'_>,
result: PyResult<Py<PyAny>>,
) -> PyResult<HookStep<Self, ()>> {
result.map(|_| HookStep::Ready(()))
}
pub(crate) fn succeeded(
&mut self,
py: Python<'_>,
@ -674,6 +818,8 @@ impl LegacyLogging {
self.response = Some(response.clone_ref(py));
match &self.stream {
Some(stream) => self.stream_success(py, stream, false)?,
None if self.messages_cache_hit() => self.dispatch_cache_hit(py)?,
None if self.unstreamed_stream_request(py)? => {}
None => self.dispatch_success(py)?,
}
Ok(HookStep::Ready(()))
@ -708,7 +854,11 @@ impl LegacyLogging {
if self.stream_billing().is_none() {
return Err(HostError::MissingState.into());
}
if let Some(key) = &self.cache_key {
if let Some(key) = self
.cache_key
.as_ref()
.filter(|_| !self.messages_cache_hit())
{
head.bind(py).set_item("cache_key", key)?;
head.bind(py).set_item("cache_hit", true)?;
}
@ -723,10 +873,18 @@ impl LegacyLogging {
self.call.kwargs(),
&self.start,
datetime(py, epoch_seconds())?,
self.messages_cache_hit(),
),
)?;
}
Streaming::Opened.call(py, (self.logger()?.object(py), head.bind(py)))?;
Streaming::Opened.call(
py,
(
self.logger()?.object(py),
head.bind(py),
self.operation == LoggingOperation::Messages,
),
)?;
self.stream = Some(DeliveredStream {
chunks: PyList::empty(py).unbind(),
first_chunk: None,
@ -747,7 +905,7 @@ impl LegacyLogging {
// A caller that goes away mid-stream cancels the call before it ends: Python's Messages
// stream iterator still bills the chunks it read, straight away rather than deferred.
fn bill_abandoned_stream(&mut self, py: Python<'_>) -> PyResult<()> {
if !self.asynchronous || self.end.is_some() {
if !self.asynchronous || self.end.is_some() || self.messages_cache_hit() {
return Ok(());
}
let Some(stream) = self.stream.take() else {
@ -918,6 +1076,24 @@ if asynchronous:
});
}
#[rstest::rstest]
fn the_deployment_pre_call_hook_receives_the_logger_as_client_injects_it() {
Python::initialize();
Python::attach(|py| {
let locals = namespace(
py,
c"
kwargs = {'logger': logger}
logger.hooks['pre'] = lambda value: seen.append(value.get('litellm_logging_obj')) or 'awaitable'
seen = []
",
);
let (_, step) = begin(py, &locals, true);
assert!(awaits_deployment_hook(&step));
run(py, &locals, c"assert seen == [logger]");
});
}
#[rstest::rstest]
fn kwargs_returned_by_the_pre_call_hook_are_what_the_call_prepares() {
Python::initialize();
@ -1038,6 +1214,82 @@ assert finalized is replacement
});
}
#[rstest]
#[case::unary(false)]
#[case::stream(true)]
fn a_host_answer_keeps_the_logger_as_set_up_and_a_streamed_one_is_finalized_unstreamed(
#[case] stream: bool,
) {
use litellm_host::interceptors::{Cost, ExecutionFacts, ProviderIdentity, ResultSource};
Python::initialize();
Python::attach(|py| {
let locals = namespace(
py,
c"
kwargs = {'logger': logger}
response = object()
logger.hooks = {'pre': lambda kwargs: kwargs}
",
);
local(&locals, "kwargs").set_item("stream", stream).unwrap();
let (mut logging, step) = begin(py, &locals, true);
logging.operation = crate::LoggingOperation::Messages;
resume(
&mut logging,
step,
py,
Ok(local(&locals, "kwargs").unbind()),
)
.unwrap();
let facts = ExecutionFacts {
provider: ProviderIdentity {
model: "claude-sonnet-4-5".into(),
provider: "anthropic".into(),
},
cost: Cost::Deferred,
source: ResultSource::Host,
};
assert!(matches!(
logging.result_ready(py, &facts).unwrap(),
HookStep::Ready(())
));
let response = local(&locals, "response").unbind();
let step = logging
.transform_response(py, response.clone_ref(py), TIMING)
.unwrap();
assert_eq!(awaits_deployment_hook(&step), !stream);
if !stream {
return;
}
let HookStep::Ready(returned) = step else {
panic!("a host-answered stream is returned as is");
};
assert!(returned.bind(py).is(response.bind(py)));
logging
.on_event(
py,
PythonCallEvent::Succeeded {
timing: TIMING,
response: &response,
},
)
.unwrap();
run(
py,
&locals,
c"
names = logger.names()
assert ('finalize', response) in logger.calls, names
assert 'update_logging' not in names, names
assert not [name for name in names if 'success' in name], names
assert logger.stream is True
assert logger.model_call_details['stream'] is True
",
);
});
}
#[rstest]
#[case::pre_call(false)]
#[case::post_call(true)]
@ -1129,7 +1381,9 @@ mod payload_tests {
use std::ffi::CStr;
use litellm_host::hooks::CallHooks;
use litellm_host::interceptors::{RawResponse, RequestContext, WireRequest};
use litellm_host::interceptors::{
RawResponse, RequestContext, RetriedRequest, RetryReason, WireRequest,
};
use litellm_host::lifecycle::ExecutionEvent;
use litellm_host_python::{HookStep, PythonCallEvent, PythonOwned, to_py};
use litellm_secret_value::SecretValue;
@ -1296,6 +1550,61 @@ assert ('custom_llm_provider' in logger.update['litellm_params']) is in_params,
json!({"type": "document_url", "document_url": source})
}
#[test]
fn a_retried_request_hands_the_logged_body_the_retried_body_and_the_reasons_to_logging() {
Python::initialize();
Python::attach(|py| {
let locals = namespace(py, PAYLOAD_LOGGER);
let mut logging = LegacyLogging {
logger: Some(PythonLogger::new(local(&locals, "logger").unbind())),
..legacy_call(py, &locals, false)
};
let context = RequestContext {
model: "model".into(),
custom_llm_provider: "provider".into(),
optional_params: json!({}),
secret_fields: Vec::new(),
api_key: None,
};
let wire = WireRequest {
url: "https://provider.invalid/v1/messages".into(),
headers: Vec::new(),
body: json!({"messages": [{"role": "user", "content": "hi"}], "thinking": {"type": "enabled"}}),
};
logging
.before_provider_request(py, Box::new(wire.clone()), &context)
.unwrap();
let retried = RetriedRequest {
reasons: vec![
RetryReason::ThinkingSignature,
RetryReason::StrictTools { dropped: 2 },
],
body: json!({"messages": [{"role": "user", "content": "hi"}]}),
};
let step = logging
.on_event(
py,
PythonCallEvent::Execution(ExecutionEvent::ProviderRequestRetried {
retry: retried,
}),
)
.unwrap();
assert!(matches!(step, HookStep::Await(_, _)));
run(
py,
&locals,
c"
[(logged, body, reasons)] = [value for name, value in logger.calls if name == 'retried_request']
assert logged is logger.pre['complete_input_dict'], logged
assert body == {'messages': [{'role': 'user', 'content': 'hi'}]}, body
assert reasons == [{'kind': 'thinking_signature'}, {'kind': 'strict_tools', 'dropped': 2}], reasons
",
);
});
}
fn before_provider_request(script: &CStr, body: Value) -> WireRequest {
before_send_with_secrets(script, json!({}), body, &[])
}
@ -2018,6 +2327,53 @@ mod terminal_tests {
.unwrap()
}
#[rstest]
#[case::requested_stream(
c"kwargs = {'logger': logger, 'stream': True}",
true,
&["finalize"]
)]
#[case::unstreamed(c"kwargs = {'logger': logger, 'stream': False}", false, &["success_hook"])]
#[case::stream_flag_not_true(c"kwargs = {'logger': logger, 'stream': 'true'}", false, &["success_hook"])]
fn a_whole_messages_result_for_a_requested_stream_is_finalized_without_success_logging(
#[case] script: &CStr,
#[case] handed_back: bool,
#[case] expected: &[&str],
) {
Python::initialize();
Python::attach(|py| {
let locals = namespace(py, c"response = object()");
run(py, &locals, script);
let mut logging = LegacyLogging {
operation: crate::LoggingOperation::Messages,
..logged(py, &locals, true)
};
let step = logging
.transform_public_response(py, local(&locals, "response").unbind(), TIMING)
.unwrap();
assert_eq!(matches!(step, HookStep::Ready(_)), handed_back);
if handed_back {
assert!(matches!(
succeed(py, &locals, &mut logging),
HookStep::Ready(())
));
}
let names: Vec<String> = local(&locals, "logger")
.call_method0("names")
.unwrap()
.extract()
.unwrap();
assert_eq!(names, expected);
if handed_back {
run(
py,
&locals,
c"assert logger.stream is True and logger.model_call_details['stream'] is True",
);
}
});
}
#[rstest]
#[case::sync_listened(false, c"", &["submit"])]
#[case::async_listened(

22
litellm-rust/crates/callbacks-legacy-python/src/callbacks.rs
View file

@ -3,13 +3,19 @@
//! duplication. All of it expires with the legacy callback contract.
use litellm_host::interceptors::{RequestContext, WireRequest};
use litellm_host_python::to_py;
use litellm_host_python::json_to_py;
use pyo3::{exceptions::PyBaseException, prelude::*, types::PyDict};
use serde_json::Value;
use crate::logger::PythonLogger;
use crate::python::{Logging, Wrapper};
pub struct SentRequest<'a> {
pub url: &'a str,
pub headers: &'a Py<PyDict>,
pub body: &'a Py<PyDict>,
}
pub trait LegacyCallbacks {
/// `Logging.update_from_kwargs`: what the logger is told about the request it is
/// about to see, with consumed credentials redacted.
@ -41,6 +47,7 @@ pub trait LegacyCallbacks {
status: u16,
headers: &[(String, String)],
body: &str,
request: Option<SentRequest<'_>>,
) -> PyResult<()>;
/// Records one provider HTTP exchange's window, as Python's timed `AsyncHTTPHandler.post`
@ -117,9 +124,7 @@ impl LegacyCallbacks for PythonLogger {
let redacted_kwargs = redact(py, kwargs, &secret_fields)?;
let optional_params = redact(
py,
&to_py(py, &context.optional_params)?
.into_bound(py)
.cast_into::<PyDict>()?,
&json_to_py(py, &context.optional_params)?.cast_into::<PyDict>()?,
&secret_fields,
)?;
let params = PyDict::new(py);
@ -138,7 +143,7 @@ impl LegacyCallbacks for PythonLogger {
}
}
for (name, value) in request_params.into_iter().flatten() {
params.set_item(name, to_py(py, value)?)?;
params.set_item(name, json_to_py(py, value)?)?;
}
Logging::Update.call(
py,
@ -177,8 +182,13 @@ impl LegacyCallbacks for PythonLogger {
status: u16,
headers: &[(String, String)],
body: &str,
request: Option<SentRequest<'_>>,
) -> PyResult<()> {
Logging::RecordResponse.call(py, (self.object(py), status, headers.to_vec(), body))?;
let request = request.map(|sent| (sent.url, sent.headers.bind(py), sent.body.bind(py)));
Logging::RecordResponse.call(
py,
(self.object(py), status, headers.to_vec(), body, request),
)?;
Ok(())
}

2
litellm-rust/crates/callbacks-legacy-python/src/lib.rs
View file

@ -17,7 +17,7 @@ mod mapping;
mod python;
pub use adapter::LegacyLogging;
pub use call::PublicCall;
pub(crate) use callbacks::{LegacyCallbacks, is_internal_call};
pub(crate) use callbacks::{LegacyCallbacks, SentRequest, is_internal_call};
pub(crate) use logger::{DeploymentHooks, PythonLogger, finalize, setup};
pub use mapping::{CallBoundary, CallbackMapping, Dispatch, callback_mappings};

6
litellm-rust/crates/callbacks-legacy-python/src/logger.rs
View file

@ -77,8 +77,12 @@ pub fn finalize(
kwargs: &Py<PyDict>,
start: &Py<PyAny>,
end: &Option<Py<PyAny>>,
cache_hit: bool,
) -> PyResult<()> {
Wrapper::Finalize.call(py, (response, logger.object(py), kwargs, start, end))?;
Wrapper::Finalize.call(
py,
(response, logger.object(py), kwargs, start, end, cache_hit),
)?;
Ok(())
}

3
litellm-rust/crates/callbacks-legacy-python/src/mapping.rs
View file

@ -179,6 +179,9 @@ impl CallHooks<PythonRuntime> for LegacyLogging {
PythonCallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw }) => {
(AFTER.invoke)(self, py, raw)
}
PythonCallEvent::Execution(ExecutionEvent::ProviderRequestRetried { retry }) => {
self.request_retried(py, &retry)
}
PythonCallEvent::Succeeded { timing, response } => {
(SUCCESS.invoke)(self, py, timing, response)
}

6
litellm-rust/crates/callbacks-legacy-python/src/python.rs
View file

@ -50,6 +50,10 @@ pub(crate) enum Logging {
RecordResponse,
#[strum(serialize = "record_api_window")]
RecordApiWindow,
#[strum(serialize = "times_provider_exchange")]
TimesProviderExchange,
#[strum(serialize = "retried_request")]
RetriedRequest,
#[strum(serialize = "defers_async_logging")]
DefersAsync,
#[strum(serialize = "claim_async_success_log")]
@ -62,6 +66,8 @@ pub(crate) enum Logging {
SubmitSuccess,
#[strum(serialize = "async_success_handler")]
AsyncSuccessHandler,
#[strum(serialize = "cache_hit_success")]
CacheHitSuccess,
#[strum(serialize = "mark_async_success_enqueued")]
MarkAsyncSuccessEnqueued,
#[strum(serialize = "enqueue_logging")]

14
litellm-rust/crates/callbacks-legacy-python/src/test_support.rs
View file

@ -53,7 +53,7 @@ def setup(call_type, args, kwargs, start, asynchronous):
FAKES = {
'setup': setup,
'finalize': lambda response, logger, kwargs, start, end: logger.record('finalize', response),
'finalize': lambda response, logger, kwargs, start, end, cache_hit: logger.record('finalize', response),
'update_logging': lambda logger, kwargs, model, optional_params, litellm_params, provider: logger.update_from_kwargs(
kwargs=kwargs,
model=model,
@ -66,12 +66,17 @@ FAKES = {
'post_call': lambda logger, original_response, api_key, additional_args: logger.post_call(
original_response, api_key, additional_args
),
'record_response': lambda logger, status, headers, body: logger.model_call_details.update(
'record_response': lambda logger, status, headers, body, sent: logger.model_call_details.update(
response_headers=dict(headers)
),
'retried_request': lambda logger, logged_body, body, reasons: (
logger.record('retried_request', (logged_body, body, reasons)),
'awaitable',
)[1],
'record_api_window': lambda logger, start, end: logger.model_call_details.update(
llm_api_duration_ms=(end - start) * 1000
),
'times_provider_exchange': lambda client, timeout: client is None,
'defers_async_logging': lambda logger: bool(getattr(logger, '_defer_async_logging', False)),
'claim_async_success_log': lambda logger: logger.claim_async_success_log(),
'defer_success': lambda logger, pending, response, start, end: (
@ -86,6 +91,7 @@ FAKES = {
)(error, ''.join(traceback.format_exception(error)), start, end),
'submit_success': lambda logger, response, start, end: logger.record('submit', (response, start, end)),
'async_success_handler': lambda logger, response, start, end: logger.async_success_handler(response, start, end),
'cache_hit_success': lambda logger, response, start, end: logger.record('cache_hit_success', response),
'mark_async_success_enqueued': lambda logger: logger.record('mark_async_success_enqueued', None),
'enqueue_logging': lambda coroutine: coroutine.enqueue(),
'restore_context': lambda logger: logger.record('restore', None),
@ -96,8 +102,8 @@ FAKES = {
'success', response, call_type
),
'after_deployment_failure': lambda kwargs, error, call_type: kwargs['logger'].hook('failure', error, call_type),
'stream_opened': lambda logger, head: logger.record('stream_opened', head),
'stamp_stream_head': lambda head, logger, kwargs, start, end: logger.record('stamp_stream_head', head),
'stream_opened': lambda logger, head, keeps_requested: logger.record('stream_opened', head),
'stamp_stream_head': lambda head, logger, kwargs, start, end, cache_hit: logger.record('stamp_stream_head', head),
'stream_success': lambda logger, url_route, endpoint_type, request_body, chunks, start, end, first_chunk, stream_teardown: logger.record(
'stream_success', list(chunks)
),

7
litellm-rust/crates/gateway-inference/tests/caching.rs
View file

@ -36,7 +36,12 @@ async fn all_inference_endpoints_share_native_cache(
} else {
"message_stop"
};
let events = format!("event: {terminal}\ndata: {{\"type\":\"{terminal}\"}}\n\n");
let output = if is_responses {
""
} else {
"event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"hello\"}}\n\n"
};
let events = format!("{output}event: {terminal}\ndata: {{\"type\":\"{terminal}\"}}\n\n");
let template = if stream {
ResponseTemplate::new(200).set_body_raw(events.clone(), "text/event-stream")
} else {

7
litellm-rust/crates/host-native/src/driver.rs
View file

@ -84,6 +84,13 @@ where
reply.send(self.interceptors.after_provider_response(raw).await);
Ok(())
}
HostRequest::Intercept(InterceptRequest::ProviderRequestRetried {
retry,
reply,
}) => {
reply.send(self.interceptors.provider_request_retried(retry).await);
Ok(())
}
HostRequest::Stream(StreamDelivery::Open(head, reply)) => {
self.demand = Some(reply);
return Ok(Boundary::Open(head));

2
litellm-rust/crates/host-python/Cargo.toml
View file

@ -14,6 +14,7 @@ bytes.workspace = true
futures-util.workspace = true
parking_lot.workspace = true
serde.workspace = true
serde_json.workspace = true
thiserror.workspace = true
tokio = { workspace = true, features = ["rt", "sync"] }
@ -23,4 +24,3 @@ pythonize = "0.29.0"
[dev-dependencies]
rstest.workspace = true
serde_json.workspace = true

4
litellm-rust/crates/host-python/src/binding.rs
View file

@ -29,8 +29,8 @@ pub trait PythonBinding: PythonOwned {
/// What the stream carries at hand-off, as the caller's stream receives it.
/// The unary response after the deployment transform, before success hooks run.
fn deliver_response(&mut self, _py: Python<'_>, _response: &Bound<'_, PyAny>) -> PyResult<()> {
Ok(())
fn deliver_response(&mut self, _py: Python<'_>, response: Py<PyAny>) -> PyResult<Py<PyAny>> {
Ok(response)
}
fn encode_stream_head(

12
litellm-rust/crates/host-python/src/driver.rs
View file

@ -310,6 +310,14 @@ where
reply.send(Ok(()));
Ok(())
}
HostRequest::Intercept(InterceptRequest::ProviderRequestRetried { retry, reply }) => {
self.notify(PythonCallEvent::Execution(
ExecutionEvent::ProviderRequestRetried { retry },
))
.await?;
reply.send(Ok(()));
Ok(())
}
HostRequest::Stream(StreamDelivery::Open(head, reply)) => {
self.streaming = true;
let head = {
@ -405,10 +413,10 @@ where
};
let delivered = {
let mut parts = self.parts();
Python::attach(|py| parts.binding.deliver_response(py, response.bind(py)))
Python::attach(|py| parts.binding.deliver_response(py, response))
};
match delivered {
Ok(()) => self.succeeded(response).await,
Ok(response) => self.succeeded(response).await,
Err(error) => self.failure(error, FailureOrigin::Host).await,
}
}

9
litellm-rust/crates/host-python/src/hooks/chain/dispatch.rs
View file

@ -241,6 +241,9 @@ fn retain_event(py: Python<'_>, event: PythonCallEvent<'_>) -> OwnedEvent {
CallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw }) => {
CallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw: raw.clone() })
}
CallEvent::Execution(ExecutionEvent::ProviderRequestRetried { retry }) => {
CallEvent::Execution(ExecutionEvent::ProviderRequestRetried { retry })
}
CallEvent::Succeeded { timing, response } => CallEvent::Succeeded {
timing,
response: response.clone_ref(py),
@ -280,6 +283,12 @@ fn dispatch(
py,
CallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw }),
),
CallEvent::Execution(ExecutionEvent::ProviderRequestRetried { retry }) => hooks.on_event(
py,
CallEvent::Execution(ExecutionEvent::ProviderRequestRetried {
retry: retry.clone(),
}),
),
CallEvent::Succeeded { timing, response } => hooks.on_event(
py,
CallEvent::Succeeded {

3
litellm-rust/crates/host-python/src/lib.rs
View file

@ -31,7 +31,8 @@ pub use gil::{PythonContext, attach_blocking, release_count, release_gil};
pub use handle::{Execution, ExecutionBody, ExecutionStep, StepTypes, Suspension};
pub use hooks::{HookChain, HookResume, HookStep, PythonCallEvent, PythonCallHooks, PythonRuntime};
pub use marshal::{
Pythonized, from_py, from_py_argument, json_loads, json_object_field, panic_to_pyerr, to_py,
Pythonized, from_py, from_py_argument, json_from_py, json_loads, json_object_field, json_to_py,
panic_to_pyerr, to_py,
};
pub use owned::PythonOwned;
pub use runtime::{

182
litellm-rust/crates/host-python/src/marshal.rs
View file

@ -4,9 +4,10 @@ use std::panic::{AssertUnwindSafe, catch_unwind};
use pyo3::exceptions::PyValueError;
use pyo3::panic::PanicException;
use pyo3::prelude::*;
use pyo3::types::PyBytes;
use pyo3::types::{PyBool, PyBytes, PyDict, PyFloat, PyInt, PyList, PyString, PyTuple};
use serde::Serialize;
use serde::de::DeserializeOwned;
use serde_json::{Map, Number, Value};
/// Converts a `#[pyo3(from_py_with = ...)]` argument, reporting failures as `ValueError`
/// so a bad argument reads as a bad argument rather than as whatever the conversion hit.
@ -71,6 +72,119 @@ pub fn panic_to_pyerr(payload: Box<dyn Any + Send>) -> PyErr {
PanicException::new_err(message.to_string())
}
pub fn json_from_py(value: &Bound<'_, PyAny>) -> PyResult<Value> {
if value.is_none() {
return Ok(Value::Null);
}
if let Ok(flag) = value.cast::<PyBool>() {
return Ok(Value::Bool(flag.is_true()));
}
if let Ok(int) = value.cast::<PyInt>() {
if let Ok(number) = int.extract::<i64>() {
return Ok(Value::from(number));
}
if let Ok(number) = int.extract::<u64>() {
return Ok(Value::from(number));
}
return raw_leaf(value);
}
if let Ok(float) = value.cast::<PyFloat>() {
return match Number::from_f64(float.value()) {
Some(number) => Ok(Value::Number(number)),
None => raw_leaf(value),
};
}
if let Ok(text) = value.cast::<PyString>() {
return json_string(text);
}
if let Ok(dict) = value.cast::<PyDict>() {
if !dict
.keys()
.iter()
.all(|key| key.is_instance_of::<PyString>())
{
return from_py(value);
}
return dict
.iter()
.map(|(key, item)| Ok((json_key(key.cast::<PyString>()?)?, json_from_py(&item)?)))
.collect::<PyResult<Map<String, Value>>>()
.map(Value::Object);
}
if let Ok(list) = value.cast::<PyList>() {
return list
.iter()
.map(|item| json_from_py(&item))
.collect::<PyResult<_>>()
.map(Value::Array);
}
if let Ok(tuple) = value.cast::<PyTuple>() {
return tuple
.iter()
.map(|item| json_from_py(&item))
.collect::<PyResult<_>>()
.map(Value::Array);
}
from_py(value)
}
pub fn json_to_py<'py>(py: Python<'py>, value: &Value) -> PyResult<Bound<'py, PyAny>> {
Ok(match value {
Value::Null => py.None().into_bound(py),
Value::Bool(flag) => PyBool::new(py, *flag).to_owned().into_any(),
Value::Number(number) => match (number.as_i64(), number.as_u64()) {
(Some(int), _) => int.into_pyobject(py)?.into_any(),
(None, Some(int)) => int.into_pyobject(py)?.into_any(),
(None, None) => PyFloat::new(py, number.as_f64().unwrap_or(f64::NAN)).into_any(),
},
Value::String(text) => match litellm_host::raw_json::decode(text) {
Some(raw) => py.import("json")?.call_method1("loads", (raw,))?,
None => PyString::new(py, text).into_any(),
},
Value::Array(items) => PyList::new(
py,
items
.iter()
.map(|item| json_to_py(py, item))
.collect::<PyResult<Vec<_>>>()?,
)?
.into_any(),
Value::Object(fields) => {
let dict = PyDict::new(py);
for (key, item) in fields {
dict.set_item(
json_to_py(py, &Value::String(key.clone()))?,
json_to_py(py, item)?,
)?;
}
dict.into_any()
}
})
}
fn json_key(text: &Bound<'_, PyString>) -> PyResult<String> {
match json_string(text)? {
Value::String(key) => Ok(key),
_ => unreachable!("a Python str converts to a JSON string"),
}
}
fn json_string(text: &Bound<'_, PyString>) -> PyResult<Value> {
match text.to_str() {
Ok(text) if !litellm_host::raw_json::is_marked(text) => Ok(Value::String(text.to_owned())),
_ => raw_leaf(text.as_any()),
}
}
fn raw_leaf(value: &Bound<'_, PyAny>) -> PyResult<Value> {
let raw: String = value
.py()
.import("json")?
.call_method1("dumps", (value,))?
.extract()?;
Ok(Value::String(litellm_host::raw_json::encode(&raw)))
}
#[cfg(test)]
mod tests {
use serde::Serializer;
@ -131,6 +245,72 @@ mod tests {
});
}
#[rstest::rstest]
#[case::plain(
c"{'a': [1, 2.5, None, True, 'x'], 'b': (1,)}",
r#"{"a":[1,2.5,null,true,"x"],"b":[1]}"#
)]
#[case::nan(c"{'t': float('nan')}", r#"{"t":NaN}"#)]
#[case::infinities(c"[float('inf'), -float('inf')]", "[Infinity,-Infinity]")]
#[case::wide_integers(
c"[2**64, -(2**70), 2**64 - 1, -(2**63)]",
"[18446744073709551616,-1180591620717411303424,18446744073709551615,-9223372036854775808]"
)]
#[case::lone_surrogates(c"{'k\\udc00': 'a\\ud800b \\u00e9'}", r#"{"k\udc00":"a\ud800b é"}"#)]
#[case::marked_text(c"['\\uf8fflitellm-raw-json:1']", r#"["litellm-raw-json:1"]"#)]
fn json_values_keep_what_serde_json_cannot_hold(
#[case] source: &std::ffi::CStr,
#[case] wire: &str,
) {
crate::initialize_python();
Python::attach(|py| {
let value = py.eval(source, None, None).unwrap();
let converted = json_from_py(&value).unwrap();
let json = py.import("json").unwrap();
let loads = |text: &str| json.call_method1("loads", (text,)).unwrap();
let wire_bytes = litellm_host::raw_json::to_vec(&converted).unwrap();
assert!(
loads(std::str::from_utf8(&wire_bytes).unwrap())
.eq(loads(wire))
.unwrap()
|| wire.contains("NaN")
);
assert_eq!(
json.call_method1("dumps", (loads(std::str::from_utf8(&wire_bytes).unwrap()),))
.unwrap()
.extract::<String>()
.unwrap(),
json.call_method1("dumps", (loads(wire),))
.unwrap()
.extract::<String>()
.unwrap()
);
let back = json_to_py(py, &converted).unwrap();
assert_eq!(
json.call_method1("dumps", (back,))
.unwrap()
.extract::<String>()
.unwrap(),
json.call_method1("dumps", (loads(wire),))
.unwrap()
.extract::<String>()
.unwrap()
);
});
}
#[test]
fn values_json_cannot_hold_keep_their_old_conversion() {
crate::initialize_python();
Python::attach(|py| {
let value = py.eval(c"{1: 'a'}", None, None).unwrap();
assert_eq!(
json_from_py(&value).unwrap_err().to_string(),
from_py::<Value>(&value).unwrap_err().to_string()
);
});
}
#[test]
fn depythonize_preserves_python_exception_identity_and_traceback() {
crate::initialize_python();

3
litellm-rust/crates/host-python/tests/support/mod.rs
View file

@ -422,6 +422,9 @@ impl CallHooks<PythonRuntime> for SyntheticHooks {
PythonCallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw }) => {
format!("response:{}", raw.body)
}
PythonCallEvent::Execution(ExecutionEvent::ProviderRequestRetried { retry }) => {
format!("retried:{}", retry.body)
}
PythonCallEvent::Succeeded { response, .. } => {
format!("succeeded:{}", response.bind(py))
}

1
litellm-rust/crates/host/Cargo.toml
View file

@ -9,6 +9,7 @@ repository.workspace = true
litellm-secret-value.workspace = true
futures-util.workspace = true
litellm-coroutine.workspace = true
serde.workspace = true
serde_json.workspace = true
thiserror.workspace = true
tokio = { workspace = true, features = ["sync"] }

28
litellm-rust/crates/host/src/interceptors.rs
View file

@ -10,6 +10,18 @@ pub struct WireRequest {
pub body: Value,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum RetryReason {
ThinkingSignature,
StrictTools { dropped: usize },
}
#[derive(Clone, Debug, PartialEq)]
pub struct RetriedRequest {
pub reasons: Vec<RetryReason>,
pub body: Value,
}
/// What the route knows about the request it is sending, for a host that logs it. The
/// route owns these facts; a host reads them beside the wire request and never rewrites
/// them.
@ -62,6 +74,7 @@ pub struct ProviderIdentity {
pub enum ResultSource {
Provider,
Cache { key: String },
Host,
}
/// What the call cost, settled once by the route. A host records it; it never re-derives it.
@ -79,6 +92,7 @@ impl Cost {
pub fn settle(source: &ResultSource, reported: Option<serde_json::Number>) -> Self {
match (source, reported) {
(ResultSource::Cache { .. }, _) => Self::CacheHit,
(ResultSource::Host, _) => Self::Deferred,
(ResultSource::Provider, Some(amount)) => Self::Reported { amount },
(ResultSource::Provider, None) => Self::Deferred,
}
@ -107,6 +121,13 @@ pub trait Interceptors<E>: Send + Sync {
&self,
raw: RawResponse,
) -> impl Future<Output = Result<(), E>> + Send;
fn provider_request_retried(
&self,
_retry: RetriedRequest,
) -> impl Future<Output = Result<(), E>> + Send {
async { Ok(()) }
}
}
impl<E, T: Interceptors<E> + ?Sized> Interceptors<E> for &T {
@ -128,6 +149,13 @@ impl<E, T: Interceptors<E> + ?Sized> Interceptors<E> for &T {
) -> impl Future<Output = Result<(), E>> + Send {
(**self).after_provider_response(raw)
}
fn provider_request_retried(
&self,
retry: RetriedRequest,
) -> impl Future<Output = Result<(), E>> + Send {
(**self).provider_request_retried(retry)
}
}
impl<E> Interceptors<E> for () {

1
litellm-rust/crates/host/src/lib.rs
View file

@ -13,3 +13,4 @@ pub mod lifecycle;
pub mod machine;
pub mod observation;
pub mod protocol;
pub mod raw_json;

15
litellm-rust/crates/host/src/lifecycle.rs
View file

@ -6,7 +6,10 @@ use std::{
use futures_util::TryStreamExt;
use crate::{call::CallOutput, interceptors::RawResponse};
use crate::{
call::CallOutput,
interceptors::{RawResponse, RetriedRequest},
};
/// Seconds since the Unix epoch, on one clock for every host.
pub fn epoch_seconds() -> f64 {
@ -50,7 +53,7 @@ pub enum CallEvent<Response = (), Error = (), Raw = RawResponse> {
},
}
#[derive(Clone, Debug, PartialEq, Eq)]
#[derive(Clone, Debug, PartialEq)]
pub enum ExecutionEvent<Raw = RawResponse> {
ResultReady {
facts: crate::interceptors::ExecutionFacts,
@ -58,6 +61,9 @@ pub enum ExecutionEvent<Raw = RawResponse> {
ProviderResponseReceived {
raw: Raw,
},
ProviderRequestRetried {
retry: RetriedRequest,
},
}
impl<Response, Error, Raw: std::borrow::Borrow<RawResponse>> CallEvent<Response, Error, Raw> {
@ -76,6 +82,11 @@ impl<Response, Error, Raw: std::borrow::Borrow<RawResponse>> CallEvent<Response,
facts: facts.clone(),
})
}
Self::Execution(ExecutionEvent::ProviderRequestRetried { retry }) => {
CallEvent::Execution(ExecutionEvent::ProviderRequestRetried {
retry: retry.clone(),
})
}
Self::Succeeded { timing, .. } => CallEvent::Succeeded {
timing: *timing,
response: (),

10
litellm-rust/crates/host/src/machine/context.rs
View file

@ -5,7 +5,7 @@ use litellm_coroutine::Co;
use super::coroutine::MachineFault;
use crate::{
interceptors::{RawResponse, RequestContext, WireRequest},
interceptors::{RawResponse, RequestContext, RetriedRequest, WireRequest},
protocol::{HostRequest, InterceptRequest, Protocol, Reply, StreamDelivery},
};
@ -119,6 +119,14 @@ where
})
.await?
}
async fn provider_request_retried(&self, retry: RetriedRequest) -> Result<(), P::Error> {
self.0
.request_reply(|reply| {
HostRequest::Intercept(InterceptRequest::ProviderRequestRetried { retry, reply })
})
.await?
}
}
pub struct StreamSender<P: Protocol>(Channel<P>);

6
litellm-rust/crates/host/src/protocol.rs
View file

@ -2,7 +2,7 @@ use std::ops::ControlFlow;
pub use litellm_coroutine::{Abandoned, Answer, Reply, reply};
use crate::interceptors::{RawResponse, RequestContext, WireRequest};
use crate::interceptors::{RawResponse, RequestContext, RetriedRequest, WireRequest};
pub trait Protocol: Send + Sync + 'static {
type Request: Send + 'static;
@ -33,6 +33,10 @@ pub enum InterceptRequest<E> {
raw: RawResponse,
reply: Reply<Result<(), E>>,
},
ProviderRequestRetried {
retry: RetriedRequest,
reply: Reply<Result<(), E>>,
},
}
pub enum StreamDelivery<P: Protocol> {

362
litellm-rust/crates/host/src/raw_json.rs Normal file
View file

@ -0,0 +1,362 @@
use std::{borrow::Cow, io};
use serde::Serialize;
use serde_json::ser::{CharEscape, CompactFormatter, Formatter};
const MARK: char = '\u{F8FF}';
const PREFIX: &str = "\u{F8FF}litellm-raw-json:";
pub fn encode(raw: &str) -> String {
format!("{PREFIX}{raw}")
}
pub fn decode(text: &str) -> Option<&str> {
text.strip_prefix(PREFIX).filter(|raw| is_scalar(raw))
}
pub fn is_marked(text: &str) -> bool {
text.starts_with(PREFIX)
}
pub fn to_vec<T: Serialize + ?Sized>(value: &T) -> serde_json::Result<Vec<u8>> {
let mut out = Vec::with_capacity(128);
let mut serializer = serde_json::Serializer::with_formatter(&mut out, RawLeaves::default());
value.serialize(&mut serializer)?;
Ok(out)
}
pub fn mark_wide_numbers(text: &str) -> Cow<'_, str> {
let bytes = text.as_bytes();
let mut out = String::new();
let mut copied = 0;
let mut index = 0;
while index < bytes.len() {
match bytes[index] {
b'"' => {
let end = string_end(bytes, index + 1);
if let Some(escaped) = escaped_marked_string(&text[index..end]) {
out.push_str(&text[copied..index]);
out.push_str(&escaped);
copied = end;
}
index = end;
}
b'-' | b'0'..=b'9' => {
let end = bytes[index..]
.iter()
.position(|byte| {
!matches!(byte, b'0'..=b'9' | b'-' | b'+' | b'.' | b'e' | b'E')
})
.map_or(bytes.len(), |length| index + length);
if let Some(raw) = wide_number(&text[index..end]) {
out.push_str(&text[copied..index]);
out.push_str(&serde_json::Value::String(encode(raw)).to_string());
copied = end;
}
index = end;
}
_ => index += 1,
}
}
if copied == 0 {
return Cow::Borrowed(text);
}
out.push_str(&text[copied..]);
Cow::Owned(out)
}
fn escaped_marked_string(token: &str) -> Option<String> {
let inner = token.strip_prefix('"')?;
let escaped_mark = inner
.get(..6)
.is_some_and(|head| head.eq_ignore_ascii_case("\\uf8ff"));
if !inner.starts_with(MARK) && !escaped_mark {
return None;
}
let text = serde_json::from_str::<String>(token).ok()?;
is_marked(&text)
.then(|| serde_json::Value::String(encode(&ascii_string_literal(&text))).to_string())
}
fn ascii_string_literal(text: &str) -> String {
let mut literal = String::with_capacity(text.len() + 2);
literal.push('"');
for character in text.chars() {
match character {
'"' => literal.push_str("\\\""),
'\\' => literal.push_str("\\\\"),
'\n' => literal.push_str("\\n"),
'\r' => literal.push_str("\\r"),
'\t' => literal.push_str("\\t"),
' '..='~' => literal.push(character),
_ => {
let mut units = [0; 2];
for unit in character.encode_utf16(&mut units) {
literal.push_str(&format!("\\u{unit:04x}"));
}
}
}
}
literal.push('"');
literal
}
fn string_end(bytes: &[u8], mut index: usize) -> usize {
while index < bytes.len() {
match bytes[index] {
b'\\' => index += 2,
b'"' => return index + 1,
_ => index += 1,
}
}
index
}
fn wide_number(raw: &str) -> Option<&str> {
if is_integer(raw) {
return (raw.parse::<i64>().is_err() && raw.parse::<u64>().is_err()).then_some(raw);
}
match raw.parse::<f64>() {
Ok(number) if number == f64::INFINITY => Some("Infinity"),
Ok(number) if number == f64::NEG_INFINITY => Some("-Infinity"),
_ => None,
}
}
fn is_scalar(raw: &str) -> bool {
matches!(raw, "NaN" | "Infinity" | "-Infinity") || is_integer(raw) || is_string_literal(raw)
}
fn is_integer(raw: &str) -> bool {
let digits = raw.strip_prefix('-').unwrap_or(raw);
match digits.as_bytes() {
[b'0'] => true,
[b'1'..=b'9', rest @ ..] => rest.iter().all(u8::is_ascii_digit),
_ => false,
}
}
fn is_string_literal(raw: &str) -> bool {
let Some(inner) = raw
.strip_prefix('"')
.and_then(|rest| rest.strip_suffix('"'))
else {
return false;
};
let mut bytes = inner.bytes();
while let Some(byte) = bytes.next() {
match byte {
b'\\' => match bytes.next() {
Some(b'"' | b'\\' | b'/' | b'b' | b'f' | b'n' | b'r' | b't') => {}
Some(b'u') => {
if !(0..4).all(|_| bytes.next().is_some_and(|hex| hex.is_ascii_hexdigit())) {
return false;
}
}
_ => return false,
},
b'"' => return false,
0x20..=0x7e => {}
_ => return false,
}
}
true
}
fn unescaped(escape: CharEscape) -> char {
match escape {
CharEscape::Quote => '"',
CharEscape::ReverseSolidus => '\\',
CharEscape::Solidus => '/',
CharEscape::Backspace => '\u{8}',
CharEscape::FormFeed => '\u{c}',
CharEscape::LineFeed => '\n',
CharEscape::CarriageReturn => '\r',
CharEscape::Tab => '\t',
CharEscape::AsciiControl(byte) => char::from(byte),
}
}
#[derive(Default)]
enum Text {
#[default]
Closed,
Opened,
Plain,
Marked(String),
}
#[derive(Default)]
struct RawLeaves {
key: bool,
text: Text,
}
impl Formatter for RawLeaves {
fn begin_object_key<W: ?Sized + io::Write>(
&mut self,
writer: &mut W,
first: bool,
) -> io::Result<()> {
self.key = true;
CompactFormatter.begin_object_key(writer, first)
}
fn end_object_key<W: ?Sized + io::Write>(&mut self, writer: &mut W) -> io::Result<()> {
self.key = false;
CompactFormatter.end_object_key(writer)
}
fn begin_string<W: ?Sized + io::Write>(&mut self, _writer: &mut W) -> io::Result<()> {
self.text = Text::Opened;
Ok(())
}
fn write_string_fragment<W: ?Sized + io::Write>(
&mut self,
writer: &mut W,
fragment: &str,
) -> io::Result<()> {
match &mut self.text {
Text::Opened if fragment.starts_with(MARK) => {
self.text = Text::Marked(fragment.to_owned());
Ok(())
}
Text::Opened => {
self.text = Text::Plain;
writer.write_all(b"\"")?;
writer.write_all(fragment.as_bytes())
}
Text::Marked(held) => {
held.push_str(fragment);
Ok(())
}
Text::Plain | Text::Closed => writer.write_all(fragment.as_bytes()),
}
}
fn write_char_escape<W: ?Sized + io::Write>(
&mut self,
writer: &mut W,
escape: CharEscape,
) -> io::Result<()> {
match &mut self.text {
Text::Opened => {
self.text = Text::Plain;
writer.write_all(b"\"")?;
CompactFormatter.write_char_escape(writer, escape)
}
Text::Marked(held) => {
held.push(unescaped(escape));
Ok(())
}
Text::Plain | Text::Closed => CompactFormatter.write_char_escape(writer, escape),
}
}
fn end_string<W: ?Sized + io::Write>(&mut self, writer: &mut W) -> io::Result<()> {
match std::mem::take(&mut self.text) {
Text::Opened => writer.write_all(b"\"\""),
Text::Plain | Text::Closed => writer.write_all(b"\""),
Text::Marked(held) => match decode(&held) {
Some(raw) if !self.key || raw.starts_with('"') => writer.write_all(raw.as_bytes()),
_ => serde_json::to_writer(writer, &held).map_err(io::Error::other),
},
}
}
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use serde_json::{Value, json};
use super::*;
#[rstest]
#[case::nan("NaN")]
#[case::infinity("Infinity")]
#[case::negative_infinity("-Infinity")]
#[case::wide_integer("18446744073709551616")]
#[case::negative_wide_integer("-1180591620717411303424")]
#[case::zero("0")]
#[case::surrogate_string(r#""sys \ud800 lone""#)]
#[case::escaped_string(r#""a\"b\\c\n\u00e9""#)]
fn scalar_texts_are_written_back_verbatim(#[case] raw: &str) {
let body = json!({"leaf": encode(raw), "list": [encode(raw)]});
assert_eq!(
String::from_utf8(to_vec(&body).unwrap()).unwrap(),
format!(r#"{{"leaf":{raw},"list":[{raw}]}}"#)
);
}
#[rstest]
#[case::structure("1, \"injected\": true")]
#[case::object("{}")]
#[case::leading_zero("012")]
#[case::float("1.5")]
#[case::unquoted_quote(r#""a"b""#)]
#[case::non_ascii("\"\u{e9}\"")]
#[case::bad_escape(r#""\x""#)]
#[case::short_unicode_escape(r#""\u12""#)]
#[case::empty("")]
fn other_marked_texts_stay_strings(#[case] raw: &str) {
let marked = encode(raw);
assert_eq!(decode(&marked), None);
assert_eq!(
to_vec(&json!([marked])).unwrap(),
serde_json::to_vec(&json!([marked])).unwrap()
);
}
#[test]
fn keys_take_only_string_texts() {
let body = Value::Object(
[
(encode(r#""k\udc00""#), json!(1)),
(encode("NaN"), json!(2)),
]
.into_iter()
.collect(),
);
assert_eq!(
String::from_utf8(to_vec(&body).unwrap()).unwrap(),
format!(
r#"{{"k\udc00":1,{}:2}}"#,
serde_json::to_string(&encode("NaN")).unwrap()
)
);
}
#[rstest]
#[case::wide_integer(r#"{"a":18446744073709551616}"#, json!({"a": encode("18446744073709551616")}))]
#[case::negative_wide_integer(r#"[-12345678901234567890123]"#, json!([encode("-12345678901234567890123")]))]
#[case::overflowing_float(r#"{"a":[1e400,-1E+400]}"#, json!({"a": [encode("Infinity"), encode("-Infinity")]}))]
#[case::numbers_that_fit(r#"{"a":[18446744073709551615,-9223372036854775808,-0.0,5e-324,1E2,0.30000000000000004]}"#, json!({"a": [18_446_744_073_709_551_615_u64, i64::MIN, -0.0, 5e-324, 100.0, 0.300_000_000_000_000_04]}))]
#[case::provider_marked_text("{\"a\":\"\u{F8FF}litellm-raw-json:42\"}", json!({"a": encode(r#""\uf8fflitellm-raw-json:42""#)}))]
#[case::provider_escaped_mark(r#"{"\uF8FFlitellm-raw-json:1":["\uf8fflitellm-raw-json:\"x\"\n\u00e9"]}"#, json!({encode(r#""\uf8fflitellm-raw-json:1""#): [encode(r#""\uf8fflitellm-raw-json:\"x\"\n\u00e9""#)]}))]
#[case::mark_not_leading(r#"["x\uf8fflitellm-raw-json:NaN"]"#, json!(["x\u{F8FF}litellm-raw-json:NaN"]))]
#[case::numbers_inside_strings(r#"{"18446744073709551616":"1e400 \" 18446744073709551616"}"#, json!({"18446744073709551616": "1e400 \" 18446744073709551616"}))]
fn wide_numbers_are_read_as_marked_texts(#[case] text: &str, #[case] read: Value) {
assert_eq!(
serde_json::from_str::<Value>(&mark_wide_numbers(text)).unwrap(),
read
);
}
#[test]
fn a_text_without_wide_numbers_is_borrowed() {
assert!(matches!(
mark_wide_numbers(r#"{"a":[1,2.5,"18446744073709551616"]}"#),
Cow::Borrowed(_)
));
}
#[rstest]
#[case::plain(json!({"a": "text", "b": ["", "\"quoted\"\n", "\u{F8FF}"], "c": 1.5, "d": null}))]
#[case::mark_inside(json!(["x\u{F8FF}litellm-raw-json:NaN"]))]
#[case::control(json!(["\u{1}\t"]))]
fn unmarked_documents_serialize_as_serde_json_does(#[case] body: Value) {
assert_eq!(to_vec(&body).unwrap(), serde_json::to_vec(&body).unwrap());
}
}

20
litellm-rust/crates/http/src/config.rs
View file

@ -94,6 +94,8 @@ impl From<&HttpSettings> for Resolution {
}
}
const MAX_REDIRECTS: usize = 20;
impl TryFrom<&HttpClientConfig> for reqwest::ClientBuilder {
type Error = Error;
@ -101,7 +103,9 @@ impl TryFrom<&HttpClientConfig> for reqwest::ClientBuilder {
let base = reqwest::Client::builder()
.tls_backend_preconfigured(rustls::ClientConfig::try_from(config)?)
.connect_timeout(config.connect_timeout)
.pool_idle_timeout(config.pool_idle_timeout);
.pool_idle_timeout(config.pool_idle_timeout)
.referer(false)
.redirect(reqwest::redirect::Policy::limited(MAX_REDIRECTS));
let with_keepalive = match config.tcp_keepalive {
None => base,
Some(keepalive) => base
@ -186,17 +190,23 @@ mod tests {
assert_eq!(resolution.unsupported, []);
}
#[test]
fn unsupported_ecdh_curve_keeps_the_defaults_and_is_reported() {
#[rstest]
#[case::openssl_lacks_in_rustls("secp521r1")]
#[case::nist_name("P-256")]
#[case::nist_p384("P-384")]
#[case::sec_name("secp256r1")]
#[case::lowercase("x25519")]
#[case::padded(" X25519 ")]
fn unsupported_ecdh_curve_keeps_the_defaults_and_is_reported(#[case] curve: &str) {
let settings = HttpSettings {
ssl_ecdh_curve: Some("secp521r1".into()),
ssl_ecdh_curve: Some(curve.into()),
..HttpSettings::default()
};
let resolution = Resolution::from(&settings);
assert_eq!(resolution.config.key_exchange_group, None);
assert_eq!(
resolution.unsupported,
[Unsupported::EcdhCurve("secp521r1".into())]
[Unsupported::EcdhCurve(curve.into())]
);
}

4
litellm-rust/crates/http/src/error.rs
View file

@ -251,6 +251,10 @@ impl Error {
Repr::ComputedHeader(name.into()).into()
}
pub fn request_body(error: serde_json::Error) -> Self {
Repr::RequestBody(error).into()
}
pub fn unsignable_body() -> Self {
Repr::UnsignableBody.into()
}

15
litellm-rust/crates/http/src/outbound.rs
View file

@ -27,7 +27,20 @@ pub fn json(
body: &impl Serialize,
timeout: Option<Duration>,
) -> Result<Request, Error> {
let body = serde_json::to_vec(body).map_err(Repr::RequestBody)?;
json_bytes(
url,
headers,
serde_json::to_vec(body).map_err(Repr::RequestBody)?,
timeout,
)
}
pub fn json_bytes(
url: &str,
headers: &[(String, String)],
body: Vec<u8>,
timeout: Option<Duration>,
) -> Result<Request, Error> {
let url = Url::parse(url).map_err(Repr::ParseUrl)?;
let with_headers = headers.iter().fold(
http::Request::post(url.as_str()),

28
litellm-rust/crates/http/src/pool.rs
View file

@ -16,6 +16,7 @@ use crate::{
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum ClientVariant {
Provider,
Undecoded,
NoRedirect,
Media,
UnpinnedMedia,
@ -103,6 +104,7 @@ impl HttpClientPool {
) -> reqwest::ClientBuilder {
match variant {
ClientVariant::Provider => builder,
ClientVariant::Undecoded => builder.no_gzip().no_deflate().no_brotli().no_zstd(),
ClientVariant::NoRedirect | ClientVariant::UnpinnedMedia | ClientVariant::WebSocket => {
builder.redirect(reqwest::redirect::Policy::none())
}
@ -132,7 +134,9 @@ fn effective(config: &HttpClientConfig, variant: ClientVariant) -> HttpClientCon
},
..config.clone()
},
ClientVariant::Provider | ClientVariant::NoRedirect => config.clone(),
ClientVariant::Provider | ClientVariant::Undecoded | ClientVariant::NoRedirect => {
config.clone()
}
}
}
@ -367,6 +371,28 @@ mod tests {
assert!(request.contains("user-agent: litellm-test/9"), "{request}");
}
#[tokio::test]
async fn provider_client_follows_twenty_redirects_without_a_referer() {
let (address, _, requests) = serve("HTTP/1.1 307 Temporary Redirect").await;
let followed = pool()
.client(&config("a"), ClientVariant::Provider)
.unwrap()
.get(format!("http://{address}/start"))
.timeout(Duration::from_secs(5))
.send()
.await;
assert!(followed.is_err());
let requests = requests.lock().unwrap().clone();
assert_eq!(requests.len(), 21, "{requests:?}");
assert!(requests[1].starts_with("GET /elsewhere "), "{requests:?}");
assert!(
requests
.iter()
.all(|request| !request.to_ascii_lowercase().contains("referer:")),
"{requests:?}"
);
}
#[tokio::test]
async fn no_redirect_variant_returns_the_redirect_instead_of_following_it() {
let (address, _, _) = serve("HTTP/1.1 302 Found").await;

8
litellm-rust/crates/http/src/tls.rs
View file

@ -23,10 +23,10 @@ impl FromStr for KeyExchangeGroup {
type Err = Unsupported;
fn from_str(name: &str) -> Result<Self, Self::Err> {
match name.trim().to_ascii_lowercase().as_str() {
"x25519" => Ok(Self::X25519),
"prime256v1" | "secp256r1" | "p-256" => Ok(Self::Secp256r1),
"secp384r1" | "p-384" => Ok(Self::Secp384r1),
match name {
"X25519" => Ok(Self::X25519),
"prime256v1" => Ok(Self::Secp256r1),
"secp384r1" => Ok(Self::Secp384r1),
_ => Err(Unsupported::EcdhCurve(name.to_owned())),
}
}

1
litellm-rust/crates/inference-chat/tests/caching.rs
View file

@ -117,6 +117,7 @@ async fn the_same_route_entrypoint_reports_facts_with_or_without_caching(
assert!(caching);
assert!(!key.is_empty());
}
ResultSource::Host => panic!("the host never answers a chat call"),
}
let summaries = traces.summaries("litellm.route");
assert_eq!(summaries.len(), 2);

4
litellm-rust/crates/inference-messages/Cargo.toml
View file

@ -8,7 +8,9 @@ repository.workspace = true
[dependencies]
litellm-secret-value.workspace = true
bytes.workspace = true
flate2 = { workspace = true, features = ["zlib-rs"] }
futures-util.workspace = true
http.workspace = true
litellm-auth = { workspace = true, features = ["aws", "azure", "gcp"] }
litellm-auth-aws.workspace = true
litellm-cache-response.workspace = true
@ -30,9 +32,11 @@ tracing.workspace = true
[dev-dependencies]
litellm-error.workspace = true
litellm-framer = { workspace = true, features = ["sse"] }
litellm-cache.workspace = true
litellm-cache-memory.workspace = true
litellm-host-native.workspace = true
litellm-http = { workspace = true, features = ["test-support"] }
litellm-inference-testing.workspace = true
rstest.workspace = true
thiserror.workspace = true
wiremock.workspace = true

3
litellm-rust/crates/inference-messages/src/constants.rs
View file

@ -2,3 +2,6 @@
/// seconds. Mirrors the Python Anthropic Messages default. The per-request
/// timeout from the caller still overrides this on the request builder.
pub(crate) const MESSAGES_TIMEOUT_SECS: u64 = 600;
pub(crate) const MESSAGES_HTTP_ERROR_ATTEMPTS: usize = 2;
pub(crate) const MESSAGES_STREAM_ERROR_BODY_READ_SECS: u64 = 5;

437
litellm-rust/crates/inference-messages/src/encoding.rs Normal file
View file

@ -0,0 +1,437 @@
use bytes::Bytes;
use flate2::{Decompress, DecompressError, FlushDecompress, Status};
use futures_util::StreamExt;
use reqwest::header::{CONTENT_ENCODING, HeaderMap};
type BoxError = Box<dyn std::error::Error + Send + Sync>;
const MAX_WINDOW_BITS: u8 = 15;
const MAX_DECOMPRESS_SIZE: usize = 1 << 16;
#[derive(Debug)]
pub struct DecodingError(String);
impl std::fmt::Display for DecodingError {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter.write_str(&self.0)
}
}
impl std::error::Error for DecodingError {}
impl From<DecompressError> for DecodingError {
fn from(error: DecompressError) -> Self {
Self(match (error.needs_dictionary(), error.message()) {
(Some(_), _) => "Error 2 while decompressing data".to_owned(),
(None, Some(message)) => format!("Error -3 while decompressing data: {message}"),
(None, None) => "Error -3 while decompressing data: invalid input data".to_owned(),
})
}
}
enum Codec {
Gzip,
Deflate,
}
enum FirstAttempt {
Unread,
Reading(Vec<u8>),
Over,
}
struct Inflater {
zlib: Decompress,
first_attempt: FirstAttempt,
pending: Vec<u8>,
saturated: bool,
stalled: bool,
finished: bool,
}
impl Inflater {
fn new(codec: Codec) -> Self {
let (zlib, first_attempt) = match codec {
Codec::Gzip => (Decompress::new_gzip(MAX_WINDOW_BITS), FirstAttempt::Over),
Codec::Deflate => (Decompress::new(true), FirstAttempt::Unread),
};
Self {
zlib,
first_attempt,
pending: Vec::new(),
saturated: false,
stalled: false,
finished: false,
}
}
fn push(&mut self, data: &[u8]) {
self.first_attempt = match self.first_attempt {
FirstAttempt::Unread => FirstAttempt::Reading(data.to_vec()),
FirstAttempt::Reading(_) | FirstAttempt::Over => FirstAttempt::Over,
};
if self.finished {
return;
}
self.pending.extend_from_slice(data);
self.stalled = false;
}
fn has_work(&self) -> bool {
!self.finished && (self.saturated || (!self.pending.is_empty() && !self.stalled))
}
fn step(&mut self) -> Result<Vec<u8>, DecodingError> {
match (self.inflate(), &mut self.first_attempt) {
(Err(_), FirstAttempt::Reading(first)) => {
self.pending = std::mem::take(first);
self.first_attempt = FirstAttempt::Over;
self.zlib = Decompress::new_with_window_bits(false, MAX_WINDOW_BITS);
self.saturated = false;
self.stalled = false;
self.step()
}
(inflated, attempt) => {
if self.pending.is_empty() && !self.saturated {
*attempt = FirstAttempt::Over;
}
inflated
}
}
}
fn inflate(&mut self) -> Result<Vec<u8>, DecodingError> {
let mut output = Vec::with_capacity(MAX_DECOMPRESS_SIZE);
let (read, written) = (self.zlib.total_in(), self.zlib.total_out());
let status = self
.zlib
.decompress_vec(&self.pending, &mut output, FlushDecompress::Sync)?;
self.pending.drain(..(self.zlib.total_in() - read) as usize);
self.finished = matches!(status, Status::StreamEnd);
if self.finished {
self.pending.clear();
}
self.saturated = !self.finished && output.len() == MAX_DECOMPRESS_SIZE;
self.stalled = self.zlib.total_in() == read && self.zlib.total_out() == written;
Ok(output)
}
}
struct Decoder(Vec<Inflater>);
impl Decoder {
fn detect(headers: &HeaderMap) -> Option<Self> {
let inflaters: Vec<Inflater> = header_pairs(headers)
.into_iter()
.filter(|(name, _)| name.as_str() == CONTENT_ENCODING.as_str())
.flat_map(|(_, value)| {
value
.split(',')
.map(|coding| coding.trim().to_lowercase())
.collect::<Vec<_>>()
})
.filter_map(|coding| match coding.as_str() {
"gzip" => Some(Codec::Gzip),
"deflate" => Some(Codec::Deflate),
_ => None,
})
.rev()
.map(Inflater::new)
.collect();
(!inflaters.is_empty()).then_some(Self(inflaters))
}
fn feed(&mut self, chunk: &[u8]) {
self.0[0].push(chunk);
}
fn next(&mut self) -> Result<Option<Bytes>, DecodingError> {
self.pull(self.0.len() - 1)
.map(|decoded| decoded.map(Bytes::from))
}
fn pull(&mut self, layer: usize) -> Result<Option<Vec<u8>>, DecodingError> {
loop {
if self.0[layer].has_work() {
let decoded = self.0[layer].step()?;
if !decoded.is_empty() {
return Ok(Some(decoded));
}
continue;
}
let Some(outer) = layer.checked_sub(1) else {
return Ok(None);
};
match self.pull(outer)? {
Some(decoded) => self.0[layer].push(&decoded),
None => return Ok(None),
}
}
}
}
pub(crate) fn header_pairs(headers: &HeaderMap) -> Vec<(String, String)> {
let utf8 = headers
.values()
.all(|value| std::str::from_utf8(value.as_bytes()).is_ok());
headers
.iter()
.map(|(name, value)| {
let bytes = value.as_bytes();
let text = if utf8 {
String::from_utf8_lossy(bytes).into_owned()
} else {
bytes.iter().copied().map(char::from).collect()
};
(name.to_string(), text)
})
.collect()
}
pub(crate) fn decoded(response: reqwest::Response) -> reqwest::Response {
let Some(decoder) = Decoder::detect(response.headers()) else {
return response;
};
let status = response.status();
let version = response.version();
let headers = response.headers().clone();
let body = futures_util::stream::unfold(
Some((response.bytes_stream().boxed(), decoder)),
|state| async move {
let (mut body, mut decoder) = state?;
loop {
match decoder.next() {
Ok(Some(decoded)) => return Some((Ok(decoded), Some((body, decoder)))),
Ok(None) => {}
Err(error) => return Some((Err(BoxError::from(error)), None)),
}
match body.next().await? {
Ok(chunk) => decoder.feed(&chunk),
Err(error) => return Some((Err(BoxError::from(error)), None)),
}
}
},
);
let mut rebuilt = http::Response::new(reqwest::Body::wrap_stream(body));
*rebuilt.status_mut() = status;
*rebuilt.version_mut() = version;
*rebuilt.headers_mut() = headers;
reqwest::Response::from(rebuilt)
}
#[cfg(test)]
mod tests {
use std::io::Write;
use flate2::{Compression, write::GzEncoder, write::ZlibEncoder};
use futures_util::TryStreamExt;
use rstest::rstest;
use reqwest::header::{HeaderMap, HeaderValue};
use super::{decoded, header_pairs};
fn encoded(encoding: &str, body: &[u8]) -> Vec<u8> {
match encoding {
"gzip" => {
let mut encoder = GzEncoder::new(Vec::new(), Compression::default());
encoder.write_all(body).unwrap();
encoder.finish().unwrap()
}
"deflate" => {
let mut encoder = ZlibEncoder::new(Vec::new(), Compression::default());
encoder.write_all(body).unwrap();
encoder.finish().unwrap()
}
_ => body.to_vec(),
}
}
fn response(encoding: &str, body: Vec<u8>) -> reqwest::Response {
let length = body.len().to_string();
let mut response = http::Response::new(reqwest::Body::from(body));
let headers = response.headers_mut();
headers.insert("content-encoding", encoding.parse().unwrap());
headers.insert("content-length", length.parse().unwrap());
reqwest::Response::from(response)
}
#[rstest]
#[case::gzip("gzip")]
#[case::deflate("deflate")]
#[case::unknown("br")]
#[tokio::test]
async fn decoded_bodies_keep_the_head_the_provider_sent(#[case] encoding: &str) {
let body = br#"{"type":"message"}"#;
let wire = encoded(encoding, body);
let length = wire.len().to_string();
let response = decoded(response(encoding, wire.clone()));
assert_eq!(response.headers()["content-encoding"], encoding);
assert_eq!(response.headers()["content-length"], length.as_str());
let expected = if encoding == "br" {
wire
} else {
body.to_vec()
};
assert_eq!(
response.bytes().await.unwrap().as_ref(),
expected.as_slice()
);
}
fn raw_deflate(body: &[u8]) -> Vec<u8> {
let mut encoder = flate2::write::DeflateEncoder::new(Vec::new(), Compression::default());
encoder.write_all(body).unwrap();
encoder.finish().unwrap()
}
fn chunked(encoding: &str, wire: Vec<u8>, size: usize) -> reqwest::Response {
let chunks: Vec<Result<Vec<u8>, std::io::Error>> =
wire.chunks(size).map(|chunk| Ok(chunk.to_vec())).collect();
let mut response = http::Response::new(reqwest::Body::wrap_stream(
futures_util::stream::iter(chunks),
));
response
.headers_mut()
.insert("content-encoding", encoding.parse().unwrap());
reqwest::Response::from(response)
}
async fn read(response: reqwest::Response) -> Result<Vec<u8>, String> {
response
.bytes()
.await
.map(|bytes| bytes.to_vec())
.map_err(|error| {
litellm_error::find::<super::DecodingError>(&error)
.map_or_else(|| error.to_string(), ToString::to_string)
})
}
const BODY: &[u8] = br#"{"type":"message","content":[{"type":"text","text":"Paris"}]}"#;
#[rstest]
#[case::padded_upper_case(" GZIP ", encoded("gzip", BODY))]
#[case::listed_with_an_unknown_coding("gzip, br", encoded("gzip", BODY))]
#[case::applied_twice("gzip, gzip", encoded("gzip", &encoded("gzip", BODY)))]
#[case::deflate_without_its_zlib_header("deflate", raw_deflate(BODY))]
#[case::gzip_with_data_past_its_end("gzip", [encoded("gzip", BODY), b"\x1f\x8b trailing".to_vec()].concat())]
#[tokio::test]
async fn bodies_decode_as_httpx_decodes_them(
#[case] encoding: &str,
#[case] wire: Vec<u8>,
#[values(7, 4096)] chunk: usize,
) {
assert_eq!(
read(decoded(chunked(encoding, wire, chunk))).await.unwrap(),
BODY
);
}
#[rstest]
#[case::x_gzip_alias(" X-Gzip ")]
#[case::unknown_only("br, identity")]
#[tokio::test]
async fn codings_httpx_has_no_decoder_for_pass_through(#[case] encoding: &str) {
let wire = encoded("gzip", BODY);
assert_eq!(
read(decoded(chunked(encoding, wire.clone(), 4096)))
.await
.unwrap(),
wire
);
}
#[rstest]
#[case::not_deflate("deflate", b"not deflate at all".to_vec(), "Error -3 while decompressing data: invalid block type")]
#[case::bad_gzip_trailer(
"gzip",
{ let wire = encoded("gzip", BODY); [&wire[..wire.len() - 8], &[0; 8]].concat() },
"Error -3 while decompressing data: incorrect data check"
)]
#[case::not_gzip("gzip", b"\x1f\x8b\x08\x00not really gzip data at all".to_vec(), "Error -3 while decompressing data: invalid stored block lengths")]
#[case::no_gzip_magic("gzip", b"plain text".to_vec(), "Error -3 while decompressing data: incorrect header check")]
#[tokio::test]
async fn corrupt_bodies_fail_with_zlib_s_error(
#[case] encoding: &str,
#[case] wire: Vec<u8>,
#[case] expected: &str,
) {
assert_eq!(
read(decoded(chunked(encoding, wire, 4096)))
.await
.unwrap_err(),
expected
);
}
#[rstest]
#[case::gzip_in_gzip("gzip, gzip")]
#[case::deflate_in_gzip("deflate, gzip")]
#[tokio::test]
async fn a_nested_bomb_is_decoded_in_bounded_reads(#[case] encoding: &str) {
let zeros = vec![0; 8 << 20];
let inner = if encoding.starts_with("deflate") {
encoded("deflate", &zeros)
} else {
encoded("gzip", &zeros)
};
let wire = encoded("gzip", &inner);
let sizes: Vec<usize> = decoded(chunked(encoding, wire, 4096))
.bytes_stream()
.map_ok(|chunk| chunk.len())
.try_collect()
.await
.unwrap();
assert!(
sizes.iter().all(|size| *size <= super::MAX_DECOMPRESS_SIZE),
"{sizes:?}"
);
assert_eq!(sizes.iter().sum::<usize>(), zeros.len());
}
#[rstest]
#[tokio::test]
async fn a_truncated_body_yields_what_it_holds_without_an_error() {
let wire = encoded("gzip", BODY);
let decoded_body = read(decoded(chunked(
"gzip",
wire[..wire.len() - 8].to_vec(),
4096,
)))
.await
.unwrap();
assert_eq!(decoded_body, BODY);
}
#[rstest]
#[case::ascii(&[b"plain".as_slice()], &["plain"])]
#[case::utf_8(&[b"plain".as_slice(), b"caf\xc3\xa9"], &["plain", "caf\u{e9}"])]
#[case::latin_1_once_any_value_is_not_utf_8(
&[b"caf\xc3\xa9".as_slice(), b"caf\xe9"],
&["caf\u{c3}\u{a9}", "caf\u{e9}"]
)]
fn header_values_decode_as_httpx_decodes_them(
#[case] values: &[&[u8]],
#[case] expected: &[&str],
) {
let headers: HeaderMap = values
.iter()
.enumerate()
.map(|(index, value)| {
(
format!("x-fz-{index}")
.parse::<reqwest::header::HeaderName>()
.unwrap(),
HeaderValue::from_bytes(value).unwrap(),
)
})
.collect();
let decoded: Vec<String> = header_pairs(&headers)
.into_iter()
.map(|(_, value)| value)
.collect();
assert_eq!(decoded, expected);
}
}

968
litellm-rust/crates/inference-messages/src/handler.rs
View file

File diff suppressed because it is too large Load diff

11
litellm-rust/crates/inference-messages/src/host_answer.rs Normal file
View file

@ -0,0 +1,11 @@
use futures_util::future::BoxFuture;
use crate::Error;
pub trait HostAnswer: Send + Sync {
fn answer(&self, stream: bool) -> BoxFuture<'_, Result<(), Error>>;
fn intercept(&self, stream: bool) -> BoxFuture<'_, Result<bool, Error>> {
Box::pin(async move { self.answer(stream).await.map(|()| true) })
}
}

141
litellm-rust/crates/inference-messages/src/lib.rs
View file

@ -1,25 +1,41 @@
mod common_utils;
mod constants;
mod encoding;
mod handler;
mod host_answer;
mod pre_request;
mod reads;
pub mod route;
mod safeguards;
mod transport;
mod types;
mod workload_identity;
use futures_util::FutureExt;
use litellm_auth::AuthServices;
use litellm_host::interceptors::Interceptors;
use litellm_inference::context::CallContext;
use litellm_inference::provider::resolve_llm_provider;
use litellm_llms_types::recognized::Recognized;
use litellm_secrets::source::SecretSource;
use std::sync::Arc;
use crate::common_utils::messages_provider;
pub use common_utils::merge_messages_header_sources;
pub use encoding::DecodingError;
pub use host_answer::HostAnswer;
pub use litellm_inference::Error;
pub use pre_request::PreRequestHooks;
pub use safeguards::Safeguards;
pub use transport::{
Delivered, Transport, TransportBody, TransportRequest, TransportResponse, UnencodableHeader,
};
pub use types::{
MessagesBody, MessagesCall, MessagesCallResponse, MessagesSettings, MessagesShaping,
litellm_params, messages_body,
};
pub use workload_identity::WorkloadIdentity;
#[derive(Clone)]
pub struct MessagesRoute {
@ -27,6 +43,11 @@ pub struct MessagesRoute {
auth: Arc<AuthServices>,
secrets: Arc<dyn SecretSource>,
cache: Option<litellm_cache_response::ScopedCache>,
workload_identity: Option<Arc<dyn WorkloadIdentity>>,
transport: Option<Arc<dyn Transport>>,
pre_request_hooks: Option<Arc<dyn PreRequestHooks>>,
safeguards: Option<Arc<dyn Safeguards>>,
host_answer: Option<Arc<dyn HostAnswer>>,
}
impl MessagesRoute {
@ -40,6 +61,51 @@ impl MessagesRoute {
auth,
secrets,
cache: None,
workload_identity: None,
transport: None,
pre_request_hooks: None,
safeguards: None,
host_answer: None,
}
}
#[must_use]
pub fn with_workload_identity(self, workload_identity: Arc<dyn WorkloadIdentity>) -> Self {
Self {
workload_identity: Some(workload_identity),
..self
}
}
#[must_use]
pub fn with_transport(self, transport: Arc<dyn Transport>) -> Self {
Self {
transport: Some(transport),
..self
}
}
#[must_use]
pub fn with_pre_request_hooks(self, hooks: Arc<dyn PreRequestHooks>) -> Self {
Self {
pre_request_hooks: Some(hooks),
..self
}
}
#[must_use]
pub fn with_safeguards(self, safeguards: Arc<dyn Safeguards>) -> Self {
Self {
safeguards: Some(safeguards),
..self
}
}
#[must_use]
pub fn with_host_answer(self, host_answer: Arc<dyn HostAnswer>) -> Self {
Self {
host_answer: Some(host_answer),
..self
}
}
@ -67,7 +133,7 @@ impl MessagesRoute {
model = %call.body.model,
provider,
resolved_model,
stream = call.body.params.stream == Some(true),
stream = call.body.params.streams(),
outcome
))]
async fn run(
@ -76,24 +142,69 @@ impl MessagesRoute {
context: CallContext<'_, impl Interceptors<Error>>,
) -> Result<MessagesCallResponse, Error> {
litellm_inference::diagnostic::call(async {
let provider = resolve_llm_provider(
&call.body.model,
call.custom_llm_provider.as_deref(),
"messages",
)?;
let config = messages_provider(provider.provider, provider.model).ok_or_else(|| {
Error::InvalidProvider(<&str>::from(provider.provider).to_string())
})?;
litellm_inference::diagnostic::provider(
provider.model,
<&'static str>::from(provider.provider),
);
let model = provider.model.to_owned();
let requested_stream = call.body.params.streams();
let client_streams = call.body.params.stream == Some(Recognized::Known(true));
let call = match &self.pre_request_hooks {
None => call,
Some(hooks) => {
let (config, model, _) = Self::provider(&call)?;
if let Some(output) = handler::cached(self, config, &model, &context).await? {
return Ok(output);
}
hooks.run(call).await?
}
};
let (config, model, provider) = Self::provider(&call)?;
litellm_inference::diagnostic::provider(&model, provider);
if self.pre_request_hooks.is_none()
&& let Some(output) = handler::cached(self, config, &model, &context).await?
{
return Ok(output);
}
if let Some(host_answer) = self
.host_answer
.as_deref()
.filter(|_| call.shaping.settings.intercepted)
&& host_answer.intercept(requested_stream).await?
{
return handler::deliver_answer(self, config, model, requested_stream, &context)
.await;
}
if let Some(mock) = call.shaping.settings.mock_response {
return Ok(
handler::mocked(self, config, model, &context, mock, requested_stream).await,
);
}
if let Some(host_answer) = self
.host_answer
.as_deref()
.filter(|_| call.shaping.settings.mcp_gateway)
{
let stream = call.body.params.streams();
return handler::answered(self, host_answer, config, model, stream, &context).await;
}
let secrets = self.secrets.resolve(config.config().secret_names()).await?;
handler::execute(self, config, model, call, secrets, &context)
handler::execute(self, config, model, call, secrets, client_streams, &context)
.boxed()
.await
})
.await
}
fn provider(
call: &MessagesCall,
) -> Result<(common_utils::MessagesProvider, String, &'static str), Error> {
let provider = resolve_llm_provider(
&call.body.model,
call.custom_llm_provider.as_deref(),
"messages",
)?;
let config = messages_provider(provider.provider, provider.model)
.ok_or_else(|| Error::InvalidProvider(<&str>::from(provider.provider).to_string()))?;
Ok((
config,
provider.model.to_owned(),
<&'static str>::from(provider.provider),
))
}
}

7
litellm-rust/crates/inference-messages/src/pre_request.rs Normal file
View file

@ -0,0 +1,7 @@
use futures_util::future::BoxFuture;
use super::{Error, MessagesCall};
pub trait PreRequestHooks: Send + Sync {
fn run(&self, call: MessagesCall) -> BoxFuture<'_, Result<MessagesCall, Error>>;
}

394
litellm-rust/crates/inference-messages/src/reads.rs Normal file
View file

@ -0,0 +1,394 @@
use std::{collections::VecDeque, sync::Arc};
use bytes::Bytes;
use futures_util::{FutureExt, Stream, StreamExt, stream::BoxStream};
use std::sync::Mutex;
use tokio::{sync::Notify, task::JoinHandle};
const READ_SIZE: usize = 16 * 1024;
const LOW_WATER: usize = 64 * 1024;
const HIGH_WATER: usize = 2 * LOW_WATER;
const SETTLE_YIELDS: usize = 64;
struct Buffer<E> {
chunks: VecDeque<Bytes>,
size: usize,
quiet: bool,
failure: Option<E>,
ended: bool,
}
impl<E> Buffer<E> {
fn take(&mut self) -> Bytes {
let mut joined = Vec::new();
while let Some(mut chunk) = self.chunks.pop_front() {
let room = READ_SIZE - joined.len();
if chunk.len() > room {
let rest = chunk.split_off(room);
self.chunks.push_front(rest);
}
if joined.is_empty() && (chunk.len() == READ_SIZE || self.chunks.is_empty()) {
self.size -= chunk.len();
return chunk;
}
joined.extend_from_slice(&chunk);
self.size -= chunk.len();
if joined.len() == READ_SIZE {
break;
}
}
Bytes::from(joined)
}
}
struct Reads<E> {
buffer: Mutex<Buffer<E>>,
received: Notify,
drained: Notify,
}
enum Read<E> {
Chunk(Bytes),
Failure(E),
End,
}
impl<E> Reads<E> {
fn readable(&self) -> bool {
let buffer = self.buffer();
buffer.failure.is_some()
|| buffer.ended
|| (!buffer.chunks.is_empty() && (buffer.quiet || buffer.size >= READ_SIZE))
}
fn quiet(&self) {
self.buffer().quiet = true;
self.received.notify_one();
}
fn read(&self) -> Option<Read<E>> {
let mut buffer = self.buffer();
if !buffer.chunks.is_empty() {
let chunk = buffer.take();
drop(buffer);
self.drained.notify_one();
return Some(Read::Chunk(chunk));
}
if let Some(failure) = buffer.failure.take() {
return Some(Read::Failure(failure));
}
buffer.ended.then_some(Read::End)
}
fn buffer(&self) -> std::sync::MutexGuard<'_, Buffer<E>> {
self.buffer
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
}
}
struct Reader(JoinHandle<()>);
impl Drop for Reader {
fn drop(&mut self) {
self.0.abort();
}
}
pub(crate) fn joined_reads<E: Send + 'static>(
chunks: impl Stream<Item = Result<Bytes, E>> + Send + 'static,
) -> BoxStream<'static, Result<Bytes, E>> {
let Ok(runtime) = tokio::runtime::Handle::try_current() else {
return chunks.boxed();
};
let reads = Arc::new(Reads {
buffer: Mutex::new(Buffer {
chunks: VecDeque::new(),
size: 0,
quiet: false,
failure: None,
ended: false,
}),
received: Notify::new(),
drained: Notify::new(),
});
let reader = Reader(runtime.spawn(fill(chunks.boxed(), Arc::clone(&reads))));
futures_util::stream::unfold(Some((reads, reader)), |state| async move {
let (reads, reader) = state?;
loop {
let received = reads.received.notified();
if reads.readable() {
break;
}
received.await;
}
let read = reads.read()?;
match read {
Read::Chunk(chunk) => Some((Ok(chunk), Some((reads, reader)))),
Read::Failure(failure) => Some((Err(failure), None)),
Read::End => None,
}
})
.boxed()
}
async fn fill<E>(mut chunks: BoxStream<'static, Result<Bytes, E>>, reads: Arc<Reads<E>>) {
loop {
let next = match settled_next(&mut chunks).await {
Some(next) => next,
None => {
reads.quiet();
chunks.next().await
}
};
let paused = {
let mut buffer = reads.buffer();
buffer.quiet = false;
match next {
Some(Ok(chunk)) => {
buffer.size += chunk.len();
if !chunk.is_empty() {
buffer.chunks.push_back(chunk);
}
}
Some(Err(failure)) => {
buffer.failure = Some(failure);
buffer.ended = true;
}
None => buffer.ended = true,
}
if buffer.ended {
drop(buffer);
reads.received.notify_one();
return;
}
buffer.size > HIGH_WATER
};
reads.received.notify_one();
if paused {
reads.quiet();
loop {
let drained = reads.drained.notified();
if reads.buffer().size < LOW_WATER {
break;
}
drained.await;
}
}
}
}
async fn settled_next<E>(
chunks: &mut BoxStream<'static, Result<Bytes, E>>,
) -> Option<Option<Result<Bytes, E>>> {
for _ in 0..SETTLE_YIELDS {
if let Some(next) = chunks.next().now_or_never() {
return Some(next);
}
tokio::task::yield_now().await;
}
None
}
#[cfg(test)]
mod tests {
use std::{
pin::Pin,
sync::atomic::{AtomicUsize, Ordering},
task::{Context, Poll},
time::Duration,
};
use futures_util::{FutureExt, StreamExt, TryStreamExt};
use tokio::sync::mpsc;
use super::*;
struct Upstream {
receiver: mpsc::UnboundedReceiver<Result<Bytes, &'static str>>,
pulled: Arc<AtomicUsize>,
}
impl Stream for Upstream {
type Item = Result<Bytes, &'static str>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let next = self.receiver.poll_recv(cx);
if let Poll::Ready(Some(_)) = next {
self.pulled.fetch_add(1, Ordering::SeqCst);
}
next
}
}
type Item = Result<Bytes, &'static str>;
fn upstream() -> (
mpsc::UnboundedSender<Item>,
BoxStream<'static, Item>,
Arc<AtomicUsize>,
) {
let (sender, receiver) = mpsc::unbounded_channel();
let pulled = Arc::new(AtomicUsize::new(0));
let reads = joined_reads(Upstream {
receiver,
pulled: Arc::clone(&pulled),
});
(sender, reads, pulled)
}
async fn settled() {
tokio::time::sleep(Duration::from_millis(20)).await;
}
#[rstest::rstest]
#[tokio::test]
async fn chunks_that_arrived_together_are_read_as_one() {
let (sender, mut reads, _) = upstream();
sender.send(Ok(Bytes::from("event: a\n\n"))).unwrap();
sender.send(Ok(Bytes::from("event: b\n\n"))).unwrap();
sender.send(Ok(Bytes::from("event: c\n\n"))).unwrap();
settled().await;
assert_eq!(
reads.try_next().await.unwrap().unwrap(),
"event: a\n\nevent: b\n\nevent: c\n\n"
);
sender.send(Ok(Bytes::from("event: d\n\n"))).unwrap();
settled().await;
assert_eq!(reads.try_next().await.unwrap().unwrap(), "event: d\n\n");
drop(sender);
assert!(reads.try_next().await.unwrap().is_none());
}
#[rstest::rstest]
#[tokio::test]
async fn chunks_still_being_handed_over_join_the_read() {
let handed_over = futures_util::stream::unfold(0, |sent| async move {
tokio::task::yield_now().await;
(sent < 50).then(|| (Ok::<_, &str>(Bytes::from("event: ping\n\n")), sent + 1))
});
let read: Vec<usize> = joined_reads(handed_over)
.map_ok(|chunk| chunk.len())
.try_collect()
.await
.unwrap();
assert_eq!(read, [50 * "event: ping\n\n".len()]);
}
#[rstest::rstest]
#[tokio::test]
async fn a_read_returns_the_first_chunk_without_waiting_for_more() {
let (sender, mut reads, _) = upstream();
let first = tokio::spawn(async move { reads.try_next().await });
settled().await;
sender.send(Ok(Bytes::from("event: a\n\n"))).unwrap();
assert_eq!(first.await.unwrap().unwrap().unwrap(), "event: a\n\n");
}
#[rstest::rstest]
#[tokio::test]
async fn reads_are_cut_at_sixteen_kibibytes() {
let (sender, reads, _) = upstream();
sender.send(Ok(Bytes::from(vec![b'a'; 10_000]))).unwrap();
sender.send(Ok(Bytes::from(vec![b'b'; 30_000]))).unwrap();
drop(sender);
settled().await;
let sizes: Vec<usize> = reads
.map_ok(|chunk| chunk.len())
.try_collect()
.await
.unwrap();
assert_eq!(sizes, [READ_SIZE, READ_SIZE, 40_000 - 2 * READ_SIZE]);
}
#[rstest::rstest]
#[tokio::test]
async fn the_bytes_still_buffered_are_read_before_a_failure() {
let (sender, mut reads, _) = upstream();
sender.send(Ok(Bytes::from("event: a\n\n"))).unwrap();
sender.send(Ok(Bytes::from("event: b\n\n"))).unwrap();
sender.send(Err("reset")).unwrap();
settled().await;
assert_eq!(
reads.try_next().await.unwrap().unwrap(),
"event: a\n\nevent: b\n\n"
);
assert_eq!(reads.try_next().await.unwrap_err(), "reset");
assert!(reads.next().await.is_none());
}
#[rstest::rstest]
#[tokio::test]
async fn reading_pauses_above_the_high_water_mark_until_below_the_low_one() {
let (sender, mut reads, pulled) = upstream();
std::iter::repeat_with(|| Bytes::from(vec![b'a'; READ_SIZE]))
.take(12)
.try_for_each(|chunk| sender.send(Ok(chunk)))
.unwrap();
settled().await;
assert_eq!(pulled.load(Ordering::SeqCst), 9);
let read: Vec<Bytes> = (&mut reads).take(5).try_collect().await.unwrap();
assert_eq!(read.len(), 5);
settled().await;
assert_eq!(pulled.load(Ordering::SeqCst), 9);
reads.try_next().await.unwrap();
settled().await;
assert_eq!(pulled.load(Ordering::SeqCst), 12);
}
async fn chunked_upstream(frames: usize) -> String {
use tokio::io::{AsyncReadExt, AsyncWriteExt};
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.unwrap();
let mut request = [0_u8; 4096];
let _ = socket.read(&mut request).await.unwrap();
let frame = "event: ping\ndata: {\"type\": \"ping\"}\n\n";
let body: String =
std::iter::repeat_n(format!("{:x}\r\n{frame}\r\n", frame.len()), frames).collect();
let response = format!(
"HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ntransfer-encoding: chunked\r\nconnection: close\r\n\r\n{body}0\r\n\r\n"
);
socket.write_all(response.as_bytes()).await.unwrap();
tokio::time::sleep(Duration::from_secs(1)).await;
});
format!("http://{address}")
}
#[rstest::rstest]
#[case::fifty(50)]
#[case::a_thousand(1000)]
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn http_chunks_written_together_are_read_together(#[case] frames: usize) {
let url = chunked_upstream(frames).await;
let response = litellm_http::Client::for_test()
.get(url)
.send()
.await
.unwrap();
tokio::time::sleep(Duration::from_millis(50)).await;
let sizes: Vec<usize> = joined_reads(response.bytes_stream())
.map_ok(|chunk| chunk.len())
.try_collect()
.await
.unwrap();
let total = frames * "event: ping\ndata: {\"type\": \"ping\"}\n\n".len();
let expected: Vec<usize> = std::iter::repeat_n(READ_SIZE, total / READ_SIZE)
.chain((!total.is_multiple_of(READ_SIZE)).then_some(total % READ_SIZE))
.collect();
assert_eq!(sizes, expected);
}
#[rstest::rstest]
fn without_a_runtime_the_chunks_pass_through() {
let chunks =
futures_util::stream::iter([Ok::<_, &str>(Bytes::from("a")), Ok(Bytes::from("b"))]);
let read: Vec<Bytes> = joined_reads(chunks)
.try_collect()
.now_or_never()
.unwrap()
.unwrap();
assert_eq!(read, [Bytes::from("a"), Bytes::from("b")]);
}
}

99
litellm-rust/crates/inference-messages/src/route.rs
View file

@ -74,4 +74,103 @@ impl litellm_inference::caching::StreamCachable for Messages {
fn bytes(chunk: &Self::Chunk) -> &[u8] {
chunk.as_ref()
}
fn storable_stream(text: &str) -> bool {
let lines: Vec<&str> = sse_lines(text).collect();
lines.contains(&"event: message_stop")
&& !carries_provider_error(&lines)
&& lines.iter().any(|line| line_carries_output(line))
}
}
fn sse_lines(text: &str) -> impl Iterator<Item = &str> {
let mut rest = Some(text);
std::iter::from_fn(move || {
let current = rest?;
match current.find(['\r', '\n']) {
Some(at) => {
let skip = if current[at..].starts_with("\r\n") {
2
} else {
1
};
rest = Some(&current[at + skip..]);
Some(&current[..at])
}
None => {
rest = None;
(!current.is_empty()).then_some(current)
}
}
})
}
fn sse_data(line: &str, stripped: bool) -> Option<serde_json::Map<String, serde_json::Value>> {
let data = line.strip_prefix("data:")?;
let data = if stripped {
data.trim_matches(|c: char| matches!(c, ' ' | '\t' | '\n' | '\r' | '\x0b' | '\x0c'))
} else {
data
};
match serde_json::from_str(data) {
Ok(serde_json::Value::Object(map)) => Some(map),
_ => None,
}
}
fn carries_provider_error(lines: &[&str]) -> bool {
lines
.iter()
.filter_map(|line| sse_data(line, true))
.find(|event| event.get("type").and_then(serde_json::Value::as_str) == Some("error"))
.is_some_and(|event| event.get("error").is_some_and(serde_json::Value::is_object))
}
fn line_carries_output(line: &str) -> bool {
let Some(event) = sse_data(line, false) else {
return false;
};
let block = match event.get("type").and_then(serde_json::Value::as_str) {
Some("content_block_start") => event.get("content_block"),
Some("content_block_delta") => event.get("delta"),
_ => None,
};
block
.and_then(serde_json::Value::as_object)
.is_some_and(|fields| {
fields.iter().any(|(name, value)| {
name != "type"
&& !matches!(value, serde_json::Value::Null)
&& value.as_str() != Some("")
&& value.as_array().is_none_or(|items| !items.is_empty())
})
})
}
#[cfg(test)]
mod tests {
use litellm_inference::caching::StreamCachable;
use rstest::rstest;
use super::Messages;
const TEXT: &str = "event: message_start\ndata: {\"type\":\"message_start\"}\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"Paris\"}}\n\n";
const EMPTY: &str = "event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"\"}}\n\n";
const STOP: &str = "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n";
#[rstest]
#[case::complete(format!("{TEXT}{STOP}"), true)]
#[case::crlf_lines(format!("{TEXT}{STOP}").replace('\n', "\r\n"), true)]
#[case::done_sentinel(format!("{TEXT}{STOP}data: [DONE]\n\n"), true)]
#[case::data_only_frames(format!("{TEXT}{STOP}").replace("event: message_stop\n", ""), false)]
#[case::event_without_space(format!("{TEXT}{STOP}").replace("event: message_stop", "event:message_stop"), false)]
#[case::provider_error(format!("{TEXT}{STOP}event: error\ndata: {{\"type\":\"error\",\"error\":{{\"type\":\"overloaded_error\"}}}}\n\n"), false)]
#[case::error_without_body(format!("{TEXT}{STOP}data: {{\"type\":\"error\",\"error\":\"overloaded\"}}\n\n"), true)]
#[case::no_output(format!("{EMPTY}{STOP}"), false)]
fn a_stream_is_stored_by_python_s_message_stop_error_and_output_tests(
#[case] text: String,
#[case] stored: bool,
) {
assert_eq!(Messages::storable_stream(&text), stored);
}
}

248
litellm-rust/crates/inference-messages/src/safeguards.rs Normal file
View file

@ -0,0 +1,248 @@
use std::{collections::VecDeque, sync::Arc};
use bytes::Bytes;
use futures_util::{StreamExt, TryStreamExt, future::BoxFuture, stream::BoxStream};
use serde_json::Value;
use super::{Error, MessagesBody, MessagesCallResponse};
const SAFEGUARD_RESULTS: &str = "safeguard_results";
pub trait Safeguards: Send + Sync {
fn results<'a>(
&'a self,
response: &'a MessagesBody,
) -> BoxFuture<'a, Result<Option<Value>, Error>>;
fn open_stream(&self) -> BoxFuture<'_, Result<bool, Error>>;
fn frames(&self, chunk: Bytes) -> BoxFuture<'_, Result<Vec<Bytes>, Error>>;
fn flush(&self) -> BoxFuture<'_, Result<Vec<Bytes>, Error>>;
}
pub(super) async fn stamped(
safeguards: Option<&Arc<dyn Safeguards>>,
output: MessagesCallResponse,
) -> Result<MessagesCallResponse, Error> {
let Some(safeguards) = safeguards else {
return Ok(output);
};
match output {
MessagesCallResponse::Complete(response) => {
let results = safeguards.results(&response).await?;
Ok(MessagesCallResponse::Complete(match results {
Some(results) => with_results(response, results),
None => response,
}))
}
MessagesCallResponse::Stream { head, chunks } => {
if !safeguards.open_stream().await? {
return Ok(MessagesCallResponse::Stream { head, chunks });
}
Ok(MessagesCallResponse::Stream {
head,
chunks: stamped_frames(Arc::clone(safeguards), chunks),
})
}
}
}
fn with_results(response: MessagesBody, results: Value) -> MessagesBody {
match response {
MessagesBody::Message(mut message) => {
message.extra.insert(SAFEGUARD_RESULTS.to_owned(), results);
MessagesBody::Message(message)
}
MessagesBody::Unchecked(mut body) => {
body.insert(SAFEGUARD_RESULTS.to_owned(), results);
MessagesBody::Unchecked(body)
}
}
}
struct Framing {
chunks: Option<BoxStream<'static, Result<Bytes, Error>>>,
safeguards: Arc<dyn Safeguards>,
ready: VecDeque<Bytes>,
}
fn stamped_frames(
safeguards: Arc<dyn Safeguards>,
chunks: BoxStream<'static, Result<Bytes, Error>>,
) -> BoxStream<'static, Result<Bytes, Error>> {
let framing = Framing {
chunks: Some(chunks),
safeguards,
ready: VecDeque::new(),
};
futures_util::stream::try_unfold(framing, |mut framing| async move {
loop {
if let Some(frame) = framing.ready.pop_front() {
return Ok(Some((frame, framing)));
}
let Some(mut chunks) = framing.chunks.take() else {
return Ok(None);
};
let frames = match chunks.try_next().await? {
Some(chunk) => {
framing.chunks = Some(chunks);
framing.safeguards.frames(chunk).await?
}
None => framing.safeguards.flush().await?,
};
framing.ready.extend(frames);
}
})
.boxed()
}
#[cfg(test)]
mod tests {
use std::sync::Mutex;
use litellm_llms_types::formats::messages::MessagesResponse;
use rstest::rstest;
use serde_json::{Map, json};
use super::*;
use crate::route::MessagesStreamHead;
struct Scripted {
results: Option<Value>,
streams: bool,
fed: Mutex<Vec<Bytes>>,
}
impl Scripted {
fn new(results: Option<Value>, streams: bool) -> Arc<Self> {
Arc::new(Self {
results,
streams,
fed: Mutex::new(Vec::new()),
})
}
}
impl Safeguards for Scripted {
fn results<'a>(
&'a self,
_: &'a MessagesBody,
) -> BoxFuture<'a, Result<Option<Value>, Error>> {
Box::pin(async move { Ok(self.results.clone()) })
}
fn open_stream(&self) -> BoxFuture<'_, Result<bool, Error>> {
Box::pin(async move { Ok(self.streams) })
}
fn frames(&self, chunk: Bytes) -> BoxFuture<'_, Result<Vec<Bytes>, Error>> {
Box::pin(async move {
self.fed.lock().unwrap().push(chunk.clone());
Ok(chunk
.split_inclusive(|byte| *byte == b';')
.map(Bytes::copy_from_slice)
.collect())
})
}
fn flush(&self) -> BoxFuture<'_, Result<Vec<Bytes>, Error>> {
Box::pin(async move { Ok(vec![Bytes::from_static(b"tail")]) })
}
}
fn message() -> MessagesBody {
let response: MessagesResponse = serde_json::from_value(json!({
"id": "msg_1", "type": "message", "role": "assistant", "model": "m",
"content": [], "stop_reason": "end_turn", "stop_sequence": null,
}))
.unwrap();
MessagesBody::Message(Box::new(response))
}
fn stream(chunks: &[&'static [u8]]) -> MessagesCallResponse {
let chunks: Vec<Result<Bytes, Error>> = chunks
.iter()
.map(|chunk| Ok(Bytes::from_static(chunk)))
.collect();
MessagesCallResponse::Stream {
head: MessagesStreamHead {
headers: Vec::new(),
},
chunks: futures_util::stream::iter(chunks).boxed(),
}
}
async fn collected(output: MessagesCallResponse) -> Vec<Bytes> {
let MessagesCallResponse::Stream { chunks, .. } = output else {
panic!("expected a stream");
};
chunks.try_collect().await.unwrap()
}
#[rstest]
#[case::message_gains_the_results(message(), Some(json!([{"type": "dangerous_tool_use"}])))]
#[case::unchecked_gains_the_results(
MessagesBody::Unchecked(Map::new()),
Some(json!([{"type": "dangerous_tool_use"}]))
)]
#[case::declined_leaves_it(message(), None)]
#[tokio::test]
async fn a_complete_response_carries_what_the_classifier_answered(
#[case] response: MessagesBody,
#[case] results: Option<Value>,
) {
let safeguards: Arc<dyn Safeguards> = Scripted::new(results.clone(), false);
let MessagesCallResponse::Complete(stamped) = stamped(
Some(&safeguards),
MessagesCallResponse::Complete(response.clone()),
)
.await
.unwrap() else {
panic!("expected a complete response");
};
let body = serde_json::to_value(&stamped).unwrap();
assert_eq!(body.get(SAFEGUARD_RESULTS), results.as_ref());
let mut unstamped = body;
unstamped.as_object_mut().unwrap().remove(SAFEGUARD_RESULTS);
assert_eq!(unstamped, serde_json::to_value(&response).unwrap());
}
#[tokio::test]
async fn a_stream_is_reframed_through_the_classifier_and_flushed_at_its_end() {
let scripted = Scripted::new(None, true);
let safeguards: Arc<dyn Safeguards> = scripted.clone();
let frames = collected(
stamped(Some(&safeguards), stream(&[b"a;b", b";c;"]))
.await
.unwrap(),
)
.await;
assert_eq!(
frames,
["a;", "b", ";", "c;", "tail"]
.map(|frame| Bytes::from_static(frame.as_bytes()))
.to_vec()
);
assert_eq!(scripted.fed.lock().unwrap().len(), 2);
}
#[rstest]
#[case::declined(Some(Scripted::new(None, false) as Arc<dyn Safeguards>))]
#[case::not_classified(None)]
#[tokio::test]
async fn a_stream_the_classifier_does_not_answer_passes_through(
#[case] safeguards: Option<Arc<dyn Safeguards>>,
) {
let frames = collected(
stamped(safeguards.as_ref(), stream(&[b"a;b", b";c;"]))
.await
.unwrap(),
)
.await;
assert_eq!(
frames,
[b"a;b".as_slice(), b";c;"].map(Bytes::from_static).to_vec()
);
}
}

59
litellm-rust/crates/inference-messages/src/transport.rs Normal file
View file

@ -0,0 +1,59 @@
use bytes::Bytes;
use futures_util::{future::BoxFuture, stream::BoxStream};
use serde_json::Value;
use crate::Error;
pub enum TransportBody {
Json(Value),
Signed(Bytes),
}
pub struct TransportRequest {
pub url: String,
pub headers: Vec<(String, String)>,
pub body: TransportBody,
pub stream: bool,
}
pub enum Delivered {
Complete(String),
Stream(BoxStream<'static, Result<Bytes, Error>>),
Rejected { text: String, error: Error },
}
pub struct TransportResponse {
pub status: u16,
pub headers: Vec<(String, String)>,
pub delivered: Delivered,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum UnencodableHeader {
Sent(Vec<u8>),
Received(Vec<u8>),
}
impl std::fmt::Display for UnencodableHeader {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter.write_str("a header value is not ASCII, which Python's HTTP client cannot encode")
}
}
impl std::error::Error for UnencodableHeader {}
impl UnencodableHeader {
pub(crate) fn received(headers: &reqwest::header::HeaderMap) -> Option<Self> {
headers
.values()
.find(|value| !value.as_bytes().is_ascii())
.map(|value| Self::Received(value.as_bytes().to_vec()))
}
}
pub trait Transport: Send + Sync {
fn send(
&self,
request: TransportRequest,
) -> BoxFuture<'static, Result<TransportResponse, Error>>;
}

43
litellm-rust/crates/inference-messages/src/types.rs
View file

@ -2,7 +2,10 @@ use std::time::Duration;
use bytes::Bytes;
use litellm_host::call::CallOutput;
use litellm_llms::{ErrorDetail, base_llm::messages::context::MessagesModelCapabilities};
use litellm_llms::{
ErrorDetail,
base_llm::messages::context::{MessagesModelCapabilities, MessagesShapeOptions},
};
use litellm_llms_types::{
formats::messages::{MessagesRequest, MessagesResponse},
headers::ProviderSpecificHeaders,
@ -71,9 +74,38 @@ pub struct MessagesSettings {
pub additional_drop_params: Vec<String>,
#[serde(default)]
pub request_timeout: Option<f64>,
#[serde(default)]
pub mock_response: Option<Map<String, Value>>,
#[serde(default)]
pub whole_sse_frames: bool,
#[serde(default)]
pub joined_reads: bool,
#[serde(default)]
pub presanitized: bool,
#[serde(default)]
pub drops_unsupported_params: bool,
#[serde(default)]
pub intercepted: bool,
#[serde(default)]
pub mcp_gateway: bool,
#[serde(default)]
pub ascii_response_headers: bool,
#[serde(default)]
pub unencodable_request_header: Option<Vec<u8>>,
/// Python's HTTP client encodes every header value it sends as ASCII, so a value it cannot
/// encode fails the call before the send.
#[serde(default)]
pub ascii_request_headers: bool,
}
impl MessagesSettings {
pub fn shape_options(&self) -> MessagesShapeOptions {
MessagesShapeOptions {
reasoning_auto_summary: self.reasoning_auto_summary,
presanitized: self.presanitized,
}
}
pub fn request_timeout(&self) -> Option<Duration> {
self.request_timeout
.and_then(|secs| Duration::try_from_secs_f64(secs).ok())
@ -140,7 +172,8 @@ mod tests {
"drop_params": true,
"reasoning_auto_summary": true,
"additional_drop_params": ["metadata.user_id", "thinking"],
"request_timeout": 1800.0
"request_timeout": 1800.0,
"drops_unsupported_params": true
}),
MessagesShaping {
settings: MessagesSettings {
@ -148,6 +181,8 @@ mod tests {
reasoning_auto_summary: true,
additional_drop_params: vec!["metadata.user_id".to_string(), "thinking".to_string()],
request_timeout: Some(1800.0),
drops_unsupported_params: true,
..MessagesSettings::default()
},
capabilities: MessagesModelCapabilities {
supports_reasoning: true,
@ -185,6 +220,10 @@ mod tests {
serialized["reasoning_auto_summary"],
json!(expected.settings.reasoning_auto_summary)
);
assert_eq!(
serialized["drops_unsupported_params"],
json!(expected.settings.drops_unsupported_params)
);
assert_eq!(
serialized["additional_drop_params"],
json!(expected.settings.additional_drop_params)

11
litellm-rust/crates/inference-messages/src/workload_identity.rs Normal file
View file

@ -0,0 +1,11 @@
use futures_util::future::BoxFuture;
use litellm_router_types::LitellmParams;
use litellm_secret_value::SecretValue;
pub trait WorkloadIdentity: Send + Sync {
fn token<'a>(
&'a self,
litellm_params: &'a LitellmParams,
api_base: Option<&'a str>,
) -> BoxFuture<'a, Result<Option<SecretValue>, litellm_auth::Error>>;
}

132
litellm-rust/crates/inference-messages/tests/caching.rs
View file

@ -301,3 +301,135 @@ async fn provider_reported_cost_is_delivered_once_and_never_charged_for_a_cache_
assert!(matches!(cached.source, ResultSource::Cache { .. }));
assert_eq!(cached.cost, Cost::CacheHit);
}
#[derive(Default)]
struct CallerKeyed {
config: ResponseCacheConfig,
entry: std::sync::Mutex<Option<Value>>,
}
impl ResponseCacheService for CallerKeyed {
fn config(&self) -> &ResponseCacheConfig {
&self.config
}
fn keyed_by_caller(&self) -> bool {
true
}
fn lookup<'a>(
&'a self,
_: &'a litellm_cache_response::ResponseCacheRequest,
_: Duration,
) -> std::pin::Pin<
Box<
dyn std::future::Future<Output = Result<Option<Value>, litellm_cache::Error>>
+ Send
+ 'a,
>,
> {
Box::pin(async move { Ok(self.entry.lock().unwrap().clone()) })
}
fn store<'a>(
&'a self,
_: &'a litellm_cache_response::ResponseCacheRequest,
value: Value,
_: Duration,
) -> std::pin::Pin<
Box<dyn std::future::Future<Output = Result<(), litellm_cache::Error>> + Send + 'a>,
> {
Box::pin(async move {
*self.entry.lock().unwrap() = Some(value);
Ok(())
})
}
}
#[tokio::test]
async fn a_caller_keyed_hit_is_served_before_the_request_is_built_or_intercepted() {
use litellm_inference_messages::{MessagesCall, MessagesCallResponse};
let response = json!({"id": "msg_caller", "type": "message", "role": "assistant", "model": "native-test",
"content": [{"type": "text", "text": "answer"}], "stop_reason": "end_turn", "stop_sequence": null,
"usage": {"input_tokens": 3, "output_tokens": 2}});
let upstream = support::upstream([support::json_response(response)]).await;
let hooks = ChangingHooks::default();
let route = support::messages_route(litellm_inference_testing::no_secrets()).with_cache(
ScopedCache::new(Arc::new(CallerKeyed::default()), CacheScope::Shared),
);
let request = |api_key: Option<&str>| {
MessagesCall {
body: serde_json::from_value(json!({"model": "native-test", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 16})).unwrap(),
api_key: api_key.map(Into::into), api_base: Some(upstream.uri()), custom_llm_provider: Some("anthropic".into()),
litellm_params: Default::default(), extra_headers: None, provider_specific_header: None,
timeout: None, shaping: Default::default(),
}
};
let first = route
.execute(request(Some("key")), &hooks, None)
.await
.unwrap();
let second = route.execute(request(None), &hooks, None).await.unwrap();
let (MessagesCallResponse::Complete(first), MessagesCallResponse::Complete(second)) =
(first, second)
else {
panic!("expected Messages responses");
};
assert_eq!(first, second);
assert_eq!(support::received(&upstream).await.len(), 1);
assert_eq!(hooks.calls.load(Ordering::SeqCst), 1);
let facts = hooks.facts.lock().unwrap();
let [provider, cached] = facts.as_slice() else {
panic!("expected provider and cache results");
};
assert_eq!(provider.source, ResultSource::Provider);
assert!(matches!(cached.source, ResultSource::Cache { .. }));
assert_eq!(cached.cost, Cost::CacheHit);
}
#[rstest]
#[case::unary(false, true)]
#[case::streamed(true, false)]
#[tokio::test]
async fn a_mock_is_stored_under_the_caller_key_only_when_unstreamed(
#[case] stream: bool,
#[case] stored: bool,
) {
use litellm_inference_messages::{
MessagesBody, MessagesCall, MessagesCallResponse, MessagesSettings, MessagesShaping,
};
let mock = json!({"id": "msg_mock", "type": "message", "role": "assistant", "model": "native-test",
"content": [{"type": "text", "text": "mocked"}], "stop_reason": "end_turn", "stop_sequence": null,
"usage": {"input_tokens": 3, "output_tokens": 2}});
let upstream = support::upstream([]).await;
let hooks = ChangingHooks::default();
let service = Arc::new(CallerKeyed::default());
let route = support::messages_route(litellm_inference_testing::no_secrets())
.with_cache(ScopedCache::new(service.clone(), CacheScope::Shared));
let call = MessagesCall {
body: serde_json::from_value(json!({"model": "native-test", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 16, "stream": stream})).unwrap(),
api_key: Some("key".into()), api_base: Some(upstream.uri()), custom_llm_provider: Some("anthropic".into()),
litellm_params: Default::default(), extra_headers: None, provider_specific_header: None,
timeout: None,
shaping: MessagesShaping {
settings: MessagesSettings {
mock_response: serde_json::from_value(mock.clone()).unwrap(),
..Default::default()
},
..Default::default()
},
};
let output = route.execute(call, &hooks, None).await.unwrap();
let MessagesCallResponse::Complete(MessagesBody::Unchecked(body)) = output else {
panic!("expected the mock as a whole response");
};
assert_eq!(Value::Object(body), mock);
assert!(support::received(&upstream).await.is_empty());
assert_eq!(hooks.calls.load(Ordering::SeqCst), 0);
assert!(hooks.facts.lock().unwrap().is_empty());
assert_eq!(service.entry.lock().unwrap().is_some(), stored);
}

99
litellm-rust/crates/inference-messages/tests/messages/gateway.rs Normal file
View file

@ -0,0 +1,99 @@
use std::sync::atomic::{AtomicUsize, Ordering};
use futures_util::future::BoxFuture;
use litellm_inference_messages::{HostAnswer, MessagesRoute};
use rstest::rstest;
use super::*;
#[derive(Default)]
struct CountingGateway(AtomicUsize);
impl HostAnswer for CountingGateway {
fn answer(&self, _stream: bool) -> BoxFuture<'_, Result<(), Error>> {
self.0.fetch_add(1, Ordering::SeqCst);
Box::pin(async { Ok(()) })
}
}
fn gateway_call(call: MessagesCall, api_base: String, mcp_gateway: bool) -> MessagesCall {
MessagesCall {
api_key: Some("sk-ant-test".into()),
api_base: Some(api_base),
shaping: MessagesShaping {
settings: MessagesSettings {
mcp_gateway,
..MessagesSettings::default()
},
..MessagesShaping::default()
},
..call
}
}
async fn run_through(route: MessagesRoute, call: MessagesCall) -> Result<MessagesOutput, Error> {
let host = LocalMessagesHost::new(call);
litellm_host_native::in_process::run_hosted(
route.machine(host.request()?, None),
host.runtime(),
)
.await
}
#[rstest]
#[tokio::test]
async fn a_call_the_host_hands_to_the_mcp_gateway_never_reaches_the_provider(call: MessagesCall) {
let upstream = upstream([message_response()]).await;
let gateway = Arc::new(CountingGateway::default());
let route =
messages_route(Arc::new(RecordingSecrets::empty())).with_host_answer(gateway.clone());
let output = run_through(route, gateway_call(call, upstream.uri(), true))
.await
.unwrap();
let MessagesOutput::Complete(MessagesBody::Unchecked(body)) = output else {
panic!("the host answers in place of the provider");
};
assert!(body.is_empty());
assert_eq!(gateway.0.load(Ordering::SeqCst), 1);
assert!(received(&upstream).await.is_empty());
}
#[rstest]
#[tokio::test]
async fn a_call_the_host_keeps_reaches_the_provider_past_the_gateway(call: MessagesCall) {
let upstream = upstream([message_response()]).await;
let gateway = Arc::new(CountingGateway::default());
let route =
messages_route(Arc::new(RecordingSecrets::empty())).with_host_answer(gateway.clone());
let output = run_through(route, gateway_call(call, upstream.uri(), false))
.await
.unwrap();
assert!(matches!(
output,
MessagesOutput::Complete(MessagesBody::Message(_))
));
assert_eq!(gateway.0.load(Ordering::SeqCst), 0);
only_request(&upstream).await;
}
#[rstest]
#[tokio::test]
async fn a_mock_answers_before_the_gateway(call: MessagesCall) {
let gateway = Arc::new(CountingGateway::default());
let route =
messages_route(Arc::new(RecordingSecrets::empty())).with_host_answer(gateway.clone());
let mut mocked = gateway_call(call, UNREACHABLE_BASE.into(), true);
mocked.shaping.settings.mock_response = Some(object(json!({"mocked": true})));
let output = run_through(route, mocked).await.unwrap();
let MessagesOutput::Complete(MessagesBody::Unchecked(body)) = output else {
panic!("the mock is the call's response");
};
assert_eq!(body, object(json!({"mocked": true})));
assert_eq!(gateway.0.load(Ordering::SeqCst), 0);
}

199
litellm-rust/crates/inference-messages/tests/messages/host.rs
View file

@ -3,7 +3,7 @@ use litellm_host::lifecycle::ExecutionEvent;
use std::sync::Mutex;
use litellm_host::{
interceptors::{ExecutionFacts, RequestContext, ResultSource, WireRequest},
interceptors::{ExecutionFacts, RawResponse, RequestContext, ResultSource, WireRequest},
lifecycle::CallEvent,
};
use litellm_inference_messages::{MessagesBody, MessagesCallResponse, route::Messages};
@ -46,19 +46,23 @@ impl RecordingHost {
Self::new(call, Box::new(Ok))
}
fn raw_responses(&self) -> Vec<String> {
fn raw_heads(&self) -> Vec<RawResponse> {
self.events
.lock()
.unwrap()
.iter()
.filter_map(|event| match event {
CallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw }) => {
Some(raw.body.clone())
Some(raw.clone())
}
_ => None,
})
.collect()
}
fn raw_responses(&self) -> Vec<String> {
self.raw_heads().into_iter().map(|raw| raw.body).collect()
}
}
impl RecordingHost {
@ -137,7 +141,7 @@ async fn rejected_results_are_not_delivered_or_cached(
let response = if streaming {
ResponseTemplate::new(200).set_body_raw(
"event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n",
"event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"Hi\"}}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n",
"text/event-stream",
)
} else {
@ -291,22 +295,23 @@ async fn what_before_send_returns_is_what_the_provider_receives(call: MessagesCa
}
#[rstest]
#[case::enable(false, json!(true), Some(true))]
#[case::disable(true, json!(false), Some(false))]
#[case::null(true, Value::Null, Some(false))]
#[case::invalid(false, json!("true"), None)]
#[case::enable(false, json!(true), true)]
#[case::disable(true, json!(false), false)]
#[case::null(true, Value::Null, false)]
#[case::truthy_string(false, json!("false"), true)]
#[case::falsy_integer(true, json!(0), false)]
#[tokio::test]
async fn response_mode_follows_the_intercepted_request(
call: MessagesCall,
traces: TraceCapture,
#[case] original_stream: bool,
#[case] rewritten_stream: Value,
#[case] expected_stream: Option<bool>,
#[case] expected_stream: bool,
) {
use futures_util::TryStreamExt;
let sse = "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n";
let response = if expected_stream == Some(true) {
let response = if expected_stream {
ResponseTemplate::new(200).set_body_raw(sse, "text/event-stream")
} else {
message_response()
@ -332,14 +337,14 @@ async fn response_mode_follows_the_intercepted_request(
.await?;
match output {
MessagesCallResponse::Stream { chunks, .. } => {
assert_eq!(expected_stream, Some(true));
assert!(expected_stream);
assert_eq!(
chunks.try_collect::<Vec<_>>().await?.concat(),
sse.as_bytes()
);
}
MessagesCallResponse::Complete(message) => {
assert_eq!(expected_stream, Some(false));
assert!(!expected_stream);
assert_eq!(
message,
MessagesBody::Message(serde_json::from_value(message_body()).unwrap())
@ -352,18 +357,12 @@ async fn response_mode_follows_the_intercepted_request(
let summaries = traces.summaries("litellm.route");
assert_eq!(summaries.len(), 1);
let Some(expected_stream) = expected_stream else {
assert!(result.unwrap_err().is_request());
assert!(received(&upstream).await.is_empty());
assert_eq!(summaries[0]["outcome"], "failure");
return;
};
result.unwrap();
assert_eq!(
only_request(&upstream).await.json()["stream"],
rewritten_stream
);
assert_eq!(host.raw_responses().len(), usize::from(!expected_stream));
assert_eq!(host.raw_responses().len(), 1);
assert_eq!(summaries[0]["stream"], expected_stream);
assert_eq!(summaries[0]["outcome"], "success");
}
@ -405,14 +404,9 @@ async fn the_raw_upstream_text_is_emitted_once_for_a_message(call: MessagesCall)
}
#[rstest]
#[case::upstream_error(ResponseTemplate::new(500).set_body_string("boom"))]
#[case::stream(ResponseTemplate::new(200).set_body_raw("event: message_stop\ndata: {}\n\n", "text/event-stream"))]
#[tokio::test]
async fn no_raw_response_is_emitted_for_a_stream_or_a_failure(
call: MessagesCall,
#[case] response: ResponseTemplate,
) {
let upstream = upstream([response]).await;
async fn no_raw_response_is_emitted_for_a_failure(call: MessagesCall) {
let upstream = upstream([ResponseTemplate::new(500).set_body_string("boom")]).await;
let host = RecordingHost::passthrough(authenticated(
with_fields(call, json!({"stream": true})),
upstream.uri(),
@ -424,6 +418,32 @@ async fn no_raw_response_is_emitted_for_a_stream_or_a_failure(
assert!(host.raw_responses().is_empty());
}
#[rstest]
#[tokio::test]
async fn a_stream_emits_its_head_without_a_body_before_any_chunk(call: MessagesCall) {
let upstream = upstream([ResponseTemplate::new(200)
.insert_header("request-id", "req_stream")
.set_body_raw("event: message_stop\ndata: {}\n\n", "text/event-stream")])
.await;
let host = RecordingHost::passthrough(authenticated(
with_fields(call, json!({"stream": true})),
upstream.uri(),
));
let _ = run_through(&host).await;
let [raw] = <[RawResponse; 1]>::try_from(host.raw_heads())
.unwrap_or_else(|raws| panic!("expected one raw response, got {}", raws.len()));
assert_eq!(raw.body, "");
let head = raw.head.expect("the stream's head is forwarded");
assert_eq!(head.status, 200);
assert!(
head.headers
.iter()
.any(|(name, value)| name == "request-id" && value == "req_stream")
);
}
/// Python logs `optional_params` after its requested-param gate, so a dropped param must
/// not resurface in callbacks.
#[rstest]
@ -456,6 +476,38 @@ async fn the_request_context_carries_the_shaped_params_without_model_or_messages
assert_eq!(optional_params, json!({"max_tokens": 16}));
}
#[rstest]
#[tokio::test]
async fn the_request_context_keeps_the_nested_fields_additional_drop_params_strip(
call: MessagesCall,
) {
let upstream = upstream([message_response()]).await;
let tool = json!({"name": "lookup", "input_schema": {"type": "object"}, "input_examples": [{"q": "x"}]});
let host = RecordingHost::passthrough(authenticated(
MessagesCall {
shaping: MessagesShaping {
settings: MessagesSettings {
additional_drop_params: vec!["tools[*].input_examples".into()],
..MessagesSettings::default()
},
..MessagesShaping::default()
},
..with_fields(call, json!({"tools": [tool.clone()]}))
},
upstream.uri(),
));
run_through(&host).await.expect("messages call succeeds");
let [optional_params] = <[Value; 1]>::try_from(host.optional_params.into_inner().unwrap())
.unwrap_or_else(|seen| panic!("before_provider_request runs once, saw {}", seen.len()));
assert_eq!(optional_params["tools"], json!([tool]));
assert_eq!(
only_request(&upstream).await.json()["tools"][0].get("input_examples"),
None
);
}
#[rstest]
#[tokio::test]
async fn the_request_context_leaves_out_a_listed_param_the_model_refuses(call: MessagesCall) {
@ -534,3 +586,98 @@ async fn the_request_context_carries_the_requested_reasoning_not_its_translation
)
);
}
#[derive(Default)]
struct RecordingAnswer(Mutex<Vec<bool>>);
impl litellm_inference_messages::HostAnswer for RecordingAnswer {
fn answer(&self, stream: bool) -> futures_util::future::BoxFuture<'_, Result<(), Error>> {
self.0.lock().unwrap().push(stream);
Box::pin(async { Ok(()) })
}
}
#[rstest]
#[case::unary(false)]
#[case::stream(true)]
#[tokio::test]
async fn an_intercepted_call_is_answered_by_the_host_before_the_mock_and_the_provider(
call: MessagesCall,
#[case] streaming: bool,
) {
let upstream = upstream([message_response()]).await;
let answer = Arc::new(RecordingAnswer::default());
let route = messages_route(no_secrets()).with_host_answer(answer.clone());
let mut request = authenticated(
with_fields(call, json!({"stream": streaming})),
upstream.uri(),
);
request.shaping.settings.intercepted = true;
request.shaping.settings.mock_response = Some(object(json!({"id": "mocked"})));
let host = RecordingHost::passthrough(request);
let output = route
.execute(host.request().unwrap(), &host, None)
.await
.unwrap();
assert!(matches!(
output,
MessagesCallResponse::Complete(MessagesBody::Unchecked(body)) if body.is_empty()
));
assert_eq!(*answer.0.lock().unwrap(), vec![streaming]);
assert!(received(&upstream).await.is_empty());
let facts = host.facts.lock().unwrap();
assert_eq!(facts.len(), 1);
assert_eq!(facts[0].source, ResultSource::Host);
}
#[rstest]
#[tokio::test]
async fn a_call_no_interceptor_owns_never_asks_the_host(call: MessagesCall) {
let upstream = upstream([message_response()]).await;
let answer = Arc::new(RecordingAnswer::default());
let route = messages_route(no_secrets()).with_host_answer(answer.clone());
let host = RecordingHost::passthrough(authenticated(call, upstream.uri()));
route
.execute(host.request().unwrap(), &host, None)
.await
.unwrap();
assert!(answer.0.lock().unwrap().is_empty());
assert_eq!(received(&upstream).await.len(), 1);
}
#[derive(Default)]
struct DecliningAnswer(Mutex<Vec<bool>>);
impl litellm_inference_messages::HostAnswer for DecliningAnswer {
fn answer(&self, _stream: bool) -> futures_util::future::BoxFuture<'_, Result<(), Error>> {
panic!("a declined interception never asks for an answer")
}
fn intercept(&self, stream: bool) -> futures_util::future::BoxFuture<'_, Result<bool, Error>> {
self.0.lock().unwrap().push(stream);
Box::pin(async { Ok(false) })
}
}
#[rstest]
#[tokio::test]
async fn an_interception_no_interceptor_takes_reaches_the_provider(call: MessagesCall) {
let upstream = upstream([message_response()]).await;
let answer = Arc::new(DecliningAnswer::default());
let route = messages_route(no_secrets()).with_host_answer(answer.clone());
let mut request = authenticated(call, upstream.uri());
request.shaping.settings.intercepted = true;
let host = RecordingHost::passthrough(request);
route
.execute(host.request().unwrap(), &host, None)
.await
.unwrap();
assert_eq!(*answer.0.lock().unwrap(), vec![false]);
assert_eq!(received(&upstream).await.len(), 1);
}

4
litellm-rust/crates/inference-messages/tests/messages/main.rs
View file

@ -19,12 +19,16 @@ use wiremock::ResponseTemplate;
mod support;
use support::*;
mod gateway;
mod host;
mod pre_request;
mod request;
mod response;
mod retry;
mod secrets;
mod stream;
mod tencent;
mod transport;
const MODEL: &str = "claude-sonnet-4-5";

141
litellm-rust/crates/inference-messages/tests/messages/pre_request.rs Normal file
View file

@ -0,0 +1,141 @@
use futures_util::future::BoxFuture;
use litellm_inference_messages::PreRequestHooks;
use rstest::rstest;
use super::*;
struct Hooks {
rewrite: Box<dyn Fn(MessagesCall) -> Result<MessagesCall, Error> + Send + Sync>,
runs: Mutex<usize>,
}
impl PreRequestHooks for Hooks {
fn run(&self, call: MessagesCall) -> BoxFuture<'_, Result<MessagesCall, Error>> {
*self.runs.lock().unwrap() += 1;
let rewritten = (self.rewrite)(call);
Box::pin(async move { rewritten })
}
}
fn hooks(
rewrite: impl Fn(MessagesCall) -> Result<MessagesCall, Error> + Send + Sync + 'static,
) -> Arc<Hooks> {
Arc::new(Hooks {
rewrite: Box::new(rewrite),
runs: Mutex::new(0),
})
}
async fn run_hooked(hooks: Arc<Hooks>, call: MessagesCall) -> Result<MessagesOutput, Error> {
let host = LocalMessagesHost::new(call);
let route = messages_route(Arc::new(RecordingSecrets::empty())).with_pre_request_hooks(hooks);
litellm_host_native::in_process::run_hosted(
route.machine(host.request()?, None),
host.runtime(),
)
.await
}
fn mocked(call: MessagesCall, mock: Map<String, Value>) -> MessagesCall {
MessagesCall {
shaping: MessagesShaping {
settings: MessagesSettings {
mock_response: Some(mock),
..call.shaping.settings
},
..call.shaping
},
..call
}
}
fn routed(call: MessagesCall, base: String) -> MessagesCall {
MessagesCall {
api_key: Some("sk-test".into()),
api_base: Some(base),
..call
}
}
#[rstest]
#[tokio::test]
async fn the_provider_receives_the_call_the_hooks_leave(call: MessagesCall) {
let upstream = upstream([message_response()]).await;
let hooks = hooks(|call| {
Ok(with_fields(
call,
json!({"system": "hooked", "max_tokens": 7}),
))
});
run_hooked(hooks.clone(), routed(call, upstream.uri()))
.await
.expect("messages call succeeds");
let sent: Value = only_request(&upstream).await.body_json().unwrap();
assert_eq!(sent["system"], "hooked");
assert_eq!(sent["max_tokens"], 7);
assert_eq!(*hooks.runs.lock().unwrap(), 1);
}
#[rstest]
#[case::hook_sets_the_mock(false, true, true)]
#[case::hook_keeps_the_callers_mock(true, false, true)]
#[case::hook_clears_the_callers_mock(true, true, false)]
#[tokio::test]
async fn the_mock_is_read_from_the_hooked_call(
call: MessagesCall,
#[case] called_with_mock: bool,
#[case] hook_rewrites_mock: bool,
#[case] mocked_response: bool,
) {
let upstream = upstream([message_response()]).await;
let mock = object(json!({"id": "mocked"}));
let call = routed(call, upstream.uri());
let call = match called_with_mock {
true => mocked(call, mock.clone()),
false => call,
};
let hooked_mock = (!called_with_mock).then(|| mock.clone());
let hooks = hooks(move |call| {
Ok(match (hook_rewrites_mock, &hooked_mock) {
(false, _) => call,
(true, Some(mock)) => mocked(call, mock.clone()),
(true, None) => MessagesCall {
shaping: MessagesShaping {
settings: MessagesSettings {
mock_response: None,
..call.shaping.settings
},
..call.shaping
},
..call
},
})
});
let output = run_hooked(hooks, call)
.await
.expect("messages call succeeds");
let mocked_output = matches!(
output,
MessagesOutput::Complete(MessagesBody::Unchecked(ref body)) if *body == mock
);
assert_eq!(mocked_output, mocked_response);
assert_eq!(received(&upstream).await.is_empty(), mocked_response);
}
#[rstest]
#[tokio::test]
async fn a_failing_hook_fails_the_call_before_the_provider(call: MessagesCall) {
let upstream = upstream([message_response()]).await;
let hooks = hooks(|_| Err(Error::InvalidProvider("refused".into())));
let Err(error) = run_hooked(hooks, routed(call, upstream.uri())).await else {
panic!("the hook's failure fails the call");
};
assert!(matches!(error, Error::InvalidProvider(ref provider) if provider == "refused"));
assert!(received(&upstream).await.is_empty());
}

119
litellm-rust/crates/inference-messages/tests/messages/response.rs
View file

@ -2,7 +2,9 @@ use litellm_host::{
interceptors::{ExecutionFacts, ResultSource},
lifecycle::ExecutionEvent,
};
use litellm_inference_messages::{MessagesBody, MessagesCallResponse, messages_body};
use litellm_inference_messages::{
MessagesBody, MessagesCallResponse, UnencodableHeader, messages_body,
};
use litellm_inference_testing::{
RecordingSecrets, http_config, no_secrets, provider_http, resources,
};
@ -153,6 +155,37 @@ async fn the_message_passes_through_losslessly(call: MessagesCall) {
assert_eq!(serde_json::to_value(&message).unwrap(), upstream_body);
}
#[rstest]
#[case::without_stop_reason(json!({"stop_sequence": null}))]
#[case::without_stop_fields(json!({}))]
#[tokio::test]
async fn a_message_without_stop_reason_keeps_it_absent(call: MessagesCall, #[case] stop: Value) {
let mut upstream_body = json!({
"id": "msg_3", "type": "message", "role": "assistant", "model": MODEL,
"content": [{"type": "text", "text": "hi"}],
"usage": {"input_tokens": 1, "output_tokens": 1}
});
upstream_body
.as_object_mut()
.unwrap()
.extend(stop.as_object().unwrap().clone());
let upstream = upstream([json_response(upstream_body.clone())]).await;
let output = run(MessagesCall {
api_key: Some("sk".into()),
api_base: Some(upstream.uri()),
..call
})
.await
.expect("messages call succeeds");
let MessagesOutput::Complete(MessagesBody::Unchecked(body)) = output else {
panic!("a message without stop_reason is returned as sent");
};
upstream_body["stop_sequence"] = Value::Null;
assert_eq!(Value::Object(body), upstream_body);
}
#[rstest]
#[tokio::test]
async fn a_json_error_envelope_is_kept_verbatim(call: MessagesCall) {
@ -310,6 +343,88 @@ async fn a_provider_slower_than_the_timeout_fails_the_call(call: MessagesCall) {
assert_eq!(error.upstream_response(), None);
}
fn with_settings(call: MessagesCall, settings: MessagesSettings) -> MessagesCall {
MessagesCall {
shaping: MessagesShaping {
settings,
..call.shaping
},
..call
}
}
#[rstest]
#[case::unary(false)]
#[case::stream(true)]
#[tokio::test]
async fn a_non_ascii_response_header_fails_the_call_when_python_reads_heads_as_ascii(
call: MessagesCall,
#[case] stream: bool,
) {
let upstream = upstream([message_response().insert_header("x-fz-unicode", "café")]).await;
let call = with_fields(
MessagesCall {
api_key: Some("sk".into()),
api_base: Some(upstream.uri()),
..call
},
json!({"stream": stream}),
);
let settings = MessagesSettings {
ascii_response_headers: true,
..MessagesSettings::default()
};
let error = run(with_settings(call, settings))
.await
.expect_err("the head cannot be encoded");
assert_eq!(
litellm_error::find::<UnencodableHeader>(&error),
Some(&UnencodableHeader::Received("café".as_bytes().to_vec()))
);
}
#[rstest]
#[tokio::test]
async fn a_non_ascii_response_header_is_served_when_python_keeps_raw_heads(call: MessagesCall) {
let upstream = upstream([message_response().insert_header("x-fz-unicode", "café")]).await;
let message = run_message(MessagesCall {
api_key: Some("sk".into()),
api_base: Some(upstream.uri()),
..call
})
.await;
assert_eq!(message.id, "msg_1");
}
#[rstest]
#[tokio::test]
async fn an_unencodable_request_header_fails_before_anything_is_sent(call: MessagesCall) {
let upstream = upstream([message_response()]).await;
let call = MessagesCall {
api_key: Some("sk".into()),
api_base: Some(upstream.uri()),
..call
};
let settings = MessagesSettings {
unencodable_request_header: Some(b"h\xed\xa0\x80".to_vec()),
..MessagesSettings::default()
};
let error = run(with_settings(call, settings))
.await
.expect_err("the header cannot be encoded");
assert_eq!(
litellm_error::find::<UnencodableHeader>(&error),
Some(&UnencodableHeader::Sent(b"h\xed\xa0\x80".to_vec()))
);
assert!(received(&upstream).await.is_empty());
}
#[rstest]
#[tokio::test]
async fn a_transport_error_does_not_echo_the_provider_url(call: MessagesCall) {
@ -365,7 +480,7 @@ async fn the_facade_sends_through_the_injected_http_pool_configuration(call: Mes
}
#[rstest]
#[case::mistyped_param(json!({"model": MODEL, "messages": [], "max_tokens": "16"}))]
#[case::mistyped_param(json!({"model": MODEL, "messages": "hello", "max_tokens": 16}))]
#[case::missing_messages(json!({"model": MODEL, "max_tokens": 16}))]
fn a_body_that_does_not_parse_is_an_invalid_request(#[case] raw: Value) {
let error = messages_body(object(raw)).expect_err("the body is rejected");

364
litellm-rust/crates/inference-messages/tests/messages/retry.rs Normal file
View file

@ -0,0 +1,364 @@
use futures_util::TryStreamExt;
use litellm_host::{
interceptors::{RetriedRequest, RetryReason},
lifecycle::{CallEvent, ExecutionEvent},
};
use litellm_inference_messages::MessagesCallResponse;
use litellm_inference_testing::RecordingSecrets;
use litellm_llms::base_llm::messages::context::MessagesModelCapabilities;
use rstest::rstest;
use super::*;
const INVALID_SIGNATURE: &str = "messages.1.content.0: Invalid `signature` in `thinking` block";
const BOUND_ELSEWHERE: &str = "messages.1.content.0: Invalid `signature` in `thinking` block. The block is bound to a different conversation.";
fn rejection(status: u16, message: &str) -> ResponseTemplate {
status_response(
status,
json!({"type": "error", "error": {"type": "invalid_request_error", "message": message}}),
)
}
fn sse_response() -> ResponseTemplate {
ResponseTemplate::new(200).set_body_raw(
"event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n",
"text/event-stream",
)
}
fn replay(call: MessagesCall, api_base: String, stream: bool, thinking: Value) -> MessagesCall {
MessagesCall {
api_key: Some("sk-ant".into()),
api_base: Some(api_base),
..with_fields(
call,
json!({
"max_tokens": 4096,
"stream": stream,
"thinking": thinking,
"messages": [
{"role": "user", "content": "hi"},
{"role": "assistant", "content": [
{"type": "thinking", "thinking": "plan", "signature": "sig"},
{"type": "text", "text": "hello"}
]},
{"role": "assistant", "content": [{"type": "redacted_thinking", "data": "opaque"}]},
{"role": "user", "content": "again"}
]
}),
)
}
}
fn enabled() -> Value {
json!({"type": "enabled", "budget_tokens": 2048})
}
async fn execute(
host: &RecordingCall<Messages>,
) -> Result<(MessagesCallResponse, Vec<bytes::Bytes>), Error> {
let route = messages_route(Arc::new(RecordingSecrets::empty()));
let response = route
.execute(host.request()?, host, Some(host.events.0.sender.clone()))
.await?;
match response {
MessagesCallResponse::Stream { head, chunks } => Ok((
MessagesCallResponse::Stream {
head,
chunks: Box::pin(futures_util::stream::empty()),
},
chunks.try_collect().await?,
)),
complete => Ok((complete, Vec::new())),
}
}
fn observed_retries(host: &RecordingCall<Messages>) -> Vec<RetriedRequest> {
host.events
.0
.lock()
.unwrap()
.iter()
.filter_map(|event| match event {
CallEvent::Execution(ExecutionEvent::ProviderRequestRetried { retry }) => {
Some(retry.clone())
}
_ => None,
})
.collect()
}
#[rstest]
#[case::invalid_signature(INVALID_SIGNATURE)]
#[case::missing_signature(
"messages.1.content.0.thinking.signature.str: Input should be a valid string"
)]
#[case::empty_thinking("messages.1.content.0.thinking: each thinking block must contain thinking")]
#[case::unpinned_binding_mismatch(BOUND_ELSEWHERE)]
#[tokio::test]
async fn a_rejected_thinking_replay_is_sent_once_more_without_thinking(
call: MessagesCall,
#[case] message: &str,
#[values(false, true)] stream: bool,
#[values("anthropic", "azure_ai", "deepseek", "minimax")] provider: &str,
) {
let success = if stream {
sse_response()
} else {
message_response()
};
let upstream = upstream([rejection(400, message), success]).await;
let host = RecordingCall::<Messages>::new(replay(
MessagesCall {
custom_llm_provider: Some(provider.into()),
..call
},
upstream.uri(),
stream,
enabled(),
));
let (response, chunks) = execute(&host).await.expect("the retry succeeds");
let [first, second] = <[_; 2]>::try_from(received(&upstream).await).unwrap();
assert_eq!(first.json()["thinking"], enabled());
assert_eq!(
second.json()["messages"],
json!([
{"role": "user", "content": "hi"},
{"role": "assistant", "content": [{"type": "text", "text": "hello"}]},
{"role": "user", "content": "again"}
])
);
let mut expected = first.json();
expected.as_object_mut().unwrap().shift_remove("thinking");
expected["messages"] = second.json()["messages"].clone();
assert_eq!(second.json(), expected);
assert_eq!(first.header("x-hook"), second.header("x-hook"));
let expected_retry = RetriedRequest {
reasons: vec![RetryReason::ThinkingSignature],
body: second.json(),
};
assert_eq!(*host.retried.lock().unwrap(), vec![expected_retry.clone()]);
assert_eq!(observed_retries(&host), vec![expected_retry]);
assert_eq!(
matches!(response, MessagesCallResponse::Stream { .. }),
stream
);
assert_eq!(chunks.is_empty(), !stream);
}
#[rstest]
#[tokio::test]
async fn a_second_rejection_is_the_error_the_caller_sees(
call: MessagesCall,
#[values(false, true)] stream: bool,
) {
let upstream = upstream([
rejection(400, INVALID_SIGNATURE),
rejection(
400,
"messages.0.content.0: Invalid `signature` in `thinking` block, again",
),
])
.await;
let host = RecordingCall::<Messages>::new(replay(call, upstream.uri(), stream, enabled()));
let error = execute(&host).await.err().expect("the retry fails");
assert_eq!(received(&upstream).await.len(), 2);
let (status, body) = error.upstream_response().unwrap();
assert_eq!(status, 400);
assert!(body.contains("again"), "{body}");
assert_eq!(host.retried.lock().unwrap().len(), 1);
}
#[rstest]
#[case::not_a_400(429, INVALID_SIGNATURE, json!({"type": "adaptive"}))]
#[case::unrelated_400(400, "max_tokens: 4096 > 64000", json!({"type": "adaptive"}))]
#[case::signature_without_thinking(400, "Invalid `signature` in `tool_use` block", json!({"type": "adaptive"}))]
#[case::caller_pinned_binding(
400,
BOUND_ELSEWHERE,
json!({"type": "adaptive", "block_binding": {"prefix_mismatch_behavior": "error"}})
)]
#[tokio::test]
async fn other_failures_reach_the_caller_without_a_retry(
call: MessagesCall,
#[case] status: u16,
#[case] message: &str,
#[case] thinking: Value,
) {
let upstream = upstream([rejection(status, message), message_response()]).await;
let adaptive = MessagesShaping {
capabilities: MessagesModelCapabilities {
supports_adaptive_thinking: true,
..MessagesModelCapabilities::default()
},
..MessagesShaping::default()
};
let host = RecordingCall::<Messages>::new(replay(
MessagesCall {
shaping: adaptive,
..call
},
upstream.uri(),
false,
thinking.clone(),
));
let error = execute(&host).await.err().expect("the call fails");
let sent = only_request(&upstream).await;
assert_eq!(sent.json()["thinking"], thinking);
assert_eq!(error.upstream_response().unwrap().0, status);
assert!(host.retried.lock().unwrap().is_empty());
assert!(observed_retries(&host).is_empty());
}
const STRICT_REJECTION: &str = "Too many strict tools (21). The maximum is 20.";
fn strict_tools(
call: MessagesCall,
api_base: String,
stream: bool,
drop_params: bool,
) -> MessagesCall {
MessagesCall {
api_key: Some("sk-ant".into()),
api_base: Some(api_base),
shaping: MessagesShaping {
settings: MessagesSettings {
drop_params,
drops_unsupported_params: drop_params,
..MessagesSettings::default()
},
..MessagesShaping::default()
},
..with_fields(
call,
json!({
"max_tokens": 64,
"stream": stream,
"messages": [{"role": "user", "content": "hi"}],
"tools": [
{"name": "a", "input_schema": {"type": "object"}, "strict": true},
{"name": "b", "input_schema": {"type": "object"}, "strict": false},
{"name": "c", "input_schema": {"type": "object"}}
]
}),
)
}
}
#[rstest]
#[tokio::test]
async fn a_rejected_strict_tool_is_sent_once_more_without_strict_under_drop_params(
call: MessagesCall,
#[values(false, true)] stream: bool,
) {
let success = if stream {
sse_response()
} else {
message_response()
};
let upstream = upstream([rejection(400, STRICT_REJECTION), success]).await;
let host = RecordingCall::<Messages>::new(strict_tools(call, upstream.uri(), stream, true));
let (response, _) = execute(&host).await.expect("the retry succeeds");
let [first, second] = <[_; 2]>::try_from(received(&upstream).await).unwrap();
assert_eq!(first.json()["tools"][0]["strict"], json!(true));
let mut expected = first.json();
expected["tools"] = json!([
{"name": "a", "input_schema": {"type": "object"}},
{"name": "b", "input_schema": {"type": "object"}},
{"name": "c", "input_schema": {"type": "object"}}
]);
assert_eq!(second.json(), expected);
let expected_retry = RetriedRequest {
reasons: vec![RetryReason::StrictTools { dropped: 1 }],
body: second.json(),
};
assert_eq!(*host.retried.lock().unwrap(), vec![expected_retry.clone()]);
assert_eq!(observed_retries(&host), vec![expected_retry]);
assert_eq!(
matches!(response, MessagesCallResponse::Stream { .. }),
stream
);
}
#[rstest]
#[tokio::test]
async fn a_rejected_strict_tool_reaches_the_caller_without_drop_params(call: MessagesCall) {
let upstream = upstream([rejection(400, STRICT_REJECTION), message_response()]).await;
let host = RecordingCall::<Messages>::new(strict_tools(call, upstream.uri(), false, false));
let error = execute(&host).await.err().expect("the call fails");
assert_eq!(
only_request(&upstream).await.json()["tools"][0]["strict"],
json!(true)
);
assert_eq!(error.upstream_response().unwrap().0, 400);
assert!(observed_retries(&host).is_empty());
}
#[rstest]
#[tokio::test]
async fn strict_tools_and_a_thinking_replay_share_the_one_retry(call: MessagesCall) {
let upstream = upstream([
rejection(400, STRICT_REJECTION),
rejection(400, INVALID_SIGNATURE),
message_response(),
])
.await;
let host = RecordingCall::<Messages>::new(replay(
strict_tools(call, upstream.uri(), false, true),
upstream.uri(),
false,
enabled(),
));
let error = execute(&host)
.await
.err()
.expect("the second rejection is raised");
let [first, second] = <[_; 2]>::try_from(received(&upstream).await).unwrap();
assert_eq!(first.json()["thinking"], enabled());
assert_eq!(second.json()["thinking"], enabled());
assert_eq!(
second.json()["tools"][0],
json!({"name": "a", "input_schema": {"type": "object"}})
);
assert_eq!(received(&upstream).await.len(), 2);
assert!(error.upstream_response().unwrap().1.contains("signature"));
}
#[rstest]
#[tokio::test]
async fn a_truthy_global_drop_params_that_is_not_true_does_not_drop_strict(call: MessagesCall) {
let upstream = upstream([rejection(400, STRICT_REJECTION), message_response()]).await;
let strict = strict_tools(call, upstream.uri(), false, true);
let host = RecordingCall::<Messages>::new(MessagesCall {
shaping: MessagesShaping {
settings: MessagesSettings {
drops_unsupported_params: false,
..strict.shaping.settings.clone()
},
..strict.shaping.clone()
},
..strict
});
let error = execute(&host).await.err().expect("the call fails");
assert_eq!(
only_request(&upstream).await.json()["tools"][0]["strict"],
json!(true)
);
assert_eq!(error.upstream_response().unwrap().0, 400);
assert!(observed_retries(&host).is_empty());
}

4
litellm-rust/crates/inference-messages/tests/messages/stream.rs
View file

@ -451,7 +451,7 @@ async fn stalling_error_upstream() -> (String, JoinHandle<()>) {
#[rstest]
#[tokio::test]
async fn the_timeout_covers_a_stalled_stream_error_body(call: MessagesCall) {
async fn a_stalled_stream_error_body_reads_as_empty(call: MessagesCall) {
let (base, connection) = stalling_error_upstream().await;
let host = RecordingStreamHost::new(
MessagesCall {
@ -467,7 +467,7 @@ async fn the_timeout_covers_a_stalled_stream_error_body(call: MessagesCall) {
.expect_err("a stalled error body fails the call");
assert!(error.dispatched(), "{error:?}");
assert_eq!(error.upstream_response(), None);
assert_eq!(error.upstream_response(), Some((429, "")));
assert!(host.seen.into_inner().unwrap().is_empty());
tokio::time::timeout(Duration::from_secs(5), connection)
.await

256
litellm-rust/crates/inference-messages/tests/messages/transport.rs Normal file
View file

@ -0,0 +1,256 @@
use std::collections::VecDeque;
use bytes::Bytes;
use futures_util::{StreamExt, TryStreamExt, future::BoxFuture};
use litellm_inference_messages::{
Delivered, MessagesCallResponse, Transport, TransportBody, TransportRequest, TransportResponse,
};
use litellm_inference_testing::RecordingSecrets;
use rstest::rstest;
use super::*;
const INVALID_SIGNATURE: &str = "messages.1.content.0: Invalid `signature` in `thinking` block";
const UNREACHABLE_BASE: &str = "http://127.0.0.1:9";
#[derive(Debug, thiserror::Error)]
#[error("raised by the host client")]
struct Raised;
enum Scripted {
Complete(Value),
Stream(&'static [&'static str]),
Rejected(u16, &'static str),
Failed,
}
struct Sent {
url: String,
headers: Vec<(String, String)>,
body: Value,
stream: bool,
}
#[derive(Default)]
struct ScriptedTransport {
script: Mutex<VecDeque<Scripted>>,
sent: Mutex<Vec<Sent>>,
}
impl ScriptedTransport {
fn new(script: impl IntoIterator<Item = Scripted>) -> Arc<Self> {
Arc::new(Self {
script: Mutex::new(script.into_iter().collect()),
sent: Mutex::default(),
})
}
fn sent(&self) -> std::sync::MutexGuard<'_, Vec<Sent>> {
self.sent.lock().unwrap()
}
}
impl Transport for ScriptedTransport {
fn send(
&self,
request: TransportRequest,
) -> BoxFuture<'static, Result<TransportResponse, Error>> {
let TransportBody::Json(body) = request.body else {
panic!("an unsigned request hands the host its JSON body");
};
self.sent.lock().unwrap().push(Sent {
url: request.url,
headers: request.headers,
body,
stream: request.stream,
});
let next = self
.script
.lock()
.unwrap()
.pop_front()
.expect("a scripted answer");
Box::pin(async move {
let (status, delivered) = match next {
Scripted::Complete(body) => (200, Delivered::Complete(body.to_string())),
Scripted::Stream(frames) => (
200,
Delivered::Stream(
futures_util::stream::iter(
frames.iter().map(|frame| Ok(Bytes::from_static(frame.as_bytes()))),
)
.boxed(),
),
),
Scripted::Rejected(status, message) => (
status,
Delivered::Rejected {
text: json!({"type": "error", "error": {"type": "invalid_request_error", "message": message}})
.to_string(),
error: Error::transport(Raised),
},
),
Scripted::Failed => return Err(Error::transport(Raised)),
};
Ok(TransportResponse {
status,
headers: vec![("request-id".into(), "req_1".into())],
delivered,
})
})
}
}
fn hosted_call(call: MessagesCall, fields: Value) -> MessagesCall {
MessagesCall {
api_key: Some("sk-ant".into()),
api_base: Some(UNREACHABLE_BASE.into()),
..with_fields(call, fields)
}
}
async fn execute(
transport: Arc<ScriptedTransport>,
call: MessagesCall,
) -> Result<(MessagesCallResponse, Vec<Bytes>), Error> {
let host = RecordingCall::<Messages>::new(call);
let route = messages_route(Arc::new(RecordingSecrets::empty())).with_transport(transport);
match route
.execute(host.request()?, &host, Some(host.events.0.sender.clone()))
.await?
{
MessagesCallResponse::Stream { head, chunks } => Ok((
MessagesCallResponse::Stream {
head,
chunks: Box::pin(futures_util::stream::empty()),
},
chunks.try_collect().await?,
)),
complete => Ok((complete, Vec::new())),
}
}
#[rstest]
#[tokio::test]
async fn the_host_transport_sends_the_wire_request_instead_of_the_native_client(
call: MessagesCall,
) {
let transport = ScriptedTransport::new([Scripted::Complete(message_body())]);
let (response, _) = execute(transport.clone(), hosted_call(call, json!({})))
.await
.expect("the host transport answers");
let MessagesCallResponse::Complete(MessagesBody::Message(message)) = response else {
panic!("a complete message");
};
assert_eq!(message.id, "msg_1");
let sent = transport.sent();
let [request] = sent.as_slice() else {
panic!("one request");
};
assert_eq!(request.url, format!("{UNREACHABLE_BASE}/v1/messages"));
assert!(!request.stream);
assert_eq!(request.body["model"], MODEL);
assert!(
request
.headers
.iter()
.any(|(name, value)| name == "x-api-key" && value == "sk-ant"),
"{:?}",
request.headers
);
}
#[rstest]
#[tokio::test]
async fn a_host_stream_reaches_the_caller_chunk_by_chunk(call: MessagesCall) {
let frames: &'static [&'static str] = &[
"event: message_start\ndata: {\"type\":\"message_start\"}\n\n",
"event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n",
];
let transport = ScriptedTransport::new([Scripted::Stream(frames)]);
let (response, chunks) = execute(
transport.clone(),
hosted_call(call, json!({"stream": true})),
)
.await
.expect("the host stream opens");
let MessagesCallResponse::Stream { head, .. } = response else {
panic!("a stream");
};
assert_eq!(
head.headers,
vec![("request-id".to_string(), "req_1".to_string())]
);
assert_eq!(
chunks,
frames
.iter()
.map(|frame| Bytes::from_static(frame.as_bytes()))
.collect::<Vec<_>>()
);
assert!(transport.sent()[0].stream);
}
fn replayed_thinking() -> Value {
json!({
"max_tokens": 4096,
"thinking": {"type": "enabled", "budget_tokens": 2048},
"messages": [
{"role": "user", "content": "hi"},
{"role": "assistant", "content": [
{"type": "thinking", "thinking": "plan", "signature": "sig"},
{"type": "text", "text": "hello"}
]},
{"role": "user", "content": "again"}
]
})
}
#[rstest]
#[tokio::test]
async fn a_recoverable_rejection_is_retried_through_the_host_transport(call: MessagesCall) {
let transport = ScriptedTransport::new([
Scripted::Rejected(400, INVALID_SIGNATURE),
Scripted::Complete(message_body()),
]);
execute(transport.clone(), hosted_call(call, replayed_thinking()))
.await
.expect("the retry succeeds");
let sent = transport.sent();
let [first, second] = sent.as_slice() else {
panic!("two requests");
};
assert!(first.body.get("thinking").is_some());
assert!(second.body.get("thinking").is_none());
}
#[rstest]
#[case::final_rejection(vec![Scripted::Rejected(429, "slow down")])]
#[case::spent_retry(vec![
Scripted::Rejected(400, INVALID_SIGNATURE),
Scripted::Rejected(400, INVALID_SIGNATURE),
])]
#[case::client_failure(vec![Scripted::Failed])]
#[tokio::test]
async fn the_caller_gets_the_error_the_host_transport_raised(
call: MessagesCall,
#[case] script: Vec<Scripted>,
) {
let attempts = script.len();
let transport = ScriptedTransport::new(script);
let error = execute(transport.clone(), hosted_call(call, replayed_thinking()))
.await
.err()
.expect("the call fails");
assert!(litellm_error::find::<Raised>(&error).is_some(), "{error:?}");
assert_eq!(error.upstream_response(), None);
assert_eq!(transport.sent().len(), attempts);
}

10
litellm-rust/crates/inference-messages/tests/support/mod.rs
View file

@ -117,6 +117,7 @@ pub struct RecordingCall<P: litellm_host::protocol::Protocol> {
pub events: Arc<CallEvents>,
pub chunks: Mutex<Vec<P::Chunk>>,
pub head: Mutex<Option<P::StreamHead>>,
pub retried: Mutex<Vec<litellm_host::interceptors::RetriedRequest>>,
}
#[derive(Default)]
@ -167,6 +168,7 @@ impl<P: litellm_host::protocol::Protocol> RecordingCall<P> {
events: Arc::new(CallEvents::default()),
chunks: Mutex::new(Vec::new()),
head: Mutex::new(None),
retried: Mutex::new(Vec::new()),
}
}
}
@ -195,6 +197,14 @@ impl<P: litellm_host::protocol::Protocol> litellm_host::interceptors::Intercepto
) -> Result<(), P::Error> {
Ok(())
}
async fn provider_request_retried(
&self,
retry: litellm_host::interceptors::RetriedRequest,
) -> Result<(), P::Error> {
self.retried.lock().unwrap().push(retry);
Ok(())
}
}
impl<P> RecordingCall<P>

1
litellm-rust/crates/inference-ocr/tests/ocr/lifecycle.rs
View file

@ -22,6 +22,7 @@ pub(crate) fn event_name(event: &CallEvent) -> &'static str {
CallEvent::Execution(ExecutionEvent::ResultReady { .. }) => "result_ready",
CallEvent::Started { .. } => "started",
CallEvent::Execution(ExecutionEvent::ProviderResponseReceived { .. }) => "response",
CallEvent::Execution(ExecutionEvent::ProviderRequestRetried { .. }) => "retried",
CallEvent::Succeeded { .. } => "success",
CallEvent::Failed { .. } => "failure",
CallEvent::Cancelled { .. } => "cancelled",

6
litellm-rust/crates/inference-ocr/tests/ocr/machine.rs
View file

@ -66,6 +66,11 @@ async fn drive_until(
reply.send(host.after_provider_response(raw).await);
Ok(())
}
HostRequest::Intercept(InterceptRequest::ProviderRequestRetried { reply, .. }) => {
ops.push("retried");
reply.send(Ok(()));
Ok(())
}
};
if let Err(failure) = answer {
break machine.interrupt(failure).await.map(completed);
@ -86,6 +91,7 @@ async fn drive_until_notified(machine: &mut OcrMachine, host: &LocalOcrHost, sto
MachineStep::Suspended(HostRequest::HostCall(op)) => host.handle_host_call(op).await.unwrap(),
MachineStep::Suspended(HostRequest::Intercept(InterceptRequest::BeforeProviderRequest { wire, reply, .. })) => reply.send(*wire),
MachineStep::Suspended(HostRequest::Intercept(InterceptRequest::AfterProviderResponse { reply, .. })) => reply.send(Ok(())),
MachineStep::Suspended(HostRequest::Intercept(InterceptRequest::ProviderRequestRetried { reply, .. })) => reply.send(Ok(())),
MachineStep::Suspended(HostRequest::Stream(stream)) => match stream {
litellm_host::protocol::StreamDelivery::Open(head, _) => match head {},
litellm_host::protocol::StreamDelivery::Chunk(chunk, _) => match chunk {},

1
litellm-rust/crates/inference/AGENTS.md
View file

@ -81,3 +81,4 @@
- Cache lookup follows provider preparation, credential resolution, and the request interceptor
- keys describe the effective provider URL, authenticated headers, and rewritten body
- signed requests bypass caching until the signing identity has a stable cache representation
- A service that keys entries by the caller's arguments (`keyed_by_caller`, the Python host's legacy cache) is looked up before provider preparation, as Python's `@client` wrapper does, and a hit skips request building and the request interceptor

2
litellm-rust/crates/inference/Cargo.toml
View file

@ -12,6 +12,7 @@ litellm-framer.workspace = true
tokio-util = { version = "0.7", features = ["codec"] }
litellm-secrets.workspace = true
litellm-core-utils.workspace = true
litellm-error.workspace = true
litellm-host.workspace = true
bytes.workspace = true
futures-util.workspace = true
@ -29,7 +30,6 @@ time.workspace = true
url.workspace = true
[dev-dependencies]
litellm-error.workspace = true
litellm-cache-memory.workspace = true
litellm-http = { workspace = true, features = ["test-support"] }
litellm-llms = { workspace = true, features = ["test-support"] }

17
litellm-rust/crates/inference/src/caching.rs
View file

@ -61,6 +61,10 @@ pub trait StreamCachable: Cachable {
fn replay(data: Bytes) -> Option<OutputOf<Self>>;
fn bytes(chunk: &Self::Chunk) -> &[u8];
fn storable_stream(text: &str) -> bool {
successful_stream(text, Self::TERMINAL_EVENT)
}
}
#[derive(Serialize, Deserialize)]
@ -238,12 +242,19 @@ impl<P: StreamCachable> CallCache<P> {
identity: &ProviderIdentity,
wire: &WireRequest,
) -> Self {
Self::new(
cache,
policy,
CacheRequest::from_wire(identity.clone(), Some(wire)),
)
}
pub fn new(cache: Option<&ScopedCache>, policy: CachePolicy, request: CacheRequest) -> Self {
let session = cache.and_then(|cache| {
if !policy.enabled() {
return None;
}
let options = cache.options(Some(policy));
let request = CacheRequest::from_wire(identity.clone(), Some(wire));
Some(CacheSession {
request: options.request(
&cache.service.config().namespace,
@ -281,7 +292,7 @@ impl<P: StreamCachable> CallCache<P> {
P::Response: Serialize,
{
let Some(session) = self.session.filter(|session| {
*source == ResultSource::Provider && session.request.controls.writes()
!matches!(source, ResultSource::Cache { .. }) && session.request.controls.writes()
}) else {
return output;
};
@ -322,7 +333,7 @@ fn capture_stream<P: StreamCachable>(
None => {
if let Some(data) = captured
&& let Ok(text) = String::from_utf8(data)
&& successful_stream(&text, P::TERMINAL_EVENT)
&& P::storable_stream(&text)
&& let Ok(entry) = serde_json::to_value(ResponseEnvelope::new(
P::SURFACE,
CachedOutput::<Value>::Stream(text),

13
litellm-rust/crates/inference/src/context.rs
View file

@ -1,6 +1,6 @@
use litellm_cache_response::CachePolicy;
use litellm_host::{
interceptors::{ExecutionFacts, Interceptors, RawResponse, ResultSource},
interceptors::{ExecutionFacts, Interceptors, RawResponse, ResultSource, RetriedRequest},
lifecycle::{CallEvent, ExecutionEvent},
observation::ObservationSender,
};
@ -57,4 +57,15 @@ impl<'a, I: Interceptors<Error>> CallContext<'a, I> {
.await
.map_err(Error::post_call)
}
pub async fn provider_request_retried(&self, retry: RetriedRequest) -> Result<(), Error> {
if let Some(observers) = &self.observers {
observers.emit(CallEvent::Execution(
ExecutionEvent::ProviderRequestRetried {
retry: retry.clone(),
},
));
}
self.interceptors.provider_request_retried(retry).await
}
}

19
litellm-rust/crates/inference/src/error.rs
View file

@ -24,6 +24,8 @@ pub enum Error {
Secret(#[from] SecretError),
#[error("the host driver failed")]
Host(#[source] MachineFault),
#[error("the host transport failed")]
Transport(#[source] litellm_error::Source),
#[error("post-call hook failed")]
PostCallHook(#[source] Arc<Error>),
}
@ -47,6 +49,10 @@ impl Error {
Self::PostCallHook(Arc::new(error))
}
pub fn transport(cause: impl std::error::Error + Send + Sync + 'static) -> Self {
Self::Transport(litellm_error::Source::new(cause))
}
/// The caller's request is what is wrong, as opposed to the environment, the wire, or
/// the provider's answer.
pub fn is_request(&self) -> bool {
@ -55,7 +61,11 @@ impl Error {
Self::Llm(error) => error.is_request(),
Self::Auth(error) => error.is_request(),
Self::Http(error) => error.is_request(),
Self::Declined(_) | Self::Secret(_) | Self::Host(_) | Self::PostCallHook(_) => false,
Self::Declined(_)
| Self::Secret(_)
| Self::Host(_)
| Self::Transport(_)
| Self::PostCallHook(_) => false,
}
}
@ -70,6 +80,7 @@ impl Error {
| Self::Headers(_)
| Self::Secret(_)
| Self::Host(_)
| Self::Transport(_)
| Self::PostCallHook(_) => false,
}
}
@ -85,6 +96,7 @@ impl Error {
| Self::Http(_)
| Self::Secret(_)
| Self::Host(_)
| Self::Transport(_)
| Self::PostCallHook(_) => false,
}
}
@ -100,6 +112,7 @@ impl Error {
| Self::Http(_)
| Self::Secret(_)
| Self::Host(_)
| Self::Transport(_)
| Self::PostCallHook(_) => None,
}
}
@ -121,7 +134,7 @@ impl Error {
match self {
Self::Llm(error) => error.dispatched(),
Self::Http(error) => error.dispatched(),
Self::Host(_) | Self::PostCallHook(_) => true,
Self::Host(_) | Self::Transport(_) | Self::PostCallHook(_) => true,
Self::InvalidProvider(_)
| Self::UnresolvedProvider { .. }
| Self::Declined(_)
@ -142,6 +155,7 @@ impl Error {
| Self::Http(_)
| Self::Secret(_)
| Self::Host(_)
| Self::Transport(_)
| Self::PostCallHook(_) => None,
}
}
@ -193,6 +207,7 @@ impl Error {
| Self::Headers(_)
| Self::Secret(_)
| Self::Host(_)
| Self::Transport(_)
| Self::PostCallHook(_) => None,
}
}

17
litellm-rust/crates/inference/src/outbound.rs
View file

@ -1,6 +1,6 @@
use std::time::Duration;
use litellm_http::outbound::{json, signed};
use litellm_http::outbound::{json, json_bytes, signed};
use litellm_llms::base_llm::auth::Authenticated;
use serde_json::Value;
@ -34,3 +34,18 @@ pub fn outbound_request(
Some(signer) => signed(request, &signer),
}
}
pub fn raw_leaf_outbound_request(
authenticated: Authenticated,
url: &str,
body: &Value,
timeout: Option<Duration>,
) -> Result<reqwest::Request, litellm_http::Error> {
let Authenticated { headers, signer } = authenticated;
let bytes = litellm_host::raw_json::to_vec(body).map_err(litellm_http::Error::request_body)?;
let request = json_bytes(url, &headers, bytes, timeout)?;
match signer {
None => Ok(request),
Some(signer) => signed(request, &signer),
}
}

64
litellm-rust/crates/llms-types/src/formats/messages/request.rs
View file

@ -320,7 +320,7 @@ impl ThinkingConfig {
#[macro_rules_attribute::apply(crate::wire_type)]
pub struct MessagesRequest {
pub model: String,
pub messages: Vec<Message>,
pub messages: Vec<Recognized<Message>>,
#[serde(flatten)]
pub params: MessagesOptionalParams,
}
@ -329,7 +329,7 @@ pub struct MessagesRequest {
#[derive(Default)]
pub struct MessagesOptionalParams {
#[serde(skip_serializing_if = "Option::is_none")]
pub max_tokens: Option<u64>,
pub max_tokens: Option<Recognized<u64>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub system: Option<SystemPrompt>,
#[serde(skip_serializing_if = "Option::is_none")]
@ -337,13 +337,21 @@ pub struct MessagesOptionalParams {
#[serde(skip_serializing_if = "Option::is_none")]
pub stop_sequences: Option<Vec<String>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub stream: Option<bool>,
pub stream: Option<Recognized<bool>>,
#[serde(
default,
skip_serializing_if = "Option::is_none",
deserialize_with = "crate::recognized::exact_float"
)]
pub temperature: Option<Recognized<f64>>,
#[serde(
default,
skip_serializing_if = "Option::is_none",
deserialize_with = "crate::recognized::exact_float"
)]
pub top_p: Option<Recognized<f64>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub temperature: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub top_p: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub top_k: Option<i64>,
pub top_k: Option<Recognized<i64>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tools: Option<Vec<Recognized<MessagesTool>>>,
#[serde(skip_serializing_if = "Option::is_none")]
@ -374,6 +382,12 @@ pub struct MessagesOptionalParams {
pub extra: Map<String, Value>,
}
impl MessagesOptionalParams {
pub fn streams(&self) -> bool {
self.stream.as_ref().is_some_and(Recognized::is_truthy)
}
}
impl Message {
pub fn blocks(&self) -> &[ContentBlock] {
match &self.content {
@ -447,8 +461,8 @@ mod tests {
assert_eq!(
(
request.params.max_tokens,
request.params.stream,
request.params.max_tokens.clone(),
request.params.stream.clone(),
request
.params
.extra
@ -456,11 +470,39 @@ mod tests {
.map(String::as_str)
.collect::<Vec<_>>(),
),
(Some(16_u64), Some(true), vec!["safeguards"])
(
Some(Recognized::Known(16_u64)),
Some(Recognized::Known(true)),
vec!["safeguards"]
)
);
assert_eq!(serde_json::to_value(request).unwrap(), body);
}
#[rstest]
#[case::numeric_string(json!({"max_tokens": "1024", "temperature": "0.5", "top_p": "x", "top_k": "5"}))]
#[case::string_stream(json!({"stream": "true"}))]
#[case::integer_stream(json!({"stream": 0}))]
#[case::floats(json!({"max_tokens": 1024.0, "top_k": 1.5}))]
#[case::booleans(json!({"max_tokens": true, "temperature": false}))]
#[case::wide_integers(json!({"max_tokens": 18_446_744_073_709_551_615_u64, "top_k": 9_223_372_036_854_775_808_u64}))]
#[case::nested(json!({"temperature": [1], "top_p": {"p": 1}}))]
fn scalar_params_of_any_type_round_trip_unchanged(#[case] params: Value) {
let body = Value::Object(
[
("model".to_string(), json!("m")),
(
"messages".to_string(),
json!([{"role": "user", "content": "hi"}]),
),
]
.into_iter()
.chain(params.as_object().unwrap().clone())
.collect(),
);
assert_eq!(round_trip::<MessagesRequest>(&body), body);
}
#[test]
fn text_constructor_serializes_as_a_text_block() {
assert_eq!(

5
litellm-rust/crates/llms-types/src/providers/anthropic/beta.rs
View file

@ -317,6 +317,11 @@ impl AnthropicBeta {
BetaProvider::VertexAi,
BetaProvider::Databricks,
],
Self::Compact20260904 => &[
BetaProvider::BedrockConverse,
BetaProvider::BedrockMantle,
BetaProvider::Databricks,
],
Self::AfkMode20260131
| Self::CePlugins20260901
| Self::CeUserManagement20260713

149
litellm-rust/crates/llms-types/src/recognized.rs
View file

@ -17,6 +17,25 @@ where
<Recognized<T> as serde::Deserialize<'de>>::deserialize(deserializer).map(Some)
}
/// `deserialize_with` for an optional float Python forwards as the caller sent it: an integer
/// stays `Unrecognized` with its exact digits, as Python sends, logs and prints it, and `null`
/// stays absent as before.
pub fn exact_float<'de, D>(deserializer: D) -> Result<Option<Recognized<f64>>, D::Error>
where
D: serde::Deserializer<'de>,
{
let value = <Value as serde::Deserialize>::deserialize(deserializer)?;
if value.is_null() {
return Ok(None);
}
if matches!(&value, Value::Number(number) if !number.is_f64()) {
return Ok(Some(Recognized::Unrecognized(value)));
}
serde_json::from_value(value)
.map(Some)
.map_err(serde::de::Error::custom)
}
impl<T> Recognized<T> {
pub fn known(&self) -> Option<&T> {
match self {
@ -24,6 +43,60 @@ impl<T> Recognized<T> {
Self::Unrecognized(_) => None,
}
}
pub fn map_known(self, f: impl FnOnce(T) -> T) -> Self {
match self {
Self::Known(value) => Self::Known(f(value)),
unrecognized => unrecognized,
}
}
}
impl Recognized<bool> {
pub fn is_truthy(&self) -> bool {
match self {
Self::Known(flag) => *flag,
Self::Unrecognized(other) => python_truthy(other),
}
}
}
pub fn python_truthy(value: &Value) -> bool {
match value {
Value::Null => false,
Value::Bool(flag) => *flag,
Value::Number(number) => number.as_f64().is_none_or(|number| number != 0.0),
Value::String(text) => !text.is_empty(),
Value::Array(items) => !items.is_empty(),
Value::Object(fields) => !fields.is_empty(),
}
}
pub fn map_known_runs<T>(
items: Vec<Recognized<T>>,
mut f: impl FnMut(Vec<T>) -> Vec<T>,
) -> Vec<Recognized<T>> {
let mut mapped = Vec::with_capacity(items.len());
let mut run = Vec::new();
for item in items {
match item {
Recognized::Known(value) => run.push(value),
unrecognized => {
if !run.is_empty() {
mapped.extend(
f(std::mem::take(&mut run))
.into_iter()
.map(Recognized::Known),
);
}
mapped.push(unrecognized);
}
}
}
if !run.is_empty() {
mapped.extend(f(run).into_iter().map(Recognized::Known));
}
mapped
}
#[cfg(test)]
@ -32,6 +105,82 @@ mod tests {
use rstest::rstest;
use serde_json::json;
#[rstest]
#[case::known_true(Recognized::Known(true), true)]
#[case::known_false(Recognized::Known(false), false)]
#[case::string_true(Recognized::Unrecognized(json!("true")), true)]
#[case::string_false(Recognized::Unrecognized(json!("false")), true)]
#[case::empty_string(Recognized::Unrecognized(json!("")), false)]
#[case::zero(Recognized::Unrecognized(json!(0)), false)]
#[case::negative_zero(Recognized::Unrecognized(json!(-0.0)), false)]
#[case::one(Recognized::Unrecognized(json!(1)), true)]
#[case::null(Recognized::Unrecognized(Value::Null), false)]
#[case::empty_list(Recognized::Unrecognized(json!([])), false)]
#[case::list(Recognized::Unrecognized(json!([0])), true)]
#[case::empty_object(Recognized::Unrecognized(json!({})), false)]
#[case::object(Recognized::Unrecognized(json!({"a": 0})), true)]
fn flags_are_read_by_python_truthiness(#[case] flag: Recognized<bool>, #[case] expected: bool) {
assert_eq!(flag.is_truthy(), expected);
}
#[test]
fn runs_of_known_items_are_mapped_around_unrecognized_ones() {
let items = vec![
Recognized::Known(1),
Recognized::Known(2),
Recognized::Unrecognized(json!("x")),
Recognized::Known(3),
];
let mut runs = Vec::new();
let mapped = map_known_runs(items, |run: Vec<u64>| {
runs.push(run.clone());
run.into_iter()
.filter(|item| *item != 2)
.map(|item| item * 10)
.collect()
});
assert_eq!(runs, vec![vec![1, 2], vec![3]]);
assert_eq!(
mapped,
vec![
Recognized::Known(10),
Recognized::Unrecognized(json!("x")),
Recognized::Known(30),
]
);
assert_eq!(
map_known_runs(Vec::<Recognized<u64>>::new(), |run| run),
vec![]
);
}
#[derive(serde::Deserialize)]
struct Sampling {
#[serde(default, deserialize_with = "exact_float")]
temperature: Option<Recognized<f64>>,
}
#[rstest]
#[case::float(json!({"temperature": 0.5}), Some(Recognized::Known(0.5)))]
#[case::small_int(json!({"temperature": 2}), Some(Recognized::Unrecognized(json!(2))))]
#[case::zero(json!({"temperature": 0}), Some(Recognized::Unrecognized(json!(0))))]
#[case::integral_float(json!({"temperature": 1.0}), Some(Recognized::Known(1.0)))]
#[case::largest_exact_int(json!({"temperature": 9_007_199_254_740_992_i64}), Some(Recognized::Unrecognized(json!(9_007_199_254_740_992_i64))))]
#[case::inexact_int(json!({"temperature": 1_152_921_504_606_846_977_i64}), Some(Recognized::Unrecognized(json!(1_152_921_504_606_846_977_i64))))]
#[case::inexact_negative_int(json!({"temperature": -9_223_372_036_854_775_807_i64}), Some(Recognized::Unrecognized(json!(-9_223_372_036_854_775_807_i64))))]
#[case::inexact_unsigned_int(json!({"temperature": 18_446_744_073_709_551_615_u64}), Some(Recognized::Unrecognized(json!(18_446_744_073_709_551_615_u64))))]
#[case::string(json!({"temperature": "0.5"}), Some(Recognized::Unrecognized(json!("0.5"))))]
#[case::null(json!({"temperature": null}), None)]
#[case::absent(json!({}), None)]
fn integers_keep_their_digits(#[case] value: Value, #[case] expected: Option<Recognized<f64>>) {
assert_eq!(
serde_json::from_value::<Sampling>(value)
.unwrap()
.temperature,
expected
);
}
#[rstest]
#[case::known(json!(7), Recognized::Known(7))]
#[case::wrong_type(json!("7"), Recognized::Unrecognized(json!("7")))]

2
litellm-rust/crates/llms/Cargo.toml
View file

@ -23,9 +23,11 @@ litellm-http.workspace = true
litellm-secrets.workspace = true
litellm-python-compat.workspace = true
litellm-router-types.workspace = true
litellm-tracing.workspace = true
base64.workspace = true
bytes.workspace = true
data-url = "0.3.2"
fancy-regex.workspace = true
futures-util.workspace = true
reqwest.workspace = true
serde.workspace = true

12
litellm-rust/crates/llms/src/anthropic/beta_headers.rs
View file

@ -1,5 +1,5 @@
use litellm_http::request::{with_header, without_headers};
use litellm_llms_types::providers::anthropic::{BETA_HEADER, BetaProvider, BetaSet};
use litellm_llms_types::providers::anthropic::{AnthropicBeta, BETA_HEADER, BetaProvider, BetaSet};
use crate::{anthropic::common_utils::existing_betas, base_llm::auth::Headers};
@ -33,6 +33,16 @@ impl BetaPolicy {
}
}
pub fn accepted_or_unlisted(betas: BetaSet, provider: BetaProvider) -> BetaSet {
betas
.into_iter()
.filter_map(|beta| match beta {
AnthropicBeta::Other(_) => Some(beta),
known => known.on(provider),
})
.collect()
}
#[cfg(test)]
mod tests {
use crate::anthropic::beta_headers::BetaPolicy;

375
litellm-rust/crates/llms/src/anthropic/common_utils.rs
View file

@ -1,4 +1,7 @@
use std::sync::LazyLock;
use crate::base_llm::messages::context::MessagesModelCapabilities;
use fancy_regex::Regex;
use litellm_auth::CredentialPlacement;
use litellm_secret_value::SecretValue;
@ -22,7 +25,7 @@ use serde_json::Value;
use crate::{
anthropic::ANTHROPIC_OAUTH_TOKEN_PREFIX,
base_llm::auth::{AuthScheme, Headers},
base_llm::auth::{AuthScheme, Headers, ValidatedEnvironment},
};
pub const ANTHROPIC_API_KEY_ENV: &str = "ANTHROPIC_API_KEY";
@ -80,10 +83,21 @@ impl<'a> OauthToken<'a> {
}
pub fn into_auth(self) -> AuthScheme {
AuthScheme::Credential {
placement: CredentialPlacement::Bearer,
secret: SecretValue::new(self.0),
}
oauth_bearer(SecretValue::new(self.0))
}
}
fn oauth_bearer(secret: SecretValue) -> AuthScheme {
AuthScheme::Credential {
placement: CredentialPlacement::Bearer,
secret,
}
}
pub fn minted_oauth_environment(headers: Headers, token: SecretValue) -> ValidatedEnvironment {
ValidatedEnvironment {
headers: with_oauth_companions(headers, &[], BrowserAccess::Omitted),
auth: oauth_bearer(token),
}
}
@ -92,11 +106,23 @@ pub fn get_api_key(
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Option<String> {
resolve_non_empty(api_key, env_lookup, &[ANTHROPIC_API_KEY_ENV])
nonempty_secret(api_key.map(str::to_owned))
.or_else(|| nonempty_secret(env_lookup(ANTHROPIC_API_KEY_ENV)))
}
pub fn get_auth_token(env_lookup: &dyn Fn(&str) -> Option<String>) -> Option<String> {
resolve_non_empty(None, env_lookup, &[ANTHROPIC_AUTH_TOKEN_ENV])
nonempty_secret(env_lookup(ANTHROPIC_AUTH_TOKEN_ENV))
}
/// Python's `normalize_nonempty_secret_str`: `str.strip()` also strips the ASCII separators.
fn nonempty_secret(value: Option<String>) -> Option<String> {
value
.map(|value| {
value
.trim_matches(|c: char| c.is_whitespace() || ('\u{1c}'..='\u{1f}').contains(&c))
.to_owned()
})
.filter(|value| !value.is_empty())
}
/// Python's `AnthropicModelInfo.get_auth_header`, naming the credential instead of building
@ -232,14 +258,24 @@ pub fn forwarded_oauth_bearer(headers: &[(String, String)]) -> Option<OauthToken
.and_then(OauthToken::parse)
}
fn with_oauth_companions(headers: Headers, dropped: &[&str]) -> Headers {
merge_beta_headers(
#[derive(Clone, Copy, PartialEq, Eq)]
enum BrowserAccess {
Granted,
Omitted,
}
fn with_oauth_companions(headers: Headers, dropped: &[&str], browser: BrowserAccess) -> Headers {
let headers = merge_beta_headers(
without_headers(headers, dropped),
BetaSet::from_iter([AnthropicBeta::Oauth20250420]),
)
.into_iter()
.chain([(DIRECT_BROWSER_ACCESS_HEADER.to_string(), "true".to_string())])
.collect()
);
match browser {
BrowserAccess::Granted => headers
.into_iter()
.chain([(DIRECT_BROWSER_ACCESS_HEADER.to_string(), "true".to_string())])
.collect(),
BrowserAccess::Omitted => headers,
}
}
pub fn optionally_handle_anthropic_oauth(headers: Headers, api_key: Option<&str>) -> OauthHandling {
@ -247,13 +283,17 @@ pub fn optionally_handle_anthropic_oauth(headers: Headers, api_key: Option<&str>
forwarded_oauth_bearer(&headers).map(|token| SecretValue::new(token.as_str()))
{
return OauthHandling::Bearer {
headers: with_oauth_companions(headers, &[API_KEY_HEADER, AUTHORIZATION]),
headers: with_oauth_companions(
headers,
&[API_KEY_HEADER, AUTHORIZATION],
BrowserAccess::Granted,
),
token,
};
}
if let Some(token) = api_key.and_then(OauthToken::parse) {
return OauthHandling::Bearer {
headers: with_oauth_companions(headers, &[API_KEY_HEADER]),
headers: with_oauth_companions(headers, &[API_KEY_HEADER], BrowserAccess::Granted),
token: SecretValue::new(token.as_str()),
};
}
@ -278,12 +318,19 @@ pub fn has_advisor_tool(tools: Option<&[Recognized<MessagesTool>]>) -> bool {
.any(|tool| matches!(tool, Recognized::Known(MessagesTool::Advisor { .. })))
}
pub fn requires_native_compaction_beta(compaction: Option<&Value>, messages: &[Message]) -> bool {
pub fn requires_native_compaction_beta(
compaction: Option<&Value>,
messages: &[Recognized<Message>],
) -> bool {
compaction.is_some()
|| messages.iter().flat_map(Message::blocks).any(|block| {
block.is_type(ContentBlockType::Compaction)
&& block.signature.as_deref().is_some_and(|s| !s.is_empty())
})
|| messages
.iter()
.filter_map(Recognized::known)
.flat_map(Message::blocks)
.any(|block| {
block.is_type(ContentBlockType::Compaction)
&& block.signature.as_deref().is_some_and(|s| !s.is_empty())
})
}
fn is_blank(text: Option<&str>) -> bool {
@ -650,6 +697,113 @@ pub fn flatten_unencrypted_web_search_results(messages: Vec<Message>) -> Vec<Mes
.collect()
}
pub fn is_anthropic_invalid_thinking_block_error(error_text: &str) -> bool {
let lower = error_text.to_lowercase();
if !lower.contains("thinking") {
return false;
}
(lower.contains("signature") && (lower.contains("invalid") || lower.contains("valid string")))
|| lower.contains("must contain thinking")
}
pub fn is_anthropic_thinking_binding_mismatch_error(error_text: &str) -> bool {
error_text
.to_lowercase()
.contains("bound to a different conversation")
}
pub fn anthropic_request_errors_on_thinking_binding_mismatch(request_body: &Value) -> bool {
request_body
.pointer("/thinking/block_binding/prefix_mismatch_behavior")
.and_then(Value::as_str)
== Some("error")
}
fn is_thinking_value_block(block: &Value) -> bool {
matches!(
block.get("type").and_then(Value::as_str),
Some("thinking" | "redacted_thinking")
)
}
fn without_thinking_value_blocks(message: Value) -> Option<Value> {
let Value::Object(mut message) = message else {
return Some(message);
};
if let Some(Value::Array(content)) = message.get_mut("content") {
content.retain(|block| !is_thinking_value_block(block));
if content.is_empty() {
return None;
}
}
Some(Value::Object(message))
}
pub fn strip_thinking_blocks_from_anthropic_messages_request(request_body: Value) -> Value {
let Value::Object(mut body) = request_body else {
return request_body;
};
if let Some(Value::Array(messages)) = body.get_mut("messages") {
*messages = std::mem::take(messages)
.into_iter()
.filter_map(without_thinking_value_blocks)
.collect();
}
body.shift_remove("thinking");
Value::Object(body)
}
static CLAUDE_STRICT_TOOLS_REJECTION: LazyLock<Regex> = LazyLock::new(|| {
Regex::new(concat!(
r"Too many strict tools \(\d+\)",
r"|Schemas contains too many (?:optional parameters|parameters with union types)",
r"|[Cc]ompiled grammar (?:is too large|size)",
r"|tools\.\d+\.custom(?:\.strict)?: ",
r"|Grammar compilation timed out",
))
.expect("valid regex")
});
pub fn is_claude_strict_tools_rejection(status: u16, error_text: &str) -> bool {
status == 400
&& CLAUDE_STRICT_TOOLS_REJECTION
.is_match(error_text)
.unwrap_or(false)
}
#[derive(Clone, Debug, PartialEq)]
pub struct ToolsWithoutStrict {
pub tools: Vec<Value>,
pub dropped_count: usize,
}
pub fn tools_without_strict(tools: Option<&Value>) -> Option<ToolsWithoutStrict> {
let tools = tools?.as_array()?;
if !tools.iter().all(Value::is_object) {
return None;
}
let dropped_count = tools
.iter()
.filter(|tool| tool.get("strict") == Some(&Value::Bool(true)))
.count();
if dropped_count == 0 {
return None;
}
Some(ToolsWithoutStrict {
tools: tools
.iter()
.cloned()
.map(|mut tool| {
if let Value::Object(fields) = &mut tool {
fields.shift_remove("strict");
}
tool
})
.collect(),
dropped_count,
})
}
#[cfg(test)]
mod tests {
use crate::base_llm::messages::context::SupportedEffortTiers;
@ -668,7 +822,7 @@ mod tests {
serde_json::from_value(value).unwrap()
}
fn history(messages: Value) -> Vec<Message> {
fn history(messages: Value) -> Vec<Recognized<Message>> {
serde_json::from_value(messages).unwrap()
}
@ -1636,6 +1790,11 @@ mod tests {
&[("ANTHROPIC_AUTH_TOKEN", "env-token")],
Some(("Authorization", "env-token")),
)]
#[case::auth_token_stripped_of_whitespace(
None,
&[("ANTHROPIC_AUTH_TOKEN", " env-token \t\n\u{1f}")],
Some(("Authorization", "env-token")),
)]
#[case::auth_token_when_the_env_key_is_blank(
None,
&[("ANTHROPIC_API_KEY", " \t"), ("ANTHROPIC_AUTH_TOKEN", "env-token")],
@ -1647,7 +1806,19 @@ mod tests {
&[("ANTHROPIC_API_KEY", OAUTH_BEARER)],
Some(("Authorization", OAUTH_TOKEN)),
)]
#[case::env_key_stripped_of_whitespace(
None,
&[("ANTHROPIC_API_KEY", " sk-env \t\n\u{1f}")],
Some(("x-api-key", "sk-env")),
)]
#[case::api_key_param_stripped_of_separators(Some("\u{1c}sk-param "), &[], Some(("x-api-key", "sk-param")))]
#[case::separator_only_env_key_falls_to_the_auth_token(
None,
&[("ANTHROPIC_API_KEY", "\u{1c}\u{1f}"), ("ANTHROPIC_AUTH_TOKEN", "env-token")],
Some(("Authorization", "env-token")),
)]
#[case::no_credentials(None, &[], None)]
#[case::separator_only_auth_token(None, &[("ANTHROPIC_AUTH_TOKEN", "\u{1c}\u{1f}")], None)]
#[case::blank_everything(Some(""), &[("ANTHROPIC_API_KEY", " "), ("ANTHROPIC_AUTH_TOKEN", " \t")], None)]
fn auth_header_prefers_the_key_then_the_auth_token(
#[case] api_key: Option<&str>,
@ -1974,4 +2145,166 @@ mod tests {
fn the_messages_url_follows_python_ensure_api_route(#[case] api_base: &str, #[case] url: &str) {
assert_eq!(complete_anthropic_url(Some(api_base), &|_| None), url);
}
fn shared_cases(table: &str) -> Vec<(String, String, bool)> {
let fixture: Value = serde_json::from_str(include_str!(concat!(
env!("CARGO_MANIFEST_DIR"),
"/tests/fixtures/anthropic_messages_http_error_retry.json"
)))
.unwrap();
fixture[table]
.as_array()
.unwrap()
.iter()
.map(|case| {
(
case["name"].as_str().unwrap().to_string(),
case["error_text"].as_str().unwrap().to_string(),
case["expected"].as_bool().unwrap(),
)
})
.collect()
}
#[rstest]
#[case::invalid_thinking_block(
"invalid_thinking_block_error",
is_anthropic_invalid_thinking_block_error
)]
#[case::thinking_binding_mismatch(
"thinking_binding_mismatch_error",
is_anthropic_thinking_binding_mismatch_error
)]
fn error_texts_match_the_cases_shared_with_python(
#[case] table: &str,
#[case] recognizes: fn(&str) -> bool,
) {
let cases = shared_cases(table);
assert!(!cases.is_empty());
let mismatched: Vec<&String> = cases
.iter()
.filter(|(_, error_text, expected)| recognizes(error_text) != *expected)
.map(|(name, _, _)| name)
.collect();
assert_eq!(mismatched, Vec::<&String>::new());
}
#[rstest]
#[case::pinned(json!({"thinking": {"block_binding": {"prefix_mismatch_behavior": "error"}}}), true)]
#[case::other_behavior(json!({"thinking": {"block_binding": {"prefix_mismatch_behavior": "drop"}}}), false)]
#[case::string_binding(json!({"thinking": {"block_binding": "strict"}}), false)]
#[case::no_thinking(json!({"messages": []}), false)]
#[case::not_an_object(json!(["thinking"]), false)]
fn a_pinned_binding_is_the_caller_asking_for_the_error(
#[case] body: Value,
#[case] expected: bool,
) {
assert_eq!(
anthropic_request_errors_on_thinking_binding_mismatch(&body),
expected
);
}
#[rstest]
#[case::mixed_turn(
json!({"model": "m", "thinking": {"type": "enabled"}, "messages": [
{"role": "assistant", "content": [
{"type": "thinking", "thinking": "t", "signature": "s"},
{"type": "text", "text": "a"},
{"type": "redacted_thinking", "data": "d"}
]}
]}),
json!({"model": "m", "messages": [{"role": "assistant", "content": [{"type": "text", "text": "a"}]}]})
)]
#[case::thinking_only_and_empty_turns_are_dropped(
json!({"messages": [
{"role": "user", "content": "hi"},
{"role": "assistant", "content": [{"type": "redacted_thinking", "data": "d"}]},
{"role": "assistant", "content": []},
{"role": "user", "content": "again"}
]}),
json!({"messages": [{"role": "user", "content": "hi"}, {"role": "user", "content": "again"}]})
)]
#[case::non_object_entries_are_kept(
json!({"messages": ["raw", {"role": "user", "content": ["raw", {"type": "thinking"}]}]}),
json!({"messages": ["raw", {"role": "user", "content": ["raw"]}]})
)]
#[case::messages_not_a_list(
json!({"messages": "hi", "thinking": {"type": "adaptive"}}),
json!({"messages": "hi"})
)]
fn retrying_strips_thinking_from_the_replay_and_the_request(
#[case] body: Value,
#[case] expected: Value,
) {
assert_eq!(
strip_thinking_blocks_from_anthropic_messages_request(body),
expected
);
}
#[test]
fn stripping_keeps_the_order_of_the_remaining_fields() {
let stripped = strip_thinking_blocks_from_anthropic_messages_request(
json!({"model": "m", "thinking": {}, "messages": [], "max_tokens": 1}),
);
let keys: Vec<&str> = stripped
.as_object()
.unwrap()
.keys()
.map(String::as_str)
.collect();
assert_eq!(keys, ["model", "messages", "max_tokens"]);
}
fn shared_table(table: &str) -> Vec<Value> {
let fixture: Value = serde_json::from_str(include_str!(concat!(
env!("CARGO_MANIFEST_DIR"),
"/tests/fixtures/anthropic_messages_http_error_retry.json"
)))
.unwrap();
fixture[table].as_array().unwrap().clone()
}
#[test]
fn strict_tools_rejections_match_the_cases_shared_with_python() {
let cases = shared_table("claude_strict_tools_rejection");
assert!(!cases.is_empty());
let mismatched: Vec<&str> = cases
.iter()
.filter(|case| {
let status = u16::try_from(case["status"].as_u64().unwrap()).unwrap();
is_claude_strict_tools_rejection(status, case["error_text"].as_str().unwrap())
!= case["expected"].as_bool().unwrap()
})
.map(|case| case["name"].as_str().unwrap())
.collect();
assert_eq!(mismatched, Vec::<&str>::new());
}
#[test]
fn tools_without_strict_matches_the_cases_shared_with_python() {
let cases = shared_table("tools_without_strict");
assert!(!cases.is_empty());
let mismatched: Vec<&str> = cases
.iter()
.filter(|case| {
let expected = case["expected"]
.as_object()
.map(|expected| ToolsWithoutStrict {
tools: expected["tools"].as_array().unwrap().clone(),
dropped_count: usize::try_from(expected["dropped_count"].as_u64().unwrap())
.unwrap(),
});
tools_without_strict(Some(&case["tools"])) != expected
})
.map(|case| case["name"].as_str().unwrap())
.collect();
assert_eq!(mismatched, Vec::<&str>::new());
}
#[test]
fn absent_tools_have_no_strict_to_drop() {
assert_eq!(tools_without_strict(None), None);
}
}

44
litellm-rust/crates/llms/src/anthropic/messages/handler.rs
View file

@ -3,7 +3,7 @@ use litellm_llms_types::{
AdaptiveThinking, EnabledThinking, Message, MessagesOptionalParams, MessagesRequest,
ThinkingConfig, ThinkingDisplay,
},
recognized::Recognized,
recognized::{Recognized, map_known_runs},
};
use serde_json::{Value, json};
@ -13,14 +13,19 @@ use crate::{
flatten_unencrypted_web_search_results, sanitize_tool_use_ids, strip_empty_content_blocks,
strip_provider_specific_fields,
},
base_llm::messages::context::MessagesShapeOptions,
};
pub fn shape_anthropic_messages_request(
request: MessagesRequest,
reasoning_auto_summary: bool,
options: MessagesShapeOptions,
) -> Result<MessagesRequest, Error> {
Ok(MessagesRequest {
messages: sanitize_anthropic_messages(request.messages),
messages: if options.presanitized {
map_known_runs(request.messages, strip_provider_specific_fields)
} else {
map_known_runs(request.messages, sanitize_anthropic_messages)
},
params: MessagesOptionalParams {
metadata: request
.params
@ -28,7 +33,10 @@ pub fn shape_anthropic_messages_request(
.as_ref()
.map(validate_anthropic_api_metadata)
.transpose()?,
thinking: with_reasoning_auto_summary(request.params.thinking, reasoning_auto_summary),
thinking: with_reasoning_auto_summary(
request.params.thinking,
options.reasoning_auto_summary,
),
..request.params
},
..request
@ -299,7 +307,10 @@ mod tests {
"thinking": {"type": "enabled", "budget_tokens": 1024},
"safeguards": [{"type": "dangerous_tool_use"}]
})),
true,
MessagesShapeOptions {
reasoning_auto_summary: true,
presanitized: false,
},
)
.unwrap();
assert_eq!(
@ -315,4 +326,27 @@ mod tests {
})
);
}
#[test]
fn presanitized_shaping_keeps_blocks_and_strips_only_provider_fields() {
let messages = json!([{"role": "user", "content": [
{"type": "text", "text": " ", "cache_control": {"type": "ephemeral"}},
{"type": "tool_result", "tool_use_id": "functions.Bash:0", "content": "ok", "provider_specific_fields": {"x": 1}}
]}]);
let sanitized = shape_anthropic_messages_request(
request(json!({"model": "m", "messages": messages})),
MessagesShapeOptions {
reasoning_auto_summary: false,
presanitized: true,
},
)
.unwrap();
assert_eq!(
serde_json::to_value(sanitized).unwrap(),
json!({"model": "m", "messages": [{"role": "user", "content": [
{"type": "text", "text": " ", "cache_control": {"type": "ephemeral"}},
{"type": "tool_result", "tool_use_id": "functions.Bash:0", "content": "ok"}
]}]})
);
}
}

1
litellm-rust/crates/llms/src/anthropic/messages/mod.rs
View file

@ -1,3 +1,4 @@
pub mod handler;
mod scalar;
pub mod thinking;
pub mod transformation;

113
litellm-rust/crates/llms/src/anthropic/messages/scalar.rs Normal file
View file

@ -0,0 +1,113 @@
use litellm_host::raw_json;
use litellm_llms_types::recognized::Recognized;
use litellm_python_compat::{
Value as PythonValue,
json::{from_json, loads},
repr::to_str,
};
use serde_json::Value;
pub(crate) fn python_value(value: &Value) -> PythonValue {
match value {
Value::String(text) => raw_json::decode(text)
.and_then(|raw| loads(raw).ok())
.unwrap_or_else(|| PythonValue::Str(text.clone())),
Value::Array(items) => PythonValue::List(items.iter().map(python_value).collect()),
Value::Object(entries) => PythonValue::Dict(
entries
.iter()
.map(|(key, item)| (PythonValue::Str(key.clone()), python_value(item)))
.collect(),
),
Value::Null | Value::Bool(_) | Value::Number(_) => from_json(value.clone()),
}
}
pub(crate) fn python_number(value: &Value) -> Option<f64> {
match value {
Value::Bool(flag) => Some(f64::from(u8::from(*flag))),
Value::Number(number) => number.as_f64(),
Value::String(text) => match raw_json::decode(text)? {
"NaN" => Some(f64::NAN),
"Infinity" => Some(f64::INFINITY),
"-Infinity" => Some(f64::NEG_INFINITY),
raw => raw.parse().ok(),
},
_ => None,
}
}
pub(crate) fn is_one(value: &Recognized<f64>) -> bool {
match value {
Recognized::Known(number) => *number == 1.0,
Recognized::Unrecognized(other) => python_number(other) == Some(1.0),
}
}
pub(crate) fn python_str<T: serde::Serialize>(value: &Recognized<T>) -> String {
match value {
Recognized::Known(known) => serde_json::to_value(known)
.map(|known| to_str(&from_json(known)))
.unwrap_or_default(),
Recognized::Unrecognized(other) => unrecognized_str(other),
}
}
fn unrecognized_str(value: &Value) -> String {
match value {
Value::Null => "None".into(),
Value::Bool(true) => "True".into(),
Value::Bool(false) => "False".into(),
Value::String(text) => match raw_json::decode(text) {
Some("NaN") => "nan".into(),
Some("Infinity") => "inf".into(),
Some("-Infinity") => "-inf".into(),
Some(raw) => raw.trim_matches('"').into(),
None => text.clone(),
},
other => other.to_string(),
}
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use serde_json::json;
use super::*;
#[rstest]
#[case::int(json!(3), Some(3.0))]
#[case::float(json!(1.5), Some(1.5))]
#[case::bool(json!(true), Some(1.0))]
#[case::wide(json!(raw_json::encode("18446744073709551616")), Some(18_446_744_073_709_551_616.0))]
#[case::negative_infinity(json!(raw_json::encode("-Infinity")), Some(f64::NEG_INFINITY))]
#[case::string(json!("1"), None)]
#[case::marked_string(json!(raw_json::encode("\"1\"")), None)]
#[case::null(Value::Null, None)]
fn numbers_compare_as_python_compares_them(
#[case] value: Value,
#[case] expected: Option<f64>,
) {
assert_eq!(python_number(&value), expected);
}
#[test]
fn nan_is_a_number_unequal_to_everything() {
let nan = json!(raw_json::encode("NaN"));
assert!(python_number(&nan).is_some_and(f64::is_nan));
assert!(!is_one(&Recognized::Unrecognized(nan)));
}
#[rstest]
#[case::known(Recognized::Known(0.5), "0.5")]
#[case::known_wide(Recognized::Known(1e16), "1e+16")]
#[case::int(Recognized::Unrecognized(json!(0)), "0")]
#[case::bool(Recognized::Unrecognized(json!(false)), "False")]
#[case::string(Recognized::Unrecognized(json!("0.5")), "0.5")]
#[case::nan(Recognized::Unrecognized(json!(raw_json::encode("NaN"))), "nan")]
#[case::wide(Recognized::Unrecognized(json!(raw_json::encode("-1180591620717411303424"))), "-1180591620717411303424")]
fn values_print_as_python_prints_them(#[case] value: Recognized<f64>, #[case] expected: &str) {
assert_eq!(python_str(&value), expected);
}
}

210
litellm-rust/crates/llms/src/anthropic/messages/thinking.rs
View file

@ -1,17 +1,20 @@
use litellm_host::raw_json;
use litellm_llms_types::{
formats::{
chat_completions::ReasoningEffort,
messages::{
AdaptiveThinking, DisabledThinking, EffortLevel, MessagesOptionalParams,
MessagesRequest, OutputConfig, ThinkingConfig, ThinkingDisplay,
AdaptiveThinking, DisabledThinking, EffortLevel, EnabledThinking,
MessagesOptionalParams, MessagesRequest, OutputConfig, ThinkingConfig, ThinkingDisplay,
},
},
recognized::Recognized,
};
use litellm_python_compat::{json::from_json, repr::repr, truthy::truthy};
use litellm_python_compat::{
Value as PythonValue, json::from_json, repr::repr, reprlib::reprlib, truthy::truthy,
};
use serde_json::Value;
use strum::VariantArray;
use super::scalar::{is_one, python_number, python_str, python_value};
use crate::base_llm::messages::context::{
MessagesModelCapabilities, ThinkingBudgets, ThinkingContext,
};
@ -56,13 +59,61 @@ fn unmapped_effort(effort: &Value) -> Error {
Error::InvalidRequest(crate::ErrorDetail::InvalidChoice {
field: "reasoning effort",
actual: repr(&from_json(effort.clone())),
choices: ReasoningEffort::VARIANTS
.iter()
.map(|effort| <&'static str>::from(*effort))
.collect(),
choices: vec!["minimal", "low", "medium", "high", "xhigh", "max", "none"],
})
}
fn invalid_adaptive_effort(effort: &Value) -> Error {
Error::InvalidRequest(crate::ErrorDetail::InvalidReasoningEffort {
actual: repr(&from_json(effort.clone())),
})
}
const NON_STRING_EFFORT_REASON: &str = "it is not a string. Send reasoning_effort as a level name such as 'high' to set a reasoning level.";
const NON_NUMERIC_BUDGET_REASON: &str =
"it is not a number. Send budget_tokens as an integer to set a reasoning level.";
fn warn_dropped_reasoning(
param: &str,
value: &PythonValue,
model: &str,
reason: &str,
drop_params_gated: bool,
) {
let gate = if drop_params_gated {
" (drop_params=True)"
} else {
""
};
litellm_tracing::warn!(
"Dropping {param}={} for model={}{gate} because {reason}",
reprlib(value),
repr(&PythonValue::Str(model.to_owned()))
);
}
fn warn_non_string_reasoning_effort(effort: &Value, model: &str) {
let value = python_value(effort);
let requested = match &value {
PythonValue::Str(_) => return,
PythonValue::Dict(entries) => entries
.iter()
.find(|(key, _)| *key == PythonValue::Str("effort".to_owned()))
.map(|(_, level)| level),
other => Some(other),
};
if matches!(requested, None | Some(PythonValue::None)) {
return;
}
warn_dropped_reasoning(
"reasoning_effort",
&value,
model,
NON_STRING_EFFORT_REASON,
false,
);
}
fn unsupported_effort(level: EffortLevel, model: &str) -> Error {
Error::InvalidRequest(crate::ErrorDetail::UnsupportedValue {
field: "effort",
@ -82,11 +133,56 @@ fn output_effort(effort: ReasoningEffort) -> Option<EffortLevel> {
}
}
fn fit_budget_to_max_tokens(budget_tokens: u64, max_tokens: Option<u64>) -> Option<u64> {
let Some(max_tokens) = max_tokens else {
return Some(budget_tokens);
/// `cap_thinking_budget_to_max_tokens`, whose first comparison (`max_tokens <= 1024`) raises
/// `TypeError` for a `max_tokens` Python cannot order against an int, such as a string.
fn fit_budget_to_max_tokens(
budget_tokens: u64,
max_tokens: Option<&Recognized<u64>>,
) -> Result<Option<Recognized<u64>>, Error> {
let max_tokens = match max_tokens {
None => return Ok(Some(Recognized::Known(budget_tokens))),
Some(Recognized::Known(max_tokens)) => *max_tokens,
Some(Recognized::Unrecognized(other)) => {
return fit_budget_to_python_number(budget_tokens, other);
}
};
(max_tokens > ANTHROPIC_MIN_THINKING_BUDGET_TOKENS).then(|| budget_tokens.min(max_tokens - 1))
Ok((max_tokens > ANTHROPIC_MIN_THINKING_BUDGET_TOKENS)
.then(|| Recognized::Known(budget_tokens.min(max_tokens - 1))))
}
fn enabled_thinking(budget_tokens: Recognized<u64>) -> ThinkingConfig {
ThinkingConfig::Enabled(EnabledThinking {
budget_tokens: Some(budget_tokens),
..EnabledThinking::default()
})
}
fn fit_budget_to_python_number(
budget_tokens: u64,
max_tokens: &Value,
) -> Result<Option<Recognized<u64>>, Error> {
let Some(max_tokens) = python_number(max_tokens) else {
if max_tokens.is_null() {
return Ok(Some(Recognized::Known(budget_tokens)));
}
return Err(Error::InvalidRequest(
crate::ErrorDetail::UnorderableTypes {
op: "<=",
left: python_type_name(max_tokens),
right: "int",
},
));
};
if max_tokens <= ANTHROPIC_MIN_THINKING_BUDGET_TOKENS as f64 {
return Ok(None);
}
if (budget_tokens as f64) < max_tokens {
return Ok(Some(Recognized::Known(budget_tokens)));
}
Ok(Some(Recognized::Unrecognized(
serde_json::Number::from_f64(max_tokens - 1.0)
.map_or_else(|| Value::String(raw_json::encode("NaN")), Value::Number),
)))
}
fn known_thinking(request: &MessagesRequest) -> Option<&ThinkingConfig> {
@ -161,10 +257,17 @@ fn translate_reasoning_effort(
};
let effort = match reasoning_effort {
Recognized::Known(effort) => effort,
Recognized::Unrecognized(value @ Value::String(_)) => {
return Err(unmapped_effort(&value));
Recognized::Unrecognized(value) if matches!(python_value(&value), PythonValue::Str(_)) => {
return Err(if context.capabilities.supports_adaptive_thinking {
invalid_adaptive_effort(&value)
} else {
unmapped_effort(&value)
});
}
Recognized::Unrecognized(value) => {
warn_non_string_reasoning_effort(&value, &request.model);
return Ok(request);
}
Recognized::Unrecognized(_) => return Ok(request),
};
let (Some(level), Some(budget)) = (
output_effort(effort),
@ -203,10 +306,20 @@ fn translate_reasoning_effort(
..request
});
}
let Some(budget) = fit_budget_to_max_tokens(budget, request.params.max_tokens) else {
let Some(budget) = fit_budget_to_max_tokens(budget, request.params.max_tokens.as_ref())? else {
litellm_tracing::warn!(
"Dropping `thinking` mapped from reasoning_effort={} for model={}: max_tokens={} is too small to fit the minimum thinking budget.",
<&'static str>::from(effort),
request.model,
request
.params
.max_tokens
.as_ref()
.map_or_else(|| "None".to_owned(), python_str)
);
return Ok(request);
};
let enabled = ThinkingConfig::enabled(budget);
let enabled = enabled_thinking(budget);
Ok(MessagesRequest {
params: MessagesOptionalParams {
thinking: Some(
@ -340,6 +453,15 @@ fn translate_legacy_thinking_for_adaptive_model(
crate::ErrorDetail::NonNumericThinkingBudget { actual },
));
}
if let (None, Some(Recognized::Unrecognized(dropped))) = (budget, &enabled.budget_tokens) {
warn_dropped_reasoning(
"thinking.budget_tokens",
&python_value(dropped),
&request.model,
NON_NUMERIC_BUDGET_REASON,
true,
);
}
let thinking = ThinkingConfig::Adaptive(AdaptiveThinking {
display: kept_display(enabled.display.as_ref(), context),
extra: carried_thinking_keys(&enabled.extra, context),
@ -396,11 +518,20 @@ fn translate_adaptive_effort_for_non_adaptive_model(
} else {
None
};
let thinking = match budget {
Some(budget) => fit_budget_to_max_tokens(budget, request.params.max_tokens.as_ref())?
.map(|budget| Recognized::Known(enabled_thinking(budget))),
None => None,
};
if thinking.is_none() {
litellm_tracing::warn!(
"Dropping adaptive `thinking`/`output_config.effort` for model={}: the model does not support extended thinking, or max_tokens is too small to fit the minimum thinking budget.",
request.model
);
}
Ok(MessagesRequest {
params: MessagesOptionalParams {
thinking: budget
.and_then(|budget| fit_budget_to_max_tokens(budget, request.params.max_tokens))
.map(|budget| Recognized::Known(ThinkingConfig::enabled(budget))),
thinking,
output_config: without_effort(request.params.output_config),
..request.params
},
@ -418,7 +549,8 @@ fn drop_incompatible_temperature_for_thinking(
let pinned = request
.params
.temperature
.is_some_and(|temperature| temperature != 1.0);
.as_ref()
.is_some_and(|temperature| !is_one(temperature));
let thinking_enabled = matches!(known_thinking(&request), Some(ThinkingConfig::Enabled(_)));
let effort_enabled = known_effort(&request).is_some();
if !pinned || !(thinking_enabled || effort_enabled) {
@ -451,7 +583,7 @@ mod tests {
use super::*;
use crate::base_llm::messages::context::SupportedEffortTiers;
const EFFORT_CHOICES: &str = "'none', 'minimal', 'low', 'medium', 'high', 'xhigh', 'max'";
const EFFORT_CHOICES: &str = "'minimal', 'low', 'medium', 'high', 'xhigh', 'max', 'none'";
fn request(fields: Value) -> MessagesRequest {
let mut body = serde_json::json!({"model": "claude", "messages": [{"role": "user", "content": "Hello"}]});
@ -819,17 +951,17 @@ mod tests {
#[case::invalid_on_adaptive_model(
opus_4_7(),
serde_json::json!({"max_tokens": 1024, "reasoning_effort": "invalid"}),
format!("Unmapped reasoning effort: 'invalid'. Must be one of: {EFFORT_CHOICES}.")
format!("Invalid reasoning_effort: 'invalid'. Must be one of: {EFFORT_CHOICES}")
)]
#[case::disabled_on_adaptive_model(
opus_4_7(),
serde_json::json!({"max_tokens": 1024, "reasoning_effort": "disabled"}),
format!("Unmapped reasoning effort: 'disabled'. Must be one of: {EFFORT_CHOICES}.")
format!("Invalid reasoning_effort: 'disabled'. Must be one of: {EFFORT_CHOICES}")
)]
#[case::empty_on_adaptive_model(
opus_4_7(),
serde_json::json!({"max_tokens": 1024, "reasoning_effort": ""}),
format!("Unmapped reasoning effort: ''. Must be one of: {EFFORT_CHOICES}.")
format!("Invalid reasoning_effort: ''. Must be one of: {EFFORT_CHOICES}")
)]
#[case::xhigh_without_xhigh_tier_on_4_6(
sonnet_4_6(),
@ -1434,6 +1566,34 @@ mod tests {
);
}
#[rstest]
#[case::reasoning_effort(json!({"max_tokens": "8192", "reasoning_effort": "high"}), "str")]
#[case::caller_thinking(
json!({"max_tokens": "8192", "reasoning_effort": "low", "thinking": {"type": "enabled", "budget_tokens": 2048}}),
"str"
)]
#[case::adaptive_interface(
json!({"max_tokens": "8192", "thinking": {"type": "adaptive"}, "output_config": {"effort": "high"}}),
"str"
)]
#[case::list(json!({"max_tokens": [8192], "reasoning_effort": "high"}), "list")]
fn an_unorderable_max_tokens_raises_python_type_error_while_capping_the_budget(
haiku_4_5: MessagesModelCapabilities,
#[case] fields: Value,
#[case] left: &'static str,
) {
assert_eq!(
translate(haiku_4_5, fields),
Err(Error::InvalidRequest(
crate::ErrorDetail::UnorderableTypes {
op: "<=",
left,
right: "int",
}
))
);
}
#[rstest]
#[case::summarized("summarized", false, Some("summarized"))]
#[case::updates_kept("updates", true, Some("updates"))]

660
litellm-rust/crates/llms/src/anthropic/messages/transformation.rs
View file

@ -1,29 +1,45 @@
use litellm_auth::CredentialPlacement;
use litellm_host::interceptors::{RetriedRequest, RetryReason};
use litellm_http::request::{with_header, without_headers};
use litellm_llms_types::{
formats::messages::{
ContentBlock, ContentBlockType, ContextEdit, ContextManagement, Message,
ContentBlock, ContentBlockType, ContextEdit, ContextManagement, Message, MessageContent,
MessagesOptionalParams, MessagesRequest, Speed, ThinkingConfig, ThinkingDisplay,
},
providers::anthropic::{AnthropicBeta, BetaProvider, BetaSet, DEFAULT_HEADERS},
recognized::Recognized,
providers::anthropic::{AnthropicBeta, BETA_HEADER, BetaProvider, BetaSet, DEFAULT_HEADERS},
recognized::{Recognized, map_known_runs},
};
use litellm_router_types::LitellmParams;
use serde_json::{Map, Value, json};
use super::{handler::shape_anthropic_messages_request, thinking::translate_thinking};
use crate::base_llm::messages::context::MessagesTransformContext;
use super::{
handler::shape_anthropic_messages_request,
scalar::{is_one, python_str, python_value},
thinking::translate_thinking,
};
use crate::anthropic::beta_headers::accepted_or_unlisted;
use crate::base_llm::messages::{
context::{MessagesShapeOptions, MessagesTransformContext},
tools::ToolCustomField,
};
use crate::{
Error,
anthropic::common_utils::{
ANTHROPIC_API_BASE_ENV, ANTHROPIC_API_KEY_ENV, ANTHROPIC_AUTH_TOKEN_ENV,
ANTHROPIC_BASE_URL_ENV, OauthHandling, complete_anthropic_url, get_auth_header,
has_advisor_tool, has_anthropic_credential, is_tool_search_used, merge_beta_headers,
optionally_handle_anthropic_oauth, requires_native_compaction_beta, strip_advisor_blocks,
strip_encrypted_reasoning_blocks,
ANTHROPIC_BASE_URL_ENV, OauthHandling,
anthropic_request_errors_on_thinking_binding_mismatch, complete_anthropic_url,
existing_betas, get_auth_header, has_advisor_tool, has_anthropic_credential,
is_anthropic_invalid_thinking_block_error, is_anthropic_thinking_binding_mismatch_error,
is_claude_strict_tools_rejection, is_tool_search_used, merge_beta_headers,
minted_oauth_environment, optionally_handle_anthropic_oauth,
requires_native_compaction_beta, strip_advisor_blocks, strip_encrypted_reasoning_blocks,
strip_thinking_blocks_from_anthropic_messages_request, tools_without_strict,
},
base_llm::{
auth::AuthScheme,
messages::transformation::{BaseMessagesConfig, Headers, ValidatedEnvironment},
messages::transformation::{
BaseMessagesConfig, FederatedEnvironment, Headers, ValidatedEnvironment,
},
},
};
@ -35,9 +51,9 @@ impl BaseMessagesConfig for AnthropicMessagesConfig {
fn shape_request(
&self,
request: MessagesRequest,
reasoning_auto_summary: bool,
options: MessagesShapeOptions,
) -> Result<MessagesRequest, Error> {
shape_anthropic_messages_request(request, reasoning_auto_summary)
shape_anthropic_messages_request(request, options)
}
fn get_complete_url(
@ -59,6 +75,12 @@ impl BaseMessagesConfig for AnthropicMessagesConfig {
transform_messages_request(request, &context.keeping_display_updates())
}
fn tool_custom_field(&self, model: &str) -> ToolCustomField {
ToolCustomField::Rewrite {
hoist_defer_loading: is_claude_model(model),
}
}
fn secret_names(&self) -> &'static [&'static str] {
&[
ANTHROPIC_API_KEY_ENV,
@ -79,31 +101,22 @@ impl BaseMessagesConfig for AnthropicMessagesConfig {
_litellm_params: &LitellmParams,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<ValidatedEnvironment, Error> {
let headers = match optionally_handle_anthropic_oauth(headers, api_key) {
OauthHandling::Bearer { headers, token } => {
return Ok(ValidatedEnvironment {
headers,
auth: AuthScheme::Credential {
placement: CredentialPlacement::Bearer,
secret: token,
},
});
}
OauthHandling::Untouched(headers) => headers,
};
if has_anthropic_credential(&headers) {
return Ok(ValidatedEnvironment {
headers,
auth: AuthScheme::Forwarded,
});
match anthropic_environment(headers, api_key, env_lookup) {
FederatedEnvironment::Validated(validated) => Ok(validated),
FederatedEnvironment::Unauthenticated { missing, .. } => Err(missing),
}
let auth = get_auth_header(api_key, env_lookup).ok_or(Error::Auth(
litellm_auth::Error::MissingApiKey {
provider: "Anthropic",
environment_variable: ANTHROPIC_API_KEY_ENV,
},
))?;
Ok(ValidatedEnvironment { headers, auth })
}
fn validate_federated_environment(
&self,
headers: Headers,
api_key: Option<&str>,
_model: &str,
_optional_params: &MessagesOptionalParams,
_litellm_params: &LitellmParams,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<FederatedEnvironment, Error> {
Ok(anthropic_environment(headers, api_key, env_lookup))
}
fn default_headers(&self) -> &'static [(&'static str, &'static str)] {
@ -116,6 +129,157 @@ impl BaseMessagesConfig for AnthropicMessagesConfig {
provider_feature_betas(request, BetaProvider::Anthropic),
)
}
fn retry_after_http_error(
&self,
status: u16,
error_text: &str,
request_body: &Value,
context: &MessagesTransformContext,
) -> Option<RetriedRequest> {
retry_anthropic_messages_after_http_error(status, error_text, request_body, context)
}
fn headers_for_url(&self, headers: Headers, url: &str) -> Headers {
if !is_foundry_url(url) {
return headers;
}
let betas = existing_betas(&headers);
if betas.is_empty() {
return headers;
}
let kept = accepted_or_unlisted(betas, BetaProvider::AzureAi);
if kept.is_empty() {
return without_headers(headers, &[BETA_HEADER]);
}
with_header(headers, BETA_HEADER, kept.to_string())
}
}
const FOUNDRY_HOST_SUFFIXES: [&str; 3] = [
"services.ai.azure.com",
"openai.azure.com",
"cognitiveservices.azure.com",
];
fn is_foundry_url(url: &str) -> bool {
let absolute = if url.contains("://") {
url.to_string()
} else {
format!("https://{url}")
};
url::Url::parse(&absolute)
.ok()
.and_then(|parsed| parsed.host_str().map(str::to_ascii_lowercase))
.is_some_and(|host| {
FOUNDRY_HOST_SUFFIXES
.iter()
.any(|suffix| host == *suffix || host.ends_with(&format!(".{suffix}")))
})
}
pub(crate) fn is_claude_model(model: &str) -> bool {
model.to_lowercase().contains("claude")
}
pub fn should_retry_anthropic_messages_on_http_error(
status: u16,
error_text: &str,
request_body: &Value,
context: &MessagesTransformContext,
) -> bool {
if status != 400 {
return false;
}
let drops_rejected_strict_tools = context.drops_unsupported_params
&& is_claude_strict_tools_rejection(status, error_text)
&& tools_without_strict(request_body.get("tools")).is_some();
if drops_rejected_strict_tools {
return true;
}
let caller_asked_for_this_error =
anthropic_request_errors_on_thinking_binding_mismatch(request_body)
&& is_anthropic_thinking_binding_mismatch_error(error_text);
!caller_asked_for_this_error && is_anthropic_invalid_thinking_block_error(error_text)
}
pub fn transform_anthropic_messages_request_on_http_error(
status: u16,
error_text: &str,
request_body: Value,
) -> RetriedRequest {
let mut reasons = Vec::new();
let mut body = if status == 400 && is_anthropic_invalid_thinking_block_error(error_text) {
reasons.push(RetryReason::ThinkingSignature);
strip_thinking_blocks_from_anthropic_messages_request(request_body)
} else {
request_body
};
let without_strict = if is_claude_strict_tools_rejection(status, error_text) {
tools_without_strict(body.get("tools"))
} else {
None
};
if let (Some(without_strict), Value::Object(fields)) = (without_strict, &mut body) {
reasons.push(RetryReason::StrictTools {
dropped: without_strict.dropped_count,
});
fields.insert("tools".to_owned(), Value::Array(without_strict.tools));
}
RetriedRequest { reasons, body }
}
pub fn retry_anthropic_messages_after_http_error(
status: u16,
error_text: &str,
request_body: &Value,
context: &MessagesTransformContext,
) -> Option<RetriedRequest> {
should_retry_anthropic_messages_on_http_error(status, error_text, request_body, context).then(
|| {
transform_anthropic_messages_request_on_http_error(
status,
error_text,
request_body.clone(),
)
},
)
}
fn anthropic_environment(
headers: Headers,
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> FederatedEnvironment {
let headers = match optionally_handle_anthropic_oauth(headers, api_key) {
OauthHandling::Bearer { headers, token } => {
return FederatedEnvironment::Validated(ValidatedEnvironment {
headers,
auth: AuthScheme::Credential {
placement: CredentialPlacement::Bearer,
secret: token,
},
});
}
OauthHandling::Untouched(headers) => headers,
};
if has_anthropic_credential(&headers) {
return FederatedEnvironment::Validated(ValidatedEnvironment {
headers,
auth: AuthScheme::Forwarded,
});
}
match get_auth_header(api_key, env_lookup) {
Some(auth) => FederatedEnvironment::Validated(ValidatedEnvironment { headers, auth }),
None => FederatedEnvironment::Unauthenticated {
headers,
missing: Error::Auth(litellm_auth::Error::MissingApiKey {
provider: "Anthropic",
environment_variable: ANTHROPIC_API_KEY_ENV,
}),
federate: minted_oauth_environment,
},
}
}
pub(crate) fn transform_messages_request(
@ -125,6 +289,7 @@ pub(crate) fn transform_messages_request(
if request.params.max_tokens.is_none() {
return Err(Error::MissingField("max_tokens"));
}
let request = without_unstructured_output_format(request, context);
let request = drop_unsupported_params(request, context)?;
let request = translate_thinking(request, &context.thinking)?;
let context_management = request
@ -132,13 +297,20 @@ pub(crate) fn transform_messages_request(
.context_management
.clone()
.map(map_openai_context_management_to_anthropic);
let messages = if has_advisor_tool(request.params.tools.as_deref()) {
let messages = if has_advisor_tool(request.params.tools.as_deref())
|| continues_message_thread(&request.params)
{
request.messages
} else {
strip_advisor_blocks(request.messages)
map_known_runs(request.messages, strip_advisor_blocks)
};
let messages = if context.drops_unsupported_params {
drop_non_string_tool_result_texts(messages)
} else {
messages
};
Ok(MessagesRequest {
messages: strip_encrypted_reasoning_blocks(messages),
messages: map_known_runs(messages, strip_encrypted_reasoning_blocks),
params: MessagesOptionalParams {
context_management,
extra: without_null_safeguards(request.params.extra),
@ -148,6 +320,84 @@ pub(crate) fn transform_messages_request(
})
}
fn without_unstructured_output_format(
request: MessagesRequest,
context: &MessagesTransformContext,
) -> MessagesRequest {
if !context.drops_unsupported_params
|| request
.params
.output_format
.as_ref()
.is_none_or(Value::is_object)
{
return request;
}
MessagesRequest {
params: MessagesOptionalParams {
output_format: None,
..request.params
},
..request
}
}
fn is_python_str(value: &Value) -> bool {
match value {
Value::String(text) => {
litellm_host::raw_json::decode(text).is_none_or(|raw| raw.starts_with('"'))
}
_ => false,
}
}
fn is_non_string_text(part: &Value) -> bool {
part.get("type").and_then(Value::as_str) == Some("text")
&& part.get("text").is_some_and(|text| !is_python_str(text))
}
fn without_non_string_texts(content: &mut Value) {
if let Value::Array(parts) = content {
parts.retain(|part| !is_non_string_text(part));
}
}
fn drop_non_string_tool_result_texts(
messages: Vec<Recognized<Message>>,
) -> Vec<Recognized<Message>> {
messages
.into_iter()
.map(|message| match message {
Recognized::Known(mut message) => {
if message.role == "user"
&& let MessageContent::Blocks(blocks) = &mut message.content
{
blocks
.iter_mut()
.filter(|block| block.is_type(ContentBlockType::ToolResult))
.filter_map(|block| block.content.as_mut())
.for_each(without_non_string_texts);
}
Recognized::Known(message)
}
Recognized::Unrecognized(mut message) => {
if message.get("role").and_then(Value::as_str) == Some("user")
&& let Some(Value::Array(blocks)) = message.get_mut("content")
{
blocks
.iter_mut()
.filter(|block| {
block.get("type").and_then(Value::as_str) == Some("tool_result")
})
.filter_map(|block| block.get_mut("content"))
.for_each(without_non_string_texts);
}
Recognized::Unrecognized(message)
}
})
.collect()
}
fn without_null_safeguards(extra: Map<String, Value>) -> Map<String, Value> {
extra
.into_iter()
@ -193,22 +443,45 @@ fn mapped_host_betas(params: &MessagesOptionalParams) -> BetaSet {
.collect()
}
fn continues_message_thread(params: &MessagesOptionalParams) -> bool {
let Some(Value::Object(thread)) = params.extra.get("thread") else {
return false;
};
thread.get("type").and_then(Value::as_str) == Some("continue")
&& thread
.get("previous_message_id")
.is_none_or(Value::is_string)
}
fn is_threaded(params: &MessagesOptionalParams) -> bool {
params
.extra
.get("thread")
.is_some_and(|thread| !thread.is_null())
}
fn feature_betas(request: &MessagesRequest) -> BetaSet {
let params = &request.params;
let tools = params.tools.as_deref();
let scanned_messages: &[Recognized<Message>] = if continues_message_thread(params) {
&[]
} else {
&request.messages
};
[
requires_native_compaction_beta(params.compaction.as_ref(), &request.messages)
requires_native_compaction_beta(params.compaction.as_ref(), scanned_messages)
.then_some(AnthropicBeta::Compact20260904),
uses_structured_output(params).then_some(AnthropicBeta::StructuredOutputs20251113),
(params.speed == Some(Recognized::Known(Speed::Fast)))
.then_some(AnthropicBeta::FastMode20260201),
messages_carry_output_config(&request.messages)
messages_carry_output_config(scanned_messages)
.then_some(AnthropicBeta::PerTurnControl20260701),
is_threaded(params).then_some(AnthropicBeta::MessageThreads20260812),
uses_thinking_display_updates(params)
.then_some(AnthropicBeta::ThinkingDisplayUpdates20260818),
uses_thinking_block_binding(params)
.then_some(AnthropicBeta::ThinkingBindingControls20260801),
messages_carry_tool_changes(&request.messages)
messages_carry_tool_changes(scanned_messages)
.then_some(AnthropicBeta::MidConversationToolChanges20260701),
has_advisor_tool(tools).then_some(AnthropicBeta::AdvisorTool20260301),
is_tool_search_used(tools).then_some(AnthropicBeta::AdvancedToolUse20251120),
@ -254,9 +527,10 @@ fn carries_safeguards(params: &MessagesOptionalParams) -> bool {
.is_some_and(|value| !value.is_null())
}
fn messages_carry_output_config(messages: &[Message]) -> bool {
fn messages_carry_output_config(messages: &[Recognized<Message>]) -> bool {
messages
.iter()
.filter_map(Recognized::known)
.any(|message| message.extra.contains_key("output_config"))
}
@ -301,9 +575,10 @@ fn is_tool_reference_change(block: &ContentBlock) -> bool {
== Some("tool_reference")
}
fn messages_carry_tool_changes(messages: &[Message]) -> bool {
fn messages_carry_tool_changes(messages: &[Recognized<Message>]) -> bool {
messages
.iter()
.filter_map(Recognized::known)
.filter(|message| message.role == "system")
.flat_map(Message::blocks)
.any(is_tool_reference_change)
@ -334,7 +609,7 @@ pub fn drop_unsupported_params(
let params = request.params;
let speed = match &params.speed {
Some(speed) if !capabilities.supports_speed => {
reject("speed", format!("'{}'", speed_text(speed)), "")?;
reject("speed", speed_repr(speed), "")?;
None
}
_ => params.speed.clone(),
@ -346,21 +621,21 @@ pub fn drop_unsupported_params(
});
}
let temperature = match params.temperature {
Some(temperature) if temperature != 1.0 => {
Some(temperature) if !is_one(&temperature) => {
reject(
"temperature",
json!(temperature).to_string(),
python_str(&temperature),
"Only temperature=1 is supported. ",
)?;
None
}
temperature => temperature,
};
if let Some(top_p) = params.top_p {
reject("top_p", json!(top_p).to_string(), "")?;
if let Some(top_p) = &params.top_p {
reject("top_p", python_str(top_p), "")?;
}
if let Some(top_k) = params.top_k {
reject("top_k", json!(top_k).to_string(), "")?;
if let Some(top_k) = &params.top_k {
reject("top_k", python_str(top_k), "")?;
}
Ok(MessagesRequest {
params: MessagesOptionalParams {
@ -374,11 +649,10 @@ pub fn drop_unsupported_params(
})
}
fn speed_text(speed: &Recognized<Speed>) -> String {
fn speed_repr(speed: &Recognized<Speed>) -> String {
match speed {
Recognized::Known(speed) => <&'static str>::from(*speed).to_string(),
Recognized::Unrecognized(Value::String(text)) => text.clone(),
Recognized::Unrecognized(other) => other.to_string(),
Recognized::Known(speed) => format!("'{}'", <&'static str>::from(*speed)),
Recognized::Unrecognized(other) => litellm_python_compat::repr::repr(&python_value(other)),
}
}
@ -437,6 +711,7 @@ mod tests {
use super::*;
use crate::anthropic::common_utils::ENCRYPTED_REASONING_SIGNATURE_PREFIX;
use litellm_http::request::header_values;
type Env = &'static [(&'static str, &'static str)];
@ -679,6 +954,21 @@ mod tests {
json!({"speed": "fast"}),
"claude does not support speed='fast'. To drop unsupported params, set `litellm.drop_params = True`."
)]
#[case::integer_temperature_keeps_no_decimal(
sampling_removed(),
json!({"temperature": 0}),
"claude does not support temperature=0. Only temperature=1 is supported. To drop unsupported params, set `litellm.drop_params = True`."
)]
#[case::integer_top_p(
sampling_removed(),
json!({"top_p": 1}),
"claude does not support top_p=1. To drop unsupported params, set `litellm.drop_params = True`."
)]
#[case::non_string_speed_is_reprd(
unmapped(),
json!({"speed": 2}),
"claude does not support speed=2. To drop unsupported params, set `litellm.drop_params = True`."
)]
#[case::speed_ahead_of_sampling_params(
sampling_removed(),
json!({"speed": "fast", "temperature": 0.5}),
@ -696,6 +986,28 @@ mod tests {
assert_eq!(detail.to_string(), message);
}
#[rstest]
#[case::string_dropped(json!("json"), true, None)]
#[case::string_kept_without_drop_params(json!("json"), false, Some(json!("json")))]
#[case::object_kept(json!({"type": "json_schema"}), true, Some(json!({"type": "json_schema"})))]
fn a_non_object_output_format_is_dropped_under_drop_params(
#[case] output_format: Value,
#[case] drops_unsupported_params: bool,
#[case] expected: Option<Value>,
) {
let context = MessagesTransformContext {
drops_unsupported_params,
..MessagesTransformContext::with_lookup(unmapped(), false, &no_env)
};
let transformed = ANTHROPIC_MESSAGES_CONFIG
.transform_anthropic_messages_request(
request(json!({"output_format": output_format})),
&context,
)
.unwrap();
assert_eq!(transformed.params.output_format, expected);
}
#[rstest]
#[case::compaction_threshold(
json!([{"type": "compaction", "compact_threshold": 200000}]),
@ -802,6 +1114,40 @@ mod tests {
);
}
#[rstest]
#[case::continue_turn(json!({"type": "continue", "previous_message_id": "msg_1"}), false)]
#[case::continue_without_a_previous_message_id(json!({"type": "continue"}), false)]
#[case::continue_with_extra_keys(json!({"type": "continue", "previous_message_id": "msg_1", "note": 1}), false)]
#[case::create_turn(json!({"type": "create"}), true)]
#[case::continue_with_a_null_previous_message_id(json!({"type": "continue", "previous_message_id": null}), true)]
#[case::continue_with_a_numeric_previous_message_id(json!({"type": "continue", "previous_message_id": 1}), true)]
#[case::not_an_object(json!("continue"), true)]
fn advisor_history_is_kept_on_a_thread_continue_turn(
#[case] thread: Value,
#[case] stripped: bool,
unmapped: MessagesModelCapabilities,
) {
let messages = if stripped {
json!([
{"role": "user", "content": "Build a worker pool."},
{"role": "assistant", "content": [
{"type": "text", "text": "Let me consult the advisor."},
{"type": "text", "text": "Here is the implementation."}
]}
])
} else {
advisor_history()
};
assert_eq!(
transform(
json!({"messages": advisor_history(), "thread": thread.clone()}),
unmapped,
false
),
Ok(body(json!({"messages": messages, "thread": thread})))
);
}
#[rstest]
fn bridge_minted_reasoning_is_stripped_from_the_wire(unmapped: MessagesModelCapabilities) {
let messages = json!([
@ -864,6 +1210,7 @@ mod tests {
keeps_display_updates: false,
},
drop_params: true,
drops_unsupported_params: false,
}
);
return;
@ -1165,6 +1512,55 @@ mod tests {
]}),
&[]
)]
#[case::thread_create(json!({"thread": {"type": "create"}}), &["message-threads-2026-08-12"])]
#[case::thread_continue(
json!({"thread": {"type": "continue", "previous_message_id": "msg_1"}}),
&["message-threads-2026-08-12"]
)]
#[case::thread_of_any_shape(json!({"thread": "opaque"}), &["message-threads-2026-08-12"])]
#[case::null_thread(json!({"thread": null}), &[])]
#[case::thread_create_still_scans_history(
json!({
"thread": {"type": "create"},
"messages": [
{"role": "assistant", "content": [{"type": "compaction", "content": "summary", "signature": "sig"}]},
{"role": "user", "content": "hi", "output_config": {"effort": "low"}},
{"role": "system", "content": [{"type": "tool_addition", "tool": {"type": "tool_reference", "name": "ping"}}]},
],
}),
&[
"compact-2026-09-04",
"message-threads-2026-08-12",
"mid-conversation-tool-changes-2026-07-01",
"per-turn-control-2026-07-01",
]
)]
#[case::thread_continue_skips_history(
json!({
"thread": {"type": "continue", "previous_message_id": "msg_1"},
"messages": [
{"role": "assistant", "content": [{"type": "compaction", "content": "summary", "signature": "sig"}]},
{"role": "user", "content": "hi", "output_config": {"effort": "low"}},
{"role": "system", "content": [{"type": "tool_addition", "tool": {"type": "tool_reference", "name": "ping"}}]},
],
}),
&["message-threads-2026-08-12"]
)]
#[case::thread_continue_keeps_param_betas(
json!({
"thread": {"type": "continue", "previous_message_id": "msg_1"},
"compaction": {},
"speed": "fast",
}),
&["compact-2026-09-04", "fast-mode-2026-02-01", "message-threads-2026-08-12"]
)]
#[case::malformed_continue_scans_history(
json!({
"thread": {"type": "continue", "previous_message_id": 7},
"messages": [{"role": "user", "content": "hi", "output_config": {"effort": "low"}}],
}),
&["message-threads-2026-08-12", "per-turn-control-2026-07-01"]
)]
fn feature_betas_follow_the_request(#[case] fields: Value, #[case] expected: &[&str]) {
assert_eq!(feature_betas(&request(fields)), betas(expected));
}
@ -1352,4 +1748,158 @@ mod tests {
.collect();
assert_eq!(undeclared, Vec::<&String>::new());
}
const FOUNDRY: &str = "https://team.services.ai.azure.com/anthropic/v1/messages";
#[rstest]
#[case::drops_a_beta_foundry_lacks(
FOUNDRY,
"advisor-tool-2026-03-01,context-1m-2025-08-07",
Some("context-1m-2025-08-07")
)]
#[case::forwards_unknown_betas_sorted(
FOUNDRY,
"zzz,files-api-2025-04-14,aaa",
Some("aaa,files-api-2025-04-14,zzz")
)]
#[case::removes_the_header_when_nothing_is_left(FOUNDRY, "advisor-tool-2026-03-01", None)]
#[case::cognitive_services_host(
"team.cognitiveservices.azure.com/anthropic",
"advisor-tool-2026-03-01",
None
)]
#[case::anthropic_host_is_untouched(
"https://api.anthropic.com/v1/messages",
"advisor-tool-2026-03-01",
Some("advisor-tool-2026-03-01")
)]
#[case::lookalike_host_is_untouched(
"https://notservices.ai.azure.com.example/v1/messages",
"advisor-tool-2026-03-01",
Some("advisor-tool-2026-03-01")
)]
fn foundry_urls_filter_betas_through_the_azure_ai_row(
#[case] url: &str,
#[case] betas: &str,
#[case] expected: Option<&str>,
) {
let filtered = ANTHROPIC_MESSAGES_CONFIG
.headers_for_url(headers(&[("x-a", "1"), (BETA_HEADER, betas)]), url);
assert_eq!(
header_values(&filtered, BETA_HEADER).collect::<Vec<_>>(),
expected.into_iter().collect::<Vec<_>>()
);
assert_eq!(header_values(&filtered, "x-a").collect::<Vec<_>>(), ["1"]);
}
#[rstest]
fn foundry_url_without_betas_keeps_the_headers() {
let sent = headers(&[("x-a", "1")]);
assert_eq!(
ANTHROPIC_MESSAGES_CONFIG.headers_for_url(sent.clone(), FOUNDRY),
sent
);
}
#[test]
fn the_http_error_retry_matches_the_cases_shared_with_python() {
let fixture: Value = serde_json::from_str(include_str!(concat!(
env!("CARGO_MANIFEST_DIR"),
"/tests/fixtures/anthropic_messages_http_error_retry.json"
)))
.unwrap();
let table = &fixture["http_error_retry"];
let cases = table["cases"].as_array().unwrap();
assert!(!cases.is_empty());
let mismatched: Vec<&str> = cases
.iter()
.filter(|case| {
let mut body = table["request"].clone();
body["thinking"] = case["thinking"].clone();
let status = u16::try_from(case["status"].as_u64().unwrap()).unwrap();
let error_text = case["error_text"].as_str().unwrap();
let retried = should_retry_anthropic_messages_on_http_error(
status,
error_text,
&body,
&MessagesTransformContext::default(),
);
let stripped = case["stripped"].as_bool().unwrap();
let transformed = transform_anthropic_messages_request_on_http_error(
status,
error_text,
body.clone(),
);
let expected = if stripped {
RetriedRequest {
reasons: vec![RetryReason::ThinkingSignature],
body: table["stripped_request"].clone(),
}
} else {
RetriedRequest {
reasons: Vec::new(),
body,
}
};
retried != case["retried"].as_bool().unwrap() || transformed != expected
})
.map(|case| case["name"].as_str().unwrap())
.collect();
assert_eq!(mismatched, Vec::<&str>::new());
}
#[test]
fn the_strict_tools_retry_matches_the_cases_shared_with_python() {
let fixture: Value = serde_json::from_str(include_str!(concat!(
env!("CARGO_MANIFEST_DIR"),
"/tests/fixtures/anthropic_messages_http_error_retry.json"
)))
.unwrap();
let replay = &fixture["http_error_retry"];
let table = &fixture["strict_tools_retry"];
let cases = table["cases"].as_array().unwrap();
assert!(!cases.is_empty());
let mismatched: Vec<&str> = cases
.iter()
.filter(|case| {
let mut body = replay["request"].clone();
body["thinking"] = case["thinking"].clone();
body["tools"] = table["tools"].clone();
let status = u16::try_from(case["status"].as_u64().unwrap()).unwrap();
let error_text = case["error_text"].as_str().unwrap();
let context = MessagesTransformContext {
drops_unsupported_params: case["drops_unsupported_params"].as_bool().unwrap(),
..MessagesTransformContext::default()
};
let mut reasons = Vec::new();
let mut expected = if case["stripped"].as_bool().unwrap() {
reasons.push(RetryReason::ThinkingSignature);
replay["stripped_request"].clone()
} else {
body.clone()
};
expected["tools"] = if case["strict_dropped"].as_bool().unwrap() {
reasons.push(RetryReason::StrictTools {
dropped: usize::try_from(table["dropped"].as_u64().unwrap()).unwrap(),
});
table["tools_without_strict"].clone()
} else {
table["tools"].clone()
};
let retried = should_retry_anthropic_messages_on_http_error(
status, error_text, &body, &context,
);
let transformed =
transform_anthropic_messages_request_on_http_error(status, error_text, body);
retried != case["retried"].as_bool().unwrap()
|| transformed
!= RetriedRequest {
reasons,
body: expected,
}
})
.map(|case| case["name"].as_str().unwrap())
.collect();
assert_eq!(mismatched, Vec::<&str>::new());
}
}

25
litellm-rust/crates/llms/src/azure_ai/messages/transformation.rs
View file

@ -1,3 +1,5 @@
use crate::anthropic::messages::transformation::retry_anthropic_messages_after_http_error;
use litellm_host::interceptors::RetriedRequest;
use litellm_http::request::{has_bearer_auth, has_header};
use litellm_llms_types::{
formats::messages::{MessagesOptionalParams, MessagesRequest},
@ -5,6 +7,7 @@ use litellm_llms_types::{
};
use litellm_router_types::LitellmParams;
use litellm_secret_value::SecretValue;
use serde_json::Value;
use crate::{
Error,
@ -21,7 +24,7 @@ use crate::{
base_llm::{
auth::{AuthScheme, Headers, ValidatedEnvironment},
messages::{
context::MessagesTransformContext,
context::{MessagesShapeOptions, MessagesTransformContext},
normalization::{normalize_system_role_messages, strip_cache_control_scope},
transformation::BaseMessagesConfig,
},
@ -36,12 +39,28 @@ pub const AZURE_ANTHROPIC_MESSAGES_CONFIG: AzureAnthropicMessagesConfig =
AzureAnthropicMessagesConfig;
impl BaseMessagesConfig for AzureAnthropicMessagesConfig {
fn tool_custom_field(&self, model: &str) -> crate::base_llm::messages::tools::ToolCustomField {
crate::base_llm::messages::tools::ToolCustomField::Rewrite {
hoist_defer_loading: crate::anthropic::messages::transformation::is_claude_model(model),
}
}
fn retry_after_http_error(
&self,
status: u16,
error_text: &str,
request_body: &Value,
context: &MessagesTransformContext,
) -> Option<RetriedRequest> {
retry_anthropic_messages_after_http_error(status, error_text, request_body, context)
}
fn shape_request(
&self,
request: MessagesRequest,
reasoning_auto_summary: bool,
options: MessagesShapeOptions,
) -> Result<MessagesRequest, Error> {
shape_anthropic_messages_request(request, reasoning_auto_summary)
shape_anthropic_messages_request(request, options)
}
fn get_complete_url(

8
litellm-rust/crates/llms/src/base_llm/messages/context.rs
View file

@ -23,6 +23,12 @@ impl SupportedEffortTiers {
}
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct MessagesShapeOptions {
pub reasoning_auto_summary: bool,
pub presanitized: bool,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct MessagesModelCapabilities {
#[serde(default)]
@ -123,6 +129,7 @@ pub struct ThinkingContext {
pub struct MessagesTransformContext {
pub thinking: ThinkingContext,
pub drop_params: bool,
pub drops_unsupported_params: bool,
}
impl MessagesTransformContext {
pub fn new(capabilities: MessagesModelCapabilities, drop_params: bool) -> Self {
@ -142,6 +149,7 @@ impl MessagesTransformContext {
keeps_display_updates: false,
},
drop_params,
drops_unsupported_params: false,
}
}

1
litellm-rust/crates/llms/src/base_llm/messages/mod.rs
View file

@ -2,4 +2,5 @@ pub mod context;
pub mod mid_conversation_system;
pub mod normalization;
pub mod streaming;
pub mod tools;
pub mod transformation;

26
litellm-rust/crates/llms/src/base_llm/messages/normalization.rs
View file

@ -1,6 +1,9 @@
use litellm_llms_types::formats::messages::{
CacheControl, ContentBlock, Message, MessageContent, MessagesOptionalParams, MessagesRequest,
SystemPrompt,
use litellm_llms_types::{
formats::messages::{
CacheControl, ContentBlock, Message, MessageContent, MessagesOptionalParams,
MessagesRequest, SystemPrompt,
},
recognized::{Recognized, map_known_runs},
};
use crate::{
@ -22,15 +25,22 @@ pub fn normalize_system_role_messages(
let leading_count = request
.messages
.iter()
.take_while(|message| is_system_role_message(message))
.take_while(|message| message.known().is_some_and(is_system_role_message))
.count();
let mut messages = request.messages.into_iter();
let hoisted: Vec<Message> = messages.by_ref().take(leading_count).collect();
let remaining: Vec<Message> = messages.collect();
let hoisted: Vec<Message> = messages
.by_ref()
.take(leading_count)
.filter_map(|message| match message {
Recognized::Known(message) => Some(message),
Recognized::Unrecognized(_) => None,
})
.collect();
let remaining: Vec<Recognized<Message>> = messages.collect();
let remaining = if supports_mid_conversation_system {
remaining
} else {
convert_mid_conversation_system_turns(remaining)
map_known_runs(remaining, convert_mid_conversation_system_turns)
};
let system = if hoisted.is_empty() {
request.params.system
@ -94,7 +104,7 @@ pub fn strip_cache_control_scope(request: MessagesRequest) -> MessagesRequest {
messages: request
.messages
.into_iter()
.map(strip_scope_from_message)
.map(|message| message.map_known(strip_scope_from_message))
.collect(),
params: MessagesOptionalParams {
system: request.params.system.map(strip_scope_from_system),

194
litellm-rust/crates/llms/src/base_llm/messages/tools.rs Normal file
View file

@ -0,0 +1,194 @@
use litellm_llms_types::{
formats::messages::{MessagesOptionalParams, MessagesRequest, MessagesTool},
recognized::Recognized,
};
use serde_json::{Map, Value};
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ToolCustomField {
Keep,
Rewrite { hoist_defer_loading: bool },
}
impl ToolCustomField {
pub fn apply(self, request: MessagesRequest, drop_params: bool) -> MessagesRequest {
match self {
Self::Keep => request,
Self::Rewrite {
hoist_defer_loading,
} => tools_without_custom_field(request, hoist_defer_loading, drop_params),
}
}
}
fn tools_without_custom_field(
request: MessagesRequest,
hoist: bool,
drop_custom: bool,
) -> MessagesRequest {
let tools = request.params.tools.map(|tools| {
tools
.into_iter()
.map(|tool| tool_without_custom_field(tool, hoist, drop_custom))
.collect()
});
MessagesRequest {
params: MessagesOptionalParams {
tools,
..request.params
},
..request
}
}
fn tool_without_custom_field(
tool: Recognized<MessagesTool>,
hoist: bool,
drop_custom: bool,
) -> Recognized<MessagesTool> {
let Ok(Value::Object(fields)) = serde_json::to_value(&tool) else {
return tool;
};
let Some(custom) = fields.get("custom").cloned() else {
return tool;
};
let defer_loading = hoist
.then(|| hoistable_defer_loading(&fields, &custom))
.flatten();
let without_custom = fields
.iter()
.filter(|(key, _)| *key != "custom")
.map(|(key, value)| (key.clone(), value.clone()));
let rebuilt: Map<String, Value> = match defer_loading {
None if drop_custom => without_custom.collect(),
None => return tool,
Some(defer_loading) => {
let unhoisted: Map<String, Value> = custom
.as_object()
.into_iter()
.flatten()
.filter(|(key, _)| *key != "defer_loading")
.map(|(key, value)| (key.clone(), value.clone()))
.collect();
let kept_custom = (!unhoisted.is_empty() && !drop_custom)
.then(|| ("custom".to_string(), Value::Object(unhoisted)));
without_custom
.chain([("defer_loading".to_string(), Value::Bool(defer_loading))])
.chain(kept_custom)
.collect()
}
};
serde_json::from_value(Value::Object(rebuilt)).unwrap_or(tool)
}
fn hoistable_defer_loading(tool: &Map<String, Value>, custom: &Value) -> Option<bool> {
if tool.contains_key("defer_loading") {
return None;
}
custom.as_object()?.get("defer_loading")?.as_bool()
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use serde_json::json;
use super::*;
const WEATHER: &str = "get_weather";
fn tool(fields: Value) -> Value {
merged(
json!({"name": WEATHER, "input_schema": {"type": "object"}}),
fields,
)
}
fn tools_after(policy: ToolCustomField, tools: Value, drop_custom: bool) -> Value {
let request: MessagesRequest = serde_json::from_value(json!({
"model": "claude-sonnet-4-5",
"max_tokens": 16,
"messages": [{"role": "user", "content": "hi"}],
"tools": [tools]
}))
.unwrap();
serde_json::to_value(policy.apply(request, drop_custom)).unwrap()["tools"][0].clone()
}
fn merged(base: Value, fields: Value) -> Value {
let mut merged = base.as_object().unwrap().clone();
merged.extend(fields.as_object().unwrap().clone());
Value::Object(merged)
}
const HOIST: ToolCustomField = ToolCustomField::Rewrite {
hoist_defer_loading: true,
};
const NO_HOIST: ToolCustomField = ToolCustomField::Rewrite {
hoist_defer_loading: false,
};
#[rstest]
#[case::hoists_defer_loading(
HOIST,
json!({"custom": {"defer_loading": true}}),
false,
json!({"defer_loading": true})
)]
#[case::keeps_other_custom_keys(
HOIST,
json!({"custom": {"defer_loading": false, "note": "kept"}}),
false,
json!({"defer_loading": false, "custom": {"note": "kept"}})
)]
#[case::drops_other_custom_keys_with_drop_params(
HOIST,
json!({"custom": {"defer_loading": true, "note": "dropped"}}),
true,
json!({"defer_loading": true})
)]
#[case::keeps_an_opaque_custom(
HOIST,
json!({"custom": "opaque"}),
false,
json!({"custom": "opaque"})
)]
#[case::drops_an_opaque_custom_with_drop_params(
HOIST,
json!({"custom": "opaque"}),
true,
json!({})
)]
#[case::own_defer_loading_wins(
HOIST,
json!({"defer_loading": false, "custom": {"defer_loading": true}}),
false,
json!({"defer_loading": false, "custom": {"defer_loading": true}})
)]
#[case::not_hoisted_off_claude_hosts(
NO_HOIST,
json!({"custom": {"defer_loading": true}}),
false,
json!({"custom": {"defer_loading": true}})
)]
fn tool_custom_field_follows_the_python_handler(
#[case] policy: ToolCustomField,
#[case] fields: Value,
#[case] drop_custom: bool,
#[case] expected: Value,
) {
assert_eq!(
tools_after(policy, tool(fields), drop_custom),
tool(expected)
);
}
#[rstest]
fn keep_leaves_the_custom_field() {
let custom = tool(json!({"custom": {"defer_loading": true}}));
assert_eq!(
tools_after(ToolCustomField::Keep, custom.clone(), true),
custom
);
}
}

71
litellm-rust/crates/llms/src/base_llm/messages/transformation.rs
View file

@ -1,19 +1,33 @@
use litellm_host::interceptors::RetriedRequest;
use litellm_llms_types::formats::messages::{
MessagesOptionalParams, MessagesRequest, MessagesResponse,
};
use litellm_router_types::LitellmParams;
use litellm_secret_value::SecretValue;
use serde_json::Value;
use super::context::MessagesTransformContext;
use super::{
context::{MessagesShapeOptions, MessagesTransformContext},
tools::ToolCustomField,
};
pub use crate::base_llm::auth::{Headers, ValidatedEnvironment};
use crate::{Error, base_llm::messages::streaming::StreamDecoder};
pub enum FederatedEnvironment {
Validated(ValidatedEnvironment),
Unauthenticated {
headers: Headers,
missing: Error,
federate: fn(Headers, SecretValue) -> ValidatedEnvironment,
},
}
pub trait BaseMessagesConfig: Sync {
fn shape_request(
&self,
request: MessagesRequest,
_reasoning_auto_summary: bool,
_options: MessagesShapeOptions,
) -> Result<MessagesRequest, Error> {
Ok(request)
}
@ -35,6 +49,10 @@ pub trait BaseMessagesConfig: Sync {
Ok(request)
}
fn tool_custom_field(&self, _model: &str) -> ToolCustomField {
ToolCustomField::Keep
}
fn transform_anthropic_messages_response(
&self,
_model: &str,
@ -62,6 +80,26 @@ pub trait BaseMessagesConfig: Sync {
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<ValidatedEnvironment, Error>;
fn validate_federated_environment(
&self,
headers: Headers,
api_key: Option<&str>,
model: &str,
optional_params: &MessagesOptionalParams,
litellm_params: &LitellmParams,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<FederatedEnvironment, Error> {
self.validate_environment(
headers,
api_key,
model,
optional_params,
litellm_params,
env_lookup,
)
.map(FederatedEnvironment::Validated)
}
/// `None` relays the upstream bytes untouched, which is right for every host that already
/// speaks Anthropic SSE. A host on another wire returns the decoder that lifts its frames
/// into Anthropic stream events, and the route re-encodes those as Anthropic SSE.
@ -77,6 +115,10 @@ pub trait BaseMessagesConfig: Sync {
headers
}
fn headers_for_url(&self, headers: Headers, _url: &str) -> Headers {
headers
}
fn request_body(
&self,
request: MessagesRequest,
@ -91,6 +133,16 @@ pub trait BaseMessagesConfig: Sync {
fn wire_body(&self, body: Value) -> Value {
body
}
fn retry_after_http_error(
&self,
_status: u16,
_error_text: &str,
_request_body: &Value,
_context: &MessagesTransformContext,
) -> Option<RetriedRequest> {
None
}
}
#[cfg(test)]
@ -149,9 +201,10 @@ mod tests {
}
#[rstest]
#[case::disabled(false)]
#[case::enabled(true)]
fn default_shaping_preserves_provider_policy_inputs(#[case] reasoning_auto_summary: bool) {
fn default_shaping_preserves_provider_policy_inputs(
#[values(false, true)] reasoning_auto_summary: bool,
#[values(false, true)] presanitized: bool,
) {
let request: MessagesRequest = serde_json::from_value(serde_json::json!({
"model": "test-model",
"metadata": {"user_id": 7, "extra": "keep"},
@ -160,7 +213,13 @@ mod tests {
}))
.unwrap();
assert_eq!(
DefaultsConfig.shape_request(request.clone(), reasoning_auto_summary),
DefaultsConfig.shape_request(
request.clone(),
MessagesShapeOptions {
reasoning_auto_summary,
presanitized
}
),
Ok(request)
);
assert_eq!(

27
litellm-rust/crates/llms/src/bedrock/claude_platform/messages/transformation.rs
View file

@ -1,3 +1,5 @@
use crate::anthropic::messages::transformation::retry_anthropic_messages_after_http_error;
use litellm_host::interceptors::RetriedRequest;
use std::sync::LazyLock;
use litellm_auth::{AwsParams, CredentialPlacement, ParamKind, ParamSpec};
@ -21,7 +23,10 @@ use crate::{
},
base_llm::{
auth::{AuthScheme, Headers, ValidatedEnvironment},
messages::{context::MessagesTransformContext, transformation::BaseMessagesConfig},
messages::{
context::{MessagesShapeOptions, MessagesTransformContext},
transformation::BaseMessagesConfig,
},
},
bedrock::common_utils::validated_region,
};
@ -50,12 +55,28 @@ pub const BEDROCK_CLAUDE_PLATFORM_MESSAGES_CONFIG: BedrockClaudePlatformMessages
BedrockClaudePlatformMessagesConfig;
impl BaseMessagesConfig for BedrockClaudePlatformMessagesConfig {
fn tool_custom_field(&self, model: &str) -> crate::base_llm::messages::tools::ToolCustomField {
crate::base_llm::messages::tools::ToolCustomField::Rewrite {
hoist_defer_loading: crate::anthropic::messages::transformation::is_claude_model(model),
}
}
fn retry_after_http_error(
&self,
status: u16,
error_text: &str,
request_body: &Value,
context: &MessagesTransformContext,
) -> Option<RetriedRequest> {
retry_anthropic_messages_after_http_error(status, error_text, request_body, context)
}
fn shape_request(
&self,
request: MessagesRequest,
reasoning_auto_summary: bool,
options: MessagesShapeOptions,
) -> Result<MessagesRequest, Error> {
shape_anthropic_messages_request(request, reasoning_auto_summary)
shape_anthropic_messages_request(request, options)
}
fn validate_environment(

18
litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude_transformation.rs
View file

@ -1,9 +1,11 @@
use crate::anthropic::messages::transformation::retry_anthropic_messages_after_http_error;
use litellm_core_utils::decline::Decline;
use litellm_host::interceptors::RetriedRequest;
use std::convert::Infallible;
use crate::{
anthropic::messages::handler::shape_anthropic_messages_request,
base_llm::messages::context::MessagesTransformContext,
base_llm::messages::context::{MessagesShapeOptions, MessagesTransformContext},
};
use futures_util::StreamExt;
use litellm_auth::AwsParams;
@ -103,12 +105,22 @@ fn invoke_url(
}
impl BaseMessagesConfig for AmazonAnthropicClaudeMessagesConfig {
fn retry_after_http_error(
&self,
status: u16,
error_text: &str,
request_body: &Value,
context: &MessagesTransformContext,
) -> Option<RetriedRequest> {
retry_anthropic_messages_after_http_error(status, error_text, request_body, context)
}
fn shape_request(
&self,
request: MessagesRequest,
reasoning_auto_summary: bool,
options: MessagesShapeOptions,
) -> Result<MessagesRequest, Error> {
shape_anthropic_messages_request(request, reasoning_auto_summary)
shape_anthropic_messages_request(request, options)
}
fn get_complete_url(

31
litellm-rust/crates/llms/src/bedrock_mantle/messages/transformation.rs
View file

@ -1,3 +1,5 @@
use crate::anthropic::messages::transformation::retry_anthropic_messages_after_http_error;
use litellm_host::interceptors::RetriedRequest;
use std::sync::LazyLock;
use litellm_auth::{AwsParams, CredentialPlacement};
@ -17,6 +19,7 @@ use serde_json::Value;
use crate::{
Error,
anthropic::{
beta_headers::accepted_or_unlisted,
common_utils::{existing_betas, merge_beta_headers},
messages::{
handler::shape_anthropic_messages_request,
@ -26,7 +29,7 @@ use crate::{
base_llm::{
auth::AuthScheme,
messages::{
context::MessagesTransformContext,
context::{MessagesShapeOptions, MessagesTransformContext},
transformation::{BaseMessagesConfig, Headers, ValidatedEnvironment},
},
},
@ -56,6 +59,16 @@ pub const BEDROCK_MANTLE_MESSAGES_CONFIG: BedrockMantleAnthropicMessagesConfig =
BedrockMantleAnthropicMessagesConfig;
impl BaseMessagesConfig for BedrockMantleAnthropicMessagesConfig {
fn retry_after_http_error(
&self,
status: u16,
error_text: &str,
request_body: &Value,
context: &MessagesTransformContext,
) -> Option<RetriedRequest> {
retry_anthropic_messages_after_http_error(status, error_text, request_body, context)
}
fn validate_environment(
&self,
headers: Headers,
@ -129,9 +142,9 @@ impl BaseMessagesConfig for BedrockMantleAnthropicMessagesConfig {
fn shape_request(
&self,
request: MessagesRequest,
reasoning_auto_summary: bool,
options: MessagesShapeOptions,
) -> Result<MessagesRequest, Error> {
shape_anthropic_messages_request(request, reasoning_auto_summary)
shape_anthropic_messages_request(request, options)
}
fn wire_body(&self, body: Value) -> Value {
@ -191,14 +204,10 @@ impl BaseMessagesConfig for BedrockMantleAnthropicMessagesConfig {
.filter_map(Value::as_str)
.filter_map(|beta| beta.parse::<AnthropicBeta>().ok())
.collect();
let caller: BetaSet = existing_betas(&headers)
.union(body_betas)
.into_iter()
.filter_map(|beta| match beta {
AnthropicBeta::Other(_) => Some(beta),
_ => beta.on(BetaProvider::BedrockMantle),
})
.collect();
let caller = accepted_or_unlisted(
existing_betas(&headers).union(body_betas),
BetaProvider::BedrockMantle,
);
let safeguards: BetaSet = request
.params
.extra

27
litellm-rust/crates/llms/src/deepseek/messages/transformation.rs
View file

@ -1,4 +1,6 @@
use crate::anthropic::messages::transformation::retry_anthropic_messages_after_http_error;
use litellm_auth::CredentialPlacement;
use litellm_host::interceptors::RetriedRequest;
use litellm_secret_value::SecretValue;
use litellm_core_utils::settings::resolve_non_empty;
@ -21,7 +23,10 @@ use crate::{
},
base_llm::{
auth::{AuthScheme, Headers, ValidatedEnvironment},
messages::{context::MessagesTransformContext, transformation::BaseMessagesConfig},
messages::{
context::{MessagesShapeOptions, MessagesTransformContext},
transformation::BaseMessagesConfig,
},
},
};
@ -42,12 +47,28 @@ pub const DEEPSEEK_ANTHROPIC_MESSAGES_CONFIG: DeepSeekAnthropicMessagesConfig =
DeepSeekAnthropicMessagesConfig;
impl BaseMessagesConfig for DeepSeekAnthropicMessagesConfig {
fn tool_custom_field(&self, _model: &str) -> crate::base_llm::messages::tools::ToolCustomField {
crate::base_llm::messages::tools::ToolCustomField::Rewrite {
hoist_defer_loading: false,
}
}
fn retry_after_http_error(
&self,
status: u16,
error_text: &str,
request_body: &Value,
context: &MessagesTransformContext,
) -> Option<RetriedRequest> {
retry_anthropic_messages_after_http_error(status, error_text, request_body, context)
}
fn shape_request(
&self,
request: MessagesRequest,
reasoning_auto_summary: bool,
options: MessagesShapeOptions,
) -> Result<MessagesRequest, Error> {
shape_anthropic_messages_request(request, reasoning_auto_summary)
shape_anthropic_messages_request(request, options)
}
fn get_complete_url(

18
litellm-rust/crates/llms/src/edenai/messages/transformation.rs
View file

@ -1,4 +1,6 @@
use crate::anthropic::messages::transformation::retry_anthropic_messages_after_http_error;
use litellm_auth::CredentialPlacement;
use litellm_host::interceptors::RetriedRequest;
use litellm_secret_value::SecretValue;
use litellm_core_utils::settings::resolve_non_empty;
@ -33,6 +35,22 @@ pub struct EdenAIAnthropicMessagesConfig;
pub const EDENAI_MESSAGES_CONFIG: EdenAIAnthropicMessagesConfig = EdenAIAnthropicMessagesConfig;
impl BaseMessagesConfig for EdenAIAnthropicMessagesConfig {
fn tool_custom_field(&self, _model: &str) -> crate::base_llm::messages::tools::ToolCustomField {
crate::base_llm::messages::tools::ToolCustomField::Rewrite {
hoist_defer_loading: false,
}
}
fn retry_after_http_error(
&self,
status: u16,
error_text: &str,
request_body: &Value,
context: &MessagesTransformContext,
) -> Option<RetriedRequest> {
retry_anthropic_messages_after_http_error(status, error_text, request_body, context)
}
fn validate_environment(
&self,
headers: Headers,

12
litellm-rust/crates/llms/src/error.rs
View file

@ -168,6 +168,10 @@ pub enum ErrorDetail {
actual: String,
choices: Vec<&'static str>,
},
#[error(
"Invalid reasoning_effort: {actual}. Must be one of: 'minimal', 'low', 'medium', 'high', 'xhigh', 'max', 'none'"
)]
InvalidReasoningEffort { actual: String },
#[error("{field}='{value}' is not supported by this model. Got model: {model}")]
UnsupportedValue {
field: &'static str,
@ -187,6 +191,14 @@ pub enum ErrorDetail {
"thinking.budget_tokens must be a number, got {actual}. To drop it, set `litellm.drop_params=True`, or `drop_params: true` under `litellm_settings` on the proxy"
)]
NonNumericThinkingBudget { actual: &'static str },
/// Python's comparison of two operands it cannot order, which raises `TypeError` with this
/// text before any request is sent.
#[error("'{op}' not supported between instances of '{left}' and '{right}'")]
UnorderableTypes {
op: &'static str,
left: &'static str,
right: &'static str,
},
#[error(
"Missing API Base - A call is being made to LLM Provider but no api base is set either in the environment variables ({{LLM_PROVIDER}}_API_KEY) or via params"
)]

28
litellm-rust/crates/llms/src/minimax/messages/transformation.rs
View file

@ -1,5 +1,8 @@
use crate::anthropic::messages::transformation::retry_anthropic_messages_after_http_error;
use litellm_auth::CredentialPlacement;
use litellm_host::interceptors::RetriedRequest;
use litellm_secret_value::SecretValue;
use serde_json::Value;
use litellm_core_utils::settings::resolve_non_empty;
use litellm_llms_types::{
@ -19,7 +22,10 @@ use crate::{
},
base_llm::{
auth::{AuthScheme, Headers, ValidatedEnvironment},
messages::{context::MessagesTransformContext, transformation::BaseMessagesConfig},
messages::{
context::{MessagesShapeOptions, MessagesTransformContext},
transformation::BaseMessagesConfig,
},
},
};
@ -31,12 +37,28 @@ pub struct MinimaxMessagesConfig;
pub const MINIMAX_MESSAGES_CONFIG: MinimaxMessagesConfig = MinimaxMessagesConfig;
impl BaseMessagesConfig for MinimaxMessagesConfig {
fn tool_custom_field(&self, _model: &str) -> crate::base_llm::messages::tools::ToolCustomField {
crate::base_llm::messages::tools::ToolCustomField::Rewrite {
hoist_defer_loading: false,
}
}
fn retry_after_http_error(
&self,
status: u16,
error_text: &str,
request_body: &Value,
context: &MessagesTransformContext,
) -> Option<RetriedRequest> {
retry_anthropic_messages_after_http_error(status, error_text, request_body, context)
}
fn shape_request(
&self,
request: MessagesRequest,
reasoning_auto_summary: bool,
options: MessagesShapeOptions,
) -> Result<MessagesRequest, Error> {
shape_anthropic_messages_request(request, reasoning_auto_summary)
shape_anthropic_messages_request(request, options)
}
fn validate_environment(

13
litellm-rust/crates/llms/src/openrouter/messages/transformation.rs
View file

@ -1,5 +1,8 @@
use crate::anthropic::messages::transformation::retry_anthropic_messages_after_http_error;
use litellm_auth::CredentialPlacement;
use litellm_host::interceptors::RetriedRequest;
use litellm_secret_value::SecretValue;
use serde_json::Value;
use litellm_core_utils::settings::resolve_non_empty;
use litellm_http::request::{has_bearer_auth, without_headers};
@ -31,6 +34,16 @@ pub const OPENROUTER_MESSAGES_CONFIG: OpenRouterAnthropicMessagesConfig =
OpenRouterAnthropicMessagesConfig;
impl BaseMessagesConfig for OpenRouterAnthropicMessagesConfig {
fn retry_after_http_error(
&self,
status: u16,
error_text: &str,
request_body: &Value,
context: &MessagesTransformContext,
) -> Option<RetriedRequest> {
retry_anthropic_messages_after_http_error(status, error_text, request_body, context)
}
fn validate_environment(
&self,
headers: Headers,

28
litellm-rust/crates/llms/src/tencent/messages/transformation.rs
View file

@ -1,5 +1,8 @@
use crate::anthropic::messages::transformation::retry_anthropic_messages_after_http_error;
use litellm_auth::CredentialPlacement;
use litellm_host::interceptors::RetriedRequest;
use litellm_secret_value::SecretValue;
use serde_json::Value;
use litellm_core_utils::settings::resolve_non_empty;
use litellm_llms_types::{
@ -19,7 +22,10 @@ use crate::{
},
base_llm::{
auth::{AuthScheme, Headers, ValidatedEnvironment},
messages::{context::MessagesTransformContext, transformation::BaseMessagesConfig},
messages::{
context::{MessagesShapeOptions, MessagesTransformContext},
transformation::BaseMessagesConfig,
},
},
};
@ -32,12 +38,28 @@ pub struct TencentAnthropicMessagesConfig;
pub const TENCENT_MESSAGES_CONFIG: TencentAnthropicMessagesConfig = TencentAnthropicMessagesConfig;
impl BaseMessagesConfig for TencentAnthropicMessagesConfig {
fn tool_custom_field(&self, _model: &str) -> crate::base_llm::messages::tools::ToolCustomField {
crate::base_llm::messages::tools::ToolCustomField::Rewrite {
hoist_defer_loading: false,
}
}
fn retry_after_http_error(
&self,
status: u16,
error_text: &str,
request_body: &Value,
context: &MessagesTransformContext,
) -> Option<RetriedRequest> {
retry_anthropic_messages_after_http_error(status, error_text, request_body, context)
}
fn shape_request(
&self,
request: MessagesRequest,
reasoning_auto_summary: bool,
options: MessagesShapeOptions,
) -> Result<MessagesRequest, Error> {
shape_anthropic_messages_request(request, reasoning_auto_summary)
shape_anthropic_messages_request(request, options)
}
fn validate_environment(

26
litellm-rust/crates/llms/src/vertex_ai/messages/transformation.rs
View file

@ -1,8 +1,10 @@
use crate::anthropic::messages::transformation::retry_anthropic_messages_after_http_error;
use litellm_auth::VertexParams;
use litellm_auth_gcp::{
VertexConfig, get_vertex_ai_location, get_vertex_ai_project,
get_vertex_ai_project_from_credentials, secret_names,
};
use litellm_host::interceptors::RetriedRequest;
use litellm_http::request::{has_header, header_value, without_headers};
use litellm_llms_types::{
formats::messages::{MessagesOptionalParams, MessagesRequest, MessagesTool, OutputConfig},
@ -24,7 +26,7 @@ use crate::{
base_llm::{
auth::{AuthScheme, Headers, ValidatedEnvironment},
messages::{
context::MessagesTransformContext,
context::{MessagesShapeOptions, MessagesTransformContext},
normalization::{normalize_system_role_messages, strip_cache_control_scope},
transformation::BaseMessagesConfig,
},
@ -47,12 +49,28 @@ pub const VERTEX_ANTHROPIC_MESSAGES_CONFIG: VertexAiPartnerModelsAnthropicMessag
VertexAiPartnerModelsAnthropicMessagesConfig;
impl BaseMessagesConfig for VertexAiPartnerModelsAnthropicMessagesConfig {
fn tool_custom_field(&self, model: &str) -> crate::base_llm::messages::tools::ToolCustomField {
crate::base_llm::messages::tools::ToolCustomField::Rewrite {
hoist_defer_loading: crate::anthropic::messages::transformation::is_claude_model(model),
}
}
fn retry_after_http_error(
&self,
status: u16,
error_text: &str,
request_body: &Value,
context: &MessagesTransformContext,
) -> Option<RetriedRequest> {
retry_anthropic_messages_after_http_error(status, error_text, request_body, context)
}
fn shape_request(
&self,
request: MessagesRequest,
reasoning_auto_summary: bool,
options: MessagesShapeOptions,
) -> Result<MessagesRequest, Error> {
shape_anthropic_messages_request(request, reasoning_auto_summary)
shape_anthropic_messages_request(request, options)
}
fn get_complete_url(
@ -700,6 +718,7 @@ mod tests {
..ThinkingContext::default()
},
drop_params: false,
drops_unsupported_params: false,
},
);
assert_eq!(body["system"], json!("be terse"));
@ -717,6 +736,7 @@ mod tests {
..ThinkingContext::default()
},
drop_params: false,
drops_unsupported_params: false,
}
}

623
litellm-rust/crates/llms/tests/fixtures/anthropic_messages_http_error_retry.json vendored Normal file
View file

@ -0,0 +1,623 @@
{
"invalid_thinking_block_error": [
{
"name": "invalid_signature",
"error_text": "messages.1.content.0: Invalid `signature` in `thinking` block",
"expected": true
},
{
"name": "invalid_signature_in_an_error_envelope",
"error_text": "{\"type\":\"error\",\"error\":{\"type\":\"invalid_request_error\",\"message\":\"messages.3.content.3: Invalid `signature` in `thinking` block\"},\"request_id\":\"req_011Ca2EtQDxp7x6RGUY2jVn9\"}",
"expected": true
},
{
"name": "signature_not_a_string",
"error_text": "{\"message\":\"messages.2.content.0.thinking.signature.str: Input should be a valid string\"}",
"expected": true
},
{
"name": "signature_not_a_string_bare",
"error_text": "messages.4.content.1.thinking.signature.str: Input should be a valid string",
"expected": true
},
{
"name": "empty_thinking",
"error_text": "messages.1.content.0.thinking: each thinking block must contain thinking",
"expected": true
},
{
"name": "bound_to_a_different_conversation",
"error_text": "messages.1.content.0: Invalid `signature` in `thinking` block. The block is bound to a different conversation.",
"expected": true
},
{
"name": "uppercase",
"error_text": "MESSAGES.1: INVALID SIGNATURE IN THINKING BLOCK",
"expected": true
},
{
"name": "signature_without_a_verdict",
"error_text": "thinking signature is malformed",
"expected": false
},
{
"name": "invalid_without_thinking",
"error_text": "Invalid `signature` in `tool_use` block",
"expected": false
},
{
"name": "model_not_found",
"error_text": "invalid_request_error: model not found",
"expected": false
},
{
"name": "rate_limited",
"error_text": "rate limit exceeded",
"expected": false
},
{
"name": "unrelated",
"error_text": "max_tokens: 4096 > 64000",
"expected": false
},
{
"name": "empty",
"error_text": "",
"expected": false
}
],
"thinking_binding_mismatch_error": [
{
"name": "bound_elsewhere",
"error_text": "messages.1.content.0: Invalid `signature` in `thinking` block. The block is bound to a different conversation.",
"expected": true
},
{
"name": "mixed_case",
"error_text": "The block is Bound To A Different Conversation.",
"expected": true
},
{
"name": "plain_signature",
"error_text": "messages.1.content.0: Invalid `signature` in `thinking` block",
"expected": false
}
],
"http_error_retry": {
"request": {
"model": "claude-sonnet-4-5",
"max_tokens": 4096,
"messages": [
{
"role": "user",
"content": "hi"
},
{
"role": "assistant",
"content": [
{
"type": "thinking",
"thinking": "plan",
"signature": "sig"
},
{
"type": "text",
"text": "hello"
}
]
},
{
"role": "assistant",
"content": [
{
"type": "redacted_thinking",
"data": "opaque"
}
]
},
{
"role": "user",
"content": "again"
}
]
},
"stripped_request": {
"model": "claude-sonnet-4-5",
"max_tokens": 4096,
"messages": [
{
"role": "user",
"content": "hi"
},
{
"role": "assistant",
"content": [
{
"type": "text",
"text": "hello"
}
]
},
{
"role": "user",
"content": "again"
}
]
},
"cases": [
{
"name": "invalid_signature",
"status": 400,
"error_text": "messages.1.content.0: Invalid `signature` in `thinking` block",
"thinking": {
"type": "enabled",
"budget_tokens": 2048
},
"retried": true,
"stripped": true
},
{
"name": "missing_signature",
"status": 400,
"error_text": "messages.1.content.0.thinking.signature.str: Input should be a valid string",
"thinking": {
"type": "enabled",
"budget_tokens": 2048
},
"retried": true,
"stripped": true
},
{
"name": "empty_thinking",
"status": 400,
"error_text": "messages.1.content.0.thinking: each thinking block must contain thinking",
"thinking": {
"type": "enabled",
"budget_tokens": 2048
},
"retried": true,
"stripped": true
},
{
"name": "unpinned_binding_mismatch",
"status": 400,
"error_text": "messages.1.content.0: Invalid `signature` in `thinking` block. The block is bound to a different conversation.",
"thinking": {
"type": "adaptive"
},
"retried": true,
"stripped": true
},
{
"name": "caller_pinned_binding",
"status": 400,
"error_text": "messages.1.content.0: Invalid `signature` in `thinking` block. The block is bound to a different conversation.",
"thinking": {
"type": "adaptive",
"block_binding": {
"prefix_mismatch_behavior": "error"
}
},
"retried": false,
"stripped": true
},
{
"name": "caller_pinned_binding_other_signature_error",
"status": 400,
"error_text": "messages.1.content.0: Invalid `signature` in `thinking` block",
"thinking": {
"type": "adaptive",
"block_binding": {
"prefix_mismatch_behavior": "error"
}
},
"retried": true,
"stripped": true
},
{
"name": "not_a_400",
"status": 429,
"error_text": "messages.1.content.0: Invalid `signature` in `thinking` block",
"thinking": {
"type": "enabled",
"budget_tokens": 2048
},
"retried": false,
"stripped": false
},
{
"name": "unrelated_400",
"status": 400,
"error_text": "max_tokens: 4096 > 64000",
"thinking": {
"type": "enabled",
"budget_tokens": 2048
},
"retried": false,
"stripped": false
},
{
"name": "signature_without_thinking",
"status": 400,
"error_text": "Invalid `signature` in `tool_use` block",
"thinking": {
"type": "enabled",
"budget_tokens": 2048
},
"retried": false,
"stripped": false
}
]
},
"claude_strict_tools_rejection": [
{
"name": "too_many_strict_tools",
"status": 400,
"error_text": "Too many strict tools (21). The maximum number of strict tools is 20.",
"expected": true
},
{
"name": "too_many_strict_tools_in_an_error_envelope",
"status": 400,
"error_text": "{\"type\":\"error\",\"error\":{\"type\":\"invalid_request_error\",\"message\":\"Too many strict tools (24). The maximum is 20.\"}}",
"expected": true
},
{
"name": "non_ascii_digits",
"status": 400,
"error_text": "Too many strict tools (٢١)",
"expected": true
},
{
"name": "optional_parameters",
"status": 400,
"error_text": "Schemas contains too many optional parameters (30)",
"expected": true
},
{
"name": "union_types",
"status": 400,
"error_text": "Schemas contains too many parameters with union types",
"expected": true
},
{
"name": "compiled_grammar_too_large",
"status": 400,
"error_text": "The compiled grammar is too large",
"expected": true
},
{
"name": "compiled_grammar_size",
"status": 400,
"error_text": "Compiled grammar size exceeds the limit",
"expected": true
},
{
"name": "custom_strict",
"status": 400,
"error_text": "tools.3.custom.strict: Extra inputs are not permitted",
"expected": true
},
{
"name": "custom",
"status": 400,
"error_text": "tools.0.custom: Extra inputs are not permitted",
"expected": true
},
{
"name": "grammar_compilation_timed_out",
"status": 400,
"error_text": "Grammar compilation timed out",
"expected": true
},
{
"name": "no_count",
"status": 400,
"error_text": "tools.0.strict: too many strict tools",
"expected": false
},
{
"name": "empty_count",
"status": 400,
"error_text": "Too many strict tools ()",
"expected": false
},
{
"name": "lowercase_grammar_timeout",
"status": 400,
"error_text": "grammar compilation timed out",
"expected": false
},
{
"name": "custom_without_separator",
"status": 400,
"error_text": "tools.0.custom.strict",
"expected": false
},
{
"name": "uppercase_schemas",
"status": 400,
"error_text": "SCHEMAS CONTAINS TOO MANY OPTIONAL PARAMETERS",
"expected": false
},
{
"name": "not_a_400",
"status": 429,
"error_text": "Too many strict tools (21). The maximum number of strict tools is 20.",
"expected": false
},
{
"name": "empty",
"status": 400,
"error_text": "",
"expected": false
}
],
"tools_without_strict": [
{
"name": "strict_dropped_from_every_tool",
"tools": [
{
"name": "a",
"strict": true,
"input_schema": {}
},
{
"name": "b",
"strict": false
},
{
"name": "c"
}
],
"expected": {
"tools": [
{
"name": "a",
"input_schema": {}
},
{
"name": "b"
},
{
"name": "c"
}
],
"dropped_count": 1
}
},
{
"name": "only_a_true_strict_counts",
"tools": [
{
"name": "a",
"strict": 1
},
{
"name": "b",
"strict": "true"
}
],
"expected": null
},
{
"name": "no_strict",
"tools": [
{
"name": "a"
}
],
"expected": null
},
{
"name": "empty",
"tools": [],
"expected": null
},
{
"name": "non_object_tool",
"tools": [
{
"name": "a",
"strict": true
},
"b"
],
"expected": null
},
{
"name": "not_a_list",
"tools": {
"name": "a",
"strict": true
},
"expected": null
},
{
"name": "null",
"tools": null,
"expected": null
}
],
"strict_tools_retry": {
"tools": [
{
"name": "a",
"input_schema": {
"type": "object"
},
"strict": true
},
{
"name": "b",
"input_schema": {
"type": "object"
},
"strict": false
},
{
"name": "c",
"input_schema": {
"type": "object"
}
}
],
"tools_without_strict": [
{
"name": "a",
"input_schema": {
"type": "object"
}
},
{
"name": "b",
"input_schema": {
"type": "object"
}
},
{
"name": "c",
"input_schema": {
"type": "object"
}
}
],
"dropped": 1,
"cases": [
{
"name": "strict_rejection_under_drop_params",
"status": 400,
"error_text": "Too many strict tools (21). The maximum number of strict tools is 20.",
"thinking": {
"type": "enabled",
"budget_tokens": 2048
},
"drops_unsupported_params": true,
"retried": true,
"stripped": false,
"strict_dropped": true
},
{
"name": "strict_rejection_without_drop_params",
"status": 400,
"error_text": "Too many strict tools (21). The maximum number of strict tools is 20.",
"thinking": {
"type": "enabled",
"budget_tokens": 2048
},
"drops_unsupported_params": false,
"retried": false,
"stripped": false,
"strict_dropped": true
},
{
"name": "compiled_grammar_under_drop_params",
"status": 400,
"error_text": "The compiled grammar is too large",
"thinking": {
"type": "enabled",
"budget_tokens": 2048
},
"drops_unsupported_params": true,
"retried": true,
"stripped": false,
"strict_dropped": true
},
{
"name": "strict_text_on_a_429",
"status": 429,
"error_text": "Too many strict tools (21). The maximum number of strict tools is 20.",
"thinking": {
"type": "enabled",
"budget_tokens": 2048
},
"drops_unsupported_params": true,
"retried": false,
"stripped": false,
"strict_dropped": false
},
{
"name": "unrelated_400_under_drop_params",
"status": 400,
"error_text": "max_tokens: 4096 > 64000",
"thinking": {
"type": "enabled",
"budget_tokens": 2048
},
"drops_unsupported_params": true,
"retried": false,
"stripped": false,
"strict_dropped": false
},
{
"name": "thinking_rejection_keeps_strict",
"status": 400,
"error_text": "messages.1.content.0: Invalid `signature` in `thinking` block",
"thinking": {
"type": "enabled",
"budget_tokens": 2048
},
"drops_unsupported_params": true,
"retried": true,
"stripped": true,
"strict_dropped": false
},
{
"name": "thinking_rejection_without_drop_params_keeps_strict",
"status": 400,
"error_text": "messages.1.content.0: Invalid `signature` in `thinking` block",
"thinking": {
"type": "enabled",
"budget_tokens": 2048
},
"drops_unsupported_params": false,
"retried": true,
"stripped": true,
"strict_dropped": false
},
{
"name": "one_error_naming_both_without_drop_params",
"status": 400,
"error_text": "messages.1.content.0: Invalid `signature` in `thinking` block. Too many strict tools (21). The maximum number of strict tools is 20.",
"thinking": {
"type": "enabled",
"budget_tokens": 2048
},
"drops_unsupported_params": false,
"retried": true,
"stripped": true,
"strict_dropped": true
},
{
"name": "pinned_binding_mismatch_with_a_strict_rejection_under_drop_params",
"status": 400,
"error_text": "messages.1.content.0: Invalid `signature` in `thinking` block. The block is bound to a different conversation. Too many strict tools (21). The maximum number of strict tools is 20.",
"thinking": {
"type": "adaptive",
"block_binding": {
"prefix_mismatch_behavior": "error"
}
},
"drops_unsupported_params": true,
"retried": true,
"stripped": true,
"strict_dropped": true
},
{
"name": "pinned_binding_mismatch_with_a_strict_rejection_without_drop_params",
"status": 400,
"error_text": "messages.1.content.0: Invalid `signature` in `thinking` block. The block is bound to a different conversation. Too many strict tools (21). The maximum number of strict tools is 20.",
"thinking": {
"type": "adaptive",
"block_binding": {
"prefix_mismatch_behavior": "error"
}
},
"drops_unsupported_params": false,
"retried": false,
"stripped": true,
"strict_dropped": true
}
]
}
}

138
litellm-rust/crates/llms/tests/messages_http_error_retry.rs Normal file
View file

@ -0,0 +1,138 @@
use litellm_host::interceptors::{RetriedRequest, RetryReason};
use litellm_llms::{
anthropic::messages::transformation::ANTHROPIC_MESSAGES_CONFIG,
azure_ai::messages::transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG,
base_llm::messages::{context::MessagesTransformContext, transformation::BaseMessagesConfig},
bedrock::{
claude_platform::messages::transformation::BEDROCK_CLAUDE_PLATFORM_MESSAGES_CONFIG,
messages::invoke_transformations::anthropic_claude_transformation::BEDROCK_ANTHROPIC_MESSAGES_CONFIG,
},
bedrock_mantle::messages::transformation::BEDROCK_MANTLE_MESSAGES_CONFIG,
deepseek::messages::transformation::DEEPSEEK_ANTHROPIC_MESSAGES_CONFIG,
edenai::messages::transformation::EDENAI_MESSAGES_CONFIG,
minimax::messages::transformation::MINIMAX_MESSAGES_CONFIG,
openrouter::messages::transformation::OPENROUTER_MESSAGES_CONFIG,
tencent::messages::transformation::TENCENT_MESSAGES_CONFIG,
vertex_ai::messages::transformation::VERTEX_ANTHROPIC_MESSAGES_CONFIG,
};
use rstest::rstest;
use serde_json::Value;
#[rstest]
#[case::anthropic(&ANTHROPIC_MESSAGES_CONFIG)]
#[case::azure_ai(&AZURE_ANTHROPIC_MESSAGES_CONFIG)]
#[case::bedrock_invoke(&BEDROCK_ANTHROPIC_MESSAGES_CONFIG)]
#[case::bedrock_claude_platform(&BEDROCK_CLAUDE_PLATFORM_MESSAGES_CONFIG)]
#[case::bedrock_mantle(&BEDROCK_MANTLE_MESSAGES_CONFIG)]
#[case::deepseek(&DEEPSEEK_ANTHROPIC_MESSAGES_CONFIG)]
#[case::edenai(&EDENAI_MESSAGES_CONFIG)]
#[case::minimax(&MINIMAX_MESSAGES_CONFIG)]
#[case::openrouter(&OPENROUTER_MESSAGES_CONFIG)]
#[case::tencent(&TENCENT_MESSAGES_CONFIG)]
#[case::vertex_ai(&VERTEX_ANTHROPIC_MESSAGES_CONFIG)]
fn every_messages_provider_retries_like_the_python_base_config(
#[case] config: &dyn BaseMessagesConfig,
) {
let fixture: Value = serde_json::from_str(include_str!(
"fixtures/anthropic_messages_http_error_retry.json"
))
.unwrap();
let table = &fixture["http_error_retry"];
let cases = table["cases"].as_array().unwrap();
assert!(!cases.is_empty());
let mismatched: Vec<&str> = cases
.iter()
.filter(|case| {
let mut body = table["request"].clone();
body["thinking"] = case["thinking"].clone();
let status = u16::try_from(case["status"].as_u64().unwrap()).unwrap();
let expected = case["retried"].as_bool().unwrap().then(|| {
if case["stripped"].as_bool().unwrap() {
RetriedRequest {
reasons: vec![RetryReason::ThinkingSignature],
body: table["stripped_request"].clone(),
}
} else {
RetriedRequest {
reasons: Vec::new(),
body: body.clone(),
}
}
});
config.retry_after_http_error(
status,
case["error_text"].as_str().unwrap(),
&body,
&MessagesTransformContext::default(),
) != expected
})
.map(|case| case["name"].as_str().unwrap())
.collect();
assert_eq!(mismatched, Vec::<&str>::new());
}
#[rstest]
#[case::anthropic(&ANTHROPIC_MESSAGES_CONFIG)]
#[case::azure_ai(&AZURE_ANTHROPIC_MESSAGES_CONFIG)]
#[case::bedrock_invoke(&BEDROCK_ANTHROPIC_MESSAGES_CONFIG)]
#[case::bedrock_claude_platform(&BEDROCK_CLAUDE_PLATFORM_MESSAGES_CONFIG)]
#[case::bedrock_mantle(&BEDROCK_MANTLE_MESSAGES_CONFIG)]
#[case::deepseek(&DEEPSEEK_ANTHROPIC_MESSAGES_CONFIG)]
#[case::edenai(&EDENAI_MESSAGES_CONFIG)]
#[case::minimax(&MINIMAX_MESSAGES_CONFIG)]
#[case::openrouter(&OPENROUTER_MESSAGES_CONFIG)]
#[case::tencent(&TENCENT_MESSAGES_CONFIG)]
#[case::vertex_ai(&VERTEX_ANTHROPIC_MESSAGES_CONFIG)]
fn every_messages_provider_drops_rejected_strict_tools_like_the_python_base_config(
#[case] config: &dyn BaseMessagesConfig,
) {
let fixture: Value = serde_json::from_str(include_str!(
"fixtures/anthropic_messages_http_error_retry.json"
))
.unwrap();
let replay = &fixture["http_error_retry"];
let table = &fixture["strict_tools_retry"];
let cases = table["cases"].as_array().unwrap();
assert!(!cases.is_empty());
let mismatched: Vec<&str> = cases
.iter()
.filter(|case| {
let mut body = replay["request"].clone();
body["thinking"] = case["thinking"].clone();
body["tools"] = table["tools"].clone();
let status = u16::try_from(case["status"].as_u64().unwrap()).unwrap();
let expected = case["retried"].as_bool().unwrap().then(|| {
let mut reasons = Vec::new();
let mut expected = if case["stripped"].as_bool().unwrap() {
reasons.push(RetryReason::ThinkingSignature);
replay["stripped_request"].clone()
} else {
body.clone()
};
expected["tools"] = if case["strict_dropped"].as_bool().unwrap() {
reasons.push(RetryReason::StrictTools {
dropped: usize::try_from(table["dropped"].as_u64().unwrap()).unwrap(),
});
table["tools_without_strict"].clone()
} else {
table["tools"].clone()
};
RetriedRequest {
reasons,
body: expected,
}
});
config.retry_after_http_error(
status,
case["error_text"].as_str().unwrap(),
&body,
&MessagesTransformContext {
drops_unsupported_params: case["drops_unsupported_params"].as_bool().unwrap(),
..MessagesTransformContext::default()
},
) != expected
})
.map(|case| case["name"].as_str().unwrap())
.collect();
assert_eq!(mismatched, Vec::<&str>::new());
}

2
litellm-rust/crates/python-bridge/src/cache/AGENTS.md vendored
View file

@ -2,6 +2,6 @@
This folder owns how Rust inference reaches the selected Python cache. Cache algorithms, storage protocols and response-cache semantics belong to their cache crates
`mod.rs` exposes the cache boundary to routes. `selection.rs` asks `litellm.rust_bridge.host.cache` for the live configured cache and owns route admission and inference protocol composition
`mod.rs` exposes the cache boundary to routes. `selection.rs` asks `litellm.rust_bridge.host.cache` for the live configured cache and owns route admission. A route's protocol carries cache operations through any host call type that converts from `CacheCall`
`python/` delegates operations to the selected Python cache without discovering configuration. Routes use only the parent module's exports

2
litellm-rust/crates/python-bridge/src/cache/mod.rs vendored
View file

@ -2,7 +2,7 @@ mod python;
mod selection;
pub(crate) use python::{CacheCall, PythonCache};
pub(crate) use selection::{Cached, Selection, configure, ensure_cache_disabled};
pub(crate) use selection::{Selection, configure, ensure_cache_disabled};
use pyo3::{exceptions::PyRuntimeError, prelude::*};

2
litellm-rust/crates/python-bridge/src/cache/python/AGENTS.md vendored
View file

@ -1,6 +1,6 @@
# Python cache delegation
This directory lets Rust inference use a selected Python cache. `service.rs` implements the injected Rust response-cache service and yields typed cache operations. `host.rs` calls the Python cache's sync or async API and delivers the result back to Rust.
This directory lets Rust inference use a selected Python cache. `service.rs` implements the injected Rust response-cache service and yields typed cache operations. `host.rs` delivers each operation's result back to Rust. Asynchronous lookups and stream stores go through `litellm.rust_bridge.host.cache`, which reads with the `LLMCachingHandler` the `@client` wrapper uses, for the legacy function the route names, and writes with its dual cache. `configure` disables the cache for synchronous calls.
Keep the shared inference protocol wrapper and adapter selection in the parent module. Receive the configured cache and prepared arguments from the parent module. Do not discover global configuration, choose native backends, or move inference to Python. Core remains independent of Python objects and cache implementation details

77
litellm-rust/crates/python-bridge/src/cache/python/host.rs vendored
View file

@ -1,6 +1,6 @@
use litellm_cache::Error;
use litellm_host::protocol::Reply;
use litellm_host_python::{from_py, to_py};
use litellm_host_python::{json_from_py, to_py};
use pyo3::{
gc::{PyTraverseError, PyVisit},
prelude::*,
@ -18,7 +18,9 @@ enum Pending {
pub(crate) struct PythonCache {
cache: Option<Py<PyAny>>,
function: Option<Py<PyAny>>,
arguments: Option<Py<PyDict>>,
read_handler: Option<Py<PyAny>>,
pending: Option<Pending>,
asynchronous: bool,
deferred: bool,
@ -28,19 +30,27 @@ impl PythonCache {
pub fn new(asynchronous: bool) -> Self {
Self {
cache: None,
function: None,
arguments: None,
read_handler: None,
pending: None,
asynchronous,
deferred: false,
}
}
pub(in crate::cache) fn asynchronous(&self) -> bool {
self.asynchronous
}
pub(in crate::cache) fn bind(
&mut self,
cache: Bound<'_, PyAny>,
function: &Bound<'_, PyAny>,
arguments: &Bound<'_, PyDict>,
) {
self.cache = Some(cache.unbind());
self.function = Some(function.clone().unbind());
self.arguments = Some(arguments.clone().unbind());
}
@ -59,7 +69,7 @@ impl PythonCache {
self.pending = Some(Pending::Lookup(reply));
(
if self.asynchronous {
"async_get_cache"
"lookup"
} else {
"get_cache"
},
@ -77,7 +87,7 @@ impl PythonCache {
self.pending = Some(Pending::Store(reply));
(
if self.asynchronous {
"async_add_cache"
"store_stream"
} else {
"add_cache"
},
@ -86,9 +96,20 @@ impl PythonCache {
}
};
let result = match result {
Some(value) if self.asynchronous => py
.import(CACHE_HOST_MODULE)?
.getattr(method)?
.call((cache, value), Some(&arguments)),
Some(value) => cache
.bind(py)
.call_method(method, (value,), Some(&arguments)),
None if self.asynchronous => py.import(CACHE_HOST_MODULE)?.getattr(method)?.call(
(
cache,
self.function.as_ref().ok_or(HandleError::Unconfigured)?,
),
Some(&arguments),
),
None => cache.bind(py).call_method(method, (), Some(&arguments)),
}
.map(Bound::unbind);
@ -98,7 +119,12 @@ impl PythonCache {
self.resume(py, result)
}
pub fn store_response(&mut self, py: Python<'_>, response: &Bound<'_, PyAny>) -> PyResult<()> {
pub fn store_response(
&mut self,
py: Python<'_>,
response: &Bound<'_, PyAny>,
verbatim: bool,
) -> PyResult<()> {
if !std::mem::take(&mut self.deferred) {
return Ok(());
}
@ -106,11 +132,30 @@ impl PythonCache {
return Err(HandleError::Unconfigured.into());
};
py.import(CACHE_HOST_MODULE)?
.getattr("store_response")?
.getattr(if verbatim {
"store_verbatim"
} else {
"store_response"
})?
.call((cache, response), Some(&arguments.bind(py).copy()?))?;
Ok(())
}
pub fn wrap_stream(&self, py: Python<'_>, stream: Py<PyAny>) -> PyResult<Py<PyAny>> {
let (Some(cache), Some(function), Some(arguments)) =
(&self.cache, &self.function, &self.arguments)
else {
return Ok(stream);
};
py.import(CACHE_HOST_MODULE)?
.getattr("wrap_stream")?
.call(
(cache, function, stream, &self.read_handler),
Some(&arguments.bind(py).copy()?),
)
.map(Bound::unbind)
}
pub fn resume(
&mut self,
py: Python<'_>,
@ -125,11 +170,14 @@ impl PythonCache {
};
match self.pending.take() {
Some(Pending::Lookup(reply)) => {
if self.asynchronous {
self.read_handler = self.caching_handler(py).ok().flatten();
}
let value = result.map_err(Error::unavailable).and_then(|value| {
if value.bind(py).is_none() {
Ok(None)
} else {
from_py(value.bind(py))
json_from_py(value.bind(py))
.map(Some)
.map_err(Error::invalid_entry)
}
@ -144,15 +192,30 @@ impl PythonCache {
Ok(None)
}
fn caching_handler(&self, py: Python<'_>) -> PyResult<Option<Py<PyAny>>> {
let Some(arguments) = self.arguments.as_ref() else {
return Ok(None);
};
let handler = py
.import(CACHE_HOST_MODULE)?
.getattr("caching_handler")?
.call((), Some(arguments.bind(py)))?;
Ok((!handler.is_none()).then(|| handler.unbind()))
}
pub fn close(&mut self) {
self.pending = None;
self.read_handler = None;
self.deferred = false;
self.cache = None;
self.function = None;
self.arguments = None;
}
pub fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
visit.call(&self.cache)?;
visit.call(&self.arguments)
visit.call(&self.function)?;
visit.call(&self.arguments)?;
visit.call(&self.read_handler)
}
}

23
litellm-rust/crates/python-bridge/src/cache/python/service.rs vendored
View file

@ -55,11 +55,12 @@ struct PythonCacheService<P: Protocol> {
config: ResponseCacheConfig,
}
pub(in crate::cache) fn service<P: Protocol<HostCall = CacheCall>>(
pub(in crate::cache) fn service<P: Protocol>(
services: HostServices<P>,
namespace: String,
) -> std::sync::Arc<dyn ResponseCacheService>
where
P::HostCall: From<CacheCall>,
P::Error: From<MachineFault> + std::error::Error,
{
std::sync::Arc::new(PythonCacheService {
@ -71,14 +72,19 @@ where
})
}
impl<P: Protocol<HostCall = CacheCall>> ResponseCacheService for PythonCacheService<P>
impl<P: Protocol> ResponseCacheService for PythonCacheService<P>
where
P::HostCall: From<CacheCall>,
P::Error: From<MachineFault> + std::error::Error,
{
fn config(&self) -> &ResponseCacheConfig {
&self.config
}
fn keyed_by_caller(&self) -> bool {
true
}
fn lookup<'a>(
&'a self,
_: &'a ResponseCacheRequest,
@ -86,7 +92,7 @@ where
) -> Pin<Box<dyn Future<Output = Result<Option<Value>, Error>> + Send + 'a>> {
Box::pin(async move {
self.services
.call(|reply| CacheCall::Lookup { reply })
.call(|reply| CacheCall::Lookup { reply }.into())
.await
.map_err(Error::unavailable)??
.map(from_python)
@ -104,10 +110,13 @@ where
let value = to_python(value)?;
let unary = value.get(STREAM_EVENTS_KEY).is_none();
self.services
.call(|reply| CacheCall::Store {
value,
unary,
reply,
.call(|reply| {
CacheCall::Store {
value,
unary,
reply,
}
.into()
})
.await
.map_err(Error::unavailable)?

46
litellm-rust/crates/python-bridge/src/cache/selection.rs vendored
View file

@ -6,17 +6,6 @@ use litellm_host::{
};
use pyo3::{prelude::*, types::PyDict};
pub(crate) struct Cached<P>(std::marker::PhantomData<P>);
impl<P: Protocol> Protocol for Cached<P> {
type Request = (P::Request, Selection);
type Response = P::Response;
type Error = P::Error;
type HostCall = python::CacheCall;
type Chunk = P::Chunk;
type StreamHead = P::StreamHead;
}
enum Backend {
Disabled,
Python { namespace: String },
@ -28,11 +17,12 @@ pub(crate) struct Selection {
}
impl Selection {
pub(crate) fn attach<P: Protocol<HostCall = python::CacheCall>>(
pub(crate) fn attach<P: Protocol>(
self,
services: HostServices<P>,
) -> (Option<ScopedCache>, CacheOptions)
where
P::HostCall: From<python::CacheCall>,
P::Error: From<MachineFault> + std::error::Error,
{
let service = match self.backend {
@ -51,15 +41,21 @@ fn selected_cache<'py>(
kwargs: &Bound<'py, PyDict>,
call_type: &str,
) -> PyResult<Option<Bound<'py, PyAny>>> {
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<Option<Bound<'py, PyAny>>> {
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")?;
@ -90,9 +86,14 @@ pub(crate) fn configure(
python: &mut python::PythonCache,
py: Python<'_>,
arguments: &Bound<'_, PyDict>,
call_type: &str,
function: &Bound<'_, PyAny>,
) -> PyResult<Selection> {
let selected = selected_cache(py, arguments, call_type)?;
let call_type: String = function.getattr("__name__")?.extract()?;
let selected = if python.asynchronous() {
configured_cache(py, &call_type)?
} else {
None
};
let Some(cache) = selected else {
return Ok(Selection {
backend: Backend::Disabled,
@ -102,6 +103,9 @@ pub(crate) fn configure(
let enabled = cache
.call_method("should_use_cache", (), Some(arguments))?
.extract::<bool>()?;
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());
@ -117,7 +121,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()
},
@ -129,7 +133,7 @@ pub(crate) fn configure(
.map(|value| value.extract())
.transpose()?
.unwrap_or_default();
python.bind(cache, arguments);
python.bind(cache, function, arguments);
Ok(Selection {
backend: if enabled {
Backend::Python { namespace }

117
litellm-rust/crates/python-bridge/src/errors.rs
View file

@ -3,7 +3,7 @@ use std::error::Error as StdError;
use litellm_error::Report;
use litellm_inference::Error;
use pyo3::{
exceptions::{PyRuntimeError, PyValueError},
exceptions::{PyRuntimeError, PyTypeError, PyValueError},
prelude::*,
};
use reqwest::header::HeaderMap;
@ -36,6 +36,37 @@ const MESSAGES_MISSING_CREDENTIAL: &str = "messages_missing_credential";
/// Python refuses the request only once its logger knows the resolved model and provider (the
/// provider config's request transform), so the failure logs them.
const MESSAGES_LATE_REJECTION: &str = "messages_late_rejection";
/// Python's request transform raises a builtin exception (`TypeError`) that its handler does not
/// map, so the caller gets the native exception as it is.
const MESSAGES_PYTHON_RAISED: &str = "messages_python_raised";
/// The text of a refusal Python's provider config raises from its request transform, which its
/// HTTP handler maps as a provider error.
const MESSAGES_PROVIDER_REJECTION: &str = "messages_provider_rejection";
fn provider_rejection(rejection: &(dyn StdError + 'static)) -> Option<String> {
match rejection.downcast_ref::<litellm_llms::Error>() {
Some(litellm_llms::Error::MissingField("max_tokens")) => {
return Some("max_tokens is required for Anthropic /v1/messages API".to_owned());
}
Some(_) => return None,
None => {}
}
match rejection.downcast_ref::<litellm_llms::ErrorDetail>() {
Some(
detail @ litellm_llms::ErrorDetail::InvalidChoice {
field: "reasoning effort",
..
},
) => Some(format!("litellm.BadRequestError: {detail}")),
Some(
detail @ (litellm_llms::ErrorDetail::InvalidReasoningEffort { .. }
| litellm_llms::ErrorDetail::UnsupportedValue {
field: "effort", ..
}),
) => Some(detail.to_string()),
_ => None,
}
}
pub(crate) fn upstream_error(
py: Python<'_>,
@ -90,6 +121,9 @@ pub(crate) fn route_failure(
if let Some(raised) = crate::secrets::python_error(py, &error) {
return Ok(raised);
}
if let Some(raised) = crate::routes::messages::workload_identity::python_error(py, &error) {
return Ok(raised);
}
if let Some((status, body)) = error.upstream_response() {
let headers = error
.upstream_headers()
@ -104,6 +138,18 @@ pub(crate) fn route_failure(
return public(upstream_error(py, status, body, &headers)?);
}
let native = match (marker, error.rejected_request()) {
(RequestMarker::Messages, Some(rejection))
if matches!(
rejection.downcast_ref::<litellm_llms::ErrorDetail>(),
Some(litellm_llms::ErrorDetail::UnorderableTypes { .. })
) =>
{
let native = PyTypeError::new_err(rejection.to_string());
let value = native.value(py);
value.setattr(MESSAGES_LATE_REJECTION, true)?;
value.setattr(MESSAGES_PYTHON_RAISED, true)?;
native
}
(RequestMarker::Messages, Some(rejection)) => {
let native = PyValueError::new_err(Report::new(rejection).to_string());
let value = native.value(py);
@ -123,6 +169,9 @@ pub(crate) fn route_failure(
) {
value.setattr(MESSAGES_LATE_REJECTION, true)?;
}
if let Some(text) = provider_rejection(rejection) {
value.setattr(MESSAGES_PROVIDER_REJECTION, text)?;
}
native
}
(RequestMarker::Messages, None) => {
@ -324,6 +373,72 @@ mod tests {
});
}
#[rstest::rstest]
#[case::missing_max_tokens(
litellm_llms::Error::MissingField("max_tokens").into(),
Some("max_tokens is required for Anthropic /v1/messages API")
)]
#[case::unmapped_effort(
Error::invalid_request(litellm_llms::ErrorDetail::InvalidChoice {
field: "reasoning effort",
actual: "'ultra'".into(),
choices: vec!["low", "none"],
}),
Some("litellm.BadRequestError: Unmapped reasoning effort: 'ultra'. Must be one of: 'low', 'none'.")
)]
#[case::other_missing_field(litellm_llms::Error::MissingField("model").into(), None)]
#[case::unsupported_param(
Error::invalid_request(litellm_llms::ErrorDetail::UnsupportedParameter {
model: "claude".into(),
param: "top_k".into(),
value: "5".into(),
hint: String::new(),
}),
None
)]
fn provider_config_rejections_carry_python_text(
#[case] error: Error,
#[case] expected: Option<&str>,
) {
Python::initialize();
Python::attach(|py| {
let native = route_failure(py, error, RequestMarker::Messages, Ok).unwrap();
let text = native
.value(py)
.getattr_opt(MESSAGES_PROVIDER_REJECTION)
.unwrap()
.map(|value| value.extract::<String>().unwrap());
assert_eq!(text.as_deref(), expected);
});
}
#[test]
fn an_unorderable_comparison_raises_python_type_error_as_it_is() {
Python::initialize();
Python::attach(|py| {
let error = Error::invalid_request(litellm_llms::ErrorDetail::UnorderableTypes {
op: "<=",
left: "str",
right: "int",
});
let native = route_failure(py, error, RequestMarker::Messages, Ok).unwrap();
assert!(native.is_instance_of::<PyTypeError>(py));
let value = native.value(py);
assert_eq!(
value.str().unwrap().to_string(),
"'<=' not supported between instances of 'str' and 'int'"
);
assert!(value.getattr_opt(MESSAGES_PYTHON_RAISED).unwrap().is_some());
assert!(
value
.getattr_opt(MESSAGES_LATE_REJECTION)
.unwrap()
.is_some()
);
assert!(value.getattr_opt(MESSAGES_REQUEST_ERROR).unwrap().is_none());
});
}
#[rstest::rstest]
#[case::rejected_request(
Error::invalid_request(litellm_llms::ErrorDetail::UnsupportedParameter {

47
litellm-rust/crates/python-bridge/src/http.rs
View file

@ -101,15 +101,22 @@ pub(crate) fn call_config(
kwargs: &Bound<'_, PyDict>,
asynchronous: bool,
) -> PyResult<HttpClientConfig> {
let settings = HttpSettings::from_layers([
Ok(resolved(py, call_settings(py, kwargs, asynchronous)?))
}
fn call_settings(
py: Python<'_>,
kwargs: &Bound<'_, PyDict>,
asynchronous: bool,
) -> PyResult<HttpSettings> {
Ok(HttpSettings::from_layers([
for_call(call_ssl_verify(kwargs)?, asynchronous),
HttpSettingsLayer::from_environment(&ProcessEnvironment),
match PythonSettings::Http.read_or_unset(py)? {
Some(snapshot) => configured(&snapshot)?,
None => HttpSettingsLayer::default(),
},
]);
Ok(resolved(py, settings))
]))
}
pub(crate) fn settings_config(
@ -132,10 +139,31 @@ pub(crate) fn settings_config(
Ok(resolved(py, settings))
}
#[pyfunction]
#[pyo3(signature = (ssl_security_level=None, ssl_ecdh_curve=None))]
pub(crate) fn unsupported_tls_settings(
ssl_security_level: Option<String>,
ssl_ecdh_curve: Option<String>,
) -> Vec<String> {
let settings = HttpSettings::from_layers([HttpSettingsLayer {
ssl_security_level,
ssl_ecdh_curve,
..HttpSettingsLayer::default()
}]);
Resolution::from(&settings)
.unsupported
.iter()
.map(ToString::to_string)
.collect()
}
fn resolution(settings: HttpSettings) -> Resolution {
Resolution::from(&settings.without_missing_files(&|path: &Path| path.exists()))
}
/// The client config `settings` resolve to, warning once per process about each setting the client cannot apply
fn resolved(py: Python<'_>, settings: HttpSettings) -> HttpClientConfig {
let resolution =
Resolution::from(&settings.without_missing_files(&|path: &Path| path.exists()));
let resolution = resolution(settings);
for unsupported in unreported(&REPORTED_UNSUPPORTED, resolution.unsupported) {
crate::logger::capture(py).scope(|| litellm_tracing::warn!("{unsupported}"));
}
@ -151,6 +179,15 @@ pub(crate) fn provider_client(
Ok(pool().client(&config, ClientVariant::Provider))
}
pub(crate) fn messages_client(
py: Python<'_>,
kwargs: &Bound<'_, PyDict>,
asynchronous: bool,
) -> PyResult<Result<Client, litellm_http::Error>> {
let config = resolution(call_settings(py, kwargs, asynchronous)?).config;
Ok(pool().client(&config, ClientVariant::Undecoded))
}
pub(crate) fn host_client(py: Python<'_>, variant: ClientVariant) -> PyResult<Client> {
let config = call_config(py, &PyDict::new(py), true)?;
pool().client(&config, variant).map_err(client_error)

Some files were not shown because too many files have changed in this diff Show more