scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-13 08:52:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 174

chore: merge main into usage cost branch

Browse source 
Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com>
This commit is contained in:
Devin AI 2026-11-10 01:44:57 +00:00
commit ecb74f6b22
242 changed files with 17016 additions and 1329 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
litellm-rust/Cargo.lock generated
View file

@ -4423,6 +4423,7 @@ dependencies = [
"pyo3-async-runtimes",
"reqwest",
"rstest",
"rustls 0.23.42",
"serde",
"serde_json",
"serde_with",

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

@ -27,6 +27,10 @@
"api_key",
"additional_args"
],
"messages_input": [
"body",
"kwargs"
],
"post_call": [
"logger",
"original_response",
@ -47,9 +51,18 @@
"defers_async_logging": [
"logger"
],
"claim_async_success_log": [
"logger"
],
"mark_async_success_enqueued": [
"logger"
],
"defer_success": [
"logger",
"pending"
"pending",
"response",
"start",
"end"
],
"sync_success_for_async_call": [
"logger",
@ -117,7 +130,8 @@
"chunks",
"start",
"end",
"first_chunk"
"first_chunk",
"stream_teardown"
],
"stream_failure": [
"logger",

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

@ -14,7 +14,7 @@ use pyo3::{
exceptions::{PyBaseException, PyException},
gc::{PyTraverseError, PyVisit},
prelude::*,
types::{PyDateTime, PyDict, PyList},
types::{PyDateTime, PyDict, PyList, PyString},
};
use serde_json::Value;
@ -22,7 +22,7 @@ use crate::{
DeploymentHooks, LegacyCallbacks, PublicCall, PythonLogger,
deferred::{PendingLogging, PendingSuccess},
finalize, is_internal_call,
python::Streaming,
python::{Logging, Streaming},
setup,
};
@ -137,6 +137,40 @@ impl LegacyLogging {
Some(params)
}
// Compared as Python compares them, so a caller's `0` is kept over a typed float's `0.0`.
fn caller_value<'py>(
&self,
py: Python<'py>,
name: &str,
sent: &Value,
) -> PyResult<Option<Bound<'py, PyAny>>> {
match self.call.lookup(py, name)? {
Some(value) if from_py::<Value>(&value).is_ok() && value.eq(to_py(py, sent)?)? => {
Ok(Some(value))
}
_ => Ok(None),
}
}
// Python's Messages handler logs the caller's own optional params, not their typed copies.
fn as_called(&self, py: Python<'_>, context: &RequestContext) -> PyResult<RequestContext> {
let mut context = context.clone();
if self.operation != LoggingOperation::Messages {
return Ok(context);
}
for (name, sent) in context
.optional_params
.as_object_mut()
.into_iter()
.flatten()
{
if let Some(value) = self.caller_value(py, name, sent)? {
*sent = from_py(&value)?;
}
}
Ok(context)
}
/// 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<()> {
@ -158,13 +192,20 @@ impl LegacyLogging {
}
}
fn input_description(&self) -> &'static str {
match self.operation {
fn logged_input<'py>(
&self,
py: Python<'py>,
body: &Bound<'py, PyDict>,
) -> PyResult<Bound<'py, PyAny>> {
let description = match self.operation {
LoggingOperation::Completion => "Chat completions",
LoggingOperation::Responses => "Responses",
LoggingOperation::Messages => "Messages",
LoggingOperation::Messages => {
return Logging::MessagesInput.call(py, (body, self.call.kwargs().bind(py)));
}
LoggingOperation::Ocr => "OCR document processing",
}
};
Ok(PyString::new(py, description).into_any())
}
/// The key `pre_call` logs. Python's Messages handler logs an empty string, never the key.
@ -251,30 +292,49 @@ impl LegacyLogging {
if !self.asynchronous {
return pending().sync(py);
}
if !self.internal
&& self
if self.internal {
return Ok(());
}
let defers = logger.defers_async_logging(py);
if logger.claim_async_success_log(py)? {
let enqueues = self
.call
.kwargs()
.bind(py)
.get_item("fallbacks")?
.is_none_or(|value| value.is_none())
{
if logger.defers_async_logging(py) {
.is_none_or(|value| value.is_none());
if defers {
let pending = Py::new(
py,
PendingLogging {
pending: Some(pending()),
enqueues,
},
)?;
logger.defer_success(py, pending.bind(py).as_any())?;
} else {
return logger.defer_success(
py,
pending.bind(py).as_any(),
&self.response,
&self.start,
&self.end,
);
}
if enqueues {
pending().asynchronous(py)?;
}
}
if defers {
return Ok(());
}
logger.sync_success_for_async_call(py, &self.response, &self.start, &self.end)
}
fn stream_success(&self, py: Python<'_>, stream: &DeliveredStream) -> PyResult<()> {
fn stream_success(
&self,
py: Python<'_>,
stream: &DeliveredStream,
teardown: bool,
) -> PyResult<()> {
let logger = self.logger()?;
let billing = self.stream_billing().ok_or(HostError::MissingState)?;
let billed = Streaming::Success.call(
@ -291,6 +351,7 @@ impl LegacyLogging {
&self.start,
&self.end,
&stream.first_chunk,
teardown,
),
);
match billed {
@ -442,6 +503,7 @@ impl LegacyLogging {
wire: Box<WireRequest>,
context: &RequestContext,
) -> PyResult<HookStep<Self, Box<WireRequest>>> {
let context = &self.as_called(py, context)?;
let logger = self.logger()?;
logger.update_from_kwargs(
py,
@ -454,12 +516,17 @@ impl LegacyLogging {
.into_bound(py)
.cast_into::<PyDict>()?;
for (name, sent) in wire.body.as_object().into_iter().flatten() {
if let Some(value) = self.call.lookup(py, name)?
&& from_py::<Value>(&value).is_ok_and(|caller| caller == *sent)
{
if let Some(value) = self.caller_value(py, name, sent)? {
body.set_item(name, value)?;
}
}
if self.operation == LoggingOperation::Messages {
// Python's Messages handler merges its transformed request body into the call details.
logger
.object(py)
.getattr("model_call_details")?
.call_method1("update", (&body,))?;
}
let headers = PyDict::new(py);
for (name, value) in &wire.headers {
headers.set_item(name, value)?;
@ -469,9 +536,10 @@ impl LegacyLogging {
headers: headers.clone().unbind(),
context: context.clone(),
});
let input = self.logged_input(py, &body)?;
self.logger()?.pre_call(
py,
self.input_description(),
&input,
self.logged_api_key(context),
&body,
&headers,
@ -525,7 +593,11 @@ impl LegacyLogging {
.getattr("litellm_params")?
.cast_into::<PyDict>()?
.copy()?;
params.set_item("custom_llm_provider", &facts.provider.provider)?;
if self.operation != LoggingOperation::Messages {
// Python's Messages handler hands the provider to the logger beside its params, never
// in them, and the proxy's logger carries over to the next fallback hop.
params.set_item("custom_llm_provider", &facts.provider.provider)?;
}
crate::python::Logging::Update.call(
py,
(
@ -548,8 +620,17 @@ impl LegacyLogging {
ResultSource::Provider => None,
ResultSource::Cache { key } => Some(key.clone()),
};
details.set_item("cache_hit", self.cache_key.is_some())?;
if self.cache_key.is_some() {
details.set_item("cache_hit", true)?;
} else {
details.call_method1("pop", ("cache_hit", py.None()))?;
}
details.set_item("cache_key", self.cache_key.as_deref())?;
if let (LoggingOperation::Messages, ResultSource::Provider, Some(request)) =
(self.operation, &facts.source, &self.request)
{
details.call_method1("update", (request.body.bind(py),))?;
}
Ok(HookStep::Ready(()))
}
@ -592,7 +673,7 @@ impl LegacyLogging {
self.end = Some(datetime(py, timing.end_time)?);
self.response = Some(response.clone_ref(py));
match &self.stream {
Some(stream) => self.stream_success(py, stream)?,
Some(stream) => self.stream_success(py, stream, false)?,
None => self.dispatch_success(py)?,
}
Ok(HookStep::Ready(()))
@ -662,8 +743,28 @@ impl LegacyLogging {
}
}
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() {
return Ok(());
}
let Some(stream) = self.stream.take() else {
return Ok(());
};
self.end = Some(datetime(py, epoch_seconds())?);
self.stream_success(py, &stream, true)
}
}
impl PythonOwned for LegacyLogging {
fn close(&mut self, py: Python<'_>) {
if self.logger.is_some()
&& let Err(error) = self.bill_abandoned_stream(py)
{
error.write_unraisable(py, None);
}
if let Some(logger) = self.logger.take()
&& let Err(error) = logger.restore_context(py)
{
@ -1114,10 +1215,10 @@ logger._provider_reported_cost = None
};
let previous = litellm_host::interceptors::ExecutionFacts {
provider: facts.provider.clone(),
cost: litellm_host::interceptors::Cost::Reported {
amount: serde_json::Number::from_f64(9.0).unwrap(),
cost: litellm_host::interceptors::Cost::CacheHit,
source: litellm_host::interceptors::ResultSource::Cache {
key: "earlier-hit".into(),
},
source: litellm_host::interceptors::ResultSource::Provider,
};
logging.result_ready(py, &previous).unwrap();
local(&locals, "logger")
@ -1135,13 +1236,59 @@ logger._provider_reported_cost = None
c"
assert logger.model_call_details.get('response_cost') is None, logger.model_call_details
assert logger._provider_reported_cost == reported, logger._provider_reported_cost
assert logger.model_call_details['cache_hit'] is cache_hit, logger.model_call_details
assert logger.model_call_details.get('cache_hit') is (True if cache_hit else None), logger.model_call_details
assert logger.update['custom_llm_provider'] == 'edenai', logger.update
",
);
});
}
#[rstest]
#[case::messages(crate::LoggingOperation::Messages, false)]
#[case::ocr(crate::LoggingOperation::Ocr, true)]
fn result_ready_puts_the_provider_in_the_logged_params_except_for_messages(
#[case] operation: crate::LoggingOperation,
#[case] in_params: bool,
) {
Python::initialize();
Python::attach(|py| {
let locals = namespace(py, PAYLOAD_LOGGER);
run(
py,
&locals,
c"
logger.litellm_params = {}
logger.optional_params = {}
logger.model_call_details = {}
logger._provider_reported_cost = None
",
);
let mut logging = LegacyLogging {
logger: Some(PythonLogger::new(local(&locals, "logger").unbind())),
operation,
..legacy_call(py, &locals, false)
};
let facts = litellm_host::interceptors::ExecutionFacts {
provider: litellm_host::interceptors::ProviderIdentity {
model: "native-test".into(),
provider: "anthropic".into(),
},
cost: litellm_host::interceptors::Cost::Deferred,
source: litellm_host::interceptors::ResultSource::Provider,
};
logging.result_ready(py, &facts).unwrap();
locals.set_item("in_params", in_params).unwrap();
run(
py,
&locals,
c"
assert logger.update['custom_llm_provider'] == 'anthropic', logger.update
assert ('custom_llm_provider' in logger.update['litellm_params']) is in_params, logger.update
",
);
});
}
const DOCUMENT: &str = "data:application/pdf;base64,YWJj";
const EDITED: &str = "data:application/pdf;base64,ZWRpdGVk";
@ -1276,7 +1423,10 @@ assert logger.update['custom_llm_provider'] == 'edenai', logger.update
#[rstest]
#[case::completion(crate::LoggingOperation::Completion, "Chat completions")]
#[case::responses(crate::LoggingOperation::Responses, "Responses")]
#[case::messages(crate::LoggingOperation::Messages, "Messages")]
#[case::messages(
crate::LoggingOperation::Messages,
r#"[{'role': 'user', 'content': '{"pages": [1]}'}]"#
)]
#[case::ocr(crate::LoggingOperation::Ocr, "OCR document processing")]
fn prepared_arguments_replace_the_legacy_view_without_losing_callback_aliases(
#[case] operation: crate::LoggingOperation,
@ -1317,7 +1467,7 @@ prepared = {'pages': replacement}
&locals,
c"
assert logger.pre['complete_input_dict']['pages'] is replacement
assert logger.pre_input == description
assert str(logger.pre_input) == description
assert original == [0]
",
);
@ -1873,10 +2023,15 @@ mod terminal_tests {
#[case::async_listened(
true,
c"",
&["async_success_handler", "enqueued", "sync_success_for_async_call"]
&["async_success_handler", "mark_async_success_enqueued", "enqueued", "sync_success_for_async_call"]
)]
#[case::async_deferred(true, c"logger._defer_async_logging = True", &["sync_success_for_async_call"])]
#[case::async_deferred(true, c"logger._defer_async_logging = True", &[])]
#[case::async_with_fallbacks(true, c"kwargs = {'fallbacks': ['other']}", &["sync_success_for_async_call"])]
#[case::async_deferred_with_fallbacks(
true,
c"logger._defer_async_logging = True\nkwargs = {'fallbacks': ['other']}",
&[]
)]
fn success_reaches_the_logging_handlers(
#[case] asynchronous: bool,
#[case] script: &CStr,
@ -1902,12 +2057,88 @@ mod terminal_tests {
&locals,
c"
assert all(value is response for name, value in logger.calls if name.endswith('_handler'))
assert hasattr(logger, '_native_pending_logging') == getattr(logger, '_defer_async_logging', False)
assert hasattr(logger, 'enqueue_deferred_logging') == getattr(logger, '_defer_async_logging', False)
assert getattr(logger, 'deferred_success_result', (response,))[0] is response
",
);
});
}
#[rstest]
#[case::accepted(c"", &["sync_success_for_async_call", "async_success_handler", "mark_async_success_enqueued", "enqueued"])]
#[case::accepted_with_fallbacks(c"kwargs = {'fallbacks': ['other']}", &["sync_success_for_async_call"])]
fn a_deferred_success_runs_sync_callbacks_only_once_released(
#[case] script: &CStr,
#[case] expected: &[&str],
) {
Python::initialize();
Python::attach(|py| {
let locals = namespace(
py,
c"response = object()\nlogger._defer_async_logging = True",
);
run(py, &locals, script);
let mut logging = logged(py, &locals, true);
succeed(py, &locals, &mut logging);
run(
py,
&locals,
c"
assert logger.names() == [], logger.calls
logger.enqueue_deferred_logging()
logger.enqueue_deferred_logging()
",
);
let names: Vec<String> = local(&locals, "logger")
.call_method0("names")
.unwrap()
.extract()
.unwrap();
assert_eq!(names, expected);
});
}
#[rstest]
#[case::deferred(
c"logger._defer_async_logging = True",
c"logger.enqueue_deferred_logging()",
&["sync_success_for_async_call", "async_success_handler", "mark_async_success_enqueued", "enqueued"]
)]
#[case::deferred_and_blocked(c"logger._defer_async_logging = True", c"", &[])]
#[case::immediate(
c"",
c"",
&["async_success_handler", "mark_async_success_enqueued", "enqueued", "sync_success_for_async_call", "sync_success_for_async_call"]
)]
fn a_fallback_attempt_on_the_same_logger_leaves_the_first_success_claimed(
#[case] script: &CStr,
#[case] release: &CStr,
#[case] expected: &[&str],
) {
Python::initialize();
Python::attach(|py| {
let locals = namespace(py, c"refused = object()\nresponse = refused");
run(py, &locals, script);
let mut first = logged(py, &locals, true);
succeed(py, &locals, &mut first);
run(py, &locals, c"response = object()");
let mut fallback = logged(py, &locals, true);
succeed(py, &locals, &mut fallback);
run(py, &locals, release);
let names: Vec<String> = local(&locals, "logger")
.call_method0("names")
.unwrap()
.extract()
.unwrap();
assert_eq!(names, expected);
run(
py,
&locals,
c"assert all(value is refused for name, value in logger.calls if name == 'async_success_handler')",
);
});
}
#[rstest]
fn dropped_observations_preserve_deferred_success_and_response_identity() {
Python::initialize();
@ -1933,10 +2164,10 @@ assert hasattr(logger, '_native_pending_logging') == getattr(logger, '_defer_asy
));
assert_eq!(sender.dropped_events(), 1);
run(py, &locals, c"
assert logger.names() == ['sync_success_for_async_call'], logger.calls
logger._native_pending_logging.release(True)
logger._native_pending_logging.release(True)
assert logger.names() == ['sync_success_for_async_call', 'async_success_handler', 'enqueued'], logger.calls
assert logger.names() == [], logger.calls
logger.enqueue_deferred_logging()
logger.enqueue_deferred_logging()
assert logger.names() == ['sync_success_for_async_call', 'async_success_handler', 'mark_async_success_enqueued', 'enqueued'], logger.calls
assert logger.calls[0][1] is response
assert logger.calls[1][1] is response
");
@ -1981,6 +2212,45 @@ assert chunks[1] is last
});
}
#[rstest]
#[case::cancelled_mid_stream(None, &["stamp_stream_head", "stream_opened", "stream_success", "restore"])]
#[case::after_success(Some(true), &["stamp_stream_head", "stream_opened", "stream_success", "restore"])]
#[case::after_failure(Some(false), &["stamp_stream_head", "stream_opened", "stream_failure", "restore"])]
fn closing_an_unfinished_stream_bills_its_delivered_chunks_once(
#[case] ended: Option<bool>,
#[case] expected: &[&str],
) {
Python::initialize();
Python::attach(|py| {
let locals = namespace(
py,
c"first = b'first'\nresponse = None\nfailure = ValueError('provider')\nhead = {}",
);
let mut logging = LegacyLogging {
operation: crate::LoggingOperation::Messages,
..logged(py, &locals, true)
};
logging
.on_stream_open(py, &local(&locals, "head").unbind())
.unwrap();
logging
.on_stream_chunk(py, &local(&locals, "first").unbind())
.unwrap();
match ended {
Some(true) => drop(succeed(py, &locals, &mut logging)),
Some(false) => drop(fail(py, &locals, &mut logging)),
None => {}
}
litellm_host_python::PythonOwned::close(&mut logging, py);
let names: Vec<String> = local(&locals, "logger")
.call_method0("names")
.unwrap()
.extract()
.unwrap();
assert_eq!(names, expected);
});
}
#[rstest]
#[case::synchronous(false, &["failure_handler"])]
#[case::asynchronous(true, &[])]
@ -2009,7 +2279,7 @@ assert chunks[1] is last
}
#[rstest::rstest]
fn internal_async_calls_skip_the_async_success_fan_out() {
fn internal_async_calls_skip_success_logging() {
Python::initialize();
Python::attach(|py| {
let locals = namespace(py, c"response = object()");
@ -2018,11 +2288,7 @@ assert chunks[1] is last
..logged(py, &locals, true)
};
succeed(py, &locals, &mut logging);
run(
py,
&locals,
c"assert logger.names() == ['sync_success_for_async_call'], logger.calls",
);
run(py, &locals, c"assert logger.calls == [], logger.calls");
});
}

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

@ -26,7 +26,7 @@ pub trait LegacyCallbacks {
fn pre_call(
&self,
py: Python<'_>,
input: &str,
input: &Bound<'_, PyAny>,
api_key: Option<&str>,
body: &Bound<'_, PyDict>,
headers: &Bound<'_, PyDict>,
@ -59,7 +59,16 @@ pub trait LegacyCallbacks {
fn defers_async_logging(&self, py: Python<'_>) -> bool;
fn defer_success(&self, py: Python<'_>, pending: &Bound<'_, PyAny>) -> PyResult<()>;
fn claim_async_success_log(&self, py: Python<'_>) -> PyResult<bool>;
fn defer_success(
&self,
py: Python<'_>,
pending: &Bound<'_, PyAny>,
response: &Option<Py<PyAny>>,
start: &Py<PyAny>,
end: &Option<Py<PyAny>>,
) -> PyResult<()>;
fn sync_success_for_async_call(
&self,
@ -148,7 +157,7 @@ impl LegacyCallbacks for PythonLogger {
fn pre_call(
&self,
py: Python<'_>,
input: &str,
input: &Bound<'_, PyAny>,
api_key: Option<&str>,
body: &Bound<'_, PyDict>,
headers: &Bound<'_, PyDict>,
@ -203,8 +212,21 @@ impl LegacyCallbacks for PythonLogger {
.unwrap_or(false)
}
fn defer_success(&self, py: Python<'_>, pending: &Bound<'_, PyAny>) -> PyResult<()> {
Logging::DeferSuccess.call(py, (self.object(py), pending))?;
fn claim_async_success_log(&self, py: Python<'_>) -> PyResult<bool> {
Logging::ClaimAsyncSuccess
.call(py, (self.object(py),))?
.extract()
}
fn defer_success(
&self,
py: Python<'_>,
pending: &Bound<'_, PyAny>,
response: &Option<Py<PyAny>>,
start: &Py<PyAny>,
end: &Option<Py<PyAny>>,
) -> PyResult<()> {
Logging::DeferSuccess.call(py, (self.object(py), pending, response, start, end))?;
Ok(())
}
@ -252,6 +274,7 @@ impl LegacyCallbacks for PythonLogger {
) -> PyResult<()> {
let coroutine =
Logging::AsyncSuccessHandler.call(py, (self.object(py), response, start, end))?;
Logging::MarkAsyncSuccessEnqueued.call(py, (self.object(py),))?;
let enqueue = Logging::Enqueue.call(py, (&coroutine,));
if enqueue.is_err()
&& let Err(error) = coroutine.call_method0("close")

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

@ -22,21 +22,35 @@ impl PendingSuccess {
self.logger
.enqueue_success(py, &self.response, &self.start, &self.end)
}
pub(crate) fn sync_for_async_call(&self, py: Python<'_>) -> PyResult<()> {
self.logger
.sync_success_for_async_call(py, &self.response, &self.start, &self.end)
}
}
#[pyclass]
pub(crate) struct PendingLogging {
pub(crate) pending: Option<PendingSuccess>,
pub(crate) enqueues: bool,
}
#[pymethods]
impl PendingLogging {
fn release(slf: &Bound<'_, Self>, py: Python<'_>, success: bool) -> PyResult<()> {
let pending = slf.borrow_mut().pending.take();
if let Some(pending) = pending
&& success
{
match pending.asynchronous(py) {
fn __call__(slf: &Bound<'_, Self>, py: Python<'_>) -> PyResult<()> {
let (pending, enqueues) = {
let mut slot = slf.borrow_mut();
(slot.pending.take(), slot.enqueues)
};
if let Some(pending) = pending {
let released = pending.sync_for_async_call(py).and_then(|()| {
if enqueues {
pending.asynchronous(py)
} else {
Ok(())
}
});
match released {
Err(error) if error.is_instance_of::<PyException>(py) => {
error.write_unraisable(py, Some(pending.logger.object(py)));
}
@ -87,6 +101,7 @@ mod tests {
start: py.None(),
end: Some(py.None()),
}),
enqueues: true,
},
)
.unwrap();
@ -108,7 +123,7 @@ observed = []
def on_enqueue(coroutine):
observed.append(marker.get())
pending.release(True)
pending()
logger.on_enqueue = on_enqueue
",
@ -118,28 +133,12 @@ logger.on_enqueue = on_enqueue
&locals,
c"
marker.set('release')
pending.release(True)
pending.release(True)
pending()
pending()
assert observed == ['release'], observed
assert logger.names() == ['async_success_handler', 'enqueued'], logger.calls
assert logger.names() == ['sync_success_for_async_call', 'async_success_handler', 'mark_async_success_enqueued', 'enqueued'], logger.calls
assert logger.calls[0][1] is response
",
);
});
}
#[test]
fn a_blocked_release_drops_the_success_for_good() {
Python::initialize();
Python::attach(|py| {
let locals = defer(py, c"");
run(
py,
&locals,
c"
pending.release(False)
pending.release(True)
assert logger.calls == [], logger.calls
assert logger.calls[1][1] is response
",
);
});
@ -168,7 +167,7 @@ logger.on_enqueue = on_enqueue
locals
.set_item("failure", py.eval(failure, None, Some(&locals)).unwrap())
.unwrap();
let released = local(&locals, "pending").call_method1("release", (true,));
let released = local(&locals, "pending").call0();
match released {
Ok(_) => assert!(!propagates),
Err(error) => {
@ -181,8 +180,8 @@ logger.on_enqueue = on_enqueue
py,
&locals,
c"
pending.release(True)
assert logger.names() == ['async_success_handler', 'enqueued', 'closed'], logger.calls
pending()
assert logger.names() == ['sync_success_for_async_call', 'async_success_handler', 'mark_async_success_enqueued', 'enqueued', 'closed'], logger.calls
assert unraisable_from(logger) == ([] if propagates else [failure])
",
);

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

@ -42,6 +42,8 @@ pub(crate) enum Logging {
Update,
#[strum(serialize = "pre_call")]
PreCall,
#[strum(serialize = "messages_input")]
MessagesInput,
#[strum(serialize = "post_call")]
PostCall,
#[strum(serialize = "record_response")]
@ -50,6 +52,8 @@ pub(crate) enum Logging {
RecordApiWindow,
#[strum(serialize = "defers_async_logging")]
DefersAsync,
#[strum(serialize = "claim_async_success_log")]
ClaimAsyncSuccess,
#[strum(serialize = "defer_success")]
DeferSuccess,
#[strum(serialize = "sync_success_for_async_call")]
@ -58,6 +62,8 @@ pub(crate) enum Logging {
SubmitSuccess,
#[strum(serialize = "async_success_handler")]
AsyncSuccessHandler,
#[strum(serialize = "mark_async_success_enqueued")]
MarkAsyncSuccessEnqueued,
#[strum(serialize = "enqueue_logging")]
Enqueue,
#[strum(serialize = "failure_handler")]

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

@ -62,6 +62,7 @@ FAKES = {
custom_llm_provider=provider,
),
'pre_call': lambda logger, input, api_key, additional_args: logger.pre_call(input, api_key, additional_args),
'messages_input': lambda body, kwargs: [{'role': 'user', 'content': json.dumps(body)}],
'post_call': lambda logger, original_response, api_key, additional_args: logger.post_call(
original_response, api_key, additional_args
),
@ -72,7 +73,11 @@ FAKES = {
llm_api_duration_ms=(end - start) * 1000
),
'defers_async_logging': lambda logger: bool(getattr(logger, '_defer_async_logging', False)),
'defer_success': lambda logger, pending: setattr(logger, '_native_pending_logging', pending),
'claim_async_success_log': lambda logger: logger.claim_async_success_log(),
'defer_success': lambda logger, pending, response, start, end: (
setattr(logger, 'deferred_success_result', (response, start, end)),
setattr(logger, 'enqueue_deferred_logging', pending),
),
'sync_success_for_async_call': lambda logger, response, start, end: logger.handle_sync_success_callbacks_for_async_calls(
response, start, end
),
@ -81,6 +86,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),
'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),
'custom_pricing_fields': lambda: ('ocr_cost_per_page',),
@ -92,7 +98,7 @@ FAKES = {
'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_success': lambda logger, url_route, endpoint_type, request_body, chunks, start, end, first_chunk: logger.record(
'stream_success': lambda logger, url_route, endpoint_type, request_body, chunks, start, end, first_chunk, stream_teardown: logger.record(
'stream_success', list(chunks)
),
'stream_failure': lambda logger, endpoint_type, request_body, chunks, error: logger.record('stream_failure', error),
@ -129,6 +135,7 @@ class StubCoroutine:
class StubLogger:
def __init__(self):
self.calls = []
self.model_call_details = {}
self.hooks = {}
self.on_enqueue = lambda coroutine: None
@ -156,6 +163,10 @@ class StubLogger:
self.record('async_success_handler', response)
return StubCoroutine(self)
def claim_async_success_log(self):
claimed, self.claimed = getattr(self, 'claimed', False), True
return not claimed
def handle_sync_success_callbacks_for_async_calls(self, response, start, end):
self.record('sync_success_for_async_call', response)

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

@ -28,6 +28,11 @@ pub trait PythonBinding: PythonOwned {
) -> PyResult<Py<PyAny>>;
/// 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 encode_stream_head(
&mut self,
py: Python<'_>,

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

@ -399,8 +399,16 @@ where
Ok(step) => step,
Err(error) => return self.failure(error, FailureOrigin::Host).await,
};
match self.settle(step).await {
Ok(response) => self.succeeded(response).await,
let response = match self.settle(step).await {
Ok(response) => response,
Err(error) => return self.failure(error, FailureOrigin::Host).await,
};
let delivered = {
let mut parts = self.parts();
Python::attach(|py| parts.binding.deliver_response(py, response.bind(py)))
};
match delivered {
Ok(()) => self.succeeded(response).await,
Err(error) => self.failure(error, FailureOrigin::Host).await,
}
}

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

@ -16,6 +16,15 @@ pub enum TlsSource {
#[error(transparent)]
pub struct Error(Arc<Repr>);
// A streamed body the connection ended early, after the response head arrived; `expected` is
// the declared `content-length`, `None` for a chunked body.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct BodyCut {
pub received: u64,
pub expected: Option<u64>,
pub reset: bool,
}
#[derive(Debug, thiserror::Error)]
pub(crate) enum Repr {
#[error("could not read {}", path.display())]
@ -56,6 +65,12 @@ pub(crate) enum Repr {
Connect(#[source] reqwest::Error),
#[error("upstream network error")]
Network(#[source] reqwest::Error),
#[error("upstream network error")]
BodyCut {
#[source]
source: reqwest::Error,
cut: BodyCut,
},
#[error("upstream request timed out")]
Timeout(#[source] reqwest::Error),
/// No bytes arrived within the read timeout, which the caller enforces per read; `body` says
@ -201,6 +216,37 @@ impl Error {
matches!(&*self.0, Repr::ReadTimeout { body: true, .. })
}
pub fn body_read(error: reqwest::Error, received: u64, expected: Option<u64>) -> Self {
let mut source: Option<&(dyn std::error::Error + 'static)> = Some(&error);
let kind = std::iter::from_fn(|| {
let current = source?;
source = current.source();
Some(current)
})
.find_map(|cause| cause.downcast_ref::<io::Error>().map(io::Error::kind));
let reset = match kind {
Some(io::ErrorKind::UnexpectedEof) => false,
Some(io::ErrorKind::ConnectionReset) => true,
_ => return error.into(),
};
Repr::BodyCut {
source: error.without_url(),
cut: BodyCut {
received,
expected,
reset,
},
}
.into()
}
pub fn body_cut(&self) -> Option<BodyCut> {
match &*self.0 {
Repr::BodyCut { cut, .. } => Some(*cut),
_ => None,
}
}
pub fn computed_header(name: impl Into<String>) -> Self {
Repr::ComputedHeader(name.into()).into()
}
@ -226,6 +272,7 @@ impl Error {
pub fn dispatched(&self) -> bool {
match &*self.0 {
Repr::Network(_)
| Repr::BodyCut { .. }
| Repr::Timeout(_)
| Repr::ReadTimeout { .. }
| Repr::Status { .. }
@ -285,6 +332,7 @@ impl Error {
| Repr::Signature(_)
| Repr::Connect(_)
| Repr::Network(_)
| Repr::BodyCut { .. }
| Repr::Timeout(_)
| Repr::ReadTimeout { .. }
| Repr::Status { .. }
@ -318,6 +366,7 @@ impl Error {
| Repr::Signature(_)
| Repr::Connect(_)
| Repr::Network(_)
| Repr::BodyCut { .. }
| Repr::Timeout(_)
| Repr::ReadTimeout { .. }
| Repr::Status { .. }
@ -355,6 +404,7 @@ impl Error {
| Repr::Signature(_)
| Repr::Connect(_)
| Repr::Network(_)
| Repr::BodyCut { .. }
| Repr::Timeout(_)
| Repr::ReadTimeout { .. }
| Repr::Status { .. }
@ -394,6 +444,7 @@ impl Error {
| Repr::UnsignableHeader(_)
| Repr::Signature(_)
| Repr::Network(_)
| Repr::BodyCut { .. }
| Repr::Timeout(_)
| Repr::ReadTimeout { .. }
| Repr::Status { .. }
@ -437,6 +488,7 @@ impl Error {
| Repr::Signature(_)
| Repr::Connect(_)
| Repr::Network(_)
| Repr::BodyCut { .. }
| Repr::Status { .. }
| Repr::WebSocketUrl(_)
| Repr::WebSocketScheme(_)
@ -473,6 +525,7 @@ impl Error {
| Repr::Signature(_)
| Repr::Connect(_)
| Repr::Network(_)
| Repr::BodyCut { .. }
| Repr::Timeout(_)
| Repr::ReadTimeout { .. }
| Repr::WebSocketUrl(_)
@ -513,6 +566,7 @@ impl Error {
| Repr::Signature(_)
| Repr::Connect(_)
| Repr::Network(_)
| Repr::BodyCut { .. }
| Repr::WebSocketUrl(_)
| Repr::WebSocketScheme(_)
| Repr::WebSocketSchemeRewrite(_)
@ -548,6 +602,7 @@ impl Error {
| Repr::Signature(_)
| Repr::Connect(_)
| Repr::Network(_)
| Repr::BodyCut { .. }
| Repr::Timeout(_)
| Repr::ReadTimeout { .. }
| Repr::Status { .. }
@ -742,4 +797,53 @@ mod tests {
assert_eq!(error.upstream_response(), Some((429, body.as_str())));
assert_eq!(error.status(), Some(429));
}
#[rstest]
#[case::chunked("transfer-encoding: chunked\r\n\r\n5\r\nhello\r\n", None)]
#[case::content_length("content-length: 100\r\n\r\nhello", Some(100))]
#[tokio::test]
async fn a_connection_that_ends_mid_body_is_a_body_cut(
#[case] tail: &'static str,
#[case] expected: Option<u64>,
) {
use std::io::{Read as _, Write as _};
let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("bind");
let address = listener.local_addr().expect("address");
let server = std::thread::spawn(move || {
let (mut stream, _) = listener.accept().expect("accept");
let mut request = [0_u8; 1024];
let _ = stream.read(&mut request).expect("request");
stream
.write_all(format!("HTTP/1.1 200 OK\r\n{tail}").as_bytes())
.expect("response");
});
let mut response = reqwest::Client::builder()
.no_proxy()
.build()
.expect("client")
.get(format!("http://{address}/"))
.send()
.await
.expect("head");
server.join().expect("server");
let mut received = 0_u64;
let error = loop {
match response.chunk().await {
Ok(Some(chunk)) => received += chunk.len() as u64,
Ok(None) => panic!("the body must end with an error"),
Err(error) => break crate::Error::body_read(error, received, expected),
}
};
assert_eq!(
error.body_cut(),
Some(super::BodyCut {
received: 5,
expected,
reset: false,
})
);
assert!(error.dispatched());
assert_eq!(error.to_string(), "upstream network error");
}
}

2
litellm-rust/crates/http/src/lib.rs
View file

@ -18,7 +18,7 @@ pub mod websocket;
pub use client::Client;
pub use config::{ClientIdentity, HttpClientConfig, Resolution, Verify};
pub use error::{Error, TlsSource};
pub use error::{BodyCut, Error, TlsSource};
pub use pool::{ClientVariant, HttpClientPool};
pub use proxy::EnvironmentProxies;
pub use settings::{HttpSettings, HttpSettingsLayer, SslVerify, TcpKeepalive};

83
litellm-rust/crates/inference-messages/src/common_utils.rs
View file

@ -1,5 +1,5 @@
use litellm_core_utils::get_llm_provider_logic::LlmProviders;
use litellm_http::request::string_headers as shared_string_headers;
use litellm_llms::ErrorDetail;
use litellm_llms::tencent::messages::transformation::TENCENT_MESSAGES_CONFIG;
use litellm_llms::{
anthropic::messages::transformation::ANTHROPIC_MESSAGES_CONFIG,
@ -19,8 +19,6 @@ use serde_json::{Map, Value};
use super::Error;
const HEADER_CONTEXT: &str = "messages";
#[derive(Clone, Copy)]
pub(crate) struct MessagesProvider {
provider: LlmProviders,
@ -75,7 +73,52 @@ pub(crate) fn messages_provider(provider: LlmProviders, model: &str) -> Option<M
pub(super) fn string_headers(
extra_headers: Option<Map<String, Value>>,
) -> Result<Vec<(String, String)>, Error> {
shared_string_headers(HEADER_CONTEXT, extra_headers).map_err(Error::from)
extra_headers
.unwrap_or_default()
.into_iter()
.map(|(name, value)| match value {
Value::String(value) => Ok((name, value)),
value => Err(Error::invalid_request(ErrorDetail::invalid(
"extra_headers",
NonStringHeader {
actual: python_type_name(&value),
name,
},
))),
})
.collect()
}
/// A header value the caller set to something other than a string, named by its Python type
/// since the caller wrote it in Python.
#[derive(Debug)]
struct NonStringHeader {
name: String,
actual: &'static str,
}
impl std::fmt::Display for NonStringHeader {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(
formatter,
"header {} must be a string, got {}",
self.name, self.actual
)
}
}
impl std::error::Error for NonStringHeader {}
fn python_type_name(value: &Value) -> &'static str {
match value {
Value::Null => "NoneType",
Value::Bool(_) => "bool",
Value::Number(number) if number.is_f64() => "float",
Value::Number(_) => "int",
Value::String(_) => "str",
Value::Array(_) => "list",
Value::Object(_) => "dict",
}
}
#[cfg(test)]
@ -85,7 +128,6 @@ mod tests {
use rstest::rstest;
use super::{messages_provider, string_headers};
use crate::Error;
use litellm_core_utils::get_llm_provider_logic::LlmProviders;
#[rstest]
@ -113,17 +155,28 @@ mod tests {
assert!(messages_provider(provider, model).is_none());
}
#[test]
fn string_headers_rejects_non_string_values() {
let headers = json!({"x-count": 3}).as_object().unwrap().clone();
let err = string_headers(Some(headers)).expect_err("non-string header rejected");
#[rstest]
#[case::int(json!(12345), "int")]
#[case::float(json!(1.5), "float")]
#[case::bool(json!(true), "bool")]
#[case::null(json!(null), "NoneType")]
#[case::list(json!([1]), "list")]
#[case::dict(json!({"a": 1}), "dict")]
fn non_string_header_is_a_bad_request_naming_the_header_and_its_type(
#[case] value: serde_json::Value,
#[case] python_type: &str,
) {
let headers = json!({"x-kept": "v", "x-count": value})
.as_object()
.unwrap()
.clone();
let error = string_headers(Some(headers)).expect_err("non-string header refused");
assert!(error.is_request(), "{error:?}");
assert_eq!(error.upstream_response(), None);
let rejection = error.rejected_request().expect("a rejected request");
assert_eq!(
err,
Error::Headers(litellm_http::request::HeaderError {
context: "messages",
name: "x-count".to_string(),
actual: "number",
})
litellm_error::Report::new(rejection).to_string(),
format!("invalid extra_headers: header x-count must be a string, got {python_type}")
);
}
}

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

@ -26,10 +26,10 @@ use litellm_llms_types::formats::messages::{
};
use litellm_secrets::source::Secrets;
use litellm_tracing::ByteChunk;
use serde_json::Value;
use serde_json::{Map, Value};
use super::{
Error, MessagesCall, MessagesCallResponse, MessagesRoute,
Error, MessagesBody, MessagesCall, MessagesCallResponse, MessagesRoute,
common_utils::{MessagesProvider, merge_messages_header_sources, string_headers},
constants::MESSAGES_TIMEOUT_SECS,
route::Messages,
@ -85,7 +85,10 @@ pub(super) async fn execute(
),
};
let reported = match &output {
MessagesCallResponse::Complete(response) => provider.config().reported_cost(response),
MessagesCallResponse::Complete(response) => match response {
MessagesBody::Message(message) => provider.config().reported_cost(message),
MessagesBody::Unchecked(_) => None,
},
MessagesCallResponse::Stream { .. } => None,
};
context
@ -130,9 +133,10 @@ async fn provider_call(
)?;
let transform_context =
MessagesTransformContext::new(shaping.capabilities, shaping.settings.drop_params);
let filtered = requested_optional_params(sanitized, provider, &transform_context)?;
let trimmed =
without_additional_drop_params(filtered, &shaping.settings.additional_drop_params)?;
let without_listed =
without_additional_drop_params(sanitized, &shaping.settings.additional_drop_params)?;
let trimmed = requested_optional_params(without_listed, provider, &transform_context)?;
let requested_params = serde_json::to_value(&trimmed.params).map_err(serialize_failure)?;
let scoped =
get_provider_specific_headers(provider_specific_header.as_ref(), provider.as_str());
@ -173,7 +177,7 @@ async fn provider_call(
let request_context = RequestContext {
model: body.model.clone(),
custom_llm_provider: provider.as_str().to_string(),
optional_params: serde_json::to_value(&body.params).map_err(serialize_failure)?,
optional_params: requested_params,
secret_fields: Vec::new(),
api_key: api_key.map(SecretValue::new),
};
@ -346,8 +350,7 @@ async fn call_provider(
// body it cannot parse, and the public mapping classifies it by that status.
return Err(litellm_http::Error::upstream(status, text).into());
}
decode_response(config, &identity.model, &text)
.map(|message| MessagesCallResponse::Complete(Box::new(message)))
decode_response(config, &identity.model, &text).map(MessagesCallResponse::Complete)
}
fn serialize_failure(err: serde_json::Error) -> Error {
@ -383,13 +386,20 @@ fn timed_chunks(
response: reqwest::Response,
timeout: Duration,
) -> impl futures_util::Stream<Item = Result<Bytes, litellm_http::Error>> + Send + 'static {
futures_util::stream::try_unfold(response, move |mut response| async move {
let chunk = tokio::time::timeout(timeout, response.chunk())
.await
.map_err(litellm_http::Error::body_read_timeout)?
.map_err(litellm_http::Error::from)?;
Ok(chunk.map(|chunk| (chunk, response)))
})
let expected = response.content_length();
futures_util::stream::try_unfold(
(response, 0_u64),
move |(mut response, received)| async move {
let chunk = tokio::time::timeout(timeout, response.chunk())
.await
.map_err(litellm_http::Error::body_read_timeout)?
.map_err(|error| litellm_http::Error::body_read(error, received, expected))?;
Ok(chunk.map(|chunk| {
let received = received + chunk.len() as u64;
(chunk, (response, received))
}))
},
)
}
/// The provider's error, its body read under `timeout`: a stream has no overall deadline once
@ -411,15 +421,23 @@ fn decode_response(
config: &dyn BaseMessagesConfig,
model: &str,
text: &str,
) -> Result<MessagesResponse, Error> {
let response = serde_json::from_str(text).map_err(|err| {
Error::invalid_response(litellm_llms::ErrorDetail::invalid(
"messages response JSON",
err,
))
})?;
) -> Result<MessagesBody, Error> {
let response = match serde_json::from_str::<MessagesResponse>(text) {
Ok(response) => response,
Err(err) => {
let Ok(mut body) = serde_json::from_str::<Map<String, Value>>(text) else {
return Err(Error::invalid_response(litellm_llms::ErrorDetail::invalid(
"messages response JSON",
err,
)));
};
body.entry("stop_sequence").or_insert(Value::Null);
return Ok(MessagesBody::Unchecked(body));
}
};
config
.transform_anthropic_messages_response(model, response)
.map(|message| MessagesBody::Message(Box::new(message)))
.map_err(Error::from)
.map_err(Error::after_send)
}

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

@ -17,8 +17,8 @@ use crate::common_utils::messages_provider;
pub use common_utils::merge_messages_header_sources;
pub use litellm_inference::Error;
pub use types::{
MessagesCall, MessagesCallResponse, MessagesSettings, MessagesShaping, litellm_params,
messages_body,
MessagesBody, MessagesCall, MessagesCallResponse, MessagesSettings, MessagesShaping,
litellm_params, messages_body,
};
#[derive(Clone)]

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

@ -5,11 +5,10 @@ use litellm_host::{
call::{HostedCompletion, HostedMachine, hosted_call},
protocol::Protocol,
};
use litellm_llms_types::formats::messages::MessagesResponse;
use super::{Error, MessagesCall};
use super::{Error, MessagesBody, MessagesCall};
pub type MessagesOutput = HostedCompletion<Box<MessagesResponse>>;
pub type MessagesOutput = HostedCompletion<MessagesBody>;
/// The upstream response as the caller sees it at stream hand-off, before any chunk.
pub struct MessagesStreamHead {
@ -19,7 +18,7 @@ pub struct MessagesStreamHead {
pub struct Messages;
impl Protocol for Messages {
type Response = Box<MessagesResponse>;
type Response = MessagesBody;
type Error = Error;
type Request = MessagesCall;
type HostCall = Infallible;

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

@ -42,7 +42,16 @@ pub(super) fn invalid_request(err: serde_json::Error) -> Error {
}
pub type MessagesCallResponse =
CallOutput<Box<MessagesResponse>, super::route::MessagesStreamHead, Bytes, Error>;
CallOutput<MessagesBody, super::route::MessagesStreamHead, Bytes, Error>;
// Python's `transform_anthropic_messages_response` returns any JSON object the provider sends
// with a 200, message or not, so one that is not a message passes through as sent.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(untagged)]
pub enum MessagesBody {
Message(Box<MessagesResponse>),
Unchecked(Map<String, Value>),
}
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
pub struct MessagesShaping {

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

@ -1,7 +1,9 @@
use std::{sync::Arc, time::Duration};
use futures_util::TryStreamExt;
use litellm_inference_messages::{Error, MessagesCall, MessagesCallResponse, MessagesShaping};
use litellm_inference_messages::{
Error, MessagesBody, MessagesCall, MessagesCallResponse, MessagesShaping,
};
use litellm_inference_testing::{RecordingSecrets, no_secrets};
use litellm_llms_types::formats::messages::{MessagesRequest, MessagesResponse};
use rstest::{fixture, rstest};
@ -74,7 +76,10 @@ async fn run(call: MessagesCall) -> Result<MessagesCallResponse, Error> {
async fn run_message(call: MessagesCall) -> MessagesResponse {
match run(call).await.unwrap() {
MessagesCallResponse::Complete(message) => *message,
MessagesCallResponse::Complete(message) => match message {
MessagesBody::Message(message) => *message,
MessagesBody::Unchecked(_) => panic!("expected a message"),
},
MessagesCallResponse::Stream { .. } => panic!("expected a completed message"),
}
}

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

@ -50,42 +50,47 @@ async fn messages_cache_identity_includes_provider_native_parameters(
#[case] original: Value,
#[case] changed: Value,
) {
use litellm_inference_messages::route::Messages;
use litellm_inference_messages::{MessagesBody, route::Messages};
use litellm_llms_types::formats::messages::MessagesResponse;
let calls = AtomicUsize::new(0);
for (value, expected_call) in [(original.clone(), 0), (changed, 1), (original, 0)] {
let response =
execute_unary::<Messages, _, _>(
CacheRequest::from_wire(
ProviderIdentity {
model: "test".into(),
provider: "anthropic".into(),
},
Some(&WireRequest {
url: "https://example.test/v1/messages".into(),
headers: vec![],
body: json!({
"model":"test", "messages":[{"role":"user","content":"hello"}],
"max_tokens":32, (field):value
}),
}),
),
Some(cache.clone()),
Some(CacheOptions::new(CacheScope::Shared)),
&(),
None,
|| async {
let call = calls.fetch_add(1, Ordering::SeqCst);
Ok(Box::new(serde_json::from_value::<MessagesResponse>(json!({
"id":call.to_string(), "type":"message", "role":"assistant", "model":"test",
"content":[{"type":"text","text":format!("answer {call}")}],
"stop_reason":"end_turn", "stop_sequence":null
})).unwrap()))
let response = execute_unary::<Messages, _, _>(
CacheRequest::from_wire(
ProviderIdentity {
model: "test".into(),
provider: "anthropic".into(),
},
)
.await
.unwrap();
Some(&WireRequest {
url: "https://example.test/v1/messages".into(),
headers: vec![],
body: json!({
"model":"test", "messages":[{"role":"user","content":"hello"}],
"max_tokens":32, (field):value
}),
}),
),
Some(cache.clone()),
Some(CacheOptions::new(CacheScope::Shared)),
&(),
None,
|| async {
let call = calls.fetch_add(1, Ordering::SeqCst);
Ok(MessagesBody::Message(Box::new(
serde_json::from_value::<MessagesResponse>(json!({
"id":call.to_string(), "type":"message", "role":"assistant", "model":"test",
"content":[{"type":"text","text":format!("answer {call}")}],
"stop_reason":"end_turn", "stop_sequence":null
}))
.unwrap(),
)))
},
)
.await
.unwrap();
let MessagesBody::Message(response) = response else {
panic!("expected a message");
};
assert_eq!(response.id, expected_call.to_string());
assert_eq!(
response.content[0]["text"],
@ -253,7 +258,7 @@ async fn provider_reported_cost_is_delivered_once_and_never_charged_for_a_cache_
#[case] cost: Value,
#[case] expected: Option<f64>,
) {
use litellm_inference_messages::{MessagesCall, MessagesCallResponse};
use litellm_inference_messages::{MessagesBody, MessagesCall, MessagesCallResponse};
let response = json!({"id": "msg_cost", "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}, "cost": cost});
@ -277,6 +282,9 @@ async fn provider_reported_cost_is_delivered_once_and_never_charged_for_a_cache_
panic!("expected Messages responses");
};
assert_eq!(first, second);
let MessagesBody::Message(first) = first else {
panic!("expected a message");
};
assert_eq!(first.extra.get("cost"), Some(&cost));
assert_eq!(support::received(&upstream).await.len(), 1);
let facts = hooks.facts.lock().unwrap();

9
litellm-rust/crates/inference-messages/tests/live/support.rs
View file

@ -4,7 +4,9 @@ use futures_util::{FutureExt, TryStreamExt};
use litellm_host::interceptors::{
ExecutionFacts, Interceptors, RawResponse, RequestContext, ResultSource, WireRequest,
};
use litellm_inference_messages::{Error, MessagesCall, MessagesCallResponse, MessagesRoute};
use litellm_inference_messages::{
Error, MessagesBody, MessagesCall, MessagesCallResponse, MessagesRoute,
};
use litellm_inference_testing::live::{LiveResources, model, within_deadline};
use litellm_llms_types::formats::messages::MessagesResponse;
use serde_json::{Value, json};
@ -111,7 +113,10 @@ pub(super) fn call(provider: &str, payload: Value) -> MessagesCall {
pub(super) fn complete(response: MessagesCallResponse) -> Box<MessagesResponse> {
match response {
MessagesCallResponse::Complete(message) => message,
MessagesCallResponse::Complete(message) => match message {
MessagesBody::Message(message) => message,
MessagesBody::Unchecked(_) => panic!("a complete request returned no message"),
},
MessagesCallResponse::Stream { .. } => panic!("a complete request returned a stream"),
}
}

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

@ -6,9 +6,11 @@ use litellm_host::{
interceptors::{ExecutionFacts, RequestContext, ResultSource, WireRequest},
lifecycle::CallEvent,
};
use litellm_inference_messages::{MessagesCallResponse, route::Messages};
use litellm_inference_messages::{MessagesBody, MessagesCallResponse, route::Messages};
use litellm_inference_testing::{RecordingSecrets, no_secrets};
use litellm_llms::base_llm::messages::context::MessagesModelCapabilities as AnthropicModelCapabilities;
use litellm_llms::base_llm::messages::context::{
MessagesModelCapabilities as AnthropicModelCapabilities, SupportedEffortTiers,
};
use rstest::rstest;
use super::*;
@ -338,7 +340,10 @@ async fn response_mode_follows_the_intercepted_request(
}
MessagesCallResponse::Complete(message) => {
assert_eq!(expected_stream, Some(false));
assert_eq!(*message, serde_json::from_value(message_body()).unwrap());
assert_eq!(
message,
MessagesBody::Message(serde_json::from_value(message_body()).unwrap())
);
}
}
Ok::<_, Error>(())
@ -419,7 +424,7 @@ async fn no_raw_response_is_emitted_for_a_stream_or_a_failure(
assert!(host.raw_responses().is_empty());
}
/// Python logs `optional_params` as what it is about to send, so a dropped param must
/// Python logs `optional_params` after its requested-param gate, so a dropped param must
/// not resurface in callbacks.
#[rstest]
#[tokio::test]
@ -450,3 +455,82 @@ async fn the_request_context_carries_the_shaped_params_without_model_or_messages
.unwrap_or_else(|seen| panic!("before_provider_request runs once, saw {}", seen.len()));
assert_eq!(optional_params, json!({"max_tokens": 16}));
}
#[rstest]
#[tokio::test]
async fn the_request_context_leaves_out_a_listed_param_the_model_refuses(call: MessagesCall) {
let upstream = upstream([message_response()]).await;
let host = RecordingHost::passthrough(authenticated(
MessagesCall {
shaping: MessagesShaping {
settings: MessagesSettings {
additional_drop_params: vec!["temperature".into()],
..MessagesSettings::default()
},
capabilities: AnthropicModelCapabilities {
supports_sampling_params: false,
..AnthropicModelCapabilities::default()
},
},
..with_fields(call, json!({"temperature": 0.2}))
},
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, json!({"max_tokens": 16}));
}
#[rstest]
#[case::reasoning_effort_on_an_adaptive_model(
AnthropicModelCapabilities {
supports_reasoning: true,
supports_adaptive_thinking: true,
supports_output_config: true,
effort_tiers: SupportedEffortTiers { low: true, ..SupportedEffortTiers::default() },
..AnthropicModelCapabilities::default()
},
json!({"reasoning_effort": "low"})
)]
#[case::adaptive_payload_on_a_legacy_model(
AnthropicModelCapabilities {
supports_reasoning: true,
..AnthropicModelCapabilities::default()
},
json!({"max_tokens": 3000, "thinking": {"type": "adaptive"}, "output_config": {"effort": "high"}})
)]
#[tokio::test]
async fn the_request_context_carries_the_requested_reasoning_not_its_translation(
call: MessagesCall,
#[case] capabilities: AnthropicModelCapabilities,
#[case] fields: Value,
) {
let upstream = upstream([message_response()]).await;
let host = RecordingHost::passthrough(authenticated(
MessagesCall {
shaping: MessagesShaping {
capabilities,
..MessagesShaping::default()
},
..with_fields(call, fields.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,
Value::Object(
object(json!({"max_tokens": 16}))
.into_iter()
.chain(object(fields))
.collect()
)
);
}

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

@ -5,7 +5,7 @@ use std::{
use litellm_http::{HttpSettings, Resolution};
use litellm_inference_messages::{
Error, MessagesCall, MessagesSettings, MessagesShaping,
Error, MessagesBody, MessagesCall, MessagesSettings, MessagesShaping,
route::{Messages, MessagesMachine, MessagesOutput},
};
use litellm_inference_testing::RecordingSecrets;
@ -119,7 +119,10 @@ async fn run(call: MessagesCall) -> Result<MessagesOutput, Error> {
async fn run_message(call: MessagesCall) -> MessagesResponse {
match run(call).await.expect("messages call succeeds") {
MessagesOutput::Complete(message) => *message,
MessagesOutput::Complete(message) => match message {
MessagesBody::Message(message) => *message,
MessagesBody::Unchecked(_) => panic!("a call returned no message"),
},
MessagesOutput::StreamEnded | MessagesOutput::Detached => {
panic!("a non-streaming call returned a stream")
}

24
litellm-rust/crates/inference-messages/tests/messages/request.rs
View file

@ -530,6 +530,30 @@ async fn unsupported_params_are_dropped_under_drop_params_and_rejected_without_i
assert_eq!(sent["max_tokens"], 16);
}
#[rstest]
#[tokio::test]
async fn additional_drop_params_clear_a_param_the_model_does_not_support(call: MessagesCall) {
let upstream = upstream([message_response()]).await;
run_message(MessagesCall {
api_key: Some("sk".into()),
api_base: Some(upstream.uri()),
shaping: MessagesShaping {
capabilities: sampling_removed(),
settings: MessagesSettings {
additional_drop_params: vec!["top_k".into()],
..MessagesSettings::default()
},
},
..with_fields(call, json!({"top_k": 5}))
})
.await;
let sent = only_request(&upstream).await.json();
assert_eq!(sent.get("top_k"), None);
assert_eq!(sent["max_tokens"], 16);
}
#[rstest]
#[case::adaptive_thinking(json!({"type": "adaptive"}), json!({"type": "adaptive", "display": "summarized"}))]
#[case::disabled_thinking(json!({"type": "disabled"}), json!({"type": "disabled"}))]

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

@ -2,7 +2,7 @@ use litellm_host::{
interceptors::{ExecutionFacts, ResultSource},
lifecycle::ExecutionEvent,
};
use litellm_inference_messages::{MessagesCallResponse, messages_body};
use litellm_inference_messages::{MessagesBody, MessagesCallResponse, messages_body};
use litellm_inference_testing::{
RecordingSecrets, http_config, no_secrets, provider_http, resources,
};
@ -48,6 +48,9 @@ async fn calls_defer_execution_until_polled(
let MessagesCallResponse::Complete(response) = future.await.unwrap() else {
panic!("expected a completed message");
};
let MessagesBody::Message(response) = response else {
panic!("expected a message");
};
assert_eq!(
response.content,
message_body()["content"].as_array().unwrap().as_slice()
@ -240,10 +243,42 @@ async fn an_unparseable_success_body_is_reported_as_the_upstream_answer(call: Me
assert_eq!(error.upstream_response(), Some((200, "not json")));
}
#[rstest]
#[case::error_object(
json!({"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}}),
json!({"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"},
"stop_sequence": null}),
)]
#[case::message_missing_id(
json!({"type": "message", "content": [], "stop_sequence": "###"}),
json!({"type": "message", "content": [], "stop_sequence": "###"}),
)]
#[tokio::test]
async fn a_success_body_that_is_not_a_message_passes_through(
call: MessagesCall,
#[case] body: Value,
#[case] expected: Value,
) {
let upstream = upstream([json_response(body)]).await;
let response = run(MessagesCall {
api_key: Some("sk".into()),
api_base: Some(upstream.uri()),
..call
})
.await
.expect("any JSON object is returned");
let MessagesOutput::Complete(response) = response else {
panic!("expected a completed response");
};
assert_eq!(serde_json::to_value(response).unwrap(), expected);
}
#[rstest]
#[tokio::test]
async fn an_unreadable_success_body_is_an_invalid_response(call: MessagesCall) {
let upstream = upstream([json_response(json!({"unexpected": true}))]).await;
let upstream = upstream([json_response(json!(["unexpected"]))]).await;
let error = run(MessagesCall {
api_key: Some("sk".into()),
@ -320,6 +355,9 @@ async fn the_facade_sends_through_the_injected_http_pool_configuration(call: Mes
let MessagesCallResponse::Complete(message) = response else {
panic!("a non-streaming request returns a message");
};
let MessagesBody::Message(message) = message else {
panic!("expected a message");
};
assert_eq!(message.id, "msg_1");
let sent = only_request(&upstream).await;
assert_eq!(sent.header("x-api-key"), Some("sk-ant"));
@ -412,6 +450,9 @@ async fn route_uses_injected_dependencies_and_optional_cache(
else {
panic!("expected a completed message");
};
let MessagesBody::Message(response) = response else {
panic!("expected a message");
};
assert_eq!(
response.content,
message_body()["content"].as_array().unwrap().as_slice()
@ -484,6 +525,9 @@ async fn cache_overrides_preserve_the_routes_isolated_scope(call: MessagesCall)
else {
panic!("expected a completed message");
};
let MessagesBody::Message(response) = response else {
panic!("expected a message");
};
assert_eq!(response.id, expected["id"].as_str().unwrap());
}
assert_eq!(received(&upstream).await.len(), 2);

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

@ -164,6 +164,14 @@ impl Error {
}
}
pub fn body_cut(&self) -> Option<litellm_http::BodyCut> {
match self {
Self::Llm(error) => error.body_cut(),
Self::Http(error) => error.body_cut(),
_ => None,
}
}
/// `Some(connect)` when the call timed out, where `connect` says the connection was never
/// established; `None` for every other error.
pub fn timeout(&self) -> Option<bool> {

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

@ -103,6 +103,13 @@ impl Error {
matches!(self, Self::Http(error) if error.is_body_read_timeout())
}
pub fn body_cut(&self) -> Option<litellm_http::BodyCut> {
match self {
Self::Http(error) => error.body_cut(),
_ => None,
}
}
pub fn is_timeout(&self) -> Option<bool> {
match self {
Self::Http(error) if error.is_timeout() => Some(error.is_connect_timeout()),

1
litellm-rust/crates/python-bridge/Cargo.toml
View file

@ -19,6 +19,7 @@ huggingface = ["litellm-token-counter/huggingface"]
tiktoken = ["litellm-token-counter/tiktoken"]
[dependencies]
rustls.workspace = true
litellm-secret-value.workspace = true
fancy-regex.workspace = true
litellm-tracing.workspace = true

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

@ -6,6 +6,8 @@ pub(crate) use selection::{Cached, Selection, configure, ensure_cache_disabled};
use pyo3::{exceptions::PyRuntimeError, prelude::*};
const CACHE_HOST_MODULE: &str = "litellm.rust_bridge.host.cache";
#[derive(Debug, thiserror::Error)]
pub(crate) enum HandleError {
#[error("this inference route does not support configured caches")]

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

@ -4,6 +4,6 @@ This directory lets Rust inference use a selected Python cache. `service.rs` imp
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
Await asynchronous cache operations through the existing host driver in the caller's task. Do not create another asyncio task or event loop. Cancellation must prevent subsequent provider requests and cache writes. Preserve ordinary cache failure handling without swallowing cancellation or other Python base exceptions
Await asynchronous cache operations through the existing host driver in the caller's task. Do not create another asyncio task or event loop. The one exception is the asynchronous unary messages store: Python's `_async_set_cache` writes the returned response from `create_cache_write_task` after the router has stamped its `_hidden_params`, so `host.rs` defers that store until the deployment transform has returned the response the caller receives and hands that response to `litellm.rust_bridge.host.cache.store_response`, which schedules the same write task. Lookups, sync stores and stream stores stay in the caller's task. Cancellation must prevent subsequent provider requests and cache writes. Preserve ordinary cache failure handling without swallowing cancellation or other Python base exceptions
Traverse retained Python references for GC and release pending replies when the call closes. Regression tests must require native inference with Python fallback disabled and assert observable hits, provider request counts, cache-key headers, task identity and cancellation

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

@ -9,7 +9,7 @@ use pyo3::{
use serde_json::Value;
use super::service::CacheCall;
use crate::cache::HandleError;
use crate::cache::{CACHE_HOST_MODULE, HandleError};
enum Pending {
Lookup(Reply<Result<Option<Value>, Error>>),
@ -21,6 +21,7 @@ pub(crate) struct PythonCache {
arguments: Option<Py<PyDict>>,
pending: Option<Pending>,
asynchronous: bool,
deferred: bool,
}
impl PythonCache {
@ -30,6 +31,7 @@ impl PythonCache {
arguments: None,
pending: None,
asynchronous,
deferred: false,
}
}
@ -64,7 +66,14 @@ impl PythonCache {
None,
)
}
CacheCall::Store { value, reply } => {
CacheCall::Store {
unary: true, reply, ..
} if self.asynchronous => {
self.deferred = true;
reply.send(Ok(()));
return Ok(None);
}
CacheCall::Store { value, reply, .. } => {
self.pending = Some(Pending::Store(reply));
(
if self.asynchronous {
@ -89,6 +98,19 @@ impl PythonCache {
self.resume(py, result)
}
pub fn store_response(&mut self, py: Python<'_>, response: &Bound<'_, PyAny>) -> PyResult<()> {
if !std::mem::take(&mut self.deferred) {
return Ok(());
}
let (Some(cache), Some(arguments)) = (self.cache.as_ref(), self.arguments.as_ref()) else {
return Err(HandleError::Unconfigured.into());
};
py.import(CACHE_HOST_MODULE)?
.getattr("store_response")?
.call((cache, response), Some(&arguments.bind(py).copy()?))?;
Ok(())
}
pub fn resume(
&mut self,
py: Python<'_>,
@ -124,6 +146,7 @@ impl PythonCache {
pub fn close(&mut self) {
self.pending = None;
self.deferred = false;
self.cache = None;
self.arguments = None;
}

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

@ -45,6 +45,7 @@ pub(crate) enum CacheCall {
},
Store {
value: Value,
unary: bool,
reply: Reply<Result<(), Error>>,
},
}
@ -101,8 +102,13 @@ where
) -> Pin<Box<dyn Future<Output = Result<(), Error>> + Send + 'a>> {
Box::pin(async move {
let value = to_python(value)?;
let unary = value.get(STREAM_EVENTS_KEY).is_none();
self.services
.call(|reply| CacheCall::Store { value, reply })
.call(|reply| CacheCall::Store {
value,
unary,
reply,
})
.await
.map_err(Error::unavailable)?
})

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

@ -1,11 +1,10 @@
use super::{HandleError, python};
use super::{CACHE_HOST_MODULE, HandleError, python};
use litellm_cache_response::{CacheOptions, CachePolicy, CacheScope, ScopedCache};
use litellm_host::{
machine::{HostServices, MachineFault},
protocol::Protocol,
};
use pyo3::{prelude::*, types::PyDict};
const CACHE_HOST_MODULE: &str = "litellm.rust_bridge.host.cache";
pub(crate) struct Cached<P>(std::marker::PhantomData<P>);

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

@ -29,6 +29,8 @@ const MESSAGES_UNSUPPORTED_PARAM: &str = "messages_unsupported_param";
/// The phase label (`Connection` or `Read`) of a timed-out call, for Python's timeout text, or
/// `Body` for a streamed body read that timed out after the head arrived.
const MESSAGES_TIMEOUT: &str = "messages_timeout";
// `(received, expected, reset)` of a streamed body the connection ended early.
const MESSAGES_BODY_CUT: &str = "messages_body_cut";
/// No credential reached the provider, which Python raises as `AuthenticationError`.
const MESSAGES_MISSING_CREDENTIAL: &str = "messages_missing_credential";
/// Python refuses the request only once its logger knows the resolved model and provider (the
@ -134,6 +136,9 @@ pub(crate) fn route_failure(
};
value.setattr(MESSAGES_TIMEOUT, phase)?;
}
if let Some(cut) = error.body_cut() {
value.setattr(MESSAGES_BODY_CUT, (cut.received, cut.expected, cut.reset))?;
}
if error.is_missing_credential() {
value.setattr(MESSAGES_MISSING_CREDENTIAL, true)?;
}

20
litellm-rust/crates/python-bridge/src/lib.rs
View file

@ -55,7 +55,9 @@ mod _native {
#[pymodule_export]
use crate::redis::store::sensitive_route_pins::NativeSensitiveRoutePins;
#[pymodule_export]
use crate::redis::store::spend_counters::{NativeSpendCounters, PySpendRead};
use crate::redis::store::spend_counters::{
NativeSpendCounters, PyPostCallSettlement, PySpendRead,
};
#[pymodule_export]
use crate::redis::store::{
NativePending, NativeRedisStoreError, router_cooldowns::NativeRouterCooldowns,
@ -90,6 +92,7 @@ mod _native {
#[pymodule_init]
fn init(module: &Bound<'_, PyModule>) -> PyResult<()> {
let _ = rustls::crypto::ring::default_provider().install_default();
let py = module.py();
let dict = module.dict();
dict.set_item(
@ -154,6 +157,7 @@ mod tests {
"NativeSensitiveRoutePins",
"NativeRedisStoreError",
"NativeSpendCounters",
"PostCallSettlement",
"SpendRead",
"LocalSizes",
"native_redis_breaker_states",
@ -186,4 +190,18 @@ mod tests {
assert_eq!(public_names, expected);
});
}
#[rstest::rstest]
fn module_import_installs_ring_as_the_process_crypto_provider() {
Python::initialize();
Python::attach(|py| {
native_module(py);
});
let installed = rustls::crypto::CryptoProvider::get_default()
.expect("importing _native should install a process crypto provider");
assert_eq!(
format!("{installed:?}"),
format!("{:?}", rustls::crypto::ring::default_provider())
);
}
}

27
litellm-rust/crates/python-bridge/src/redis/store/spend_counters.rs
View file

@ -2,12 +2,12 @@ use std::{collections::HashMap, sync::Arc};
use litellm_redis_store::{
Number,
spend_counters::{LostSeed, Reading, SpendCounter, SpendCounters},
spend_counters::{LostSeed, PostCallSettlement, Reading, SpendCounter, SpendCounters},
};
use pyo3::prelude::*;
use super::{
NativePending, PyData, PyNumber, kind_error, pending, redis_async, remote,
NativePending, PyData, PyNumber, duration, kind_error, pending, redis_async, remote,
request::{NativeRequestRedis, NativeSend},
ttl, unwrap_all,
};
@ -29,6 +29,18 @@ impl PySpendRead {
}
}
/// Whether every increment one `increment_many_post_call` declared has its reply.
#[pyclass(frozen, name = "PostCallSettlement")]
pub(crate) struct PyPostCallSettlement(PostCallSettlement);
#[pymethods]
impl PyPostCallSettlement {
#[getter]
fn settled(&self) -> bool {
self.0.is_settled()
}
}
fn increments(increments: Vec<(PyData<SpendCounter>, PyNumber)>) -> Vec<(SpendCounter, Number)> {
increments
.into_iter()
@ -73,6 +85,12 @@ impl NativeSpendCounters {
self.0.redis_ttl().map(|ttl| ttl.as_secs_f64())
}
/// The least TTL a Redis write gives a committed spend twin.
fn set_committed_ttl(&self, ttl_seconds: f64) -> PyResult<()> {
self.0.set_committed_ttl(duration(ttl_seconds)?);
Ok(())
}
/// `("redis" | "memory", value)`: where the answer came from, a Redis miss included.
fn read<'py>(
&self,
@ -125,14 +143,15 @@ impl NativeSpendCounters {
})
}
/// Declares the increments on the request's post-call batch; `False` when there is no Redis.
/// Declares the increments on the request's post-call batch; `None` when there is no Redis.
fn increment_many_post_call(
&self,
request: &NativeRequestRedis,
increments: Vec<(PyData<SpendCounter>, PyNumber)>,
) -> bool {
) -> Option<PyPostCallSettlement> {
self.0
.increment_many_post_call(&request.0, &self::increments(increments))
.map(PyPostCallSettlement)
}
fn set_max<'py>(

46
litellm-rust/crates/python-bridge/src/router/host.rs
View file

@ -357,7 +357,11 @@ impl HostCallResponse {
Self::Select(ids) | Self::Filtered(Ok((ids, _))),
HostCall::Select { candidates, .. },
) = (self, call)
&& let Some(id) = ids.iter().find(|id| !candidates.contains(id))
&& let Some(id) = {
let known: std::collections::HashSet<&str> =
candidates.iter().map(String::as_str).collect();
ids.iter().find(|id| !known.contains(id.as_str()))
}
{
return Err(DriveError::UnknownCandidate { id: id.clone() });
}
@ -688,14 +692,22 @@ fn shares_ids(candidates: &[Arc<Deployment>]) -> bool {
/// The candidates a host answer names, in its order. An id names every candidate that carries
/// it, so deployments sharing an id stay together through a filter, as Python's list keeps them.
fn answered_by_id(candidates: &[Arc<Deployment>], ids: &[String]) -> Vec<Arc<Deployment>> {
let mut taken = vec![false; candidates.len()];
let mut positions: std::collections::HashMap<&str, Vec<usize>> =
std::collections::HashMap::new();
for (position, deployment) in candidates.iter().enumerate() {
positions
.entry(deployment.id.as_str())
.or_default()
.push(position);
}
let mut kept = Vec::new();
for id in ids {
for (position, deployment) in candidates.iter().enumerate() {
if !taken[position] && deployment.id == *id {
taken[position] = true;
kept.push(deployment.clone());
}
if let Some(found) = positions.remove(id.as_str()) {
kept.extend(
found
.into_iter()
.map(|position| candidates[position].clone()),
);
}
}
kept
@ -1540,6 +1552,26 @@ mod tests {
);
}
#[rstest]
fn an_answer_by_id_keeps_its_order_and_every_candidate_sharing_an_id() {
let candidates: Vec<Arc<Deployment>> = ["a", "b", "a", "c"]
.into_iter()
.map(|id| {
let spec: litellm_router_core::DeploymentSpec =
serde_json::from_value(deployment("g", id)).unwrap();
Arc::new(Deployment::from_spec(spec).unwrap())
})
.collect();
let ids = ["c", "a", "x", "a"].map(String::from);
let kept = answered_by_id(&candidates, &ids);
let positions: Vec<usize> = kept
.iter()
.map(|d| candidates.iter().position(|c| Arc::ptr_eq(c, d)).unwrap())
.collect();
assert_eq!(positions, [3, 0, 2]);
}
#[rstest]
fn a_pick_among_deployments_sharing_an_id_is_answered_by_position() {
let mut second = deployment("primary", "a1");

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

@ -239,9 +239,13 @@ impl PythonBinding for MessagesPythonHost {
fn encode_response(
&mut self,
py: Python<'_>,
response: Box<litellm_llms_types::formats::messages::MessagesResponse>,
response: litellm_inference_messages::MessagesBody,
) -> PyResult<Py<PyAny>> {
public_response(py, ROUTE_HOST_MODULE, response.as_ref())
public_response(py, ROUTE_HOST_MODULE, &response)
}
fn deliver_response(&mut self, py: Python<'_>, response: &Bound<'_, PyAny>) -> PyResult<()> {
self.cache.store_response(py, response)
}
fn encode_stream_head(

1
litellm-rust/crates/redis-connect/Cargo.toml
View file

@ -17,7 +17,6 @@ litellm-circuit-breaker.workspace = true
litellm-redis-types.workspace = true
litellm-core-utils.workspace = true
litellm-http.workspace = true
rustls.workspace = true
litellm-redis-client.workspace = true
futures-util.workspace = true
serde.workspace = true

3
litellm-rust/crates/redis-connect/src/connect.rs
View file

@ -42,9 +42,6 @@ pub fn connect(
host: HostServices,
breaker: Option<Breaker>,
) -> Result<RedisClient, Error> {
if rustls::crypto::CryptoProvider::get_default().is_none() {
let _ = rustls::crypto::ring::default_provider().install_default();
}
tracing::info!(target = ?resolved.target, namespace = ?resolved.namespace, "redis client configured");
let credentials = credentials(resolved.auth, &host)?;
let events = host.events;

36
litellm-rust/crates/redis-connect/src/sources/gcp.rs
View file

@ -41,6 +41,15 @@ impl GcpIam {
}
}
fn token_url(service_account: &str) -> String {
match service_account.starts_with("projects/") {
true => format!("{IAM_CREDENTIALS}/{service_account}:generateAccessToken"),
false => format!(
"{IAM_CREDENTIALS}/projects/-/serviceAccounts/{service_account}:generateAccessToken"
),
}
}
impl CredentialSource for GcpIam {
fn fetch(&self) -> SourceFuture<'_> {
Box::pin(async move {
@ -52,10 +61,7 @@ impl CredentialSource for GcpIam {
.map_err(Error::credential(CredentialStage::Identity))?;
let response = self
.http
.post(format!(
"{IAM_CREDENTIALS}/{}:generateAccessToken",
self.service_account
))
.post(token_url(&self.service_account))
.bearer_auth(caller.expose())
.json(&serde_json::json!({ "scope": [CLOUD_PLATFORM_SCOPE] }))
.send()
@ -92,9 +98,29 @@ impl CredentialSource for GcpIam {
mod tests {
use rstest::rstest;
use super::GcpIam;
use super::{GcpIam, token_url};
use crate::CredentialSource;
#[rstest]
#[case::bare_email(
"redis@project.iam.gserviceaccount.com",
"https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/redis@project.iam.gserviceaccount.com:generateAccessToken"
)]
#[case::resource_name(
"projects/-/serviceAccounts/redis@project.iam.gserviceaccount.com",
"https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/redis@project.iam.gserviceaccount.com:generateAccessToken"
)]
#[case::project_scoped_resource_name(
"projects/project/serviceAccounts/redis@project.iam.gserviceaccount.com",
"https://iamcredentials.googleapis.com/v1/projects/project/serviceAccounts/redis@project.iam.gserviceaccount.com:generateAccessToken"
)]
fn a_bare_service_account_email_is_addressed_as_a_resource_name(
#[case] service_account: &str,
#[case] expected: &str,
) {
assert_eq!(token_url(service_account), expected);
}
#[rstest]
#[tokio::test]
async fn the_caller_identity_comes_from_application_default_credentials_not_vertex_settings() {

2
litellm-rust/crates/redis-store/AGENTS.md
View file

@ -1,5 +1,5 @@
- One folder per Redis family (`router_cooldowns/`, `router_usage/`, `spend_counters/`, `auth_cache/`, `rate_limits/`), each with `mod.rs` (the store and its ops), `keys.rs` (key types and spellings) and `types.rs` (the family's values, with the `FromValue` / `IntoValue` codecs for what it stores in Redis). It is the only place that family's keys, value codec and scripts are written, and it takes domain ids, never keys. Shared pieces live once at the crate root: `Error` and `ErrorKind` (`error.rs`), `Memory` and `Clock` (`memory.rs`), `Pending` (`request.rs`), `Remote` and its attach-once `RemoteSlot`, which knows nothing of memory (`remote.rs`), `DualCache` (`dual_cache/`) and the Python codecs (`codec.rs`, `value.rs`)
- An op declared on a `Send` returns `Pending<T>`, a lazy future built when the op is declared. An op that only means something on Redis (an `MGET` the caller falls back from, a script) returns `None` without one, and an op that replaced a `DualCache` call answers from memory. A post-call method returns `false` when there was no Redis to declare on
- An op declared on a `Send` returns `Pending<T>`, a lazy future built when the op is declared. An op that only means something on Redis (an `MGET` the caller falls back from, a script) returns `None` without one, and an op that replaced a `DualCache` call answers from memory. A post-call method returns `false` when there was no Redis to declare on, or `None` when it returns a settlement the caller waits on, like `SpendCounterStore::increment_many_post_call`
- A store that sends a pipeline of its own uses `RedisBatch::pipelined`, never `Pipeline` and raw `Value`s
- Python and Rust pods share these keys during a rolling deploy, so a module writes the bytes `main`'s Python writes and reads everything `main` reads. Python formats go through `litellm-python-compat` (`repr`, `literal_eval`, `json`), never a hand-rolled codec
- A family holds one `DualCache` per `DualCache` its Python held on `main`, and each `DualCache` method is the Python method it names, so a family op that was a `DualCache` call is one call here. Where `main`'s Python reached past `DualCache` to `in_memory_cache` or `redis_cache`, the family does the same through `memory()` and `redis()`, which keeps that reach visible. Caches a family's Python attached to one `RedisCache` share one `Arc<RemoteSlot>`. Only a family that batch-reads builds its cache with `DualCache::batch_reading` from a `TierPolicy` the host passes in, and the Python host fills it from what `main` reads from the environment or `litellm` globals (`litellm/rust_bridge/host/settings.py`, `router_tiers`). Callers pass a `Remote` (or none), never a cache. Redis-only families are plain functions over `&RedisClient`

16
litellm-rust/crates/redis-store/src/spend_counters/keys.rs
View file

@ -1,6 +1,9 @@
/// One budget's spend counter. `parse` and `key` round-trip `main`'s `spend:*` keys exactly.
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub enum SpendCounter {
/// The charges-only twin of a key, team, user, org or end user lifetime counter, which budget
/// holds never touch: `spend:committed:` and the twin's key past its `spend:`.
Committed(Box<SpendCounter>),
Key {
token: String,
window: Option<String>,
@ -33,6 +36,9 @@ impl SpendCounter {
None => format!("spend:{scope}:{id}"),
};
match self {
Self::Committed(counter) => {
format!("spend:committed:{}", &counter.key()["spend:".len()..])
}
Self::Key { token, window } => windowed("key", token, window),
Self::Team { team_id, window } => windowed("team", team_id, window),
Self::TeamMember { user_id, team_id } => {
@ -51,6 +57,16 @@ impl SpendCounter {
pub fn parse(key: &str) -> Result<Self, NotASpendCounter> {
let invalid = || NotASpendCounter(key.to_owned());
let rest = key.strip_prefix("spend:").ok_or_else(invalid)?;
if let Some(twin) = rest.strip_prefix("committed:") {
return match Self::parse(&format!("spend:{twin}")).map_err(|_| invalid())? {
counter @ (Self::Key { window: None, .. }
| Self::Team { window: None, .. }
| Self::User(_)
| Self::Org(_)
| Self::EndUser(_)) => Ok(Self::Committed(Box::new(counter))),
_ => Err(invalid()),
};
}
let (scope, id) = rest.split_once(':').ok_or_else(invalid)?;
let windowed = || match id.rsplit_once(":window:") {
Some((id, window)) if !window.is_empty() => (id.to_owned(), Some(window.to_owned())),

72
litellm-rust/crates/redis-store/src/spend_counters/mod.rs
View file

@ -1,13 +1,19 @@
mod keys;
mod types;
use std::{sync::Arc, time::Duration};
use std::{
sync::{
Arc,
atomic::{AtomicBool, AtomicU64, Ordering},
},
time::Duration,
};
use litellm_python_compat::repr::to_str;
use litellm_redis_client::{Declared, Expiry, RedisBatch, commands};
pub use self::keys::{NotASpendCounter, SpendCounter};
pub use self::types::{LostSeed, Reading};
pub use self::types::{LostSeed, PostCallSettlement, Reading};
use crate::{
Error, Number, Pending, RequestRedis, Send,
dual_cache::{DualCache, MemorySettings},
@ -31,6 +37,7 @@ const SET_MAX: &str = concat!(
/// `DualCache` to its Redis and its memory.
pub struct SpendCounters {
cache: DualCache<Number>,
committed_ttl_millis: AtomicU64,
}
impl SpendCounters {
@ -45,6 +52,33 @@ impl SpendCounters {
MemorySettings::default(),
now,
),
committed_ttl_millis: AtomicU64::new(0),
}
}
/// The least TTL a Redis write gives a [`SpendCounter::Committed`] twin, so it outlives the DB
/// write that holds its charges and its next seed never reads a total still missing them.
pub fn set_committed_ttl(&self, ttl: Duration) {
let millis = u64::try_from(ttl.as_millis()).unwrap_or(u64::MAX);
self.committed_ttl_millis.store(millis, Ordering::Relaxed);
}
fn ttl(&self, counter: &SpendCounter, default_ttl: Duration) -> Duration {
match counter {
SpendCounter::Committed(_) => default_ttl.max(Duration::from_millis(
self.committed_ttl_millis.load(Ordering::Relaxed),
)),
_ => default_ttl,
}
}
/// A committed twin's expiry moves out on every charge, on a cluster too, so it never expires
/// while a charge it holds is still on its way to the DB.
fn post_call_expiry(&self, counter: &SpendCounter, default_ttl: Duration) -> Expiry {
let ttl = self.ttl(counter, default_ttl);
match counter {
SpendCounter::Committed(_) => Expiry::Refresh(ttl),
_ => Expiry::KeepExisting(ttl),
}
}
@ -127,7 +161,7 @@ impl SpendCounters {
let increment = commands::incr_by_float(
&key,
amount.into(),
Some(Expiry::Refresh(remote.default_ttl)),
Some(Expiry::Refresh(self.ttl(counter, remote.default_ttl))),
);
match remote.client.send(increment).await {
Ok(total) => {
@ -159,7 +193,7 @@ impl SpendCounters {
pipeline.add(commands::incr_by_float(
&counter.key(),
(*amount).into(),
Some(Expiry::Refresh(remote.default_ttl)),
Some(Expiry::Refresh(self.ttl(counter, remote.default_ttl))),
))
});
let totals = Declared::all(totals.collect::<Vec<_>>()).await;
@ -175,15 +209,14 @@ impl SpendCounters {
}
/// `increment_many` on the request's post-call batch. Each total reaches memory once the batch
/// goes out, and a counter whose increment failed is dropped everywhere. `false` without Redis.
/// goes out, and a counter whose increment failed is dropped everywhere. `None` without Redis,
/// else a settlement that reports when every reply is back, whoever flushed the batch.
pub fn increment_many_post_call(
self: &Arc<Self>,
request: &RequestRedis,
increments: &[(SpendCounter, Number)],
) -> bool {
let Some(remote) = self.cache.redis().get() else {
return false;
};
) -> Option<PostCallSettlement> {
let remote = self.cache.redis().get()?;
let batch = Send::PostCall(request).batch(&remote.client);
let declared = increments
.iter()
@ -191,12 +224,14 @@ impl SpendCounters {
let total = batch.add(commands::incr_by_float(
&counter.key(),
(*amount).into(),
Some(Expiry::KeepExisting(remote.default_ttl)),
Some(self.post_call_expiry(counter, remote.default_ttl)),
));
(counter.clone(), total)
})
.collect::<Vec<_>>();
let store = Arc::clone(self);
let settled = Arc::new(AtomicBool::new(false));
let settlement = PostCallSettlement(Arc::clone(&settled));
request.spawn(async move {
for (counter, total) in declared {
match total.settled().await {
@ -215,8 +250,9 @@ impl SpendCounters {
}
}
}
settled.store(true, Ordering::Release);
});
true
Some(settlement)
}
/// `_repair_stale_spend_counter`: raise memory and Redis to `spend`, never lowering either.
@ -243,7 +279,7 @@ impl SpendCounters {
};
let args = (
to_str(&spend.to_value()),
remote.default_ttl.as_secs().to_string(),
self.ttl(counter, remote.default_ttl).as_secs().to_string(),
);
let stored: Option<f64> = remote
.client
@ -276,8 +312,9 @@ impl SpendCounters {
})
})?);
};
let ttl = self.ttl(counter, *default_ttl);
let current = if client
.send(commands::set_nx(&key, spend.to_json(), *default_ttl))
.send(commands::set_nx(&key, spend.to_json(), ttl))
.await?
{
spend
@ -290,7 +327,7 @@ impl SpendCounters {
let increment = commands::incr_by_float(
&key,
spend.into(),
Some(Expiry::KeepExisting(*default_ttl)),
Some(Expiry::KeepExisting(ttl)),
);
Number::Float(client.send(increment).await?)
}
@ -313,7 +350,7 @@ impl SpendCounters {
let Some(remote) = self.cache.redis().get() else {
return Ok(());
};
let ttl = ttl.unwrap_or(remote.default_ttl);
let ttl = ttl.unwrap_or_else(|| self.ttl(counter, remote.default_ttl));
remote
.client
.send(commands::set(&key, spend.to_json(), ttl))
@ -334,7 +371,10 @@ impl SpendCounters {
};
Ok(remote
.client
.send(commands::expire(&counter.key(), remote.default_ttl))
.send(commands::expire(
&counter.key(),
self.ttl(counter, remote.default_ttl),
))
.await?)
}

16
litellm-rust/crates/redis-store/src/spend_counters/types.rs
View file

@ -1,3 +1,8 @@
use std::sync::{
Arc,
atomic::{AtomicBool, Ordering},
};
use litellm_python_compat::Value;
use super::keys::SpendCounter;
@ -32,6 +37,17 @@ impl Reading {
}
}
/// One `increment_many_post_call`: settled once every increment it declared has its reply and
/// memory mirrors it, whether `flush_post` or the post-call deadline sent the batch.
#[derive(Clone, Debug)]
pub struct PostCallSettlement(pub(super) Arc<AtomicBool>);
impl PostCallSettlement {
pub fn is_settled(&self) -> bool {
self.0.load(Ordering::Acquire)
}
}
/// Who wins when a seed's `SET NX` lost and the winner's value is gone by the time it is read.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum LostSeed {

165
litellm-rust/crates/redis-store/tests/spend_counters.rs
View file

@ -16,6 +16,7 @@ use support::{COOLDOWN_POLICY, connect, connect_observed, remote, within};
use tokio::runtime::Handle;
const TTL: Duration = Duration::from_secs(60);
const COMMITTED_TTL: Duration = Duration::from_secs(300);
const LONG_DEADLINE: Duration = Duration::from_secs(60);
const READS: [(&str, Number); 4] = [
@ -45,6 +46,10 @@ fn team(id: &str) -> SpendCounter {
}
}
fn committed(counter: SpendCounter) -> SpendCounter {
SpendCounter::Committed(Box::new(counter))
}
fn store(client: Option<RedisClient>) -> Arc<SpendCounters> {
Arc::new(SpendCounters::new(client.map(remote)))
}
@ -114,6 +119,26 @@ fn a_reading_names_where_it_was_answered(
"spend:model_access_group:gpt-group"
)]
#[case(SpendCounter::ShadowEval("job-7".to_owned()), "spend:shadow_eval:job-7")]
#[case::committed_key(
committed(SpendCounter::Key {
token: "88dc28d0".to_owned(),
window: None,
}),
"spend:committed:key:88dc28d0"
)]
#[case::committed_team(committed(team("team-a")), "spend:committed:team:team-a")]
#[case::committed_user(
committed(SpendCounter::User("user@example.com".to_owned())),
"spend:committed:user:user@example.com"
)]
#[case::committed_org(
committed(SpendCounter::Org("org-1".to_owned())),
"spend:committed:org:org-1"
)]
#[case::committed_end_user(
committed(SpendCounter::EndUser("customer é 1".to_owned())),
"spend:committed:end_user:customer é 1"
)]
fn every_counter_formats_and_parses_back_to_mains_key(
#[case] counter: SpendCounter,
#[case] key: &str,
@ -128,6 +153,11 @@ fn every_counter_formats_and_parses_back_to_mains_key(
#[case("spend:unknown:abc")]
#[case("spend:team_member:no-team")]
#[case("spend:")]
#[case::committed_window("spend:committed:key:abc:window:1d")]
#[case::committed_tag("spend:committed:tag:prod")]
#[case::committed_team_member("spend:committed:team_member:user:team-a")]
#[case::committed_twice("spend:committed:committed:key:abc")]
#[case::committed_nothing("spend:committed:")]
fn keys_outside_the_spend_counter_keyspace_are_refused(#[case] key: &str) {
assert!(SpendCounter::parse(key).is_err(), "{key}");
}
@ -307,7 +337,7 @@ async fn a_post_call_increment_from_a_thread_without_a_runtime_lands_in_redis_an
.expect("a Python thread has no Tokio context, and declaring there must not panic");
request.flush_post().await;
assert!(declared);
assert!(declared.is_some());
let mirrored =
litellm_redis_testing::eventually(async || store.get_local(&counter).unwrap()).await;
assert_eq!(mirrored, Number::Float(3.5));
@ -331,9 +361,12 @@ async fn post_call_increments_wait_for_the_post_batch_then_land_in_memory(
let request = RequestRedis::new(Handle::current(), LONG_DEADLINE);
let counter = team("t");
assert!(store.increment_many_post_call(&request, &[(counter.clone(), Number::Float(3.5))]));
let settlement = store
.increment_many_post_call(&request, &[(counter.clone(), Number::Float(3.5))])
.unwrap();
request.flush_pre().await;
assert!(!settlement.is_settled());
assert!(trips.seen().is_empty());
assert_eq!(
raw.send(commands::get::<String>(&counter.key()))
@ -345,6 +378,7 @@ async fn post_call_increments_wait_for_the_post_batch_then_land_in_memory(
request.flush_post().await;
tokio::time::sleep(Duration::from_millis(100)).await;
assert!(settlement.is_settled());
let increment = match topology {
Topology::Standalone => Operation::Pipeline { ops: 1 },
Topology::Cluster => Operation::Command("INCRBYFLOAT"),
@ -557,13 +591,132 @@ async fn a_failed_post_call_increment_drops_the_counter_from_memory() {
let request = RequestRedis::new(Handle::current(), LONG_DEADLINE);
let counter = team("t");
store.increment(&counter, Number::Float(1.0)).await.unwrap();
assert!(store.increment_many_post_call(&request, &[(counter.clone(), Number::Float(1.0))]));
let settlement = store
.increment_many_post_call(&request, &[(counter.clone(), Number::Float(1.0))])
.unwrap();
deployment.stop().await.unwrap();
request.flush_post().await;
tokio::time::sleep(Duration::from_millis(200)).await;
assert_eq!(store.get_local(&counter).unwrap(), None);
assert!(settlement.is_settled());
}
#[rstest]
#[tokio::test]
async fn a_post_call_settlement_settles_when_the_deadline_sends_the_batch(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let store = store(Some(connect(&deployment)));
let request = RequestRedis::new(Handle::current(), Duration::from_millis(50));
let counter = committed(team("t"));
let settlement = store
.increment_many_post_call(&request, &[(counter.clone(), Number::Float(2.5))])
.unwrap();
assert!(!settlement.is_settled());
litellm_redis_testing::eventually(async || settlement.is_settled().then_some(())).await;
assert_eq!(
raw.send(commands::get::<String>(&counter.key()))
.await
.unwrap(),
Some("2.5".to_owned())
);
}
#[rstest]
#[case::committed_twin(committed, 300)]
#[case::enforcement_counter(std::convert::identity, 60)]
#[tokio::test]
async fn every_write_gives_a_committed_twin_the_committed_ttl_and_other_counters_the_redis_ttl(
#[case] counter_of: fn(SpendCounter) -> SpendCounter,
#[case] ttl_seconds: u64,
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let store = store(Some(connect(&deployment)));
store.set_committed_ttl(COMMITTED_TTL);
let request = RequestRedis::new(Handle::current(), LONG_DEADLINE);
let [
incremented,
pipelined,
post_call,
seeded,
repaired,
refreshed,
] = ["a", "b", "c", "d", "e", "f"].map(|id| counter_of(team(id)));
raw.send(commands::set(
&refreshed.key(),
"1.0",
Duration::from_secs(5),
))
.await
.unwrap();
store
.increment(&incremented, Number::Float(1.0))
.await
.unwrap();
store
.increment_many(&[(pipelined.clone(), Number::Float(1.0))])
.await
.unwrap();
let settlement = store
.increment_many_post_call(&request, &[(post_call.clone(), Number::Float(1.0))])
.unwrap();
request.flush_post().await;
store
.seed(&seeded, Number::Float(1.0), LostSeed::Answer)
.await
.unwrap();
store.set_max(&repaired, Number::Float(1.0)).await.unwrap();
assert!(store.refresh_ttl(&refreshed).await.unwrap());
assert!(settlement.is_settled());
let remaining =
async |counter: &SpendCounter| match raw.send(commands::ttl(&counter.key())).await.unwrap()
{
Ttl::Expires(remaining) => remaining,
other => panic!("{} has no TTL: {other:?}", counter.key()),
};
assert!(within(remaining(&incremented).await, ttl_seconds));
assert!(within(remaining(&pipelined).await, ttl_seconds));
assert!(within(remaining(&post_call).await, ttl_seconds));
assert!(within(remaining(&seeded).await, ttl_seconds));
assert!(within(remaining(&repaired).await, ttl_seconds));
assert!(within(remaining(&refreshed).await, ttl_seconds));
}
#[rstest]
#[tokio::test]
async fn a_post_call_charge_moves_a_committed_twins_expiry_out_to_the_committed_ttl(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let store = store(Some(connect(&deployment)));
store.set_committed_ttl(COMMITTED_TTL);
let request = RequestRedis::new(Handle::current(), LONG_DEADLINE);
let counter = committed(team("t"));
raw.send(commands::set(&counter.key(), "1.0", Duration::from_secs(5)))
.await
.unwrap();
let settlement = store
.increment_many_post_call(&request, &[(counter.clone(), Number::Float(1.0))])
.unwrap();
request.flush_post().await;
assert!(settlement.is_settled());
match raw.send(commands::ttl(&counter.key())).await.unwrap() {
Ttl::Expires(remaining) => assert!(within(remaining, COMMITTED_TTL.as_secs())),
other => panic!("{} has no TTL: {other:?}", counter.key()),
}
}
#[tokio::test]
@ -592,7 +745,11 @@ async fn without_redis_counters_live_in_memory() {
store.read(&counter).await.unwrap(),
Reading::Memory(Some(Number::Float(7.0)))
);
assert!(!store.increment_many_post_call(&request, &[(counter.clone(), Number::Int(1))]));
assert!(
store
.increment_many_post_call(&request, &[(counter.clone(), Number::Int(1))])
.is_none()
);
assert!(
store
.read_many(RedisSend::PreCall(&request), &[counter])

17
litellm/caching/redis_batch.py
View file

@ -452,6 +452,13 @@ class NativePending(Protocol[_T_co]):
def resolve(self) -> Awaitable[_T_co]: ...
class PostCallSettlement(Protocol):
"""Post-call writes declared together: settled once every reply is back, whoever sent their batch."""
@property
def settled(self) -> bool: ...
class NativeBatchResult(Generic[_T]):
"""Awaitable handle for a read declared on the request's native pre-call batch; awaiting it runs the
native flush hooks, flushes that batch, then answers from what it read."""
@ -526,6 +533,7 @@ class RequestRedisBatches:
"_native_misses",
"_native_send_factory",
"_post_call",
"_post_call_settlements",
"post_call_deadline",
"prefetched",
)
@ -538,6 +546,7 @@ class RequestRedisBatches:
) -> None:
self._batches: Final[dict[object, RedisBatch]] = {} # mutable-ok: lazily filled per backend
self._post_call: Final[dict[object, RedisBatch]] = {} # mutable-ok: lazily filled per backend
self._post_call_settlements: tuple[PostCallSettlement, ...] = ()
self.post_call_deadline: Final = post_call_deadline
self._native_factory: Final = native_factory
self._native_send_factory: Final = native_send_factory
@ -641,6 +650,14 @@ class RequestRedisBatches:
if not any(batch.pending for batch in self._post_call.values()):
_open_post_call.discard(self)
def track_post_call_settlement(self, settlement: PostCallSettlement) -> None:
"""Post-call writes a later reader in this request must know have landed, such as its committed spend."""
self._post_call_settlements = (*self._post_call_settlements, settlement)
@property
def post_call_settlements(self) -> tuple[PostCallSettlement, ...]:
return self._post_call_settlements
@property
def batches(self) -> tuple[RedisBatch, ...]:
return tuple(self._batches.values())

3
litellm/chat_completions/dispatch.py
View file

@ -16,6 +16,7 @@ from litellm.rust_bridge.public_call import (
AsyncFlag,
NativeCall,
bind,
native_acall_hook,
native_call,
native_call_hook,
optional_sequence,
@ -131,7 +132,7 @@ async def acompletion(*args: object, **kwargs: object) -> ChatResult: # kwargs-
_FLAG.dropped(kwargs),
python=python,
binding=NATIVE_ACOMPLETION,
native=native_call_hook,
native=native_acall_hook,
)

3
litellm/constants.py
View file

@ -1533,6 +1533,7 @@ BEDROCK_CONVERSE_MODELS: Final = [
"openai.gpt-oss-120b-1:0",
"anthropic.claude-haiku-4-5-20251001-v1:0",
"anthropic.claude-sonnet-4-5-20250929-v1:0",
"anthropic.claude-fable-5-2",
"anthropic.claude-fable-5-1",
"anthropic.claude-fable-5",
"anthropic.claude-sonnet-5",
@ -1787,6 +1788,7 @@ AZURE_SPEECH_PASS_THROUGH_ROUTE_PREFIX: Final = "/azure_speech"
AZURE_SPEECH_SHORT_AUDIO_PATH_PREFIX: Final = "/speech/"
AZURE_SPEECH_BATCH_PATH_PREFIX: Final = "/speechtotext/"
AZURE_SPEECH_FAST_TRANSCRIPTION_PATH: Final = "/speechtotext/transcriptions:transcribe"
OPENSEARCH_PASS_THROUGH_ROUTE_PREFIX: Final = "/opensearch"
AZURE_SPEECH_STT_DOMAIN: Final = "stt.speech.microsoft.com"
AZURE_SPEECH_COGNITIVE_SERVICES_DOMAIN: Final = "api.cognitive.microsoft.com"
AZURE_SPEECH_SUBSCRIPTION_KEY_HEADER: Final = "Ocp-Apim-Subscription-Key"
@ -1940,6 +1942,7 @@ SPEND_ROLLUP_LOCK_TIMEOUT_MS: Final = max(1, get_env_int("SPEND_ROLLUP_LOCK_TIME
SPEND_LOG_QUEUE_SIZE_THRESHOLD: Final = get_env_int("SPEND_LOG_QUEUE_SIZE_THRESHOLD", 100)
SPEND_LOG_QUEUE_MAX_BYTES: Final = max(1, get_env_int("SPEND_LOG_QUEUE_MAX_BYTES", 64000000))
SPEND_LOG_QUEUE_POLL_INTERVAL: Final = get_env_float("SPEND_LOG_QUEUE_POLL_INTERVAL", 2.0)
SPEND_LOG_QUEUE_MAX_BACKOFF_SECONDS: Final = 30.0
RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS: Final = max(1, get_env_int("RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS", 3))
RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL: Final = get_env_float("RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL", 0.2)
SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE: Final = get_env_int("SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE", 10000)

40
litellm/cost_calculator.py
View file

@ -103,9 +103,6 @@ from litellm.llms.openai.cost_calculation import (
from litellm.llms.openai.cost_calculation import (
video_output_cost_per_second,
)
from litellm.llms.openai.transcriptions.gpt_transformation import (
OpenAIGPTAudioTranscriptionConfig,
)
from litellm.llms.perplexity.cost_calculator import (
cost_per_token as perplexity_cost_per_token,
)
@ -381,38 +378,15 @@ def _transcription_uses_token_pricing(
return (prompt_tokens_val > 0) or (completion_tokens_val > 0)
def _bills_transcription_usage_duration(pricing_model: str | None, custom_llm_provider: str | None) -> bool:
if pricing_model is None or custom_llm_provider is None:
return False
transcription_info: Final = _lookup_model_info_or_none(pricing_model, custom_llm_provider)
transcription_metadata: Final = (
as_str_mapping(transcription_info.get("provider_specific_entry")) if transcription_info is not None else None
)
return transcription_metadata is not None and transcription_metadata.get("transcription_usage_duration") == 1
def _transcription_duration_for_cost(
completion_response: object, pricing_model: str | None, model: str | None, custom_llm_provider: str | None
) -> float:
def _transcription_duration_for_cost(completion_response: object) -> float:
billed_duration: Final = _get_transcription_usage_duration(completion_response)
if billed_duration is not None and billed_duration > 0:
return billed_duration
hidden_params: Final = as_str_mapping(attribute_or_none(completion_response, "_hidden_params")) or {}
measured_duration: Final = normalized_audio_duration_seconds(
hidden_params.get("audio_transcription_duration", attribute_or_none(completion_response, "duration"))
)
if measured_duration is not None and measured_duration > 0:
return measured_duration
provider_duration: Final = _get_transcription_usage_duration(completion_response)
if (
provider_duration is not None
and provider_duration > 0
and any(_bills_transcription_usage_duration(name, custom_llm_provider) for name in (pricing_model, model))
):
return provider_duration
if model is None or custom_llm_provider != LlmProviders.OPENAI.value:
return 0.0
transcription_config: Final = ProviderConfigManager.get_provider_audio_transcription_config(
model=model.removeprefix(f"{LlmProviders.OPENAI.value}/"), provider=LlmProviders.OPENAI
)
return (provider_duration or 0.0) if isinstance(transcription_config, OpenAIGPTAudioTranscriptionConfig) else 0.0
return measured_duration if measured_duration is not None and measured_duration > 0 else 0.0
OCRPricingField = Literal["ocr_cost_per_page", "ocr_cost_per_credit", "annotation_cost_per_page"]
@ -2132,9 +2106,7 @@ def completion_cost(
prompt,
)
elif call_type in _TRANSCRIPTION_CALL_TYPES:
audio_transcription_file_duration = _transcription_duration_for_cost(
completion_response, model, requested_model, custom_llm_provider
)
audio_transcription_file_duration = _transcription_duration_for_cost(completion_response)
elif call_type in _RERANK_CALL_TYPES:
if completion_response is not None and isinstance(completion_response, RerankResponse):
meta_obj = completion_response.meta

3
litellm/embeddings/dispatch.py
View file

@ -13,6 +13,7 @@ from litellm.rust_bridge.public_call import (
AsyncFlag,
NativeCall,
bind,
native_acall_hook,
native_call,
native_call_hook,
optional_str,
@ -89,7 +90,7 @@ async def aembedding(*args: object, **kwargs: object) -> EmbeddingResponse: # k
_FLAG.dropped(kwargs),
python=_PYTHON_AEMBEDDING,
binding=NATIVE_AEMBEDDING,
native=native_call_hook,
native=native_acall_hook,
)

238
litellm/integrations/prometheus.py
View file

@ -21,6 +21,7 @@ from typing_extensions import ReadOnly, TypedDict, assert_never
import litellm
from litellm._internal_context import with_service_target
from litellm._logging import print_verbose, verbose_logger
from litellm.caching.redis_batch import active_request_redis_batches
from litellm.constants import (
AUTH_OBJECTS_TARGET,
LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL,
@ -75,8 +76,10 @@ from litellm.types.integrations.prometheus import (
sanitize_prometheus_label_value,
validate_prometheus_deployment_and_latency_caller_identity,
)
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth, committed_budget_spend_of_row
from litellm.types.proxy.carried_budget_state import (
CommittedBudgetSpend,
EntitySpendCounterKeys,
KeyBudgetSnapshot,
OrgBudgetSnapshot,
TeamBudgetSnapshot,
@ -281,6 +284,11 @@ def _logged_fields(value: object) -> Mapping[str, object]:
return _NO_LOGGED_FIELDS
def _response_cost_of(logged: Mapping[str, object]) -> float:
cost: Final = logged.get("response_cost")
return float(cost) if isinstance(cost, (int, float)) else 0.0
_GOOGLE_PASS_THROUGH_PROVIDERS: Final = frozenset(
(litellm.LlmProviders.GEMINI.value, litellm.LlmProviders.VERTEX_AI.value)
)
@ -365,6 +373,104 @@ def _get_budget_metrics_per_request_timeout() -> float:
return parsed
def _stored_spend(value: object) -> float:
return float(value) if isinstance(value, (int, float)) else 0.0
_OWN_CHARGE_MATCH_RATIO: Final = 1 - 1e-9
def _committed_floor(stored: float, committed: float | None) -> float:
return stored if committed is None else max(stored, committed)
def _customer_spend_after_request(cached_spend: float, committed_spend: float | None, response_cost: float) -> float:
if committed_spend is None:
return cached_spend + response_cost
return max(committed_spend + response_cost, cached_spend)
def _budgeted_org_counter(org_id: str | None, org: OrgBudgetSnapshot | None) -> EntitySpendCounterKeys:
if org is None or org.max_budget is None:
return EntitySpendCounterKeys()
return EntitySpendCounterKeys.of(token=None, team_id=None, user_id=None, org_id=org_id, end_user_id=None)
async def _committed_spend_now(counter_keys: EntitySpendCounterKeys, timeout: float) -> CommittedBudgetSpend:
if not counter_keys.tracked():
return CommittedBudgetSpend()
try:
from litellm.proxy.proxy_server import spend_counter_store
from litellm.proxy.spend_tracking.committed_spend import read_committed_budget_spend
return await asyncio.wait_for(
read_committed_budget_spend(spend_counter_store, counter_keys),
timeout=timeout,
)
except Exception as e: # noqa: BLE001 # a gauge read never fails the request; the gauges keep the stored spend
verbose_logger.debug("[Non-Blocking] Prometheus: committed spend read failed: %s", e)
return CommittedBudgetSpend()
OwnCharge: TypeAlias = Literal["unknown", "pending", "counted"]
def _own_charge() -> OwnCharge:
"""Whether this request's committed spend charge is in Redis yet, known only inside a request with Redis."""
try:
from litellm.proxy.proxy_server import spend_counter_store
except ImportError:
return "unknown"
batches: Final = active_request_redis_batches()
if batches is None or not spend_counter_store.has_redis:
return "unknown"
settlements: Final = batches.post_call_settlements
return "counted" if settlements and all(settlement.settled for settlement in settlements) else "pending"
def _committed_spend_before_request(
at_auth: CommittedBudgetSpend, now: CommittedBudgetSpend, response_cost: float, own_charge: OwnCharge
) -> CommittedBudgetSpend:
def before(at_auth_spend: float | None, now_spend: float | None) -> float | None:
if now_spend is None:
return at_auth_spend
match own_charge:
case "pending":
return now_spend
case "counted":
without_own_charge: Final = now_spend - response_cost
return without_own_charge if at_auth_spend is None else max(at_auth_spend, without_own_charge)
case "unknown":
if at_auth_spend is None:
return now_spend - response_cost
own_charge_counted: Final = now_spend - at_auth_spend >= response_cost * _OWN_CHARGE_MATCH_RATIO
return max(at_auth_spend, now_spend - response_cost) if own_charge_counted else now_spend
case _:
assert_never(own_charge)
return CommittedBudgetSpend(
key=before(at_auth.key, now.key),
team=before(at_auth.team, now.team),
user=before(at_auth.user, now.user),
org=before(at_auth.org, now.org),
end_user=before(at_auth.end_user, now.end_user),
)
async def _committed_spend_without_own_charge(
counter_keys: EntitySpendCounterKeys,
at_auth: CommittedBudgetSpend | None,
response_cost: float,
read_timeout: float,
) -> CommittedBudgetSpend:
own_charge_before_read: Final = _own_charge()
now: Final = await _committed_spend_now(counter_keys, read_timeout)
own_charge: Final = own_charge_before_read if _own_charge() == own_charge_before_read else "unknown"
return _committed_spend_before_request(
at_auth=now if at_auth is None else at_auth, now=now, response_cost=response_cost, own_charge=own_charge
)
def _get_proxy_llm_router() -> Router | None:
try:
from litellm.proxy.proxy_server import llm_router
@ -2229,6 +2335,10 @@ class PrometheusLogger(CustomLogger):
_user_max_budget: Final = _metadata.get("user_api_key_user_max_budget", None)
_user_email: Final = _metadata.get("user_api_key_user_email", None)
_org_alias: Final = _metadata.get("user_api_key_org_alias", None)
_carried_metadata: Final = _REQUEST_METADATA_ADAPTER.validate_python(_metadata)
_committed_at_auth: Final = committed_budget_spend_of_row(
_carried_metadata, _REQUEST_METADATA_ADAPTER.validate_python(litellm_params)
)
# Bound the per-request budget-metric emission so that slow Redis/DB
# lookups under load cannot consume the whole LoggingWorker watchdog
@ -2238,45 +2348,65 @@ class PrometheusLogger(CustomLogger):
# so dropping one slow per-request emission only loses sub-cron real-time
# detail, not correctness.
budget_metrics_timeout: Final = _get_budget_metrics_per_request_timeout()
gather_coro: Final = asyncio.gather(
self._set_api_key_budget_metrics_after_api_request(
user_api_key=user_api_key,
user_api_key_alias=user_api_key_alias,
committed_read_timeout: Final = budget_metrics_timeout / 2
async def emit_budget_metrics() -> tuple[object, ...]:
committed: Final = await _committed_spend_without_own_charge(
counter_keys=_committed_at_auth.counter_keys
if _committed_at_auth is not None
else EntitySpendCounterKeys.of(
token=user_api_key,
team_id=user_api_team,
user_id=user_id,
org_id=user_api_key_org_id,
end_user_id=end_user_id,
),
at_auth=_committed_at_auth,
response_cost=response_cost,
key_max_budget=_api_key_max_budget,
key_spend=_api_key_spend,
carried=KeyBudgetSnapshot.from_metadata(_metadata),
),
self._set_team_budget_metrics_after_api_request(
user_api_team=user_api_team,
user_api_team_alias=user_api_team_alias,
team_spend=_team_spend,
team_max_budget=_team_max_budget,
response_cost=response_cost,
carried=TeamBudgetSnapshot.from_metadata(_metadata),
),
self._set_user_budget_metrics_after_api_request(
user_id=user_id,
user_spend=_user_spend,
user_max_budget=_user_max_budget,
response_cost=response_cost,
carried=UserBudgetSnapshot.from_metadata(_metadata),
user_email=_user_email if isinstance(_user_email, str) else None,
),
self._set_org_budget_metrics_after_api_request(
org_id=user_api_key_org_id,
response_cost=response_cost,
carried=OrgBudgetSnapshot.from_metadata(_metadata),
org_alias=_org_alias if isinstance(_org_alias, str) else None,
),
self._set_customer_budget_metrics_after_api_request(
end_user_id=end_user_id,
response_cost=response_cost,
),
return_exceptions=True,
)
read_timeout=committed_read_timeout,
)
return await asyncio.gather(
self._set_api_key_budget_metrics_after_api_request(
user_api_key=user_api_key,
user_api_key_alias=user_api_key_alias,
response_cost=response_cost,
key_max_budget=_api_key_max_budget,
key_spend=_committed_floor(_stored_spend(_api_key_spend), committed.key),
carried=KeyBudgetSnapshot.from_metadata(_carried_metadata),
),
self._set_team_budget_metrics_after_api_request(
user_api_team=user_api_team,
user_api_team_alias=user_api_team_alias,
team_spend=_committed_floor(_stored_spend(_team_spend), committed.team),
team_max_budget=_team_max_budget,
response_cost=response_cost,
carried=TeamBudgetSnapshot.from_metadata(_carried_metadata),
),
self._set_user_budget_metrics_after_api_request(
user_id=user_id,
user_spend=_committed_floor(_stored_spend(_user_spend), committed.user),
user_max_budget=_user_max_budget,
response_cost=response_cost,
carried=UserBudgetSnapshot.from_metadata(_carried_metadata),
user_email=_user_email if isinstance(_user_email, str) else None,
),
self._set_org_budget_metrics_after_api_request(
org_id=user_api_key_org_id,
response_cost=response_cost,
carried=OrgBudgetSnapshot.from_metadata(_carried_metadata),
org_alias=_org_alias if isinstance(_org_alias, str) else None,
committed_spend=committed.org,
),
self._set_customer_budget_metrics_after_api_request(
end_user_id=end_user_id,
response_cost=response_cost,
committed_spend=committed.end_user,
),
return_exceptions=True,
)
try:
results: Final = await asyncio.wait_for(gather_coro, timeout=budget_metrics_timeout)
results: Final = await asyncio.wait_for(emit_budget_metrics(), timeout=budget_metrics_timeout)
except asyncio.TimeoutError:
verbose_logger.debug(
"[Non-Blocking] Prometheus: per-request budget metric emission "
@ -2745,7 +2875,7 @@ class PrometheusLogger(CustomLogger):
_labels,
)
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time) -> None:
verbose_logger.debug(
"prometheus Logging - Enters failure logging function (kwargs keys: %s)",
list(kwargs.keys()) if isinstance(kwargs, dict) else type(kwargs).__name__,
@ -2799,9 +2929,27 @@ class PrometheusLogger(CustomLogger):
label_context=PrometheusLabelFactoryContext(enum_values),
)
self.set_llm_deployment_failure_metrics(kwargs)
failure_metadata: Final = _REQUEST_METADATA_ADAPTER.validate_python(
get_litellm_metadata_from_kwargs(kwargs)
)
failure_org_alias: Final = failure_metadata.get("user_api_key_org_alias")
failure_org: Final = OrgBudgetSnapshot.from_metadata(failure_metadata)
failure_cost: Final = _response_cost_of(_LOGGED_FIELDS.validate_python(standard_logging_payload))
failure_committed: Final = await _committed_spend_without_own_charge(
counter_keys=_budgeted_org_counter(user_api_key_org_id, failure_org),
at_auth=committed_budget_spend_of_row(
failure_metadata, _REQUEST_METADATA_ADAPTER.validate_python(litellm_params)
)
or CommittedBudgetSpend(),
response_cost=failure_cost,
read_timeout=_get_budget_metrics_per_request_timeout(),
)
await self._set_org_budget_metrics_after_api_request(
org_id=user_api_key_org_id,
response_cost=0,
response_cost=failure_cost,
carried=failure_org,
org_alias=failure_org_alias if isinstance(failure_org_alias, str) else None,
committed_spend=failure_committed.org,
)
except Exception as e:
verbose_logger.exception("prometheus Layer Error(): Exception occured - %s", e)
@ -4516,6 +4664,7 @@ class PrometheusLogger(CustomLogger):
response_cost: float,
carried: OrgBudgetSnapshot | None = None,
org_alias: str | None = None,
committed_spend: float | None = None,
):
"""
Set org budget metrics after an LLM API request
@ -4533,7 +4682,7 @@ class PrometheusLogger(CustomLogger):
self._set_org_budget_metrics(
org_id=org_id,
org_alias=org_alias or "",
spend=carried.spend + response_cost,
spend=_committed_floor(carried.spend, committed_spend) + response_cost,
max_budget=carried.max_budget,
budget_reset_at=None,
)
@ -4559,7 +4708,7 @@ class PrometheusLogger(CustomLogger):
if org_info is None:
return
_total_org_spend: Final = (org_info.spend or 0.0) + response_cost
_total_org_spend: Final = _committed_floor(org_info.spend or 0.0, committed_spend) + response_cost
budget_table: Final = org_info.litellm_budget_table
max_budget: Final = budget_table.max_budget if budget_table else None
budget_reset_at: Final = getattr(budget_table, "budget_reset_at", None) if budget_table else None
@ -4626,6 +4775,7 @@ class PrometheusLogger(CustomLogger):
self,
end_user_id: str | None,
response_cost: float,
committed_spend: float | None = None,
):
if self._customer_budget_gauges_are_noop() or not _customer_budget_metrics_enabled():
return
@ -4651,7 +4801,9 @@ class PrometheusLogger(CustomLogger):
budget_table: Final = cached_customer.litellm_budget_table
self._set_customer_budget_metrics(
end_user_id=end_user_id,
spend=cached_customer.spend + response_cost,
spend=_customer_spend_after_request(
cached_spend=cached_customer.spend, committed_spend=committed_spend, response_cost=response_cost
),
max_budget=budget_table.max_budget if budget_table is not None else None,
budget_reset_at=None,
)

9
litellm/litellm_core_utils/audio_utils/utils.py
View file

@ -301,9 +301,12 @@ def calculate_request_duration(file: FileTypes) -> float | None:
"Pass bytes, an open file handle, a (filename, content) "
"tuple, or a pathlib.Path."
)
elif isinstance(file, os.PathLike):
# File path (PathLike): SDK convenience.
with open(str(file), "rb") as f:
elif isinstance(file, os.PathLike) or (
isinstance(file, tuple) and len(file) >= 2 and isinstance(file[1], os.PathLike)
):
# File path (PathLike), bare or as the content of a (filename, path) tuple: SDK convenience.
path: Final = file if isinstance(file, os.PathLike) else file[1]
with open(str(path), "rb") as f:
file_content = f.read()
elif isinstance(file, tuple):
# Tuple format: (filename, content, optional content_type)

21
litellm/litellm_core_utils/litellm_logging.py
View file

@ -2761,6 +2761,8 @@ class Logging(LiteLLMLoggingBaseClass):
def carry_failed_attempt_usage(self) -> None:
"""Move what a failed earlier attempt stashed off the retry sharing this logging object, so the retry logs its
own usage while the request's failure log still bills every attempt."""
if self._carry_parked_success():
return
usage: Final = self._model_call_details_view.get("combined_usage_object")
response_cost: Final = self._model_call_details_view.get("response_cost")
billed_by_continuation: Final = self._stash_billed_by_continuation
@ -2783,6 +2785,25 @@ class Logging(LiteLLMLoggingBaseClass):
return
self._failed_attempts_owed_on_success = (*self._failed_attempts_owed_on_success, attempt)
def _carry_parked_success(self) -> bool:
"""A billed response the router discards after its success log was parked (a safeguard refusal it falls
back from) is owed by whichever attempt ends the request, instead of being logged in its place."""
parked: Final = self.deferred_success_result
if parked is None or self.enqueue_deferred_logging is None:
return False
self.deferred_success_result = None
self.enqueue_deferred_logging = None
self._async_success_scheduled = False
served: Final = self.served_round(parked[0])
if served is None:
return True
self.model_call_details["failed_attempts_usage"] = (
*self.model_call_details.get("failed_attempts_usage", ()),
served,
)
self._failed_attempts_owed_on_success = (*self._failed_attempts_owed_on_success, served)
return True
def mark_stash_billed_by_continuation(self) -> None:
self._stash_billed_by_continuation = True

175
litellm/litellm_core_utils/prompt_templates/common_utils.py
View file

@ -2,10 +2,14 @@
Common utility functions used for translating messages across providers
"""
import importlib
import json
import sys
import warnings
from collections.abc import Hashable, Iterable, Mapping, Sequence
from dataclasses import dataclass
from itertools import chain, groupby
from types import MappingProxyType, ModuleType
from typing import TYPE_CHECKING, Final, Literal, TypeAlias, cast
from pydantic import TypeAdapter
@ -222,6 +226,108 @@ _extract_reasoning_content: Final = extract_reasoning_content
_get_image_mime_type_from_url: Final = get_image_mime_type_from_url
_parse_content_for_reasoning: Final = parse_content_for_reasoning
_MOVED_PRIVATE_NAME_MODULES: Final[Mapping[str, str]] = MappingProxyType(
{
**dict.fromkeys(
("_decode_vector_store_ids_in_tools",),
"file_id_utils",
),
**dict.fromkeys(
(
"_EMPTY_SCHEMA",
"_LOCAL_SCHEMA_REF_PREFIXES",
"_MAX_SCHEMA_FLATTEN_DEPTH",
"_MAX_SCHEMA_NESTING",
"_OPENAI_REJECTED_TOP_LEVEL_SCHEMA_KEYS",
"_REGEX_LOOKAROUND_RE",
"_SCHEMA_LIST",
"_SUBSCHEMA_KEYWORDS",
"_SUBSCHEMA_LIST_KEYWORDS",
"_SUBSCHEMA_MAP_KEYWORDS",
"_TOP_LEVEL_SCHEMA_COMBINATORS",
"_any_of",
"_combinator_required_names",
"_dropped_pattern_properties",
"_flatten_schema_against_root",
"_grammar_branches",
"_is_not_python_regex",
"_is_object_schema",
"_is_open_object",
"_is_union",
"_keyword_value_rebuilt",
"_list_at",
"_may_be_string",
"_mergeable_branch",
"_named_schemas",
"_node_without_rejected_regex",
"_optional_property_names",
"_property_schemas",
"_refs_in",
"_resolve_local_schema_ref",
"_root_definitions",
"_schema_at",
"_schema_branches",
"_schema_levels",
"_schema_nodes",
"_schema_properties",
"_schema_required_names",
"_schema_without_rejected_regex",
"_subschemas",
"_uses_regex_lookaround",
),
"json_schema_utils",
),
**dict.fromkeys(
(
"_BREAKPOINT_BLOCKS",
"_BREAKPOINT_BLOCK_LIST_KEYS",
"_BREAKPOINT_FIELDS",
"_MarkedT",
"_PROMPT_CACHE_BREAKPOINT",
"_breakpoints_stripped_from",
"_without_prompt_cache_breakpoint",
),
"prompt_cache_breakpoint_utils",
),
**dict.fromkeys(
(
"_carries_encrypted_reasoning",
"_converse_wrapped_payload",
"_encrypted_reasoning_field",
"_provider_text_or_none",
"_readable_thinking_text",
"_reasoning_item_from_block_group",
"_reasoning_replay_group_key",
"_strip_encrypted_reasoning_from_blocks",
),
"reasoning_utils",
),
**dict.fromkeys(
(
"_DECODED_JSON",
"_attempt_json_repair",
),
"tool_call_json_repair",
),
**dict.fromkeys(
(
"_TPMMediaProjection",
"_tpm_blank_fields",
"_tpm_code_output_media",
"_tpm_document_media",
"_tpm_file_media",
"_tpm_image_part",
"_tpm_input_media",
"_tpm_media_leaf",
"_tpm_media_list",
"_tpm_replace_field",
"_tpm_tool_result_media",
),
"tpm_media_utils",
),
}
)
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LoggingClass
@ -1578,3 +1684,72 @@ def invalid_type_message(field_path: str, value: object) -> str:
def is_unhashable_value_behind(error: TypeError, value: object) -> bool:
return not isinstance(value, Hashable) and f"unhashable type: '{type(value).__name__}'" in str(error)
_UNBOUND: Final = object()
def _module(module_name: str) -> ModuleType:
return importlib.import_module(f"litellm.litellm_core_utils.prompt_templates.{module_name}")
def _namespace(module: ModuleType) -> Mapping[str, object]:
return vars(module)
def _moved_name_home(name: str) -> ModuleType:
module_name: Final = _MOVED_PRIVATE_NAME_MODULES.get(name)
if module_name is None:
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
return _module(module_name)
def _split_modules_sharing(name: str, value: object) -> tuple[ModuleType, ...]:
if value is _UNBOUND:
return ()
split_modules: Final = (_module(module_name) for module_name in frozenset(_MOVED_PRIVATE_NAME_MODULES.values()))
return tuple(module for module in split_modules if _namespace(module).get(name, _UNBOUND) is value)
def _warn_old_path(name: str, home: ModuleType) -> None:
warnings.warn(
f"{__name__}.{name} moved to {home.__name__}, import or patch it there. "
"The old path will be removed in a future release",
DeprecationWarning,
stacklevel=3,
)
_RELEASED_MOVED_VALUES: Final = MappingProxyType(
{name: _namespace(_moved_name_home(name))[name] for name in _MOVED_PRIVATE_NAME_MODULES}
)
class _CommonUtilsModule(ModuleType):
def __getattr__(self, name: str) -> object:
home: Final = _moved_name_home(name)
value: Final = _namespace(home).get(name, _UNBOUND)
if value is _UNBOUND:
raise AttributeError(f"module {home.__name__!r} has no attribute {name!r}")
_warn_old_path(name, home)
return value
def __setattr__(self, name: str, value: object) -> None:
if name in _MOVED_PRIVATE_NAME_MODULES:
home: Final = _moved_name_home(name)
_warn_old_path(name, home)
setattr(home, name, value)
return
for module in _split_modules_sharing(name, _namespace(self).get(name, _UNBOUND)):
setattr(module, name, value)
super().__setattr__(name, value)
def __delattr__(self, name: str) -> None:
released: Final = _RELEASED_MOVED_VALUES.get(name, _UNBOUND)
if released is _UNBOUND:
super().__delattr__(name)
return
setattr(_moved_name_home(name), name, released)
sys.modules[__name__].__class__ = _CommonUtilsModule

16
litellm/litellm_core_utils/prompt_templates/tool_call_json_repair.py
View file

@ -3,10 +3,21 @@ from typing import Any, Final # noqa: TID251 # Existing annotations retained u
from pydantic import TypeAdapter
import litellm
from litellm import verbose_logger
from litellm.constants import REDACTED_BY_LITELLM
from litellm.litellm_core_utils.core_helpers import as_str_mapping
_DECODED_JSON: Final = TypeAdapter(object)
_ARGUMENT_PREVIEW_CHARS: Final = 200
def _logged_arguments(arguments: str) -> str:
if litellm.turn_off_message_logging:
return REDACTED_BY_LITELLM
if len(arguments) <= _ARGUMENT_PREVIEW_CHARS:
return arguments
return f"{arguments[:_ARGUMENT_PREVIEW_CHARS]}..."
def _attempt_json_repair(s: str) -> object | None:
@ -104,12 +115,11 @@ def parse_tool_call_arguments(
repaired: Final = _attempt_json_repair(arguments)
if repaired is not None:
verbose_logger.warning(
"Repaired truncated tool call arguments for tool '%s' (%s). Original (%d chars): %.200s%s",
"Repaired truncated tool call arguments for tool '%s' (%s). Original (%d chars): %s",
tool_name or "<unknown>",
context or "unknown context",
len(arguments),
arguments,
"..." if len(arguments) > 200 else "",
_logged_arguments(arguments),
)
return repaired

71
litellm/litellm_core_utils/routing_common_utils.py
View file

@ -1,6 +1,7 @@
import hashlib
import json
from collections.abc import Mapping, Sequence
from itertools import chain
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, TypeVar
@ -13,10 +14,12 @@ import litellm
from litellm._logging import verbose_logger, verbose_router_logger
from litellm.constants import ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS
from litellm.exceptions import BadRequestError
from litellm.litellm_core_utils.core_helpers import as_object_list, as_str_mapping
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap
from litellm.litellm_core_utils.sensitive_data_masker import mask_sensitive_structure
from litellm.llms.sagemaker.common_utils import deployment_inference_component
from litellm.secret_managers.main import str_to_bool
from litellm.types.router import (
CredentialLiteLLMParams,
Deployment,
@ -62,6 +65,74 @@ def without_router_only_kwargs(kwargs: Mapping[str, _V]) -> dict[str, _V]:
return {key: value for key, value in kwargs.items() if key not in ROUTER_ONLY_CALL_KWARGS}
_ROUTER_OWNED_METADATA_KEYS: Final = frozenset({"caching_groups", "redis_namespace"})
_PRIVATE_SETTING_BY_FLAG: Final = MappingProxyType({"mask_input": True, "mask_output": True, "trace_public": False})
_JOINED_LIST_METADATA_KEYS: Final = frozenset({"tags", "guardrails"})
def metadata_defaults_under_request(
defaults: Mapping[str, object], request_metadata: Mapping[str, object]
) -> Mapping[str, object]:
return {
key: _default_under_request(key=key, default=default, requested=request_metadata.get(key))
for key, default in defaults.items()
}
def _default_under_request(key: str, default: object, requested: object) -> object:
if key in _ROUTER_OWNED_METADATA_KEYS or callable(default):
return default
if _PRIVATE_SETTING_BY_FLAG.get(key) is _is_flag_on(default):
return default
if key in _JOINED_LIST_METADATA_KEYS:
return _request_items_then_default_items(default=default, requested=requested)
return _request_value_over_default(default=default, requested=requested)
def _is_flag_on(value: object) -> bool:
parsed: Final = str_to_bool(value) if isinstance(value, str) else None
return bool(value) if parsed is None else parsed
def _request_items_then_default_items(default: object, requested: object) -> object:
default_items: Final = as_object_list(default)
if default_items is None:
return _request_value_over_default(default=default, requested=requested)
default_names: Final = frozenset(chain.from_iterable(_entry_names(item) for item in default_items))
request_items: Final = tuple(
chain.from_iterable(_without_entries_named(item, default_names) for item in as_object_list(requested) or ())
)
return [*request_items, *(item for item in default_items if item not in request_items)]
def _entry_names(entry: object) -> frozenset[str]:
mapping: Final = as_str_mapping(entry)
if mapping is not None:
return frozenset(mapping)
return frozenset((entry,)) if isinstance(entry, str) else frozenset()
def _without_entries_named(entry: object, names: frozenset[str]) -> tuple[object, ...]:
mapping: Final = as_str_mapping(entry)
if mapping is None or names.isdisjoint(mapping):
return (entry,)
if names.issuperset(mapping):
return ()
return ({key: value for key, value in mapping.items() if key not in names},)
def _request_value_over_default(default: object, requested: object) -> object:
if requested is None:
return default
default_mapping: Final = as_str_mapping(default)
if default_mapping is None:
return requested
requested_mapping: Final = as_str_mapping(requested)
if requested_mapping is None:
return default
return {**default_mapping, **requested_mapping}
def as_retry_skipped_deployment_ids(value: object) -> tuple[str, ...]:
if not isinstance(value, tuple):
return ()

12
litellm/llms/anthropic/messages/native/handler.py
View file

@ -12,7 +12,7 @@ from dataclasses import dataclass
from functools import partial, wraps
from typing import Final, ParamSpec, TypeAlias, cast
from pydantic import TypeAdapter
from pydantic import TypeAdapter, ValidationError
import litellm
from litellm.anthropic_beta_headers_manager import aload_beta_headers_config
@ -67,6 +67,7 @@ __all__ = ("anthropic_messages", "anthropic_messages_handler")
_RESPONSES_API_PROVIDERS: Final = frozenset({"openai"})
_REQUEST_KWARGS: Final = TypeAdapter(Mapping[str, object])
_SAFEGUARD_MESSAGES: Final = TypeAdapter(Sequence[Mapping[str, object]])
_DROP_PATHS: Final = TypeAdapter(list[object])
_MessagesParams: Final = ParamSpec("_MessagesParams")
_MessagesResult: TypeAlias = AnthropicMessagesResponse | Iterator[bytes] | AsyncIterator[object]
@ -90,6 +91,14 @@ def _with_usage_cost(
return wrapped
def additional_drop_paths(kwargs: Mapping[str, object]) -> tuple[str, ...]:
try:
configured: Final = _DROP_PATHS.validate_python(kwargs.get("additional_drop_params"))
except ValidationError:
return ()
return tuple(path for path in configured if isinstance(path, str))
def _bridges_to_responses_api(model: str, custom_llm_provider: str, deployment_model_info: object) -> bool:
from litellm.responses.bridge_check import responses_api_bridge_check, str_keyed_mapping
@ -859,6 +868,7 @@ def anthropic_messages_handler(
model=model,
drop_params=litellm_params.get("drop_params") is True,
custom_llm_provider=custom_llm_provider,
additional_drop_params=additional_drop_paths(request_kwargs),
)
)
resolved_api_base: Final = (

5
litellm/llms/anthropic/messages/native/utils.py
View file

@ -299,6 +299,7 @@ class AnthropicMessagesRequestUtils:
model: str | None = None,
drop_params: bool = False,
custom_llm_provider: str | None = None,
additional_drop_params: Iterable[str] = (),
) -> AnthropicMessagesRequestOptionalParams:
"""
Filter parameters to only include those defined in AnthropicMessagesRequestOptionalParams.
@ -308,11 +309,13 @@ class AnthropicMessagesRequestUtils:
model: Resolved model id; when set, unsupported params may be dropped
drop_params: Per-request drop_params flag (also respects litellm.drop_params)
custom_llm_provider: Routed provider; fast mode is gated to direct Anthropic
additional_drop_params: Paths the deployment drops. A listed top-level key is left out
before the unsupported-param checks run, so it never raises
Returns:
AnthropicMessagesRequestOptionalParams instance with only the valid parameters
"""
valid_keys: Final = anthropic_messages_optional_param_keys()
valid_keys: Final = anthropic_messages_optional_param_keys().difference(additional_drop_params)
filtered_params: Final = {k: v for k, v in params.items() if k in valid_keys and v is not None}
if model is not None:
from litellm.llms.anthropic.chat.transformation import AnthropicConfig

8
litellm/llms/anthropic/prompt_cache_prediction.py
View file

@ -654,17 +654,15 @@ def prepare_native_baseline_body(request: Mapping[str, object], model: str) -> M
custom_llm_provider="anthropic",
tools=_MESSAGES.validate_python(owned.get("tools") or []),
)
drop_paths: Final = TypeAdapter(tuple[str, ...]).validate_python(owned.get("additional_drop_params") or ())
options: Final = AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param(
{**owned, "system": system},
model=resolved_model,
custom_llm_provider="anthropic",
drop_params=owned.get("drop_params") is True,
additional_drop_params=drop_paths,
)
filtered: Final = reduce(
delete_nested_value,
TypeAdapter(tuple[str, ...]).validate_python(owned.get("additional_drop_params") or ()),
dict(options),
)
filtered: Final = reduce(delete_nested_value, drop_paths, dict(options))
config: Final = _BaselineAnthropicMessagesConfig()
body: Final = config.transform_anthropic_messages_request(
model=resolved_model,

26
litellm/llms/bedrock/base_aws_llm.py
View file

@ -1712,6 +1712,24 @@ class BaseAWSLLM(SignsRequestsWithAWS):
return request_headers_dict, request.body
def sign_aws_request(
method: str,
get_credentials: Callable[[], Credentials],
service_name: str,
aws_region_name: str | None,
url: str,
body: str | bytes,
headers: Mapping[str, str],
) -> AWSPreparedRequest:
ensure_optional_import("botocore")
from botocore.auth import SigV4Auth
from botocore.awsrequest import AWSRequest
aws_request: Final = AWSRequest(method=method, url=url, data=body, headers=headers)
SigV4Auth(get_credentials(), service_name, aws_region_name).add_auth(aws_request)
return aws_request.prepare()
def sign_aws_json_post(
get_credentials: Callable[[], Credentials],
service_name: str,
@ -1720,13 +1738,7 @@ def sign_aws_json_post(
body: str,
headers: Mapping[str, str],
) -> AWSPreparedRequest:
ensure_optional_import("botocore")
from botocore.auth import SigV4Auth
from botocore.awsrequest import AWSRequest
aws_request: Final = AWSRequest(method="POST", url=url, data=body, headers=headers)
SigV4Auth(get_credentials(), service_name, aws_region_name).add_auth(aws_request)
return aws_request.prepare()
return sign_aws_request("POST", get_credentials, service_name, aws_region_name, url, body, headers)
_SignParams = ParamSpec("_SignParams")

12
litellm/llms/bedrock/count_tokens/bedrock_token_counter.py
View file

@ -21,7 +21,7 @@ from litellm.llms.bedrock.common_utils import (
BedrockError,
get_bedrock_base_model,
)
from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler
from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler, BedrockCountTokensRequestError
from litellm.llms.bedrock.count_tokens.mantle_handler import BedrockMantleCountTokensHandler
from litellm.llms.bedrock_mantle.common_utils import is_mantle_claude_model
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
@ -54,6 +54,7 @@ class CountTokensFailure:
status_code: int
message: str
tokenizer_type: str
rejected_before_provider_call: bool = False
CountTokensOutcome = CountedTokens | CountTokensFailure
@ -74,6 +75,7 @@ def _claude_count_translator(litellm_params: Mapping[str, object]) -> AnthropicC
def _runtime_rejected_claude_model(outcome: CountTokensOutcome, resolved_model: str) -> bool:
return (
isinstance(outcome, CountTokensFailure)
and not outcome.rejected_before_provider_call
and outcome.status_code in _RUNTIME_FALLBACK_STATUSES
and is_mantle_claude_model(resolved_model)
)
@ -162,6 +164,13 @@ class BedrockTokenCounter(BaseTokenCounter):
client=self._client,
)
reply: Final = _CountTokensReply.model_validate(result)
except BedrockCountTokensRequestError as e:
return CountTokensFailure(
status_code=e.status_code,
message=e.message,
tokenizer_type=tokenizer_type,
rejected_before_provider_call=True,
)
except BedrockError as e:
verbose_logger.debug(
"%s CountTokens API error: status=%s, message=%s", tokenizer_type, e.status_code, e.message
@ -268,6 +277,7 @@ class BedrockTokenCounter(BaseTokenCounter):
error=True,
error_message=outcome.message,
status_code=outcome.status_code,
rejected_before_provider_call=outcome.rejected_before_provider_call,
)
case _:
assert_never(outcome)

11
litellm/llms/bedrock/count_tokens/handler.py
View file

@ -20,6 +20,11 @@ from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_h
_DECODED_JSON: Final = TypeAdapter(object)
class BedrockCountTokensRequestError(BedrockError):
def __init__(self, problem: str) -> None:
super().__init__(status_code=400, message=problem)
class BedrockCountTokensHandler(BedrockCountTokensConfig):
"""
Simplified handler for AWS Bedrock CountTokens API requests.
@ -45,10 +50,10 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig):
Returns:
Dictionary containing token count response
"""
problem: Final = self.count_tokens_request_problem(request_data)
if problem is not None:
raise BedrockCountTokensRequestError(problem)
try:
# Validate the request
self.validate_count_tokens_request(request_data)
verbose_logger.debug("Processing CountTokens request for resolved model: %s", resolved_model)
# Get AWS region using existing LiteLLM function

5
litellm/llms/bedrock/count_tokens/mantle_handler.py
View file

@ -12,6 +12,8 @@ from litellm.llms.anthropic.common_utils import context_management_betas, tools_
from litellm.llms.anthropic.count_tokens.transformation import COUNT_TOKEN_OPTION_NAMES, AnthropicCountTokensConfig
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing
from litellm.llms.bedrock.common_utils import BedrockError, build_mantle_messages_url
from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensRequestError
from litellm.llms.bedrock.count_tokens.transformation import count_tokens_messages_problem
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client
MANTLE_COUNT_TOKENS_SUFFIX: Final = "/count_tokens"
@ -58,6 +60,9 @@ class BedrockMantleCountTokensHandler(BaseAWSLLM):
) -> dict[str, JsonValue]:
try:
request: Final = _COUNT_REQUEST.validate_python(request_data)
problem: Final = count_tokens_messages_problem(request["messages"])
if problem is not None:
raise BedrockCountTokensRequestError(problem)
aws_region_name: Final = await run_aws_signing(
self.get_aws_region_name, optional_params=litellm_params, model=resolved_model, model_id=None
)

74
litellm/llms/bedrock/count_tokens/transformation.py
View file

@ -6,7 +6,7 @@ to AWS Bedrock's CountTokens API format and vice versa.
"""
import re
from collections.abc import Mapping
from collections.abc import Mapping, Sequence
from typing import Final, Literal
from pydantic import JsonValue
@ -14,7 +14,11 @@ from pydantic import JsonValue
from litellm.litellm_core_utils.url_utils import join_api_base
from litellm.llms.anthropic.common_utils import tools_without_custom_field
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.llms.bedrock.common_utils import get_bedrock_base_model
from litellm.llms.bedrock.common_utils import (
get_bedrock_base_model,
get_bedrock_cross_region_inference_regions,
strip_bedrock_routing_prefix,
)
# Placeholder satisfying the Anthropic InvokeModel schema's required
# max_tokens field; CountTokens only counts input, so it has no effect
@ -46,6 +50,23 @@ def _to_converse_message(message: JsonValue) -> dict[str, JsonValue]:
}
def _converse_message_problem(index: int, message: JsonValue) -> str | None:
if not isinstance(message, dict):
return f"Message {index} must be a dictionary"
if "role" not in message:
return f"Message {index} must have a 'role' field"
if "content" not in message and "tool_calls" in message:
return f"Message {index} has tool_calls and no content, which token counting cannot convert for this model"
if "content" not in message:
return f"Message {index} must have a 'content' field"
return None
def count_tokens_messages_problem(messages: Sequence[JsonValue]) -> str | None:
problems: Final = (_converse_message_problem(index, message) for index, message in enumerate(messages))
return next((problem for problem in problems if problem is not None), None)
def _sanitized_bedrock_tool_name(raw_name: JsonValue) -> str:
name: Final = re.sub(r"[^a-zA-Z0-9_]", "_", raw_name if isinstance(raw_name, str) else "")
prefixed: Final = name if not name or name[0].isalpha() else f"t_{name}"
@ -217,6 +238,9 @@ class BedrockCountTokensConfig(BaseAWSLLM):
@staticmethod
def count_tokens_model_id(model: str) -> str:
model_id: Final = strip_bedrock_routing_prefix(model)
if model_id in get_bedrock_cross_region_inference_regions():
return model_id
return get_bedrock_base_model(model).removeprefix("bedrock/")
def get_bedrock_count_tokens_endpoint(
@ -265,42 +289,12 @@ class BedrockCountTokensConfig(BaseAWSLLM):
return {"input_tokens": input_tokens}
def validate_count_tokens_request(self, request_data: Mapping[str, JsonValue]) -> None:
"""
Validate the incoming count tokens request.
Supports both Converse and InvokeModel input formats.
Args:
request_data: The request payload
Raises:
ValueError: If the request is invalid
"""
def count_tokens_request_problem(self, request_data: Mapping[str, JsonValue]) -> str | None:
if not request_data.get("model"):
raise ValueError("model parameter is required")
input_type: Final = self._detect_input_type(request_data)
if input_type == "converse":
# Validate Converse format (messages-based)
messages: Final = request_data.get("messages", [])
if not messages:
raise ValueError("messages parameter is required for Converse input")
if not isinstance(messages, list):
raise ValueError("messages must be a list")
for i, message in enumerate(messages):
if not isinstance(message, dict):
raise ValueError(f"Message {i} must be a dictionary")
if "role" not in message:
raise ValueError(f"Message {i} must have a 'role' field")
if "content" not in message:
raise ValueError(f"Message {i} must have a 'content' field")
else:
# For InvokeModel format, we need at least some content to count tokens
# The content structure varies by model, so we do minimal validation
if len(request_data) <= 1: # Only has 'model' field
raise ValueError("Request must contain content to count tokens")
return "model parameter is required"
if self._detect_input_type(request_data) == "invokeModel":
return "Request must contain content to count tokens" if len(request_data) <= 1 else None
messages: Final = _json_list(request_data.get("messages"))
if not messages:
return "messages parameter is required for Converse input"
return count_tokens_messages_problem(messages)

15
litellm/llms/bedrock/web_identity_session_policy.py
View file

@ -266,6 +266,21 @@ POLLY_SESSION_POLICY: Final = _policy(_allow("PollyLiteLLM", ("polly:SynthesizeS
OPENSEARCH_SESSION_POLICY: Final = _policy(_allow("OpenSearchLiteLLM", ("es:ESHttpPost", "aoss:APIAccessAll")))
OPENSEARCH_INDEX_SESSION_POLICY: Final = _policy(
_allow(
"OpenSearchIndexLiteLLM",
(
"es:ESHttpGet",
"es:ESHttpHead",
"es:ESHttpPost",
"es:ESHttpPut",
"es:ESHttpDelete",
"es:ESHttpPatch",
"aoss:APIAccessAll",
),
),
)
S3_VECTORS_QUERY_SESSION_POLICY: Final = _policy(
_allow("S3VectorsLiteLLM", ("s3vectors:QueryVectors", "s3vectors:GetVectors")),
)

23
litellm/llms/minimax/chat/transformation.py
View file

@ -11,6 +11,13 @@ from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam
def _supports_reasoning(model: str) -> bool:
try:
return litellm.supports_reasoning(model=model, custom_llm_provider="minimax")
except Exception:
return False
class MinimaxChatConfig(OpenAIGPTConfig):
"""
MiniMax OpenAI configuration that extends OpenAIGPTConfig.
@ -81,16 +88,8 @@ class MinimaxChatConfig(OpenAIGPTConfig):
def get_supported_openai_params(self, model: str) -> list:
"""
Get supported OpenAI parameters for MiniMax.
Adds reasoning_split and thinking to the list of supported params.
Adds reasoning_split, plus thinking and reasoning_effort when the cost map, or its
minimax/ fallback rule for unlisted models, says the model supports reasoning.
"""
base_params: Final = super().get_supported_openai_params(model=model)
additional_params: Final = ["reasoning_split"]
# Add thinking parameter if model supports reasoning
try:
if litellm.supports_reasoning(model=model, custom_llm_provider="minimax"):
additional_params.append("thinking")
except Exception:
pass
return base_params + additional_params
reasoning_params: Final = ("thinking", "reasoning_effort") if _supports_reasoning(model) else ()
return [*super().get_supported_openai_params(model=model), "reasoning_split", *reasoning_params]

34
litellm/llms/ollama/chat/transformation.py
View file

@ -16,6 +16,9 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
from litellm.litellm_core_utils.prompt_templates.reasoning_utils import (
extract_reasoning_content,
)
from litellm.litellm_core_utils.prompt_templates.tool_call_json_repair import (
parse_tool_call_arguments,
)
from litellm.litellm_core_utils.reasoning_effort_utils import OPAQUE_VALUE
from litellm.litellm_core_utils.tool_fields import tool_params_without_hosted_mcp
from litellm.litellm_core_utils.url_utils import ensure_api_route
@ -44,6 +47,27 @@ else:
LiteLLMLoggingObj = Any
_REQUEST_PARAMS: Final = TypeAdapter(dict[str, object])
_TOOL_CALL_ARGUMENTS: Final = TypeAdapter(dict[str, object])
def _to_ollama_tool_call_function(model: str, tool_call: ChatCompletionAssistantToolCall) -> OllamaToolCallFunction:
function: Final = tool_call["function"]
name: Final = function.get("name") or ""
raw_arguments: Final[object] = function.get("arguments")
try:
parsed: Final[object] = (
parse_tool_call_arguments(raw_arguments, tool_name=name, context="Ollama chat request")
if isinstance(raw_arguments, str) or raw_arguments is None
else raw_arguments
)
arguments: Final = _TOOL_CALL_ARGUMENTS.validate_python({} if parsed is None else parsed)
except ValueError:
raise litellm.BadRequestError(
message=f"OllamaException - arguments of assistant tool call {tool_call.get('id')!r} are not a JSON object",
model=model,
llm_provider="ollama_chat",
) from None
return OllamaToolCallFunction(name=name, arguments=arguments)
class OllamaChatConfig(BaseConfig):
@ -247,12 +271,7 @@ class OllamaChatConfig(BaseConfig):
tools: Final = optional_params.pop("tools", None)
def transform_tool_call(typed_tool: ChatCompletionAssistantToolCall) -> OllamaToolCall:
function: Final = OllamaToolCallFunction(
name=typed_tool["function"].get("name") or "",
arguments=json.loads(typed_tool["function"]["arguments"])
if "arguments" in typed_tool["function"]
else {},
)
function: Final = _to_ollama_tool_call_function(model, typed_tool)
tool_id: Final = typed_tool.get("id")
ollama_tool_call: Final = (
OllamaToolCall(function=function) if tool_id is None else OllamaToolCall(id=tool_id, function=function)
@ -442,6 +461,7 @@ class OllamaChatCompletionResponseIterator(BaseModelResponseIterator):
started_reasoning_content: bool = False
finished_reasoning_content: bool = False
seen_tool_calls: bool = False
_next_tool_call_index: int = 0
def _is_function_call_complete(self, function_args: str | dict) -> bool:
if isinstance(function_args, dict):
@ -489,6 +509,8 @@ class OllamaChatCompletionResponseIterator(BaseModelResponseIterator):
if tool_calls is not None:
self.seen_tool_calls = True
for tool_call in tool_calls:
tool_call["index"] = self._next_tool_call_index
self._next_tool_call_index += 1
function_args = tool_call.get("function").get("arguments")
if function_args is not None and len(function_args) > 0:
is_function_call_complete = self._is_function_call_complete(function_args)

379
litellm/llms/opensearch/vector_stores/index_passthrough.py Normal file
View file

@ -0,0 +1,379 @@
import json
from collections.abc import Iterable, Iterator, Mapping, Sequence
from dataclasses import dataclass
from itertools import accumulate, chain, repeat, takewhile
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, Literal, TypeAlias
from urllib.parse import quote
from pydantic import TypeAdapter
from litellm.exceptions import AuthenticationError, BadRequestError, InternalServerError
from litellm.llms.bedrock.base_aws_llm import run_aws_signing
if TYPE_CHECKING:
from litellm.llms.opensearch.vector_stores.transformation import OpenSearchVectorStoreConfig
IndexHttpMethod: TypeAlias = Literal["GET", "HEAD", "POST", "PUT", "DELETE", "PATCH"]
IndexPermission: TypeAlias = Literal["read", "write"]
BodyKind: TypeAlias = Literal["query", "mget", "bulk", "update", "unscanned"]
DEFAULT_CONTENT_TYPE: Final = "application/json"
NDJSON_CONTENT_TYPE: Final = "application/x-ndjson"
ACCEPTED_MEDIA_TYPES: Final = frozenset({DEFAULT_CONTENT_TYPE, NDJSON_CONTENT_TYPE})
_ACCEPTED_MEDIA_TYPE_PARAMETERS: Final = frozenset({"", "charset=utf-8"})
_INDEX_KEYS: Final = frozenset({"index", "_index", "indices"})
_OUT_OF_INDEX_KEYS: Final = _INDEX_KEYS | {"indexed_shape", "wrapper", "search_pipeline"}
_BULK_METADATA_KEYS: Final = _INDEX_KEYS | {"pipeline"}
_OUT_OF_INDEX_QUERY_PARAMETERS: Final = frozenset({"pipeline", "source", "source_content_type"})
_BULK_ACTIONS: Final = frozenset({"index", "create", "update", "delete"})
_DOCUMENT_READING_KEYS: Final = frozenset({"_source", "script"})
_JSON_BLANK: Final = b" \t\r\n"
_PATH_HELP: Final = "Use /opensearch/<index>/<api>, for example /opensearch/<index>/_search."
@dataclass(frozen=True, slots=True)
class IndexRefusal:
status_code: int
message: str
@dataclass(frozen=True, slots=True)
class IndexPath:
index_name: str
api: tuple[str, ...]
@dataclass(frozen=True, slots=True)
class IndexOperation:
permission: IndexPermission
methods: tuple[IndexHttpMethod, ...]
api: str
takes_id: bool
body: BodyKind
def matches(self, method: str, api: tuple[str, ...]) -> bool:
named: Final = (self.api,) if self.api else ()
return method in self.methods and api[: len(named)] == named and len(api) == len(named) + self.takes_id
@dataclass(frozen=True, slots=True)
class SignedIndexRequest:
url: str
media_type: str
headers: Mapping[str, str]
body: bytes
INDEX_OPERATIONS: Final = (
IndexOperation("read", ("GET", "HEAD"), "", False, "unscanned"),
IndexOperation("write", ("POST",), "_doc", False, "unscanned"),
IndexOperation("write", ("PUT", "POST", "DELETE"), "_doc", True, "unscanned"),
IndexOperation("write", ("PUT", "POST"), "_create", True, "unscanned"),
IndexOperation("write", ("POST",), "_update", True, "update"),
IndexOperation("write", ("POST", "PUT"), "_bulk", False, "bulk"),
IndexOperation("write", ("POST",), "_delete_by_query", False, "query"),
IndexOperation("write", ("POST",), "_update_by_query", False, "query"),
IndexOperation("read", ("GET", "POST"), "_search", False, "query"),
IndexOperation("read", ("GET", "POST"), "_count", False, "query"),
IndexOperation("read", ("GET", "HEAD"), "_doc", True, "unscanned"),
IndexOperation("read", ("GET", "HEAD"), "_source", True, "unscanned"),
IndexOperation("read", ("GET", "POST"), "_mget", False, "mget"),
IndexOperation("read", ("GET", "POST"), "_explain", True, "query"),
IndexOperation("read", ("GET",), "_mapping", False, "unscanned"),
)
def index_operation(method: str, api: tuple[str, ...]) -> IndexOperation | None:
return next((operation for operation in INDEX_OPERATIONS if operation.matches(method, api)), None)
def parse_index_path(endpoint: str) -> IndexPath | IndexRefusal:
segments: Final = tuple(endpoint.split("/"))
if any(segment in ("", ".", "..") for segment in segments):
return IndexRefusal(400, f"The OpenSearch path '{endpoint}' has an empty, '.' or '..' segment. {_PATH_HELP}")
return IndexPath(index_name=segments[0], api=segments[1:])
@dataclass(frozen=True, slots=True)
class _JsonObject:
pairs: tuple[tuple[str, object], ...]
@dataclass(frozen=True, slots=True)
class _Json:
value: object
@dataclass(frozen=True, slots=True)
class _BulkAction:
name: str
metadata: _JsonObject
has_source: bool
@dataclass(frozen=True, slots=True)
class _IndexScope:
name: str
provider_name: str
can_read: bool
def _json_object(pairs: Sequence[tuple[str, object]]) -> _JsonObject:
return _JsonObject(tuple(pairs))
_DECODER: Final = json.JSONDecoder(object_pairs_hook=_json_object)
_JSON_ARRAY: Final = TypeAdapter(list[object])
_JSON_VALUE: Final = TypeAdapter(object)
def _parse(raw: bytes) -> _Json | None:
try:
return _Json(_JSON_VALUE.validate_python(_DECODER.decode(raw.decode("utf-8"))))
except ValueError:
return None
def _own_keys(value: object) -> tuple[str, ...]:
return tuple(key for key, _ in value.pairs) if isinstance(value, _JsonObject) else ()
def _children(value: object) -> Sequence[object]:
if isinstance(value, _JsonObject):
return tuple(child for _, child in value.pairs)
if isinstance(value, list):
return _JSON_ARRAY.validate_python(value)
return ()
def _next_level(level: Sequence[object], _: None) -> Sequence[object]:
return tuple(chain.from_iterable(map(_children, level)))
def _forbidden_key(value: object, forbidden: frozenset[str]) -> str | None:
levels: Final = takewhile(bool, accumulate(repeat(None), _next_level, initial=(value,)))
keys: Final[Iterable[str]] = chain.from_iterable(map(_own_keys, chain.from_iterable(levels)))
return next((key for key in keys if key in forbidden), None)
def _plain(value: _JsonObject) -> Mapping[str, object]:
return dict(value.pairs)
def _object_json(pairs: Sequence[tuple[str, object]]) -> str:
fields: Final = ",".join(
f"{json.dumps(key)}:{json.dumps(child, default=_plain, separators=(',', ':'))}" for key, child in pairs
)
return f"{{{fields}}}"
def _unparseable(api: str) -> IndexRefusal:
return IndexRefusal(400, f"LiteLLM could not parse the {api} request body as JSON, so it did not send it.")
def _index_reference(value: object, forbidden: frozenset[str]) -> IndexRefusal | None:
key: Final = _forbidden_key(value, forbidden)
if key is None:
return None
return IndexRefusal(
400,
f"The request body uses '{key}', which can reach beyond the OpenSearch index in the URL path. "
f"Remove '{key}' from the body.",
)
def _read_needed(cause: str, index_name: str) -> IndexRefusal:
return IndexRefusal(
403,
f"{cause}, so the key needs read on index '{index_name}' as well as write. "
"Ask your administrator to add read to the key or its team.",
)
def _document_reading_payload(api: str, payload: bytes, index: _IndexScope) -> IndexRefusal | None:
if index.can_read:
return None
parsed: Final = _parse(payload)
if parsed is None:
return _unparseable(api)
key: Final = next((key for key in _own_keys(parsed.value) if key in _DOCUMENT_READING_KEYS), None)
return None if key is None else _read_needed(f"'{key}' in an {api} body reads stored documents", index.name)
def _line_refusal(api: str, line: bytes, forbidden: frozenset[str]) -> IndexRefusal | None:
parsed: Final = _parse(line)
if parsed is None:
return _unparseable(api)
return _index_reference(parsed.value, forbidden)
def _whole_body_refusal(api: str, body: bytes) -> IndexRefusal | None:
if not body.strip(_JSON_BLANK):
return None
return _line_refusal(api, body, _OUT_OF_INDEX_KEYS)
def _pinned_index(metadata: _JsonObject, index: _IndexScope) -> _JsonObject | IndexRefusal | None:
unnamed: Final = _JsonObject(tuple(pair for pair in metadata.pairs if pair[0] != "_index"))
refusal: Final = _index_reference(unnamed, _BULK_METADATA_KEYS)
if refusal is not None:
return refusal
named: Final = tuple(value for name, value in metadata.pairs if name == "_index")
if not named:
return None
if any(value != index.name for value in named):
return IndexRefusal(
400,
f"A _bulk action line sets '_index' to another index than '{index.name}' in the URL path. "
f"Remove '_index' or set it to '{index.name}'.",
)
return _JsonObject(
tuple((name, index.provider_name if name == "_index" else value) for name, value in metadata.pairs)
)
def _bulk_action(line: bytes) -> _BulkAction | None:
parsed: Final = _parse(line)
if parsed is None:
return None
match parsed.value:
case _JsonObject(pairs=((name, _JsonObject() as metadata),)) if name in _BULK_ACTIONS:
return _BulkAction(name=name, metadata=metadata, has_source=name != "delete")
case _:
return None
def _bulk_action_line(action: _BulkAction | None, line: bytes, index: _IndexScope) -> bytes | IndexRefusal:
if action is None:
return IndexRefusal(
400,
"A _bulk action line must be one JSON object naming index, create, update, or delete. "
"LiteLLM did not send the request.",
)
if not index.can_read and "_source" in _own_keys(action.metadata):
return _read_needed("'_source' in a _bulk action line returns stored documents", index.name)
pinned: Final = _pinned_index(action.metadata, index)
if pinned is None:
return line
if isinstance(pinned, IndexRefusal):
return pinned
return f"{{{json.dumps(action.name)}:{_object_json(pinned.pairs)}}}".encode()
def _bulk_entry(line: bytes, lines: Iterator[bytes], index: _IndexScope) -> tuple[bytes | IndexRefusal, ...]:
if not line.strip(_JSON_BLANK):
return (line,)
action: Final = _bulk_action(line)
rendered: Final = _bulk_action_line(action, line, index)
if isinstance(rendered, IndexRefusal) or action is None or not action.has_source:
return (rendered,)
source: Final = next(lines, None)
if source is None:
return (rendered,)
if action.name != "update":
return (rendered, source)
return (rendered, _document_reading_payload("_bulk update", source, index) or source)
def _bulk_lines(lines: Iterator[bytes], index: _IndexScope) -> Iterator[bytes | IndexRefusal]:
for line in lines:
yield from (entry := _bulk_entry(line, lines, index))
if any(isinstance(rendered, IndexRefusal) for rendered in entry):
return
def _rendered_body(lines: Iterator[bytes | IndexRefusal]) -> bytes | IndexRefusal:
rendered: Final = tuple(lines)
refusal: Final = next((line for line in rendered if isinstance(line, IndexRefusal)), None)
if refusal is not None:
return refusal
return b"\n".join(line for line in rendered if isinstance(line, bytes))
def _checked_body(operation: IndexOperation, body: bytes, index: _IndexScope) -> bytes | IndexRefusal:
match operation.body:
case "query" if operation.permission == "write" and not index.can_read:
return _read_needed(f"{operation.api} searches stored documents", index.name)
case "query" | "mget":
return _whole_body_refusal(operation.api, body) or body
case "update":
return _document_reading_payload(operation.api, body, index) or body
case "bulk":
return _rendered_body(_bulk_lines(iter(body.split(b"\n")), index))
case "unscanned":
return body
def accepted_media_type(content_type: str | None) -> str | None:
media, _, parameters = (content_type or DEFAULT_CONTENT_TYPE).partition(";")
normalized: Final = media.strip().lower()
if normalized in ACCEPTED_MEDIA_TYPES and parameters.strip().lower() in _ACCEPTED_MEDIA_TYPE_PARAMETERS:
return normalized
return None
def _checked_request_body(
operation: IndexOperation | None,
query: Sequence[tuple[str, str]],
body: bytes,
index: _IndexScope,
) -> bytes | IndexRefusal:
if operation is None:
return body
parameter: Final = next((key for key, _ in query if key in _OUT_OF_INDEX_QUERY_PARAMETERS), None)
if parameter is not None:
return IndexRefusal(
400, f"The {parameter} query parameter can reach beyond the OpenSearch index in the URL path. Remove it."
)
source_parameter: Final = next((key for key, _ in query if key.startswith("_source")), None)
if operation.permission == "write" and not index.can_read and source_parameter is not None:
return _read_needed(f"The {source_parameter} query parameter returns stored documents", index.name)
try:
return _checked_body(operation, body, index)
except RecursionError:
return IndexRefusal(400, f"The {operation.api} request body is nested too deeply for LiteLLM to check.")
def upstream_url(api_base: str, provider_index: str, api: tuple[str, ...], query: Sequence[tuple[str, str]]) -> str:
path: Final = "/".join(quote(segment, safe="") for segment in (provider_index, *api))
query_string: Final = "&".join(f"{quote(key, safe='')}={quote(value, safe='')}" for key, value in query)
return f"{api_base}/{path}?{query_string}" if query_string else f"{api_base}/{path}"
async def prepare_index_request(
config: "OpenSearchVectorStoreConfig",
method: str,
path: IndexPath,
operation: IndexOperation | None,
provider_index: str,
content_type: str | None,
query: Sequence[tuple[str, str]],
body: bytes,
store_params: Mapping[str, object],
can_read: bool,
) -> SignedIndexRequest | IndexRefusal:
media: Final = accepted_media_type(content_type)
if media is None:
return IndexRefusal(
400, "The OpenSearch pass-through takes application/json or application/x-ndjson bodies in UTF-8."
)
checked: Final = _checked_request_body(
operation, query, body, _IndexScope(name=path.index_name, provider_name=provider_index, can_read=can_read)
)
if isinstance(checked, IndexRefusal):
return checked
api_base: Final = store_params.get("api_base")
try:
url: Final = upstream_url(
config.get_complete_url(api_base if isinstance(api_base, str) else None, store_params),
provider_index,
path.api,
query,
)
signed: Final = await run_aws_signing(config.sign_http_request, method, url, checked, media, store_params)
except (AuthenticationError, BadRequestError, InternalServerError) as error:
return IndexRefusal(error.status_code, error.message)
return SignedIndexRequest(
url=url, media_type=media, headers=MappingProxyType(dict(signed.headers.items())), body=checked
)

68
litellm/llms/opensearch/vector_stores/transformation.py
View file

@ -10,7 +10,7 @@ import httpx
from pydantic import ConfigDict, Field, TypeAdapter, ValidationError
from litellm._logging import verbose_logger
from litellm.exceptions import AuthenticationError, BadRequestError
from litellm.exceptions import AuthenticationError, BadRequestError, InternalServerError
from litellm.litellm_core_utils.core_helpers import drops_unsupported_params
from litellm.litellm_core_utils.optional_imports import ensure_optional_import
from litellm.llms.base_llm.chat.transformation import BaseLLMException
@ -18,8 +18,12 @@ from litellm.llms.base_llm.vector_store.transformation import (
BaseQueryEmbeddingVectorStoreConfig,
VectorStoreEmbeddingExecutor,
)
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, SignsRequestsWithAWS, sign_aws_json_post
from litellm.llms.bedrock.web_identity_session_policy import OPENSEARCH_SESSION_POLICY, WebIdentitySessionPolicy
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, SignsRequestsWithAWS, sign_aws_request
from litellm.llms.bedrock.web_identity_session_policy import (
OPENSEARCH_INDEX_SESSION_POLICY,
OPENSEARCH_SESSION_POLICY,
WebIdentitySessionPolicy,
)
from litellm.types.llms.base import LiteLLMBaseModel
from litellm.types.llms.bedrock import AwsAuthParams
from litellm.types.router import GenericLiteLLMParams
@ -35,6 +39,7 @@ from litellm.types.vector_stores import (
)
if TYPE_CHECKING:
from botocore.awsrequest import AWSPreparedRequest
from botocore.credentials import Credentials
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
@ -81,6 +86,16 @@ def _unauthenticated(reason: str) -> AuthenticationError:
)
def _unavailable_credentials(error: Exception) -> InternalServerError:
verbose_logger.warning("OpenSearch vector store AWS credentials could not be loaded or refreshed: %s", error)
return InternalServerError(
message="LiteLLM could not load or refresh the AWS credentials that sign the OpenSearch request, "
"see the LiteLLM logs.",
llm_provider="opensearch",
model=None,
)
def _usable_credentials(credentials: object) -> "Credentials":
from botocore.credentials import Credentials
@ -213,6 +228,10 @@ class _OpenSearchAWS(BaseAWSLLM):
_web_identity_session_policy: ClassVar[WebIdentitySessionPolicy] = OPENSEARCH_SESSION_POLICY
class OpenSearchIndexAWS(BaseAWSLLM):
_web_identity_session_policy: ClassVar[WebIdentitySessionPolicy] = OPENSEARCH_INDEX_SESSION_POLICY
class OpenSearchVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig, SignsRequestsWithAWS):
def __init__(self, aws: BaseAWSLLM | None = None) -> None:
self._aws: Final = aws if aws is not None else _OpenSearchAWS()
@ -243,6 +262,7 @@ class OpenSearchVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig, SignsRequ
if bool(auth.aws_access_key_id) != bool(auth.aws_secret_access_key):
raise _unauthenticated("aws_access_key_id and aws_secret_access_key must be set together")
ensure_optional_import("botocore")
from botocore.credentials import Credentials
from botocore.exceptions import (
CredentialRetrievalError,
NoCredentialsError,
@ -251,12 +271,21 @@ class OpenSearchVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig, SignsRequ
)
try:
return _usable_credentials(self._aws.resolve_credentials(auth, region))
credentials: Final = _usable_credentials(self._aws.resolve_credentials(auth, region))
except CredentialRetrievalError as error:
verbose_logger.warning("OpenSearch vector store AWS credential source failed: %s", error)
raise _unauthenticated("the AWS credential source failed, see the LiteLLM logs") from None
except OSError as error:
raise _unavailable_credentials(error) from None
except (NoCredentialsError, PartialCredentialsError, ProfileNotFound) as error:
raise _unauthenticated(str(error)) from None
try:
frozen: Final = credentials.get_frozen_credentials()
except (CredentialRetrievalError, OSError, RuntimeError) as error:
raise _unavailable_credentials(error) from None
if frozen.access_key is None or frozen.secret_key is None:
raise _unauthenticated("none were found")
return Credentials(frozen.access_key, frozen.secret_key, frozen.token, account_id=frozen.account_id)
def get_complete_url(self, api_base: str | None, litellm_params: Mapping[str, object]) -> str:
parsed: Final = urlsplit(api_base or "")
@ -368,25 +397,28 @@ class OpenSearchVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig, SignsRequ
api_base: str,
api_key: str | None = None,
) -> tuple[dict[str, str], bytes]: # mutable-ok: the shared HTTP handler requires writable headers
region: Final = self._signing_region(optional_params, api_base)
body: Final = json.dumps(request_data)
signed: Final = sign_aws_json_post(
get_credentials=partial(self._signing_credentials, optional_params, region),
service_name=aws_signing_target(api_base)[0],
aws_region_name=region,
url=api_base,
body=body,
headers={
"Content-Type": "application/json",
"x-amz-content-sha256": hashlib.sha256(body.encode()).hexdigest(),
},
)
body: Final = json.dumps(request_data).encode()
signed: Final = self.sign_http_request("POST", api_base, body, "application/json", optional_params)
signed_headers: Final = dict(signed.headers.items())
signed_names: Final = frozenset(name.lower() for name in signed_headers)
return {
**{name: value for name, value in headers.items() if name.lower() not in signed_names},
**signed_headers,
}, body.encode()
}, body
def sign_http_request(
self, method: str, url: str, body: bytes, content_type: str, litellm_params: Mapping[str, object]
) -> "AWSPreparedRequest":
region: Final = self._signing_region(litellm_params, url)
return sign_aws_request(
method=method,
get_credentials=partial(self._signing_credentials, litellm_params, region),
service_name=aws_signing_target(url)[0],
aws_region_name=region,
url=url,
body=body,
headers={"Content-Type": content_type, "x-amz-content-sha256": hashlib.sha256(body).hexdigest()},
)
def transform_search_vector_store_response(
self, response: httpx.Response, litellm_logging_obj: "LiteLLMLoggingObj"

23
litellm/messages/dispatch.py
View file

@ -6,6 +6,7 @@ from typing import Final, TypeAlias, cast # noqa: TID251 # native binding sele
from litellm.exceptions import BadRequestError
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
from litellm.llms.anthropic.messages.native import handler as main
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
from litellm.rust_bridge import detectors
from litellm.rust_bridge.catalog import Feature, Route, RouteContext, Target
from litellm.rust_bridge.dispatch import PublicDispatch
@ -18,6 +19,7 @@ from litellm.rust_bridge.public_call import (
AsyncFlag,
NativeCall,
bind,
native_acall_hook,
native_call,
native_call_hook,
optional_sequence,
@ -25,6 +27,7 @@ from litellm.rust_bridge.public_call import (
signature,
)
from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse
from litellm.types.router import GenericLiteLLMParams
__all__ = ("anthropic_messages", "anthropic_messages_handler")
@ -55,10 +58,20 @@ _AMESSAGES: Final = signature(_PYTHON_AMESSAGES)
def _native_kwargs(fields: Mapping[str, object], kwargs: Mapping[str, object]) -> Mapping[str, object]:
stream_timeout: Final = kwargs.get("stream_timeout")
if fields.get("stream") is not True or stream_timeout is None:
"""The native call takes one timeout, resolved from timeout, request_timeout and stream_timeout as Python's."""
_, provider = _resolved_model(str(fields["model"]), optional_str(fields.get("custom_llm_provider")))
timeout: Final = BaseLLMHTTPHandler._resolve_anthropic_messages_timeout( # pyright: ignore[reportPrivateUsage] # the handler's own timeout resolution
litellm_params=GenericLiteLLMParams(
timeout=kwargs.get("timeout"),
request_timeout=kwargs.get("request_timeout"),
stream_timeout=kwargs.get("stream_timeout"),
),
stream=fields.get("stream") is True,
custom_llm_provider=provider or "",
)
if timeout is None:
return kwargs
return {**kwargs, "timeout": stream_timeout}
return {**kwargs, "timeout": timeout}
def _public_request(
@ -122,7 +135,7 @@ def messages_target(model: str, custom_llm_provider: str | None, model_info: obj
_FEATURES: Final[detectors.Detectors] = MappingProxyType(
{
**detectors.COMMON,
**{feature: detect for feature, detect in detectors.COMMON.items() if feature is not Feature.TIMEOUT_VARIANTS},
Feature.PROVIDER_AFFINITY: route_detectors.provider_affinity,
Feature.SAFEGUARDS_CLASSIFIER: route_detectors.safeguards_classifier,
Feature.RESPONSE_CACHE: detectors.response_cache,
@ -195,7 +208,7 @@ async def anthropic_messages(*args: object, **kwargs: object) -> MessagesResult:
_FLAG.dropped(kwargs),
python=python,
binding=NATIVE_AMESSAGES,
native=native_call_hook,
native=native_acall_hook,
)

850
litellm/model_prices_and_context_window_backup.json
View file

File diff suppressed because it is too large Load diff

3
litellm/ocr/dispatch.py
View file

@ -12,6 +12,7 @@ from litellm.rust_bridge.ocr.entrypoints import NATIVE_AOCR, NATIVE_OCR
from litellm.rust_bridge.public_call import (
AsyncFlag,
NativeCall,
native_acall_hook_on_loop,
native_call,
native_call_hook,
optional_str,
@ -97,5 +98,5 @@ async def aocr(*args: object, **kwargs: object) -> OCRResponse: # kwargs-ok: pr
_FLAG.dropped(kwargs),
python=runtime.NO_PYTHON,
binding=NATIVE_AOCR,
native=native_call_hook,
native=native_acall_hook_on_loop,
)

1
litellm/proxy/_lazy_features.py
View file

@ -243,6 +243,7 @@ LAZY_FEATURES: Final[tuple[LazyFeature, ...]] = (
"/mistral/",
"/nvidia_nim/",
"/openai/",
"/opensearch/",
"/openai_passthrough/",
"/tinyfish/",
"/transcribe",

380
litellm/proxy/_lazy_openapi_snapshot.json
View file

@ -35987,6 +35987,386 @@
]
}
},
"/opensearch/{endpoint}": {
"delete": {
"description": "Call an OpenSearch index through the proxy, e.g. `POST /opensearch/<index>/_search`.\n\n`<index>` is a LiteLLM vector store index on an OpenSearch vector store. The key or its team needs `read` or\n`write` on it in `allowed_vector_store_indexes`, and LiteLLM signs the call with the vector store's AWS identity.",
"operationId": "opensearch_proxy_route_opensearch__endpoint__delete",
"parameters": [
{
"in": "path",
"name": "endpoint",
"required": true,
"schema": {
"title": "Endpoint",
"type": "string"
}
}
],
"requestBody": {
"content": {
"application/json": {
"schema": {
"additionalProperties": true,
"type": "object"
}
},
"application/x-ndjson": {
"schema": {
"type": "string"
}
}
},
"required": false
},
"responses": {
"200": {
"content": {
"application/json": {
"schema": {
"additionalProperties": true,
"type": "object"
}
}
},
"description": "OpenSearch's response, unmodified"
},
"422": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/HTTPValidationError"
}
}
},
"description": "Validation Error"
}
},
"security": [
{
"APIKeyHeader": []
}
],
"summary": "Opensearch Proxy Route",
"tags": [
"llm_passthrough"
]
},
"get": {
"description": "Call an OpenSearch index through the proxy, e.g. `POST /opensearch/<index>/_search`.\n\n`<index>` is a LiteLLM vector store index on an OpenSearch vector store. The key or its team needs `read` or\n`write` on it in `allowed_vector_store_indexes`, and LiteLLM signs the call with the vector store's AWS identity.",
"operationId": "opensearch_proxy_route_opensearch__endpoint__get",
"parameters": [
{
"in": "path",
"name": "endpoint",
"required": true,
"schema": {
"title": "Endpoint",
"type": "string"
}
}
],
"requestBody": {
"content": {
"application/json": {
"schema": {
"additionalProperties": true,
"type": "object"
}
},
"application/x-ndjson": {
"schema": {
"type": "string"
}
}
},
"required": false
},
"responses": {
"200": {
"content": {
"application/json": {
"schema": {
"additionalProperties": true,
"type": "object"
}
}
},
"description": "OpenSearch's response, unmodified"
},
"422": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/HTTPValidationError"
}
}
},
"description": "Validation Error"
}
},
"security": [
{
"APIKeyHeader": []
}
],
"summary": "Opensearch Proxy Route",
"tags": [
"llm_passthrough"
]
},
"head": {
"description": "Call an OpenSearch index through the proxy, e.g. `POST /opensearch/<index>/_search`.\n\n`<index>` is a LiteLLM vector store index on an OpenSearch vector store. The key or its team needs `read` or\n`write` on it in `allowed_vector_store_indexes`, and LiteLLM signs the call with the vector store's AWS identity.",
"operationId": "opensearch_proxy_route_opensearch__endpoint__head",
"parameters": [
{
"in": "path",
"name": "endpoint",
"required": true,
"schema": {
"title": "Endpoint",
"type": "string"
}
}
],
"requestBody": {
"content": {
"application/json": {
"schema": {
"additionalProperties": true,
"type": "object"
}
},
"application/x-ndjson": {
"schema": {
"type": "string"
}
}
},
"required": false
},
"responses": {
"200": {
"content": {
"application/json": {
"schema": {
"additionalProperties": true,
"type": "object"
}
}
},
"description": "OpenSearch's response, unmodified"
},
"422": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/HTTPValidationError"
}
}
},
"description": "Validation Error"
}
},
"security": [
{
"APIKeyHeader": []
}
],
"summary": "Opensearch Proxy Route",
"tags": [
"llm_passthrough"
]
},
"patch": {
"description": "Call an OpenSearch index through the proxy, e.g. `POST /opensearch/<index>/_search`.\n\n`<index>` is a LiteLLM vector store index on an OpenSearch vector store. The key or its team needs `read` or\n`write` on it in `allowed_vector_store_indexes`, and LiteLLM signs the call with the vector store's AWS identity.",
"operationId": "opensearch_proxy_route_opensearch__endpoint__patch",
"parameters": [
{
"in": "path",
"name": "endpoint",
"required": true,
"schema": {
"title": "Endpoint",
"type": "string"
}
}
],
"requestBody": {
"content": {
"application/json": {
"schema": {
"additionalProperties": true,
"type": "object"
}
},
"application/x-ndjson": {
"schema": {
"type": "string"
}
}
},
"required": false
},
"responses": {
"200": {
"content": {
"application/json": {
"schema": {
"additionalProperties": true,
"type": "object"
}
}
},
"description": "OpenSearch's response, unmodified"
},
"422": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/HTTPValidationError"
}
}
},
"description": "Validation Error"
}
},
"security": [
{
"APIKeyHeader": []
}
],
"summary": "Opensearch Proxy Route",
"tags": [
"llm_passthrough"
]
},
"post": {
"description": "Call an OpenSearch index through the proxy, e.g. `POST /opensearch/<index>/_search`.\n\n`<index>` is a LiteLLM vector store index on an OpenSearch vector store. The key or its team needs `read` or\n`write` on it in `allowed_vector_store_indexes`, and LiteLLM signs the call with the vector store's AWS identity.",
"operationId": "opensearch_proxy_route_opensearch__endpoint__post",
"parameters": [
{
"in": "path",
"name": "endpoint",
"required": true,
"schema": {
"title": "Endpoint",
"type": "string"
}
}
],
"requestBody": {
"content": {
"application/json": {
"schema": {
"additionalProperties": true,
"type": "object"
}
},
"application/x-ndjson": {
"schema": {
"type": "string"
}
}
},
"required": false
},
"responses": {
"200": {
"content": {
"application/json": {
"schema": {
"additionalProperties": true,
"type": "object"
}
}
},
"description": "OpenSearch's response, unmodified"
},
"422": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/HTTPValidationError"
}
}
},
"description": "Validation Error"
}
},
"security": [
{
"APIKeyHeader": []
}
],
"summary": "Opensearch Proxy Route",
"tags": [
"llm_passthrough"
]
},
"put": {
"description": "Call an OpenSearch index through the proxy, e.g. `POST /opensearch/<index>/_search`.\n\n`<index>` is a LiteLLM vector store index on an OpenSearch vector store. The key or its team needs `read` or\n`write` on it in `allowed_vector_store_indexes`, and LiteLLM signs the call with the vector store's AWS identity.",
"operationId": "opensearch_proxy_route_opensearch__endpoint__put",
"parameters": [
{
"in": "path",
"name": "endpoint",
"required": true,
"schema": {
"title": "Endpoint",
"type": "string"
}
}
],
"requestBody": {
"content": {
"application/json": {
"schema": {
"additionalProperties": true,
"type": "object"
}
},
"application/x-ndjson": {
"schema": {
"type": "string"
}
}
},
"required": false
},
"responses": {
"200": {
"content": {
"application/json": {
"schema": {
"additionalProperties": true,
"type": "object"
}
}
},
"description": "OpenSearch's response, unmodified"
},
"422": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/HTTPValidationError"
}
}
},
"description": "Validation Error"
}
},
"security": [
{
"APIKeyHeader": []
}
],
"summary": "Opensearch Proxy Route",
"tags": [
"llm_passthrough"
]
}
},
"/tinyfish/{endpoint}": {
"get": {
"description": "Pass-through for the TinyFish Agent API (goal-based web automation).\n\nForwarded endpoints:\n- POST /v1/automation/run \u2014 run to completion (blocking)\n- POST /v1/automation/run-async \u2014 submit a run, poll GET /v1/runs/{id} for the result\n- POST /v1/automation/run-sse \u2014 run with SSE progress events\n- GET /v1/runs/{id} \u2014 run status / result\n- POST /v1/runs/{id}/cancel \u2014 cancel a run\n\nEvery other Agent API endpoint (vault, wallet, browser profiles, and the GET /v1/runs\nlisting, which would let any caller discover other callers' run ids) returns 403: all\nproxy callers share one upstream key.\n\nCredential lookup order:\n1. passthrough_endpoint_router (config.yaml deployments with use_in_pass_through)\n2. TINYFISH_API_KEY environment variable\n\n[Docs](https://docs.litellm.ai/docs/pass_through/tinyfish)",

23
litellm/proxy/_types.py
View file

@ -37,6 +37,7 @@ from litellm.litellm_core_utils.initialize_dynamic_callback_params import (
validate_langfuse_span_scope_value,
validate_no_callback_env_reference,
)
from litellm.proxy.auth.access_group_rules import AccessGroupRule, _validate_rules
from litellm.proxy.mcp_server.stdio_gate import MCP_STDIO_DISABLED_MESSAGE, is_mcp_stdio_enabled
from litellm.types.agents import (
AgentCaller as AgentCaller, # noqa: PLC0414 # public re-export
@ -427,6 +428,7 @@ class LiteLLMRoutes(enum.Enum):
"/bespoke",
"/openrouter",
"/milvus",
"/opensearch",
"/gigachat",
"/watsonx",
"/nvidia_nim",
@ -2099,6 +2101,11 @@ class BlockModelRequest(LiteLLMPydanticObjectBase):
model_id: str # required
CALLBACK_VAR_NAMES: Final = frozenset(
{*StandardCallbackDynamicParams.__annotations__, CAPTURE_MESSAGE_CONTENT_VAR, EXCLUDED_SERVICES_VAR}
)
class AddTeamCallback(LiteLLMPydanticObjectBase):
callback_name: str
callback_type: Literal["success", "failure", "success_and_failure"] | None = "success_and_failure"
@ -2108,14 +2115,9 @@ class AddTeamCallback(LiteLLMPydanticObjectBase):
@classmethod
def validate_callback_vars(cls, values):
callback_vars: Final = values.get("callback_vars", {})
valid_keys: Final = {
*StandardCallbackDynamicParams.__annotations__,
CAPTURE_MESSAGE_CONTENT_VAR,
EXCLUDED_SERVICES_VAR,
}
for key, value in callback_vars.items():
if key not in valid_keys:
raise ValueError(f"Invalid callback variable: {key}. Must be one of {valid_keys}")
if key not in CALLBACK_VAR_NAMES:
raise ValueError(f"Invalid callback variable: {key}. Must be one of {set(CALLBACK_VAR_NAMES)}")
if key in (CAPTURE_MESSAGE_CONTENT_VAR, EXCLUDED_SERVICES_VAR) and not isinstance(value, str):
raise ValueError(f"{key} must be a string, got {type(value).__name__}")
callback_vars[key] = str(value)
@ -2506,6 +2508,13 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
Documents all the fields supported by `general_settings` in config.yaml
"""
access_group_rules: tuple[AccessGroupRule, ...] | None = None
@field_validator("access_group_rules", mode="before")
@classmethod
def validate_access_group_rules(cls, value: object) -> tuple[AccessGroupRule, ...] | None:
return _validate_rules(value) if value is not None else None
use_spend_logs_ledger: bool = Field(
True,
description="Apply request accounting from stored SpendLogs deliveries. Set to false to roll back to the previous in-memory and Redis spend queues.",

201
litellm/proxy/auth/access_group_rules.py Normal file
View file

@ -0,0 +1,201 @@
from __future__ import annotations
import json
from collections.abc import Callable, Mapping, Sequence
from functools import lru_cache, reduce
from importlib import import_module
from typing import TYPE_CHECKING, Final, Literal, Protocol, runtime_checkable
from pydantic import BaseModel, ConfigDict, TypeAdapter, model_validator
from litellm._internal_context import with_service_target
from litellm._logging import verbose_proxy_logger
from litellm.constants import AUTH_OBJECTS_TARGET
_TRUSTED_ATTRIBUTES: Final = frozenset(("jwt_claims", "user_id", "user_role"))
if TYPE_CHECKING:
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.repositories.prisma_protocols import DatabaseClient
class AccessGroupRule(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid", strict=True)
attribute: str
contains: str | int | float | bool | None = None
equals: str | int | float | bool | None = None
access_group: str
@model_validator(mode="after")
def validate_operator(self) -> AccessGroupRule:
if ("contains" in self.model_fields_set) == ("equals" in self.model_fields_set):
raise ValueError("An access group rule requires exactly one of contains or equals")
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
if not self.attribute or any(not part for part in self.attribute.split(".")) or not self.access_group:
raise ValueError("An access group rule requires a nonempty attribute path and access_group")
if self.attribute.split(".")[0] not in _TRUSTED_ATTRIBUTES:
raise ValueError(
"An access group rule attribute must start with an administrator- or IdP-controlled field: "
+ ", ".join(sorted(_TRUSTED_ATTRIBUTES))
)
if self.attribute.split(".")[0] not in UserAPIKeyAuth.model_fields:
raise ValueError("An access group rule attribute must start with a UserAPIKeyAuth field")
return self
class _Enforcer(Protocol):
def enforce(self, subject: str, resource: str) -> bool: ...
def add_policy(self, subject: str, resource: str) -> bool: ...
class _CasbinModel(Protocol):
def load_model_from_text(self, text: str) -> None: ...
@runtime_checkable
class _Casbin(Protocol):
Model: Callable[[], _CasbinModel]
Enforcer: Callable[[_CasbinModel], _Enforcer]
_MODEL: Final = """
[request_definition]
r = sub, obj
[policy_definition]
p = sub, obj
[policy_effect]
e = some(where (p.eft == allow))
[matchers]
m = r.sub == p.sub && r.obj == p.obj
"""
_RULES: Final = TypeAdapter(tuple[AccessGroupRule, ...])
_ATTRIBUTES: Final = TypeAdapter(dict[str, object])
_LIST: Final = TypeAdapter(list[object])
_CASBIN: Final = TypeAdapter(_Casbin, config=ConfigDict(arbitrary_types_allowed=True))
_MISSING: Final = object()
def _subject(operator: Literal["contains", "equals"], claim: str, value: object) -> str:
return json.dumps((operator, claim, value), ensure_ascii=True, separators=(",", ":"))
def _attribute_segment(value: object, segment: str) -> object:
if not isinstance(value, Mapping):
return _MISSING
return _ATTRIBUTES.validate_python(value).get(segment, _MISSING)
def _claim_value(attributes: Mapping[str, object], path: Sequence[str]) -> object:
return reduce(_attribute_segment, path, attributes)
def _subjects(rule: AccessGroupRule, identity: UserAPIKeyAuth) -> tuple[str, ...]:
first, *nested = rule.attribute.split(".")
root: Final[object] = getattr(identity, first, _MISSING) if first in type(identity).model_fields else _MISSING
value: Final = (
_claim_value(_ATTRIBUTES.validate_python(root), nested)
if nested and isinstance(root, Mapping)
else _MISSING
if nested
else root
)
if "contains" in rule.model_fields_set:
if not isinstance(value, list):
return ()
values: Final = _LIST.validate_python(value)
return tuple(_subject("contains", rule.attribute, item) for item in values if _is_scalar(item))
if not _is_scalar(value):
return ()
return (_subject("equals", rule.attribute, value),)
def _is_scalar(value: object) -> bool:
return isinstance(value, (str, int, float, bool)) or value is None
@lru_cache(maxsize=1)
def _enforcer(policies: tuple[tuple[str, str], ...]) -> _Enforcer:
casbin: Final = _CASBIN.validate_python(import_module("casbin"))
model: Final = casbin.Model()
model.load_model_from_text(_MODEL)
enforcer: Final = casbin.Enforcer(model)
for subject, resource in policies:
enforcer.add_policy(subject, resource)
return enforcer
def _policy(rule: AccessGroupRule) -> tuple[str, str]:
operator: Final = "contains" if "contains" in rule.model_fields_set else "equals"
value: Final = rule.contains if operator == "contains" else rule.equals
return _subject(operator, rule.attribute, value), rule.access_group
def _validate_rules(value: object) -> tuple[AccessGroupRule, ...]:
rules: Final = _RULES.validate_python(value)
if rules:
_enforcer(tuple(_policy(rule) for rule in rules))
return rules
@with_service_target(AUTH_OBJECTS_TARGET)
async def _resolve_name(name: str, prisma_client: DatabaseClient | None, cache: UserApiKeyCache) -> str | None:
from litellm.constants import DEFAULT_ACCESS_GROUP_CACHE_TTL
from litellm.proxy._types import LiteLLM_AccessGroupTable
from litellm.proxy.auth.auth_checks import get_access_group_tables
from litellm.repositories.table_repositories import AccessGroupRepository
if prisma_client is None:
verbose_proxy_logger.debug("Access group rule cannot resolve %s without a database", name)
return None
cache_key: Final = f"access_group_rule_name:{name}"
cached: Final = await cache.async_get_cache(key=cache_key, model_type=LiteLLM_AccessGroupTable)
if cached is not None:
groups: Final = await get_access_group_tables(
(cached.access_group_id,), prisma_client=prisma_client, user_api_key_cache=cache
)
if groups and groups[0].access_group_name == name:
return cached.access_group_id
group: Final = await AccessGroupRepository(prisma_client).find_by_name(name)
if group is None:
verbose_proxy_logger.debug("Access group rule references unknown access group %s", name)
return None
await cache.async_set_cache(
key=cache_key, value=group, model_type=LiteLLM_AccessGroupTable, ttl=DEFAULT_ACCESS_GROUP_CACHE_TTL
)
return group.access_group_id
async def access_group_ids_for_identity(
identity: UserAPIKeyAuth,
*,
rules: tuple[AccessGroupRule, ...] | None = None,
prisma_client: DatabaseClient | None = None,
cache: UserApiKeyCache | None = None,
) -> tuple[str, ...]:
if rules is None:
from litellm.proxy import proxy_server
settings: Final = _ATTRIBUTES.validate_python(proxy_server.general_settings) # pyright: ignore[reportUnknownMemberType] # legacy settings are validated at this boundary
configured: Final = settings.get("access_group_rules")
if not configured:
return ()
return await access_group_ids_for_identity(
identity,
rules=_validate_rules(configured),
prisma_client=proxy_server.prisma_client,
cache=proxy_server.user_api_key_cache,
)
if not rules or cache is None:
return ()
enforcer: Final = _enforcer(tuple(_policy(rule) for rule in rules))
names: Final = frozenset(
rule.access_group
for rule in rules
if any(enforcer.enforce(subject, rule.access_group) for subject in _subjects(rule, identity))
)
ids: Final = tuple([await _resolve_name(name, prisma_client, cache) for name in sorted(names)])
return tuple(group_id for group_id in ids if group_id is not None)

30
litellm/proxy/auth/auth_checks.py
View file

@ -5839,7 +5839,7 @@ async def get_authorized_resources_from_key_access_groups(
continue
team_authorized = bool(key_team_id and key_team_id in (ag.assigned_team_ids or []))
key_authorized = bool(key_token and key_token in (ag.assigned_key_ids or []))
if team_authorized or key_authorized:
if team_authorized or key_authorized or ag_id in valid_token.access_group_rule_grants:
authorized_resources.extend(getattr(ag, resource_field, []) or [])
return list(set(authorized_resources))
@ -7076,6 +7076,27 @@ def team_member_max_budget_alert_check(
_team_member_max_budget_alert_check: Final = team_member_max_budget_alert_check
def _raise_team_member_without_membership(user_id: str, team_id: str) -> None:
from litellm.proxy.litellm_pre_call_utils import sanitize_for_log
verbose_proxy_logger.warning(
"fail_closed_budget_enforcement: rejecting request, user_id=%s has no membership row in team_id=%s "
"so spend against the team member budget cannot be verified",
sanitize_for_log(user_id),
sanitize_for_log(team_id),
)
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail={
"error": (
f"User={user_id} has no membership in Team={team_id}, so their spend against the team member "
"budget cannot be verified and fail_closed_budget_enforcement is enabled. Add them to the team "
"with POST /team/member_add."
)
},
)
async def _check_team_member_budget(
team_object: LiteLLM_TeamTable | None,
user_object: LiteLLM_UserTable | None,
@ -7125,11 +7146,14 @@ async def _check_team_member_budget(
)
if team_member_budget is not None:
from litellm.proxy.proxy_server import fail_closed_budget_enforcement_enabled, get_current_spend
if loaded_membership is None and fail_closed_budget_enforcement_enabled():
_raise_team_member_without_membership(user_id=valid_token.user_id, team_id=team_object.team_id)
team_member_spend = (loaded_membership.spend if loaded_membership is not None else 0.0) or 0.0
# Read from cross-pod counter (Redis-first) if available
from litellm.proxy.proxy_server import get_current_spend
team_member_spend = await get_current_spend(
counter_key=team_member_counter(valid_token.user_id, team_object.team_id),
fallback_spend=team_member_spend,

44
litellm/proxy/auth/auth_utils.py
View file

@ -2,7 +2,7 @@ import importlib.util
import os
import re
import sys
from collections.abc import Collection, Iterator, Mapping
from collections.abc import Collection, Iterator, Mapping, Sequence
from functools import lru_cache
from itertools import accumulate, chain, repeat
from logging import Logger
@ -1556,6 +1556,48 @@ def enforce_batch_limits_are_admin_only(
)
VECTOR_STORE_INDEX_GRANTS_METADATA_KEY: Final = "allowed_vector_store_indexes"
def _stored_vector_store_index_grants(existing_metadata: Mapping[str, object] | None) -> object:
if existing_metadata is None:
return None
return existing_metadata.get(VECTOR_STORE_INDEX_GRANTS_METADATA_KEY) or None
def _requested_vector_store_index_grants(
data: BatchLimitRequest,
top_level_grants: Sequence[AllowedVectorStoreIndexItem] | None,
stored_grants: object,
) -> object:
if top_level_grants is not None:
return [grant.model_dump() for grant in top_level_grants] or None
if "metadata" in data.model_fields_set:
return (data.metadata or EMPTY_MAPPING).get(VECTOR_STORE_INDEX_GRANTS_METADATA_KEY) or None
return stored_grants
def enforce_vector_store_index_grants_are_admin_only(
data: BatchLimitRequest,
top_level_grants: Sequence[AllowedVectorStoreIndexItem] | None,
existing_metadata: Mapping[str, object] | None,
user_api_key_dict: UserAPIKeyAuth,
entity: Literal["key", "team"],
) -> None:
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value:
return
stored_grants: Final = _stored_vector_store_index_grants(existing_metadata)
if _requested_vector_store_index_grants(data, top_level_grants, stored_grants) == stored_grants:
return
raise HTTPException(
status_code=403,
detail={
"error": f"Only proxy admins can set {VECTOR_STORE_INDEX_GRANTS_METADATA_KEY} on a {entity}. "
"It grants read and write access to vector store indexes, so the holder cannot widen it."
},
)
def get_model_rate_limit_from_metadata(
user_api_key_dict: UserAPIKeyAuth,
metadata_accessor_key: Literal["team_metadata", "organization_metadata", "project_metadata"],

48
litellm/proxy/auth/handle_jwt.py
View file

@ -1768,6 +1768,7 @@ class JWTAuthManager:
denied_pass_through_team: LiteLLM_TeamTable | None = None
route_denied_team_id: str | None = None # rebind-ok: retain route-denied team while checking all claims
fallback_team: tuple[str, LiteLLM_TeamTable] | None = None # rebind-ok: retain first route-allowed candidate
if not team_ids:
if (
@ -1796,6 +1797,21 @@ class JWTAuthManager:
if team_object and team_object.models is not None:
team_models = team_object.models
if (
fallback_team is None
and allowed_routes_check(
user_role=LitellmUserRoles.TEAM,
user_route=route,
litellm_proxy_roles=jwt_handler.litellm_jwtauth,
)
and JWTAuthManager._team_has_passthrough_route_access(
team_object=team_object,
route=route,
request_method=request_method,
team_allowed_routes=jwt_handler.litellm_jwtauth.team_allowed_routes,
)
):
fallback_team = (team_id, team_object)
if isinstance(team_models, list) and (
not requested_model
or await can_team_access_model(
@ -1828,6 +1844,9 @@ class JWTAuthManager:
except Exception:
continue
if fallback_team is not None:
return fallback_team
if denied_pass_through_team is not None:
JWTAuthManager._raise_team_passthrough_route_denial(route=route, team_object=denied_pass_through_team)
@ -2380,6 +2399,7 @@ class JWTAuthManager:
route=route, request_method=request_method, jwt_handler=jwt_handler
)
any_team_resolved = False
fallback_team: tuple[str, LiteLLM_TeamTable] | None = None # rebind-ok: retain first route-allowed membership
for candidate_team_id in user_team_ids:
try:
team_object = await get_team_object(
@ -2393,6 +2413,17 @@ class JWTAuthManager:
except HTTPException:
continue
any_team_resolved = True
if (
team_route_allowed
and fallback_team is None
and JWTAuthManager._team_has_passthrough_route_access(
team_object=team_object,
route=route,
request_method=request_method,
team_allowed_routes=jwt_handler.litellm_jwtauth.team_allowed_routes,
)
):
fallback_team = (candidate_team_id, team_object)
if requested_model:
try:
await can_team_access_model(
@ -2424,6 +2455,23 @@ class JWTAuthManager:
)
return candidate_team_id, team_object, None
if fallback_team is not None:
fallback_id, fallback_object = fallback_team
return (
fallback_id,
fallback_object,
await get_team_membership(
user_id=user_id,
team_id=fallback_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
)
if user_id
else None,
)
if enforce_team_based_model_access:
if requested_model and any_team_resolved:
raise HTTPException(

103
litellm/proxy/auth/user_api_key_auth.py
View file

@ -14,6 +14,7 @@ import secrets
from collections.abc import Mapping, Sequence
from datetime import datetime, timezone
from enum import Enum
from functools import partial
from types import MappingProxyType
from typing import Final, NamedTuple, Protocol, Union, cast
@ -44,6 +45,7 @@ from litellm.litellm_core_utils.routing_common_utils import resolve_model_group_
from litellm.models.end_user import LiteLLM_EndUserTable
from litellm.proxy._types import *
from litellm.proxy.agent_endpoints.auth.agent_caller import agent_caller_from_headers
from litellm.proxy.auth.access_group_rules import access_group_ids_for_identity
from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module exports
ExperimentalUIJWTToken,
TeamNotFoundError,
@ -144,7 +146,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # leg
_safe_get_request_headers, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
_safe_get_request_query_params, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
_safe_set_request_parsed_body, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
is_opaque_audio_pass_through_request,
is_opaque_pass_through_request,
populate_request_with_path_params,
read_raw_json_body,
read_request_body,
@ -170,7 +172,15 @@ from litellm.proxy.db.db_lookup_gate import bounded_db_lookup
from litellm.proxy.db.db_span import db_span
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
from litellm.proxy.spend_tracking.carried_budget_state import carry_team_and_user_budget_state
from litellm.proxy.spend_tracking.carried_budget_state import (
carry_committed_budget_spend,
carry_team_and_user_budget_state,
)
from litellm.proxy.spend_tracking.committed_spend import (
bind_committed_admission_counters,
committed_spend_tracked,
read_admission_committed_spend,
)
from litellm.proxy.spend_tracking.spend_counter_batch import (
bind_admission_counter_keys,
release_spend_counter_batch,
@ -397,6 +407,17 @@ def _get_model_names_for_budget_checks(
return model
async def _normalize_model_for_auth(
request_data: dict[str, object], # mutable-ok: normalization updates dispatch data
valid_token: UserAPIKeyAuth,
request: Request | None,
route: str,
) -> None:
await _normalize_claude_model(request_data, valid_token, request, route)
await _resolve_router_settings_model_group_alias(request_data, valid_token, request, route)
_resolve_codex_native_chatgpt_model(request_data, valid_token, request, route)
def _budget_fallback_target_model(target: object) -> str | None:
if isinstance(target, str):
return target
@ -2128,7 +2149,7 @@ async def _read_request_body_deferring_parse_failure(
must run (resolving identity onto the request's trace) before the 400 goes
out; the caller re-raises the returned exception once identity is seeded.
"""
if is_opaque_audio_pass_through_request(
if is_opaque_pass_through_request(
route=get_request_route(request=request),
content_type=safe_get_request_headers(request=request).get("content-type", ""),
):
@ -3044,14 +3065,7 @@ async def validate_resolved_virtual_key( # noqa: C901 # Preserve ordering of e
param=None,
code=status.HTTP_401_UNAUTHORIZED,
)
await enforce_key_and_fallback_model_access(
valid_token=valid_token,
request_data=request_data,
route=route,
request=request,
llm_model_list=llm_model_list,
llm_router=llm_router,
)
await _normalize_model_for_auth(request_data, valid_token, request, route)
await _prefetch_referenced_auth_objects(
valid_token, end_user_id=end_user_id, user_api_key_cache=user_api_key_cache, prisma_client=prisma_client
)
@ -3405,6 +3419,7 @@ async def validate_resolved_virtual_key( # noqa: C901 # Preserve ordering of e
start_time=start_time,
)
virtual_key_auth_obj.via_virtual_key = True
request.state.requires_key_model_check = True # rebind-ok: defer until central grants
return virtual_key_auth_obj
@ -3650,6 +3665,7 @@ async def run_centralized_common_checks(
model_max_budget_limiter,
prisma_client,
proxy_logging_obj,
spend_counter_store,
user_api_key_cache,
user_custom_auth,
)
@ -3960,6 +3976,7 @@ async def run_centralized_common_checks(
)
bind_admission_counter_keys(user_api_key_auth_obj, end_user_id=end_user_id)
bind_committed_admission_counters(user_api_key_auth_obj, end_user_id=end_user_id)
try:
_ = await common_checks(
request=request,
@ -4035,6 +4052,18 @@ async def run_centralized_common_checks(
user_obj=user_object,
)
if committed_spend_tracked():
await carry_committed_budget_spend(
valid_token=user_api_key_auth_obj,
team_object=typed_team_object,
user_object=user_object,
end_user_object=end_user_object,
read_committed=(
None
if route_skips_budget_checks(route=route)
else partial(read_admission_committed_spend, spend_counter_store)
),
)
await _reserve_budget_after_common_checks(
user_api_key_auth_obj=user_api_key_auth_obj,
request=request,
@ -4228,7 +4257,33 @@ async def _authorize_authenticated_request(
prepare_agent_invocation,
)
from litellm.proxy.agent_endpoints.identity_store import AgentIdentityStore
from litellm.proxy.proxy_server import general_settings, prisma_client, user_model
from litellm.proxy.proxy_server import (
general_settings,
llm_model_list, # pyright: ignore[reportUnknownVariableType] # proxy global has untyped deployment entries
llm_router,
prisma_client,
user_model,
)
grants: Final = await access_group_ids_for_identity(user_api_key_auth_obj)
user_api_key_auth_obj.apply_access_group_rule_grants(grants)
key_check_required: Final = (
getattr(request.state, "requires_key_model_check", False) or force_virtual_key_checks
)
if key_check_required or grants:
model_identity: Final = (
user_api_key_auth_obj
if key_check_required
else user_api_key_auth_obj.model_copy(update={"models": list(user_api_key_auth_obj.team_models or ())})
)
await enforce_key_and_fallback_model_access(
valid_token=model_identity,
request_data=request_data,
route=route,
request=request,
llm_model_list=llm_model_list,
llm_router=llm_router,
)
store: Final = AgentIdentityStore.from_client(prisma_client) if prisma_client is not None else None
if user_api_key_auth_obj.agent_id is not None:
@ -4345,9 +4400,9 @@ async def _prefetch_referenced_auth_objects(
user_api_key_cache: UserApiKeyCache,
prisma_client: PrismaClient | None,
) -> None:
"""Warm every object and spend counter the checks below will read, in one MGET each (one DB query when cold).
Runs after the key's model access check so a denied request costs no more than it did before."""
"""Warm every object and spend counter the checks below will read, in one MGET each (one DB query when cold)."""
bind_admission_counter_keys(valid_token, end_user_id=end_user_id or None)
bind_committed_admission_counters(valid_token, end_user_id=end_user_id or None)
await prefetch_auth_objects(
refs=AuthObjectRefs.from_token(valid_token),
user_api_key_cache=user_api_key_cache,
@ -4438,7 +4493,7 @@ async def user_api_key_auth(
spend_counter_batch_scope(_spend_counter_store()),
):
try:
user_api_key_auth_obj: Final = await user_api_key_auth_builder(
authenticated: Final = await user_api_key_auth_builder(
request=request,
api_key=api_key,
azure_api_key_header=azure_api_key_header,
@ -4455,6 +4510,14 @@ async def user_api_key_auth(
if body_parse_exception is not None:
raise body_parse_exception
raise
from litellm.proxy.proxy_server import (
general_settings, # pyright: ignore[reportUnknownVariableType] # proxy global is validated below
)
settings: Final = _GENERAL_SETTINGS_ADAPTER.validate_python(general_settings)
user_api_key_auth_obj: Final = (
authenticated.model_copy() if settings.get("access_group_rules") else authenticated
)
user_api_key_auth_obj.budget_reservation = None
user_api_key_auth_obj.agent_caller = agent_caller_from_headers(
safe_get_request_headers(request), user_api_key_auth_obj
@ -4793,14 +4856,8 @@ async def _run_post_custom_auth_checks(
)
if general_settings.get("custom_auth_run_common_checks", False):
await enforce_key_and_fallback_model_access(
valid_token=valid_token,
request_data=request_data,
route=route,
request=request,
llm_model_list=llm_model_list,
llm_router=llm_router,
)
await _normalize_model_for_auth(request_data, valid_token, request, route)
request.state.requires_key_model_check = True # rebind-ok: defer until central grants
current_model = _get_model_from_request_context(
request_data=request_data,

11
litellm/proxy/common_request_processing.py
View file

@ -64,7 +64,6 @@ from litellm.litellm_core_utils.bug_report import (
strip_bug_report_notice,
)
from litellm.litellm_core_utils.core_helpers import (
attribute_or_none,
finite_cost,
get_or_create_metadata_bucket,
has_attributes,
@ -72,7 +71,6 @@ from litellm.litellm_core_utils.core_helpers import (
is_expected_client_error,
proxy_request_headers,
redact_nested_match_and_regex_keys,
require_attributes,
stream_cost_injection_requested,
)
from litellm.litellm_core_utils.dd_tracing import NullTracer, tracer
@ -724,10 +722,6 @@ class _ReadsHeaderValues(Protocol):
def get(self, key: str, default: str = "") -> str: ...
class _NativeLoggingRelease(Protocol):
def release(self, success: bool, /) -> object: ...
def _as_header_reader(headers: _ReadsHeaderValues) -> _ReadsHeaderValues:
return headers
@ -4384,11 +4378,6 @@ class ProxyBaseLLMRequestProcessing:
Extracted as a static method so tests can exercise the production
gating logic directly rather than reimplementing the finally block.
"""
if getattr(logging_obj, "call_type", None) in ("ocr", "aocr"):
pending: Final = attribute_or_none(logging_obj, "_native_pending_logging")
if pending is not None:
logging_obj._native_pending_logging = None # rebind-ok: consume the native OCR release signal once
require_attributes(pending, _NativeLoggingRelease, "release").release(not exception_raised)
_enqueue_fn: Final = getattr(logging_obj, "enqueue_deferred_logging", None)
if _enqueue_fn is None:
return

81
litellm/proxy/common_utils/callback_config_validation.py
View file

@ -7,9 +7,13 @@ per-integration checks here.
import math
from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass
from types import MappingProxyType
from typing import Final, cast # noqa: TID251 # preserve arbitrary keys from validated callback mappings
from typing import Final, Literal, cast # noqa: TID251 # preserve arbitrary keys from validated callback mappings
from pydantic import TypeAdapter, ValidationError
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
from litellm.types.utils import CAPTURE_MESSAGE_CONTENT_VAR, EXCLUDED_SERVICES_VAR
_NEWRELIC_CALLBACK: Final = "newrelic"
@ -276,7 +280,7 @@ def logging_metadata_config_error(metadata: Mapping[str, object] | None) -> str
"""Validate every ``logging`` entry of a team/key metadata payload."""
if not metadata:
return None
entries: Final = _object_sequence(metadata.get("logging"))
entries: Final = object_sequence(metadata.get("logging"))
if entries is None:
return None
entry_error: Final = next(
@ -302,11 +306,11 @@ def logging_metadata_config_error(metadata: Mapping[str, object] | None) -> str
def stored_callback_entries(logging_entries: object) -> tuple[tuple[str | None, Mapping[str, str]], ...]:
"""The callback name and vars of every stored ``logging`` entry."""
entries: Final = _object_sequence(logging_entries) or ()
entries: Final = object_sequence(logging_entries) or ()
return tuple((_entry_callback_name(entry), _entry_callback_vars(entry)) for entry in entries)
def _object_sequence(value: object) -> Sequence[object] | None:
def object_sequence(value: object) -> Sequence[object] | None:
return value if isinstance(value, Sequence) and not isinstance(value, (str, bytes)) else None
@ -445,3 +449,72 @@ def _newrelic_config_error(callback_vars: Mapping[str, str]) -> str | None:
if region is not None and (not api_key or api_key == "None"):
return "newrelic_region requires newrelic_api_key; the region rides the team's own key."
return None
@dataclass(frozen=True, slots=True)
class CredentialReferenceRejection:
status_code: Literal[400, 403]
message: str
def credential_is_defined(credential_name: str) -> bool:
return CredentialAccessor.find_credential(credential_name) is not None
_STRING_KEYED: Final = TypeAdapter(Mapping[str, object])
def string_keyed(value: object) -> Mapping[str, object]:
try:
return _STRING_KEYED.validate_python(value)
except ValidationError:
return MappingProxyType({})
def _logging_entries(metadata: object) -> object:
return string_keyed(metadata).get("logging")
def _entry_credential_name(entry: object) -> str | None:
match entry:
case {"litellm_credential_name": str() as credential_name}:
return credential_name
case _:
return None
def logging_credential_names(logging_entries: object) -> tuple[str, ...]:
return tuple(
name for entry in object_sequence(logging_entries) or () if (name := _entry_credential_name(entry)) is not None
)
def credential_reference_rejection(
requested_metadata: object,
stored_logging: object,
*,
is_proxy_admin: bool,
is_defined: Callable[[str], bool] = credential_is_defined,
) -> CredentialReferenceRejection | None:
"""Gate a key metadata write whose logging entries reference a proxy credential.
``requested_metadata`` is None when the write leaves the key's metadata untouched. The admin check runs before
the existence check so a non-admin cannot probe which credential names are defined.
"""
if requested_metadata is None:
return None
requested_logging: Final = _logging_entries(requested_metadata)
requested_names: Final = logging_credential_names(requested_logging)
stored_names: Final = logging_credential_names(stored_logging)
references_involved: Final = bool(requested_names or stored_names)
if references_involved and not is_proxy_admin and requested_logging != stored_logging:
return CredentialReferenceRejection(
403,
"Only proxy admins can change the logging entries of a key that uses litellm_credential_name",
)
undefined: Final = sorted({name for name in set(requested_names) - set(stored_names) if not is_defined(name)})
if undefined:
return CredentialReferenceRejection(
400, f"metadata.logging references undefined credential(s): {', '.join(undefined)}"
)
return None

45
litellm/proxy/common_utils/callback_utils.py
View file

@ -26,8 +26,10 @@ from litellm.litellm_core_utils.core_helpers import (
get_metadata_variable_name_from_kwargs,
get_or_create_metadata_bucket,
)
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker
from litellm.proxy._types import CommonProxyErrors, LiteLLMPromptInjectionParams
from litellm.proxy._types import CALLBACK_VAR_NAMES, CommonProxyErrors, LiteLLMPromptInjectionParams
from litellm.proxy.common_utils.callback_config_validation import object_sequence, string_keyed
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
decrypt_value_helper,
encrypt_value_helper,
@ -736,6 +738,47 @@ def decrypt_callback_vars(metadata: object) -> Any:
return _transform_callback_vars(metadata, _decrypt_or_passthrough)
def decrypted_callback_metadata(metadata: object) -> object:
return _transform_callback_vars(metadata, _decrypt_or_passthrough)
def decrypted_logging_entries(metadata: object) -> object:
return string_keyed(decrypted_callback_metadata(metadata)).get("logging")
def key_logging_entries(metadata: object) -> tuple[Mapping[str, object], ...] | None:
entries: Final = decrypted_logging_entries(metadata)
if entries is None:
return None
return tuple(resolved for entry in _logging_objects(entries) if (resolved := _with_credential(entry)) is not None)
def _logging_objects(entries: object) -> tuple[Mapping[str, object], ...]:
items: Final = () if entries in ("", {}) else object_sequence(entries)
if items is None or not all(isinstance(item, Mapping) for item in items):
raise TypeError(f"key metadata.logging must be a list of objects, got {type(entries).__name__}")
return tuple(string_keyed(item) for item in items)
def _with_credential(entry: Mapping[str, object]) -> Mapping[str, object] | None:
credential_name: Final = entry.get("litellm_credential_name")
if not isinstance(credential_name, str):
return entry
credential_vars: Final = {
name: str(value)
for name, value in CredentialAccessor.get_credential_values(credential_name).items()
if name in CALLBACK_VAR_NAMES and value is not None
}
if not credential_vars:
verbose_proxy_logger.error(
"Ignoring key logging entry for %s: credential %r is not defined or holds no logging values",
entry.get("callback_name"),
credential_name,
)
return None
return {**entry, "callback_vars": {**credential_vars, **string_keyed(entry.get("callback_vars"))}}
def _transform_callback_vars(metadata: object, transform: Callable[[str, object], object]) -> object:
if not isinstance(metadata, dict):
return metadata

7
litellm/proxy/common_utils/http_parsing_utils.py
View file

@ -19,6 +19,7 @@ from litellm.constants import (
LITELLM_CLIENT_DISCONNECT_DETAIL,
LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED,
MAX_REQUEST_BODY_SIZE_TO_REPAIR_MB,
OPENSEARCH_PASS_THROUGH_ROUTE_PREFIX,
)
from litellm.integrations.otel.runtime import phase_event
from litellm.litellm_core_utils.core_helpers import as_object_list, has_attributes, require_attributes
@ -404,8 +405,10 @@ async def read_request_body(request: Request | None) -> dict:
_read_request_body: Final = read_request_body
def is_opaque_audio_pass_through_request(route: str, content_type: str) -> bool:
"""Azure Speech bodies (raw audio, multipart uploads) are forwarded byte for byte, so auth must not consume them."""
def is_opaque_pass_through_request(route: str, content_type: str) -> bool:
"""Azure Speech audio and every OpenSearch body are forwarded byte for byte, so auth must not parse them."""
if route.startswith(f"{OPENSEARCH_PASS_THROUGH_ROUTE_PREFIX}/"):
return True
media_type: Final = _normalize_media_type(content_type)
return route.startswith(f"{AZURE_SPEECH_PASS_THROUGH_ROUTE_PREFIX}/") and (
media_type.startswith("audio/") or media_type == "multipart/form-data"

5
litellm/proxy/common_utils/reset_budget_job.py
View file

@ -48,6 +48,7 @@ from litellm.proxy.db.budget_window_spend_writer import roll_window_spend_row
from litellm.proxy.db.db_span import db_span, db_spanned
from litellm.proxy.db.db_transaction_queue.pod_lock_manager import POD_LOCK_TARGET, PodLockManager
from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry
from litellm.proxy.spend_tracking.committed_spend import committed_spend_keys
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.redis.stores.spend_counters import (
SPEND_COUNTERS_TARGET,
@ -621,7 +622,7 @@ class ResetBudgetJob:
from litellm.proxy.proxy_server import spend_counter_store
try:
await spend_counter_store.delete((counter_key,))
await spend_counter_store.delete((counter_key, *committed_spend_keys((counter_key,))))
except Exception as e:
verbose_proxy_logger.warning(
"Failed to reset spend counter %s in Redis: %s. Budget may be over-enforced until counter expires.",
@ -680,7 +681,7 @@ class ResetBudgetJob:
match cache:
case "spend counter":
await spend_counter_store.delete(keys)
await spend_counter_store.delete((*keys, *committed_spend_keys(keys)))
case "user_api_key_cache":
await user_api_key_cache.async_delete_cache_keys(keys)
case _:

4
litellm/proxy/db/spend_counter_reseed.py
View file

@ -94,6 +94,10 @@ class SpendCounterReseed:
_locks: ClassVar["OrderedDict[str, asyncio.Lock]"] = OrderedDict()
_registry_lock: ClassVar[asyncio.Lock | None] = None
@staticmethod
async def counter_lock(counter_key: str) -> asyncio.Lock:
return await SpendCounterReseed._get_lock(counter_key)
@staticmethod
async def _get_lock(counter_key: str) -> asyncio.Lock:
if SpendCounterReseed._registry_lock is None:

67
litellm/proxy/hooks/proxy_track_cost_callback.py
View file

@ -47,6 +47,7 @@ from litellm.proxy.db.db_spend_update_writer import (
from litellm.proxy.hooks.search_free_tier import apply_search_free_tier
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup, metadata_variable_name_for_request
from litellm.proxy.route_llm_request import ProxyModelNotFoundError
from litellm.proxy.spend_tracking.committed_spend import committed_spend_tracked, warm_committed_counter_keys
from litellm.proxy.spend_tracking.spend_counter_batch import post_call_counter_keys, spend_counter_batch_scope
from litellm.proxy.spend_tracking.spend_event import (
ObjectMapping,
@ -75,6 +76,8 @@ from litellm.proxy.utils import (
BilledFailedAttempt,
ProxyUpdateSpend,
)
from litellm.types.proxy.auth.user_api_key_auth import committed_budget_spend_of_row
from litellm.types.proxy.carried_budget_state import EntitySpendCounterKeys
from litellm.types.router import Deployment
from litellm.types.utils import (
CallTypes,
@ -268,6 +271,11 @@ def _proxy_spend_writer() -> DBSpendUpdateWriter:
return proxy_logging_obj.db_spend_update_writer
def _committed_spend_scope(user_api_key_dict: UserAPIKeyAuth) -> EntitySpendCounterKeys:
committed: Final = user_api_key_dict.committed_budget_spend
return committed.counter_keys if committed is not None else EntitySpendCounterKeys()
def _proxy_spend_counters() -> "_IncrementSpendCounters":
from litellm.proxy.proxy_server import increment_spend_counters
@ -471,15 +479,20 @@ class _ProxyDBLogger(CustomLogger):
return
request_route: Final = user_api_key_dict.request_route
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import InitPassThroughEndpointHelpers
route_is_trackable: Final = (
request_route is None
or RouteChecks.is_llm_api_route(route=request_route)
or RouteChecks.is_info_route(route=request_route)
or InitPassThroughEndpointHelpers.has_configured_route_ignoring_trailing_slash(route=request_route)
)
if (
ProxyDBLogger._should_track_errors_in_db() is False
and not completed_cost
or ProxyDBLogger._skips_pre_routing_rejections()
and is_pre_routing_rejection(original_exception, _REQUEST_DATA.validate_python(request_data))
or request_route is not None
and not (
RouteChecks.is_llm_api_route(route=request_route) or RouteChecks.is_info_route(route=request_route)
)
or not route_is_trackable
or request_route is not None
and RouteChecks.is_dashboard_session_info_route(request_route, user_api_key_dict.team_id)
):
@ -639,6 +652,7 @@ class _ProxyDBLogger(CustomLogger):
team_id=user_api_key_dict.team_id,
org_id=user_api_key_dict.org_id,
project_id=user_api_key_dict.project_id,
committed_spend_scope=_committed_spend_scope(user_api_key_dict),
kwargs=dict(row_kwargs),
completion_response=original_exception,
start_time=start_time,
@ -757,6 +771,7 @@ class _ProxyDBLogger(CustomLogger):
router=get_llm_router(),
),
project_id=user_api_key_dict.project_id,
committed_spend_scope=_committed_spend_scope(user_api_key_dict),
**_ledger_counter_options(row_kwargs.get(ACCOUNTING_LEDGER_KWARG)),
)
except Exception: # noqa: BLE001 # one row's counter failure must not keep the other rows' counters from landing
@ -1269,10 +1284,38 @@ class _IncrementSpendCounters(Protocol):
request_started_at: datetime | None = None,
model_access_groups: Sequence[str] | None = None,
project_id: str | None = None,
committed_spend_scope: EntitySpendCounterKeys | None = None,
windows_from_ledger: bool = False,
) -> None: ...
async def _committed_spend_scope_of_row(
kwargs: Mapping[str, object],
response_cost: float,
token: str | None,
team_id: str | None,
user_id: str | None,
org_id: str | None,
end_user_id: str | None,
) -> EntitySpendCounterKeys:
metadata: Final = _OBJECT_MAPPING.validate_python(get_litellm_metadata_from_kwargs(kwargs=dict(kwargs)))
carried: Final = committed_budget_spend_of_row(
metadata, _OBJECT_MAPPING.validate_python(kwargs.get("litellm_params") or {})
)
if carried is not None:
return carried.counter_keys
if response_cost <= 0 or not committed_spend_tracked():
return EntitySpendCounterKeys()
from litellm.proxy.proxy_server import spend_counter_store
return await warm_committed_counter_keys(
spend_counter_store,
EntitySpendCounterKeys.of(
token=token, team_id=team_id, user_id=user_id, org_id=org_id, end_user_id=end_user_id
),
)
async def _update_database_and_spend_counters(
proxy_logging_obj: "ProxyLogging",
increment_spend_counters: _IncrementSpendCounters,
@ -1291,6 +1334,7 @@ async def _update_database_and_spend_counters(
model_access_groups: Sequence[str] | None = None,
project_id: str | None = None,
spend_writer: DBSpendUpdateWriter | None = None,
committed_spend_scope: EntitySpendCounterKeys | None = None,
) -> bool:
"""The reservation is reconciled before the spend is persisted, from its own read. One spend counter batch then
spans the database write and the counter update, so the post-call counters are read with a single MGET after the
@ -1333,6 +1377,7 @@ async def _update_database_and_spend_counters(
request_tags=request_tags,
model_access_groups=model_access_groups,
project_id=project_id,
committed_spend_scope=committed_spend_scope,
)
@ -1354,6 +1399,7 @@ async def _update_database_and_spend_counters_in_batch(
model_access_groups: Sequence[str] | None,
project_id: str | None,
spend_writer: DBSpendUpdateWriter | None = None,
committed_spend_scope: EntitySpendCounterKeys | None = None,
) -> bool:
try:
writer: Final = spend_writer if spend_writer is not None else proxy_logging_obj.db_spend_update_writer
@ -1407,6 +1453,19 @@ async def _update_database_and_spend_counters_in_batch(
request_started_at=start_time,
model_access_groups=model_access_groups,
project_id=project_id,
committed_spend_scope=(
committed_spend_scope
if committed_spend_scope is not None
else await _committed_spend_scope_of_row(
kwargs=kwargs,
response_cost=response_cost,
token=user_api_key,
team_id=team_id,
user_id=user_id,
org_id=org_id,
end_user_id=end_user_id,
)
),
**_ledger_counter_options(_BATCH_CALLBACK_FIELDS.validate_python(kwargs).get(ACCOUNTING_LEDGER_KWARG)),
)
except Exception:

38
litellm/proxy/litellm_pre_call_utils.py
View file

@ -88,6 +88,7 @@ from litellm.proxy.auth.route_checks import RouteChecks
from litellm.proxy.common_utils.callback_utils import (
decrypt_callback_vars,
get_metadata_variable_name_from_kwargs,
key_logging_entries,
strip_callback_config,
)
from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # legacy module exports
@ -105,7 +106,7 @@ _REDACTED_HEADER_VALUE: Final = "***REDACTED***"
_DECODED_JSON: Final = TypeAdapter(object)
_OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object])
_CREDENTIAL_HEADER_NAMES: Final = SpecialHeaders.litellm_credential_header_names() | frozenset(
{"cookie", "proxy-authorization"}
{"cookie", "proxy-authorization", "x-amz-security-token"}
)
_TRANSPORT_ONLY_CREDENTIAL_KEYS: Final = frozenset({"provider_specific_header", "headers", "api_key"})
_REQUEST_FIELDS: Final = TypeAdapter(Mapping[str, object])
@ -223,12 +224,21 @@ def _caller_session_beside_trace_header(
return get_chain_id_from_headers(session_headers) or _caller_body_session_id(data, metadata_variable_name)
def _claude_code_subagent_tier_session(normalized_headers: Mapping[str, str]) -> str | None:
session: Final = normalized_headers.get("x-claude-code-session-id", "")
agent: Final = normalized_headers.get("x-claude-code-agent-id", "")
if not _SESSION_ID_VALUE_RE.fullmatch(session) or not _SESSION_ID_VALUE_RE.fullmatch(agent):
return None
return f"{session}:{agent}"
def _child_tier_session_metadata(
normalized_headers: Mapping[str, str], data: Mapping[str, object], metadata_variable_name: str
) -> Mapping[str, str | None]:
parent: Final = normalized_headers.get("x-parent-session-id")
if not parent or not _SESSION_ID_VALUE_RE.fullmatch(parent):
return {}
claude_code_subagent: Final = _claude_code_subagent_tier_session(normalized_headers)
return {COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY: claude_code_subagent} if claude_code_subagent else {}
own_headers: Final = {
key: value
for key, value in normalized_headers.items()
@ -1237,9 +1247,11 @@ def convert_key_logging_metadata_to_callback(
return team_callback_settings_obj
def get_validated_callback_metadata(item: dict, *, source: str) -> AddTeamCallback | None:
def get_validated_callback_metadata(item: object, *, source: str) -> AddTeamCallback | None:
if not isinstance(item, Mapping):
raise TypeError(f"{source} metadata.logging must be a list of objects, got an item of {type(item).__name__}")
try:
return AddTeamCallback(**item)
return AddTeamCallback.model_validate(item)
except (PydanticValidationError, ValueError) as e:
verbose_proxy_logger.warning(
"Ignoring invalid %s callback metadata: %s",
@ -1258,9 +1270,9 @@ class KeyAndTeamLoggingSettings:
"""
@staticmethod
def get_key_dynamic_logging_settings(user_api_key_dict: UserAPIKeyAuth):
def get_key_dynamic_logging_settings(user_api_key_dict: UserAPIKeyAuth) -> tuple[Mapping[str, object], ...] | None:
if user_api_key_dict.metadata is not None and "logging" in user_api_key_dict.metadata:
return decrypt_callback_vars(user_api_key_dict.metadata).get("logging")
return key_logging_entries(user_api_key_dict.metadata)
return None
@staticmethod
@ -1274,9 +1286,7 @@ def get_dynamic_logging_metadata(
user_api_key_dict: UserAPIKeyAuth, proxy_config: ProxyConfig
) -> TeamCallbackMetadata | None:
callback_settings_obj: TeamCallbackMetadata | None = None
key_dynamic_logging_settings: Final[dict | None] = KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(
user_api_key_dict
)
key_dynamic_logging_settings: Final = KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(user_api_key_dict)
team_dynamic_logging_settings: Final[dict | None] = KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings(
user_api_key_dict
)
@ -2156,7 +2166,14 @@ class LiteLLMProxyRequestSetup:
user_api_key_dict.matched_model_access_groups
)
# UserAPIKeyAuth object for MCP server access control
data[_metadata_variable_name]["user_api_key_auth"] = user_api_key_dict.model_copy(
data[_metadata_variable_name]["user_api_key_auth"] = LiteLLMProxyRequestSetup.auth_for_request_metadata(
user_api_key_dict
)
return data
@staticmethod
def auth_for_request_metadata(user_api_key_dict: UserAPIKeyAuth) -> UserAPIKeyAuth:
return user_api_key_dict.model_copy(
update={
"metadata": strip_callback_config(user_api_key_dict.metadata),
"team_metadata": strip_callback_config(user_api_key_dict.team_metadata),
@ -2164,7 +2181,6 @@ class LiteLLMProxyRequestSetup:
"organization_metadata": strip_callback_config(user_api_key_dict.organization_metadata),
}
)
return data
@staticmethod
def add_management_endpoint_metadata_to_request_metadata(

7
litellm/proxy/management_endpoints/customer_endpoints.py
View file

@ -42,7 +42,9 @@ from litellm.proxy.management_helpers.object_permission_utils import ( # noqa:
handle_update_object_permission_common,
set_object_permission,
)
from litellm.proxy.spend_tracking.committed_spend import drop_committed_spend
from litellm.proxy.utils import handle_exception_on_proxy
from litellm.redis.stores.spend_counters import end_user_counter
from litellm.repositories.budget_repository import BudgetRepository
from litellm.repositories.chunked_in import find_many_in
from litellm.repositories.table_repositories import EndUserRepository
@ -812,7 +814,7 @@ async def delete_end_user(
See below for all params
```
"""
from litellm.proxy.proxy_server import prisma_client
from litellm.proxy.proxy_server import prisma_client, spend_counter_store
try:
if prisma_client is None:
@ -842,6 +844,9 @@ async def delete_end_user(
verbose_proxy_logger.debug("received response from updating prisma client. response=%s", response)
await _evict_end_user_cache_keys(_end_user_cache_keys(data.user_ids))
await drop_committed_spend(
spend_counter_store, tuple(end_user_counter(user_id) for user_id in data.user_ids)
)
return DeleteCustomersResponse(
deleted_customers=response,

13
litellm/proxy/management_endpoints/internal_user_endpoints.py
View file

@ -39,7 +39,10 @@ from litellm.proxy.auth.auth_checks import (
get_team_object,
get_user_object,
)
from litellm.proxy.auth.auth_utils import enforce_batch_limits_are_admin_only
from litellm.proxy.auth.auth_utils import (
enforce_batch_limits_are_admin_only,
enforce_vector_store_index_grants_are_admin_only,
)
from litellm.proxy.auth.password_policy import (
validate_password_not_breached,
validate_password_policy,
@ -76,6 +79,7 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( # noq
check_permissions_caller_permission,
generate_key_helper_fn,
prepare_metadata_fields,
raise_on_logging_credential_change,
)
from litellm.proxy.management_helpers.bulk_user_deletion import delete_users_from_current_teams
from litellm.proxy.management_helpers.object_permission_utils import ( # noqa: F401 # legacy module exports
@ -594,6 +598,8 @@ async def new_user(
if data.auto_create_key and isinstance(user_api_key_dict, UserAPIKeyAuth):
enforce_batch_limits_are_admin_only(data, None, user_api_key_dict, "key")
enforce_vector_store_index_grants_are_admin_only(data, None, None, user_api_key_dict, "key")
raise_on_logging_credential_change(data.metadata, None, user_api_key_dict)
check_permissions_caller_permission(
data=data,
@ -1479,8 +1485,11 @@ async def _invalidate_user_spend_counter_if_changed(
"""
if non_default_values.get("spend") is not None:
from litellm.proxy.proxy_server import invalidate_spend_counter
from litellm.proxy.spend_tracking.committed_spend import committed_spend_key
await invalidate_spend_counter(counter_key=user_counter(str(non_default_values["user_id"])))
counter_key: Final = user_counter(str(non_default_values["user_id"]))
await invalidate_spend_counter(counter_key=counter_key)
await invalidate_spend_counter(counter_key=committed_spend_key(counter_key))
def _clears_object_permission(user_request: UpdateUserRequest) -> bool:

77
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

@ -26,7 +26,7 @@ from typing import TYPE_CHECKING, Final, Literal, Optional, Protocol, TypeVar, c
import fastapi
import yaml
from fastapi import APIRouter, Depends, Header, HTTPException, Query, Request, status
from pydantic import TypeAdapter
from pydantic import ConfigDict, TypeAdapter
from typing_extensions import ReadOnly, TypedDict
import litellm
@ -62,14 +62,20 @@ from litellm.proxy.auth.auth_utils import (
abbreviate_api_key,
enforce_batch_limits_are_admin_only,
enforce_output_token_estimates_are_admin_only,
enforce_vector_store_index_grants_are_admin_only,
)
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import (
evict_and_broadcast,
)
from litellm.proxy.common_utils.callback_config_validation import logging_metadata_config_error
from litellm.proxy.common_utils.callback_config_validation import (
credential_reference_rejection,
logging_metadata_config_error,
)
from litellm.proxy.common_utils.callback_utils import (
decrypt_callback_vars,
decrypted_callback_metadata,
decrypted_logging_entries,
encrypt_callback_vars,
)
from litellm.proxy.common_utils.config_sync_pubsub import (
@ -157,6 +163,7 @@ from litellm.proxy.mcp_server.outbound_credentials.sso_assertion_store import (
)
from litellm.proxy.search_endpoints.search_tool_registry import rotate_search_tools_master_key
from litellm.proxy.spend_tracking.budget_reservation import get_budget_window_start
from litellm.proxy.spend_tracking.committed_spend import committed_spend_key
from litellm.proxy.spend_tracking.spend_tracking_utils import ( # noqa: F401 # legacy module exports
_is_master_key, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
is_master_key,
@ -264,6 +271,12 @@ class _StoredKeyRouterSettings(LiteLLMBaseModel):
router_settings: Mapping[str, object] | None = None
class _StoredKeyMetadata(LiteLLMBaseModel):
model_config = ConfigDict(frozen=True)
metadata: object = None
class _KeyRowWhere(TypedDict):
token: ReadOnly[str]
@ -932,6 +945,23 @@ def raise_on_invalid_key_logging_config(metadata: Mapping[str, object] | None) -
raise HTTPException(status_code=400, detail={"error": error})
def raise_on_logging_credential_change(
requested_metadata: object, existing_key_row: LiteLLM_VerificationToken | None, user_api_key_dict: UserAPIKeyAuth
) -> None:
stored_metadata: Final = (
None
if existing_key_row is None
else _StoredKeyMetadata.model_validate(existing_key_row, from_attributes=True).metadata
)
rejection: Final = credential_reference_rejection(
decrypted_callback_metadata(requested_metadata),
decrypted_logging_entries(stored_metadata),
is_proxy_admin=user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN,
)
if rejection is not None:
raise HTTPException(status_code=rejection.status_code, detail={"error": rejection.message})
def common_key_access_checks(
user_api_key_dict: UserAPIKeyAuth,
data: GenerateKeyRequest | UpdateKeyRequest,
@ -1483,6 +1513,7 @@ async def _common_key_generation_helper(
validate_budget_duration(data.budget_duration)
raise_on_invalid_key_logging_config(data.metadata)
raise_on_logging_credential_change(data.metadata, None, user_api_key_dict)
if data.throttle_on_budget_exceeded is True and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value:
raise HTTPException(
@ -1509,6 +1540,13 @@ async def _common_key_generation_helper(
user_api_key_dict=user_api_key_dict,
entity="key",
)
enforce_vector_store_index_grants_are_admin_only(
data=data,
top_level_grants=data.allowed_vector_store_indexes,
existing_metadata=None,
user_api_key_dict=user_api_key_dict,
entity="key",
)
if data.metadata is not None and data.metadata.get("service_account_id") is not None and data.team_id is None:
await validate_team_id_used_in_service_account_request(
@ -2418,7 +2456,7 @@ async def generate_key_fn(
- prompts: Optional[List[str]] - List of allowed prompts for the key. If specified, the key will only be able to use these specific prompts.
- auto_rotate: Optional[bool] - Whether this key should be automatically rotated (regenerated)
- rotation_interval: Optional[str] - How often to auto-rotate this key (e.g., '30s', '30m', '30h', '30d'). Required if auto_rotate=True.
- allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Create index, using `/v1/indexes` endpoint.
- allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Only a proxy admin can set or change it. Create index, using `/v1/indexes` endpoint.
- router_settings: Optional[UpdateRouterConfig] - key-specific router settings. Example - {"model_group_retry_policy": {"gpt-4": {"RateLimitErrorRetries": 5}}}. IF null or {} then no router settings.
- access_group_ids: Optional[List[str]] - List of access group IDs to associate with the key. Access groups define which models a key can access. Example - ["access_group_1", "access_group_2"].
- budget_limits: Optional[list] - List of concurrent budget windows for the key. Each window specifies a budget_limit, time_period, and optional budget_duration. Example - [{"budget_limit": 10.0, "time_period": "1d"}, {"budget_limit": 50.0, "time_period": "7d"}].
@ -2620,7 +2658,7 @@ async def generate_service_account_key_fn(
- allowed_routes: Optional[list] - List of allowed routes for the key. Store the actual route or store a wildcard pattern for a set of routes. Example - ["/chat/completions", "/embeddings", "/keys/*"]
- object_permission: Optional[LiteLLM_ObjectPermissionBase] - key-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"], "agents": ["agent_1", "agent_2"], "agent_access_groups": ["dev_group"]}. IF null or {} then no object permission.
Examples:
- allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Create index, using `/v1/indexes` endpoint.
- allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Only a proxy admin can set or change it. Create index, using `/v1/indexes` endpoint.
1. Allow users to turn on/off pii masking
@ -3339,6 +3377,13 @@ async def _process_single_key_update(
user_api_key_dict=user_api_key_dict,
entity="key",
)
enforce_vector_store_index_grants_are_admin_only(
data=update_key_request,
top_level_grants=update_key_request.allowed_vector_store_indexes,
existing_metadata=existing_key_row.metadata, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # LiteLLM_VerificationToken.metadata is a bare dict
user_api_key_dict=user_api_key_dict,
entity="key",
)
await _check_single_key_update_team_permissions(
user_api_key_dict=user_api_key_dict,
@ -3414,6 +3459,7 @@ async def _process_single_key_update(
)
# Prepare update data
raise_on_logging_credential_change(key_request.metadata, existing_key_row, user_api_key_dict)
non_default_values = await prepare_key_update_data(
data=key_request, existing_key_row=existing_key_row, prisma_client=prisma_client, llm_router=llm_router
)
@ -3847,6 +3893,13 @@ async def _validate_update_key_data(
user_api_key_dict=user_api_key_dict,
entity="key",
)
enforce_vector_store_index_grants_are_admin_only(
data=data,
top_level_grants=data.allowed_vector_store_indexes,
existing_metadata=_existing_metadata if isinstance(_existing_metadata, dict) else None, # pyright: ignore[reportUnknownArgumentType] # stored metadata is a bare dict
user_api_key_dict=user_api_key_dict,
entity="key",
)
# Personal-key bypass: the caller both created the key AND still owns it
# (user_id == caller). Checking only created_by would let a demoted admin
@ -4090,7 +4143,7 @@ async def update_key_fn(
- object_permission: Optional[LiteLLM_ObjectPermissionBase] - key-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"], "agents": ["agent_1", "agent_2"], "agent_access_groups": ["dev_group"]}. IF null or {} then no object permission.
- auto_rotate: Optional[bool] - Whether this key should be automatically rotated
- rotation_interval: Optional[str] - How often to rotate this key (e.g., '30d', '90d'). Required if auto_rotate=True
- allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Create index, using `/v1/indexes` endpoint.
- allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Only a proxy admin can set or change it. Create index, using `/v1/indexes` endpoint.
- router_settings: Optional[UpdateRouterConfig] - key-specific router settings. Example - {"model_group_retry_policy": {"gpt-4": {"RateLimitErrorRetries": 5}}}. IF null or {} then no router settings.
- access_group_ids: Optional[List[str]] - List of access group IDs to associate with the key. Access groups define which models a key can access. Example - ["access_group_1", "access_group_2"].
- budget_limits: Optional[list] - List of concurrent budget windows for the key. Each window specifies a budget_limit, time_period, and optional budget_duration. Example - [{"budget_limit": 10.0, "time_period": "1d"}, {"budget_limit": 50.0, "time_period": "7d"}].
@ -4153,6 +4206,7 @@ async def update_key_fn(
# Enforce upperbound key params on update (don't fill defaults)
_enforce_upperbound_key_params(data, fill_defaults=False)
raise_on_logging_credential_change(data.metadata, existing_key_row, user_api_key_dict)
non_default_values: Final = await prepare_key_update_data(
data=data, existing_key_row=existing_key_row, prisma_client=prisma_client, llm_router=llm_router
)
@ -4266,7 +4320,7 @@ async def update_key_fn(
)
if data.spend is not None:
from litellm.proxy.proxy_server import spend_counter_store
from litellm.proxy.proxy_server import invalidate_spend_counter, spend_counter_store
counter_key: Final = key_counter(hash_token_if_needed(key))
try:
@ -4278,6 +4332,7 @@ async def update_key_fn(
counter_key,
redis_err,
)
await invalidate_spend_counter(counter_key=committed_spend_key(counter_key))
asyncio.create_task(
KeyManagementEventHooks.async_key_updated_hook(
@ -6312,6 +6367,13 @@ async def _execute_virtual_key_regeneration(
user_api_key_dict=user_api_key_dict,
entity="key",
)
enforce_vector_store_index_grants_are_admin_only(
data=data,
top_level_grants=data.allowed_vector_store_indexes,
existing_metadata=_existing_key_metadata if isinstance(_existing_key_metadata, dict) else None, # pyright: ignore[reportUnknownArgumentType] # stored metadata is a bare dict
user_api_key_dict=user_api_key_dict,
entity="key",
)
await _validate_end_user_budget_id_change(
requested_budget_id=_requested_end_user_budget_id(data),
existing_budget_id=get_key_end_user_budget_id(
@ -6332,6 +6394,7 @@ async def _execute_virtual_key_regeneration(
await _enforce_custom_key_update_policy(hook=_custom_key_update_hook(proxy_server), data=update_request)
# Enforce upperbound key params on regenerate (don't fill defaults)
_enforce_upperbound_key_params(data, fill_defaults=False)
raise_on_logging_credential_change(data.metadata, key_in_db, user_api_key_dict)
non_default_values = await prepare_key_update_data(
data=data, existing_key_row=key_in_db, prisma_client=prisma_client, llm_router=llm_router
)
@ -6938,6 +7001,7 @@ async def reset_key_spend_fn(
try:
from litellm.proxy.proxy_server import (
hash_token,
invalidate_spend_counter,
prisma_client,
proxy_logging_obj,
set_spend_counter_with_floor_and_broadcast,
@ -6993,6 +7057,7 @@ async def reset_key_spend_fn(
# pre-reset value.
_counter_key: Final = key_counter(hashed_api_key)
await set_spend_counter_with_floor_and_broadcast(counter_key=_counter_key, value=reset_to)
await invalidate_spend_counter(counter_key=committed_spend_key(_counter_key))
await _reset_key_budget_windows(
prisma_client=prisma_client,
hashed_api_key=hashed_api_key,

22
litellm/proxy/management_endpoints/organization_endpoints.py
View file

@ -66,7 +66,9 @@ from litellm.proxy.management_helpers.utils import (
get_new_internal_user_defaults,
management_endpoint_wrapper,
)
from litellm.proxy.spend_tracking.committed_spend import drop_committed_spend
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.redis.stores.spend_counters import org_counter, team_counter
from litellm.repositories.budget_repository import BudgetRepository
from litellm.repositories.chunked_in import find_many_in
from litellm.repositories.object_permission_repository import ObjectPermissionRepository
@ -93,6 +95,7 @@ if TYPE_CHECKING:
LiteLLM_OrganizationMembership as PrismaOrganizationMembership,
)
from prisma.models import LiteLLM_OrganizationTable as PrismaOrganizationTable
from prisma.models import LiteLLM_TeamTable as PrismaTeamTable
from prisma.models import LiteLLM_UserTable as PrismaUserTable
from prisma.models import LiteLLM_VerificationToken as PrismaVerificationToken
@ -180,6 +183,8 @@ class _OrganizationMembershipTableClient(Protocol):
class _TeamTableClient(Protocol):
async def find_many(self, *, where: Mapping[str, object]) -> "Sequence[PrismaTeamTable]": ...
async def delete_many(self, where: Mapping[str, object]) -> int: ...
@ -819,6 +824,11 @@ async def update_organization(
include={"members": True, "teams": True, "litellm_budget_table": True},
)
if data.spend is not None:
from litellm.proxy.proxy_server import spend_counter_store
await drop_committed_spend(spend_counter_store, (org_counter(data.organization_id),))
return response
@ -1045,7 +1055,8 @@ async def delete_organization(
teams_in_deleted_orgs: Final = await find_many_in(
_table(TeamRepository(prisma_client)), "organization_id", data.organization_ids
)
raise_if_teams_hold_ptu_shares(tuple(team.team_id for team in teams_in_deleted_orgs), llm_router)
team_ids_in_deleted_orgs: Final = tuple(team.team_id for team in teams_in_deleted_orgs)
raise_if_teams_hold_ptu_shares(team_ids_in_deleted_orgs, llm_router)
deleted_orgs: Final = []
for organization_id in data.organization_ids:
@ -1073,6 +1084,15 @@ async def delete_organization(
)
deleted_orgs.append(deleted_org)
from litellm.proxy.proxy_server import spend_counter_store
await drop_committed_spend(
spend_counter_store,
(
*(org_counter(organization_id) for organization_id in data.organization_ids),
*(team_counter(team_id) for team_id in team_ids_in_deleted_orgs),
),
)
return deleted_orgs

26
litellm/proxy/management_endpoints/team_endpoints.py
View file

@ -122,6 +122,7 @@ from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module expo
from litellm.proxy.auth.auth_utils import (
enforce_batch_limits_are_admin_only,
enforce_output_token_estimates_are_admin_only,
enforce_vector_store_index_grants_are_admin_only,
)
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import (
@ -225,12 +226,14 @@ from litellm.proxy.management_helpers.utils import (
add_new_member,
management_endpoint_wrapper,
)
from litellm.proxy.spend_tracking.committed_spend import drop_committed_spend
from litellm.proxy.utils import (
PrismaClient,
ProxyLogging,
handle_exception_on_proxy,
query_raw_within_statement_timeout,
)
from litellm.redis.stores.spend_counters import team_counter
from litellm.repositories.budget_repository import BudgetRepository
from litellm.repositories.chunked_in import delete_many_in, find_many_in
from litellm.repositories.daily_activity_buckets import api_key_from_bucket
@ -1674,7 +1677,7 @@ async def new_team(
- team_member_key_duration: Optional[str] - The duration for a team member's key. e.g. "1d", "1w", "1mo"
- allowed_passthrough_routes: Optional[List[str]] - List of allowed pass through routes for the team.
- denied_passthrough_routes: Optional[List[str]] - List of pass through routes the team's keys may not call, even if allowed. Applies together with each key's `denied_passthrough_routes`.
- allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Create index, using `/v1/indexes` endpoint.
- allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Only a proxy admin can set or change it. Create index, using `/v1/indexes` endpoint.
- secret_manager_settings: Optional[dict] - Secret manager settings for the team. [Docs](https://docs.litellm.ai/docs/secret_managers/overview)
- router_settings: Optional[UpdateRouterConfig] - team-specific router settings. Example - {"model_group_retry_policy": {"gpt-4": {"RateLimitErrorRetries": 5}}}. IF null or {} then no router settings.
- access_group_ids: Optional[List[str]] - List of access group IDs to associate with the team. Access groups define which models the team can access. Example - ["access_group_1", "access_group_2"].
@ -1785,6 +1788,13 @@ async def new_team(
user_api_key_dict=user_api_key_dict,
entity="team",
)
enforce_vector_store_index_grants_are_admin_only(
data=data,
top_level_grants=data.allowed_vector_store_indexes,
existing_metadata=None,
user_api_key_dict=user_api_key_dict,
entity="team",
)
# Check if license is over limit
total_teams: Final = await _team_db(prisma_client).count()
@ -2459,7 +2469,7 @@ async def update_team(
- default_estimated_output_tokens_per_model: Optional[Dict[str, int]] - Per-model override of the above. Example: {"gpt-4": 4096, "gpt-3.5-turbo": 1024}
- mcp_rpm_limit: Optional[Dict[str, int]] - Per-MCP-server RPM limit for this team, keyed by MCP server name (alias if set, else the configured name). Example: {"github": 100, "slack": 200}. Applied across all keys for this team.
Example - update team TPM Limit
- allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Create index, using `/v1/indexes` endpoint.
- allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Only a proxy admin can set or change it. Create index, using `/v1/indexes` endpoint.
- secret_manager_settings: Optional[dict] - Secret manager settings for the team. [Docs](https://docs.litellm.ai/docs/secret_managers/overview)
- router_settings: Optional[UpdateRouterConfig] - team-specific router settings. Example - {"model_group_retry_policy": {"gpt-4": {"RateLimitErrorRetries": 5}}}. IF null or {} then no router settings.
- access_group_ids: Optional[List[str]] - List of access group IDs to associate with the team. Access groups define which models the team can access. Example - ["access_group_1", "access_group_2"].
@ -2597,6 +2607,13 @@ async def update_team(
user_api_key_dict=user_api_key_dict,
entity="team",
)
enforce_vector_store_index_grants_are_admin_only(
data=data,
top_level_grants=data.allowed_vector_store_indexes,
existing_metadata=_existing_team_metadata if isinstance(_existing_team_metadata, dict) else None, # pyright: ignore[reportUnknownArgumentType] # stored metadata is a bare dict
user_api_key_dict=user_api_key_dict,
entity="team",
)
check_passthrough_routes_caller_permission(
data,
@ -4848,6 +4865,7 @@ async def delete_validated_team_rows(
"""
from litellm.proxy.proxy_server import (
proxy_logging_obj,
spend_counter_store,
user_api_key_cache,
)
@ -4919,7 +4937,9 @@ async def delete_validated_team_rows(
_emit_team_members_metric(deleted_team.model_copy(update={"members_with_roles": ()}))
await sync_team_access_group_membership(prisma_client=prisma_client, team_id=deleted_team.team_id)
return DeletedTeams(team_ids=tuple(team.team_id for team in locked))
deleted_team_ids: Final = tuple(team.team_id for team in locked)
await drop_committed_spend(spend_counter_store, tuple(team_counter(team_id) for team_id in deleted_team_ids))
return DeletedTeams(team_ids=deleted_team_ids)
def _locked_roster_user_ids(locked: Sequence[LiteLLM_TeamTable]) -> tuple[str, ...]:

16
litellm/proxy/management_helpers/bulk_user_creation.py
View file

@ -29,8 +29,12 @@ from litellm.proxy._types import (
UserAPIKeyAuth,
)
from litellm.proxy.auth.auth_checks import invalidate_team_member_spend_state
from litellm.proxy.auth.auth_utils import enforce_batch_limits_are_admin_only
from litellm.proxy.auth.auth_utils import (
enforce_batch_limits_are_admin_only,
enforce_vector_store_index_grants_are_admin_only,
)
from litellm.proxy.auth.litellm_license import LicenseCheck
from litellm.proxy.common_utils.callback_config_validation import credential_reference_rejection
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
from litellm.proxy.hooks.user_management_event_hooks import UserManagementEventHooks
@ -219,9 +223,17 @@ def _row_error(item: BulkNewUserItem, user_api_key_dict: UserAPIKeyAuth) -> str
check_permissions_caller_permission(data=item, user_api_key_dict=user_api_key_dict)
if item.auto_create_key:
enforce_batch_limits_are_admin_only(item, None, user_api_key_dict, "key")
enforce_vector_store_index_grants_are_admin_only(item, None, None, user_api_key_dict, "key")
except Exception as exc: # noqa: BLE001 # any validation failure is reported on this row only
return _error_message(exc)
return None
rejection: Final = (
credential_reference_rejection(
item.metadata, None, is_proxy_admin=user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
)
if item.auto_create_key
else None
)
return None if rejection is None else rejection.message
def _normalized_email(email: str | None) -> str | None:

5
litellm/proxy/management_helpers/bulk_user_deletion.py
View file

@ -46,7 +46,9 @@ from litellm.proxy.management_helpers.access_group_team_sync import (
TEAM_ADVISORY_LOCK_SQL,
USER_ROWS_LOCK_SQL,
)
from litellm.proxy.spend_tracking.committed_spend import drop_committed_spend
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.redis.stores.spend_counters import user_counter
from litellm.repositories.table_repositories import (
OrganizationMembershipRepository,
TeamMembershipRepository,
@ -585,6 +587,9 @@ async def _after_users_deleted(
)
await evict_and_broadcast(cache_keys=deletion.jwt_mapping_cache_keys, user_api_key_cache=user_api_key_cache)
await evict_and_broadcast(cache_keys=sorted(user_ids), user_api_key_cache=user_api_key_cache)
from litellm.proxy.proxy_server import spend_counter_store
await drop_committed_spend(spend_counter_store, tuple(user_counter(user_id) for user_id in sorted(user_ids)))
for removal in deletion.removals.values():
_emit_team_members_metric(removal.team)

10
litellm/proxy/middleware/billable_request_metrics_middleware.py
View file

@ -24,6 +24,7 @@ from typing import Final, Protocol, runtime_checkable
from starlette.types import ASGIApp, Message, Receive, Scope, Send
from litellm._logging import verbose_proxy_logger
from litellm.constants import OPENSEARCH_PASS_THROUGH_ROUTE_PREFIX
from litellm.proxy._types import LiteLLMRoutes
from litellm.proxy.auth.route_checks import RouteChecks
@ -101,16 +102,21 @@ _LLM_ROUTE_EXACT: Final[tuple[str, ...]] = (
# inference calls that write SpendLogs rows, so they bill. Anchored to the
# routes enum so new providers are picked up without touching this module.
# /langfuse forwards observability traffic, not inference: it writes no
# SpendLogs row and must not bill.
_NON_BILLABLE_PASSTHROUGH_PREFIXES: Final = frozenset({"/langfuse"})
# SpendLogs row and must not bill. /opensearch bills only index searches, like
# /v1/vector_stores/{id}/search, and never its document writes or lookups.
_NON_BILLABLE_PASSTHROUGH_PREFIXES: Final = frozenset({"/langfuse", OPENSEARCH_PASS_THROUGH_ROUTE_PREFIX})
_PASSTHROUGH_PREFIXES: Final[tuple[str, ...]] = tuple(
prefix
for prefix in LiteLLMRoutes.mapped_pass_through_routes.value
if prefix not in _NON_BILLABLE_PASSTHROUGH_PREFIXES
)
_OPENSEARCH_INDEX_SEARCH: Final = re.compile(rf"{re.escape(OPENSEARCH_PASS_THROUGH_ROUTE_PREFIX)}/[^/]+/_search")
def _classify_llm_route(path: str) -> str | None:
if path.startswith(f"{OPENSEARCH_PASS_THROUGH_ROUTE_PREFIX}/"):
return OPENSEARCH_PASS_THROUGH_ROUTE_PREFIX if _OPENSEARCH_INDEX_SEARCH.fullmatch(path) else None
exact_match: Final = next((route for route in _LLM_ROUTE_EXACT if path == route), None)
if exact_match is not None:
return exact_match

133
litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
View file

@ -159,11 +159,15 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
)
from litellm.proxy.utils import ProxyLogging as ProxyLoggingType
from litellm.proxy.utils import is_known_model, normalize_route_for_root_path
from litellm.proxy.vector_store_endpoints.endpoints import build_request_data_from_managed_vector_store
from litellm.proxy.vector_store_endpoints.utils import (
assert_proxy_admin_for_vector_store_index_management,
assert_user_can_access_vector_store,
assert_vector_store_index_access,
get_litellm_managed_vector_store,
has_vector_store_index_access,
is_allowed_to_call_vector_store_endpoint,
normalize_litellm_params,
)
from litellm.secret_managers.main import get_secret_str, normalize_nonempty_secret_str, str_to_bool
from litellm.types.passthrough_endpoints.eu_jev import EuJevDecisionRequest, EuJevDecisionResponse
@ -182,8 +186,8 @@ from litellm.types.passthrough_endpoints.tinyfish import (
from litellm.types.passthrough_endpoints.vertex_ai import VertexPassThroughCredentials
from litellm.types.router import LiteLLMParamsTypedDict
from litellm.types.utils import CallTypes, LlmProviders
from litellm.types.vector_stores import LiteLLM_ManagedVectorStore
from litellm.utils import ProviderConfigManager
from litellm.types.vector_stores import LiteLLM_ManagedVectorStore, LiteLLM_ManagedVectorStoreIndex
from litellm.utils import ProviderConfigManager, with_named_credential
from .passthrough_endpoint_router import PassthroughEndpointRouter
@ -1056,6 +1060,131 @@ async def milvus_proxy_route(
return received_value
def _registered_vector_store_index(index_name: str) -> LiteLLM_ManagedVectorStoreIndex:
registry: Final = litellm.vector_store_index_registry
index_object: Final = registry.get_vector_store_index_by_name(index_name) if registry is not None else None
if index_object is None:
raise HTTPException(
status_code=400,
detail=f"Index {index_name} is not a litellm managed vector store index. "
"Only litellm managed vector store indexes are supported.",
)
return index_object
async def _opensearch_store_params(
index_object: LiteLLM_ManagedVectorStoreIndex, user_api_key_dict: UserAPIKeyAuth
) -> Mapping[str, object]:
registry: Final = litellm.vector_store_registry
store_name: Final = index_object.litellm_params.vector_store_name
named_stores: Final = tuple(
store
for store in (registry.vector_stores if registry is not None else ())
if store.get("vector_store_name") == store_name
)
if len(named_stores) > 1:
raise HTTPException(
status_code=400,
detail=f"{len(named_stores)} vector stores are named {store_name}, so LiteLLM cannot tell which one "
f"index {index_object.index_name} is on. Give each vector store a unique name.",
)
if not named_stores:
raise HTTPException(
status_code=400,
detail=f"Vector store {store_name} of index {index_object.index_name} is not registered.",
)
vector_store: Final = named_stores[0]
params: Final = with_named_credential(
build_request_data_from_managed_vector_store(normalize_litellm_params(vector_store))
)
provider: Final = params.get("custom_llm_provider")
if not isinstance(provider, str) or provider.split("/", 1)[0] != LlmProviders.OPENSEARCH.value:
raise HTTPException(
status_code=400,
detail=f"Index {index_object.index_name} is on vector store {store_name}, "
"which is not an OpenSearch vector store.",
)
await assert_user_can_access_vector_store(vector_store=vector_store, user_api_key_dict=user_api_key_dict)
return params
@router.api_route(
"/opensearch/{endpoint:path}",
methods=["GET", "HEAD", "POST", "PUT", "DELETE", "PATCH"],
tags=["OpenSearch Pass-through", "pass-through"],
openapi_extra={
"requestBody": {
"required": False,
"content": {
"application/json": {"schema": {"type": "object", "additionalProperties": True}},
"application/x-ndjson": {"schema": {"type": "string"}},
},
}
},
responses={
200: {
"description": "OpenSearch's response, unmodified",
"content": {"application/json": {"schema": {"type": "object", "additionalProperties": True}}},
}
},
)
async def opensearch_proxy_route( # noqa: ANN201 # returns create_pass_through_route's untyped handler result
endpoint: str,
request: Request,
fastapi_response: Response,
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
):
"""
Call an OpenSearch index through the proxy, e.g. `POST /opensearch/<index>/_search`.
`<index>` is a LiteLLM vector store index on an OpenSearch vector store. The key or its team needs `read` or
`write` on it in `allowed_vector_store_indexes`, and LiteLLM signs the call with the vector store's AWS identity.
"""
from litellm.llms.opensearch.vector_stores.index_passthrough import (
IndexRefusal,
index_operation,
parse_index_path,
prepare_index_request,
)
from litellm.llms.opensearch.vector_stores.transformation import OpenSearchIndexAWS, OpenSearchVectorStoreConfig
path: Final = parse_index_path(endpoint)
if isinstance(path, IndexRefusal):
raise HTTPException(status_code=path.status_code, detail=path.message)
operation: Final = index_operation(request.method, path.api)
assert_vector_store_index_access(
user_api_key_dict, path.index_name, operation.permission if operation is not None else None
)
index_object: Final = _registered_vector_store_index(path.index_name)
store_params: Final = await _opensearch_store_params(index_object, user_api_key_dict)
prepared: Final = await prepare_index_request(
config=OpenSearchVectorStoreConfig(aws=OpenSearchIndexAWS()),
method=request.method,
path=path,
operation=operation,
provider_index=index_object.litellm_params.vector_store_index,
content_type=request.headers.get("content-type"),
query=tuple(request.query_params.multi_items()),
body=await request.body(),
store_params=store_params,
can_read=has_vector_store_index_access(user_api_key_dict, path.index_name, "read"),
)
if isinstance(prepared, IndexRefusal):
raise HTTPException(status_code=prepared.status_code, detail=prepared.message)
endpoint_func: Final = create_pass_through_route(
endpoint=endpoint,
target=prepared.url,
custom_headers=prepared.headers,
custom_llm_provider=LlmProviders.OPENSEARCH.value,
keep_target_query=True,
upstream_headers_filter=lambda _forwarded: prepared.headers,
)
setattr(request.state, LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY, prepared.body)
setattr(request.state, LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY, {})
safe_set_request_parsed_body(request, {})
return await endpoint_func(request, fastapi_response, user_api_key_dict)
async def is_streaming_request_fn(request: Request) -> bool:
if request.method == "POST":
content_type: Final = request.headers.get("content-type", None)

36
litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
View file

@ -119,6 +119,7 @@ from litellm.proxy.common_utils.error_body_call_id import JSON_OBJECT, error_bod
from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # legacy module exports
_read_request_body, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
_safe_get_request_headers, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
is_opaque_pass_through_request,
read_request_body,
safe_get_request_headers,
)
@ -164,6 +165,8 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import (
WebSocketSessionUsage,
WebsocketUpstreamFrameKind,
)
from litellm.types.proxy.auth.user_api_key_auth import COMMITTED_BUDGET_SPEND_PARAM
from litellm.types.proxy.carried_budget_state import CommittedBudgetSpend
from litellm.types.utils import TRUSTED_CALLBACK_VARS_FIELD, CallTypes, Usage
from .llm_provider_handlers.tinyfish_passthrough_logging_handler import (
@ -423,6 +426,10 @@ async def chat_completion_pass_through_endpoint(
)
def _is_provider_payload(request: Request) -> bool:
return is_opaque_pass_through_request(get_request_route(request), request.headers.get("content-type", ""))
class HttpPassThroughEndpointHelpers(BasePassthroughUtils):
@staticmethod
def get_response_headers(
@ -774,6 +781,7 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils):
"litellm_params": {
**litellm_params_in_body,
"metadata": _metadata,
COMMITTED_BUDGET_SPEND_PARAM: user_api_key_dict.committed_budget_spend or CommittedBudgetSpend(),
"proxy_server_request": {
"url": str(request.url),
"method": request.method,
@ -956,11 +964,16 @@ def _build_passthrough_failure_request_payload(
model and provider. A user-defined route never gets that, since its URL
names an upstream model rather than a LiteLLM-managed one.
"""
request_payload: Final[dict] = dict(cast(Mapping[str, object], parsed_body or {})) # cast-ok: json body
request_payload: Final = dict(cast(Mapping[str, object], parsed_body or {})) # cast-ok: json body
if kwargs:
request_payload.update(kwargs)
request_payload.update(_json_object(_DECODED_JSON.validate_python(kwargs)))
if logging_obj is not None:
request_payload["litellm_logging_obj"] = logging_obj
logging_litellm_params: Final = _json_object(logging_obj.model_call_details.get("litellm_params"))
existing_litellm_params: Final = _json_object(request_payload.get("litellm_params"))
api_base: Final = logging_litellm_params.get("api_base")
if isinstance(api_base, str) and api_base and "api_base" not in existing_litellm_params:
request_payload["litellm_params"] = {**existing_litellm_params, "api_base": api_base}
url_route: Final = str(url) if url is not None else ""
google_provider: Final = (
_passthrough_google_provider(url_route, custom_llm_provider) if builtin_provider_route else None
@ -1209,6 +1222,14 @@ def _unfiltered_headers(headers: Mapping[str, str]) -> Mapping[str, str]:
return headers
def _caller_query_source(
keep_target_query: bool, query_params: Mapping[str, object] | None, request: Request
) -> Mapping[str, object]:
if keep_target_query:
return MappingProxyType({})
return query_params or request.query_params
def _json_object(value: object) -> Mapping[str, object]:
return cast("Mapping[str, object]", value) if isinstance(value, Mapping) else MappingProxyType({}) # cast-ok: json
@ -1266,6 +1287,7 @@ async def pass_through_request(
timeout: float | None = None,
upstream_headers_filter: Callable[[Mapping[str, str]], Mapping[str, str]] = _unfiltered_headers,
operation: GenAIOperation | None = None,
keep_target_query: bool = False,
):
"""
Pass through endpoint handler, makes the httpx request for pass-through endpoints and ensures logging hooks are called
@ -1288,6 +1310,7 @@ async def pass_through_request(
general_settings.pass_through_request_timeout, then 600s.
upstream_headers_filter: Optional last rewrite of the headers sent upstream, after x-pass- expansion
operation: Optional GenAI operation the route reports, unless the upstream names its own
keep_target_query: Send the target URL's query exactly as given, adding none of the caller's query params
"""
from litellm.exceptions import ModifyResponseException
from litellm.litellm_core_utils.litellm_logging import Logging
@ -1326,7 +1349,10 @@ async def pass_through_request(
parent_span=user_api_key_dict.parent_otel_span,
)
requested_query_params: dict | None = query_params or dict(request.query_params) or None
requested_query_params: dict | None = (
dict(_caller_query_source(keep_target_query, query_params, request)) # pyright: ignore[reportUnknownArgumentType] # query_params keeps its legacy untyped dict annotation
or None
)
endpoint_type: Final[EndpointType] = HttpPassThroughEndpointHelpers.get_endpoint_type(
str(url), custom_llm_provider
@ -2263,6 +2289,8 @@ async def _parse_request_data_by_content_type(
Returns:
Tuple of (query_params_data, custom_body_data, file_data, stream)
"""
if _is_provider_payload(request):
return None, None, None, None
content_type: Final = request.headers.get("content-type", "")
query_params_data = None
@ -2333,6 +2361,7 @@ def create_pass_through_route(
timeout: float | None = None,
upstream_headers_filter: Callable[[Mapping[str, str]], Mapping[str, str]] = _unfiltered_headers,
operation: GenAIOperation | None = None,
keep_target_query: bool = False,
):
# check if target is an adapter.py or a url
from litellm._uuid import uuid
@ -2482,6 +2511,7 @@ def create_pass_through_route(
timeout=cast(float | None, param_timeout),
upstream_headers_filter=upstream_headers_filter,
operation=operation,
keep_target_query=keep_target_query,
)
finally:
if hasattr(request.state, LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY):

30
litellm/proxy/plugin_routes.py
View file

@ -18,6 +18,10 @@ Plugin iframe auth:
litellm bearer token, so a compromised plugin learns only the caller's
identity, never their credential. LITELLM_SALT_KEY itself is never shared
with plugins — each plugin holds only its own derived key.
With general_settings.fips_mode on, the per-plugin key is instead
derive_key_from_salt_key(b"litellm-plugin:session-claim:", plugin_name), and
the Fernet token is built from the FIPS provider's DRBG, AES-CBC and HMAC.
"""
import base64
@ -25,11 +29,14 @@ import hashlib
import hmac as _hmac
import json
import os
import ssl
import time
from collections.abc import Mapping
from typing import Final
from cryptography.fernet import Fernet, InvalidToken
from cryptography.hazmat.primitives import padding
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from fastapi import APIRouter, Depends, HTTPException, Request, Response
from pydantic import ConfigDict, TypeAdapter, with_config
from typing_extensions import NotRequired, ReadOnly, TypedDict
@ -37,6 +44,8 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
from litellm.proxy._types import PluginConfig, SpecialHeaders, UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.encrypt_decrypt_utils import derive_key_from_salt_key
from litellm.proxy.common_utils.fips import fips_mode_active
from litellm.types.llms.custom_http import httpxSpecialProvider
router: Final = APIRouter()
@ -145,11 +154,30 @@ def _plugin_fernet(plugin_name: str) -> Fernet:
A plugin possessing its own key cannot derive the master salt or
forge claims intended for a different plugin.
"""
if fips_mode_active():
return Fernet(base64.urlsafe_b64encode(_fips_plugin_key(plugin_name)))
salt: Final = os.getenv("LITELLM_SALT_KEY", "").encode()
derived: Final = _hmac.new(salt, plugin_name.encode(), hashlib.sha256).digest()
return Fernet(base64.urlsafe_b64encode(derived))
_FIPS_PLUGIN_KEY_LABEL: Final = b"litellm-plugin:session-claim:"
def _fips_plugin_key(plugin_name: str) -> bytes:
return bytes.fromhex(derive_key_from_salt_key(_FIPS_PLUGIN_KEY_LABEL, plugin_name.encode()))
def _fips_fernet_token(key: bytes, data: bytes) -> str:
"""Encrypt data into a standard Fernet token using only the FIPS provider's DRBG, AES-128-CBC and HMAC-SHA256."""
iv: Final = ssl.RAND_bytes(16)
padder: Final = padding.PKCS7(128).padder()
encryptor: Final = Cipher(algorithms.AES(key[16:]), modes.CBC(iv)).encryptor()
ciphertext: Final = encryptor.update(padder.update(data) + padder.finalize()) + encryptor.finalize()
body: Final = b"\x80" + int(time.time()).to_bytes(8, "big") + iv + ciphertext
return base64.urlsafe_b64encode(body + _hmac.digest(key[:16], body, "sha256")).decode()
_CLAIM_TTL_SECONDS: Final = 30 # identity claims expire after 30 s
@ -166,6 +194,8 @@ def issue_plugin_session_claim(plugin_name: str, user_id: str | None, user_role:
"user_role": user_role or "",
"exp": int(time.time()) + _CLAIM_TTL_SECONDS,
}
if fips_mode_active():
return _fips_fernet_token(_fips_plugin_key(plugin_name), json.dumps(claim).encode())
return _plugin_fernet(plugin_name).encrypt(json.dumps(claim).encode()).decode()

62
litellm/proxy/proxy_server.py
View file

@ -93,6 +93,7 @@ from litellm.constants import (
LITELLM_UI_SESSION_DURATION,
OPENAI_LIVE_SESSION_START_TIMEOUT_SECONDS,
RUNTIME_UPDATABLE_ROUTER_SETTINGS,
SPEND_LOG_QUEUE_MAX_BACKOFF_SECONDS,
VLLM_CHOICE_FIELDS,
)
from litellm.litellm_core_utils.access_windows import access_windows_config_error
@ -338,7 +339,7 @@ from litellm._logging import (
verbose_router_logger,
)
from litellm.caching.caching import DualCache, RedisCache
from litellm.caching.redis_batch import drain_post_call_redis_batches
from litellm.caching.redis_batch import active_request_redis_batches, drain_post_call_redis_batches
from litellm.caching.redis_cache import RedisCircuitBreakerOpenError, is_redis_timeout_failure
from litellm.caching.redis_cluster_cache import RedisClusterCache
from litellm.constants import (
@ -995,12 +996,20 @@ from litellm.proxy.spend_tracking.budget_reservation import (
release_unbound_budget_reservation,
stamp_budget_reservation_actual_cost,
)
from litellm.proxy.spend_tracking.committed_spend import (
committed_spend_key,
committed_spend_keys,
has_committed_twin,
is_committed_spend_key,
seed_committed_spend,
)
from litellm.proxy.spend_tracking.spend_capture_rate import (
run_scheduled_spend_capture_rate_check,
)
from litellm.proxy.spend_tracking.spend_counter_batch import (
PendingSpendIncrement,
active_spend_counter_batch,
bind_spend_counter_keys,
forget_spend_counter,
post_call_counter_keys,
read_batched_spend_counter,
@ -1105,6 +1114,7 @@ from litellm.types.llms.openai import (
ResponseInProgressEvent,
ResponsesAPIResponse,
)
from litellm.types.proxy.carried_budget_state import EntitySpendCounterKeys
from litellm.types.proxy.control_plane_endpoints import WorkerRegistryEntry
from litellm.types.proxy.management_endpoints.model_management_endpoints import (
ModelGroupInfoProxy,
@ -3240,6 +3250,7 @@ proxy_budget_rescheduler_min_time = PROXY_BUDGET_RESCHEDULER_MIN_TIME
proxy_budget_rescheduler_max_time = PROXY_BUDGET_RESCHEDULER_MAX_TIME
proxy_batch_polling_interval = PROXY_BATCH_POLLING_INTERVAL
proxy_batch_write_at = PROXY_BATCH_WRITE_AT
_BATCH_WRITE_JITTER_MAX_SECONDS: Final = 5
proxy_config_reload_interval_seconds = PROXY_CONFIG_RELOAD_INTERVAL_SECONDS
litellm_master_key_hash = None
disable_spend_logs = False
@ -3378,6 +3389,10 @@ def _fail_closed_budget_enforcement() -> bool:
return general_settings.get("fail_closed_budget_enforcement") is True
def fail_closed_budget_enforcement_enabled() -> bool:
return _fail_closed_budget_enforcement()
def _raise_budget_unverifiable(counter_key: str) -> None:
verbose_proxy_logger.warning(
"fail_closed_budget_enforcement: rejecting request — spend for %s could "
@ -3638,6 +3653,7 @@ async def increment_spend_counters(
request_started_at: datetime | None = None,
model_access_groups: Sequence[str] | None = None,
project_id: str | None = None,
committed_spend_scope: EntitySpendCounterKeys | None = None,
windows_from_ledger: bool = False,
):
"""
@ -3675,6 +3691,7 @@ async def increment_spend_counters(
request_started_at=request_started_at,
model_access_groups=model_access_groups,
project_id=project_id,
committed_spend_scope=committed_spend_scope,
windows_from_ledger=windows_from_ledger,
)
@ -3790,10 +3807,17 @@ async def _increment_spend_counters_batched(
request_started_at: datetime | None,
model_access_groups: Sequence[str] | None,
project_id: str | None = None,
committed_spend_scope: EntitySpendCounterKeys | None = None,
windows_from_ledger: bool = False,
):
"""Runs inside one spend counter batch: the reservation reconcile and the warm checks share a single MGET, and
the reconcile adjustments go out in the same INCRBYFLOAT pipeline as the counter increments."""
committed_counter_keys: Final = (
tuple(counter_key for counter_key in committed_spend_scope.tracked() if has_committed_twin(counter_key))
if committed_spend_scope is not None and response_cost is not None and response_cost > 0
else ()
)
bind_spend_counter_keys(frozenset(committed_spend_keys(committed_counter_keys)))
reservation_update: Final = await _reconcile_budget_reservation_for_counter_update(
budget_reservation=budget_reservation,
response_cost=response_cost,
@ -3980,6 +4004,9 @@ async def _increment_spend_counters_batched(
)
if project_id is not None
else None,
_prepare_committed_spend_increments(counter_keys=committed_counter_keys, increment=cost)
if committed_counter_keys
else None,
)
if coro is not None
)
@ -4082,6 +4109,23 @@ async def _reconcile_budget_reservation_for_counter_update(
return _ReservationCounterUpdate(reserved_counter_keys=frozenset(reserved_counter_keys), pending=pending)
async def _committed_spend_seeded(counter_key: str) -> bool:
return await _is_spend_counter_cache_warm(committed_spend_key(counter_key)) or await seed_committed_spend(
spend_counter_store=spend_counter_store, prisma_client=prisma_client, counter_key=counter_key
)
async def _prepare_committed_spend_increments(
counter_keys: tuple[str, ...], increment: float
) -> tuple[PendingSpendIncrement, ...]:
seeded: Final = await asyncio.gather(*(_committed_spend_seeded(counter_key) for counter_key in counter_keys))
return tuple(
PendingSpendIncrement(counter_key=committed_spend_key(counter_key), increment=increment)
for counter_key, is_seeded in zip(counter_keys, seeded)
if is_seeded
)
async def _prepare_end_user_and_tag_spend_increments(
end_user_id: str | None,
tags: list[str] | None,
@ -4487,7 +4531,15 @@ def _defer_spend_counter_increments(pending: Sequence[PendingSpendIncrement]) ->
"""Post-call increments ride the request's native post-call pipeline with the other counters. Each counter's
new value lands in memory when the pipeline settles; a failed one is dropped everywhere so no reader trusts a
counter whose increment may not have applied, as ``increment_spend_counters_pipeline`` does."""
return spend_counter_store.increment_many_post_call(tuple((item.counter_key, item.increment) for item in pending))
settlement: Final = spend_counter_store.increment_many_post_call(
tuple((item.counter_key, item.increment) for item in pending)
)
if settlement is None:
return False
batches: Final = active_request_redis_batches()
if batches is not None and any(is_committed_spend_key(item.counter_key) for item in pending):
batches.track_post_call_settlement(settlement)
return True
async def increment_spend_counters_pipeline(pending: Sequence[PendingSpendIncrement]) -> tuple[float | None, ...]:
@ -7224,6 +7276,7 @@ class ProxyConfig:
"claude_code_marketplace_auth_required",
"vertex_ai_pass_through_drop_anthropic_betas",
"request_tag_mappings",
"access_group_rules",
)
if name in typed_general_settings
}
@ -11651,7 +11704,10 @@ class ProxyStartupEvent:
)
# Ensure minimum interval of 30 seconds for batch writing to prevent memory issues
batch_writing_interval: Final = proxy_batch_write_at + random.randint(0, 5)
batch_writing_interval: Final = proxy_batch_write_at + random.randint(0, _BATCH_WRITE_JITTER_MAX_SECONDS)
spend_counter_store.set_committed_ttl(
2 * (SPEND_LOG_QUEUE_MAX_BACKOFF_SECONDS + proxy_batch_write_at + _BATCH_WRITE_JITTER_MAX_SECONDS)
)
### PROXY WORKER HEARTBEAT ###
worker_heartbeat: Final = ProxyWorkerHeartbeat(prisma_client=prisma_client)

40
litellm/proxy/spend_tracking/carried_budget_state.py
View file

@ -1,14 +1,17 @@
"""Pins the budget state auth resolved onto ``UserAPIKeyAuth`` and emits it as request metadata."""
from collections.abc import Mapping
from collections.abc import Awaitable, Callable, Mapping
from types import MappingProxyType
from typing import Final
from litellm.models.end_user import LiteLLM_EndUserTable
from litellm.models.organization import LiteLLM_OrganizationTable
from litellm.models.team import LiteLLM_TeamTable
from litellm.models.user import LiteLLM_UserTable
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.proxy.carried_budget_state import (
CommittedBudgetSpend,
EntitySpendCounterKeys,
OrgBudgetSnapshot,
TeamBudgetSnapshot,
UserBudgetSnapshot,
@ -34,6 +37,41 @@ def carry_team_and_user_budget_state(
)
def _budgeted_org_id(valid_token: UserAPIKeyAuth, team_object: LiteLLM_TeamTable | None) -> str | None:
org_snapshot: Final = valid_token.org_budget_snapshot
if org_snapshot is None or org_snapshot.max_budget is None:
return None
return valid_token.org_id or (team_object.organization_id if team_object is not None else None)
def _budgeted_end_user_id(end_user_object: LiteLLM_EndUserTable | None) -> str | None:
budget_table: Final = end_user_object.litellm_budget_table if end_user_object is not None else None
if end_user_object is None or budget_table is None or budget_table.max_budget is None:
return None
return end_user_object.user_id
async def carry_committed_budget_spend(
valid_token: UserAPIKeyAuth,
team_object: LiteLLM_TeamTable | None,
user_object: LiteLLM_UserTable | None,
end_user_object: LiteLLM_EndUserTable | None,
read_committed: Callable[[EntitySpendCounterKeys], Awaitable[CommittedBudgetSpend]] | None,
) -> None:
counter_keys: Final = EntitySpendCounterKeys.of(
token=valid_token.token if valid_token.max_budget is not None else None,
team_id=team_object.team_id if team_object is not None and team_object.max_budget is not None else None,
user_id=user_object.user_id if user_object is not None and user_object.max_budget is not None else None,
org_id=_budgeted_org_id(valid_token=valid_token, team_object=team_object),
end_user_id=_budgeted_end_user_id(end_user_object=end_user_object),
)
valid_token.committed_budget_spend = ( # rebind-ok: request-scoped credential
CommittedBudgetSpend(counter_keys=counter_keys)
if read_committed is None
else await read_committed(counter_keys)
)
def carry_organization_budget_state(valid_token: UserAPIKeyAuth, org_table: LiteLLM_OrganizationTable) -> None:
budget_table: Final = org_table.litellm_budget_table
valid_token.organization_alias = org_table.organization_alias

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